simrun 0.1.0__py3-none-any.whl

simrun/service/logs.py

@@ -0,0 +1,91 @@
+"""The logs of the pods of a task, kept in the run store.
+
+The logs of a pod are gone with the pod: Kueue deletes a preempted pod, Argo deletes
+the pods of a workflow with it after its time to live, and stopping a workflow ends
+its pods. So the controller copies the logs of the main containers of the pods of a
+task into `tasks/<task>/logs/` once the task is final and its pods terminated, the
+ones which ended before it stops or cancels a workflow (the last `KEPT_TAIL_LINES`
+lines, at most what the backend reads of a log): a file `<pod>.log` per pod (every
+attempt the backend can still read) and `index.json`, the pods in their order,
+written last, which marks the logs as kept. An ended run is followed until its logs
+are kept or its workflow is gone (`RunRecord.pending_logs`); a run which never had a
+workflow keeps none. The API serves the kept logs first and the logs of the live pods
+otherwise.
+"""
+
+import json
+import os
+import uuid
+from collections import deque
+from pathlib import Path
+
+from simrun.service.backend import PodLog
+
+LOGS_DIR = "logs"
+INDEX_FILE = "index.json"
+# the lines of every pod an agent or a person gets by default
+DEFAULT_TAIL_LINES = 300
+# the last lines of every pod kept in the store
+KEPT_TAIL_LINES = 100_000
+
+
+def _write(path: Path, text: str) -> None:
+    """Write a file atomically (a temporary file in its directory, renamed)."""
+    tmp = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
+    tmp.write_text(text)
+    os.replace(tmp, path)
+
+
+def kept(task_dir: Path) -> bool:
+    """Whether the logs of a task are kept in its directory."""
+    return (task_dir / LOGS_DIR / INDEX_FILE).is_file()
+
+
+def keep(task_dir: Path, logs: list[PodLog]) -> None:
+    """Keep the logs of the pods of a task in `task_dir/logs`.
+
+    Args:
+        task_dir: The directory of the task.
+        logs: The logs of its pods, oldest first.
+    """
+    directory = task_dir / LOGS_DIR
+    directory.mkdir(parents=True, exist_ok=True)
+    index = []
+    for log in logs:
+        name = f"{log.pod}.log"
+        _write(directory / name, log.text)
+        index.append({"node": log.node, "pod": log.pod, "file": name})
+    _write(directory / INDEX_FILE, json.dumps(index, indent=2))
+
+
+def _tail(path: Path, lines: int | None) -> str:
+    with path.open(encoding="utf-8", errors="replace") as handle:
+        if lines is None:
+            return handle.read()
+        return "".join(deque(handle, maxlen=lines))
+
+
+def read_kept(task_dir: Path, tail_lines: int | None = None) -> list[PodLog] | None:
+    """The kept logs of a task, the last `tail_lines` lines of each.
+
+    Returns:
+        The logs, None when they are not kept (yet).
+    """
+    directory = task_dir / LOGS_DIR
+    try:
+        index = json.loads((directory / INDEX_FILE).read_text())
+    except FileNotFoundError:
+        return None
+    logs = []
+    for entry in index:
+        try:
+            text = _tail(directory / entry["file"], tail_lines)
+        except FileNotFoundError:
+            text = "(the kept log is missing)\n"
+        logs.append(PodLog(node=entry["node"], pod=entry["pod"], text=text))
+    return logs
+
+
+def render(logs: list[PodLog]) -> str:
+    """The logs of the pods of a task as one text, each pod under a header."""
+    return "\n".join(f"==> {log.node} ({log.pod}) <==\n{log.text}" for log in logs)

simrun/service/submission.py

