Refactor JobStatus

Author: AryanAhadinia

src/swiss_ai_model_launch/launchers/firecrest_launcher.py

@@ -170,7 +170,7 @@ async def get_job_status(self, job_id: int) -> JobStatus:
             jobid=str(job_id),
             # account=self.account,  # TODO
         )
-        return JobStatus(str(job_info[0]["status"]["state"]))
+        return JobStatus.from_str(str(job_info[0]["status"]["state"]))
 
     async def get_job_logs(self, job_id: int) -> tuple[str, str]:
         log_dir = Path(self._get_working_dir()) / "logs" / str(job_id)

src/swiss_ai_model_launch/launchers/launcher.py

@@ -11,6 +11,13 @@ class JobStatus(Enum):
     TIMEOUT = "TIMEOUT"
     UNKNOWN = "UNKNOWN"
 
+    @classmethod
+    def from_str(cls, state: str) -> "JobStatus":
+        try:
+            return cls(state)
+        except ValueError:
+            return cls.UNKNOWN
+
 
 class Launcher(ABC):
     def __init__(

src/swiss_ai_model_launch/launchers/slurm_launcher.py

@@ -19,14 +19,6 @@
 
 _APP_WORKING_DIRECTORY = ".sml"
 
-_SLURM_STATE_MAP: dict[str, JobStatus] = {
-    "PENDING": JobStatus.PENDING,
-    "CONFIGURING": JobStatus.PENDING,
-    "RUNNING": JobStatus.RUNNING,
-    "COMPLETING": JobStatus.RUNNING,
-    "TIMEOUT": JobStatus.TIMEOUT,
-}
-
 
 class SlurmLauncher(Launcher):
     def __init__(
@@ -177,7 +169,7 @@ async def get_job_status(self, job_id: int) -> JobStatus:
         state = stdout.decode().strip()
 
         if state:
-            return _SLURM_STATE_MAP.get(state, JobStatus.UNKNOWN)
+            return JobStatus.from_str(state)
 
         # Job not in squeue — check sacct for terminal state
         proc = await asyncio.create_subprocess_exec(
@@ -194,7 +186,7 @@ async def get_job_status(self, job_id: int) -> JobStatus:
         stdout, _ = await proc.communicate()
         lines = [line.strip() for line in stdout.decode().splitlines() if line.strip()]
         if lines:
-            return _SLURM_STATE_MAP.get(lines[0].split()[0], JobStatus.UNKNOWN)
+            return JobStatus.from_str(lines[0].split()[0])
 
         return JobStatus.UNKNOWN