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

cluster_api/__init__.py

@@ -0,0 +1,58 @@
+"""py-cluster-api: Generic Python library for running jobs on HPC clusters."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ._types import ArrayElement, JobExitCondition, JobRecord, JobStatus, ResourceSpec
+from .config import ClusterConfig, load_config
+from .core import Executor
+from .exceptions import (
+    ClusterAPIError,
+    CommandFailedError,
+    CommandTimeoutError,
+    SubmitError,
+)
+from .executors import get_executor_class
+from .executors.local import LocalExecutor
+from .executors.lsf import LSFExecutor
+from .monitor import JobMonitor
+
+__all__ = [
+    "create_executor",
+    "ArrayElement",
+    "Executor",
+    "LSFExecutor",
+    "LocalExecutor",
+    "JobRecord",
+    "JobStatus",
+    "JobExitCondition",
+    "ResourceSpec",
+    "ClusterConfig",
+    "JobMonitor",
+    "load_config",
+    "ClusterAPIError",
+    "CommandFailedError",
+    "CommandTimeoutError",
+    "SubmitError",
+]
+
+
+def create_executor(
+    profile: str | None = None,
+    config_path: str | None = None,
+    **overrides: Any,
+) -> Executor:
+    """Create an executor from config.
+
+    Args:
+        profile: Config profile name to use.
+        config_path: Explicit path to config YAML.
+        **overrides: Override individual config values.
+
+    Returns:
+        An Executor instance configured for the specified backend.
+    """
+    config = load_config(path=config_path, profile=profile, overrides=overrides or None)
+    cls = get_executor_class(config.executor)
+    return cls(config)

cluster_api/_types.py

@@ -0,0 +1,153 @@
+"""Shared types for cluster_api."""
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable
+
+
+class JobStatus(enum.Enum):
+    """Status of a cluster job."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    DONE = "done"
+    FAILED = "failed"
+    KILLED = "killed"
+    UNKNOWN = "unknown"
+
+
+class JobExitCondition(enum.Enum):
+    """Conditions for callback dispatch."""
+
+    SUCCESS = "success"
+    FAILURE = "failure"
+    KILLED = "killed"
+    ANY = "any"
+
+
+_TERMINAL_STATUSES = frozenset({JobStatus.DONE, JobStatus.FAILED, JobStatus.KILLED})
+
+
+@dataclass
+class ResourceSpec:
+    """Resource requirements for a job."""
+
+    cpus: int | None = None
+    gpus: int | None = None
+    memory: str | None = None
+    walltime: str | None = None
+    queue: str | None = None
+    account: str | None = None
+    work_dir: str | None = None
+    cluster_options: list[str] | None = None
+
+
+@dataclass
+class ArrayElement:
+    """Tracks a single element of a job array."""
+
+    index: int
+    status: JobStatus = JobStatus.PENDING
+    exit_code: int | None = None
+    exec_host: str | None = None
+    max_mem: str | None = None
+    submit_time: datetime | None = None
+    start_time: datetime | None = None
+    finish_time: datetime | None = None
+
+
+@dataclass
+class JobRecord:
+    """Tracks a submitted job and its metadata."""
+
+    job_id: str
+    name: str
+    command: str
+    status: JobStatus = JobStatus.PENDING
+    exit_code: int | None = None
+    resources: ResourceSpec | None = None
+    script_path: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+    exec_host: str | None = None
+    max_mem: str | None = None
+    submit_time: datetime | None = None
+    start_time: datetime | None = None
+    finish_time: datetime | None = None
+    array_elements: dict[int, ArrayElement] = field(default_factory=dict)
+    _callbacks: list[tuple[JobExitCondition, Callable]] = field(
+        default_factory=list, repr=False
+    )
+    _last_seen: datetime | None = field(default=None, repr=False)
+
+    def on_exit(self, callback: Callable, condition: JobExitCondition = JobExitCondition.ANY) -> JobRecord:
+        """Register a callback for the given exit condition. Returns self for chaining."""
+        self._callbacks.append((condition, callback))
+        return self
+
+    def on_success(self, callback: Callable) -> JobRecord:
+        """Register a callback for successful completion."""
+        return self.on_exit(callback, JobExitCondition.SUCCESS)
+
+    def on_failure(self, callback: Callable) -> JobRecord:
+        """Register a callback for failure."""
+        return self.on_exit(callback, JobExitCondition.FAILURE)
+
+    @property
+    def is_terminal(self) -> bool:
+        """Whether the job has reached a terminal state."""
+        return self.status in _TERMINAL_STATUSES
+
+    @property
+    def is_array(self) -> bool:
+        """Whether this is an array job."""
+        return "array_range" in self.metadata
+
+    @property
+    def element_count(self) -> int:
+        """Total number of expected array elements (0 if not an array)."""
+        if not self.is_array:
+            return 0
+        start, end = self.metadata["array_range"]
+        return end - start + 1
+
+    @property
+    def completed_elements(self) -> int:
+        """Number of elements in a terminal state."""
+        return sum(1 for e in self.array_elements.values() if e.status in _TERMINAL_STATUSES)
+
+    @property
+    def failed_element_indices(self) -> list[int]:
+        """Indices of failed or killed elements."""
+        return [
+            e.index for e in self.array_elements.values()
+            if e.status in {JobStatus.FAILED, JobStatus.KILLED}
+        ]
+
+    def compute_array_status(self) -> JobStatus:
+        """Derive overall status from element statuses.
+
+        Conservative: only returns a terminal status when ALL expected
+        elements have been seen and are themselves terminal.
+        """
+        if not self.array_elements:
+            return self.status
+
+        statuses = {e.status for e in self.array_elements.values()}
+
+        # Any non-terminal element → still in progress
+        if statuses & {JobStatus.RUNNING, JobStatus.PENDING, JobStatus.UNKNOWN}:
+            return JobStatus.RUNNING
+
+        # All seen elements are terminal — but have we seen them all?
+        if len(self.array_elements) < self.element_count:
+            return JobStatus.RUNNING
+
+        # All expected elements accounted for and terminal
+        if JobStatus.KILLED in statuses:
+            return JobStatus.KILLED
+        if JobStatus.FAILED in statuses:
+            return JobStatus.FAILED
+        return JobStatus.DONE