@@ -0,0 +1,349 @@
+"""The submission of a run in the background: its environment, image and workflow.
+
+`POST /runs` validates and plans a run, writes its record `submitting` (with the
+deterministic name of its workflow and the settings of its tasks) and returns its id
+at once. The controller then runs `Submitter.run` in its pool of submissions
+(`Controller.start_submission`): it resolves the environment of the run, builds its
+image when the registry does not have it (once per image: a second submission of the
+same new environment waits for the first build and finds the image), records the
+image and the time it was ready (`built`, from which the grace of a submission
+counts), reserves the node of the run and submits its workflow; the run is queued
+then.
+
+A failure lands in the record: the run fails with a message, its node is released and
+a workflow a failed submission may have created is cancelled. A run cancelled while it
+is resolved or built is not submitted; one cancelled while its workflow is submitted
+has the workflow cancelled; the workflow of one cancelled and deleted meanwhile is
+cancelled as an orphan.
+"""
+
+import dataclasses
+import logging
+import shutil
+import tempfile
+import threading
+from collections.abc import Generator
+from contextlib import contextmanager, suppress
+from datetime import UTC, datetime
+from pathlib import Path
+
+import simrun
+from simrun.argo import to_argo
+from simrun.config import ClusterConfig
+from simrun.environment import (
+    BUILD_CONSTRAINTS,
+    Environment,
+    dockerfile,
+    image_ref,
+    resolve_environment,
+)
+from simrun.planner import PodShape, TaskPlan
+from simrun.plugins import TaskType
+from simrun.service.backend import Backend
+from simrun.service.controller import TERMINAL_RUN_STATES, Controller, closed
+from simrun.snapshot import SourceProvenance
+from simrun.store import RunRecord, RunState, RunStore
+
+logger = logging.getLogger(__name__)
+
+ENVIRONMENT_LOCK = "environment.lock"
+
+
+class SubmissionError(Exception):
+    """A submission which cannot go on; the message is the one of the record."""
+
+
+def has_model_project(store: Path, sources: list[SourceProvenance]) -> bool:
+    """Whether the model of a run (its first source) has a `pyproject.toml`."""
+    return bool(sources) and (store / sources[0].snapshot / "pyproject.toml").is_file()
+
+
+def task_plans(record: RunRecord) -> dict[str, TaskPlan]:
+    """The plans of the tasks of a run from its record."""
+    return {
+        task_id: dataclasses.replace(TaskPlan(**plan), shape=PodShape(**plan["shape"]))
+        for task_id, plan in record.plans.items()
+    }
+
+
+def _undo(current: RunRecord, record: RunRecord) -> dict[str, object]:
+    """What a run which ended during its submission still has to undo.
+
+    The workflow is cancelled, and the node released: the submission may have
+    reserved it after the end of the run released it (a release removes only the
+    reservation of its run, so it may be repeated).
+    """
+    return {
+        "pending_cancel": True,
+        "pending_release": current.pending_release or record.reserve,
+        "pending_logs": True,
+    }
+
+
+class Submitter:
+    """Resolves, builds and submits the runs written `submitting`."""
+
+    def __init__(
+        self,
+        store: RunStore,
+        backend: Backend,
+        types: dict[str, TaskType],
+        controller: Controller,
+        cluster: ClusterConfig,
+        python: str,
+        simrun_wheel: Path | None,
+        build_constraints: Path | None = None,
+    ) -> None:
+        """Create the submitter.
+
+        Args:
+            store: The run store.
+            backend: The backend the workflows run on.
+            types: The task types by name.
+            controller: The controller, whose lock guards the records.
+            cluster: The cluster.
+            python: Python version of the images.
+            simrun_wheel: The simrun wheel the images install.
+            build_constraints: The pinned versions of the build backends, the ones
+                of simrun by default.
+        """
+        self.store = store
+        self.backend = backend
+        self.types = types
+        self.controller = controller
+        self.cluster = cluster
+        self.python = python
+        self.simrun_wheel = simrun_wheel
+        self.build_constraints = build_constraints or BUILD_CONSTRAINTS
+        self._reserve_lock = threading.Lock()
+        self._builds_lock = threading.Lock()
+        self._builds: dict[str, threading.Lock] = {}
+
+    @contextmanager
+    def _building(self, ref: str) -> Generator[None]:
+        """Hold the lock of an image, so it is built once at a time."""
+        with self._builds_lock:
+            lock = self._builds.setdefault(ref, threading.Lock())
+        with lock:
+            yield
+
+    def environment(self, record: RunRecord) -> tuple[str | None, str | None]:
+        """The image of a run and the lock of its environment, built when missing.
+
+        A run whose model has no `pyproject.toml` keeps the image it was written
+        with (the default image of the service).
+
+        Raises:
+            SubmissionError: The environment cannot be resolved, the service has no
+                simrun wheel to build with, or the build fails.
+        """
+        if not has_model_project(self.store.root, record.sources):
+            return record.image, None
+        model, *libraries = record.sources
+        try:
+            env = resolve_environment(
+                self.store.root,
+                model,
+                libraries,
+                python=self.python,
+                simrun_version=simrun.__version__,
+                build_constraints=self.build_constraints,
+            )
+        except (ValueError, RuntimeError) as error:
+            raise SubmissionError(
+                f"the environment cannot be resolved: {error}"
+            ) from error
+        ref = image_ref(env, self.cluster.registry)
+        with self._building(ref):
+            try:
+                missing = not self.backend.has_image(ref)
+            except Exception as error:
+                raise SubmissionError(
+                    f"the registry cannot tell whether it has {ref}: {error}"
+                ) from error
+            if missing:
+                self._build(env, ref)
+        return ref, env.lock
+
+    def _build(self, env: Environment, ref: str) -> None:
+        wheel = self.simrun_wheel
+        if wheel is None:
+            raise SubmissionError("the service has no simrun_wheel to build with")
+        with tempfile.TemporaryDirectory(prefix="simrun-build-") as tmp:
+            context = Path(tmp)
+            (context / "Dockerfile").write_text(
+                dockerfile(env, self.cluster.base_image_ref(), wheel.name)
+            )
+            shutil.copy2(wheel, context / wheel.name)
+            logger.info("building the image %s", ref)
+            try:
+                self.backend.build_image(context, ref)
+            except Exception as error:
+                raise SubmissionError(
+                    f"building the image {ref} failed: {error}"
+                ) from error
+
+    def run(self, run_id: str) -> None:
+        """Resolve, build and submit a run written `submitting`.
+
+        Args:
+            run_id: The run.
+        """
+        try:
+            record = self.store.read(run_id)
+        except KeyError:
+            return
+        if record.state != RunState.SUBMITTING:
+            return
+        try:
+            image, lock = self.environment(record)
+            built = record.model_copy(
+                update={"image": image, "built": datetime.now(UTC)}
+            )
+            if not self.controller.write_if_unchanged(record, built):
+                logger.info("run %s changed while it was built: not submitted", run_id)
+                return
+            record = built
+            if lock is not None:
+                (self.store.run_dir(run_id) / ENVIRONMENT_LOCK).write_text(lock)
+            try:
+                manifest = to_argo(record, task_plans(record), self.types, self.cluster)
+            except ValueError as error:
+                raise SubmissionError(str(error)) from error
+            if record.reserve is not None:
+                self._reserve(record.reserve, run_id)
+        except SubmissionError as error:
+            self._fail(record, str(error))
+            return
+        try:
+            self.backend.submit(manifest)
+        except Exception as error:
+            self._abort(record, error)
+            return
+        self._submitted(record)
+
+    def _reserve(self, node: str, run_id: str) -> None:
+        """Reserve the node of a run, unless another run took it meanwhile.
+
+        Raises:
+            SubmissionError: The node is unknown, reserved by another run, or its
+                reservation failed.
+        """
+        with self._reserve_lock:
+            try:
+                found = next((n for n in self.backend.nodes() if n.name == node), None)
+                if found is None:
+                    raise SubmissionError(f"unknown node {node!r} to reserve")
+                if found.reserved_by not in (None, run_id):
+                    raise SubmissionError(
+                        f"node {node!r} is reserved by run {found.reserved_by}"
+                    )
+                self.backend.reserve(node, run_id)
+            except SubmissionError:
+                raise
+            except Exception as error:
+                raise SubmissionError(
+                    f"reserving node {node!r} failed: {error}"
+                ) from error
+
+    def _fail(self, record: RunRecord, message: str) -> None:
+        """Fail a run whose submission cannot go on; nothing was submitted."""
+        logger.error("the submission of run %s failed: %s", record.id, message)
+        failed = closed(
+            record, RunState.FAILED, record.tasks, message, pending_logs=False
+        )
+        if self.controller.write_if_unchanged(record, failed):
+            self.controller.settle(record.id)
+
+    def _abort(self, record: RunRecord, error: Exception) -> None:
+        """Fail a run whose submission of the workflow failed, then settle it.
+
+        The workflow may have been created although the submission raised, so it is
+        cancelled like the node is released, retried by the controller on failure.
+        """
+        logger.error("submitting run %s failed: %s", record.id, error)
+        failed = closed(
+            record,
+            RunState.FAILED,
+            record.tasks,
+            f"submitting failed: {error}",
+            pending_cancel=True,
+        )
+        try:
+            if not self.controller.write_if_unchanged(record, failed):
+                # cancelled meanwhile: the cancellation is retried by the controller
+                self._cancel_if_pending(record)
+                return
+        except OSError:
+            logger.exception("writing the failed record of run %s failed", record.id)
+            if record.reserve is not None:
+                with suppress(Exception):
+                    self.backend.unreserve(record.reserve, record.id)
+            return
+        self.controller.settle(record.id)
+
+    def _cancel_if_pending(self, record: RunRecord) -> None:
+        """Cancel the workflow of a run which ended or was deleted meanwhile."""
+        with self.controller.lock:
+            try:
+                current: RunRecord | None = self.store.read(record.id)
+            except KeyError:
+                current = None
+            if current is not None and current.state in TERMINAL_RUN_STATES:
+                self.store.write(current.model_copy(update=_undo(current, record)))
+        if current is None:
+            self._cancel_orphan(record)
+        else:
+            self.controller.settle(record.id)
+
+    def _submitted(self, record: RunRecord) -> None:
+        """Queue a submitted run, or cancel its workflow when it ended meanwhile."""
+        ended = False
+        with self.controller.lock:
+            try:
+                current: RunRecord | None = self.store.read(record.id)
+            except KeyError:
+                current = None
+            if current is not None:
+                ended = current.state in TERMINAL_RUN_STATES
+                if ended:  # cancelled while it was submitted: cancel the workflow
+                    current = current.model_copy(update=_undo(current, record))
+                else:
+                    current = current.model_copy(update={"state": RunState.QUEUED})
+                self.store.write(current)
+        if current is None:
+            self._cancel_orphan(record)
+            return
+        if ended:
+            self.controller.settle(record.id)
+        logger.info("run %s submitted by %s", record.id, record.submitted_by)
+
+    def _cancel_orphan(self, record: RunRecord) -> None:
+        """Undo the submission of a run cancelled and deleted while it was submitted.
+
+        Its workflow is cancelled and its node released (the job may have reserved
+        it after the run was released and deleted); no record is left to retry
+        either, so both are best effort and logged.
+        """
+        name = record.backend_name
+        logger.warning(
+            "run %s was cancelled and deleted while it was submitted", record.id
+        )
+        if name is not None:
+            try:
+                self.backend.cancel(name)
+            except KeyError:
+                pass
+            except Exception:
+                logger.exception("cancelling the orphan workflow %s failed", name)
+        if record.reserve is not None:
+            try:
+                self.backend.unreserve(record.reserve, record.id)
+            except KeyError:
+                pass
+            except Exception:
+                logger.exception(
+                    "releasing node %s of the orphan run %s failed",
+                    record.reserve,
+                    record.id,
+                )

