feat(job-future): Add JobFuture implementation for monitoring running slurm jobs.

slurm_manager/job_future.py

@@ -0,0 +1,168 @@
+"""
+Module defining a JobFuture class that represents a Slurm job and allows users to check
+its status, register callbacks for events and wait for its completion.
+
+The JobFuture class provides a high-level interface for interacting with Slurm jobs.
+"""
+
+from typing import Callable
+import threading
+import time
+import logging
+from uuid import uuid4
+from concurrent.futures import Future
+
+from dataclasses import dataclass
+from enum import StrEnum
+from slurm_manager.utils.utils import SLURM_JOB_STATE_CODES
+
+logger = logging.getLogger(__name__)
+
+
+class JobStatus(StrEnum):
+    PENDING = "PENDING"
+    RUNNING = "RUNNING"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    CANCELLED = "CANCELLED"
+
+
+@dataclass
+class RegisteredCallback:
+    id: str
+    event: str
+    callback: Callable[..., None]
+
+
+class JobFuture:
+    """
+    A future-like object that represents a Slurm job. It allows users to check the status of the job,
+    retrieve the job's result, and register callbacks for events (e.g., 'status', 'stdout', 'stderr').
+    """
+
+    def __init__(self, job_id: str, manager, registry: dict | None = None, name: str | None = None):
+        self._job_id = job_id
+        self._manager = manager
+        self._registry = registry if registry is not None else {}
+        self._pubsub_client = manager.pubsub_client
+        self._future = Future()
+        self.name = name if name else f"Job-{job_id}"
+
+        self._status: str = JobStatus.PENDING.value
+        self._slurm_state: str | None = None
+
+        self._shutdown_event = threading.Event()
+        self._callback_registry: dict[str, list[RegisteredCallback]] = {"status": []}
+        self._subscribe_to_event("status", self._update_status)
+
+    #####################
+    ## Public API
+    #####################
+
+    @property
+    def status(self) -> str | None:
+        return self._status
+
+    # TBD - do we want to expose the raw slurm state in addition?
+    # Maybe this can be more hidden.
+    @property
+    def slurm_state(self) -> str | None:
+        return self._slurm_state
+
+    def cancel(self):
+        self._manager.cancel_job(self.job_id)
+
+    def exception(self) -> Exception | None:
+        if self._future.done():
+            return self._future.exception()
+        return None
+
+    def success(self) -> bool:
+        return self._future.done() and self._future.exception() is None
+
+    def done(self) -> bool:
+        return self._future.done()
+
+    def set_finished(self):
+        self._shutdown_event.set()
+        if self._future.done():
+            # TBD - have we decided
+            print(
+                f"Warning: Job {self.name} is already marked as done. Ignoring set_finished call."
+            )
+            logger.warning(
+                f"Job {self.name} is already marked as done. Ignoring set_finished call."
+            )
+            return
+        self._future.set_result(None)
+
+    def set_exception(self, status, exc_info=None):
+        if self._future.done():
+            print(
+                f"Warning: Job {self.name} is already marked as done. Ignoring set_exception call."
+            )
+            logger.warning(
+                f"Job {self.name} is already marked as done. Ignoring set_exception call."
+            )
+            return
+        self._shutdown_event.set()
+        # Here you could log the exception info or store it for later retrieval
+        if exc_info:
+            info = f"Job {self.name} failed with status {status}. Exception info: {exc_info}"
+        else:
+            info = f"Job {self.name} failed with status {status}."
+        # You could log the info here, e.g., using logging.error(info)
+        self._future.set_exception(Exception(info))
+
+    def wait(self, timeout: float | None = None):
+        if self.done():
+            self._raise_if_failed()
+            return self
+
+        start_time = time.time()
+        poll_interval = 0.05  # seconds
+        while not self.done():
+            if timeout is not None and (time.time() - start_time) > timeout:
+                raise TimeoutError(f"Job {self.job_id} did not complete within {timeout} seconds.")
+            if self._shutdown_event.wait(timeout=poll_interval):
+                break
+
+    def add_callback(self, event: str, callback: Callable[..., None]) -> str:
+        cb = RegisteredCallback(id=str(uuid4()), event=event, callback=callback)
+        if event not in self._callback_registry:
+            self._callback_registry[event] = []
+            self._subscribe_to_event(event)
+        self._callback_registry[event].append(cb)
+        return cb.id
+
+    def remove_callback(self, callback_id: str):
+        for event, callbacks in self._callback_registry.items():
+            for cb in callbacks:
+                if cb.id == callback_id:
+                    callbacks.remove(cb)
+                    return
+
+    ###################
+    ## Internal Methods
+    ###################
+
+    def _update_status(self, status: str, exc_info=None):
+        self._status = JobStatus(status).value
+        if status == JobStatus.COMPLETED:
+            self.set_finished()
+        elif status in (JobStatus.FAILED, JobStatus.CANCELLED):
+            # Consider using exception info here to provide more context about the failure
+            self.set_exception(status, exc_info)
+
+        for cb in self._callback_registry["status"]:
+            cb.callback(status)
+
+    def _update_slurm_state(self, slurm_state: str):
+        self._slurm_state = SLURM_JOB_STATE_CODES.get(slurm_state, slurm_state)
+        # Depending on which event, we may like to trigger in addition here that the status changes
+        # at least for events when the SBatch job is terminated, and we don't expect any more updates. TBD.
+
+    def _subscribe_to_event(self, event: str):
+        # Subscribe to the relevant Redis channel for the event
+        channel = f"job/{self.job_id}/{event}"
+        self._pubsub_client.subscribe(channel, self._handle_event)

