From 9385ad0c8dee6da5e2023dbfd356f51197dd4ac0 Mon Sep 17 00:00:00 2001 From: Mahdi Shafiei Date: Wed, 28 Oct 2026 15:45:03 -0700 Subject: [PATCH 1/4] fix(compute): stop a task's whole process tree when it ends or is cancelled On a timeout the worker killed the task's process alone, so a descendant it had started kept running with the task's log pipes open, and the worker waited on them. Each task now runs in its own session, and the worker kills that process group when the task exits, times out or is cancelled. The worker records SIGTERM and SIGINT and acts on them between polls of the task, so a cancellation finishes the cleanup before the worker exits, including for a task that ignores SIGTERM. The local backend's cancel signals the worker, which now reaps the task group it no longer shares. --- CHANGELOG.md | 4 + src/substrax/compute/local.py | 2 +- src/substrax/compute/worker.py | 98 +++++++++++++++++--- tests/compute/test_worker_process_tree.py | 108 ++++++++++++++++++++++ 4 files changed, 200 insertions(+), 12 deletions(-) create mode 100644 tests/compute/test_worker_process_tree.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ff4ace..9eed4b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -76,6 +76,10 @@ and uses semantic versioning while the public API stabilizes. ### Fixed +- A compute task runs in its own process group, and the worker kills the whole group when the + task ends, times out or is cancelled, so a descendant no longer outlives its task holding its + log pipes open. The worker turns SIGTERM and SIGINT into that cleanup, so cancelling a local + run stops its task tree, a task that ignores SIGTERM included. - The lockfile resolves the fixed releases of nine packages pip-audit reports: aiohttp 3.14.5, certifi 2026.10.23, filelock 4.2.4, pillow 13.0.0 and virtualenv 21.14.13, and gitpython 3.2.1, mlflow 3.17.2, opentelemetry-api 1.45.1 and werkzeug 3.1.9, which only the `mlflow` and `wandb` diff --git a/src/substrax/compute/local.py b/src/substrax/compute/local.py index cc6aa5b..5de0a8c 100644 --- a/src/substrax/compute/local.py +++ b/src/substrax/compute/local.py @@ -184,7 +184,7 @@ def cancel(self, handle: RunHandle) -> None: (Path(handle.details["run_dir"]) / _CANCELLED).touch() pid = int(handle.details["pid"]) if _is_worker(pid, handle.run_id): - # The worker leads its own process group, which its tasks inherit. + # The worker handles TERM and reaps the separate task process group. os.killpg(pid, signal.SIGTERM) def _interpreter(self, project: Path) -> list[str]: diff --git a/src/substrax/compute/worker.py b/src/substrax/compute/worker.py index 5dbe094..9ace992 100644 --- a/src/substrax/compute/worker.py +++ b/src/substrax/compute/worker.py @@ -12,16 +12,21 @@ from __future__ import annotations import argparse +import contextlib import dataclasses import json +import os +import signal import subprocess # nosec B404 - runs the job's own commands, never through a shell import sys import threading import time -from collections.abc import Sequence +from collections.abc import Generator, Sequence +from contextvars import ContextVar from dataclasses import dataclass from enum import StrEnum from pathlib import Path +from types import FrameType from typing import BinaryIO, IO from substrax.artifacts import OUTPUT_DIR_ENV @@ -44,6 +49,8 @@ "print(json.dumps({'platform': info.platform, 'kinds': list(info.device_kinds)}))" ) _DEVICE_PROBE_SECONDS = 300.0 +# How often a waiting task checks for a cancellation; the delay a cancel can take to land. +_POLL_SECONDS = 0.1 class TaskStatus(StrEnum): @@ -194,6 +201,43 @@ def run_job( return manifest +@dataclass +class _Cancellation: + signum: int | None = None + + +_CANCELLATION: ContextVar[_Cancellation | None] = ContextVar("worker_cancellation", default=None) + + +@contextlib.contextmanager +def _cancellation_scope() -> Generator[None]: + """Remember termination until the worker can reap its owned task group. + + Recording the signal instead of raising inside Popen also covers a signal + arriving between child creation and assignment of its process handle. + """ + state = _Cancellation() + token = _CANCELLATION.set(state) + + def remember(signum: int, frame: FrameType | None) -> None: + del frame + state.signum = signum + + previous = {sig: signal.signal(sig, remember) for sig in (signal.SIGTERM, signal.SIGINT)} + try: + yield + finally: + for sig, handler in previous.items(): + signal.signal(sig, handler) + _CANCELLATION.reset(token) + + +def _check_cancelled() -> None: + state = _CANCELLATION.get() + if state is not None and state.signum is not None: + raise SystemExit(128 + state.signum) + + def main(argv: Sequence[str] | None = None) -> int: """Run the job a spec file describes. @@ -209,7 +253,8 @@ def main(argv: Sequence[str] | None = None) -> int: parser.add_argument("--outputs", type=Path, required=True, help="where the run's files go") arguments = parser.parse_args(argv) spec = read_job_spec(json.loads(arguments.spec.read_text(encoding="utf-8"))) - manifest = run_job(spec, project=arguments.project, outputs=arguments.outputs) + with _cancellation_scope(): + manifest = run_job(spec, project=arguments.project, outputs=arguments.outputs) return 0 if manifest.succeeded else 1 @@ -255,6 +300,7 @@ def _probe_devices( def _run_task( spec: JobSpec, task: Task, *, project: Path, outputs: Path, echo: bool, python: str ) -> TaskOutcome: + _check_cancelled() directory = outputs / task.name directory.mkdir(parents=True, exist_ok=True) argv = (python, *task.argv[1:]) if task.argv[0] == _PYTHON else task.argv @@ -266,7 +312,12 @@ def _run_task( ): try: process = subprocess.Popen( # noqa: S603 # nosec B603 - the job's command, no shell - argv, cwd=project, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE + argv, + cwd=project, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + start_new_session=os.name == "posix", ) except OSError as error: err_log.write(f"{argv[0]} could not start: {error}\n".encode()) @@ -320,17 +371,42 @@ def _wait( for copier in copiers: copier.start() try: - returncode = process.wait(timeout=task.timeout_seconds) - status = TaskStatus.SUCCEEDED if returncode == 0 else TaskStatus.FAILED - except subprocess.TimeoutExpired: - process.kill() - returncode = process.wait() - status = TaskStatus.TIMED_OUT - for copier in copiers: - copier.join() + status = _wait_for_exit(process, deadline=time.monotonic() + task.timeout_seconds) + finally: + returncode = _stop_group(process) + for copier in copiers: + copier.join() return status, returncode +def _wait_for_exit(process: subprocess.Popen[bytes], *, deadline: float) -> TaskStatus: + """Poll until the task exits or its deadline passes; a cancellation raises between polls.""" + while True: + _check_cancelled() + remaining = deadline - time.monotonic() + if remaining <= 0: + return TaskStatus.TIMED_OUT + try: + returncode = process.wait(timeout=min(_POLL_SECONDS, remaining)) + except subprocess.TimeoutExpired: + continue + return TaskStatus.SUCCEEDED if returncode == 0 else TaskStatus.FAILED + + +def _stop_group(process: subprocess.Popen[bytes]) -> int: + """Kill the task's process group and return the task's exit code. + + The group is killed after a normal exit too: a descendant can outlive the task and hold its + log pipes open. A cancellation finishes this before the worker exits. + """ + with contextlib.suppress(ProcessLookupError): + if os.name == "posix": + os.killpg(process.pid, signal.SIGKILL) + else: + process.kill() + return process.wait() + + def _copy(source: IO[bytes] | None, log: BinaryIO, mirror: BinaryIO | None, prefix: bytes) -> None: """Write each line of ``source`` to the log and, prefixed, to the mirror.""" if source is None: diff --git a/tests/compute/test_worker_process_tree.py b/tests/compute/test_worker_process_tree.py new file mode 100644 index 0000000..f6819c7 --- /dev/null +++ b/tests/compute/test_worker_process_tree.py @@ -0,0 +1,108 @@ +"""A task timeout must not leave descendants holding the worker's pipes.""" + +import contextlib +import json +import os +import signal +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +from substrax.compute import JobSpec, Task +from substrax.compute.local import LocalBackend + + +def test_timeout_reaps_the_task_process_group(tmp_path: Path) -> None: + """An inherited stdout pipe must not extend the task's lifetime.""" + if os.name != "posix": + pytest.skip("POSIX process groups") + pid_file = tmp_path / "descendant.pid" + script = tmp_path / "check.py" + ready = tmp_path / "ready" + script.write_text( + "from pathlib import Path\n" + "from substrax.compute import JobSpec, Task\n" + "from substrax.compute.worker import run_job\n" + "from substrax.runtime import JaxRuntime\n" + f"root=Path({str(tmp_path)!r})\n" + 'body="import subprocess,sys,time; from pathlib import Path; ' + "p=subprocess.Popen([sys.executable,'-c','import time; time.sleep(60)']); " + f'Path({str(pid_file)!r}).write_text(str(p.pid)); time.sleep(60)"\n' + "job=JobSpec(name='tree',tasks=(Task(name='child',argv=('python','-c',body)," + "timeout_seconds=0.5),),timeout_seconds=20,runtime=JaxRuntime(platforms=('cpu',)))\n" + f"Path({str(ready)!r}).write_text('ready')\n" + "result=run_job(job,project=root,outputs=root/'outputs',echo=False)\n" + "assert result.tasks[0].status.value=='timed-out'\n", + encoding="utf-8", + ) + process = subprocess.Popen( # noqa: S603 - fixed local regression program + [sys.executable, str(script)], stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + try: + # Import/instrumentation startup has its own bound; the task still gets + # 0.5 seconds and the pipe-drain assertion still gets 12 seconds. + deadline = time.monotonic() + 30 + while not ready.exists() and process.poll() is None and time.monotonic() < deadline: + time.sleep(0.05) + assert ready.exists(), "worker imports did not complete within the startup budget" + _stdout, stderr = process.communicate(timeout=12) + assert process.returncode == 0, stderr.decode() + assert pid_file.exists(), "timeout must be exercised after the descendant starts" + finally: + if process.poll() is None: + process.kill() + process.communicate() + if pid_file.exists(): + with contextlib.suppress(ProcessLookupError): + os.kill(int(pid_file.read_text()), signal.SIGKILL) + + +@pytest.mark.parametrize("ignore_term", [False, True]) +def test_local_cancel_stops_task_and_descendant(tmp_path: Path, ignore_term: bool) -> None: + """Cancellation must stop actual processes, including TERM-resistant children.""" + + if os.name != "posix": + pytest.skip("POSIX local backend") + marker = tmp_path / "pids.json" + body = ( + "import json,os,signal,subprocess,sys,time; from pathlib import Path; " + + ("signal.signal(signal.SIGTERM,signal.SIG_IGN); " if ignore_term else "") + + "child=subprocess.Popen([sys.executable,'-c','import time; time.sleep(60)']); " + + f"Path({str(marker)!r}).write_text(json.dumps([os.getpid(),child.pid])); time.sleep(60)" + ) + backend = LocalBackend({"python": sys.executable, "state_dir": str(tmp_path / "state")}) + spec = JobSpec( + name="cancel-tree", + timeout_seconds=30, + tasks=(Task(name="tree", argv=("python", "-c", body), timeout_seconds=25),), + ) + handle = backend.submit(spec, project=tmp_path) + pids = [] + try: + deadline = time.monotonic() + 10 + while not marker.exists() and time.monotonic() < deadline: + time.sleep(0.05) + assert marker.exists(), "task did not start" + pids = json.loads(marker.read_text()) + backend.cancel(handle) + deadline = time.monotonic() + 3 + while any(_running(pid) for pid in pids) and time.monotonic() < deadline: + time.sleep(0.05) + assert not any(_running(pid) for pid in pids), "task tree survived cancellation" + finally: + for pid in [*pids, int(handle.details["pid"])]: + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + + +def _running(pid: int) -> bool: + state = subprocess.run( # noqa: S603 - numeric PID of this test's child + ["ps", "-o", "stat=", "-p", str(pid)], # noqa: S607 - POSIX ps + capture_output=True, + text=True, + check=False, + ).stdout.strip() + return bool(state) and not state.startswith("Z") From e3923c45871cbc4f27511301f1ad85b585278eaf Mon Sep 17 00:00:00 2001 From: Mahdi Shafiei Date: Wed, 28 Oct 2026 15:46:55 -0700 Subject: [PATCH 2/4] feat(compute)!: a run handle keeps the settings it was submitted with substrax-compute rebuilt a stored run's backend from the project's configuration as it is now, so editing or removing a backend's settings redirected or broke every run submitted before the edit. RunHandle gains a required settings field, which every backend fills through new_run(spec, project, settings=...) with its validated settings, its state directory resolved to an absolute path; the CLI rebuilds the backend from it alone. Handles are version 2, and a version-1 handle is refused: it holds no settings to rebuild from. RunStore.read_all() replaces handles(): it returns the readable handles and each unreadable record with the reason, so one old or damaged record no longer hides every run. substrax-compute status lists what it can read, names each record it cannot, and exits 1. BREAKING CHANGE: new_run takes settings and RunHandle requires them; version-1 handles are refused; RunStore.handles() is now read_all(). --- CHANGELOG.md | 9 +++ docs/api/compute.md | 9 +++ src/substrax/compute/backend.py | 78 +++++++++++++++++++----- src/substrax/compute/cli.py | 24 ++++---- src/substrax/compute/local.py | 9 +-- src/substrax/compute/modal_backend.py | 7 ++- src/substrax/compute/skypilot_backend.py | 9 +-- tests/compute/test_cli.py | 74 +++++++++++++++++++++- tests/compute/test_local.py | 1 + 9 files changed, 183 insertions(+), 37 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9eed4b3..5a4a73b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -57,6 +57,15 @@ and uses semantic versioning while the public API stabilizes. ### Changed +- **Breaking:** a run handle records the backend settings it was submitted with + (`RunHandle.settings`, required), and `substrax-compute` rebuilds the backend from them, never + from the project's current configuration, so editing a project no longer redirects its earlier + runs. A backend passes its validated settings to `new_run(spec, project, settings=...)`. + Handles are of version 2, and a version-1 handle is refused, since it holds no settings to + rebuild its backend from: remove the records of runs submitted before this release from the + state directory's `compute/runs/`. `substrax-compute status` lists every run it can read, + names each record it cannot with the reason, and exits 1. `RunStore.handles()` is replaced by + `RunStore.read_all()`, which returns both. - substrax's CI is one workflow, `ci.yml`. It calls `python-gates.yml`, adds the Orbax-floor checkpoint tests and the derived-status check, and ends in the `required` check. The Quality Checks, Security, Documentation and Build Verification workflows are folded into it, so a pull diff --git a/docs/api/compute.md b/docs/api/compute.md index 4dfd71a..5f3cedd 100644 --- a/docs/api/compute.md +++ b/docs/api/compute.md @@ -89,6 +89,15 @@ which the serialized worker requires, and which `uv sync` is told to use so a pr the backend writes a task file and runs `sky launch`, `sky queue -o json`, `sky logs`, `sky cancel` and `sky down`. +## Run handles + +A run's handle records the backend settings it was submitted with (`RunHandle.settings`), and +`status`, `logs`, `fetch` and `cancel` rebuild the backend from them, never from the project's +current configuration; a backend passes them to `new_run(spec, project, settings=...)`. Handles +are of version 2. One of another version is refused: `substrax-compute status RUN_ID` exits 2, +and `substrax-compute status` lists every run it can read, names each record it cannot with the +reason, and exits 1. + ## Adding a provider A backend is any object that satisfies `ComputeBackend` and is built from its settings table. diff --git a/src/substrax/compute/backend.py b/src/substrax/compute/backend.py index 67d3984..61242ea 100644 --- a/src/substrax/compute/backend.py +++ b/src/substrax/compute/backend.py @@ -23,7 +23,7 @@ from substrax.typing import JsonValue -HANDLE_VERSION = 1 +HANDLE_VERSION = 2 SPEC_NAME = "spec.json" """The file a staged run's spec is written to.""" STATE_DIR_ENV = "SUBSTRAX_STATE_DIR" @@ -56,6 +56,7 @@ class RunHandle: job: The job's name. project: The project root it was submitted from. submitted_at: When it was submitted, in UTC. + settings: Frozen provider configuration used for every recovery operation. details: The provider's own references to the run, such as an app or cluster id. """ @@ -67,6 +68,7 @@ class RunHandle: job: str project: str submitted_at: datetime + settings: Mapping[str, JsonValue] details: Mapping[str, str] = field(default_factory=dict) @@ -78,11 +80,13 @@ class NewRun: run_id: The run's id, from :func:`new_run_id`. spec: The job. project: The resolved project root. + settings: Frozen provider configuration used to construct the run handle. """ run_id: str spec: JobSpec project: Path + settings: Mapping[str, JsonValue] def handle(self, backend: str, details: Mapping[str, str]) -> RunHandle: """The handle of this run, submitted now. @@ -100,6 +104,7 @@ def handle(self, backend: str, details: Mapping[str, str]) -> RunHandle: job=self.spec.name, project=str(self.project), submitted_at=datetime.now(UTC), + settings=self.settings, details=details, ) @@ -181,17 +186,23 @@ def backend_state_dir(state_dir: str | None, backend: str) -> Path: return root / "compute" / backend -def new_run(spec: JobSpec, project: Path) -> NewRun: +def new_run(spec: JobSpec, project: Path, *, settings: Mapping[str, JsonValue]) -> NewRun: """Start submitting a run of ``spec`` from ``project``. Args: spec: The job. project: The project root. + settings: Validated provider settings to freeze into the handle. Returns: The run, named by :func:`new_run_id`. """ - return NewRun(run_id=new_run_id(spec.name), spec=spec, project=project.resolve()) + return NewRun( + run_id=new_run_id(spec.name), + spec=spec, + project=project.resolve(), + settings=json.loads(json.dumps(settings)), + ) def stage_run(runs: Path, run: NewRun) -> Path: @@ -248,21 +259,60 @@ def load(self, run_id: str) -> RunHandle: if not path.is_file(): msg = f"no run {run_id!r} in {self.root}" raise KeyError(msg) - return read_record(RunHandle, json.loads(path.read_text(encoding="utf-8"))) + return _read_handle(path) - def handles(self) -> list[RunHandle]: - """Every kept handle, newest first. + def read_all(self) -> StoredRuns: + """Every kept handle, newest first, and each record that cannot be read, with the reason. + + A record an older substrax wrote, or a damaged one, is named on its own rather than + hiding the runs that can still be read. Returns: - The handles. + The readable handles and the unreadable records. """ - if not self.root.is_dir(): - return [] - found = [ - read_record(RunHandle, json.loads(path.read_text(encoding="utf-8"))) - for path in self.root.glob("*.json") - ] - return sorted(found, key=lambda handle: handle.submitted_at, reverse=True) + handles: list[RunHandle] = [] + unreadable: dict[Path, str] = {} + for path in sorted(self.root.glob("*.json")) if self.root.is_dir() else []: + try: + handles.append(_read_handle(path)) + except ValueError as error: + unreadable[path] = str(error) + handles.sort(key=lambda handle: handle.submitted_at, reverse=True) + return StoredRuns(handles=tuple(handles), unreadable=unreadable) def _path(self, run_id: str) -> Path: return self.root / f"{run_id}.json" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class StoredRuns: + """What a :class:`RunStore` holds: the handles it can read, and why it cannot read the rest. + + Attributes: + handles: The readable handles, newest first. + unreadable: Each record that cannot be read, by path, with the reason. + """ + + handles: tuple[RunHandle, ...] + unreadable: Mapping[Path, str] + + +def _read_handle(path: Path) -> RunHandle: + """Read one stored handle; one of another handle version is refused, never reinterpreted. + + Args: + path: The record. + + Returns: + The handle. + + Raises: + ValueError: If the record is of another handle version; its JSON and its fields raise + ``ValueError`` subclasses too (``json.JSONDecodeError``, pydantic's validation error). + """ + data = json.loads(path.read_text(encoding="utf-8")) + version = data.get("schema_version") if isinstance(data, dict) else None + if version != HANDLE_VERSION: + msg = f"run handle version {version} cannot be read; this substrax reads {HANDLE_VERSION}" + raise ValueError(msg) + return read_record(RunHandle, data) diff --git a/src/substrax/compute/cli.py b/src/substrax/compute/cli.py index 90c144e..9472a22 100644 --- a/src/substrax/compute/cli.py +++ b/src/substrax/compute/cli.py @@ -12,7 +12,8 @@ ``run`` submits the job, follows its log, waits for it and fetches its outputs; with ``--detach`` it returns once the job is submitted, and the other commands find the run again by the id it prints. Fetched outputs go to ``--into``, or under the state directory, never into the project. -Exit codes: 0 for success, 1 for a run that did not succeed, 2 for a refused command. +Exit codes: 0 for success, 1 for a run that did not succeed or a stored run ``status`` cannot +read, 2 for a refused command. """ from __future__ import annotations @@ -28,9 +29,9 @@ RunHandle, RunState, RunStore, + StoredRuns, ) from substrax.compute.config import ( - ComputeConfig, parse_accelerator, read_compute_config, require_locked_substrax, @@ -83,11 +84,16 @@ def _run(arguments: argparse.Namespace, store: RunStore) -> int: def _status(arguments: argparse.Namespace, store: RunStore) -> int: - handles = [store.load(arguments.run_id)] if arguments.run_id else store.handles() - for handle in handles: + if arguments.run_id: + stored = StoredRuns(handles=(store.load(arguments.run_id),), unreadable={}) + else: + stored = store.read_all() + for handle in stored.handles: state = _backend_for(handle).status(handle) sys.stdout.write(f"{handle.run_id} {handle.backend} {handle.job} {state}\n") - return _SUCCESS + for path, reason in stored.unreadable.items(): + sys.stderr.write(f"substrax-compute: {path}: {reason}\n") + return _UNSUCCESSFUL if stored.unreadable else _SUCCESS def _logs(arguments: argparse.Namespace, store: RunStore) -> int: @@ -119,12 +125,8 @@ def _finish(backend: ComputeBackend, handle: RunHandle, into: Path | None) -> in def _backend_for(handle: RunHandle) -> ComputeBackend: - """Rebuild the backend of a stored run with the settings its project configures now.""" - project = Path(handle.project) - config = ( - read_compute_config(project) if (project / "pyproject.toml").is_file() else ComputeConfig() - ) - return load_backend(handle.backend, config.backends.get(handle.backend, {})) + """Recover every provider from the settings bound to this run.""" + return load_backend(handle.backend, handle.settings) def _destination(into: Path | None) -> Path: diff --git a/src/substrax/compute/local.py b/src/substrax/compute/local.py index 5de0a8c..faf46ce 100644 --- a/src/substrax/compute/local.py +++ b/src/substrax/compute/local.py @@ -21,7 +21,7 @@ import subprocess # nosec B404 - runs ps to confirm a pid is still this run's worker import time from collections.abc import Iterator, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from substrax.compute.backend import ( @@ -34,7 +34,7 @@ ) from substrax.compute.spec import JobSpec from substrax.compute.worker import JobManifest, MANIFEST_NAME, read_manifest -from substrax.records import read_record, UNKNOWN_FIELDS_REFUSED +from substrax.records import dump_record, read_record, UNKNOWN_FIELDS_REFUSED from substrax.typing import JsonValue @@ -77,7 +77,8 @@ def __init__(self, settings: Mapping[str, JsonValue] | None = None) -> None: # """ chosen = read_record(LocalSettings, settings or {}) self._python = chosen.python - self._runs = backend_state_dir(chosen.state_dir, NAME) + self._runs = backend_state_dir(chosen.state_dir, NAME).resolve() + self._settings = replace(chosen, state_dir=str(self._runs.parent.parent)) def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: """Start the worker on ``spec`` in the background and return its handle. @@ -89,7 +90,7 @@ def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: Returns: The run's handle. """ - run = new_run(spec, project) + run = new_run(spec, project, settings=dump_record(self._settings)) root, run_dir = run.project, stage_run(self._runs, run) argv = [ *self._interpreter(root), diff --git a/src/substrax/compute/modal_backend.py b/src/substrax/compute/modal_backend.py index f1ae5b5..7bba992 100644 --- a/src/substrax/compute/modal_backend.py +++ b/src/substrax/compute/modal_backend.py @@ -34,7 +34,7 @@ import threading import time from collections.abc import Iterator, Mapping -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path, PurePosixPath @@ -133,7 +133,8 @@ def __init__(self, settings: Mapping[str, JsonValue]) -> None: # noqa: DOC502 ValueError: If a setting is missing, unknown or of the wrong type. """ self._settings = read_record(ModalSettings, settings) - self._cancelled = backend_state_dir(self._settings.state_dir, NAME) / "cancelled" + self._cancelled = backend_state_dir(self._settings.state_dir, NAME).resolve() / "cancelled" + self._settings = replace(self._settings, state_dir=str(self._cancelled.parents[2])) def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: """Build the image and spawn the worker in a detached app. @@ -145,7 +146,7 @@ def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: Returns: The run's handle. """ - run = new_run(spec, project) + run = new_run(spec, project, settings=dump_record(self._settings)) root, run_id = run.project, run.run_id app = modal.App(f"{root.name}-{spec.name}") worker = app.function( diff --git a/src/substrax/compute/skypilot_backend.py b/src/substrax/compute/skypilot_backend.py index 62ed5b5..2046484 100644 --- a/src/substrax/compute/skypilot_backend.py +++ b/src/substrax/compute/skypilot_backend.py @@ -35,7 +35,7 @@ import subprocess # nosec B404 - runs the sky and bucket command line tools, no shell import tempfile from collections.abc import Iterator, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from substrax.compute.backend import ( @@ -48,7 +48,7 @@ ) from substrax.compute.spec import Accelerator, JobSpec from substrax.compute.worker import MANIFEST_NAME, read_manifest -from substrax.records import read_record, UNKNOWN_FIELDS_REFUSED +from substrax.records import dump_record, read_record, UNKNOWN_FIELDS_REFUSED from substrax.typing import JsonValue @@ -113,7 +113,8 @@ def __init__(self, settings: Mapping[str, JsonValue]) -> None: # noqa: DOC502 """ self._settings = read_record(SkyPilotSettings, settings) _downloader(self._settings.outputs_bucket) - self._runs = backend_state_dir(self._settings.state_dir, NAME) + self._runs = backend_state_dir(self._settings.state_dir, NAME).resolve() + self._settings = replace(self._settings, state_dir=str(self._runs.parent.parent)) def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: """Write the task file and launch it on a cluster of its own. @@ -125,7 +126,7 @@ def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: Returns: The run's handle, once the job is submitted on the cluster. """ - run = new_run(spec, project) + run = new_run(spec, project, settings=dump_record(self._settings)) root, run_id, run_dir = run.project, run.run_id, stage_run(self._runs, run) task = run_dir / "task.yaml" # YAML is a superset of JSON task.write_text(json.dumps(self._task(spec, root, run_id, run_dir), indent=2), "utf-8") diff --git a/tests/compute/test_cli.py b/tests/compute/test_cli.py index 9633cbd..eab5521 100644 --- a/tests/compute/test_cli.py +++ b/tests/compute/test_cli.py @@ -2,15 +2,20 @@ from __future__ import annotations +import json import sys import textwrap import time +from collections.abc import Mapping from pathlib import Path import pytest -from substrax.compute.backend import STATE_DIR_ENV +from substrax.compute import cli +from substrax.compute.backend import new_run, RunHandle, RunState, RunStore, STATE_DIR_ENV from substrax.compute.cli import main +from substrax.compute.spec import JobSpec, Task +from substrax.typing import JsonValue @pytest.fixture @@ -132,3 +137,70 @@ def test_a_project_that_does_not_lock_substrax_is_refused_before_submitting( assert code == 2 assert "substrax is not in" in capsys.readouterr().err + + +_SAVED = JobSpec( + name="saved", timeout_seconds=10, tasks=(Task(name="t", argv=("true",), timeout_seconds=5),) +) + + +class _RecordingBackend: + """A backend the CLI rebuilds for a stored run; it records the settings it was given.""" + + def __init__(self, name: str, settings: Mapping[str, JsonValue], built: list[object]) -> None: + built.append((name, settings)) + + def status(self, handle: RunHandle) -> RunState: + del handle + return RunState.SUCCEEDED + + +def _record_backends(monkeypatch: pytest.MonkeyPatch) -> list[object]: + built: list[object] = [] + monkeypatch.setattr( + cli, "load_backend", lambda name, settings: _RecordingBackend(name, settings, built) + ) + return built + + +@pytest.mark.parametrize("backend_name", ["skypilot", "local", "modal", "third-party"]) +def test_a_stored_run_is_recovered_with_the_settings_it_was_submitted_with( + project: Path, monkeypatch: pytest.MonkeyPatch, backend_name: str +) -> None: + """Editing a project's provider settings must not redirect an existing run.""" + settings: dict[str, JsonValue] = {"state_dir": "/original/state"} + handle = new_run(_SAVED, project, settings=settings).handle(backend_name, {"id": "owned"}) + RunStore().save(handle) + (project / "pyproject.toml").write_text("[tool.substrax.compute]\n", "utf-8") + built = _record_backends(monkeypatch) + + assert main(["status", handle.run_id]) == 0 + assert built == [(backend_name, settings)] + + +def test_listing_names_each_unsupported_record_and_still_lists_the_rest( + project: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + """One record an older substrax wrote does not hide the runs that can still be read.""" + handle = new_run(_SAVED, project, settings={}).handle("local", {"pid": "1"}) + store = RunStore() + store.save(handle) + old = store.root / "saved-20260901T000000Z-0000.json" + old.write_text(json.dumps({"schema_version": 1, "run_id": "saved-old"}), "utf-8") + _record_backends(monkeypatch) + + code = main(["status"]) + + captured = capsys.readouterr() + assert code == 1 + assert f"{handle.run_id} local saved succeeded" in captured.out + assert str(old) in captured.err + assert "version 1" in captured.err + + +def test_new_run_binds_typed_settings(project: Path) -> None: + """Every provider uses the same required typed settings field.""" + run = new_run(_SAVED, project, settings={"state_dir": "/original"}) + handle = run.handle("local", {"pid": "123"}) + assert handle.settings == {"state_dir": "/original"} + assert handle.schema_version == 2 diff --git a/tests/compute/test_local.py b/tests/compute/test_local.py index db4fd22..36336b5 100644 --- a/tests/compute/test_local.py +++ b/tests/compute/test_local.py @@ -48,6 +48,7 @@ def test_a_run_that_finishes_between_the_status_checks_reads_as_finished( job="job", project=str(tmp_path), submitted_at=datetime.now(UTC), + settings={}, details={"pid": "1", "run_dir": str(run_dir)}, ) From c7f381d75f11d994ed440276c47cb83b762f7e1a Mon Sep 17 00:00:00 2001 From: Mahdi Shafiei Date: Wed, 28 Oct 2026 15:52:26 -0700 Subject: [PATCH 3/4] feat(compute)!: run SkyPilot through its public SDK The skypilot backend drove the sky command and parsed what it printed: `sky queue -o json` mixed notices into its output, a job status it did not know read as a failed job, and a launch interrupted before it returned left no record of the request SkyPilot was still provisioning. Every operation now runs substrax/compute/sky_sdk.py in SkyPilot's own interpreter (sky_python, by default the python beside the sky command), which calls the public SDK and returns its native records as JSON. Before it waits, a launch writes the run's intent and SkyPilot's request id beside the task file. cancel leaves a finished run alone; otherwise it cancels that request and the job and tears the cluster down, attempting every step and raising an ExceptionGroup of the failures, which substrax-compute reports with exit code 1. teardown() removes a cluster on its own. An unknown job status raises; FAILED, FAILED_SETUP and FAILED_DRIVER read as failed. How a run's files move is a transport's, in skypilot_transports: the bucket transport holds what the backend did with buckets, so the backend no longer branches on it. BREAKING CHANGE: the backend needs SkyPilot's interpreter to import sky; the `sky` command alone is not used to run operations. --- CHANGELOG.md | 10 + docs/api/compute.md | 32 +- src/substrax/compute/cli.py | 15 +- src/substrax/compute/sky_sdk.py | 162 ++++++++++ src/substrax/compute/skypilot_backend.py | 338 ++++++++++++-------- src/substrax/compute/skypilot_transports.py | 218 +++++++++++++ tests/compute/fake_sky/sky.py | 118 ++++--- tests/compute/test_cli.py | 23 ++ tests/compute/test_sky_sdk.py | 110 +++++++ tests/compute/test_skypilot_backend.py | 183 +++++++++-- 10 files changed, 1000 insertions(+), 209 deletions(-) create mode 100644 src/substrax/compute/sky_sdk.py create mode 100644 src/substrax/compute/skypilot_transports.py create mode 100644 tests/compute/test_sky_sdk.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a4a73b..34cfc62 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -66,6 +66,16 @@ and uses semantic versioning while the public API stabilizes. state directory's `compute/runs/`. `substrax-compute status` lists every run it can read, names each record it cannot with the reason, and exits 1. `RunStore.handles()` is replaced by `RunStore.read_all()`, which returns both. +- **Breaking:** the `skypilot` backend runs SkyPilot through its public SDK, in SkyPilot's own + interpreter (`sky_python`, by default the `python` beside the `sky` command), and reads its + native request and queue records instead of the `sky` command's output. A launch records its + intent and SkyPilot's request id before it waits, so an interrupted one is reconciled rather + than submitted again. `cancel` cancels that request and the job and tears the cluster down, + attempting every step and raising an `ExceptionGroup` of the failures, which + `substrax-compute cancel` reports with exit code 1; a run that already ended is left as it is, + and `teardown(handle)` removes its cluster. A job status the backend does not know raises + instead of reading as failed; `FAILED_DRIVER` reads as failed. + - substrax's CI is one workflow, `ci.yml`. It calls `python-gates.yml`, adds the Orbax-floor checkpoint tests and the derived-status check, and ends in the `required` check. The Quality Checks, Security, Documentation and Build Verification workflows are folded into it, so a pull diff --git a/docs/api/compute.md b/docs/api/compute.md index 5f3cedd..a251f69 100644 --- a/docs/api/compute.md +++ b/docs/api/compute.md @@ -85,9 +85,33 @@ builder 2025.06 or later, set for the workspace in Modal's dashboard under Setti config; an older builder fails the build with `Image builder version <= 2024.10 requires modal to be specified in your pyproject.toml file`. The image runs the Python this interpreter runs, which the serialized worker requires, and which `uv sync` is told to use so a project's own -`python-preference` cannot choose another. SkyPilot is never imported: -the backend writes a task file and runs `sky launch`, `sky queue -o json`, `sky logs`, -`sky cancel` and `sky down`. +`python-preference` cannot choose another. + +### SkyPilot + +SkyPilot is never imported into the project's environment. Every operation (launch, queue, logs, +cancel, teardown) runs `substrax/compute/sky_sdk.py` in SkyPilot's own interpreter, which calls +SkyPilot's public SDK (0.11.1) and reads its native request and queue records. Before a launch +waits for provisioning, it writes the run's intent and SkyPilot's request id beside the task +file; an interrupted launch is reconciled from them, not submitted again. `cancel` leaves a run +that already ended as it is; otherwise it cancels the request and the job and tears the cluster +down, attempting every step and reporting each failure (`substrax-compute cancel` exits 1). + +The settings under `[tool.substrax.compute.backends.skypilot]`: + +| Setting | Default | Meaning | +| --- | --- | --- | +| `outputs_bucket` | required | the `gs://` or `s3://` bucket the runs' outputs go to | +| `buckets` | `{}` | the bucket behind each of a job's mounts, by mount name | +| `infra` | any enabled cloud | where to run, in SkyPilot's `infra` syntax: `gcp`, `gcp/us-central1` | +| `use_spot` | `false` | run on spot instances | +| `sky` | `"sky"` | the `sky` command, used to find SkyPilot's interpreter beside it | +| `sky_python` | the `python` beside `sky` | SkyPilot's interpreter | +| `worker_python` | `"uv run --no-sync python"` | the command that starts Python in the synced environment | +| `state_dir` | the default state directory | where task files, launch records and cancellations are kept | + +SkyPilot mounts the buckets, so the outputs outlive the cluster, which SkyPilot removes when the +job ends. ## Run handles @@ -139,3 +163,5 @@ class TestMyBackend(BackendContract): ::: substrax.compute.modal_backend ::: substrax.compute.skypilot_backend + +::: substrax.compute.skypilot_transports diff --git a/src/substrax/compute/cli.py b/src/substrax/compute/cli.py index 9472a22..8f5d83f 100644 --- a/src/substrax/compute/cli.py +++ b/src/substrax/compute/cli.py @@ -12,8 +12,8 @@ ``run`` submits the job, follows its log, waits for it and fetches its outputs; with ``--detach`` it returns once the job is submitted, and the other commands find the run again by the id it prints. Fetched outputs go to ``--into``, or under the state directory, never into the project. -Exit codes: 0 for success, 1 for a run that did not succeed or a stored run ``status`` cannot -read, 2 for a refused command. +Exit codes: 0 for success, 1 for unsuccessful execution or cleanup, or a stored run ``status`` +cannot read, 2 for a refused command. """ from __future__ import annotations @@ -56,12 +56,23 @@ def main(argv: Sequence[str] | None = None) -> int: store = RunStore() try: return arguments.command(arguments, store) + except ExceptionGroup as error: + for message in _failure_messages(error): + sys.stderr.write(f"substrax-compute: {message}\n") + return _UNSUCCESSFUL except (LookupError, FileNotFoundError, ValueError) as error: message = error.args[0] if isinstance(error, KeyError) and error.args else str(error) sys.stderr.write(f"substrax-compute: {message}\n") return _REFUSED +def _failure_messages(error: Exception) -> list[str]: + """Retain each cleanup cause without an unhandled exception traceback.""" + if isinstance(error, ExceptionGroup): + return [message for child in error.exceptions for message in _failure_messages(child)] + return [str(error)] + + def _run(arguments: argparse.Namespace, store: RunStore) -> int: project = arguments.project.resolve() config = read_compute_config(project) diff --git a/src/substrax/compute/sky_sdk.py b/src/substrax/compute/sky_sdk.py new file mode 100644 index 0000000..dc08627 --- /dev/null +++ b/src/substrax/compute/sky_sdk.py @@ -0,0 +1,162 @@ +"""Small public-SDK bridge executed in the separate SkyPilot tool environment. + +This file intentionally imports only the standard library until its entry point +runs. The scientific environment does not need SkyPilot installed. +""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib +import json +import os +import sys +from pathlib import Path +from types import ModuleType +from typing import TextIO + + +def queue_snapshot(api: ModuleType, cluster: str) -> list[dict[str, str]]: + """Read native queue records through the public asynchronous SDK. + + Args: + api: The SkyPilot module, or an SDK-contract test double. + cluster: The exact owned cluster name. + + Returns: + Job identifiers and status names, without provider-specific objects. + """ + jobs = api.get(api.queue(cluster)) + return [{"job_id": str(job["job_id"]), "status": job["status"].value} for job in jobs] + + +def launch_recorded( + api: ModuleType, task: str, cluster: str, receipt: Path, *, down: bool +) -> dict[str, str]: + """Persist the provisioning request before waiting for its result. + + Args: + api: The public SkyPilot SDK module. + task: The staged task file. + cluster: The owned cluster name. + receipt: New receipt path; existing attempts cannot be overwritten. + down: Whether SkyPilot should remove the cluster when the job ends. + + Returns: + Exact request and job identifiers. + """ + # Reserve the local receipt before any request. A crash leaves an explicit + # unresolved attempt; callers must reconcile it rather than submit again. + with receipt.open("x", encoding="utf-8") as stream: + request = api.launch( + api.Task.from_yaml(task), cluster_name=cluster, retry_until_up=False, down=down + ) + identity = {"request_id": str(request), "cluster": cluster} + json.dump(identity, stream) + stream.flush() + os.fsync(stream.fileno()) + job_id, _handle = api.get(request) + return {**identity, "job_id": str(job_id)} + + +def cancel_request(api: ModuleType, request_id: str) -> None: + """Wait for cancellation of one owned server request. + + Args: + api: The public SkyPilot SDK module. + request_id: The exact durable launch request identity. + """ + api.get(api.api_cancel(request_ids=[request_id])) + + +def cancel_job(api: ModuleType, cluster: str, job_id: int) -> None: + """Cancel one job through the same SDK used for submission. + + Args: + api: The public SkyPilot SDK module. + cluster: The owned cluster. + job_id: The exact job identifier. + """ + api.get(api.cancel(cluster, job_ids=[job_id])) + + +def down_cluster(api: ModuleType, cluster: str) -> None: + """Wait for SDK teardown without purging failed provider state. + + Args: + api: The public SkyPilot SDK module. + cluster: The owned cluster. + """ + api.get(api.down(cluster)) + + +def stream_logs( + api: ModuleType, cluster: str, job_id: int, *, follow: bool, output: TextIO +) -> None: + """Stream native log chunks, separate from SDK diagnostics. + + Args: + api: The public SkyPilot SDK module. + cluster: The owned cluster. + job_id: The exact job identifier. + follow: Wait for subsequent log output. + output: Destination for job log text. + """ + for chunk in api.tail_logs(cluster, job_id, follow, preload_content=False): + if chunk is not None: + output.write(chunk) + output.flush() + + +def main() -> None: + """Execute a public SDK operation; diagnostics go to standard error.""" + parser = argparse.ArgumentParser() + sub = parser.add_subparsers(dest="operation", required=True) + queue = sub.add_parser("queue") + queue.add_argument("cluster") + launch = sub.add_parser("launch") + launch.add_argument("cluster") + launch.add_argument("task") + launch.add_argument("receipt", type=Path) + launch.add_argument("--down", action="store_true") + cancel = sub.add_parser("cancel-request") + cancel.add_argument("request_id") + cancel_job_parser = sub.add_parser("cancel-job") + cancel_job_parser.add_argument("cluster") + cancel_job_parser.add_argument("job_id", type=int) + down = sub.add_parser("down") + down.add_argument("cluster") + logs = sub.add_parser("logs") + logs.add_argument("cluster") + logs.add_argument("job_id", type=int) + logs.add_argument("--follow", action="store_true") + args = parser.parse_args() + output = sys.stdout + with contextlib.redirect_stdout(sys.stderr): + api = importlib.import_module("sky") + if args.operation == "queue": + records = queue_snapshot(api, args.cluster) + elif args.operation == "cancel-request": + cancel_request(api, args.request_id) + records = {"cancelled_request": args.request_id} + elif args.operation == "cancel-job": + cancel_job(api, args.cluster, args.job_id) + records = {"cancelled_job": str(args.job_id)} + elif args.operation == "down": + down_cluster(api, args.cluster) + records = {"down": args.cluster} + elif args.operation == "logs": + stream_logs(api, args.cluster, args.job_id, follow=args.follow, output=output) + return + else: + records = launch_recorded(api, args.task, args.cluster, args.receipt, down=args.down) + sys.stdout.write(json.dumps(records) + "\n") + + +if __name__ == "__main__": + try: + main() + except Exception as error: # External process boundary reports every SDK failure. + sys.stderr.write(f"{type(error).__name__}: {error}\n") + sys.exit(1) diff --git a/src/substrax/compute/skypilot_backend.py b/src/substrax/compute/skypilot_backend.py index 2046484..d047bbc 100644 --- a/src/substrax/compute/skypilot_backend.py +++ b/src/substrax/compute/skypilot_backend.py @@ -1,17 +1,17 @@ """The ``skypilot`` backend: runs a job through SkyPilot (https://docs.skypilot.ai) on its clouds. SkyPilot is a separate tool, installed the way its documentation says, in an environment of its -own (``uv tool install --with pip "skypilot[gcp]"``). This backend never imports it. It writes a -SkyPilot task file and runs the ``sky`` command: +own (``uv tool install --with pip "skypilot[gcp]"``). The scientific environment never imports it: +every operation runs ``substrax/compute/sky_sdk.py`` in SkyPilot's own interpreter, which calls +SkyPilot's public SDK and reads its native request and queue records. -- ``sky launch -c -d --down -y`` provisions a cluster for the run. The cluster syncs the - project as its workdir and installs the locked environment (``uv sync --frozen``). It runs the - worker, then tears itself down. -- ``sky queue -o json`` gives the job's state while the cluster is up. -- ``sky logs`` streams its output, and ``sky cancel`` then ``sky down`` stop it. +Each run gets a cluster of its own. The task syncs the project as its workdir, installs the locked +environment (``uv sync --frozen``) and runs the worker. Before it waits for provisioning, a launch +records the run's intent and SkyPilot's request id beside the task file, so an interrupted launch +is reconciled from them rather than submitted again. Cancelling cancels that request and the job, +then tears the cluster down, keeping every failure. -The outputs go to a bucket mounted at ``/outputs``, so they outlive the cluster. They are fetched -with the bucket's own command line tool: ``gcloud storage`` for ``gs://``, ``aws s3`` for ``s3://``. +Inputs and outputs are buckets SkyPilot mounts; see :mod:`substrax.compute.skypilot_transports`. Settings (``[tool.substrax.compute.backends.skypilot]``): @@ -21,23 +21,25 @@ - ``infra``: where to run, in SkyPilot's ``infra`` syntax, such as ``gcp`` or ``gcp/us-central1``; any enabled cloud when unset. - ``use_spot``: run on spot instances. -- ``sky``: the ``sky`` command; ``sky`` on ``PATH`` by default. +- ``sky``: the ``sky`` command, used to find SkyPilot's interpreter; ``sky`` on ``PATH``. +- ``sky_python``: SkyPilot's interpreter, when it is not beside the ``sky`` command. - ``worker_python``: the command that starts Python in the synced environment; ``uv run --no-sync python`` by default. -- ``state_dir``: where task files and cancellations are kept; the default state directory when - unset. +- ``state_dir``: where task files, launch records and cancellations are kept; the default state + directory when unset. """ from __future__ import annotations import json import shlex -import subprocess # nosec B404 - runs the sky and bucket command line tools, no shell -import tempfile -from collections.abc import Iterator, Mapping, Sequence +import shutil +import subprocess # nosec B404 - runs SkyPilot's interpreter on the bridge, no shell +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass, field, replace from pathlib import Path +from substrax.compute import skypilot_transports as transports from substrax.compute.backend import ( backend_state_dir, new_run, @@ -46,17 +48,21 @@ SPEC_NAME, stage_run, ) +from substrax.compute.skypilot_transports import ( + BucketSettings, + BucketTransport, + OUTPUTS_PATH, + SPEC_PATH, + Transport, + UV_ON_PATH, +) from substrax.compute.spec import Accelerator, JobSpec -from substrax.compute.worker import MANIFEST_NAME, read_manifest from substrax.records import dump_record, read_record, UNKNOWN_FIELDS_REFUSED from substrax.typing import JsonValue NAME = "skypilot" -OUTPUTS_PATH = "/outputs" -SPEC_PATH = "/tmp/substrax-job-spec.json" # noqa: S108 # nosec B108 - a path on the run's own cluster _JOB_ID = "1" # the first and only job on a cluster made for the run -_UV_ON_PATH = 'export PATH="$HOME/.local/bin:$PATH"' _STATES = { "INIT": RunState.PENDING, "PENDING": RunState.PENDING, @@ -64,11 +70,11 @@ "RUNNING": RunState.RUNNING, "SUCCEEDED": RunState.SUCCEEDED, "CANCELLED": RunState.CANCELLED, + "FAILED": RunState.FAILED, + "FAILED_SETUP": RunState.FAILED, + "FAILED_DRIVER": RunState.FAILED, } -_DOWNLOADERS = { - "gs://": ("gcloud", "storage", "rsync", "--recursive"), - "s3://": ("aws", "s3", "sync"), -} +_CLEANUP_ERRORS = (RuntimeError, OSError, ValueError, KeyError) @dataclass(frozen=True, slots=True, kw_only=True) @@ -80,9 +86,10 @@ class SkyPilotSettings: buckets: The bucket URI behind each mount, by mount name. infra: Where to run, in SkyPilot's ``infra`` syntax. use_spot: Run on spot instances. - sky: The ``sky`` command. + sky: The ``sky`` command, used to find SkyPilot's interpreter. + sky_python: SkyPilot's interpreter; the ``python`` beside ``sky`` when unset. worker_python: The command that starts Python in the synced environment. - state_dir: Where task files and cancellations are kept. + state_dir: Where task files, launch records and cancellations are kept. """ __pydantic_config__ = UNKNOWN_FIELDS_REFUSED @@ -92,12 +99,13 @@ class SkyPilotSettings: infra: str | None = None use_spot: bool = False sky: str = "sky" + sky_python: str | None = None worker_python: str = "uv run --no-sync python" state_dir: str | None = None class SkyPilotBackend: - """Runs jobs through SkyPilot's ``sky`` command.""" + """Runs jobs through SkyPilot's public SDK, in SkyPilot's own interpreter.""" name = NAME @@ -109,15 +117,15 @@ def __init__(self, settings: Mapping[str, JsonValue]) -> None: # noqa: DOC502 Raises: ValueError: If a setting is missing, unknown or of the wrong type, or the outputs - bucket's scheme has no known download tool. + bucket's scheme has no download tool. """ - self._settings = read_record(SkyPilotSettings, settings) - _downloader(self._settings.outputs_bucket) - self._runs = backend_state_dir(self._settings.state_dir, NAME).resolve() - self._settings = replace(self._settings, state_dir=str(self._runs.parent.parent)) + chosen = read_record(SkyPilotSettings, settings) + self._runs = backend_state_dir(chosen.state_dir, NAME).resolve() + self._settings = replace(chosen, state_dir=str(self._runs.parent.parent)) + self._transport = _transport(self._settings) def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: - """Write the task file and launch it on a cluster of its own. + """Write the task file, record the launch, and launch it on a cluster of its own. Args: spec: The job. @@ -127,33 +135,48 @@ def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: The run's handle, once the job is submitted on the cluster. """ run = new_run(spec, project, settings=dump_record(self._settings)) - root, run_id, run_dir = run.project, run.run_id, stage_run(self._runs, run) + run_dir = stage_run(self._runs, run) task = run_dir / "task.yaml" # YAML is a superset of JSON - task.write_text(json.dumps(self._task(spec, root, run_id, run_dir), indent=2), "utf-8") - cluster = _cluster_name(run_id) - self._sky("launch", "-c", cluster, "-d", "--down", "-y", str(task)) - return run.handle(NAME, {"cluster": cluster, "job_id": _JOB_ID}) + content = self.task_for(spec, project=run.project, run_id=run.run_id, run_dir=run_dir) + task.write_text(json.dumps(content, indent=2), "utf-8") + details = {"cluster": _cluster_name(run.run_id), "job_id": _JOB_ID} + intent = dump_record(run.handle(NAME, details)) + (run_dir / "intent.json").write_text(json.dumps(intent), encoding="utf-8") + launch = [details["cluster"], str(task), str(run_dir / "request.json")] + down = ["--down"] if self._transport.down_with_job else [] + launched = json.loads(self._sdk("launch", *launch, *down)) + return run.handle( + NAME, {**details, "job_id": launched["job_id"], "request_id": launched["request_id"]} + ) - def status(self, handle: RunHandle) -> RunState: - """The job's state from ``sky queue`` while its cluster is up, else from its manifest. + def status(self, handle: RunHandle) -> RunState: # noqa: DOC503 # the ValueError is raised by _state + """The job's state from the cluster's queue while it is up, else from its outputs. Args: handle: The run. Returns: Its state. + + Raises: + RuntimeError: If the queue cannot be read and the outputs cannot answer for the run. + ValueError: If SkyPilot reports a job status this backend does not know. """ if (self._runs / handle.run_id / "cancelled").exists(): return RunState.CANCELLED - queue = _json_value(self._sky("queue", handle.details["cluster"], "-o", "json")) - jobs = queue.get(handle.details["cluster"], []) if isinstance(queue, dict) else [] - for job in jobs if isinstance(jobs, list) else []: - if isinstance(job, dict) and str(job.get("job_id")) == handle.details["job_id"]: - return _STATES.get(str(job.get("status")), RunState.FAILED) - return self._state_from_manifest(handle) + try: + jobs = json.loads(self._sdk("queue", handle.details["cluster"])) + except RuntimeError: + if not self._transport.recovers_without_queue: + raise + return self._transport.state_from_outputs(handle) + for job in jobs: + if str(job["job_id"]) == handle.details["job_id"]: + return _state(str(job["status"])) + return self._transport.state_from_outputs(handle) def logs(self, handle: RunHandle, *, follow: bool) -> Iterator[str]: - """The job's output lines, from ``sky logs``. + """The job's output lines, from SkyPilot's log stream. Args: handle: The run. @@ -161,17 +184,22 @@ def logs(self, handle: RunHandle, *, follow: bool) -> Iterator[str]: Yields: str: Each line, without its line ending. + + Raises: + RuntimeError: If the log stream fails. """ - flag = "--follow" if follow else "--no-follow" - command = [*self._command(), "logs", handle.details["cluster"], handle.details["job_id"]] - with subprocess.Popen( # noqa: S603 # nosec B603 - the sky command, no shell - [*command, flag], stdout=subprocess.PIPE, text=True + command = [*self._bridge(), "logs", handle.details["cluster"], handle.details["job_id"]] + with subprocess.Popen( # noqa: S603 # nosec B603 - SkyPilot's interpreter, no shell + [*command, *(["--follow"] if follow else [])], stdout=subprocess.PIPE, text=True ) as streamer: if streamer.stdout is not None: yield from (line.rstrip("\n") for line in streamer.stdout) + if streamer.wait() != 0: + msg = f"the log stream of {handle.run_id} failed" + raise RuntimeError(msg) def fetch(self, handle: RunHandle, destination: Path) -> Path: - """Download the run's directory of the outputs bucket to ``destination / handle.run_id``. + """Copy the run's outputs to ``destination / handle.run_id``. Args: handle: The run. @@ -182,47 +210,61 @@ def fetch(self, handle: RunHandle, destination: Path) -> Path: """ target = destination / handle.run_id target.mkdir(parents=True, exist_ok=True) - source = f"{self._settings.outputs_bucket.rstrip('/')}/{handle.run_id}" - _run([*_downloader(source), source, str(target)]) + self._transport.fetch(handle, target) return target def cancel(self, handle: RunHandle) -> None: - """Cancel the job and tear its cluster down; a run that already ended is left as it is. + """Cancel unfinished work and tear its cluster down; a run that already ended is left. + + Every step is attempted, whatever the one before it did. Args: handle: The run. + + Raises: + ExceptionGroup: If any step failed; it holds each failure. """ - if self.status(handle).is_final: + errors: list[Exception] = [] + state = _attempt(errors, self.status, handle) + if state is not None and state.is_final: return - (self._runs / handle.run_id / "cancelled").touch() - self._sky("cancel", handle.details["cluster"], "-a", "-y") - self._sky("down", handle.details["cluster"], "-y") + _attempt(errors, self._cancel_request, handle) + _attempt(errors, self._cancel_job, handle) + _attempt(errors, self.teardown, handle) + if errors: + msg = f"the cleanup of {handle.run_id} is incomplete" + raise ExceptionGroup(msg, errors) + if state is not None: + (self._runs / handle.run_id / "cancelled").touch() + + def teardown(self, handle: RunHandle) -> None: + """Remove the run's cluster, and with it anything left on it. - def _task( - self, spec: JobSpec, project: Path, run_id: str, run_dir: Path + Unlike :meth:`cancel`, this removes the cluster of a run that already ended too; fetch a + copy transport's outputs first. + + Args: + handle: The run. + """ + self._sdk("down", handle.details["cluster"]) + + def task_for( + self, spec: JobSpec, *, project: Path, run_id: str, run_dir: Path ) -> dict[str, JsonValue]: - resources: dict[str, JsonValue] = {"use_spot": self._settings.use_spot} - requested = spec.resources - if requested.accelerator is not None: - resources["accelerators"] = _accelerators(requested.accelerator) - # SkyPilot reads "+" as at least n and the unit "MB" as 2**20 bytes - # (sky/resources.py, sky/utils/resources_utils.py parse_memory_resource). - if requested.cpu is not None: - resources["cpus"] = f"{requested.cpu:g}+" - if requested.memory_mib is not None: - resources["memory"] = f"{requested.memory_mib}MB+" - if self._settings.infra is not None: - resources["infra"] = self._settings.infra - mounts: dict[str, JsonValue] = { - OUTPUTS_PATH: {"source": self._settings.outputs_bucket, "mode": "MOUNT"}, - SPEC_PATH: str(run_dir / SPEC_NAME), - } - for mount in spec.mounts: - if mount.name not in self._settings.buckets: - msg = f"mount {mount.name!r} has no bucket in the backend's `buckets`" - raise LookupError(msg) - mounts[mount.path] = {"source": self._settings.buckets[mount.name], "mode": "MOUNT"} - extras = " ".join(f"--extra {shlex.quote(extra)}" for extra in spec.extras) + """The SkyPilot task that runs ``spec``, as its task file holds it. + + Args: + spec: The job. + project: The project root, synced as the task's workdir. + run_id: The run. + run_dir: The run's staged directory, which holds its spec. + + Returns: + The task. + """ + resources = _resources(spec, self._settings) + mounts: dict[str, JsonValue] = {SPEC_PATH: str(run_dir / SPEC_NAME)} + self._transport.add_to_task(spec, resources, mounts) worker = shlex.join( [ *shlex.split(self._settings.worker_python), @@ -236,70 +278,90 @@ def _task( f"{OUTPUTS_PATH}/{run_id}", ] ) + extras = " ".join(f"--extra {shlex.quote(extra)}" for extra in spec.extras) + setup = [ + "set -e", + "command -v uv >/dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh", + UV_ON_PATH, + f"uv sync --frozen {extras}".rstrip(), + ] + run = [UV_ON_PATH, self._transport.wrap_worker(worker, spec)] return { "name": _cluster_name(run_id), "workdir": str(project), "resources": resources, "file_mounts": mounts, - "setup": "\n".join( - [ - "set -e", - "command -v uv >/dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh", - _UV_ON_PATH, - f"uv sync --frozen {extras}".rstrip(), - ] - ), - "run": "\n".join([_UV_ON_PATH, worker]), + "setup": "\n".join(setup), + "run": "\n".join(run), } - def _state_from_manifest(self, handle: RunHandle) -> RunState: - """The state of a run whose cluster is gone, from the manifest it left in the bucket.""" - with tempfile.TemporaryDirectory() as scratch: - source = f"{self._settings.outputs_bucket.rstrip('/')}/{handle.run_id}" - completed = subprocess.run( # noqa: S603 # nosec B603 - the bucket tool, no shell - [*_downloader(source), source, scratch], capture_output=True, check=False - ) - manifest_path = Path(scratch) / MANIFEST_NAME - if completed.returncode != 0 or not manifest_path.is_file(): - return RunState.FAILED - manifest = read_manifest(json.loads(manifest_path.read_text(encoding="utf-8"))) - if not manifest.finished: - return RunState.FAILED - return RunState.SUCCEEDED if manifest.succeeded else RunState.FAILED - - def _command(self) -> list[str]: - return shlex.split(self._settings.sky) - - def _sky(self, *arguments: str) -> str: - return _run([*self._command(), *arguments]) - - -def _run(command: Sequence[str]) -> str: - """Run a command line tool and return its standard output; a failure raises with its stderr.""" - completed = subprocess.run( # noqa: S603 # nosec B603 - a configured tool, no shell - list(command), capture_output=True, text=True, check=False + def _cancel_request(self, handle: RunHandle) -> None: + """Cancel the launch's request with SkyPilot's API server, when one was recorded.""" + request_id = handle.details.get("request_id") + receipt = self._runs / handle.run_id / "request.json" + if request_id is None and receipt.exists(): + request_id = json.loads(receipt.read_text(encoding="utf-8"))["request_id"] + if request_id is not None: + self._sdk("cancel-request", request_id) + + def _cancel_job(self, handle: RunHandle) -> None: + self._sdk("cancel-job", handle.details["cluster"], handle.details["job_id"]) + + def _bridge(self) -> list[str]: + """The command that runs ``sky_sdk.py`` in SkyPilot's interpreter.""" + python = self._settings.sky_python + if python is None: + sky = shutil.which(shlex.split(self._settings.sky)[0]) + if sky is None: + msg = f"{self._settings.sky!r} is not on PATH; set `sky` or `sky_python`" + raise FileNotFoundError(msg) + python = str(Path(sky).resolve().with_name("python")) + return [python, str(Path(__file__).with_name("sky_sdk.py"))] + + def _sdk(self, *arguments: str) -> str: + return transports.run_tool([*self._bridge(), *arguments]) + + +def _transport(settings: SkyPilotSettings) -> Transport: + """The run's transport: the buckets the settings name.""" + return BucketTransport( + BucketSettings(outputs_bucket=settings.outputs_bucket, buckets=settings.buckets) ) - if completed.returncode != 0: - msg = f"{shlex.join(command)} failed ({completed.returncode}): {completed.stderr.strip()}" - raise RuntimeError(msg) - return completed.stdout - - -def _json_value(output: str) -> JsonValue: - """Decode the JSON a ``sky ... -o json`` command prints after any notices on the same stream.""" - lines = output.splitlines() - for index, line in enumerate(lines): - if line.startswith(("{", "[")): - return json.loads("\n".join(lines[index:])) - return None - - -def _downloader(uri: str) -> tuple[str, ...]: - for scheme, command in _DOWNLOADERS.items(): - if uri.startswith(scheme): - return command - msg = f"no download tool for {uri!r}; use a gs:// or s3:// bucket" - raise ValueError(msg) + + +def _resources(spec: JobSpec, settings: SkyPilotSettings) -> dict[str, JsonValue]: + resources: dict[str, JsonValue] = {"use_spot": settings.use_spot} + requested = spec.resources + if requested.accelerator is not None: + resources["accelerators"] = _accelerators(requested.accelerator) + # SkyPilot reads "+" as at least n and the unit "MB" as 2**20 bytes + # (sky/resources.py, sky/utils/resources_utils.py parse_memory_resource). + if requested.cpu is not None: + resources["cpus"] = f"{requested.cpu:g}+" + if requested.memory_mib is not None: + resources["memory"] = f"{requested.memory_mib}MB+" + if settings.infra is not None: + resources["infra"] = settings.infra + return resources + + +def _state(status: str) -> RunState: + """A SkyPilot job status as a run state; an unknown one is an error, never a failed job.""" + if status not in _STATES: + msg = f"Unknown SkyPilot job status: {status}" + raise ValueError(msg) + return _STATES[status] + + +def _attempt[T]( + errors: list[Exception], operation: Callable[..., T], *arguments: object +) -> T | None: + """Run one cleanup step, keeping its failure instead of stopping the ones after it.""" + try: + return operation(*arguments) + except _CLEANUP_ERRORS as error: + errors.append(error) + return None def _cluster_name(run_id: str) -> str: diff --git a/src/substrax/compute/skypilot_transports.py b/src/substrax/compute/skypilot_transports.py new file mode 100644 index 0000000..1d73ad0 --- /dev/null +++ b/src/substrax/compute/skypilot_transports.py @@ -0,0 +1,218 @@ +"""How a SkyPilot run's inputs reach its cluster and its outputs come back. + +The ``skypilot`` backend runs every job the same way and leaves this to its transport. The one +transport, :class:`BucketTransport`, keeps inputs and outputs on object-store buckets SkyPilot +mounts (``mode: MOUNT``), so the outputs outlive the cluster, which SkyPilot tears down when the +job ends. Outputs are fetched with the bucket's own tool: ``gcloud storage`` for ``gs://``, +``aws s3`` for ``s3://``. +""" + +from __future__ import annotations + +import json +import shlex +import subprocess # nosec B404 - runs the bucket tools, no shell +import tempfile +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Protocol + +from substrax.compute.backend import RunHandle, RunState +from substrax.compute.spec import JobSpec +from substrax.compute.worker import JobManifest, MANIFEST_NAME, read_manifest +from substrax.records import UNKNOWN_FIELDS_REFUSED +from substrax.typing import JsonValue + + +OUTPUTS_PATH = "/outputs" +SPEC_PATH = "/tmp/substrax-job-spec.json" # noqa: S108 # nosec B108 - a path on the run's own cluster +UV_ON_PATH = 'export PATH="$HOME/.local/bin:$PATH"' +_DOWNLOADERS = { + "gs://": ("gcloud", "storage", "rsync", "--recursive"), + "s3://": ("aws", "s3", "sync"), +} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class BucketSettings: + """``[tool.substrax.compute.backends.skypilot.bucket]``: buckets SkyPilot mounts. + + Attributes: + outputs_bucket: The ``gs://`` or ``s3://`` URI the runs' outputs go to. + buckets: The bucket URI behind each of a job's mounts, by mount name. + """ + + __pydantic_config__ = UNKNOWN_FIELDS_REFUSED + + outputs_bucket: str + buckets: Mapping[str, str] = field(default_factory=dict) + + +class Transport(Protocol): + """The part of a SkyPilot run that depends on how its files move.""" + + @property + def down_with_job(self) -> bool: + """Whether SkyPilot removes the cluster as soon as its job ends.""" + ... + + @property + def recovers_without_queue(self) -> bool: + """Whether the outputs answer for the run when its cluster's queue cannot be read.""" + ... + + def add_to_task( + self, spec: JobSpec, resources: dict[str, JsonValue], mounts: dict[str, JsonValue] + ) -> None: + """Add the transport's resources and the job's input and output mounts to the task.""" + ... + + def wrap_worker(self, worker: str, spec: JobSpec) -> str: + """The command that runs ``worker`` on the cluster.""" + ... + + def fetch(self, handle: RunHandle, target: Path) -> None: + """Copy the run's outputs into ``target``.""" + ... + + def state_from_outputs(self, handle: RunHandle) -> RunState: + """The run's state once its cluster no longer reports it, from what it left behind.""" + ... + + +class BucketTransport: + """Inputs and outputs on buckets SkyPilot mounts; see the module documentation.""" + + down_with_job = True + # The cluster may still be running the job while its queue cannot be read. + recovers_without_queue = False + + def __init__(self, settings: BucketSettings) -> None: # noqa: DOC502 # the ValueError is raised by downloader + """Check the settings. + + Args: + settings: The bucket settings. + + Raises: + ValueError: If the outputs bucket's scheme has no download tool. + """ + downloader(settings.outputs_bucket) + self._settings = settings + + def add_to_task( + self, spec: JobSpec, resources: dict[str, JsonValue], mounts: dict[str, JsonValue] + ) -> None: + """Mount the outputs bucket and each input mount's bucket. + + Args: + spec: The job. + resources: The task's resources; unchanged. + mounts: The task's file mounts. + + Raises: + LookupError: If a mount has no configured bucket. + """ + del resources + mounts[OUTPUTS_PATH] = {"source": self._settings.outputs_bucket, "mode": "MOUNT"} + for mount in spec.mounts: + if mount.name not in self._settings.buckets: + msg = f"mount {mount.name!r} has no bucket in the backend's `buckets`" + raise LookupError(msg) + mounts[mount.path] = {"source": self._settings.buckets[mount.name], "mode": "MOUNT"} + + def wrap_worker(self, worker: str, spec: JobSpec) -> str: + """The worker runs as it is. + + Args: + worker: The worker command. + spec: The job. + + Returns: + ``worker``. + """ + del spec + return worker + + def fetch(self, handle: RunHandle, target: Path) -> None: + """Download the run's directory of the outputs bucket. + + Args: + handle: The run. + target: The directory to copy into. + """ + source = self._outputs(handle) + run_tool([*downloader(source), source, str(target)]) + + def state_from_outputs(self, handle: RunHandle) -> RunState: + """The state the run's manifest in the outputs bucket records. + + Args: + handle: The run. + + Returns: + Its state; ``failed`` when the run left no finished manifest. + """ + with tempfile.TemporaryDirectory() as scratch: + source = self._outputs(handle) + completed = subprocess.run( # noqa: S603 # nosec B603 - the bucket tool, no shell + [*downloader(source), source, scratch], capture_output=True, check=False + ) + manifest_path = Path(scratch) / MANIFEST_NAME + if completed.returncode != 0 or not manifest_path.is_file(): + return RunState.FAILED + return _finished_state(_read(manifest_path)) or RunState.FAILED + + def _outputs(self, handle: RunHandle) -> str: + return f"{self._settings.outputs_bucket.rstrip('/')}/{handle.run_id}" + + +def run_tool(command: Sequence[str]) -> str: + """Run a command line tool and return its standard output. + + Args: + command: The command. + + Returns: + Its standard output. + + Raises: + RuntimeError: If it exits non-zero; the message holds its standard error. + """ + completed = subprocess.run( # noqa: S603 # nosec B603 - a configured tool, no shell + list(command), capture_output=True, text=True, check=False + ) + if completed.returncode != 0: + msg = f"{shlex.join(command)} failed ({completed.returncode}): {completed.stderr.strip()}" + raise RuntimeError(msg) + return completed.stdout + + +def downloader(uri: str) -> tuple[str, ...]: + """The command that downloads a bucket URI's directory. + + Args: + uri: A ``gs://`` or ``s3://`` URI. + + Returns: + The command, to which the source and the destination are appended. + + Raises: + ValueError: If the scheme has no download tool. + """ + for scheme, command in _DOWNLOADERS.items(): + if uri.startswith(scheme): + return command + msg = f"no download tool for {uri!r}; use a gs:// or s3:// bucket" + raise ValueError(msg) + + +def _read(path: Path) -> JobManifest: + return read_manifest(json.loads(path.read_text(encoding="utf-8"))) + + +def _finished_state(manifest: JobManifest) -> RunState | None: + """``succeeded`` or ``failed`` for a finished job's manifest, ``None`` for an unfinished one.""" + if not manifest.finished: + return None + return RunState.SUCCEEDED if manifest.succeeded else RunState.FAILED diff --git a/tests/compute/fake_sky/sky.py b/tests/compute/fake_sky/sky.py index 723123d..289f5e4 100644 --- a/tests/compute/fake_sky/sky.py +++ b/tests/compute/fake_sky/sky.py @@ -1,8 +1,8 @@ -"""A stand-in for SkyPilot's ``sky`` command, running each task's ``run`` script locally. +"""A stand-in for SkyPilot's public SDK, running each task's ``run`` script locally. State lives under ``$FAKE_SKY_ROOT``: a bucket ``gs://NAME`` is ``buckets/NAME``, and a cluster is ``clusters/`` holding its job's pid, log and exit code. Paths the task mounts are rewritten -to their local directories in the script. ``--down`` removes the cluster once its job ends, as +to their local directories in the script. The launch ``down`` flag removes the cluster once its job ends, as SkyPilot's autodown does, so a finished run is known only from its bucket. """ @@ -10,6 +10,7 @@ import argparse import contextlib +import io import json import os import shlex @@ -17,7 +18,9 @@ import signal import sys import time +import uuid from pathlib import Path +from types import SimpleNamespace ROOT = Path(os.environ["FAKE_SKY_ROOT"]) @@ -31,7 +34,7 @@ def cluster_dir(name: str) -> Path: return ROOT / "clusters" / name -def launch(arguments: argparse.Namespace) -> int: +def launch_command(arguments: argparse.Namespace) -> int: task = json.loads(Path(arguments.task).read_text()) local: dict[str, str] = {} for remote, source in task["file_mounts"].items(): @@ -62,6 +65,18 @@ def launch(arguments: argparse.Namespace) -> int: return 0 +class Task: + @staticmethod + def from_yaml(path: str) -> str: + return path + + +def launch(task: str, *, cluster_name: str, retry_until_up: bool, down: bool) -> str: + assert not retry_until_up + launch_command(argparse.Namespace(task=task, cluster=cluster_name, down=down)) + return _request("launch", [1, None]) + + def _ended(cluster: Path) -> bool: return (cluster / "exit").exists() or (cluster / "cancelled").exists() @@ -74,7 +89,7 @@ def _gone(name: str) -> bool: return down and _ended(cluster) and (cluster / "logged").exists() -def queue(arguments: argparse.Namespace) -> int: +def queue_command(arguments: argparse.Namespace) -> int: records = {} for name in arguments.clusters: cluster = cluster_dir(name) @@ -92,6 +107,60 @@ def queue(arguments: argparse.Namespace) -> int: return 0 +def queue(cluster: str) -> str: + """Return native SDK-style records for the adapter contract.""" + output = io.StringIO() + with contextlib.redirect_stdout(output): + queue_command(argparse.Namespace(clusters=[cluster])) + text = output.getvalue() + start = text.index("{") + jobs = json.loads(text[start:]).get(cluster, []) + return _request("queue", jobs) + + +def _request(operation: str, result: object) -> str: + requests = ROOT / "requests" + requests.mkdir(parents=True, exist_ok=True) + identity = "request-" + uuid.uuid4().hex + (requests / identity).write_text(json.dumps({"operation": operation, "result": result})) + return identity + + +def api_cancel(*, request_ids: list[str]) -> str: + for identity in request_ids: + assert (ROOT / "requests" / identity).is_file() + return _request("api_cancel", request_ids) + + +def get(request: str) -> object: + """Resolve a persistent request identifier, never the result itself.""" + record = json.loads((ROOT / "requests" / request).read_text()) + result = record["result"] + if record["operation"] == "queue": + for job in result: + job["status"] = SimpleNamespace(value=job["status"]) + return result + + +def tail_logs(cluster_name: str, job_id: int, follow: bool, *, preload_content: bool) -> object: + assert not preload_content + output = io.StringIO() + with contextlib.redirect_stdout(output): + logs(argparse.Namespace(cluster=cluster_name, job_id=job_id, follow=follow)) + return iter([output.getvalue(), None]) + + +def cancel(cluster_name: str, *, job_ids: list[int]) -> str: + assert job_ids == [1] + cancel_command(argparse.Namespace(cluster=cluster_name)) + return _request("cancel", None) + + +def down(cluster_name: str) -> str: + down_command(argparse.Namespace(cluster=cluster_name)) + return _request("down", None) + + def logs(arguments: argparse.Namespace) -> int: cluster = cluster_dir(arguments.cluster) with (cluster / "job.log").open() as log: @@ -109,7 +178,7 @@ def logs(arguments: argparse.Namespace) -> int: return 0 -def cancel(arguments: argparse.Namespace) -> int: +def cancel_command(arguments: argparse.Namespace) -> int: cluster = cluster_dir(arguments.cluster) (cluster / "cancelled").touch() with contextlib.suppress(ProcessLookupError): @@ -117,43 +186,6 @@ def cancel(arguments: argparse.Namespace) -> int: return 0 -def down(arguments: argparse.Namespace) -> int: +def down_command(arguments: argparse.Namespace) -> int: shutil.rmtree(cluster_dir(arguments.cluster), ignore_errors=True) return 0 - - -def main() -> int: - parser = argparse.ArgumentParser(prog="sky") - commands = parser.add_subparsers(required=True) - p = commands.add_parser("launch") - p.add_argument("task") - p.add_argument("-c", "--cluster", required=True) - p.add_argument("-d", "--detach-run", action="store_true") - p.add_argument("--down", action="store_true") - p.add_argument("-y", "--yes", action="store_true") - p.set_defaults(command=launch) - p = commands.add_parser("queue") - p.add_argument("clusters", nargs="+") - p.add_argument("-o", "--output", choices=["json"], required=True) - p.set_defaults(command=queue) - p = commands.add_parser("logs") - p.add_argument("cluster") - p.add_argument("job_id") - p.add_argument("--follow", action="store_true") - p.add_argument("--no-follow", dest="follow", action="store_false") - p.set_defaults(command=logs) - p = commands.add_parser("cancel") - p.add_argument("cluster") - p.add_argument("-a", "--all", action="store_true") - p.add_argument("-y", "--yes", action="store_true") - p.set_defaults(command=cancel) - p = commands.add_parser("down") - p.add_argument("cluster") - p.add_argument("-y", "--yes", action="store_true") - p.set_defaults(command=down) - arguments = parser.parse_args() - return arguments.command(arguments) - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/compute/test_cli.py b/tests/compute/test_cli.py index eab5521..0eb650f 100644 --- a/tests/compute/test_cli.py +++ b/tests/compute/test_cli.py @@ -198,6 +198,29 @@ def test_listing_names_each_unsupported_record_and_still_lists_the_rest( assert "version 1" in captured.err +def test_cleanup_error_group_is_reported_without_traceback( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + """Every cleanup failure stays visible with a documented unsuccessful exit.""" + + def cancel(arguments: object, store: object) -> None: + del arguments, store + raise ExceptionGroup( + "cleanup", + [ + RuntimeError("queue unavailable"), + ExceptionGroup("teardown", [RuntimeError("down failed")]), + ], + ) + + monkeypatch.setattr(cli, "_cancel", cancel) + assert main(["cancel", "owned-run"]) == 1 + error = capsys.readouterr().err + assert "queue unavailable" in error + assert "down failed" in error + assert "Traceback" not in error + + def test_new_run_binds_typed_settings(project: Path) -> None: """Every provider uses the same required typed settings field.""" run = new_run(_SAVED, project, settings={"state_dir": "/original"}) diff --git a/tests/compute/test_sky_sdk.py b/tests/compute/test_sky_sdk.py new file mode 100644 index 0000000..a99efb6 --- /dev/null +++ b/tests/compute/test_sky_sdk.py @@ -0,0 +1,110 @@ +"""The bridge calls SkyPilot's public SDK and reads its native records, never its CLI's text.""" + +from enum import Enum +from pathlib import Path +from types import ModuleType + +import pytest + +from substrax.compute.sky_sdk import ( + cancel_job, + cancel_request, + down_cluster, + launch_recorded, + queue_snapshot, +) + + +class Status(Enum): + """Provider record spelling.""" + + RUNNING = "RUNNING" + + +def _sdk(**members: object) -> ModuleType: + """A stand-in for the ``sky`` module holding exactly ``members``.""" + api = ModuleType("sky") + vars(api).update(members) + return api + + +def test_queue_uses_sdk_get_and_serializes_native_status() -> None: + """SDK queue returns a request, whose result is a list of records.""" + calls: list[str] = [] + + def queue(cluster: str) -> str: + calls.append(cluster) + return "request-id" + + def get(request: str) -> list[dict[str, object]]: + assert request == "request-id" + return [{"job_id": 4, "status": Status.RUNNING}] + + api = _sdk(queue=queue, get=get) + assert queue_snapshot(api, "owned-cluster") == [{"job_id": "4", "status": "RUNNING"}] + assert calls == ["owned-cluster"] + + +def test_queue_errors_are_not_reported_as_job_failure() -> None: + """A transport error must remain visible to the caller.""" + + def queue(cluster: str) -> str: + raise ConnectionError(cluster) + + with pytest.raises(ConnectionError, match="owned-cluster"): + queue_snapshot(_sdk(queue=queue, get=lambda request: request), "owned-cluster") + + +def test_launch_persists_request_before_waiting(tmp_path: Path) -> None: + """A failed wait must leave the exact cancellable server request on disk.""" + receipt = tmp_path / "request.json" + + class Task: + @staticmethod + def from_yaml(path: str) -> str: + return path + + def launch(task: str, **kwargs: object) -> str: + assert task == "task.yaml" + assert kwargs == {"cluster_name": "owned", "retry_until_up": False, "down": False} + return "exact-request" + + def get(request: str) -> None: + assert request == "exact-request" + assert '"request_id": "exact-request"' in receipt.read_text() + raise ConnectionError("client disconnected") + + api = _sdk(Task=Task, launch=launch, get=get) + with pytest.raises(ConnectionError, match="disconnected"): + launch_recorded(api, "task.yaml", "owned", receipt, down=False) + assert receipt.exists() + + +def test_cancel_targets_only_the_recorded_request() -> None: + """Cancelling a client process is insufficient to stop server provisioning.""" + calls: list[object] = [] + + def api_cancel(**kwargs: object) -> str: + calls.append(kwargs) + return "cancel-operation" + + cancel_request(_sdk(api_cancel=api_cancel, get=calls.append), "owned-request") + assert calls == [{"request_ids": ["owned-request"]}, "cancel-operation"] + + +def test_job_cleanup_uses_one_sdk_and_exact_job_id() -> None: + """Cleanup uses public request APIs without switching to CLI display commands.""" + calls: list[object] = [] + + def cancel(cluster: str, *, job_ids: list[int]) -> str: + calls.append((cluster, job_ids)) + return "cancel-id" + + def down(cluster: str) -> str: + calls.append(cluster) + return "down-id" + + api = _sdk(cancel=cancel, down=down, get=calls.append) + cancel_job(api, "owned", 7) + down_cluster(api, "owned") + assert calls == [("owned", [7]), "cancel-id", "owned", "down-id"] diff --git a/tests/compute/test_skypilot_backend.py b/tests/compute/test_skypilot_backend.py index b2eb9f4..4f9475e 100644 --- a/tests/compute/test_skypilot_backend.py +++ b/tests/compute/test_skypilot_backend.py @@ -1,8 +1,9 @@ -"""The ``skypilot`` backend, against a local stand-in for the ``sky`` and ``gcloud`` commands. +"""The ``skypilot`` backend, against a local stand-in for SkyPilot's SDK and ``gcloud``. -``fake_sky/sky.py`` runs each task's ``run`` script on this machine with its mounts rewritten to -local directories, and ``fake_sky/bin/gcloud`` downloads from those directories, so the contract -runs here without a cloud account. +``fake_sky/sky.py`` is imported as SkyPilot's ``sky`` module by the SDK bridge; it runs each +task's ``run`` script on this machine with its mounts rewritten to local directories, and +``fake_sky/bin/gcloud`` downloads from those directories, so the contract runs here without a +cloud account. """ from __future__ import annotations @@ -11,13 +12,30 @@ import json import os import shlex +import subprocess import sys from pathlib import Path +from typing import Any import pytest -from substrax.compute import Accelerator, JobSpec, Mount, MountAccess, Resources, Task -from substrax.compute.backend import ComputeBackend +from substrax.compute import ( + Accelerator, + JobSpec, + Mount, + MountAccess, + Resources, + skypilot_backend, + skypilot_transports, + Task, +) +from substrax.compute.backend import ( + backend_state_dir, + ComputeBackend, + new_run, + RunHandle, + RunState, +) from substrax.compute.skypilot_backend import SkyPilotBackend from substrax.testing.compute import BackendContract from substrax.typing import JsonValue @@ -28,32 +46,39 @@ @pytest.fixture(autouse=True) def fake_sky(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: - """Point the fake ``sky`` and ``gcloud`` at a directory of the test's own.""" + """Point the fake SDK and ``gcloud`` at a directory of the test's own.""" root = tmp_path / "fake-sky" monkeypatch.setenv("FAKE_SKY_ROOT", str(root)) + monkeypatch.setenv("PYTHONPATH", str(FAKE) + os.pathsep + os.environ.get("PYTHONPATH", "")) monkeypatch.setenv("PATH", f"{FAKE / 'bin'}{os.pathsep}{os.environ['PATH']}") return root def _backend(state_dir: Path, **settings: JsonValue) -> SkyPilotBackend: + """A backend on the fake SDK.""" return SkyPilotBackend( { "outputs_bucket": "gs://outputs", "sky": shlex.join([sys.executable, str(FAKE / "sky.py")]), "worker_python": sys.executable, + "sky_python": sys.executable, "state_dir": str(state_dir), **settings, } ) +def _runs(state_dir: Path) -> Path: + return backend_state_dir(str(state_dir), "skypilot").resolve() + + class TestSkyPilotBackend(BackendContract): def make_backend(self, state_dir: Path) -> ComputeBackend: return _backend(state_dir) -def _task_file(state_dir: Path) -> dict[str, object]: - (path,) = (state_dir / "compute" / "skypilot").glob("*/task.yaml") +def _task_file(state_dir: Path) -> dict[str, Any]: + (path,) = _runs(state_dir).glob("*/task.yaml") return json.loads(path.read_text(encoding="utf-8")) @@ -67,6 +92,12 @@ def _spec(**changes: object) -> JobSpec: return dataclasses.replace(base, **changes) +def _handle(tmp_path: Path) -> RunHandle: + return new_run(_spec(), tmp_path, settings={}).handle( + "skypilot", {"cluster": "owned", "job_id": "1"} + ) + + def test_the_task_syncs_the_lock_with_the_extras_and_runs_the_worker(tmp_path: Path) -> None: project = tmp_path / "project" project.mkdir() @@ -82,14 +113,11 @@ def test_the_task_syncs_the_lock_with_the_extras_and_runs_the_worker(tmp_path: P "accelerators": "L4:1", "infra": "gcp/us-central1", } - assert "uv sync --frozen --extra cuda12 --extra docs" in str(task["setup"]) - run = str(task["run"]) - assert "-m substrax.compute.worker" in run - assert f"--outputs /outputs/{handle.run_id}" in run - assert "timeout" not in run - mounts = task["file_mounts"] - assert isinstance(mounts, dict) - assert mounts["/outputs"] == {"source": "gs://outputs", "mode": "MOUNT"} + assert "uv sync --frozen --extra cuda12 --extra docs" in task["setup"] + assert "-m substrax.compute.worker" in task["run"] + assert f"--outputs /outputs/{handle.run_id}" in task["run"] + assert "timeout" not in task["run"] + assert task["file_mounts"]["/outputs"] == {"source": "gs://outputs", "mode": "MOUNT"} def test_the_resources_ask_for_at_least_the_cpu_and_memory_requested(tmp_path: Path) -> None: @@ -100,7 +128,6 @@ def test_the_resources_ask_for_at_least_the_cpu_and_memory_requested(tmp_path: P ) resources = _task_file(tmp_path / "state")["resources"] - assert isinstance(resources, dict) assert (resources["cpus"], resources["memory"]) == ("8+", "65536MB+") @@ -109,7 +136,6 @@ def test_no_cpu_or_memory_is_asked_for_when_the_job_requests_none(tmp_path: Path _backend(tmp_path / "state").submit(_spec(), project=tmp_path / "project") resources = _task_file(tmp_path / "state")["resources"] - assert isinstance(resources, dict) assert "cpus" not in resources assert "memory" not in resources @@ -117,12 +143,11 @@ def test_no_cpu_or_memory_is_asked_for_when_the_job_requests_none(tmp_path: Path def test_each_mount_is_its_configured_bucket(tmp_path: Path) -> None: (tmp_path / "project").mkdir() mount = Mount(name="datasets", path="/data", access=MountAccess.READ_ONLY) - backend = _backend(tmp_path / "state", buckets={"datasets": "gs://demo-data"}) - - backend.submit(_spec(mounts=(mount,)), project=tmp_path / "project") + _backend(tmp_path / "state", buckets={"datasets": "gs://demo-data"}).submit( + _spec(mounts=(mount,)), project=tmp_path / "project" + ) mounts = _task_file(tmp_path / "state")["file_mounts"] - assert isinstance(mounts, dict) assert mounts["/data"] == {"source": "gs://demo-data", "mode": "MOUNT"} @@ -137,3 +162,115 @@ def test_a_mount_without_a_configured_bucket_is_refused(tmp_path: Path) -> None: def test_an_outputs_bucket_without_a_download_tool_is_refused(tmp_path: Path) -> None: with pytest.raises(ValueError, match="gs:// or s3://"): _backend(tmp_path / "state", outputs_bucket="azure://outputs") + + +def test_cancel_attempts_teardown_when_status_is_unavailable( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An unavailable status endpoint must not bypass cleanup.""" + backend = _backend(tmp_path / "state") + calls: list[tuple[str, ...]] = [] + + def status(handle: RunHandle) -> RunState: + del handle + raise RuntimeError("queue unavailable") + + monkeypatch.setattr(backend, "status", status) + monkeypatch.setattr(backend, "_sdk", lambda *args: calls.append(args)) + with pytest.raises(ExceptionGroup, match="cleanup"): + backend.cancel(_handle(tmp_path)) + assert ("down", "owned") in calls + + +def test_cancel_failure_still_attempts_down( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Cancellation and instance teardown are separate operations.""" + backend = _backend(tmp_path / "state") + handle = _handle(tmp_path) + (_runs(tmp_path / "state") / handle.run_id).mkdir(parents=True) + calls: list[tuple[str, ...]] = [] + + def sdk(*args: str) -> str: + calls.append(args) + if args[0] == "cancel-job": + raise RuntimeError("cancel failed") + return "" + + monkeypatch.setattr(backend, "status", lambda _: RunState.RUNNING) + monkeypatch.setattr(backend, "_sdk", sdk) + with pytest.raises(ExceptionGroup, match="cleanup"): + backend.cancel(handle) + assert ("down", "owned") in calls + assert not (_runs(tmp_path / "state") / handle.run_id / "cancelled").exists() + + +def test_cancel_of_a_finished_run_leaves_its_cluster_and_outputs( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The published cancel contract is a no-op for terminal runs.""" + backend = _backend(tmp_path / "state") + calls: list[tuple[str, ...]] = [] + monkeypatch.setattr(backend, "status", lambda _: RunState.SUCCEEDED) + monkeypatch.setattr(backend, "_sdk", lambda *args: calls.append(args)) + backend.cancel(_handle(tmp_path)) + assert not calls + + +def test_unknown_queue_state_is_not_a_scientific_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A provider schema change must remain an infrastructure error.""" + backend = _backend(tmp_path / "state") + reply = '[{"job_id":"1","status":"NEW_STATE"}]' + monkeypatch.setattr(skypilot_transports, "run_tool", lambda _: reply) + with pytest.raises(ValueError, match="Unknown SkyPilot"): + backend.status(_handle(tmp_path)) + + +def test_a_bucket_run_whose_queue_cannot_be_read_is_an_error_not_a_state( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The cluster may still be running the job; the queue's own error is kept.""" + + def unreadable(command: list[str]) -> str: + del command + raise RuntimeError("API server unreachable") + + monkeypatch.setattr(skypilot_transports, "run_tool", unreadable) + with pytest.raises(RuntimeError, match="API server unreachable"): + _backend(tmp_path / "state").status(_handle(tmp_path)) + + +@pytest.mark.parametrize("provider_status", ["FAILED", "FAILED_SETUP", "FAILED_DRIVER"]) +def test_native_failure_statuses_are_terminal( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, provider_status: str +) -> None: + """Known SDK failures remain failures; future unknown states remain errors.""" + reply = json.dumps([{"job_id": 1, "status": provider_status}]) + monkeypatch.setattr(skypilot_transports, "run_tool", lambda _: reply) + assert _backend(tmp_path / "state").status(_handle(tmp_path)) is RunState.FAILED + + +def test_fake_launch_receipt_contains_a_request_not_its_result(tmp_path: Path) -> None: + """The provider double must model asynchronous request identities.""" + handle = _backend(tmp_path / "state").submit(_spec(), project=tmp_path) + assert handle.details["request_id"].startswith("request-") + + +def test_sdk_failure_diagnostic_has_no_embedded_traceback() -> None: + """A bridge failure must not reintroduce a traceback inside a grouped CLI error.""" + result = subprocess.run( # noqa: S603 - isolated fake SDK and fixed request + [ + sys.executable, + str(Path(skypilot_backend.__file__).with_name("sky_sdk.py")), + "cancel-request", + "unknown-request", + ], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 1 + assert "AssertionError" in result.stderr + assert "Traceback" not in result.stderr From 82e4ec2dab72c9709a1ca13c051b1d4dc651ef74 Mon Sep 17 00:00:00 2001 From: Mahdi Shafiei Date: Wed, 28 Oct 2026 15:53:53 -0700 Subject: [PATCH 4/4] feat(compute)!: SkyPilot's copy transport, and one settings table per transport Providers such as RunPod and Vast cannot mount the buckets the skypilot backend relied on, so it could not run there. The copy transport has SkyPilot copy each mount's local source onto the cluster, read-only; a writable mount is refused, since a copy cannot be written back. The worker runs under bubblewrap with the system, the project and the inputs read-only, and network_isolation also takes its network away. fetch copies the outputs over rsync and verifies them against manifests of the remote directory (sky_transfer) taken before and after; a difference keeps the cluster up. The cluster stays up until teardown() or its lifetime_minutes timer, which starts after setup, and status reads a verified copy once the cluster's queue is gone. Each transport's settings are a table of their own, and the backend configures exactly one: [...skypilot.bucket] (outputs_bucket, buckets) or [...skypilot.copy] (sources, disk_size, lifetime_minutes, network_isolation). setup_commands prepares an image before the locked environment is synced. Bucket transport refuses a provider that cannot mount buckets. BREAKING CHANGE: outputs_bucket and buckets move under [tool.substrax.compute.backends.skypilot.bucket]. --- CHANGELOG.md | 12 +- docs/api/compute.md | 39 +++- src/substrax/compute/sky_transfer.py | 92 ++++++++++ src/substrax/compute/skypilot_backend.py | 55 +++--- src/substrax/compute/skypilot_transports.py | 189 ++++++++++++++++++-- tests/compute/test_sky_transfer.py | 36 ++++ tests/compute/test_skypilot_backend.py | 173 +++++++++++++++++- tests/compute/test_skypilot_isolation.py | 61 +++++++ 8 files changed, 614 insertions(+), 43 deletions(-) create mode 100644 src/substrax/compute/sky_transfer.py create mode 100644 tests/compute/test_sky_transfer.py create mode 100644 tests/compute/test_skypilot_isolation.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 34cfc62..860f5f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -54,6 +54,13 @@ and uses semantic versioning while the public API stabilizes. The other inputs cover what differs between repositories: `sync-args`, `docs-sync-args`, `python-versions`, `pytest-args`, `test-env`, `import-modules`, `os` and `uv-version`. +- The `skypilot` backend's copy transport, `[tool.substrax.compute.backends.skypilot.copy]`, for + providers that cannot mount buckets, such as RunPod and Vast. SkyPilot copies each mount's + local source onto the cluster, read-only, and the worker runs under bubblewrap with the system, + the project and the inputs read-only (`network_isolation = true` also takes its network away). + `fetch` copies the outputs over rsync and verifies them against manifests of the remote + directory taken before and after the copy; the cluster stays up until + `SkyPilotBackend.teardown(handle)` or its `lifetime_minutes` timer removes it. ### Changed @@ -74,7 +81,10 @@ and uses semantic versioning while the public API stabilizes. attempting every step and raising an `ExceptionGroup` of the failures, which `substrax-compute cancel` reports with exit code 1; a run that already ended is left as it is, and `teardown(handle)` removes its cluster. A job status the backend does not know raises - instead of reading as failed; `FAILED_DRIVER` reads as failed. + instead of reading as failed; `FAILED_DRIVER` reads as failed. The transport's settings move + into a table of their own: `outputs_bucket` and `buckets` go under + `[tool.substrax.compute.backends.skypilot.bucket]`, and a backend configures exactly one of + `bucket` and `copy`. - substrax's CI is one workflow, `ci.yml`. It calls `python-gates.yml`, adds the Orbax-floor checkpoint tests and the derived-status check, and ends in the `required` check. The Quality diff --git a/docs/api/compute.md b/docs/api/compute.md index a251f69..3054b8c 100644 --- a/docs/api/compute.md +++ b/docs/api/compute.md @@ -76,7 +76,7 @@ itself never starts a JAX backend. | --- | --- | --- | | `local` | this machine, in the project's current environment | nothing | | `modal` | Modal, one detached function per run, outputs on a Volume | `substrax[modal]`, a Modal account, and a workspace image builder of 2025.06 or later | -| `skypilot` | a cluster SkyPilot provisions for the run on any cloud it supports, torn down after it; outputs on a `gs://` or `s3://` bucket | SkyPilot installed as its documentation says (`uv tool install --with pip "skypilot[gcp]"`), and `gcloud` or `aws` for the bucket | +| `skypilot` | a cluster SkyPilot provisions for the run on any cloud it supports; inputs and outputs on buckets, or copied | SkyPilot installed as its documentation says (`uv tool install --with pip "skypilot[gcp]"`, or the provider's extra), and the transport's tools: `gcloud` or `aws` for buckets, `ssh` and `rsync` for copies | The Modal image is built in layers keyed on what they depend on: the locked dependencies (`Image.uv_sync`), then the project's source, then the project installed without @@ -97,21 +97,46 @@ file; an interrupted launch is reconciled from them, not submitted again. `cance that already ended as it is; otherwise it cancels the request and the job and tears the cluster down, attempting every step and reporting each failure (`substrax-compute cancel` exits 1). -The settings under `[tool.substrax.compute.backends.skypilot]`: +The settings under `[tool.substrax.compute.backends.skypilot]`, all optional: | Setting | Default | Meaning | | --- | --- | --- | -| `outputs_bucket` | required | the `gs://` or `s3://` bucket the runs' outputs go to | -| `buckets` | `{}` | the bucket behind each of a job's mounts, by mount name | -| `infra` | any enabled cloud | where to run, in SkyPilot's `infra` syntax: `gcp`, `gcp/us-central1` | +| `infra` | any enabled cloud | where to run, in SkyPilot's `infra` syntax: `gcp`, `gcp/us-central1`, `runpod` | | `use_spot` | `false` | run on spot instances | +| `setup_commands` | `[]` | commands that prepare the image before the locked environment is synced | | `sky` | `"sky"` | the `sky` command, used to find SkyPilot's interpreter beside it | | `sky_python` | the `python` beside `sky` | SkyPilot's interpreter | | `worker_python` | `"uv run --no-sync python"` | the command that starts Python in the synced environment | | `state_dir` | the default state directory | where task files, launch records and cancellations are kept | -SkyPilot mounts the buckets, so the outputs outlive the cluster, which SkyPilot removes when the -job ends. +A run's files move one of two ways, configured as exactly one of these tables. + +**Bucket transport** (`[...skypilot.bucket]`): SkyPilot mounts the buckets, so the outputs +outlive the cluster, which SkyPilot removes when the job ends. Providers that cannot mount +buckets (RunPod, Vast) are refused. + +```toml +[tool.substrax.compute.backends.skypilot.bucket] +outputs_bucket = "gs://avitai-runs" # required; gs:// or s3:// +buckets = { datasets = "gs://avitai-datasets" } # each mount's bucket, by mount name +``` + +**Copy transport** (`[...skypilot.copy]`): SkyPilot copies each mount's local source onto the +cluster, read-only; a writable mount is refused, since a copy cannot be written back. The worker +runs under bubblewrap, so the image has to provide `bwrap` with working user namespaces: the +system, the project and the inputs are read-only, `/tmp` and the outputs writable. `fetch` +copies the outputs over `rsync` and checks them against manifests of the remote directory taken +before and after; a difference keeps the cluster up. Fetch the outputs before +`SkyPilotBackend.teardown(handle)`, which removes the cluster, and with it the outputs on it. +`lifetime_minutes` starts once setup ends; it is not a provisioning deadline or a cost cap. + +```toml +[tool.substrax.compute.backends.skypilot.copy] +sources = { datasets = "/data/datasets" } # each mount's local path, by mount name +disk_size = 60 # GB +lifetime_minutes = 60 # idle time after setup before the cluster removes itself +network_isolation = false # true runs the worker without network access +``` ## Run handles diff --git a/src/substrax/compute/sky_transfer.py b/src/substrax/compute/sky_transfer.py new file mode 100644 index 0000000..ecc3ee1 --- /dev/null +++ b/src/substrax/compute/sky_transfer.py @@ -0,0 +1,92 @@ +"""Manifests of a run's output directory, for SkyPilot's copy transport. + +The cluster prints the manifest of its outputs (``python -m substrax.compute.sky_transfer DIR``) +before and after they are copied, and the copy is verified against it. A link is recorded by its +target and never followed. +""" + +from __future__ import annotations + +import hashlib +import json +import sys +from pathlib import Path, PurePosixPath + + +type Manifest = dict[str, dict[str, str | int]] +"""Each relative path's identity: a file's size and digest, a directory, or a link's target.""" + + +def snapshot(root: Path) -> Manifest: + """The manifest of every entry under ``root``. + + Args: + root: The output directory. + + Returns: + Each relative path's identity. + + Raises: + ValueError: If ``root`` is not a directory, is a link, or holds an entry that is neither a + file, a directory nor a link. + """ + if not root.is_dir() or root.is_symlink(): + msg = f"{root} is not an output directory" + raise ValueError(msg) + return {path.relative_to(root).as_posix(): _identity(path) for path in sorted(root.rglob("*"))} + + +def verify(root: Path, expected: Manifest) -> None: + """Check that ``root`` holds exactly what ``expected`` describes. + + Args: + root: The copied directory. + expected: The manifest of the directory it was copied from. + + Raises: + ValueError: If a path in ``expected`` leaves ``root``, or the copy differs from it. + """ + for name in expected: + path = PurePosixPath(name) + if path.is_absolute() or ".." in path.parts or path.as_posix() != name or name == ".": + msg = f"unsafe path {name!r} in the manifest" + raise ValueError(msg) + if snapshot(root) != expected: + msg = f"the copy in {root} does not match the manifest it was copied from" + raise ValueError(msg) + + +def _identity(path: Path) -> dict[str, str | int]: + if path.is_symlink(): + return {"kind": "symlink", "target": str(path.readlink())} + if path.is_dir(): + return {"kind": "directory"} + if path.is_file(): + return _file_identity(path) + msg = f"{path} is neither a file, a directory nor a link" + raise ValueError(msg) + + +def _file_identity(path: Path) -> dict[str, str | int]: + """A file's size and SHA-256; a file that changes while it is read is refused.""" + before = path.stat() + with path.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() + after = path.stat() + if (before.st_ino, before.st_size, before.st_mtime_ns) != ( + after.st_ino, + after.st_size, + after.st_mtime_ns, + ): + msg = f"{path} changed while it was read" + raise ValueError(msg) + return {"kind": "file", "bytes": after.st_size, "sha256": digest} + + +def main() -> None: + """Print the manifest of the directory named on the command line.""" + sys.stdout.write(json.dumps(snapshot(Path(sys.argv[1])), sort_keys=True) + "\n") + + +if __name__ == "__main__": + main() diff --git a/src/substrax/compute/skypilot_backend.py b/src/substrax/compute/skypilot_backend.py index d047bbc..e3470c5 100644 --- a/src/substrax/compute/skypilot_backend.py +++ b/src/substrax/compute/skypilot_backend.py @@ -11,15 +11,15 @@ is reconciled from them rather than submitted again. Cancelling cancels that request and the job, then tears the cluster down, keeping every failure. -Inputs and outputs are buckets SkyPilot mounts; see :mod:`substrax.compute.skypilot_transports`. +How inputs reach the cluster and outputs come back is the transport's, configured as exactly one +of ``[...skypilot.bucket]`` or ``[...skypilot.copy]``; see :mod:`substrax.compute.skypilot_transports`. Settings (``[tool.substrax.compute.backends.skypilot]``): -- ``outputs_bucket`` (required): the bucket URI the runs' outputs go to, such as - ``gs://avitai-runs``. -- ``buckets``: the bucket URI behind each of a job's mounts, by mount name. -- ``infra``: where to run, in SkyPilot's ``infra`` syntax, such as ``gcp`` or ``gcp/us-central1``; - any enabled cloud when unset. +- ``bucket`` or ``copy``: the transport's table. +- ``setup_commands``: commands that prepare the image before the locked environment is synced. +- ``infra``: where to run, in SkyPilot's ``infra`` syntax, such as ``gcp`` or ``runpod``; any + enabled cloud when unset. - ``use_spot``: run on spot instances. - ``sky``: the ``sky`` command, used to find SkyPilot's interpreter; ``sky`` on ``PATH``. - ``sky_python``: SkyPilot's interpreter, when it is not beside the ``sky`` command. @@ -36,7 +36,7 @@ import shutil import subprocess # nosec B404 - runs SkyPilot's interpreter on the bridge, no shell from collections.abc import Callable, Iterator, Mapping -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, replace from pathlib import Path from substrax.compute import skypilot_transports as transports @@ -51,6 +51,8 @@ from substrax.compute.skypilot_transports import ( BucketSettings, BucketTransport, + CopySettings, + CopyTransport, OUTPUTS_PATH, SPEC_PATH, Transport, @@ -82,8 +84,9 @@ class SkyPilotSettings: """The ``skypilot`` backend's settings; see the module documentation. Attributes: - outputs_bucket: The bucket URI the runs' outputs go to. - buckets: The bucket URI behind each mount, by mount name. + bucket: Bucket transport's settings; set exactly one of ``bucket`` and ``copy``. + copy: Copy transport's settings. + setup_commands: Commands run before the locked environment is synced. infra: Where to run, in SkyPilot's ``infra`` syntax. use_spot: Run on spot instances. sky: The ``sky`` command, used to find SkyPilot's interpreter. @@ -94,8 +97,9 @@ class SkyPilotSettings: __pydantic_config__ = UNKNOWN_FIELDS_REFUSED - outputs_bucket: str - buckets: Mapping[str, str] = field(default_factory=dict) + bucket: BucketSettings | None = None + copy: CopySettings | None = None + setup_commands: tuple[str, ...] = () infra: str | None = None use_spot: bool = False sky: str = "sky" @@ -113,16 +117,16 @@ def __init__(self, settings: Mapping[str, JsonValue]) -> None: # noqa: DOC502 """Read the backend's settings. Args: - settings: The settings table; ``outputs_bucket`` is required. + settings: The settings table, with exactly one of ``bucket`` and ``copy``. Raises: - ValueError: If a setting is missing, unknown or of the wrong type, or the outputs - bucket's scheme has no download tool. + ValueError: If a setting is missing, unknown or of the wrong type, or the transport + is not exactly one of ``bucket`` and ``copy``, or its own checks fail. """ chosen = read_record(SkyPilotSettings, settings) self._runs = backend_state_dir(chosen.state_dir, NAME).resolve() self._settings = replace(chosen, state_dir=str(self._runs.parent.parent)) - self._transport = _transport(self._settings) + self._transport = _transport(self._settings, self._runs) def submit(self, spec: JobSpec, *, project: Path) -> RunHandle: """Write the task file, record the launch, and launch it on a cluster of its own. @@ -281,11 +285,17 @@ def task_for( extras = " ".join(f"--extra {shlex.quote(extra)}" for extra in spec.extras) setup = [ "set -e", + *self._settings.setup_commands, + f"mkdir -p {OUTPUTS_PATH}", "command -v uv >/dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh", UV_ON_PATH, f"uv sync --frozen {extras}".rstrip(), ] - run = [UV_ON_PATH, self._transport.wrap_worker(worker, spec)] + run = [ + UV_ON_PATH, + "export PYTHONDONTWRITEBYTECODE=1", + self._transport.wrap_worker(worker, spec), + ] return { "name": _cluster_name(run_id), "workdir": str(project), @@ -322,11 +332,14 @@ def _sdk(self, *arguments: str) -> str: return transports.run_tool([*self._bridge(), *arguments]) -def _transport(settings: SkyPilotSettings) -> Transport: - """The run's transport: the buckets the settings name.""" - return BucketTransport( - BucketSettings(outputs_bucket=settings.outputs_bucket, buckets=settings.buckets) - ) +def _transport(settings: SkyPilotSettings, runs: Path) -> Transport: + """The configured transport; exactly one of ``bucket`` and ``copy`` is set.""" + if settings.bucket is not None and settings.copy is None: + return BucketTransport(settings.bucket, infra=settings.infra) + if settings.copy is not None and settings.bucket is None: + return CopyTransport(settings.copy, worker_python=settings.worker_python, runs=runs) + msg = "configure exactly one transport: [...skypilot.bucket] or [...skypilot.copy]" + raise ValueError(msg) def _resources(spec: JobSpec, settings: SkyPilotSettings) -> dict[str, JsonValue]: diff --git a/src/substrax/compute/skypilot_transports.py b/src/substrax/compute/skypilot_transports.py index 1d73ad0..4ffce3a 100644 --- a/src/substrax/compute/skypilot_transports.py +++ b/src/substrax/compute/skypilot_transports.py @@ -1,25 +1,33 @@ """How a SkyPilot run's inputs reach its cluster and its outputs come back. -The ``skypilot`` backend runs every job the same way and leaves this to its transport. The one -transport, :class:`BucketTransport`, keeps inputs and outputs on object-store buckets SkyPilot -mounts (``mode: MOUNT``), so the outputs outlive the cluster, which SkyPilot tears down when the -job ends. Outputs are fetched with the bucket's own tool: ``gcloud storage`` for ``gs://``, -``aws s3`` for ``s3://``. +The ``skypilot`` backend runs every job the same way and leaves this to its transport, one of: + +- :class:`BucketTransport`: inputs and outputs are object-store buckets SkyPilot mounts + (``mode: MOUNT``), so the outputs outlive the cluster, which SkyPilot tears down when the job + ends. Outputs are fetched with the bucket's own tool: ``gcloud storage`` for ``gs://``, + ``aws s3`` for ``s3://``. +- :class:`CopyTransport`: SkyPilot copies explicit local inputs onto the cluster, read-only, and + the outputs are fetched over SSH and rsync, verified against complete manifests of the remote + directory taken before and after the copy. The cluster has to stay up until then, so it is + removed by an explicit teardown, or by its autostop timer, which starts after setup and is not + a provisioning deadline or a cost cap. The worker runs under bubblewrap, which the image has to + provide with working user namespaces. Providers without bucket mounts (RunPod, Vast) use it. """ from __future__ import annotations import json import shlex -import subprocess # nosec B404 - runs the bucket tools, no shell +import subprocess # nosec B404 - runs ssh, rsync and the bucket tools, no shell import tempfile from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path from typing import Protocol +from substrax.compute import sky_transfer from substrax.compute.backend import RunHandle, RunState -from substrax.compute.spec import JobSpec +from substrax.compute.spec import JobSpec, MountAccess from substrax.compute.worker import JobManifest, MANIFEST_NAME, read_manifest from substrax.records import UNKNOWN_FIELDS_REFUSED from substrax.typing import JsonValue @@ -32,6 +40,9 @@ "gs://": ("gcloud", "storage", "rsync", "--recursive"), "s3://": ("aws", "s3", "sync"), } +# Providers SkyPilot cannot mount a bucket on. +_COPY_ONLY_PROVIDERS = frozenset({"runpod", "vast"}) +_RECOVERED = "recovered.json" @dataclass(frozen=True, slots=True, kw_only=True) @@ -49,6 +60,25 @@ class BucketSettings: buckets: Mapping[str, str] = field(default_factory=dict) +@dataclass(frozen=True, slots=True, kw_only=True) +class CopySettings: + """``[tool.substrax.compute.backends.skypilot.copy]``: local inputs copied to the cluster. + + Attributes: + sources: The local path behind each of a job's mounts, by mount name; read-only. + disk_size: The disk the cluster is asked for, in GB. + lifetime_minutes: How long the cluster may idle after setup before it removes itself. + network_isolation: Run the worker without network access. + """ + + __pydantic_config__ = UNKNOWN_FIELDS_REFUSED + + sources: Mapping[str, str] = field(default_factory=dict) + disk_size: int = 60 + lifetime_minutes: int = 60 + network_isolation: bool = False + + class Transport(Protocol): """The part of a SkyPilot run that depends on how its files move.""" @@ -88,15 +118,20 @@ class BucketTransport: # The cluster may still be running the job while its queue cannot be read. recovers_without_queue = False - def __init__(self, settings: BucketSettings) -> None: # noqa: DOC502 # the ValueError is raised by downloader - """Check the settings. + def __init__(self, settings: BucketSettings, *, infra: str | None) -> None: + """Check the settings against where the run goes. Args: settings: The bucket settings. + infra: The backend's ``infra``. Raises: - ValueError: If the outputs bucket's scheme has no download tool. + ValueError: If the provider cannot mount buckets, or the outputs bucket's scheme has + no download tool. """ + if (infra or "").split("/")[0] in _COPY_ONLY_PROVIDERS: + msg = f"{infra} cannot mount buckets; configure copy transport instead" + raise ValueError(msg) downloader(settings.outputs_bucket) self._settings = settings @@ -167,6 +202,140 @@ def _outputs(self, handle: RunHandle) -> str: return f"{self._settings.outputs_bucket.rstrip('/')}/{handle.run_id}" +class CopyTransport: + """Local inputs copied to the cluster, outputs fetched over SSH; see the module documentation.""" + + down_with_job = False + # A verified copy of the outputs answers for the run once it is taken. + recovers_without_queue = True + + def __init__(self, settings: CopySettings, *, worker_python: str, runs: Path) -> None: + """Check the settings. + + Args: + settings: The copy settings. + worker_python: The command that starts Python in the cluster's synced environment. + runs: Where the backend keeps its runs' local files. + + Raises: + ValueError: If the disk size or the lifetime is not positive. + """ + if settings.disk_size <= 0 or settings.lifetime_minutes <= 0: + msg = "copy transport's disk_size and lifetime_minutes must be positive" + raise ValueError(msg) + self._settings = settings + self._worker_python = worker_python + self._runs = runs + + def add_to_task( + self, spec: JobSpec, resources: dict[str, JsonValue], mounts: dict[str, JsonValue] + ) -> None: + """Ask for the disk and the autostop timer, and copy each input mount's local source. + + Args: + spec: The job. + resources: The task's resources. + mounts: The task's file mounts. + + Raises: + ValueError: If a mount is writable; a copy cannot be written back. + LookupError: If a mount has no configured source. + """ + resources["disk_size"] = self._settings.disk_size + resources["autostop"] = { + "idle_minutes": self._settings.lifetime_minutes, + "down": True, + "wait_for": "none", + } + for mount in spec.mounts: + if mount.access != MountAccess.READ_ONLY: + msg = f"mount {mount.name!r} is writable; copy transport copies inputs read-only" + raise ValueError(msg) + if mount.name not in self._settings.sources: + msg = f"mount {mount.name!r} has no local source in the backend's `sources`" + raise LookupError(msg) + mounts[mount.path] = str(Path(self._settings.sources[mount.name]).resolve(strict=True)) + + def wrap_worker(self, worker: str, spec: JobSpec) -> str: + """Run the worker under bubblewrap: the system, the project and the inputs read-only. + + Args: + worker: The worker command. + spec: The job; each mount is bound read-only. + + Returns: + The bubblewrap command around ``worker``. + """ + flags = ["--die-with-parent", "--unshare-user", "--unshare-pid"] + if self._settings.network_isolation: + flags.append("--unshare-net") + binds = [ + ["--ro-bind", "/", "/"], + ["--dev-bind", "/dev", "/dev"], + ["--proc", "/proc"], + ["--bind", OUTPUTS_PATH, OUTPUTS_PATH], + ["--bind", "/tmp", "/tmp"], # noqa: S108 # nosec B108 - the sandbox's scratch space + *(["--ro-bind", mount.path, mount.path] for mount in spec.mounts), + ["--ro-bind", SPEC_PATH, SPEC_PATH], + ] + bound = " ".join([shlex.join(["bwrap", *flags]), *(shlex.join(b) for b in binds)]) + # The project is the cluster's working directory, which only the shell knows. + return f'{bound} --ro-bind "$PWD" "$PWD" -- {worker}' + + def fetch(self, handle: RunHandle, target: Path) -> None: + """Copy the outputs over rsync and verify them against the remote directory's manifest. + + Args: + handle: The run. + target: The directory to copy into. + + Raises: + ValueError: If the remote outputs changed during the copy, or the copy differs from + them; the cluster is left up. + """ + cluster, remote = handle.details["cluster"], f"{OUTPUTS_PATH}/{handle.run_id}" + before = self._remote_manifest(cluster, remote) + run_tool( + ["rsync", "--archive", "--protect-args", "--", f"{cluster}:{remote}/", str(target)] + ) + if self._remote_manifest(cluster, remote) != before: + msg = f"the outputs of {handle.run_id} changed during the copy; the cluster is left up" + raise ValueError(msg) + sky_transfer.verify(target, before) + receipt = {"path": str(target.resolve()), "manifest": before} + (self._runs / handle.run_id / _RECOVERED).write_text(json.dumps(receipt), encoding="utf-8") + + def state_from_outputs(self, handle: RunHandle) -> RunState: + """The state of the run's verified local copy of its outputs. + + Args: + handle: The run. + + Returns: + Its state. + + Raises: + RuntimeError: If no verified copy exists, or it records an unfinished job. + """ + receipt = self._runs / handle.run_id / _RECOVERED + if not receipt.is_file(): + msg = f"{handle.run_id} has no live status and no verified copy of its outputs" + raise RuntimeError(msg) + recovered = json.loads(receipt.read_text(encoding="utf-8")) + root = Path(recovered["path"]) + sky_transfer.verify(root, recovered["manifest"]) + state = _finished_state(_read(root / MANIFEST_NAME)) + if state is None: + msg = f"the verified copy of {handle.run_id} records an unfinished job" + raise RuntimeError(msg) + return state + + def _remote_manifest(self, cluster: str, remote: str) -> sky_transfer.Manifest: + program = [*shlex.split(self._worker_python), "-m", "substrax.compute.sky_transfer", remote] + command = f"{UV_ON_PATH}; cd ~/sky_workdir && {shlex.join(program)}" + return json.loads(run_tool(["ssh", cluster, command])) + + def run_tool(command: Sequence[str]) -> str: """Run a command line tool and return its standard output. diff --git a/tests/compute/test_sky_transfer.py b/tests/compute/test_sky_transfer.py new file mode 100644 index 0000000..27642c5 --- /dev/null +++ b/tests/compute/test_sky_transfer.py @@ -0,0 +1,36 @@ +"""Copied artifacts must match a complete remote manifest before teardown.""" + +from pathlib import Path + +import pytest + +from substrax.compute.sky_transfer import snapshot, verify + + +def test_snapshot_and_verified_copy(tmp_path: Path) -> None: + """Files and links are distinct identities; links are never followed.""" + root = tmp_path / "output" + root.mkdir() + (root / "value").write_bytes(b"result") + (root / "link").symlink_to("/unavailable/input") + manifest = snapshot(root) + verify(root, manifest) + assert manifest["link"] == {"kind": "symlink", "target": "/unavailable/input"} + (root / "value").write_bytes(b"changed") + with pytest.raises(ValueError, match="does not match"): + verify(root, manifest) + + +@pytest.mark.parametrize("name", ["../escape", "/absolute", "x/../escape"]) +def test_unsafe_remote_manifest_is_rejected(tmp_path: Path, name: str) -> None: + """Remote metadata cannot select arbitrary local verification targets.""" + with pytest.raises(ValueError, match="unsafe path"): + verify(tmp_path, {name: {"kind": "symlink", "target": "anything"}}) + + +def test_extra_or_missing_file_is_rejected(tmp_path: Path) -> None: + """A subset comparison is insufficient for a complete recovery receipt.""" + manifest = snapshot(tmp_path) + (tmp_path / "unexpected").write_text("extra") + with pytest.raises(ValueError, match="does not match"): + verify(tmp_path, manifest) diff --git a/tests/compute/test_skypilot_backend.py b/tests/compute/test_skypilot_backend.py index 4f9475e..4261a51 100644 --- a/tests/compute/test_skypilot_backend.py +++ b/tests/compute/test_skypilot_backend.py @@ -12,6 +12,7 @@ import json import os import shlex +import shutil import subprocess import sys from pathlib import Path @@ -35,6 +36,7 @@ new_run, RunHandle, RunState, + stage_run, ) from substrax.compute.skypilot_backend import SkyPilotBackend from substrax.testing.compute import BackendContract @@ -42,6 +44,7 @@ FAKE = Path(__file__).with_name("fake_sky") +BUCKET: dict[str, JsonValue] = {"outputs_bucket": "gs://outputs"} @pytest.fixture(autouse=True) @@ -55,14 +58,15 @@ def fake_sky(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: def _backend(state_dir: Path, **settings: JsonValue) -> SkyPilotBackend: - """A backend on the fake SDK.""" + """A backend on the fake SDK; bucket transport unless ``copy`` is given.""" + transport: dict[str, JsonValue] = {} if "copy" in settings else {"bucket": BUCKET} return SkyPilotBackend( { - "outputs_bucket": "gs://outputs", "sky": shlex.join([sys.executable, str(FAKE / "sky.py")]), "worker_python": sys.executable, "sky_python": sys.executable, "state_dir": str(state_dir), + **transport, **settings, } ) @@ -82,6 +86,14 @@ def _task_file(state_dir: Path) -> dict[str, Any]: return json.loads(path.read_text(encoding="utf-8")) +def _task(backend: SkyPilotBackend, state_dir: Path, spec: JobSpec, project: Path) -> Any: + """The task ``backend`` builds for ``spec``, as its task file would hold it.""" + run = new_run(spec, project, settings={}) + run_dir = stage_run(_runs(state_dir), run) + task = backend.task_for(spec, project=run.project, run_id=run.run_id, run_dir=run_dir) + return json.loads(json.dumps(task)) + + def _spec(**changes: object) -> JobSpec: base = JobSpec( name="examples", @@ -143,7 +155,9 @@ def test_no_cpu_or_memory_is_asked_for_when_the_job_requests_none(tmp_path: Path def test_each_mount_is_its_configured_bucket(tmp_path: Path) -> None: (tmp_path / "project").mkdir() mount = Mount(name="datasets", path="/data", access=MountAccess.READ_ONLY) - _backend(tmp_path / "state", buckets={"datasets": "gs://demo-data"}).submit( + bucket: dict[str, JsonValue] = {**BUCKET, "buckets": {"datasets": "gs://demo-data"}} + + _backend(tmp_path / "state", bucket=bucket).submit( _spec(mounts=(mount,)), project=tmp_path / "project" ) @@ -161,7 +175,31 @@ def test_a_mount_without_a_configured_bucket_is_refused(tmp_path: Path) -> None: def test_an_outputs_bucket_without_a_download_tool_is_refused(tmp_path: Path) -> None: with pytest.raises(ValueError, match="gs:// or s3://"): - _backend(tmp_path / "state", outputs_bucket="azure://outputs") + _backend(tmp_path / "state", bucket={"outputs_bucket": "azure://outputs"}) + + +@pytest.mark.parametrize( + "transports", + [{}, {"bucket": BUCKET, "copy": {}}], + ids=["neither", "both"], +) +def test_exactly_one_transport_is_configured( + tmp_path: Path, transports: dict[str, JsonValue] +) -> None: + with pytest.raises(ValueError, match="exactly one transport"): + SkyPilotBackend({"state_dir": str(tmp_path / "state"), **transports}) + + +def test_network_isolation_is_a_copy_transport_setting(tmp_path: Path) -> None: + """Bucket transport's runtime is left as it is: it has no isolation setting to turn on.""" + with pytest.raises(ValueError, match="network_isolation"): + _backend(tmp_path / "state", bucket={**BUCKET, "network_isolation": True}) + + +def test_a_provider_that_cannot_mount_buckets_is_refused_bucket_transport(tmp_path: Path) -> None: + """RunPod cannot mount the object stores bucket transport needs.""" + with pytest.raises(ValueError, match="copy transport"): + _backend(tmp_path / "state", infra="runpod") def test_cancel_attempts_teardown_when_status_is_unavailable( @@ -217,6 +255,45 @@ def test_cancel_of_a_finished_run_leaves_its_cluster_and_outputs( assert not calls +def test_copy_task_has_only_explicit_sources_and_read_only_worker(tmp_path: Path) -> None: + """Local staging must not masquerade as writable mounted object storage.""" + data = tmp_path / "inputs" + data.mkdir() + copy: dict[str, JsonValue] = {"sources": {"data": str(data)}, "disk_size": 60} + backend = _backend(tmp_path / "state", infra="runpod", copy=copy) + spec = _spec(mounts=(Mount(name="data", path="/input", access=MountAccess.READ_ONLY),)) + + task = _task(backend, tmp_path / "state", spec, tmp_path) + + assert task["file_mounts"]["/input"] == str(data) + assert "/outputs" not in task["file_mounts"] + assert "--ro-bind / /" in task["run"] + assert "bwrap" in task["run"] + assert task["resources"]["disk_size"] == 60 + assert task["resources"]["autostop"]["wait_for"] == "none" + + +def test_copy_refuses_writable_inputs(tmp_path: Path) -> None: + """COPY cannot provide durable write-back for a declared writable mount.""" + backend = _backend(tmp_path / "state", infra="runpod", copy={}) + spec = _spec(mounts=(Mount(name="data", path="/data", access=MountAccess.READ_WRITE),)) + with pytest.raises(ValueError, match="read-only"): + _task(backend, tmp_path / "state", spec, tmp_path) + + +def test_copy_refuses_a_mount_without_a_local_source(tmp_path: Path) -> None: + backend = _backend(tmp_path / "state", copy={}) + spec = _spec(mounts=(Mount(name="data", path="/data", access=MountAccess.READ_ONLY),)) + with pytest.raises(LookupError, match="'data'"): + _task(backend, tmp_path / "state", spec, tmp_path) + + +@pytest.mark.parametrize("limit", ["disk_size", "lifetime_minutes"]) +def test_copy_limits_must_be_positive(tmp_path: Path, limit: str) -> None: + with pytest.raises(ValueError, match="positive"): + _backend(tmp_path / "state", copy={limit: 0}) + + def test_unknown_queue_state_is_not_a_scientific_failure( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -242,6 +319,94 @@ def unreadable(command: list[str]) -> str: _backend(tmp_path / "state").status(_handle(tmp_path)) +def _finished_manifest(directory: Path) -> None: + directory.mkdir() + (directory / "manifest.json").write_text( + json.dumps( + { + "schema_version": 2, + "job": "examples", + "finished": True, + "tasks": [ + { + "name": "t", + "argv": ["python", "-c", "pass"], + "status": "succeeded", + "returncode": 0, + "seconds": 1.0, + "outputs": [], + } + ], + "devices": None, + } + ) + ) + + +def test_copy_recovery_survives_queue_outage( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Verified local outputs remain usable after the cluster disappears.""" + backend = _backend(tmp_path / "state", copy={}) + run = new_run(_spec(), tmp_path, settings={}) + stage_run(_runs(tmp_path / "state"), run) + handle = run.handle("skypilot", {"cluster": "owned", "job_id": "1"}) + remote = tmp_path / "remote" + _finished_manifest(remote) + remote_manifest = json.dumps(skypilot_transports.sky_transfer.snapshot(remote)) + calls: list[list[str]] = [] + + def command(argv: list[str]) -> str: + calls.append(argv) + if argv[0] == "ssh": + return remote_manifest + if argv[0] == "rsync": + shutil.copytree(remote, argv[-1], dirs_exist_ok=True) + return "" + raise RuntimeError("cluster gone") + + monkeypatch.setattr(skypilot_transports, "run_tool", command) + recovered = backend.fetch(handle, tmp_path / "recovered") + assert recovered.is_dir() + assert [argv[0] for argv in calls] == ["ssh", "rsync", "ssh"] + assert backend.status(handle) == RunState.SUCCEEDED + + +def test_copy_changed_export_does_not_write_recovery_receipt( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A changing remote output cannot authorize cleanup as recovered.""" + backend = _backend(tmp_path / "state", copy={}) + run = new_run(_spec(), tmp_path, settings={}) + stage_run(_runs(tmp_path / "state"), run) + handle = run.handle("skypilot", {"cluster": "owned", "job_id": "1"}) + answers = iter(["{}", "", '{"new":{"kind":"directory"}}']) + monkeypatch.setattr(skypilot_transports, "run_tool", lambda _: next(answers)) + with pytest.raises(ValueError, match="changed"): + backend.fetch(handle, tmp_path / "recovered") + assert not (_runs(tmp_path / "state") / run.run_id / "recovered.json").exists() + + +def test_bucket_read_only_mount_does_not_enable_sandbox(tmp_path: Path) -> None: + """Existing bucket users must not acquire an undeclared runtime or network restriction.""" + bucket: dict[str, JsonValue] = {**BUCKET, "buckets": {"data": "gs://data"}} + backend = _backend(tmp_path / "state", bucket=bucket) + spec = _spec(mounts=(Mount(name="data", path="/data", access=MountAccess.READ_ONLY),)) + task = _task(backend, tmp_path / "state", spec, tmp_path) + assert "bwrap" not in task["run"] + assert "--unshare-net" not in task["run"] + + +@pytest.mark.parametrize("isolated", [False, True]) +def test_copy_network_isolation_is_explicit(tmp_path: Path, isolated: bool) -> None: + """Read-only input access alone does not imply network isolation.""" + copy: dict[str, JsonValue] = {"sources": {"data": str(tmp_path)}, "network_isolation": isolated} + backend = _backend(tmp_path / "state", copy=copy) + spec = _spec(mounts=(Mount(name="data", path="/data", access=MountAccess.READ_ONLY),)) + task = _task(backend, tmp_path / "state", spec, tmp_path) + assert ("--unshare-net" in task["run"]) is isolated + + @pytest.mark.parametrize("provider_status", ["FAILED", "FAILED_SETUP", "FAILED_DRIVER"]) def test_native_failure_statuses_are_terminal( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, provider_status: str diff --git a/tests/compute/test_skypilot_isolation.py b/tests/compute/test_skypilot_isolation.py new file mode 100644 index 0000000..98e2d99 --- /dev/null +++ b/tests/compute/test_skypilot_isolation.py @@ -0,0 +1,61 @@ +"""Exercise the generated mount boundary with the real local bubblewrap binary.""" + +import errno +import json +import shlex +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +from substrax.compute import JobSpec, Mount, MountAccess, Task +from substrax.compute.backend import backend_state_dir, new_run, stage_run +from substrax.compute.skypilot_backend import SkyPilotBackend +from substrax.compute.skypilot_transports import SPEC_PATH + + +def test_generated_read_only_boundary(tmp_path: Path) -> None: + """The worker sees read-only inputs, writable outputs and a new network namespace.""" + if sys.platform != "linux" or shutil.which("bwrap") is None: + pytest.skip("Linux bubblewrap integration") + inputs, outputs = tmp_path / "input", tmp_path / "output" + inputs.mkdir() + outputs.mkdir() + backend = SkyPilotBackend( + { + "state_dir": str(tmp_path / "state"), + "copy": {"sources": {"data": str(inputs)}, "network_isolation": True}, + } + ) + spec = JobSpec( + name="isolation", + timeout_seconds=10, + tasks=(Task(name="probe", argv=("true",), timeout_seconds=5),), + mounts=(Mount(name="data", path=str(inputs), access=MountAccess.READ_ONLY),), + ) + run = new_run(spec, Path.cwd(), settings={}) + run_dir = stage_run(backend_state_dir(str(tmp_path / "state"), "skypilot"), run) + task = backend.task_for(spec, project=run.project, run_id=run.run_id, run_dir=run_dir) + boundary = str(task["run"]).splitlines()[-1].split(" -- ", 1)[0] + boundary = boundary.replace("/outputs", str(outputs)).replace( + SPEC_PATH, str(run_dir / "spec.json") + ) + body = ( + "import errno,os; from pathlib import Path; " + f"root=Path({str(inputs)!r}); out=Path({str(outputs)!r})\n" + "try: (root/'forbidden').write_text('bad')\n" + "except OSError as e: assert e.errno in (errno.EROFS,errno.EACCES,errno.EPERM)\n" + "else: raise AssertionError('input writable')\n" + "(out/'proof').write_text(os.readlink('/proc/self/ns/net'))\n" + ) + command = boundary + " -- " + shlex.join([sys.executable, "-c", body]) + result = subprocess.run( # noqa: S603 - generated boundary with quoted fixture arguments + ["/bin/bash", "-c", command], capture_output=True, timeout=10, check=False + ) + assert result.returncode == 0, json.dumps( + {"stderr": result.stderr.decode(), "errno": errno.EPERM} + ) + assert (outputs / "proof").read_text() != str(Path("/proc/self/ns/net").readlink()) + assert not (inputs / "forbidden").exists()