Source code for matmelt.backends.local

"""The ``local`` backend: run a ``JobSpec`` as a subprocess on this machine.

The zero-setup default -- and the reference implementation of the durability
contract. Registered under ``kind = "local"`` at import time.

Execution model (the *why* behind the shape):

* ``submit`` is **non-blocking**: it stages inputs into a fresh per-job
  directory, launches the command through a tiny ``job.sh`` wrapper, and returns
  a ``Handle`` immediately. Progress is observed via ``poll``; results via
  ``fetch``.
* **The ``.rc`` sentinel is the durability mechanism.** The wrapper runs the
  command, then atomically writes its exit code to ``job.rc``
  (tmp-file + ``mv``). Everything ``poll``/``fetch`` need afterwards lives on
  disk (``job.rc``, ``stdout.log``/``stderr.log``, a ``job.json`` manifest), so
  a **restarted tool resolves the same handle with no in-memory state**: rc file
  present -> terminal with a faithful code; rc absent + process alive ->
  ``RUNNING``; rc absent + process gone -> ``FAILED`` with ``exit_code=None``
  ("ran, code unknown" -- never defaulted to 0).
* **PID reuse is defeated by a start-time check.** The handle records the
  wrapper's PID *and* its ``/proc/<pid>/stat`` start time; a live PID only
  counts as *our* job if the start time still matches (best-effort: on hosts
  without ``/proc`` the check degrades to PID-aliveness).
* **Two failure tiers.** A malformed spec raises ``BackendError`` (programming
  error, nothing launched). A job-level problem is never raised: a missing input
  artifact becomes a durable already-failed job (detail in ``stderr.log``, no
  process), and a command-not-found simply exits 127 through the wrapper --
  one bad item fails while the backend stays healthy.

Stdlib-only (``subprocess``/``pathlib``/``shutil``/``shlex``/``tempfile``).
No compute detail is hardcoded: ``work_dir``/``gpu``/``max_parallel`` come from
the backend's configured options.
"""

from __future__ import annotations

import contextlib
import json
import logging
import os
import shlex
import shutil
import signal
import subprocess
import tempfile
import uuid
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any

from matmelt.backends.base import ComputeBackend
from matmelt.backends.registry import register
from matmelt.core.cpusample import (
    LocalCpuTracker,
    gpu_percent_for_roots,
    parse_pmon,
    scan_proc,
)
from matmelt.core.jobspec import (
    BackendError,
    Capabilities,
    Handle,
    HealthStatus,
    JobResult,
    JobSpec,
    JobStatus,
)

_log = logging.getLogger("matmelt.backends.local")

_STDOUT = "stdout.log"
_STDERR = "stderr.log"
_RC = "job.rc"
_SCRIPT = "job.sh"
_MANIFEST = "job.json"