simrun/snapshot.py

@@ -0,0 +1,196 @@
+"""Snapshots of working trees with provenance.
+
+A snapshot is a copy of the files git knows of a working tree (tracked and
+untracked, not ignored, uncommitted changes included) in the run store. Files
+unchanged since the latest snapshot of the same name are hard links to it. The files
+of a snapshot are read-only (`a-w`, part of the hash of its content): a snapshot is
+named by its content and never changes, and a build backend which writes through a
+link of a shadow tree into one fails loudly instead.
+Submodules are out of scope: their content is not copied.
+"""
+
+import hashlib
+import logging
+import os
+import secrets
+import shutil
+import stat
+import subprocess
+import tempfile
+from pathlib import Path
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class SourceProvenance(BaseModel):
+    """Where the files of a snapshot come from.
+
+    Attributes:
+        name: Name of the source.
+        branch: Branch of the working tree, `None` for a detached HEAD.
+        commit: Commit of HEAD.
+        dirty: Whether the working tree differs from HEAD.
+        changed_files: Porcelain status lines, e.g. `M src/a.py`.
+        diff_sha256: Hash of the snapshot tree (paths, types, modes, contents);
+            the patch stays as provenance.
+        snapshot: Directory of the snapshot, relative to the store root.
+        patch: Stored `git diff HEAD --binary`, relative to the store root.
+    """
+
+    name: str
+    branch: str | None
+    commit: str
+    dirty: bool
+    changed_files: list[str]
+    diff_sha256: str
+    snapshot: str
+    patch: str | None
+
+
+def _run(cmd: list[str], input: bytes | None = None, ok: tuple[int, ...] = ()) -> bytes:
+    result = subprocess.run(cmd, input=input, capture_output=True)
+    if result.returncode == 23 and _only_missing_files(result.stderr):
+        return result.stdout  # a file vanished between the listing and the copy
+    if result.returncode != 0 and result.returncode not in ok:
+        raise RuntimeError(
+            f"{cmd[0]} failed with exit code {result.returncode}: "
+            f"{result.stderr.decode(errors='replace').strip()}"
+        )
+    return result.stdout
+
+
+def _only_missing_files(stderr: bytes) -> bool:
+    """Whether every error of rsync (exit 23) is a source file which is gone."""
+    errors = [
+        line
+        for line in stderr.decode(errors="replace").splitlines()
+        if line.startswith("rsync:")
+    ]
+    return bool(errors) and all(
+        "link_stat" in line and "No such file" in line for line in errors
+    )
+
+
+def _git(repo: Path, *args: str) -> bytes:
+    return _run(["git", "-C", str(repo), *args])
+
+
+def _tree_hash(root: Path) -> str:
+    """Hash the entries of a tree: path, type, mode bits, content or target."""
+    digest = hashlib.sha256()
+    for dirpath, dirnames, filenames in os.walk(root):
+        dirnames.sort()
+        # a symlink to a directory is listed in dirnames and not descended into
+        links = [d for d in dirnames if (Path(dirpath) / d).is_symlink()]
+        for fname in sorted(filenames + links):
+            path = Path(dirpath) / fname
+            st = path.lstat()
+            kind = b"l" if stat.S_ISLNK(st.st_mode) else b"f"
+            content = (
+                os.fsencode(os.readlink(path)) if kind == b"l" else path.read_bytes()
+            )
+            digest.update(os.fsencode(str(path.relative_to(root))) + b"\0")
+            digest.update(kind + b"%o" % stat.S_IMODE(st.st_mode) + b"\0")
+            digest.update(b"%d\0" % len(content) + content)
+    return digest.hexdigest()
+
+
+def _split(data: bytes) -> list[str]:
+    return [os.fsdecode(p) for p in data.split(b"\0") if p]
+
+
+def _write_atomic(path: Path, data: bytes, mode: int | None = None) -> None:
+    tmp = path.with_name(f".tmp-{secrets.token_hex(6)}-{path.name}")
+    tmp.write_bytes(data)
+    if mode is not None:
+        tmp.chmod(mode)
+    os.replace(tmp, path)
+
+
+def snapshot(repo: Path, store: Path, name: str) -> SourceProvenance:
+    """Copy a working tree into `store/snapshots/<name>`.
+
+    Args:
+        repo: A git working tree (a worktree with a `.git` file works too).
+        store: Root of the run store.
+        name: Name of the source.
+
+    Returns:
+        The provenance of the snapshot.
+
+    Raises:
+        RuntimeError: When rsync is not installed.
+    """
+    if shutil.which("rsync") is None:
+        raise RuntimeError("rsync is required to snapshot a working tree")
+    commit = _git(repo, "rev-parse", "HEAD").decode().strip()
+    branch = _git(repo, "branch", "--show-current").decode().strip() or None
+    status = _git(repo, "status", "--porcelain", "--untracked-files=all")
+    changed = [line for line in status.decode().splitlines() if line]
+    patch = _git(repo, "diff", "HEAD", "--binary")
+    listed = _split(
+        _git(repo, "ls-files", "-z", "--cached", "--others", "--exclude-standard")
+    )
+    dirty = bool(changed)
+
+    # skip deleted files and the directories of submodules
+    files = sorted(
+        {
+            f
+            for f in listed
+            if os.path.lexists(repo / f)
+            and (os.path.islink(repo / f) or not os.path.isdir(repo / f))
+        }
+    )
+
+    parent = store / "snapshots" / name
+    parent.mkdir(parents=True, exist_ok=True)
+    latest = parent / "latest"
+    tmp = Path(tempfile.mkdtemp(prefix=".tmp-", dir=parent))
+    try:
+        # --checksum: link only files of identical content; exit 24 (a file
+        # vanished before the copy) leaves it out of the snapshot and its hash
+        # --chmod: the files read-only, also the ones linked to the latest snapshot,
+        # whose modes then match
+        cmd = ["rsync", "-a", "--checksum", "--chmod=F-w", "--files-from=-"]
+        cmd.append("--from0")
+        if latest.is_dir():
+            cmd.append(f"--link-dest={latest.resolve()}")
+        cmd += [f"{repo}/", f"{tmp}/"]
+        _run(cmd, input=b"".join(os.fsencode(f) + b"\0" for f in files), ok=(24,))
+        diff_sha256 = _tree_hash(tmp)
+        dirname = f"{commit[:12]}-{diff_sha256[:12]}"
+        final = parent / dirname
+        tmp.chmod(0o755)
+        if not final.is_dir():
+            try:
+                os.rename(tmp, final)
+            except OSError:
+                if not final.is_dir():
+                    raise
+                logger.info("snapshot %s was created concurrently", final)
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)
+
+    patch_rel: str | None = None
+    if dirty:
+        patch_path = parent / f"{dirname}.patch"
+        if not patch_path.exists():
+            _write_atomic(patch_path, patch, mode=0o444)
+        patch_rel = str(patch_path.relative_to(store))
+
+    link = parent / f".tmp-link-{secrets.token_hex(6)}"
+    link.symlink_to(dirname)
+    os.replace(link, latest)
+    return SourceProvenance(
+        name=name,
+        branch=branch,
+        commit=commit,
+        dirty=dirty,
+        changed_files=changed,
+        diff_sha256=diff_sha256,
+        snapshot=str(final.relative_to(store)),
+        patch=patch_rel,
+    )