From aa863c80f3b6272cfd638deef8fa467ffe48b727 Mon Sep 17 00:00:00 2001 From: appel_c Date: Tue, 24 Mar 2026 16:21:45 +0100 Subject: [PATCH] feat(job-future): Add JobFuture implementation for monitoring running slurm jobs. --- slurm_manager/job_future.py | 168 ++++++++++++++++++++++++++++++++ slurm_manager/utils/__init__.py | 0 slurm_manager/utils/utils.py | 26 +++++ 3 files changed, 194 insertions(+) create mode 100644 slurm_manager/job_future.py create mode 100644 slurm_manager/utils/__init__.py create mode 100644 slurm_manager/utils/utils.py diff --git a/slurm_manager/job_future.py b/slurm_manager/job_future.py new file mode 100644 index 0000000..534418c --- /dev/null +++ b/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) diff --git a/slurm_manager/utils/__init__.py b/slurm_manager/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/slurm_manager/utils/utils.py b/slurm_manager/utils/utils.py new file mode 100644 index 0000000..dc08437 --- /dev/null +++ b/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 +}