py-cluster-api 0.1.0__py3-none-any.whl

cluster_api/exceptions.py

@@ -0,0 +1,17 @@
+"""Exception hierarchy for cluster_api."""
+
+
+class ClusterAPIError(Exception):
+    """Base exception for all cluster-api errors."""
+
+
+class CommandTimeoutError(ClusterAPIError):
+    """A subprocess command exceeded its timeout."""
+
+
+class CommandFailedError(ClusterAPIError):
+    """A subprocess command returned a non-zero exit code."""
+
+
+class SubmitError(ClusterAPIError):
+    """Failed to submit a job or parse its ID from scheduler output."""

cluster_api/executors/__init__.py

@@ -0,0 +1,36 @@
+"""Executor registry."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..core import Executor
+
+_REGISTRY: dict[str, type[Executor]] = {}
+
+
+def _ensure_builtins() -> None:
+    """Lazily register built-in executors."""
+    if "lsf" in _REGISTRY and "local" in _REGISTRY:
+        return
+    from .lsf import LSFExecutor
+    from .local import LocalExecutor
+
+    _REGISTRY.setdefault("lsf", LSFExecutor)
+    _REGISTRY.setdefault("local", LocalExecutor)
+
+
+def get_executor_class(name: str) -> type[Executor]:
+    """Get an executor class by name."""
+    _ensure_builtins()
+    if name not in _REGISTRY:
+        raise ValueError(
+            f"Unknown executor: {name!r}. Available: {list(_REGISTRY.keys())}"
+        )
+    return _REGISTRY[name]
+
+
+def register_executor(name: str, cls: type[Executor]) -> None:
+    """Register a custom executor class."""
+    _REGISTRY[name] = cls

cluster_api/executors/local.py

@@ -0,0 +1,107 @@
+"""Local subprocess executor for testing without a real scheduler."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from datetime import datetime, timezone
+from typing import Any
+
+from .._types import JobStatus, ResourceSpec
+from ..config import ClusterConfig
+from ..core import Executor
+
+logger = logging.getLogger(__name__)
+
+
+class LocalExecutor(Executor):
+    """Runs jobs as local bash subprocesses. Useful for testing."""
+
+    submit_command = "bash"
+    cancel_command = "kill"
+    status_command = "ps"
+    directive_prefix = "# LOCAL"
+
+    def __init__(self, config: ClusterConfig) -> None:
+        super().__init__(config)
+        self._processes: dict[str, asyncio.subprocess.Process] = {}
+        self._next_id = 1
+
+    def build_header(
+        self, name: str, resources: ResourceSpec | None = None
+    ) -> list[str]:
+        """Local executor doesn't need scheduler directives."""
+        return [f"# LOCAL Job: {name}"]
+
+    async def _submit_job(
+        self,
+        script_path: str,
+        name: str,
+        env: dict[str, str] | None = None,
+    ) -> str:
+        """Run script as a background subprocess."""
+        full_env = {**os.environ, **(env or {})}
+
+        proc = await asyncio.create_subprocess_exec(
+            "bash", script_path,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=full_env,
+        )
+
+        job_id = str(self._next_id)
+        self._next_id += 1
+        self._processes[job_id] = proc
+        return job_id
+
+    def _build_status_args(self) -> list[str]:
+        # Not used for local executor; poll() is overridden
+        return []
+
+    def _parse_job_statuses(
+        self, output: str
+    ) -> dict[str, tuple[JobStatus, dict[str, Any]]]:
+        # Not used for local executor; poll() is overridden
+        return {}
+
+    async def poll(self) -> dict[str, JobStatus]:
+        """Check subprocess return codes."""
+        for job_id, record in self._jobs.items():
+            if record.is_terminal:
+                continue
+
+            proc = self._processes.get(job_id)
+            if proc is None:
+                continue
+
+            if proc.returncode is not None:
+                # Process finished
+                now = datetime.now(timezone.utc)
+                record.finish_time = now
+                record._last_seen = now
+                if proc.returncode == 0:
+                    record.status = JobStatus.DONE
+                    record.exit_code = 0
+                else:
+                    record.status = JobStatus.FAILED
+                    record.exit_code = proc.returncode
+            else:
+                record.status = JobStatus.RUNNING
+                record._last_seen = datetime.now(timezone.utc)
+
+        return {jid: r.status for jid, r in self._jobs.items()}
+
+    async def cancel(self, job_id: str) -> None:
+        """Terminate a local subprocess."""
+        proc = self._processes.get(job_id)
+        if proc and proc.returncode is None:
+            proc.terminate()
+            try:
+                await asyncio.wait_for(proc.wait(), timeout=5.0)
+            except asyncio.TimeoutError:
+                proc.kill()
+
+        if job_id in self._jobs:
+            self._jobs[job_id].status = JobStatus.KILLED
+        logger.info("Cancelled local job %s", job_id)