slurm_manager/utils/utils.py

@@ -0,0 +1,26 @@
+SLURM_JOB_STATE_CODES = {
+    "BF": "BOOT_FAIL",  # Job terminated due to launch failure, typically due to a hardware failure (e.g. unable to boot the node or block and the job can not be requeued).
+    "CA": "CANCELLED",  # Job was explicitly cancelled by the user or system administrator. The job may or may not have been initiated.
+    "CD": "COMPLETED",  # Job has terminated all processes on all nodes with an exit code of zero.
+    "CF": "CONFIGURING",  # Job has been allocated resources, but are waiting for them to become ready for use (e.g. booting).
+    "CG": "COMPLETING",  # Job is in the process of completing. Some processes on some nodes may still be active.
+    "DL": "DEADLINE",  # Job terminated on deadline.
+    "F": "FAILED",  # Job terminated with non-zero exit code or other failure condition.
+    "NF": "NODE_FAIL",  # Job terminated due to failure of one or more allocated nodes.
+    "OOM": "OUT_OF_MEMORY",  # Job experienced out of memory error.
+    "PD": "PENDING",  # Job is awaiting resource allocation.
+    "PR": "PREEMPTED",  # Job terminated due to preemption.
+    "R": "RUNNING",  #   Job currently has an allocation.
+    "RD": "RESV_DEL_HOLD",  # Job is being held after requested reservation was deleted.
+    "RF": "REQUEUE_FED",  # Job is being requeued by a federation.
+    "RH": "REQUEUE_HOLD",  # Held job is being requeued.
+    "RQ": "REQUEUED",  # Completing job is being requeued.
+    "RS": "RESIZING",  # Job is about to change size.
+    "RV": "REVOKED",  # Sibling was removed from cluster due to other cluster starting the job.
+    "SI": "SIGNALING",  # Job is being signaled.
+    "SE": "SPECIAL_EXIT",  # The job was requeued in a special state. This state can be set by users, typically in EpilogSlurmctld, if the job has terminated with a particular exit value.
+    "SO": "STAGE_OUT",  # Job is staging out files.
+    "ST": "STOPPED",  # Job has an allocation, but execution has been stopped with SIGSTOP signal. CPUS have been retained by this job.
+    "S": "SUSPENDED",  # Job has an allocation, but execution has been suspended and CPUs have been released for other jobs.
+    "TO": "TIMEOUT",  # Job terminated upon reaching its time limit
+}