experimaestro 2.0.0b4__py3-none-any.whl → 2.0.0b17__py3-none-any.whl

experimaestro/launcherfinder/registry.py

@@ -59,9 +59,9 @@ class LauncherRegistry:
             ).expanduser()
 
         if LauncherRegistry.CURRENT_CONFIG_DIR not in LauncherRegistry.INSTANCES:
-            LauncherRegistry.INSTANCES[
-                LauncherRegistry.CURRENT_CONFIG_DIR
-            ] = LauncherRegistry(LauncherRegistry.CURRENT_CONFIG_DIR)
+            LauncherRegistry.INSTANCES[LauncherRegistry.CURRENT_CONFIG_DIR] = (
+                LauncherRegistry(LauncherRegistry.CURRENT_CONFIG_DIR)
+            )
 
         return LauncherRegistry.INSTANCES[LauncherRegistry.CURRENT_CONFIG_DIR]
 
@@ -83,7 +83,6 @@ class LauncherRegistry:
 
         # Register the find launcher function if it exists
         launchers_py = basepath / "launchers.py"
-        print(f"basepath {launchers_py}")
         if launchers_py.is_file():
             logger.info("Loading %s", launchers_py)
 
@@ -163,9 +162,9 @@ class LauncherRegistry:
         if self.find_launcher_fn is not None:
             for spec in specs.requirements:
                 if launcher := self.find_launcher_fn(spec, tags):
-                    assert isinstance(
-                        launcher, Launcher
-                    ), "f{self.find_launcher_fn} did not return a Launcher but {type(launcher)}"
+                    assert isinstance(launcher, Launcher), (
+                        "f{self.find_launcher_fn} did not return a Launcher but {type(launcher)}"
+                    )
                     return launcher
 
         return None

experimaestro/launcherfinder/specs.py

@@ -27,10 +27,7 @@ class CudaSpecification:
         return (self.memory >= spec.memory) and (self.min_memory <= spec.memory)
 
     def __repr__(self):
-        return (
-            f"CUDA({self.model} "
-            f"max={format_size(self.memory, binary=True)}/min={format_size(self.min_memory, binary=True)})"
-        )
+        return f"CUDA({self.model} max={format_size(self.memory, binary=True)}/min={format_size(self.min_memory, binary=True)})"
 
 
 @dataclass
@@ -48,11 +45,7 @@ class CPUSpecification:
     """Number of CPU per GPU (0 if not defined)"""
 
     def __repr__(self):
-        return (
-            f"CPU("
-            f"mem={format_size(self.memory, binary=True)}, cores={self.cores}"
-            ")"
-        )
+        return f"CPU(mem={format_size(self.memory, binary=True)}, cores={self.cores})"
 
     def match(self, other: "CPUSpecification"):
         return (self.memory >= other.memory) and (self.cores >= other.cores)

experimaestro/launchers/slurm/__init__.py

@@ -1,2 +1,2 @@
-# flake8: noqa: F401
-from .base import *
+# ruff: noqa: F401
+from .base import *  # noqa: F403

experimaestro/launchers/slurm/base.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import threading
 from pathlib import Path
@@ -219,6 +220,13 @@ class SlurmProcessWatcher(threading.Thread):
         self.cv = ThreadingCondition()
         self.fetched_event = threading.Event()
         self.updating_jobs = threading.Lock()
+
+        # Async waiters: jobid -> list of (asyncio.Event, event_loop)
+        self.async_waiters: Dict[
+            str, List[Tuple[asyncio.Event, asyncio.AbstractEventLoop]]
+        ] = {}
+        self.async_waiters_lock = threading.Lock()
+
         self.start()
 
     @staticmethod
@@ -250,6 +258,35 @@ class SlurmProcessWatcher(threading.Thread):
         with self.updating_jobs:
             return self.jobs.get(jobid)
 
+    def register_async_waiter(
+        self, jobid: str, loop: asyncio.AbstractEventLoop
+    ) -> asyncio.Event:
+        """Register an async waiter for a job.
+
+        Returns an asyncio.Event that will be set when the job finishes.
+        """
+        event = loop.create_future()
+        with self.async_waiters_lock:
+            if jobid not in self.async_waiters:
+                self.async_waiters[jobid] = []
+            self.async_waiters[jobid].append((event, loop))
+        return event
+
+    def _notify_async_waiters(self):
+        """Notify async waiters for finished jobs"""
+        with self.async_waiters_lock:
+            finished_jobs = []
+            for jobid, waiters in self.async_waiters.items():
+                state = self.jobs.get(jobid)
+                if state and state.finished():
+                    finished_jobs.append(jobid)
+                    for future, loop in waiters:
+                        # Set the result from watcher thread to asyncio loop
+                        loop.call_soon_threadsafe(future.set_result, state)
+
+            for jobid in finished_jobs:
+                del self.async_waiters[jobid]
+
     def run(self):
         while self.count > 0:
             builder = self.launcher.connector.processbuilder()
@@ -280,6 +317,9 @@ class SlurmProcessWatcher(threading.Thread):
                             logger.error("Could not parse line %s", line)
             process.kill()
 
+            # Notify async waiters for finished jobs
+            self._notify_async_waiters()
+
             with self.cv:
                 logger.debug("Jobs %s", self.jobs)
                 self.fetched_event.set()
@@ -310,6 +350,28 @@ class BatchSlurmProcess(Process):
                     self._last_state = state
                     return 0 if state.slurm_state == "COMPLETED" else 1
 
+    async def aio_wait(self) -> int:
+        """Asynchronously wait for SLURM job to finish (event-driven)"""
+        logger.debug("Async waiting for SLURM job %s", self.jobid)
+        loop = asyncio.get_running_loop()
+
+        with SlurmProcessWatcher.get(self.launcher) as watcher:
+            # Check if already finished
+            state = watcher.getjob(self.jobid)
+            if state and state.finished():
+                self._last_state = state
+                return 0 if state.slurm_state == "COMPLETED" else 1
+
+            # Register and wait for the job to finish
+            future = watcher.register_async_waiter(self.jobid, loop)
+            self._last_state = await future
+
+            code = 0 if self._last_state.slurm_state == "COMPLETED" else 1
+            logger.debug(
+                "Finished async wait for SLURM job %s: code %s", self.jobid, code
+            )
+            return code
+
     def get_job_state(self, code: int) -> "JobState":
         """Convert SLURM exit code to JobState, detecting timeouts"""
         from experimaestro.scheduler.jobs import (