[docs] class LocalBackend(ComputeBackend): """Run jobs as child processes on this machine. The default backend, and the one that needs no setup: a bare ``{kind: local}`` works. Parameters ---------- name The backend's configured name, carried into every ``Handle`` it issues. options The config block. All keys are optional: ``work_dir`` Base directory for per-job directories. Defaults to a private temp directory -- zero setup, but those job directories die with the OS temp cleaner, so configure a persistent path to keep in-flight jobs resumable across restarts. ``gpu`` Whether this host exposes a GPU to jobs. Default ``False``, and config-driven rather than auto-detected, so capabilities stay predictable. ``max_parallel`` How many jobs this backend advertises it can run at once. Default ``1``; advertised here, enforced by the orchestrator. Raises ------ BackendError For an option of the wrong type or an unusable ``work_dir``. Notes ----- Durability is what distinguishes this from a plain ``Popen`` wrapper. Each job gets a directory holding a wrapper script, an atomically-written exit-code sentinel, both captured streams, and a manifest of what the job promised to produce. ``poll`` and ``fetch`` read *only* those, so a restarted tool resolves jobs it never launched -- no in-memory map is required, and an interrupted run resumes rather than orphaning its work. """ def __init__(self, name: str, options: dict[str, Any]) -> None: super().__init__(name, options) self._gpu = self._opt_bool("gpu", default=False) self._max_parallel = self._opt_positive_int("max_parallel", default=1) self._job_root = self._resolve_job_root(self.options.get("work_dir")) # PID -> Popen, kept ONLY to reap children (avoid zombies) and to # terminate jobs we launched. Never required by poll/fetch -- the # durable truth lives in the job dir. self._owned: dict[str, subprocess.Popen[bytes]] = {} # Delta state for the monitor's CPU hook (sample_cpu_percents); the # instance lives as long as the monitor process, which is the seam # the tracker's between-calls baseline needs. self._cpu_tracker = LocalCpuTracker() # -- option parsing (sanity-check the config boundary) ---------------------- def _opt_bool(self, key: str, *, default: bool) -> bool: value = self.options.get(key, default) if not isinstance(value, bool): raise BackendError(f"local backend option {key!r} must be a bool, got {value!r}") return value def _opt_positive_int(self, key: str, *, default: int) -> int: value = self.options.get(key, default) # bool subclasses int -- reject explicitly so a typo never becomes 0/1. if isinstance(value, bool) or not isinstance(value, int) or value < 1: raise BackendError( f"local backend option {key!r} must be a positive int, got {value!r}" ) return value def _resolve_job_root(self, work_dir: object) -> Path: if work_dir is None: root = Path(tempfile.mkdtemp(prefix="matmelt-local-")) elif isinstance(work_dir, str): root = Path(work_dir).expanduser() else: raise BackendError( f"local backend option 'work_dir' must be a string path, got {work_dir!r}" ) try: root.mkdir(parents=True, exist_ok=True) except OSError as exc: raise BackendError(f"local backend work_dir {root} is not usable: {exc}") from exc return root.resolve() # -- ComputeBackend interface ------------------------------------------------
[docs] def capabilities(self) -> Capabilities: return Capabilities(cpu=True, gpu=self._gpu, slurm=False, max_parallel=self._max_parallel)
[docs] def health_check(self) -> HealthStatus: """Probe that the job root is writable. Never raises. Returns ------- HealthStatus Whether a file can be created and removed under ``work_dir``, with the reason when it cannot. """ probe = self._job_root / ".healthcheck" try: probe.write_text("ok", encoding="utf-8") probe.unlink() except OSError as exc: return HealthStatus(ok=False, detail=f"work_dir {self._job_root} not writable: {exc}") return HealthStatus(ok=True, detail=f"work_dir {self._job_root} writable")
[docs] def sample_cpu_percents(self, handles: Sequence[Mapping[str, Any]]) -> list[float | None]: """Sample CPU usage per job -- the monitor's duck-typed CPU hook. Parameters ---------- handles The ledger's serialized handles, in order. The wrapper PID lives at ``extra.pid``. Returns ------- list One reading per handle, in the same order, as a percentage of one core. Readings come from the shared ``LocalCpuTracker`` -- ``/proc`` tick deltas over each job's whole process subtree -- so the first call after startup answers ``None`` for every job, that being the delta baseline. A handle without an integer PID, or a PID no longer running, is also ``None``. Notes ----- Never raises: this is cosmetic telemetry, and its contract is that a broken reading can never take the dashboard down. """ pids = self._handle_pids(handles) readings = self._cpu_tracker.sample([pid for pid in pids if pid is not None]) return [readings.get(pid) if pid is not None else None for pid in pids]
[docs] def sample_gpu_percents(self, handles: Sequence[Mapping[str, Any]]) -> list[float | None]: """Sample GPU usage per job -- the monitor's duck-typed GPU hook. Parameters ---------- handles The ledger's serialized handles, in order. Returns ------- list One SM-utilisation reading per handle, summed over that job's process subtree. A machine without ``nvidia-smi``, or one where it fails or times out, answers all ``None`` -- unknown, never zero. A live job with no pmon entry is an honest ``0.0``: running, but not on the GPU. Notes ----- One ``nvidia-smi pmon -c 1 -s u`` call answers the whole batch. Never raises. """ pids = self._handle_pids(handles) if not any(pid is not None for pid in pids): return [None] * len(pids) pmon_text = self._read_pmon() if pmon_text is None: return [None] * len(pids) return gpu_percent_for_roots(scan_proc(), parse_pmon(pmon_text), pids)
@staticmethod def _handle_pids(handles: Sequence[Mapping[str, Any]]) -> list[int | None]: """Read the wrapper PID out of each serialized ledger handle. Parameters ---------- handles Serialized handles. Returns ------- list ``extra.pid`` per handle, or ``None`` where it is absent or not an integer. """ pids: list[int | None] = [] for handle in handles: pid = dict(handle.get("extra") or {}).get("pid") pids.append(pid if isinstance(pid, int) else None) return pids @staticmethod def _read_pmon() -> str | None: """Read raw ``nvidia-smi pmon`` output. The seam the GPU hook -- and its tests -- go through. Returns ------- str or None The command's stdout, or ``None`` when the binary is absent, exits non-zero, or hangs past its timeout. """ if shutil.which("nvidia-smi") is None: return None try: proc = subprocess.run( ["nvidia-smi", "pmon", "-c", "1", "-s", "u"], capture_output=True, text=True, timeout=5.0, ) except (OSError, subprocess.TimeoutExpired): return None return proc.stdout if proc.returncode == 0 else None
[docs] def submit(self, spec: JobSpec) -> Handle: """Start ``spec`` through a shell wrapper and return a durable handle. Parameters ---------- spec The job to run. Its ``cwd`` is the material's directory, so outputs land beside the science data; ``scratch_inputs`` stage into this backend's private job directory instead, and the matching command token is rewritten to point there -- which is what keeps runner scripts out of material directories. Returns ------- Handle Carrying the job directory, the wrapper PID, and that PID's ``/proc`` start time. Raises ------ BackendError Only for a structurally malformed spec, which is a programming error. A *job-level* problem -- an unstageable input, for instance -- becomes a durable already-failed job instead, whose handle still resolves: one bad item must not look like a sick backend. Notes ----- The wrapper writes the exit code to a sentinel file via tmp-then-rename after the command returns. That sentinel, not the process, is what ``poll`` and ``fetch`` believe -- and recording the PID start time alongside the PID is what makes re-attachment safe against PID reuse. """ self._validate_spec(spec) job_id = uuid.uuid4().hex job_dir = self._job_root / job_id job_dir.mkdir(parents=True, exist_ok=True) cwd = spec.cwd if spec.cwd is not None else str(job_dir) extra: dict[str, Any] = {"job_dir": str(job_dir), "pid": None, "pid_starttime": None} handle = Handle(backend_name=self.name, job_id=job_id, extra=extra) self._write_manifest(job_dir, spec, cwd) # Split declared inputs: the item's science *data* stages into the # effective cwd (the material dir), while *scratch* inputs -- the job's # own self-contained runner script -- stage into this backend's private # job dir, so the material dir never accumulates our implementation files. data_inputs = { name: src for name, src in spec.input_artifacts.items() if name not in spec.scratch_inputs } scratch_inputs = { name: src for name, src in spec.input_artifacts.items() if name in spec.scratch_inputs } try: Path(cwd).mkdir(parents=True, exist_ok=True) self._stage_inputs(Path(cwd), data_inputs) self._stage_inputs(job_dir, scratch_inputs) except OSError as exc: # Job-level failure: durable, never raised (one bad item, healthy backend). detail = f"failed to stage inputs: {exc}" (job_dir / _STDERR).write_text(detail, encoding="utf-8") _log.warning("local submit failed for job %s: %s", job_id, detail) return handle # A scratch artifact referenced by bare name in the command is rewritten # to its staged path under the job dir (the cwd stays the material dir). command = self._rewrite_scratch_command(spec.command, job_dir, scratch_inputs) script = job_dir / _SCRIPT rc_path = job_dir / _RC script.write_text( "#!/bin/sh\n" f"{shlex.join(command)}\n" "rc=$?\n" f"printf '%s' \"$rc\" > {shlex.quote(str(rc_path) + '.tmp')} && " f"mv {shlex.quote(str(rc_path) + '.tmp')} {shlex.quote(str(rc_path))}\n" 'exit "$rc"\n', encoding="utf-8", ) env = os.environ.copy() env.update(spec.env) with ( (job_dir / _STDOUT).open("wb") as stdout_fh, (job_dir / _STDERR).open("wb") as stderr_fh, ): # argv from the stage, run via the wrapper script -- never a shell string popen = subprocess.Popen( ["/bin/sh", str(script)], cwd=cwd, env=env, stdout=stdout_fh, stderr=stderr_fh, ) extra["pid"] = popen.pid extra["pid_starttime"] = self._proc_starttime(popen.pid) self._owned[job_id] = popen _log.info("local submit job %s pid %d: %s", job_id, popen.pid, spec.command) return handle
[docs] def poll(self, handle: Handle) -> JobStatus: """Report status from durable evidence alone. Parameters ---------- handle The job to poll, possibly reloaded from the ledger in a later process. Returns ------- JobStatus ``DONE`` or ``FAILED`` when the exit-code sentinel exists; ``RUNNING`` when the wrapper process is alive **and still ours**; ``FAILED`` otherwise -- no exit code and no live process means the job never started, died before writing the sentinel, or could not write it. Raises ------ BackendError When the handle's job directory does not exist: unknown, or already released. Notes ----- Works identically for a job this instance launched and one re-attached after a restart, because it consults the sentinel and ``/proc`` rather than any in-memory state. """ job_dir = self._job_dir(handle) self._reap(handle.job_id) rc = self._read_rc(job_dir) if rc is not None: return JobStatus.DONE if rc == 0 else JobStatus.FAILED if self._job_alive(handle): return JobStatus.RUNNING # No exit code and no live process: never started, died before the # sentinel, or the sentinel itself was unwritable. return JobStatus.FAILED
[docs] def fetch(self, handle: Handle) -> JobResult: """Collect the faithful result from the job directory. Restart-safe. Parameters ---------- handle The job to collect. Returns ------- JobResult The sentinel's exit code -- ``None`` when absent or garbled, which means "ran, code unknown" and is never defaulted to ``0`` -- both captured streams, and the declared outputs that were actually produced, resolved against the job's recorded cwd. Raises ------ BackendError When the job directory is gone, or when the job is still running and was launched by a *previous* process, which cannot be waited on. Poll until terminal instead. Notes ----- A still-running job this instance owns is waited for, so ``fetch`` never returns a half-finished result. """ job_dir = self._job_dir(handle) owned = self._owned.get(handle.job_id) if owned is not None: owned.wait() elif self._read_rc(job_dir) is None and self._job_alive(handle): raise BackendError( f"job {handle.job_id!r} is still running; poll until terminal before fetch" ) exit_code = self._read_rc(job_dir) result = JobResult( exit_code=exit_code, stdout=self._read_text(job_dir / _STDOUT), stderr=self._read_text(job_dir / _STDERR), output_artifacts=self._collect_outputs(job_dir), ) _log.info( "local fetch job %s exit_code=%s outputs=%d", handle.job_id, exit_code, len(result.output_artifacts), ) return result
[docs] def release(self, handle: Handle) -> None: """Stop the job if it is ours, then remove its directory. Parameters ---------- handle The job to clean up. Releasing an unknown or already-released handle is a no-op, so restart-time cleanup is always safe. Notes ----- Idempotent and never raises. A process this instance launched is SIGTERMed then SIGKILLed; a re-attached one is killed only after the start-time check has proven the PID is still our wrapper and not a recycled stranger. The job directory holds this backend's scratch, including the staged runner script, so removing it is what keeps our implementation files out of the material directory -- which only ever received the item's science data and outputs. """ popen = self._owned.pop(handle.job_id, None) if popen is not None and popen.poll() is None: self._terminate(popen) elif popen is None and self._job_alive(handle): # Re-attached live job: the start-time check inside _job_alive just # proved the PID is still our wrapper, so it is safe to stop it. self._kill_pid(int(handle.extra["pid"])) job_dir = handle.extra.get("job_dir") if isinstance(job_dir, str): shutil.rmtree(job_dir, ignore_errors=True)
# -- submit helpers ----------------------------------------------------------- @staticmethod def _validate_spec(spec: JobSpec) -> None: """Reject a structurally malformed spec up front. Parameters ---------- spec The spec to check. Raises ------ BackendError When ``command`` is not a non-empty list of strings, ``env`` does not map strings to strings, or ``cwd`` is neither a string nor ``None``. These are programming errors, not job failures, which is why they raise rather than becoming a failed job. """ command = spec.command if not isinstance(command, list) or not command: raise BackendError( f"JobSpec.command must be a non-empty list of strings, got {command!r}" ) if not all(isinstance(arg, str) for arg in command): raise BackendError( f"JobSpec.command must be a non-empty list of strings, got {command!r}" ) if not all(isinstance(k, str) and isinstance(v, str) for k, v in spec.env.items()): raise BackendError("JobSpec.env must map str -> str") if spec.cwd is not None and not isinstance(spec.cwd, str): raise BackendError(f"JobSpec.cwd must be a string path or None, got {spec.cwd!r}") @staticmethod def _write_manifest(job_dir: Path, spec: JobSpec, cwd: str) -> None: """Persist what ``fetch`` needs to work after a restart. Parameters ---------- job_dir This backend's private directory for the job. spec The job being started. cwd The effective working directory, against which output names resolve. Notes ----- Without this file a restarted process would know the job exists but not what it promised to produce. Writing it is what lets ``fetch`` collect outputs for a job it never launched. """ manifest = { "output_artifacts": list(spec.output_artifacts), "cwd": cwd, } (job_dir / _MANIFEST).write_text(json.dumps(manifest), encoding="utf-8") @staticmethod def _stage_inputs(dest_dir: Path, input_artifacts: dict[str, str]) -> None: """Make each named input available inside ``dest_dir``. Parameters ---------- dest_dir Where the inputs must appear: the effective cwd for data inputs, the private job directory for scratch inputs. input_artifacts Name to source path. The *name* is where the file appears, relative to wherever the command reads it from. Raises ------ OSError Including ``FileNotFoundError``, for a missing or unstageable source. ``submit`` turns that into a durable job-level failure rather than a raised error. Notes ----- Symlink first, copy as a fallback. Idempotent: an existing destination, as a retry reusing the same directory produces, is replaced -- so restaging can never fail on leftovers, nor silently write *through* a stale symlink into the source. One case is deliberately skipped: an input that already resolves to its own destination. A data file living in the job's cwd is exactly that (every input must be declared, for remote staging), and replacing it would delete the source it links to. """ for name, src_str in input_artifacts.items(): src = Path(src_str) if not src.exists(): raise FileNotFoundError(f"input artifact {name!r} not found at {src_str}") dest = dest_dir / name if dest.exists() and dest.resolve() == src.resolve(): continue # already in place; unlink-and-relink would destroy the source dest.parent.mkdir(parents=True, exist_ok=True) if dest.is_symlink() or dest.is_file(): dest.unlink() elif dest.is_dir(): shutil.rmtree(dest) try: dest.symlink_to(src.resolve()) except OSError: if src.is_dir(): shutil.copytree(src, dest) else: shutil.copy2(src, dest) @staticmethod def _rewrite_scratch_command( command: list[str], job_dir: Path, scratch_inputs: dict[str, str] ) -> list[str]: """Point bare scratch-artifact references at their staged path. Parameters ---------- command The stage's argv. job_dir Where scratch inputs were staged. scratch_inputs The scratch artifacts, by name. Returns ------- list of str The command with exactly those tokens rewritten to their absolute job-directory path, so the job finds its runner while running with the material directory as cwd. A no-op when the spec declares no scratch inputs. """ if not scratch_inputs: return list(command) return [str(job_dir / arg) if arg in scratch_inputs else arg for arg in command] # -- durable-evidence helpers (poll/fetch work from these alone) --------------- def _job_dir(self, handle: Handle) -> Path: """Resolve the handle's job directory. Parameters ---------- handle The job to locate. Returns ------- pathlib.Path The job directory. Raises ------ BackendError When the handle belongs to a different backend, carries no job directory, or names one that no longer exists -- released, most likely. """ job_dir_str = handle.extra.get("job_dir") if handle.backend_name != self.name or not isinstance(job_dir_str, str): raise BackendError(f"unknown job handle {handle.job_id!r} for backend {self.name!r}") job_dir = Path(job_dir_str) if not job_dir.is_dir(): raise BackendError( f"job {handle.job_id!r}: job dir {job_dir} no longer exists (released?)" ) return job_dir @staticmethod def _read_rc(job_dir: Path) -> int | None: """Read the exit-code sentinel the wrapper wrote. Parameters ---------- job_dir The job's directory. Returns ------- int or None The exit code, or ``None`` when the sentinel is absent or garbled. ``None`` means "ran, code unknown" and must never be read as success. """ try: text = (job_dir / _RC).read_text(encoding="utf-8").strip() except OSError: return None try: return int(text) except ValueError: return None def _job_alive(self, handle: Handle) -> bool: """Say whether the wrapper process is still running, and still ours. Parameters ---------- handle The job to check. Returns ------- bool ``True`` only when the recorded PID is alive **and** its ``/proc`` start time still matches the one recorded at submission. That second test is what defeats PID reuse: a recycled PID belonging to an unrelated process must never be mistaken for our job -- nor killed by ``release``. On a host without ``/proc`` the check degrades to plain PID-aliveness. """ pid = handle.extra.get("pid") if not isinstance(pid, int): return False # never started if not self._pid_alive(pid): return False recorded = handle.extra.get("pid_starttime") if recorded is None: return True # no /proc on this host: degrade to PID-aliveness return self._proc_starttime(pid) == str(recorded) @staticmethod def _pid_alive(pid: int) -> bool: """Say whether a PID currently exists. Parameters ---------- pid The process id to probe with signal 0. Returns ------- bool ``True`` when the process exists, including when it belongs to another user -- a ``PermissionError`` proves existence just as well as success does. """ try: os.kill(pid, 0) except ProcessLookupError: return False except PermissionError: return True return True @staticmethod def _proc_starttime(pid: int) -> str | None: """Read a process's start time -- its identity across PID reuse. Parameters ---------- pid The process id. Returns ------- str or None Field 22 of ``/proc/<pid>/stat``, which is constant for the life of the process, or ``None`` where ``/proc`` is unavailable. """ try: stat = Path(f"/proc/{pid}/stat").read_text(encoding="ascii") except OSError: return None # comm (field 2) may contain spaces/parens; split after its closing ')'. fields = stat.rsplit(")", 1)[-1].split() return fields[19] if len(fields) > 19 else None def _reap(self, job_id: str) -> None: """Collect an owned child's status if it has exited, avoiding zombies. Parameters ---------- job_id Which owned job to reap. Unknown or still-running ids are no-ops. """ popen = self._owned.get(job_id) if popen is not None and popen.poll() is not None: del self._owned[job_id] @staticmethod def _read_text(path: Path) -> str: """Read a captured stream file, best-effort. Parameters ---------- path The stream file. Returns ------- str Its contents, with undecodable bytes replaced, or ``""`` when it cannot be read. Never raises: a missing stream must not cost a result its exit code. """ try: return path.read_text(encoding="utf-8", errors="replace") except OSError: return "" @staticmethod def _collect_outputs(job_dir: Path) -> dict[str, str]: """Collect the declared outputs that were actually produced. Parameters ---------- job_dir The job's directory, holding the manifest written at submission. Returns ------- dict Name to path for each declared output that exists, resolved against the job's recorded cwd. A missing output is simply absent, for the stage to judge; an unreadable manifest yields nothing rather than raising. """ try: manifest = json.loads((job_dir / _MANIFEST).read_text(encoding="utf-8")) except (OSError, ValueError): return {} base = Path(manifest.get("cwd", str(job_dir))) collected: dict[str, str] = {} for name in manifest.get("output_artifacts", []): path = base / name if path.exists(): collected[name] = str(path) return collected # -- termination --------------------------------------------------------------- @staticmethod def _terminate(popen: subprocess.Popen[bytes]) -> None: """Stop an owned process, best-effort. Parameters ---------- popen The child to stop: SIGTERM first, then SIGKILL if it outlasts the grace period. Every failure is swallowed -- cleanup must never raise. """ try: popen.terminate() except OSError: return try: popen.wait(timeout=5) except subprocess.TimeoutExpired: try: popen.kill() popen.wait(timeout=5) except (OSError, subprocess.TimeoutExpired): pass @staticmethod def _kill_pid(pid: int) -> None: """Send SIGTERM to a re-attached wrapper, best-effort. Parameters ---------- pid The process to signal. The caller must already have verified the PID's start time, so this can never kill a recycled stranger. """ with contextlib.suppress(OSError): os.kill(pid, signal.SIGTERM)
# Register at import time so create_backend can resolve kind="local" once this # module is imported (the discovery/wiring path imports the backends it offers). register("local", LocalBackend)