cluster_api/executors/lsf.py

@@ -0,0 +1,307 @@
+"""LSF executor using bsub/bjobs/bkill."""
+
+from __future__ import annotations
+
+import fnmatch
+import json
+import logging
+import math
+import re
+from datetime import datetime, timezone
+from typing import Any
+
+from .._types import JobStatus, ResourceSpec
+from ..config import ClusterConfig, parse_memory_bytes
+from ..core import Executor
+from ..exceptions import ClusterAPIError
+
+logger = logging.getLogger(__name__)
+
+
+_LSF_STATUS_MAP: dict[str, JobStatus] = {
+    "PEND": JobStatus.PENDING,
+    "RUN": JobStatus.RUNNING,
+    "DONE": JobStatus.DONE,
+    "EXIT": JobStatus.FAILED,
+    "ZOMBI": JobStatus.FAILED,
+    "USUSP": JobStatus.PENDING,
+    "PSUSP": JobStatus.PENDING,
+    "SSUSP": JobStatus.PENDING,
+}
+
+_BJOBS_FIELDS = (
+    "jobid stat exit_code exec_host max_mem "
+    "submit_time start_time finish_time"
+)
+
+
+def lsf_format_bytes_ceil(n_bytes: int, lsf_units: str = "MB") -> str:
+    """Format bytes into LSF memory units, rounding up.
+
+    Inspired by dask-jobqueue's lsf_format_bytes_ceil.
+    """
+    units = {"KB": 1024, "MB": 1024**2, "GB": 1024**3, "TB": 1024**4}
+    if lsf_units not in units:
+        raise ValueError(f"Unknown LSF units: {lsf_units}")
+    return str(math.ceil(n_bytes / units[lsf_units]))
+
+
+async def lsf_detect_units(
+    timeout: float = 100.0,
+) -> str:
+    """Detect LSF memory units from lsadmin output.
+
+    Inspired by dask-jobqueue's approach.
+    """
+    try:
+        out = await Executor._call(
+            ["lsadmin", "showconf", "lim"],
+            timeout=timeout,
+        )
+        for line in out.splitlines():
+            if "LSF_UNIT_FOR_LIMITS" in line:
+                return line.split("=")[-1].strip().upper()
+    except (ClusterAPIError, OSError):
+        pass
+    return "KB"  # LSF default
+
+
+class LSFExecutor(Executor):
+    """LSF executor using bsub, bjobs, bkill."""
+
+    submit_command = "bsub"
+    cancel_command = "bkill"
+    status_command = "bjobs"
+    directive_prefix = "#BSUB"
+    job_id_regexp = r"Job <(?P<job_id>\d+)>"
+
+    def __init__(self, config: ClusterConfig) -> None:
+        super().__init__(config)
+        self._lsf_units = config.lsf_units
+
+    def build_header(
+        self, name: str, resources: ResourceSpec | None = None
+    ) -> list[str]:
+        """Build #BSUB directive lines."""
+        lines: list[str] = []
+        p = self.directive_prefix
+
+        lines.append(f"{p} -J {name}")
+
+        log_dir = self.config.log_directory
+        lines.append(f"{p} -o {log_dir}/{name}.out")
+        lines.append(f"{p} -e {log_dir}/{name}.err")
+
+        # Queue
+        queue = (resources and resources.queue) or self.config.queue
+        if queue:
+            lines.append(f"{p} -q {queue}")
+
+        # Account/project
+        account = (resources and resources.account) or self.config.account
+        if account:
+            lines.append(f"{p} -P {account}")
+
+        # CPUs
+        cpus = (resources and resources.cpus) or self.config.cpus
+        if cpus:
+            lines.append(f"{p} -n {cpus}")
+            if cpus > 1:
+                lines.append(f'{p} -R "span[hosts=1]"')
+
+        # GPUs
+        gpus = (resources and resources.gpus) or self.config.gpus
+        if gpus:
+            lines.append(f'{p} -gpu "num={gpus}"')
+
+        # Memory
+        memory_str = (resources and resources.memory) or self.config.memory
+        if memory_str:
+            mem_bytes = parse_memory_bytes(memory_str)
+            mem_val = lsf_format_bytes_ceil(mem_bytes, self._lsf_units)
+            lines.append(f"{p} -M {mem_val}")
+            lines.append(f'{p} -R "rusage[mem={mem_val}]"')
+
+        # Walltime
+        walltime = (resources and resources.walltime) or self.config.walltime
+        if walltime:
+            lines.append(f"{p} -W {walltime}")
+
+        # Working directory
+        work_dir = resources and resources.work_dir
+        if work_dir:
+            lines.append(f"{p} -cwd {work_dir}")
+
+        # Custom cluster options
+        if resources and resources.cluster_options:
+            for opt in resources.cluster_options:
+                lines.append(f"{p} {opt}")
+
+        return lines
+
+    def _build_submit_env(self, env: dict[str, str] | None) -> dict[str, str] | None:
+        """Build environment dict for bsub, applying email suppression."""
+        submit_env = dict(env) if env else {}
+        if self.config.suppress_job_email:
+            submit_env["LSB_JOB_REPORT_MAIL"] = "N"
+        return submit_env or None
+
+    async def _bsub(
+        self, script_path: str, content: str | None, env: dict[str, str] | None,
+    ) -> str:
+        """Run bsub via stdin or file and return raw output."""
+        submit_env = self._build_submit_env(env)
+        if self.config.use_stdin:
+            if content is None:
+                with open(script_path) as f:
+                    content = f.read()
+            return await self._call(
+                [self.submit_command],
+                env=submit_env,
+                timeout=self.config.command_timeout,
+                stdin_data=content,
+            )
+        return await self._call(
+            [self.submit_command, script_path],
+            env=submit_env,
+            timeout=self.config.command_timeout,
+        )
+
+    async def _submit_job(
+        self,
+        script_path: str,
+        name: str,
+        env: dict[str, str] | None = None,
+    ) -> str:
+        """Submit via bsub with stdin mode support."""
+        out = await self._bsub(script_path, None, env)
+        return self._job_id_from_submit_output(out)
+
+    async def _submit_array_job(
+        self,
+        script_path: str,
+        name: str,
+        array_range: tuple[int, int],
+        env: dict[str, str] | None = None,
+        max_concurrent: int | None = None,
+    ) -> str:
+        """Submit an array job with -J 'name[start-end]'."""
+        array_spec = f"{array_range[0]}-{array_range[1]}"
+        if max_concurrent is not None:
+            array_spec += f"%{max_concurrent}"
+        array_name = f"{name}[{array_spec}]"
+
+        # Rewrite only #BSUB directive lines for array syntax
+        with open(script_path) as f:
+            lines = f.readlines()
+        new_lines = []
+        for line in lines:
+            if line.startswith(self.directive_prefix):
+                line = line.replace(f"-J {name}", f"-J {array_name}")
+                line = line.replace(f"{name}.out", f"{name}.%I.out")
+                line = line.replace(f"{name}.err", f"{name}.%I.err")
+            new_lines.append(line)
+        content = "".join(new_lines)
+        with open(script_path, "w") as f:
+            f.write(content)
+
+        out = await self._bsub(script_path, content, env)
+        return self._job_id_from_submit_output(out)
+
+    def _build_status_args(self) -> list[str]:
+        """Build bjobs command with JSON output."""
+        prefix = self._prefix
+        args = [
+            self.status_command,
+            "-J", f"{prefix}-*",
+            "-a",
+            "-o", _BJOBS_FIELDS,
+            "-json",
+        ]
+        return args
+
+    def _parse_job_statuses(
+        self, output: str
+    ) -> dict[str, tuple[JobStatus, dict[str, Any]]]:
+        """Parse bjobs JSON output into status + metadata dicts."""
+        result: dict[str, tuple[JobStatus, dict[str, Any]]] = {}
+
+        if not output.strip():
+            return result
+
+        try:
+            data = json.loads(output)
+        except json.JSONDecodeError:
+            logger.warning("Failed to parse bjobs JSON output")
+            return result
+
+        records = data.get("RECORDS", [])
+        for rec in records:
+            job_id = str(rec.get("JOBID", "")).strip()
+            if not job_id:
+                continue
+
+            stat = rec.get("STAT", "").strip()
+            status = _LSF_STATUS_MAP.get(stat, JobStatus.UNKNOWN)
+
+            exit_code_str = str(rec.get("EXIT_CODE", "")).strip()
+            exit_code = None
+            if exit_code_str and exit_code_str != "-":
+                try:
+                    exit_code = int(exit_code_str)
+                except ValueError:
+                    pass
+            # LSF returns "" for exit_code on DONE jobs — infer 0
+            if exit_code is None and status == JobStatus.DONE:
+                exit_code = 0
+
+            meta: dict[str, Any] = {
+                "exec_host": _clean_field(rec.get("EXEC_HOST")),
+                "max_mem": _clean_field(rec.get("MAX_MEM")),
+                "exit_code": exit_code,
+                "submit_time": _parse_lsf_time(rec.get("SUBMIT_TIME")),
+                "start_time": _parse_lsf_time(rec.get("START_TIME")),
+                "finish_time": _parse_lsf_time(rec.get("FINISH_TIME")),
+            }
+
+            result[job_id] = (status, meta)
+
+        return result
+
+    async def cancel_by_name(self, name_pattern: str) -> None:
+        """Cancel jobs matching name pattern via bkill -J."""
+        await self._call(
+            [self.cancel_command, "-J", name_pattern],
+            timeout=self.config.command_timeout,
+        )
+        # Update in-memory state for matching jobs
+        for record in self._jobs.values():
+            if not record.is_terminal and fnmatch.fnmatch(record.name, name_pattern):
+                record.status = JobStatus.KILLED
+        logger.info("Cancelled jobs matching %s", name_pattern)
+
+
+def _clean_field(value: Any) -> str | None:
+    """Clean a bjobs field value, returning None for empty/dash values."""
+    if value is None:
+        return None
+    s = str(value).strip()
+    if s in ("", "-"):
+        return None
+    return s
+
+
+def _parse_lsf_time(value: Any) -> datetime | None:
+    """Parse an LSF timestamp string."""
+    s = _clean_field(value)
+    if s is None:
+        return None
+    # Strip trailing timezone indicator (e.g. " L" in "Feb  8 10:31:17 2026 L")
+    s = re.sub(r"\s+[A-Z]$", "", s)
+    # LSF timestamps are typically like "Jan  1 12:00:00 2024"
+    for fmt in ("%b %d %H:%M:%S %Y", "%b  %d %H:%M:%S %Y", "%Y/%m/%d-%H:%M:%S"):
+        try:
+            return datetime.strptime(s, fmt).replace(tzinfo=timezone.utc)
+        except ValueError:
+            continue
+    return None