cluster_api/config.py

@@ -0,0 +1,121 @@
+"""YAML config loader with Nextflow-style profiles."""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+_MEMORY_UNITS = {
+    "B": 1,
+    "KB": 1024,
+    "MB": 1024**2,
+    "GB": 1024**3,
+    "TB": 1024**4,
+}
+
+_MEMORY_RE = re.compile(r"^\s*([\d.]+)\s*(B|KB|MB|GB|TB)\s*$", re.IGNORECASE)
+
+
+def parse_memory_bytes(s: str) -> int:
+    """Parse a memory string like '8 GB' into bytes."""
+    m = _MEMORY_RE.match(s)
+    if not m:
+        raise ValueError(f"Cannot parse memory string: {s!r}")
+    value = float(m.group(1))
+    unit = m.group(2).upper()
+    return int(value * _MEMORY_UNITS[unit])
+
+
+@dataclass
+class ClusterConfig:
+    """Configuration for cluster job execution."""
+
+    executor: str = "local"
+    cpus: int | None = None
+    gpus: int | None = None
+    memory: str | None = None
+    walltime: str | None = None
+    queue: str | None = None
+    account: str | None = None
+    poll_interval: float = 10.0
+    shebang: str = "#!/bin/bash"
+    script_prologue: list[str] = field(default_factory=list)
+    script_epilogue: list[str] = field(default_factory=list)
+    extra_directives: list[str] = field(default_factory=list)
+    directives_skip: list[str] = field(default_factory=list)
+    log_directory: str = "./logs"
+    lsf_units: str = "MB"
+    use_stdin: bool = False
+    job_name_prefix: str | None = None
+    zombie_timeout_minutes: float = 30.0
+    completed_retention_minutes: float = 10.0
+    command_timeout: float = 100.0
+    suppress_job_email: bool = True
+
+
+_CONFIG_SEARCH_PATHS = [
+    "cluster_api.yaml",
+    "~/.config/cluster_api/config.yaml",
+]
+
+
+def _find_config_path() -> Path | None:
+    """Search for config file in standard locations."""
+    env_path = os.environ.get("CLUSTER_API_CONFIG")
+    if env_path:
+        p = Path(env_path).expanduser()
+        if p.exists():
+            return p
+
+    for candidate in _CONFIG_SEARCH_PATHS:
+        p = Path(candidate).expanduser()
+        if p.exists():
+            return p
+
+    return None
+
+
+def load_config(
+    path: str | Path | None = None,
+    profile: str | None = None,
+    overrides: dict[str, Any] | None = None,
+) -> ClusterConfig:
+    """Load configuration from YAML with optional profile and overrides.
+
+    Merges: base config → profile → overrides.
+
+    Raises FileNotFoundError if an explicit path is given but doesn't exist.
+    """
+    raw: dict[str, Any] = {}
+
+    if path is not None:
+        config_path = Path(path)
+        if not config_path.exists():
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+        with open(config_path) as f:
+            raw = yaml.safe_load(f) or {}
+    else:
+        config_path = _find_config_path()
+        if config_path and config_path.exists():
+            with open(config_path) as f:
+                raw = yaml.safe_load(f) or {}
+
+    profiles = raw.pop("profiles", {})
+
+    if profile and profile in profiles:
+        raw = {**raw, **profiles[profile]}
+
+    if overrides:
+        raw = {**raw, **overrides}
+
+    # Build ClusterConfig from the merged dict, ignoring unknown keys
+    known_fields = {f.name for f in ClusterConfig.__dataclass_fields__.values()}
+    filtered = {k: v for k, v in raw.items() if k in known_fields}
+
+    return ClusterConfig(**filtered)

