Batch scheduler executors and the local executor now export PSIJ_NODEFILE,

which contains a list of nodes allocated to the job.

Author: hategan

src/psij/executors/batch/batch_scheduler_executor.py

@@ -152,6 +152,10 @@ class BatchSchedulerExecutor(JobExecutor):
         2. store the exit code of the launch command in the *exit code file* named
         `<native_id>.ec`, also inside `<script_dir>`.
 
+    Additionally, where appropriate, the submit script should set the environment variable named
+    ``PSIJ_NODEFILE`` to point to a file containing a list of nodes that are allocated for the job,
+    one per line, with a total number of lines matching the process count of the job.
+
     Once the submit script is generated, the executor renders the submit command using
     :func:`~get_submit_command` and executes it. Its output is then parsed using
     :func:`~job_id_from_submit_output` to retrieve the `native_id` of the job. Subsequently, the

src/psij/executors/batch/cobalt/cobalt.mustache

@@ -37,6 +37,9 @@ only results in empty files that are not cleaned up}}
 #COBALT -e /dev/null
 #COBALT -o /dev/null
 
+{{!like PBS, this is also cheap and there is not need to check setting}}
+PSIJ_NODEFILE="$COBALT_NODEFILE"
+export PSIJ_NODEFILE
 
 {{!redirect output here instead of through #COBALT directive since COBALT_JOB_ID is not available
 when the directives are evaluated; the reason for using the job id in the first place being the

src/psij/executors/batch/lsf/lsf.mustache

@@ -64,6 +64,9 @@ only results in empty files that are not cleaned up}}
 #BSUB -e /dev/null
 #BSUB -o /dev/null
 
+PSIJ_NODEFILE="$LSB_HOSTS"
+export PSIJ_NODEFILE
+
 {{!redirect output here instead of through #BSUB directive since LSB_JOBID is not available
 when the directives are evaluated; the reason for using the job id in the first place being the
 same as for the exit code file.}}

src/psij/executors/batch/pbspro/pbspro.mustache

@@ -40,6 +40,9 @@ only results in empty files that are not cleaned up}}
 #PBS -v {{name}}={{value}}
 {{/env}}
 
