Merge pull request #384 from ExaWorks/implement_list

A preliminary implementation for the list() method for batch schedule…

Author: hategan

src/psij/executors/batch/batch_scheduler_executor.py

@@ -420,6 +420,38 @@ def parse_status_output(self, exit_code: int, out: str) -> Dict[str, JobStatus]:
         """
         pass
 
+    @abstractmethod
+    def get_list_command(self) -> List[str]:
+        """Constructs a command to retrieve the list of jobs known to the LRM for the current user.
+
+        Concrete implementations of batch scheduler executors must override this method. Upon
+        running the command, the output can be parsed with :func:`~parse_list_output`.
+
+        Returns
+        -------
+        A list of strings representing the executable and arguments to invoke in order to obtain
+        the list of jobs the LRM knows for the current user.
+        """
+        pass
+
+    def parse_list_output(self, out: str) -> List[str]:
+        """Parses the output of the command obtained from :func:`~get_list_command`.
+
+        The default implementation of this method assumes that the output has no header and
+        consists of native IDs, one per line, possibly surrounded by whitespace. Concrete
+        implementations should override this method if a different format is expected.
+
+        Parameters
+        ----------
+        out
+            The output from the "list" command as returned by :func:`~get_list_command`.
+        Returns
+        -------
+        A list of strings representing the native IDs of the jobs known to the LRM for the current
+        user.
+        """
+        return [s.strip() for s in out.splitlines()]
+
     def _create_script_context(self, job: Job) -> Dict[str, object]:
         launcher = self._get_launcher_from_job(job)
         if isinstance(launcher, ScriptBasedLauncher) and logger.isEnabledFor(logging.DEBUG):
@@ -551,7 +583,10 @@ def list(self) -> List[str]:
         Implementations are encouraged to restrict the results to jobs accessible by the current
         user.
         """
-        raise NotImplementedError()
+        return self.parse_list_output(self._run_command(self.get_list_command()))
+
+    def _current_user(self) -> str:
+        return os.getlogin()
 
 
 class _QueuePollThread(Thread):

src/psij/executors/batch/cobalt.py

@@ -111,6 +111,10 @@ def parse_status_output(self, exit_code: int, out: str) -> Dict[str, JobStatus]:
                 index += 1
         return job_statuses
 
+    def get_list_command(self) -> List[str]:
+        """See :meth:`~.BatchSchedulerExecutor.get_list_command`."""
+        return [_QSTAT_COMMAND, '-u', self._current_user(), '--header', 'JobId']
+
     def job_id_from_submit_output(self, out: str) -> str:
         """See :meth:`~.BatchSchedulerExecutor.job_id_from_submit_output`."""
         match = _QSUB_REGEX.search(out)

src/psij/executors/batch/lsf.py

@@ -127,3 +127,7 @@ def job_id_from_submit_output(self, out: str) -> str:
         if match is None:
             raise SubmitException(out)
         return match.group(0)[5:-1]
+
+    def get_list_command(self) -> List[str]:
+        """See :meth:`~.BatchSchedulerExecutor.get_list_command`."""
+        return [_BJOBS_COMMAND, '-a', '-noheader', '-o', 'jobid', '-u', self._current_user()]

src/psij/executors/batch/pbspro.py

@@ -134,6 +134,14 @@ def parse_status_output(self, exit_code: int, out: str) -> Dict[str, JobStatus]:
 
         return r
 
+    def get_list_command(self) -> List[str]:
+        """See :meth:`~.BatchSchedulerExecutor.get_list_command`."""
+        return ['qstat', '-u', self._current_user()]
+
+    def parse_list_output(self, out: str) -> List[str]:
+        """See :meth:`~.BatchSchedulerExecutor.parse_list_output`."""
+        return [s.split()[0].strip() for s in out.splitlines()[2:]]
+
     def _get_state(self, state: str) -> JobState:
         assert state in _STATE_MAP, f"PBS state {state} is not known to PSI/J"
         return _STATE_MAP[state]

src/psij/executors/batch/slurm.py

@@ -160,6 +160,10 @@ def parse_status_output(self, exit_code: int, out: str) -> Dict[str, JobStatus]:
 
         return r
 
+    def get_list_command(self) -> List[str]:
+        """See :meth:`~.BatchSchedulerExecutor.get_list_command`."""
+        return ['squeue', '--me', '-o', '%i', '-h', '-r', '-t', 'all']
+
     def _get_state(self, state: str) -> JobState:
         assert state in SlurmJobExecutor._STATE_MAP
         return SlurmJobExecutor._STATE_MAP[state]

src/psij/job_executor.py

@@ -119,6 +119,8 @@ def submit(self, job: Job) -> None:
         has not been successfully sent to the underlying implementation, the job status remains
         unchanged, and no status notifications about the job will be fired.
 
+        A successful return of this method guarantees that the job's `native_id` property is set.
+
         :raises ~psij.InvalidJobException: Thrown if the job specification cannot be understood.
             This exception is fatal in that submitting another job with the exact same details will
             also fail with an `~psij.InvalidJobException`. In principle, the underlying

tests/plugins1/_batch_test/_batch_test.py

@@ -84,6 +84,9 @@ def parse_status_output(self, exit_code: int, out: str) -> Dict[str, JobStatus]:
             r[native_id] = JobStatus(state, message=msg)
         return r
 
+    def get_list_command(self) -> List[str]:
+        return [sys.executable, QSTAT_PATH]
+
     def _get_state(self, state: str) -> JobState:
         assert state in _TestJobExecutor._STATE_MAP
         return _TestJobExecutor._STATE_MAP[state]

tests/plugins1/_batch_test/test/qstat

@@ -6,16 +6,23 @@ from qlib import *
 from filelock import FileLock
 
 
-ids = sys.argv[1]
+if len(sys.argv) > 1:
+    ids = sys.argv[1]
+else:
+    ids = None
 
 
 with FileLock(lock_file):
     if not state_file.exists():
         sys.exit(0)
     state = read_state()
 
-idlist = ids.split(',')
-for id in idlist:
-    if id in state['jobs']:
-        j = state['jobs'][id]
-        print('%s %s %s' % (id, j['state'], j['error']))
+if ids is not None:
+    idlist = ids.split(',')
+    for id in idlist:
+        if id in state['jobs']:
+            j = state['jobs'][id]
+            print('%s %s %s' % (id, j['state'], j['error']))
+else:
+    for id in state['jobs'].keys():
+        print(id)

tests/test_executor.py

@@ -152,3 +152,12 @@ def test_stdin_redirect(execparams: ExecutorTestParams) -> None:
             contents = outf.read()
 
         assert contents == rnd_str
+
+
+def test_list(execparams: ExecutorTestParams) -> None:
+    job = Job(JobSpec(executable='/bin/sleep', arguments=['4']))
+    ex = _get_executor_instance(execparams, job)
+    ex.submit(job)
+    assert job.native_id is not None
+    ids = ex.list()
+    assert job.native_id in ids