cluster_api/core.py

@@ -0,0 +1,375 @@
+"""Abstract Executor base class."""
+
+from __future__ import annotations
+
+import abc
+import asyncio
+import logging
+import os
+import re
+import secrets
+import string
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from .config import ClusterConfig
+from .exceptions import ClusterAPIError, CommandFailedError, CommandTimeoutError, SubmitError
+from ._types import ArrayElement, JobRecord, JobStatus, ResourceSpec
+
+logger = logging.getLogger(__name__)
+
+_SCRIPT_TEMPLATE = """\
+%(shebang)s
+%(job_header)s
+%(prologue)s
+%(command)s
+%(epilogue)s
+"""
+
+_ARRAY_ELEMENT_RE = re.compile(r"^(.+)\[(\d+)\]$")
+
+
+class Executor(abc.ABC):
+    """Abstract base for cluster job executors."""
+
+    submit_command: str
+    cancel_command: str
+    status_command: str
+    job_id_regexp: str = r"(?P<job_id>\d+)"
+    directive_prefix: str = ""
+
+    def __init__(self, config: ClusterConfig) -> None:
+        self.config = config
+        self._jobs: dict[str, JobRecord] = {}
+        self._script_counter = 0
+        self._log_dir = Path(config.log_directory).expanduser()
+        self._log_dir.mkdir(parents=True, exist_ok=True)
+        if config.job_name_prefix:
+            self._prefix = config.job_name_prefix
+        else:
+            # Generate a random prefix so concurrent users/sessions don't
+            # see each other's jobs when polling by name.
+            alphabet = string.ascii_lowercase + string.digits
+            self._prefix = "".join(secrets.choice(alphabet) for _ in range(5))
+
+    # --- Script rendering ---
+
+    def render_script(
+        self,
+        command: str,
+        name: str,
+        resources: ResourceSpec | None = None,
+        prologue: list[str] | None = None,
+        epilogue: list[str] | None = None,
+    ) -> str:
+        """Render a job script from the template."""
+        header_lines = self.build_header(name, resources)
+        # Filter via directives_skip
+        skip = set(self.config.directives_skip)
+        if skip:
+            header_lines = [
+                line
+                for line in header_lines
+                if not any(s in line for s in skip)
+            ]
+        # Extend with extra_directives
+        header_lines.extend(self.config.extra_directives)
+
+        all_prologue = list(self.config.script_prologue)
+        if prologue:
+            all_prologue.extend(prologue)
+
+        all_epilogue = list(self.config.script_epilogue)
+        if epilogue:
+            all_epilogue.extend(epilogue)
+
+        return _SCRIPT_TEMPLATE % {
+            "shebang": self.config.shebang,
+            "job_header": "\n".join(header_lines),
+            "prologue": "\n".join(all_prologue),
+            "command": command,
+            "epilogue": "\n".join(all_epilogue),
+        }
+
+    @abc.abstractmethod
+    def build_header(
+        self, name: str, resources: ResourceSpec | None = None
+    ) -> list[str]:
+        """Build scheduler-specific directive lines."""
+        ...
+
+    # --- Submission ---
+
+    async def submit(
+        self,
+        command: str,
+        name: str,
+        resources: ResourceSpec | None = None,
+        prologue: list[str] | None = None,
+        epilogue: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> JobRecord:
+        """Submit a job to the scheduler."""
+        full_name = f"{self._prefix}-{name}"
+        script = self.render_script(command, full_name, resources, prologue, epilogue)
+        script_path = self._write_script(script, full_name)
+
+        job_id = await self._submit_job(script_path, full_name, env)
+
+        record = JobRecord(
+            job_id=job_id,
+            name=full_name,
+            command=command,
+            status=JobStatus.PENDING,
+            resources=resources,
+            script_path=script_path,
+            metadata=metadata or {},
+            _last_seen=datetime.now(timezone.utc),
+        )
+        self._jobs[job_id] = record
+        logger.info("Submitted job %s (%s)", job_id, full_name)
+        return record
+
+    async def submit_array(
+        self,
+        command: str,
+        name: str,
+        array_range: tuple[int, int],
+        resources: ResourceSpec | None = None,
+        prologue: list[str] | None = None,
+        epilogue: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        max_concurrent: int | None = None,
+    ) -> JobRecord:
+        """Submit a job array to the scheduler."""
+        full_name = f"{self._prefix}-{name}"
+        script = self.render_script(command, full_name, resources, prologue, epilogue)
+        script_path = self._write_script(script, full_name)
+
+        job_id = await self._submit_array_job(
+            script_path, full_name, array_range, env, max_concurrent
+        )
+
+        meta = {**(metadata or {}), "array_range": array_range}
+        if max_concurrent is not None:
+            meta["max_concurrent"] = max_concurrent
+
+        record = JobRecord(
+            job_id=job_id,
+            name=full_name,
+            command=command,
+            status=JobStatus.PENDING,
+            resources=resources,
+            script_path=script_path,
+            metadata=meta,
+            _last_seen=datetime.now(timezone.utc),
+        )
+        self._jobs[job_id] = record
+        logger.info(
+            "Submitted array job %s (%s[%d-%d])",
+            job_id, full_name, array_range[0], array_range[1],
+        )
+        return record
+
+    async def _submit_job(
+        self,
+        script_path: str,
+        name: str,
+        env: dict[str, str] | None = None,
+    ) -> str:
+        """Submit a script and return the job ID. Override for stdin submission."""
+        out = await self._call(
+            [self.submit_command, script_path],
+            env=env,
+            timeout=self.config.command_timeout,
+        )
+        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. Override in subclasses."""
+        return await self._submit_job(script_path, name, env)
+
+    def _write_script(self, script_content: str, name: str) -> str:
+        """Write job script to log directory and return its path."""
+        safe_name = re.sub(r"[^\w\-.]", "_", name)
+        self._script_counter += 1
+        script_path = self._log_dir / f"{safe_name}.{self._script_counter}.sh"
+        script_path.write_text(script_content)
+        script_path.chmod(0o755)
+        return str(script_path)
+
+    def _job_id_from_submit_output(self, out: str) -> str:
+        """Extract job ID from submission output using regex."""
+        match = re.search(self.job_id_regexp, out)
+        if not match:
+            raise SubmitError(
+                f"Could not parse job ID from output: {out!r}"
+            )
+        return match.group("job_id")
+
+    # --- Cancellation ---
+
+    async def cancel(self, job_id: str) -> None:
+        """Cancel a job by ID."""
+        await self._call(
+            [self.cancel_command, job_id],
+            timeout=self.config.command_timeout,
+        )
+        if job_id in self._jobs:
+            self._jobs[job_id].status = JobStatus.KILLED
+        logger.info("Cancelled job %s", job_id)
+
+    async def cancel_by_name(self, name_pattern: str) -> None:
+        """Cancel jobs by name pattern. Override in subclasses for native support."""
+        raise NotImplementedError("cancel_by_name not supported by this executor")
+
+    async def cancel_all(self) -> None:
+        """Cancel all tracked jobs."""
+        to_cancel = [jid for jid, r in self._jobs.items() if not r.is_terminal]
+        await asyncio.gather(*(self.cancel(jid) for jid in to_cancel))
+
+    # --- Status polling ---
+
+    @abc.abstractmethod
+    def _build_status_args(self) -> list[str]:
+        """Build args for the status query command."""
+        ...
+
+    @abc.abstractmethod
+    def _parse_job_statuses(
+        self, output: str
+    ) -> dict[str, tuple[JobStatus, dict[str, Any]]]:
+        """Parse status command output into {job_id: (status, metadata_dict)}."""
+        ...
+
+    async def poll(self) -> dict[str, JobStatus]:
+        """Query scheduler, update job records, detect zombies. Returns current statuses."""
+        active = [r for r in self._jobs.values() if not r.is_terminal]
+        if not active:
+            return {jid: r.status for jid, r in self._jobs.items()}
+
+        args = self._build_status_args()
+        try:
+            out = await self._call(args, timeout=self.config.command_timeout)
+        except (ClusterAPIError, OSError):
+            logger.warning("Status query failed, skipping poll cycle")
+            return {jid: r.status for jid, r in self._jobs.items()}
+
+        statuses = self._parse_job_statuses(out)
+        now = datetime.now(timezone.utc)
+        array_jobs_updated: set[str] = set()
+
+        for raw_id, (new_status, meta) in statuses.items():
+            # Check if this is an array element ID like "12345[1]"
+            m = _ARRAY_ELEMENT_RE.match(raw_id)
+            if m:
+                parent_id, element_index = m.group(1), int(m.group(2))
+                record = self._jobs.get(parent_id)
+                if record and not record.is_terminal and record.is_array:
+                    if element_index not in record.array_elements:
+                        record.array_elements[element_index] = ArrayElement(index=element_index)
+                    elem = record.array_elements[element_index]
+                    elem.status = new_status
+                    for key in ("exec_host", "max_mem", "exit_code",
+                                "submit_time", "start_time", "finish_time"):
+                        if key in meta and meta[key] is not None:
+                            setattr(elem, key, meta[key])
+                    record._last_seen = now
+                    array_jobs_updated.add(parent_id)
+            else:
+                record = self._jobs.get(raw_id)
+                if record and not record.is_terminal:
+                    record.status = new_status
+                    record._last_seen = now
+                    for key in ("exec_host", "max_mem", "exit_code",
+                                "submit_time", "start_time", "finish_time"):
+                        if key in meta and meta[key] is not None:
+                            setattr(record, key, meta[key])
+
+        # Aggregate parent status for array jobs that got element updates
+        for parent_id in array_jobs_updated:
+            record = self._jobs[parent_id]
+            record.status = record.compute_array_status()
+
+        return {jid: r.status for jid, r in self._jobs.items()}
+
+    # --- Subprocess helper ---
+
+    @staticmethod
+    async def _call(
+        cmd: list[str],
+        shell: bool = False,
+        timeout: float = 100.0,
+        env: dict[str, str] | None = None,
+        stdin_data: str | None = None,
+    ) -> str:
+        """Run a subprocess and return stdout.
+
+        Inspired by dask-jobqueue's _call(), with added timeout support.
+        """
+        full_env = None
+        if env:
+            full_env = {**os.environ, **env}
+
+        if shell:
+            proc = await asyncio.create_subprocess_shell(
+                cmd if isinstance(cmd, str) else " ".join(cmd),
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                env=full_env,
+            )
+        else:
+            proc = await asyncio.create_subprocess_exec(
+                *cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                env=full_env,
+                stdin=asyncio.subprocess.PIPE if stdin_data else None,
+            )
+
+        try:
+            stdout, stderr = await asyncio.wait_for(
+                proc.communicate(stdin_data.encode() if stdin_data else None),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            proc.kill()
+            raise CommandTimeoutError(
+                f"Command timed out after {timeout}s: {cmd}"
+            )
+
+        out = stdout.decode().strip()
+        err = stderr.decode().strip()
+
+        if proc.returncode != 0:
+            raise CommandFailedError(
+                f"Command failed (exit {proc.returncode}): {cmd}\nstderr: {err}"
+            )
+
+        return out
+
+    # --- Properties ---
+
+    def remove_job(self, job_id: str) -> None:
+        """Remove a job from tracking."""
+        self._jobs.pop(job_id, None)
+
+    @property
+    def jobs(self) -> dict[str, JobRecord]:
+        """All tracked jobs."""
+        return dict(self._jobs)
+
+    @property
+    def active_jobs(self) -> dict[str, JobRecord]:
+        """Non-terminal jobs."""
+        return {jid: r for jid, r in self._jobs.items() if not r.is_terminal}