File staging

.github/workflows/python-package.yml

@@ -37,6 +37,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
+    #- uses: mxschmitt/action-tmate@v3
+    #  with:
+    #    detached: true
+    #    limit-access-to-actor: true
     - name: Test with pytest
       uses: addnab/docker-run-action@v3
       with:
@@ -46,7 +50,7 @@ jobs:
           cd /workspace
           echo "Running in ${PWD}"
           sudo apt update
-          sudo apt install -y openssh-server openssh-client
+          sudo apt install -y openssh-server openssh-client netcat-traditional
           mkdir -p "$HOME/.ssh"
           chmod 0755 "$HOME"
           chmod 0700 "$HOME/.ssh"

.gitignore

@@ -17,5 +17,6 @@ psi_j_python.egg-info/
 venv*
 .venv*
 build/
+.packages/
 docs/.web-build
-web-build/
+web-build/

src/psij-descriptors/cobalt_descriptor.py

@@ -3,5 +3,5 @@
 from psij.descriptor import Descriptor
 
 
-__PSI_J_EXECUTORS__ = [Descriptor(name="cobalt", nice_name='Cobalt', version=StrictVersion("0.0.1"),
+__PSI_J_EXECUTORS__ = [Descriptor(name="cobalt", nice_name='Cobalt', version=StrictVersion("0.2.0"),
                                   cls='psij.executors.batch.cobalt.CobaltJobExecutor')]

src/psij-descriptors/core_descriptors.py

@@ -2,15 +2,15 @@
 from psij.descriptor import Descriptor
 
 __PSI_J_EXECUTORS__ = [
-    Descriptor(name='local', nice_name='Local', version=StrictVersion('0.0.1'),
+    Descriptor(name='local', nice_name='Local', version=StrictVersion('0.2.0'),
                cls='psij.executors.local.LocalJobExecutor')
 ]
 
 __PSI_J_LAUNCHERS__ = [
-    Descriptor(name='single', version=StrictVersion('0.0.1'),
+    Descriptor(name='single', version=StrictVersion('0.2.0'),
                cls='psij.launchers.single.SingleLauncher'),
-    Descriptor(name='multiple', version=StrictVersion('0.0.1'),
+    Descriptor(name='multiple', version=StrictVersion('0.2.0'),
                cls='psij.launchers.multiple.MultipleLauncher'),
-    Descriptor(name='mpirun', version=StrictVersion('0.0.1'),
+    Descriptor(name='mpirun', version=StrictVersion('0.2.0'),
                cls='psij.launchers.mpirun.MPILauncher'),
 ]

src/psij-descriptors/lsf_descriptor.py

@@ -3,5 +3,5 @@
 from psij.descriptor import Descriptor
 
 
-__PSI_J_EXECUTORS__ = [Descriptor(name='lsf', nice_name='LSF', version=StrictVersion('0.0.1'),
+__PSI_J_EXECUTORS__ = [Descriptor(name='lsf', nice_name='LSF', version=StrictVersion('0.2.0'),
                                   cls='psij.executors.batch.lsf.LsfJobExecutor')]

src/psij-descriptors/pbs_descriptor.py

@@ -4,8 +4,8 @@
 
 
 __PSI_J_EXECUTORS__ = [Descriptor(name='pbs', nice_name='PBS Pro', aliases=['pbspro'],
-                                  version=StrictVersion('0.0.2'),
+                                  version=StrictVersion('0.2.0'),
                                   cls='psij.executors.batch.pbs.PBSJobExecutor'),
                        Descriptor(name='pbs_classic', nice_name='PBS Classic', aliases=['torque'],
-                                  version=StrictVersion('0.0.2'),
+                                  version=StrictVersion('0.2.0'),
                                   cls='psij.executors.batch.pbs_classic.PBSClassicJobExecutor')]

src/psij-descriptors/slurm_descriptor.py

@@ -3,5 +3,5 @@
 from psij.descriptor import Descriptor
 
 
-__PSI_J_EXECUTORS__ = [Descriptor(name='slurm', nice_name='Slurm', version=StrictVersion('0.0.1'),
+__PSI_J_EXECUTORS__ = [Descriptor(name='slurm', nice_name='Slurm', version=StrictVersion('0.2.0'),
                                   cls='psij.executors.batch.slurm.SlurmJobExecutor')]

src/psij/executors/batch/batch_scheduler_executor.py

@@ -1,19 +1,22 @@
 import logging
 import os
+import weakref
+
 import subprocess
 import time
 import traceback
 from abc import abstractmethod
 from datetime import timedelta
 from pathlib import Path
 from threading import Thread, RLock
-from typing import Optional, List, Dict, Collection, cast, Union, IO
+from typing import Optional, List, Dict, Collection, cast, Union, IO, Set
 
 from psij.launchers.script_based_launcher import ScriptBasedLauncher
 
 from psij import JobExecutor, JobExecutorConfig, Launcher, Job, SubmitException, \
     JobStatus, JobState
 from psij.executors.batch.template_function_library import ALL as FUNCTION_LIBRARY
+from psij.utils import _StatusUpdater
 
 UNKNOWN_ERROR = 'PSIJ: Unknown error'
 
@@ -199,18 +202,14 @@
             configuration is used.
         """
         super().__init__(url=url, config=config if config else BatchSchedulerExecutorConfig())
+        self._queue_poll_thread = self._start_queue_poll_thread()
         assert config
         self.work_directory = config.work_directory / self.name
-        self._queue_poll_thread = self._start_queue_poll_thread()
-
-    def _ensure_work_dir(self) -> None:
         self.work_directory.mkdir(parents=True, exist_ok=True)
 
     def submit(self, job: Job) -> None:
         """See :func:`~psij.JobExecutor.submit`."""
         logger.info('Job %s: submitting', job.id)
-        self._ensure_work_dir()
-
         self._check_job(job)
 
         context = self._create_script_context(job)
@@ -489,7 +488,10 @@
             'psij': {
                 'lib': FUNCTION_LIBRARY,
                 'launch_command': launch_command,
-                'script_dir': str(self.work_directory)
+                'script_dir': str(self.work_directory),
+                'us_file': self._queue_poll_thread.status_updater.update_file_name,
+                'us_port': self._queue_poll_thread.status_updater.update_port,
+                'us_addrs': ', '.join(self._queue_poll_thread.status_updater.ips)
             }
         }
         assert job.spec is not None
@@ -531,9 +533,11 @@
             # is_greater_than returns T/F if the states are comparable and None if not, so
             # we have to check explicitly for the boolean value rather than truthiness
             return
-        if status.state.final and job.native_id:
-            self._clean_submit_script(job)
-            self._read_aux_files(job, status)
+        if status.state.final:
+            self._queue_poll_thread.unregister_job(job)
+            if job.native_id:
+                self._clean_submit_script(job)
+                self._read_aux_files(job, status)
         super()._set_job_status(job, status)
 
     def _clean_submit_script(self, job: Job) -> None:
@@ -573,9 +577,8 @@
                     # already present
                     out = self._read_aux_file(job, '.out')
                     if out:
-                        launcher = self._get_launcher_from_job(job)
-                        if launcher.is_launcher_failure(out):
-                            status.message = launcher.get_launcher_failure_message(out)
+                        if '_PSIJ_SCRIPT_DONE' not in out:
+                            status.message = out
                     logger.debug('Output from launcher: %s', status.message)
                 else:
                     self._delete_aux_file(job, '.out')
@@ -638,29 +641,46 @@
         super().__init__()
         self.name = name
         self.daemon = True
-        self.config = config
-        self.executor = executor
+        # We don't at this time cache executor instances. Even if we did, it may be wise
+        # to shut down queue polling threads when their executors (the only entities that
+        # use them) are garbage collected. So we wrap the references to the executor and
+        # config in a weak ref and exit when the ref becomes invalid.
+        self.config = weakref.ref(config)
+        self.executor = weakref.ref(executor)
         # native_id -> job
-        self._jobs: Dict[str, List[Job]] = {}
+        self._jobs: Dict[str, Set[Job]] = {}
         # counts consecutive errors while invoking qstat or equivalent
         self._poll_error_count = 0
         self._jobs_lock = RLock()
+        self.status_updater = cast(_StatusUpdater, _StatusUpdater.get_instance())
+        self.active = True
 
     def run(self) -> None:
         logger.debug('Executor %s: queue poll thread started', self.executor)
-        time.sleep(self.config.initial_queue_polling_delay)
-        while True:
-            self._poll()
-            time.sleep(self.config.queue_polling_interval)
+        try:
+            time.sleep(self.get_config().initial_queue_polling_delay)
+            while self.active:
+                self._poll()
+                start = time.time()
+                now = start
+                while now - start < self.get_config().queue_polling_interval:
+                    time.sleep(1)
+                    now = time.time()
+        except StopIteration:
+            logger.info('Thread %s exiting due to executor collection' % self)
+
+    def stop(self) -> None:
+        self.active = False
 
     def _poll(self) -> None:
+        executor = self.get_executor()
         with self._jobs_lock:
             if len(self._jobs) == 0:
                 return
             jobs_copy = dict(self._jobs)
         logger.info('Polling for %s jobs', len(jobs_copy))
         try:
-            out = self.executor._run_command(self.executor.get_status_command(jobs_copy.keys()))
+            out = executor._run_command(executor.get_status_command(jobs_copy.keys()))
         except subprocess.CalledProcessError as ex:
             out = ex.output
             exit_code = ex.returncode
@@ -674,25 +694,23 @@
             self._poll_error_count = 0
         logger.debug('Output from status command: %s', out)
         try:
-            status_map = self.executor.parse_status_output(exit_code, out)
+            status_map = executor.parse_status_output(exit_code, out)
         except Exception as ex:
             self._handle_poll_error(False,
                                     ex,
                                     f'Failed to poll for job status: {traceback.format_exc()}')
             return
         try:
-            for native_id, job_list in jobs_copy.items():
+            for native_id, job_set in jobs_copy.items():
                 try:
                     status = self._get_job_status(native_id, status_map)
                 except Exception:
                     status = JobStatus(JobState.FAILED,
                                        message='Failed to update job status: %s' %
                                                traceback.format_exc())
-                for job in job_list:
-                    self.executor._set_job_status(job, status)
-                if status.state.final:
-                    with self._jobs_lock:
-                        del self._jobs[native_id]
+
+                for job in job_set:
+                    executor._set_job_status(job, status)
         except Exception as ex:
             msg = traceback.format_exc()
             self._handle_poll_error(True, ex, 'Error updating job statuses {}'.format(msg))
@@ -706,7 +724,7 @@
     def _handle_poll_error(self, immediate: bool, ex: Exception, msg: str) -> None:
         logger.warning('Polling error: %s', msg)
         self._poll_error_count += 1
-        if immediate or (self._poll_error_count > self.config.queue_polling_error_threshold):
+        if immediate or (self._poll_error_count > self.get_config().queue_polling_error_threshold):
             self._poll_error_count = 0
             # fail all jobs
             with self._jobs_lock:
@@ -718,16 +736,47 @@
                 assert len(self._jobs) > 0
                 jobs_copy = dict(self._jobs)
                 self._jobs.clear()
-            for job_list in jobs_copy.values():
-                for job in job_list:
-                    self.executor._set_job_status(job, JobStatus(JobState.FAILED, message=msg))
+            for job_set in jobs_copy.values():
+                for job in job_set:
+                    self.unregister_job(job)
+                    self.get_executor()._set_job_status(job, JobStatus(JobState.FAILED,
+                                                                       message=msg))
 
     def register_job(self, job: Job) -> None:
+        self.status_updater.register_job(job, self.get_executor())
         assert job.native_id
         logger.info('Job %s: registering', job.id)
         with self._jobs_lock:
             native_id = job.native_id
-            if native_id not in self._jobs:
-                self._jobs[native_id] = [job]
-            else:
-                self._jobs[job.native_id].append(job)
+            try:
+                self._jobs[native_id].add(job)
+            except KeyError:
+                self._jobs[native_id] = {job}
+
+    def unregister_job(self, job: Job) -> None:
+        self.status_updater.unregister_job(job)
+        assert job.native_id
+        logger.info('Job %s: unregistering', job.id)
+        with self._jobs_lock:
+            native_id = job.native_id
+            try:
+                del self._jobs[native_id]
+            except KeyError:
+                # If two or more jobs are attached to the same native ID, the
+                # first one being unregistered would already have removed
+                # the dict entry
+                pass
+
+    def get_config(self) -> BatchSchedulerExecutorConfig:
+        config = self.config()
+        if config:
+            return config
+        else:
+            raise StopIteration()
+
+    def get_executor(self) -> BatchSchedulerExecutor:
+        ex = self.executor()
+        if ex:
+            return ex
+        else:
+            raise StopIteration()

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

@@ -46,16 +46,27 @@ 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
 same as for the exit code file.}}
 exec &>> "{{psij.script_dir}}/$COBALT_JOBID.out"
 
+
+{{> batch_lib}}
+
+{{!like PBS, this is also cheap and there is not need to check setting}}
+PSIJ_NODEFILE="$COBALT_NODEFILE"
+export PSIJ_NODEFILE
+
+{{> stagein}}
+_psij_update_status ACTIVE
+
 {{#psij.launch_command}}{{.}} {{/psij.launch_command}}
+_PSIJ_JOB_EC=$?
+
+{{> stageout}}
+{{> cleanup}}
 
 {{!we redirect to a file tied to the native ID so that we can reach the file with attach().}}
-echo "$?" > "{{psij.script_dir}}/$COBALT_JOBID.ec"
+echo $_PSIJ_JOB_EC > "{{psij.script_dir}}/$COBALT_JOBID.ec"
+echo "_PSIJ_SCRIPT_DONE"