cluster_api/monitor.py

@@ -0,0 +1,168 @@
+"""Async polling loop and callback dispatch."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+from ._types import JobExitCondition, JobRecord, JobStatus
+
+if TYPE_CHECKING:
+    from .core import Executor
+
+logger = logging.getLogger(__name__)
+
+
+class JobMonitor:
+    """Monitors job status via polling and dispatches callbacks."""
+
+    def __init__(self, executor: Executor, poll_interval: float | None = None) -> None:
+        self.executor = executor
+        self.poll_interval = poll_interval or executor.config.poll_interval
+        self._task: asyncio.Task | None = None
+        self._stopped = asyncio.Event()
+        self._completion_events: dict[str, asyncio.Event] = {}
+
+    async def start(self) -> None:
+        """Start the polling loop."""
+        self._stopped.clear()
+        self._task = asyncio.create_task(self._poll_loop())
+        logger.info("Monitor started (interval=%.1fs)", self.poll_interval)
+
+    async def stop(self) -> None:
+        """Stop the polling loop and wait for it to finish."""
+        self._stopped.set()
+        if self._task:
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+            self._task = None
+        logger.info("Monitor stopped")
+
+    async def _poll_loop(self) -> None:
+        """Main polling loop."""
+        while not self._stopped.is_set():
+            try:
+                await self.executor.poll()
+                await self._dispatch_all_callbacks()
+                await self._check_zombies()
+                self._notify_waiters()
+                await self._purge_completed()
+            except Exception:
+                logger.exception("Error in poll loop")
+
+            try:
+                await asyncio.wait_for(self._stopped.wait(), timeout=self.poll_interval)
+                break  # stopped was set
+            except asyncio.TimeoutError:
+                pass  # normal: just means poll_interval elapsed
+
+    async def _dispatch_all_callbacks(self) -> None:
+        """Check all jobs for pending callbacks."""
+        for record in list(self.executor.jobs.values()):
+            if record.is_terminal and record._callbacks:
+                await self._dispatch_callbacks(record)
+
+    async def _dispatch_callbacks(self, record: JobRecord) -> None:
+        """Fire matching callbacks for a terminal job, then clear them."""
+        condition_map = {
+            JobStatus.DONE: JobExitCondition.SUCCESS,
+            JobStatus.FAILED: JobExitCondition.FAILURE,
+            JobStatus.KILLED: JobExitCondition.KILLED,
+        }
+        job_condition = condition_map.get(record.status)
+
+        fired: list[int] = []
+        for i, (condition, callback) in enumerate(record._callbacks):
+            if condition == JobExitCondition.ANY or condition == job_condition:
+                try:
+                    result = callback(record)
+                    if inspect.isawaitable(result):
+                        await result
+                except Exception:
+                    logger.exception(
+                        "Callback error for job %s", record.job_id
+                    )
+                fired.append(i)
+
+        # Remove fired callbacks in reverse order to preserve indices
+        for i in reversed(fired):
+            record._callbacks.pop(i)
+
+    async def _check_zombies(self) -> None:
+        """Mark jobs as FAILED if not seen by scheduler for > zombie_timeout."""
+        timeout_minutes = self.executor.config.zombie_timeout_minutes
+        now = datetime.now(timezone.utc)
+
+        for record in self.executor.jobs.values():
+            if record.is_terminal:
+                continue
+            if record._last_seen is None:
+                continue
+            elapsed = (now - record._last_seen).total_seconds() / 60.0
+            if elapsed > timeout_minutes:
+                logger.warning(
+                    "Zombie detected: job %s not seen for %.1f minutes",
+                    record.job_id, elapsed,
+                )
+                record.status = JobStatus.FAILED
+                record.metadata["zombie"] = True
+
+    async def _purge_completed(self) -> None:
+        """Remove terminal jobs older than completed_retention."""
+        retention_minutes = self.executor.config.completed_retention_minutes
+        now = datetime.now(timezone.utc)
+
+        to_remove = []
+        for job_id, record in self.executor.jobs.items():
+            if not record.is_terminal:
+                continue
+            if not record._callbacks:
+                # Use finish_time if available, else _last_seen
+                ref_time = record.finish_time or record._last_seen
+                if ref_time is not None:
+                    elapsed = (now - ref_time).total_seconds() / 60.0
+                    if elapsed > retention_minutes:
+                        to_remove.append(job_id)
+
+        for job_id in to_remove:
+            self.executor.remove_job(job_id)
+            logger.debug("Purged completed job %s", job_id)
+
+    def _notify_waiters(self) -> None:
+        """Set completion events for terminal jobs that have waiters."""
+        done = []
+        for job_id, event in self._completion_events.items():
+            record = self.executor.jobs.get(job_id)
+            if record is not None and record.is_terminal:
+                event.set()
+                done.append(job_id)
+        for job_id in done:
+            del self._completion_events[job_id]
+
+    async def wait_for(
+        self, *records: JobRecord, timeout: float | None = None
+    ) -> None:
+        """Wait until all given jobs reach a terminal state."""
+        events = []
+        for r in records:
+            if r.is_terminal:
+                continue
+            if r.job_id not in self._completion_events:
+                self._completion_events[r.job_id] = asyncio.Event()
+            events.append(self._completion_events[r.job_id])
+
+        if not events:
+            return
+
+        async def _wait() -> None:
+            await asyncio.gather(*(e.wait() for e in events))
+
+        if timeout is not None:
+            await asyncio.wait_for(_wait(), timeout=timeout)
+        else:
+            await _wait()