Adding persistent workers for 3dvar_atmos suite

src/swell/deployment/create_experiment.py

@@ -591,6 +591,14 @@ def prepare_cylc_suite_jinja2(
     stall timeout = {{environ['SWELL_CYLC_TIMEOUT']}}
     {% endif %}"""
 
+    # Persistent worker mode
+    # -----------------------
+    persistent_workers = 'salloc_directives' in render_dictionary['scheduling']
+    render_dictionary['persistent_workers'] = persistent_workers
+    if persistent_workers:
+        render_dictionary['salloc_directives'] = \
+            render_dictionary['scheduling']['salloc_directives']
+
     # Render the template
     # -------------------
     new_suite_file = template_string_jinja2(logger, suite_file, render_dictionary,

src/swell/deployment/platforms/nccs_discover_cascade/slurm.yaml

@@ -3,3 +3,4 @@ nodes: 1
 ntasks-per-node: 24
 constraint: cas
 no-requeue: ''
+persistent_workers: true

src/swell/deployment/platforms/nccs_discover_sles15/slurm.yaml

@@ -3,3 +3,4 @@ nodes: 1
 ntasks-per-node: 64
 constraint: mil
 no-requeue: ''
+persistent_workers: true

src/swell/suites/3dvar_atmos/flow.cylc

@@ -14,6 +14,10 @@
 [scheduler]
     UTC mode = True
     allow implicit tasks = False
+    {% if persistent_workers %}
+    [[events]]
+        shutdown handlers = sh -c 'JOB_ID=$(cat "$CYLC_WORKFLOW_RUN_DIR/share/persistent_jobid" 2>/dev/null) && scancel "$JOB_ID" || true'
+    {% endif %}
 
 # --------------------------------------------------------------------------------------------------
 
@@ -27,8 +31,14 @@
         R1 = """
             # Triggers for non cycle time dependent tasks
             # -------------------------------------------
+            {% if persistent_workers %}
+            # Acquire a persistent Slurm allocation; all heavy tasks run inside it.
+            Allocate
+            Allocate => CloneJedi
+            {% else %}
             # Clone JEDI source code
             CloneJedi
+            {% endif %}
 
             # Build JEDI source code by linking
             CloneJedi => BuildJediByLinking?
@@ -40,6 +50,15 @@
             # Clone geos ana for generating observing system records
             CloneGeosMksi-{{model_component}}
             {% endfor %}
+
+            {% if persistent_workers %}
+            # Deallocate only after every cycle's CleanCycle has finished
+            {% for ct in cycle_times %}
+            {% for mc in model_components %}
+            CleanCycle-{{mc}}[{{ct.cycle_time}}] => Deallocate
+            {% endfor %}
+            {% endfor %}
+            {% endif %}
         """
 
         {% for cycle_time in cycle_times %}
@@ -124,13 +143,23 @@
         script = "swell task BuildJediByLinking $config"
 
     [[BuildJedi]]
+        {% if persistent_workers %}
+        platform = localhost
+        script = """
+            {% set _build_cpus = scheduling["BuildJedi"]["srun_directives"].get('ntasks-per-node', 1) -%}
+            srun --jobid=$(cat $CYLC_WORKFLOW_RUN_DIR/share/persistent_jobid) --ntasks=1 \
+            --cpus-per-task={{ _build_cpus }} \
+            swell task BuildJedi $config
+        """
+        {% else %}
         script = "swell task BuildJedi $config"
         platform = {{platform}}
         execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}}
         [[[directives]]]
         {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %}
             --{{key}} = {{value}}
         {%- endfor %}
+        {% endif %}
 
     {% for model_component in model_components %}
 
@@ -159,13 +188,25 @@
         script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}"
 
     [[RunJediVariationalExecutable-{{model_component}}]]
+        {% if persistent_workers %}
+        platform = localhost
+        script = """
+            export SWELL_PERSISTENT_JOB_ID=$(cat $CYLC_WORKFLOW_RUN_DIR/share/persistent_jobid)
+            {% set _srun = scheduling["RunJediVariationalExecutable"].get("srun_directives_by_model", {}).get(model_component, scheduling["RunJediVariationalExecutable"]["srun_directives"]) -%}
+            {% if 'nodes' in _srun -%}
+            export SWELL_SRUN_NODES={{ _srun['nodes'] }}
+            {% endif -%}
+            swell task RunJediVariationalExecutable $config -d $datetime -m {{model_component}}
+        """
+        {% else %}
         script = "swell task RunJediVariationalExecutable $config -d $datetime -m {{model_component}}"
         platform = {{platform}}
         execution time limit = {{scheduling["RunJediVariationalExecutable"]["execution_time_limit"]}}
         [[[directives]]]
         {%- for key, value in scheduling["RunJediVariationalExecutable"]["directives"][model_component].items() %}
             --{{key}} = {{value}}
         {%- endfor %}
+        {% endif %}
 
     [[EvaJediLog-{{model_component}}]]
         script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}"
@@ -174,13 +215,23 @@
         script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}"
 
     [[EvaObservations-{{model_component}}]]
+        {% if persistent_workers %}
+        platform = localhost
+        script = """
+            {% set _eva_cpus = scheduling["EvaObservations"]["srun_directives_by_model"].get(model_component, scheduling["EvaObservations"]["srun_directives"]).get('ntasks-per-node', 1) -%}
+            srun --jobid=$(cat $CYLC_WORKFLOW_RUN_DIR/share/persistent_jobid) \
+            --ntasks=1 --cpus-per-task={{ _eva_cpus }} \
+            swell task EvaObservations $config -d $datetime -m {{model_component}}
+        """
+        {% else %}
         script = "swell task EvaObservations $config -d $datetime -m {{model_component}}"
         platform = {{platform}}
         execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}}
         [[[directives]]]
         {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %}
             --{{key}} = {{value}}
         {%- endfor %}
+        {% endif %}
 
     [[SaveObsDiags-{{model_component}}]]
         script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}"
@@ -189,4 +240,39 @@
         script = "swell task CleanCycle $config -d $datetime -m {{model_component}}"
     {% endfor %}
 
+    {% if persistent_workers %}
+    [[Allocate]]
+        script = """
+            SALLOC_ARGS="{% for key, value in salloc_directives.items() %}{% if value %} --{{key}}={{value}}{% else %} --{{key}}{% endif %}{% endfor %}"
+            SALLOC_LOG=$(mktemp)
+            salloc --no-shell -J swell_${CYLC_WORKFLOW_NAME} $SALLOC_ARGS > "$SALLOC_LOG" 2>&1 &
+            SALLOC_PID=$!
+            # Read job ID from salloc's own output instead of polling squeue by
+            # name — eliminates cross-run interference when multiple experiments
+            # with the same workflow name run concurrently.
+            WAITED=0
+            while true; do
+                JOB_ID=$(sed -n 's/.*job allocation \([0-9]*\).*/\1/p' "$SALLOC_LOG" 2>/dev/null | head -n1)
+                [ -n "$JOB_ID" ] && break
+                if ! kill -0 $SALLOC_PID 2>/dev/null; then
+                    echo "salloc exited unexpectedly:" >&2 && cat "$SALLOC_LOG" >&2
+                    rm -f "$SALLOC_LOG" && exit 1
+                fi
+                sleep 2; WAITED=$((WAITED+2))
+                [ $WAITED -ge 60 ] && echo "Timed out getting salloc job ID" >&2 \
+                    && cat "$SALLOC_LOG" >&2 && rm -f "$SALLOC_LOG" && exit 1
+            done
+            rm -f "$SALLOC_LOG"
+            mkdir -p "$CYLC_WORKFLOW_RUN_DIR/share"
+            echo "$JOB_ID" > "$CYLC_WORKFLOW_RUN_DIR/share/persistent_jobid"
+            wait $SALLOC_PID || true
+        """
+
+    [[Deallocate]]
+        script = """
+            JOB_ID=$(cat "$CYLC_WORKFLOW_RUN_DIR/share/persistent_jobid" 2>/dev/null)
+            [ -n "$JOB_ID" ] && scancel "$JOB_ID" || true
+        """
+    {% endif %}
+
 # --------------------------------------------------------------------------------------------------

src/swell/utilities/run_jedi_executables.py

@@ -74,25 +74,37 @@ def run_executable(
 
     # Run the JEDI executable
     # -----------------------
+    persistent_job_id = os.environ.get('SWELL_PERSISTENT_JOB_ID')
+
     if (perhost is None or perhost == "None"):
         logger.info(f"Running {jedi_executable_path} with {str(np)} processors.")
-        command = [
-            'mpirun',
-            '-np', str(np),
-            jedi_executable_path,
-            jedi_config_file
-        ]
+        if persistent_job_id:
+            nodes_env = os.environ.get('SWELL_SRUN_NODES')
+            command = ['srun', '--mpi=pmi2', '--jobid', persistent_job_id,
+                       '--exclusive', '-n', str(np)]
+            if nodes_env:
+                nodes = int(nodes_env)
+                command += ['--nodes', nodes_env, '--ntasks-per-node', str(np // nodes)]
+            command += [jedi_executable_path, jedi_config_file]
+        else:
+            command = ['mpirun', '-np', str(np), jedi_executable_path, jedi_config_file]
     else:
         logger.info(
             f"Running {jedi_executable_path} with {str(np)} processors & perhost {str(perhost)}"
         )
-        command = [
-            'mpirun',
-            '-np', str(np),
-            '-perhost', str(perhost),
-            jedi_executable_path,
-            jedi_config_file
-        ]
+        if persistent_job_id:
+            nodes_env = os.environ.get('SWELL_SRUN_NODES')
+            nodes = int(nodes_env) if nodes_env else (np // perhost if np % perhost == 0 else None)
+            command = ['srun', '--mpi=pmi2', '--jobid', persistent_job_id,
+                       '--exclusive', '-n', str(np), '--ntasks-per-node', str(perhost)]
+            if nodes:
+                command += ['--nodes', str(nodes)]
+            command += [jedi_executable_path, jedi_config_file]
+        else:
+            command = [
+                'mpirun', '-np', str(np), '-perhost', str(perhost),
+                jedi_executable_path, jedi_config_file
+            ]
         s = ('mpi_command='+" ".join(command)+' '+output_log)
         logger.debug(s)
 

src/swell/utilities/slurm.py

@@ -16,6 +16,15 @@
 from swell.utilities.logger import Logger
 
 
+_SRUN_RESOURCE_KEYS = {
+    'nodes', 'ntasks', 'ntasks-per-node', 'cpus-per-task',
+    'mem', 'mem-per-cpu', 'gres',
+}
+
+# sbatch-only flags that salloc does not accept.
+_SALLOC_EXCLUDED_KEYS = {'no-requeue'}
+
+
 def prepare_scheduling_dict(
     logger: Logger,
     experiment_dict: dict,
@@ -41,6 +50,8 @@ def prepare_scheduling_dict(
     with resources.open_text(path_import, 'slurm.yaml') as yaml_file:
         global_defaults = yaml.load(yaml_file)
 
+    platform_persistent_workers = global_defaults.pop('persistent_workers', False)
+
     # Hard-coded SLURM defaults for certain tasks
     # -------------------------------------------
     task_defaults = {
@@ -187,6 +198,45 @@ def prepare_scheduling_dict(
         if slurm_task in experiment_task_directives.keys():
             x = experiment_task_directives[slurm_task].get('execution_time_limit', x)
         scheduling_dict[slurm_task]['execution_time_limit'] = x
+
+    # Build persistent-worker when enabled.
+    persistent_workers = experiment_dict.get('persistent_workers', platform_persistent_workers)
+    if persistent_workers:
+        salloc = {}
+        for task, task_info in scheduling_dict.items():
+            for key, val in task_info['directives']['all'].items():
+                if key in ('job-name', *_SALLOC_EXCLUDED_KEYS):
+                    continue
+                if key in _SRUN_RESOURCE_KEYS and isinstance(val, (int, float)):
+                    salloc[key] = max(salloc.get(key, 0), val)
+                elif key not in salloc:
+                    salloc[key] = val
+
+        # Without an explicit --time, SLURM uses the partition max (e.g. 12h)
+        # and excludes the job from backfill scheduling, causing long queue waits.
+        # Default to 40 min based on observed workflow runtimes; you can override via
+        # slurm_directives_global.time if runs need longer.
+        if 'time' not in salloc:
+            salloc['time'] = '00:40:00'
+
+        scheduling_dict['salloc_directives'] = salloc
+
+        for task in slurm_tasks:
+            task_dir = scheduling_dict[task]['directives']['all']
+            srun_directives = {}
+            for k, v in task_dir.items():
+                if k in _SRUN_RESOURCE_KEYS:
+                    srun_directives[k] = v
+            scheduling_dict[task]['srun_directives'] = srun_directives
+            for model_component in model_components:
+                mc_dir = scheduling_dict[task]['directives'].get(model_component, {})
+                srun_directives = {}
+                for k, v in mc_dir.items():
+                    if k in _SRUN_RESOURCE_KEYS:
+                        srun_directives[k] = v
+                scheduling_dict[task].setdefault('srun_directives_by_model', {})
+                scheduling_dict[task]['srun_directives_by_model'][model_component] = srun_directives
+
     return scheduling_dict