Properly implement wait according to spec (i.e., wait for all states

greater than `target_states` not just for `target_states`).

This goes a bit further by also implementing the yet-to-be-merged updates
from ExaWorks/job-api-spec#178

Author: hategan

src/psij/__init__.py

@@ -7,7 +7,7 @@
 from typing import Callable, TypeVar
 
 from psij.descriptor import Descriptor
-from .exceptions import SubmitException, InvalidJobException, UnreachableStateException
+from .exceptions import SubmitException, InvalidJobException
 from .job import Job, JobStatusCallback
 from .job_attributes import JobAttributes
 from .job_executor import JobExecutor

src/psij/exceptions.py

@@ -66,26 +66,3 @@ def __init__(self, message: str, exception: Optional[Exception] = None,
         conditions such an error would persist across subsequent re-tries until correct credentials
         are used.
         """
-
-
-class UnreachableStateException(Exception):
-    """
-    Indicates that a job state being waited for cannot be reached.
-
-    This exception is thrown when the :func:`~psij.Job.wait` method is called with a set of
-    states that cannot be reached by the job when the call is made.
-    """
-
-    def __init__(self, status: 'psij.JobStatus') -> None:
-        """
-        Constructs an `UnreachableStateException`.
-
-        :param status: The :class:`~psij.JobStatus` that the job was in when
-            :func:`~psij.Job.wait` was called and which prevents the desired states to be
-            reached.
-        """
-        self.status = status
-        """
-        Returns the job status that has caused an implementation to determine that the desired
-        states passed to the :func:`~psij.Job.wait` method cannot be reached.
-        """

src/psij/job.py

@@ -2,11 +2,11 @@
 import threading
 from abc import ABC, abstractmethod
 from datetime import timedelta, datetime
-from typing import Optional, Sequence, Union, Callable
+from typing import Optional, Sequence, Union, Callable, Set
 from uuid import uuid4
 
 import psij
-from psij.exceptions import SubmitException, UnreachableStateException
+from psij.exceptions import SubmitException
 from psij.job_spec import JobSpec
 from psij.job_state import JobState, JobStateOrder
 from psij.job_status import JobStatus
@@ -161,8 +161,23 @@ def cancel(self) -> None:
         else:
             self.executor.cancel(self)
 
+    def _all_greater(self, states: Optional[Union[JobState, Sequence[JobState]]]) \
+            -> Optional[Set[JobState]]:
+        if states is None:
+            return None
+        if isinstance(states, JobState):
+            states = [states]
+        ts = set()
+        for state1 in states:
+            ts.add(state1)
+            for state2 in JobState:
+                if state2.is_greater_than(state1):
+                    ts.add(state2)
+        return ts
+
     def wait(self, timeout: Optional[timedelta] = None,
-             target_states: Optional[Sequence[JobState]] = None) -> Optional[JobStatus]:
+             target_states: Optional[Union[JobState, Sequence[JobState]]] = None) \
+            -> Optional[JobStatus]:
         """
         Waits for the job to reach certain states.
 
@@ -186,15 +201,15 @@ def wait(self, timeout: Optional[timedelta] = None,
             timeout = LARGE_TIMEOUT
         end = start + timeout
 
+        ts = self._all_greater(target_states)
+
         while True:
             with self._status_cv:
                 status = self._status
                 state = status.state
-                if target_states:
-                    if state in target_states:
+                if ts:
+                    if state.final or state in ts:
                         return status
-                    elif state.final:
-                        raise UnreachableStateException(status)
                     else:
                         pass  # wait
                 else:

src/psij/job_state.py

@@ -108,6 +108,10 @@ def __bool__(self) -> bool:
         """All states are consider true-ish."""
         return True
 
+    def __hash__(self) -> int:
+        """Returns a hash for this object."""
+        return self._value_  # type: ignore
+
 
 class JobStateOrder:
     """A class that can be used to reconstruct missing states."""

tests/test_wait.py

@@ -0,0 +1,31 @@
+from datetime import timedelta
+
+from psij import Job, JobExecutor, JobSpec, JobState
+
+
+def _test_wait() -> None:
+    ex = JobExecutor.get_instance('local')
+    job = Job(JobSpec('/bin/sleep', ['4']))
+    ex.submit(job)
+    status = job.wait(target_states=JobState.ACTIVE)
+    assert status is not None
+    assert status.state == JobState.ACTIVE
+
+    status = job.wait(target_states=JobState.QUEUED)
+    assert status is not None
+    assert status.state == JobState.ACTIVE
+
+    status = job.wait(timedelta(milliseconds=100))
+    assert status is None
+
+    status = job.wait()
+    assert status is not None
+    assert status.state == JobState.COMPLETED
+
+    status = job.wait(target_states=JobState.QUEUED)
+    assert status is not None
+    assert status.state == JobState.COMPLETED
+
+    status = job.wait(target_states=JobState.FAILED)
+    assert status is not None
+    assert status.state == JobState.COMPLETED