"""The ``ssh_exec`` backend: run a ``JobSpec`` on a remote host over SSH (4.1).
Modelled on MVP_V0's ``ssh_exec`` (its durable-native-id + ``.rc``-sentinel +
single-SSH-seam shape is proven) but with the **data model inverted** per the
2026-07-08 decisions: MVP_V0 assumes a shared filesystem and stages nothing;
here **declarative staging is the default** and ``shared_filesystem: true`` is
the opt-out optimization.
* **Remotes never receive our package code.** A remote is a pure calculation
executor: the only things there are the user's environment (``remote_python``
with the science deps) and what we push per job. Every job's declared
``input_artifacts`` -- which, by the Phase-3 stage discipline, always include
the stage's self-contained runner script -- are staged into a fresh per-job
remote directory; the command runs there (artifact *names* resolve in the
cwd, same as ``local``); declared ``output_artifacts`` are pulled back into
the job's local ``cwd`` (the item dir), where the stage judges them. Our
runners are stdlib+science only, so there is no PYTHONPATH plumbing at all.
* **The staging transport is a seam.** ``_stage_in`` / ``_stage_out`` wrap
rsync-over-SSH today; an S3/Globus transport plugs in behind the same two
calls with zero stage changes. With ``shared_filesystem: true`` both become
no-ops: the job runs in ``spec.cwd`` (the same absolute path on both sides)
and outputs are collected locally, exactly like ``local``.
* **A durable native id.** ``Handle.job_id`` is the remote PID (or the SLURM
job id under ``use_slurm``); everything ``poll``/``fetch``/``release`` need
(scheme, remote job dir, log/sentinel paths, output names, local cwd) rides
in ``Handle.extra`` as plain JSON-safe strings -- no in-memory job map, so a
restarted orchestrator re-attaches from the persisted handle alone.
* **Exit-code fidelity via the ``.rc`` sentinel.** The remote wrapper appends
``echo $? > <rc>``; ``fetch`` reads it back. A missing/garbled sentinel is
``exit_code=None`` ("ran, code unknown") -- never a silent ``0``. The SLURM
path prefers ``sacct``'s ExitCode and falls back to the sentinel.
Stdlib + the system ``ssh``/``rsync`` (and remote ``sbatch``/``sacct``) only.
Key-based auth must already work: every SSH runs ``BatchMode=yes`` and never
prompts. All compute details come from options -- no host, user, path, or
queue literal in code, and ``health_check`` never leaks a secret.
"""
from __future__ import annotations
import contextlib
import shlex
import subprocess
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 gpu_percent_for_roots, parse_pmon, subtree_total
from matmelt.core.jobspec import (
BackendError,
Capabilities,
Handle,
HealthStatus,
JobResult,
JobSpec,
JobStatus,
)
#: rsync's "some files vanished / partial transfer" family: for output
#: collection this means "that output was not produced" -- a normal per-file
#: absence handled by omission, never an infrastructure fault.
_RSYNC_PARTIAL_CODES = (23, 24)
[docs]
class SshExecBackend(ComputeBackend):
"""Run jobs on a remote host over SSH, staging files per job by default.
This is the backend that makes "run the screening stage on my GPU box" a
config change rather than a code change.
Parameters
----------
name
The backend's configured name, carried into every ``Handle`` it issues.
options
The config block. Only ``host`` is required:
``host``
An ``~/.ssh/config`` alias, a hostname, or ``user@host``.
``user``
SSH user, when ``host`` carries none.
``ssh_binary``, ``rsync_binary``
Client binaries, defaulting to ``ssh`` and ``rsync``.
``remote_python``
Interpreter to substitute for ``argv[0]`` remotely. Normally
unnecessary -- the stage's own ``python`` option already names the
remote interpreter -- but this override wins when both are set.
``remote_work_dir``
Root for per-job remote directories, default
``/tmp/matmelt-jobs``. Each job gets a fresh subdirectory,
removed by ``release``.
``shared_filesystem``
``true`` skips *data* staging, assuming the remote sees the job's
``cwd`` at the same absolute path over NFS or SSHFS. Scratch inputs
-- the runner script -- are still pushed into the per-job directory
and referenced by absolute path, because they are deliberately kept
out of the material directory. Default ``false``.
``use_slurm``
Submit through ``sbatch`` and track the SLURM job id instead of a
bare PID. ``partition`` sets the queue, and ``JobSpec.resources``
becomes ``--cpus-per-task`` and ``--time``. Default ``false``.
``gpu``, ``max_parallel``
Advertised in ``capabilities``.
Raises
------
BackendError
For a missing ``host`` or any option of the wrong type.
Notes
-----
**Remotes never receive MatMelt's code.** A remote is a pure calculation
executor: what travels is the job's self-contained runner script plus its
declared input files, and what comes back are the declared outputs. There
is no ``pip install`` on the remote, no ``PYTHONPATH``, and no assumption
that our package is importable there.
Declarative staging is therefore the default, and ``shared_filesystem`` is
an optimization that skips it. Both paths run through the same two seams,
``_stage_in`` and ``_stage_out``, so a future S3 or Globus transport needs
no stage changes at all.
Everything ``poll``, ``fetch`` and ``release`` need rides in
``Handle.extra`` as JSON-safe values -- there is no in-memory job map, so a
restarted run re-attaches to live remote jobs from the ledger alone. Every
token crossing the SSH boundary is quoted with ``shlex.quote``.
"""
def __init__(self, name: str, options: dict[str, Any]) -> None:
super().__init__(name, options)
opts = self.options
host = self._req_str(opts, "host")
user = self._opt_str(opts, "user")
self._target = f"{user}@{host}" if user and "@" not in host else host
self._ssh = self._opt_str(opts, "ssh_binary") or "ssh"
self._rsync = self._opt_str(opts, "rsync_binary") or "rsync"
self._remote_python = self._opt_str(opts, "remote_python")
# A documented, per-backend overridable default (`remote_work_dir`), not a temp
# file this process creates; the job dirs made below it are per-job and unique.
default_work_dir = "/tmp/matmelt-jobs" # nosec B108
self._work_dir = self._opt_str(opts, "remote_work_dir") or default_work_dir
self._shared_fs = self._opt_bool(opts, "shared_filesystem", default=False)
self._use_slurm = self._opt_bool(opts, "use_slurm", default=False)
self._partition = self._opt_str(opts, "partition")
self._gpu = self._opt_bool(opts, "gpu", default=False)
self._max_parallel = self._opt_int(opts, "max_parallel", default=1)
# -- option parsing (the config boundary; all errors are BackendError) ----
@staticmethod
def _req_str(opts: dict[str, Any], key: str) -> str:
value = opts.get(key)
if not isinstance(value, str) or not value.strip():
raise BackendError(
f"ssh_exec option {key!r} is required and must be a non-empty string"
)
return value
@staticmethod
def _opt_str(opts: dict[str, Any], key: str) -> str | None:
value = opts.get(key)
if value is None:
return None
if not isinstance(value, str) or not value.strip():
raise BackendError(f"ssh_exec option {key!r} must be a non-empty string, got {value!r}")
return value
@staticmethod
def _opt_bool(opts: dict[str, Any], key: str, *, default: bool) -> bool:
value = opts.get(key, default)
if not isinstance(value, bool):
raise BackendError(f"ssh_exec option {key!r} must be a bool, got {value!r}")
return value
@staticmethod
def _opt_int(opts: dict[str, Any], key: str, *, default: int) -> int:
value = opts.get(key, default)
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise BackendError(f"ssh_exec option {key!r} must be a positive int, got {value!r}")
return value
# -- transport seams (the ONLY places ssh/rsync are invoked) --------------
def _ssh_argv(self, remote_cmd: str) -> list[str]:
"""Build the ssh argv for one remote command.
Parameters
----------
remote_cmd
The command string to run on the remote.
Returns
-------
list of str
The argv, always carrying ``BatchMode=yes`` so a host that wants a
password fails fast instead of hanging a run on a prompt nobody
will ever see. Key auth must already work; when it does not,
``health_check`` reports the backend down.
"""
return [self._ssh, "-o", "BatchMode=yes", self._target, remote_cmd]
def _run_ssh(
self, remote_cmd: str, *, timeout: float = 30.0
) -> subprocess.CompletedProcess[str]:
"""Run one command on the remote. The single SSH invocation point.
Parameters
----------
remote_cmd
The command string.
timeout
Seconds to allow.
Returns
-------
subprocess.CompletedProcess
The result. The remote command's own **non-zero exit is returned
faithfully**, never raised -- that is the two-failure-tier contract.
Raises
------
BackendError
Only for *transport* failure: a missing ssh binary, an unreachable
host, a timeout.
Notes
-----
Tests replace this seam and never touch a real host.
"""
try:
return subprocess.run(
self._ssh_argv(remote_cmd), capture_output=True, text=True, timeout=timeout
)
except (OSError, subprocess.TimeoutExpired) as exc:
raise BackendError(f"ssh to {self._target} failed: {exc}") from exc
def _run_rsync(self, src: str, dest: str, *, timeout: float = 600.0) -> int:
"""Transfer one path over rsync. The single rsync invocation point.
Parameters
----------
src, dest
rsync endpoints, either of which may carry a ``target:`` prefix.
timeout
Seconds to allow. Generous by default: a staged input can be large.
Returns
-------
int
rsync's own exit code, for the caller to judge. A partial-transfer
code means "the file was not there", which on collection is normal
-- an output the job never produced.
Raises
------
BackendError
For transport failure (a missing binary, a timeout) or any rsync
exit outside the expected set.
"""
argv = [
self._rsync,
"-az",
"-e",
f"{self._ssh} -o BatchMode=yes",
src,
dest,
]
try:
proc = subprocess.run(argv, capture_output=True, text=True, timeout=timeout)
except (OSError, subprocess.TimeoutExpired) as exc:
raise BackendError(f"rsync via {self._target} failed: {exc}") from exc
if proc.returncode not in (0, *_RSYNC_PARTIAL_CODES):
detail = (proc.stderr or "").strip()[-300:]
raise BackendError(f"rsync {src} -> {dest} failed (exit {proc.returncode}): {detail}")
return proc.returncode
# -- staging (declarative by default; no-ops under shared_filesystem) -----
def _stage_in(self, spec: JobSpec, job_dir: str, *, only: frozenset[str] | None = None) -> None:
"""Push declared inputs into the fresh remote job directory.
Parameters
----------
spec
The job whose ``input_artifacts`` are pushed. Artifact *names*
become paths relative to the remote cwd, mirroring how ``local``
stages into the effective cwd, so a command may reference a staged
file by bare name either way.
job_dir
The remote per-job directory.
only
Restrict staging to this subset of artifact names. Used under
``shared_filesystem`` to push *only* the scratch inputs into the
per-job directory while data inputs stay on the shared mount.
``None`` -- the staged-model default -- pushes every declared
input.
Raises
------
BackendError
When a local source is missing, or rsync could not place it.
Nothing has run yet at this point, and launching a job known to
lack its inputs would only produce a confusing remote failure.
"""
for artifact_name, src_str in spec.input_artifacts.items():
if only is not None and artifact_name not in only:
continue
src = Path(src_str)
if not src.exists():
raise BackendError(f"input artifact {artifact_name!r} not found at {src_str}")
source = f"{src}/" if src.is_dir() else str(src)
code = self._run_rsync(source, f"{self._target}:{job_dir}/{artifact_name}")
if code != 0:
raise BackendError(
f"staging input {artifact_name!r} to {self._target} failed (rsync exit {code})"
)
def _stage_out(self, handle: Handle) -> dict[str, str]:
"""Pull each declared output from the remote job directory into the local cwd.
Parameters
----------
handle
The finished job, whose ``extra`` carries the remote job directory,
the local cwd, and the declared output names.
Returns
-------
dict
Name to local path for each output that actually came back. An
output the job never produced is simply omitted -- rsync's
partial-transfer exit -- matching ``local``'s fetch semantics. The
stage's ``validate_outputs`` is the judge of what was promised.
Notes
-----
Each output is pulled **basename-preserving into its parent
directory**, so a file lands at ``<cwd>/<name>`` and a directory tree
at ``<cwd>/<name>/...``. That is what lets a fan-out stage bring its
whole ``children/`` subtree home, and it means this method never needs
to know whether the remote artifact is a file or a directory.
"""
job_dir = str(handle.extra.get("job_dir", ""))
local_cwd = Path(str(handle.extra.get("local_cwd", ".")))
collected: dict[str, str] = {}
for artifact_name in handle.extra.get("outputs", []):
dest = local_cwd / artifact_name
dest.parent.mkdir(parents=True, exist_ok=True)
code = self._run_rsync(f"{self._target}:{job_dir}/{artifact_name}", f"{dest.parent}/")
if code == 0 and dest.exists():
collected[str(artifact_name)] = str(dest)
return collected
def _collect_local(self, handle: Handle) -> dict[str, str]:
"""Collect outputs under ``shared_filesystem``, where they are already local.
Parameters
----------
handle
The finished job.
Returns
-------
dict
Name to local path for each declared output that exists. Existence
rather than file-ness is the test, so a fan-out stage's
``children/`` directory output is collected too.
"""
local_cwd = Path(str(handle.extra.get("local_cwd", ".")))
return {
str(name): str(local_cwd / name)
for name in handle.extra.get("outputs", [])
if (local_cwd / str(name)).exists()
}
# -- ComputeBackend interface ---------------------------------------------
[docs]
def capabilities(self) -> Capabilities:
return Capabilities(
cpu=True, gpu=self._gpu, slurm=self._use_slurm, max_parallel=self._max_parallel
)
[docs]
def health_check(self) -> HealthStatus:
"""Check that the remote is reachable and its work root writable. Never raises.
Returns
-------
HealthStatus
Down is a normal answer, reported with a short, **secret-free**
detail. The orchestrator pauses this backend's stages on a down
report and resumes them by itself once it recovers.
"""
probe = f"mkdir -p {shlex.quote(self._work_dir)} && test -w {shlex.quote(self._work_dir)}"
try:
proc = self._run_ssh(probe, timeout=15.0)
except BackendError as exc:
return HealthStatus(ok=False, detail=str(exc)[:200])
if proc.returncode != 0:
tail = (proc.stderr or "").strip()[:160]
return HealthStatus(
ok=False,
detail=f"{self._target}: work dir {self._work_dir} not usable "
f"(ssh exit {proc.returncode}) {tail}".strip(),
)
return HealthStatus(ok=True, detail=f"{self._target} reachable; {self._work_dir} writable")
# -- diagnosis hooks (duck-typed; consumed by core/preflight, not the ABC) --
@property
def probe_target(self) -> str:
"""Where this backend's remote probes run.
Returns
-------
str
The SSH target. ``doctor`` reports it as the *location* of a
readiness answer, and builds the ssh-prefixed install command from
it -- an answer is only as good as where it was obtained.
"""
return self._target
[docs]
def check_remote_import(self, python: str, module: str) -> tuple[bool, str]:
"""Ask whether the **remote** interpreter can import a module. Never raises.
Closes the preflight's blind spot for remotely-wired stages: the
science stack must exist in the interpreter that will actually run the
staged runner, not in ours.
Parameters
----------
python
The remote interpreter path.
module
The import name to try.
Returns
-------
tuple
``(satisfied, detail)``. Transport failure, a missing remote
interpreter, and a failing import all come back as
``(False, why)``.
Notes
-----
Duck-typed rather than part of the six-method ABC: ``core/preflight``
looks for this hook, so any future remote backend opts in simply by
providing it and ``probe_target``.
"""
probe = f"{shlex.quote(python)} -c {shlex.quote(f'import {module}')}"
try:
proc = self._run_ssh(probe, timeout=60.0)
except BackendError as exc:
return False, str(exc)[:200]
if proc.returncode == 0:
return True, f"{module} importable on {self._target}"
stderr_lines = (proc.stderr or "").strip().splitlines()
return False, stderr_lines[-1] if stderr_lines else f"import exited {proc.returncode}"
[docs]
def check_remote_file(self, path: str) -> tuple[bool, str]:
"""Ask whether a path exists as a regular file on the remote. Never raises.
The file-existence analogue of ``check_remote_import``. A stage option
naming a path -- the screen stage's ``model_path`` -- points at where
the stage *runs*, so stat-ing our own disk would report a false
problem and could refuse a perfectly valid run.
Parameters
----------
path
The remote path to test.
Returns
-------
tuple
``(present, detail)``. Transport failure and a missing file both
come back as ``(False, why)``.
"""
probe = f"test -f {shlex.quote(path)}"
try:
proc = self._run_ssh(probe, timeout=30.0)
except BackendError as exc:
return False, str(exc)[:200]
if proc.returncode == 0:
return True, f"{path} present on {self._target}"
return False, f"no such file on {self._target}"
[docs]
def sample_cpu_percents(self, handles: Sequence[Mapping[str, Any]]) -> list[float | None]:
"""Sample remote CPU usage per job -- the monitor's duck-typed CPU hook.
Parameters
----------
handles
The ledger's serialized handles, in order.
Returns
-------
list
One reading per handle, summed over that job's remote process
subtree. These are ``ps`` ``pcpu`` semantics -- a lifetime average
rather than an instantaneous rate, which is honest enough for
long-running science jobs. ``None`` for a SLURM-scheme handle,
where the scheduler owns the process and we hold only its job id;
for a PID no longer in the table; and for any transport failure.
Notes
-----
One batched ``ps`` round-trip answers every handle at once, and the
caller caches the result behind a TTL so a dashboard refresh does not
mean an SSH round-trip. Never raises.
"""
pids = self._handle_pids(handles)
if not any(pid is not None for pid in pids):
return [None] * len(pids)
try:
proc = self._run_ssh("ps -e -o pid=,ppid=,pcpu=", timeout=30.0)
except BackendError:
return [None] * len(pids)
if proc.returncode != 0:
return [None] * len(pids)
table = self._parse_ps_table(proc.stdout or "")
return [subtree_total(table, 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 remote 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
remote process subtree. A remote without ``nvidia-smi`` answers all
``None`` -- unknown, never zero -- as do SLURM-scheme handles and
any transport failure. A live job with no pmon entry is an honest
``0.0``.
Notes
-----
One combined round-trip fetches the process table and the pmon table
together, separated by a sentinel line. Never raises.
"""
pids = self._handle_pids(handles)
if not any(pid is not None for pid in pids):
return [None] * len(pids)
probe = "ps -e -o pid=,ppid=; echo ---; nvidia-smi pmon -c 1 -s u 2>/dev/null"
try:
proc = self._run_ssh(probe, timeout=30.0)
except BackendError:
return [None] * len(pids)
ps_part, sep, pmon_part = (proc.stdout or "").partition("---")
if not sep or not pmon_part.strip():
return [None] * len(pids) # no separator or no nvidia-smi on the remote
table: dict[int, tuple[int, float]] = {}
for line in ps_part.splitlines():
parts = line.split()
try:
table[int(parts[0])] = (int(parts[1]), 0.0)
except (ValueError, IndexError):
continue
return gpu_percent_for_roots(table, parse_pmon(pmon_part), pids)
@staticmethod
def _handle_pids(handles: Sequence[Mapping[str, Any]]) -> list[int | None]:
"""Read the remote PID out of each serialized handle.
Parameters
----------
handles
Serialized handles.
Returns
-------
list
The remote PID per handle, or ``None`` for a SLURM-scheme handle --
whose job id is a scheduler id, not a process -- and for anything
malformed.
"""
pids: list[int | None] = []
for handle in handles:
extra = dict(handle.get("extra") or {})
job_id = str(handle.get("job_id", ""))
is_pid = extra.get("scheme") == "pid" and job_id.isdigit()
pids.append(int(job_id) if is_pid else None)
return pids
@staticmethod
def _parse_ps_table(text: str) -> dict[int, tuple[int, float]]:
"""Parse a remote ``ps`` table into the shape the subtree helpers take.
Parameters
----------
text
Output of ``ps -e -o pid=,ppid=,pcpu=``.
Returns
-------
dict
``{pid: (ppid, pcpu)}``. Malformed lines are skipped rather than
raising.
"""
table: dict[int, tuple[int, float]] = {}
for line in text.splitlines():
parts = line.split()
if len(parts) != 3:
continue
try:
table[int(parts[0])] = (int(parts[1]), float(parts[2]))
except ValueError:
continue
return table
[docs]
def submit(self, spec: JobSpec) -> Handle:
"""Stage inputs, launch remotely, and return a durable handle. Non-blocking.
Parameters
----------
spec
The job to run.
Returns
-------
Handle
Carrying the remote job directory, the stream and sentinel paths,
the declared outputs, the local cwd, and the tracking scheme --
all JSON-safe, so the ledger alone is enough to resolve the job
after a restart.
Raises
------
BackendError
For a malformed spec, an unstageable input, an uncreatable remote
directory, a failed launch, or a launch that yielded no trackable
id.
Notes
-----
Under the **staged model** (the default) a fresh per-job directory is
created, every declared input is pushed into it, and the command runs
there -- so artifact names, the runner included, resolve in the cwd
exactly as they do under ``local``.
Under **shared_filesystem**, data inputs are not staged at all, since
the remote already sees them at ``spec.cwd``. The *scratch* inputs are
still pushed into the per-job directory, because the runner is
deliberately kept out of the material directory, and the matching
command tokens are rewritten to that absolute remote path.
Either way the launch wrapper appends an exit-code sentinel, and the
job is tracked by remote PID -- or by SLURM job id under ``use_slurm``,
where ``JobSpec.resources`` becomes scheduler directives.
"""
self._validate_spec(spec)
job_id = uuid.uuid4().hex[:12]
job_dir = f"{self._work_dir}/mcli-{job_id}"
remote_cwd = (spec.cwd or job_dir) if self._shared_fs else job_dir
out, err, rc = (f"{job_dir}/job.{suffix}" for suffix in ("out", "err", "rc"))
mkdir = self._run_ssh(f"mkdir -p {shlex.quote(job_dir)}", timeout=30.0)
if mkdir.returncode != 0:
raise BackendError(
f"cannot create {job_dir} on {self._target}: {(mkdir.stderr or '')[:200]}"
)
if not self._shared_fs:
# Staged model: everything (data + scratch) lands in the job dir,
# which is also the cwd -- bare names resolve, no rewrite needed.
self._stage_in(spec, job_dir)
scratch_dir: str | None = None
else:
# Shared model: the runner lives out of the material dir, so stage
# only the scratch inputs into the per-job dir and point the command
# at their absolute path; data inputs ride the shared mount.
if spec.scratch_inputs:
self._stage_in(spec, job_dir, only=spec.scratch_inputs)
scratch_dir = job_dir
inner = self._remote_command_string(spec, remote_cwd, scratch_dir=scratch_dir)
wrapped = f"{inner}; echo $? > {shlex.quote(rc)}"
extra: dict[str, Any] = {
"job_dir": job_dir,
"out": out,
"err": err,
"rc": rc,
"outputs": list(spec.output_artifacts),
"local_cwd": spec.cwd or ".",
"staged": not self._shared_fs,
}
if self._use_slurm:
directives = (
f"--job-name=mcli-{job_id} --output={shlex.quote(out)} --error={shlex.quote(err)}"
)
if self._partition:
directives += f" --partition={shlex.quote(self._partition)}"
if spec.resources.cores > 1:
directives += f" --cpus-per-task={spec.resources.cores}"
if spec.resources.walltime_s:
directives += f" --time={max(1, -(-spec.resources.walltime_s // 60))}"
launch = f"sbatch --parsable {directives} --wrap {shlex.quote(wrapped)}"
proc = self._run_ssh(launch, timeout=60.0)
if proc.returncode != 0:
raise BackendError(f"sbatch on {self._target} failed: {(proc.stderr or '')[:300]}")
slurm_id = (proc.stdout or "").strip().split(";")[0]
if not slurm_id.isdigit():
raise BackendError(f"sbatch on {self._target} returned no job id: {proc.stdout!r}")
return Handle(
backend_name=self.name, job_id=slurm_id, extra={"scheme": "slurm", **extra}
)
launch = (
f"nohup sh -c {shlex.quote(wrapped)} >{shlex.quote(out)} 2>{shlex.quote(err)} & echo $!"
)
proc = self._run_ssh(launch, timeout=60.0)
if proc.returncode != 0:
raise BackendError(
f"remote launch on {self._target} failed: {(proc.stderr or '')[:300]}"
)
pid = (proc.stdout or "").strip()
if not pid.isdigit():
raise BackendError(f"could not obtain a remote PID from {self._target}, got {pid!r}")
return Handle(backend_name=self.name, job_id=pid, extra={"scheme": "pid", **extra})
def _remote_command_string(
self, spec: JobSpec, remote_cwd: str, *, scratch_dir: str | None = None
) -> str:
"""Render the spec as one remote shell command: cd, env overlay, argv.
Parameters
----------
spec
The job.
remote_cwd
Where the command runs.
scratch_dir
Where scratch inputs were staged, under the shared-filesystem
model. ``None`` under the staged model, where bare names already
resolve in the cwd.
Returns
-------
str
The command string.
Notes
-----
Two independent rewrites can apply. A configured ``remote_python``
replaces ``argv[0]``, since the stage may have baked in a local
interpreter path. And with a ``scratch_dir``, any argv token naming a
scratch input is rewritten to its absolute remote path, because the cwd
is then the material directory rather than the job directory. They
coexist safely: an interpreter is never a scratch name.
**Every token is ``shlex.quote``d.** Nothing the stage or the config
wrote can splice shell across the SSH boundary.
"""
argv = list(spec.command)
if self._remote_python and argv:
argv[0] = self._remote_python
if scratch_dir is not None and spec.scratch_inputs:
argv = [f"{scratch_dir}/{arg}" if arg in spec.scratch_inputs else arg for arg in argv]
env_prefix = "".join(f"{k}={shlex.quote(v)} " for k, v in spec.env.items())
command = " ".join(shlex.quote(arg) for arg in argv)
return f"cd {shlex.quote(remote_cwd)} && {env_prefix}{command}"
@staticmethod
def _validate_spec(spec: JobSpec) -> None:
"""Reject a structurally malformed spec up front.
Parameters
----------
spec
The spec to check.
Raises
------
BackendError
For a bad ``command``, ``env`` or ``cwd``, and for an artifact name
that is not a plain relative path -- a name carrying ``..`` or a
leading slash would escape the per-job directory on staging.
"""
if (
not isinstance(spec.command, list)
or not spec.command
or not all(isinstance(arg, str) and arg for arg in spec.command)
):
raise BackendError(
f"JobSpec.command must be a non-empty list of non-empty strings, "
f"got {spec.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")
bad_names = [
name
for name in (*spec.input_artifacts, *spec.output_artifacts)
if name.startswith("/") or ".." in name
]
if bad_names:
raise BackendError(f"artifact names must be cwd-relative without '..', got {bad_names}")
[docs]
def poll(self, handle: Handle) -> JobStatus:
"""Map the remote job's state to a portable ``JobStatus``. Idempotent.
Parameters
----------
handle
The job to poll, possibly reloaded from the ledger in a later
process.
Returns
-------
JobStatus
Under the PID scheme, a live process is ``RUNNING`` and anything
else is ``DONE`` -- the true outcome is ``fetch``'s business, read
from the sentinel. Under SLURM, ``sacct`` states map
conservatively, and an **empty accounting line is ``RUNNING``**:
a job not yet visible to the accounting database must never be
guessed into a terminal state.
Raises
------
BackendError
For a handle this backend did not issue, or transport failure.
"""
scheme = handle.extra.get("scheme")
if handle.backend_name != self.name or scheme not in ("pid", "slurm"):
raise BackendError(f"handle {handle.job_id!r} is not a resolvable {self.name!r} job")
if scheme == "pid":
proc = self._run_ssh(
f"kill -0 {shlex.quote(handle.job_id)} 2>/dev/null && echo UP || echo DOWN"
)
return JobStatus.RUNNING if "UP" in (proc.stdout or "") else JobStatus.DONE
proc = self._run_ssh(
f"sacct -j {shlex.quote(handle.job_id)} --format=State --noheader --parsable2 | head -1"
)
state = (proc.stdout or "").strip().upper()
if state.startswith(("PENDING", "CONFIGURING")):
return JobStatus.PENDING
if state.startswith(("RUNNING", "COMPLETING")) or not state:
return JobStatus.RUNNING
if state.startswith("COMPLETED"):
return JobStatus.DONE
return JobStatus.FAILED
[docs]
def fetch(self, handle: Handle) -> JobResult:
"""Collect the finished job: streams, true exit code, pulled outputs.
Parameters
----------
handle
The job to collect.
Returns
-------
JobResult
The exit code comes from the sentinel, or from ``sacct`` first
under SLURM with the sentinel as fallback. A missing or garbled
sentinel is ``None`` -- "ran, code unknown", never a silent ``0``
-- and a non-zero code is reported faithfully rather than raised.
Outputs are pulled back through the staging seam, or collected
locally under ``shared_filesystem``; absent ones are omitted for
the stage to judge.
Raises
------
BackendError
For a handle this backend did not issue, or transport failure.
"""
scheme = handle.extra.get("scheme")
if handle.backend_name != self.name or scheme not in ("pid", "slurm"):
raise BackendError(f"handle {handle.job_id!r} is not a resolvable {self.name!r} job")
stdout = self._cat(str(handle.extra.get("out", "")))
stderr = self._cat(str(handle.extra.get("err", "")))
exit_code: int | None = None
if scheme == "slurm":
exit_code = self._slurm_exit_code(handle.job_id)
if exit_code is None:
exit_code = self._read_rc(str(handle.extra.get("rc", "")))
if handle.extra.get("staged"):
outputs = self._stage_out(handle)
else:
outputs = self._collect_local(handle)
return JobResult(
exit_code=exit_code, stdout=stdout, stderr=stderr, output_artifacts=outputs
)
[docs]
def release(self, handle: Handle) -> None:
"""Terminate a still-running job and remove its remote directory.
Parameters
----------
handle
The job to clean up: ``kill`` under the PID scheme, ``scancel``
under SLURM, then the per-job directory removed.
Notes
-----
Idempotent and best-effort by contract. An unreachable host or an
already-gone job never raises, because restart-time cleanup must always
be safe. Never touches material state -- that is the caller's business.
"""
scheme = handle.extra.get("scheme")
job_dir = str(handle.extra.get("job_dir", ""))
commands: list[str] = []
if scheme == "pid" and handle.job_id.isdigit():
commands.append(f"kill {shlex.quote(handle.job_id)} 2>/dev/null")
if scheme == "slurm" and handle.job_id.isdigit():
commands.append(f"scancel {shlex.quote(handle.job_id)} 2>/dev/null")
if job_dir:
commands.append(f"rm -rf {shlex.quote(job_dir)}")
if not commands:
return
# Cleanup is best-effort; a down host is not an error here.
with contextlib.suppress(BackendError):
self._run_ssh("; ".join(commands), timeout=15.0)
# -- small remote-read helpers --------------------------------------------
def _cat(self, path: str) -> str:
"""Read a remote file, best-effort.
Parameters
----------
path
The remote path.
Returns
-------
str
Its contents, or ``""`` when the path is empty or the file is
absent.
"""
if not path:
return ""
proc = self._run_ssh(f"cat {shlex.quote(path)} 2>/dev/null")
return proc.stdout or ""
def _read_rc(self, rc_path: str) -> int | None:
"""Read the exit-code sentinel the launch wrapper wrote.
Parameters
----------
rc_path
The remote sentinel path.
Returns
-------
int or None
The true exit code, or ``None`` when the sentinel is missing or
garbled -- which means "ran, code unknown" and is never read as
success.
"""
text = self._cat(rc_path).strip()
return int(text) if text.isdigit() else None
def _slurm_exit_code(self, job_id: str) -> int | None:
"""Ask SLURM's accounting database for a job's exit code.
Parameters
----------
job_id
The SLURM job id.
Returns
-------
int or None
The ``code`` half of ``sacct``'s ``code:signal``, or ``None`` when
accounting has nothing usable -- in which case ``fetch`` falls back
to the sentinel.
"""
proc = self._run_ssh(
f"sacct -j {shlex.quote(job_id)} --format=ExitCode --noheader --parsable2 | head -1"
)
code = (proc.stdout or "").strip().split(":")[0]
return int(code) if code.isdigit() else None
register("ssh_exec", SshExecBackend)