+PSIJ_NODEFILE="$PBS_NODEFILE"
+export PSIJ_NODEFILE
+
 
 {{#job.spec.directory}}
 cd "{{.}}"

src/psij/executors/batch/slurm.py

@@ -185,3 +185,7 @@ def _format_duration(self, d: timedelta) -> str:
         if d.days > 0:
             days = str(d.days) + '-'
         return days + "%s:%s:%s" % (d.seconds // 3600, (d.seconds // 60) % 60, d.seconds % 60)
+
+    def _clean_submit_script(self, job: Job) -> None:
+        super()._clean_submit_script(job)
+        self._delete_aux_file(job, '.nodefile')

src/psij/executors/batch/slurm/slurm.mustache

@@ -73,6 +73,11 @@ only results in empty files that are not cleaned up}}
 #SBATCH -e /dev/null
 #SBATCH -o /dev/null
 
+PSIJ_NODEFILE="{{psij.script_dir}}/$SLURM_JOB_ID.nodefile"
+scontrol show hostnames >"$PSIJ_NODEFILE"
+export PSIJ_NODEFILE
+
+
 {{#env}}
 #SBATCH --export={{name}}={{value}}
 {{/env}}

src/psij/executors/local.py

@@ -6,12 +6,13 @@
 import threading
 import time
 from abc import ABC, abstractmethod
+from tempfile import mkstemp
 from types import FrameType
 from typing import Optional, Dict, List, Tuple, Type, cast
 
 import psutil
 
-from psij import InvalidJobException, SubmitException, Launcher
+from psij import InvalidJobException, SubmitException, Launcher, ResourceSpecV1
 from psij import Job, JobSpec, JobExecutorConfig, JobState, JobStatus
 from psij import JobExecutor
 from psij.utils import SingletonThread
@@ -67,6 +68,7 @@ class _ChildProcessEntry(_ProcessEntry):
     def __init__(self, job: Job, executor: 'LocalJobExecutor',
                  launcher: Optional[Launcher]) -> None:
         super().__init__(job, executor, launcher)
+        self.nodefile: Optional[str] = None
 
     def kill(self) -> None:
         super().kill()
@@ -75,6 +77,8 @@ def poll(self) -> Tuple[Optional[int], Optional[str]]:
         assert self.process is not None
         exit_code = self.process.poll()
         if exit_code is not None:
+            if self.nodefile:
+                os.unlink(self.nodefile)
             if self.process.stdout:
                 return exit_code, self.process.stdout.read().decode('utf-8')
             else:
@@ -103,19 +107,30 @@ def poll(self) -> Tuple[Optional[int], Optional[str]]:
             return None, None
 
 
-def _get_env(spec: JobSpec) -> Optional[Dict[str, str]]:
+def _get_env(spec: JobSpec, nodefile: Optional[str]) -> Optional[Dict[str, str]]:
+    env: Optional[Dict[str, str]] = None
     if spec.inherit_environment:
-        if not spec.environment:
+        if spec.environment is None and nodefile is None:
             # if env is none in Popen, it inherits env from parent
             return None
         else:
             # merge current env with spec env
             env = os.environ.copy()
-            env.update(spec.environment)
+            if spec.environment:
+                env.update(spec.environment)
+            if nodefile is not None:
+                env['PSIJ_NODEFILE'] = nodefile
             return env
     else:
         # only spec env
-        return spec.environment
+        if nodefile is None:
+            env = spec.environment
+        else:
+            env = {'PSIJ_NODEFILE': nodefile}
+            if spec.environment:
+                env.update(spec.environment)
+
+        return env
 
 
 class _ProcessReaper(SingletonThread):
@@ -222,6 +237,26 @@ def __init__(self, url: Optional[str] = None,
         super().__init__(url=url, config=config if config else JobExecutorConfig())
         self._reaper = _ProcessReaper.get_instance()
 
+    def _generate_nodefile(self, job: Job, p: _ChildProcessEntry) -> Optional[str]:
+        assert job.spec is not None
+        if job.spec.resources is None:
+            return None
+        if job.spec.resources.version == 1:
+            assert isinstance(job.spec.resources, ResourceSpecV1)
+            n = job.spec.resources.computed_process_count
+            if n == 1:
+                # as a bit of an optimization, we don't generate a nodefile when doing "single
+                # node" jobs on local.
+                return None
+            (file, p.nodefile) = mkstemp(suffix='.nodelist')
+            for i in range(n):
+                os.write(file, 'localhost\n'.encode())
+            os.close(file)
+            return p.nodefile
+        else:
+            raise SubmitException('Cannot handle resource specification with version %s'
+                                  % job.spec.resources.version)
+
     def submit(self, job: Job) -> None:
         """
         Submits the specified :class:`~psij.Job` to be run locally.
@@ -245,8 +280,10 @@ def submit(self, job: Job) -> None:
                 if job.status.state == JobState.CANCELED:
                     raise SubmitException('Job canceled')
             logger.debug('Running %s,  out=%s, err=%s', args, spec.stdout_path, spec.stderr_path)
+            nodefile = self._generate_nodefile(job, p)
+            env = _get_env(spec, nodefile)
             p.process = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
-                                         close_fds=True, cwd=spec.directory, env=_get_env(spec))
+                                         close_fds=True, cwd=spec.directory, env=env)
             self._reaper.register(p)
             job._native_id = p.process.pid
             self._set_job_status(job, JobStatus(JobState.QUEUED, time=time.time(),

tests/plugins1/_batch_test/_batch_test.py

@@ -91,6 +91,10 @@ def _get_state(self, state: str) -> JobState:
         assert state in _TestJobExecutor._STATE_MAP
         return _TestJobExecutor._STATE_MAP[state]
 
+    def _clean_submit_script(self, job: Job):
+        super()._clean_submit_script(job)
+        self._delete_aux_file(job, '.nodefile')
+
 
 class _TestLauncher(MultipleLauncher):
     def __init__(self, config: Optional[JobExecutorConfig] = None):

tests/plugins1/_batch_test/test/test.mustache

@@ -22,6 +22,14 @@ export {{key}}="{{value}}"
     {{/custom_attributes.test}}
 {{/job.spec.attributes}}
 
+HOSTNAME=`hostname`
+PSIJ_NODEFILE="{{psij.script_dir}}/$PSIJ_BATCH_TEST_JOB_ID.nodefile"
+rm -f "$PSIJ_NODEFILE"
+for NODE in $(seq 1 1 "$PSIJ_TEST_BATCH_EXEC_COUNT"); do
+    echo  "$HOSTNAME-$NODE" >> "$PSIJ_NODEFILE"
+done
+
+export PSIJ_NODEFILE
 {{#job.spec.inherit_environment}}env \{{/job.spec.inherit_environment}}{{^job.spec.inherit_environment}}env --ignore-environment \{{/job.spec.inherit_environment}}{{#env}}
 {{name}}="{{value}}" \
 {{/env}}{{#psij.launch_command}}{{.}} {{/psij.launch_command}}

tests/test_nodefile.py

@@ -0,0 +1,30 @@
+import os
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import pytest
+
+from _test_tools import assert_completed
+from executor_test_params import ExecutorTestParams
+from psij import Job, JobSpec, JobExecutor, ResourceSpecV1
+
+NOT_TESTED = set(['rp', 'flux'])
+
+
+def test_nodefile(execparams: ExecutorTestParams) -> None:
+    if execparams.executor in NOT_TESTED:
+        pytest.skip('This test does not work with %s' % execparams.executor)
+
+    my_path = os.path.dirname(os.path.realpath(__file__))
+
+    N_PROC = 4
+    with TemporaryDirectory(dir=Path.home() / '.psij' / 'test') as td:
+        outp = Path(td, 'stdout.txt')
+        spec = JobSpec('/bin/bash', [os.path.join(my_path, 'test_nodefile.sh'), str(N_PROC)],
+                       stdout_path=outp)
+        job = Job(spec)
+        spec.resources = ResourceSpecV1(process_count=N_PROC)
+        ex = JobExecutor.get_instance(execparams.executor)
+        ex.submit(job)
+        status = job.wait()
+        assert_completed(job, status)