pygpu-scheduling 0.1.0__py3-none-any.whl

pygpu_scheduling-0.1.0.dist-info/METADATA

@@ -0,0 +1,93 @@
+Metadata-Version: 2.5
+Name: pygpu-scheduling
+Version: 0.1.0
+Summary: Reusable local CPU and NVIDIA GPU task scheduling with dependency graphs and recovery hooks
+Project-URL: Source, https://github.com/AllanJinYosoro/PyGPUScheduling
+Project-URL: Issues, https://github.com/AllanJinYosoro/PyGPUScheduling/issues
+Project-URL: Changelog, https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/CHANGELOG.md
+Author: PyGPUScheduling contributors
+License-Expression: MIT
+License-File: LICENSE
+Keywords: gpu,nvidia,scheduler,subprocess,workflow
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: POSIX
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.11
+Provides-Extra: progress
+Requires-Dist: tqdm>=4.66; extra == 'progress'
+Description-Content-Type: text/markdown
+
+# PyGPUScheduling
+
+单机 CPU 与 NVIDIA GPU 任务调度库。业务注册任务、依赖和生命周期处理器，调度器管理资源准入、子进程、重试及恢复。运行时仅依赖 Python 标准库，交互式进度条可选安装 `tqdm`。
+
+分发包名为 `pygpu-scheduling`，Python 导入名为 `task_scheduling`。发布状态见[维护指南](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/docs/maintaining.md#当前交接状态)。
+
+## 环境与安装
+
+需要 Python 3.11 或更新版本及 POSIX 系统。GPU 功能通过 `nvidia-smi` 工作，需要 NVIDIA 驱动；调度库不要求安装 CUDA Toolkit、Torch 或 Optuna。每个 GPU 任务使用一张卡，CPU 任务不查询 GPU。
+
+在本仓库开发：
+
+```bash
+uv sync --locked
+uv run pytest
+uv run python examples/basic.py
+```
+
+需要交互式进度条时：
+
+```bash
+uv sync --locked --extra progress
+```
+
+在其他项目中安装：
+
+```bash
+uv add "pygpu-scheduling>=0.1.0,<0.2.0"
+```
+
+GitHub 固定 tag 接入及本地联调方式见[维护指南](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/docs/maintaining.md#业务仓库接入)。
+
+## 注册任务
+
+```python
+import sys
+
+from task_scheduling import Scheduler, Scheduling, TaskSpec
+
+scheduler = Scheduler("runs/example", Scheduling(max_parallel_tasks=2))
+scheduler.submit_many([
+    TaskSpec("prepare", "command", payload={"args": [sys.executable, "prepare.py"]}),
+    TaskSpec("train", "command", dependencies=("prepare",),
+             payload={"args": [sys.executable, "train.py"]}),
+])
+records = scheduler.run()
+```
+
+`command` 是按退出码判定结果的内置处理器。`run()` 返回全部任务记录，调用方检查 `complete`、`failed`、`blocked` 等状态并决定业务退出码。
+
+GPU 任务通过 `resources=ResourceRequest("cuda", 4096)` 声明 4096 MiB 预算。预算用于启动准入，不是显存硬限制。调度器不解释模型、训练轮数、业务参数或产物文件。
+
+## 功能与边界
+
+- 静态依赖图与动态追加任务，拒绝缺失依赖、重复定义和依赖环。
+- 按优先级、显存预算和提交顺序选择任务，支持共享或独占 GPU。
+- 启动前预留显存，结合实际进程占用更新预算，保留每卡安全余量。
+- 子进程日志、超时、有限重试、GPU 冷却和故障隔离。
+- 注册 `prepare`、`assess`、`commit`、`recover`，由业务提供校验及产物流程。
+- 持久化任务状态，完成事件按至少一次语义重放。
+
+当前提供单机本地子进程执行，不包含多机、抢占或单任务多卡调度。预算账本属于单个调度实例；跨实例不会合并预算。更完整的约束见[API 与恢复协议](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/docs/api.md)。
+
+## 文档与示例
+
+- [API、调度策略和恢复协议](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/docs/api.md)
+- [开发与发布](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/docs/maintaining.md)
+- [变更记录](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/CHANGELOG.md)
+- [依赖顺序示例](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/examples/basic.py)
+- [业务产物处理示例](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/examples/artifacts.py)
+
+采用 [MIT 许可证](https://github.com/AllanJinYosoro/PyGPUScheduling/blob/main/LICENSE)。

pygpu_scheduling-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+task_scheduling/__init__.py,sha256=E3Sk8j90tz1ny5EN1XbLTJVsTkxcnhhg8SBXYTeGhsU,513
+task_scheduling/gpu.py,sha256=3rX6tfD0aLviDbdDx1osBTzDFSHDHeUSPGmPgnoC84w,3679
+task_scheduling/memory.py,sha256=JTF-B5HtmnISzrmv_Cc91FljCtix4nAkKkl1cGgPROc,1940
+task_scheduling/policy.py,sha256=-sMVPieIVx5qPxq6fhM-pT3ZyH6jHUF4mJn5IKeYZk0,2087
+task_scheduling/progress.py,sha256=kS4x7ZDMwZSxwGrnh_x_9SR_HX0GQyKcfWhTcPqy3p4,5555
+task_scheduling/runtime.py,sha256=YnHPljoZL-XaWBJWGvTQ2eZ5k-ma-0-q0Ilf-Y2JFI4,1332
+task_scheduling/scheduler.py,sha256=Y3gzKyD64DlubsMnAqcWnQ0cMOtewEjIX6QvV9GMwmc,21607
+task_scheduling/storage.py,sha256=q3gk8zt2uChTOPRlACQC0Z3fd8FAGl4GP7R0QD1f03g,1988
+task_scheduling/task.py,sha256=Mo0pxnImk12_S_1pnGOX0vdWxM3BalcgfXjKHwHz6MI,6848
+pygpu_scheduling-0.1.0.dist-info/METADATA,sha256=pADOEm7CSc3beNY5fwH5-kyZ14uTTVvODWVUbrzUJNM,4357
+pygpu_scheduling-0.1.0.dist-info/WHEEL,sha256=2iZrJkDIkeDh0k9SikKMqFuRP-ua6mb4Ir1AsER4BEw,87
+pygpu_scheduling-0.1.0.dist-info/licenses/LICENSE,sha256=Ni80rrOgBgnqG6v0XpmDABF5_g-w05-_v7mvlwrv6Hs,1085
+pygpu_scheduling-0.1.0.dist-info/RECORD,,

pygpu_scheduling-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.33.0
+Root-Is-Purelib: true
+Tag: py3-none-any

pygpu_scheduling-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 PyGPUScheduling contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

task_scheduling/__init__.py

@@ -0,0 +1,14 @@
+from .policy import Scheduling
+from .runtime import LocalRuntime
+from .scheduler import Scheduler
+from .task import (
+    Assessment, Attempt, Command, CommandHandler, Recovery, ResourceRequest,
+    TaskContext, TaskEvent, TaskHandler, TaskOutcome, TaskRecord, TaskSpec,
+)
+
+
+__all__ = [
+    "Assessment", "Attempt", "Command", "CommandHandler", "LocalRuntime", "Recovery",
+    "ResourceRequest", "Scheduler", "Scheduling", "TaskContext", "TaskEvent", "TaskHandler",
+    "TaskOutcome", "TaskRecord", "TaskSpec",
+]

task_scheduling/gpu.py

@@ -0,0 +1,90 @@
+from dataclasses import dataclass, field
+import csv
+import io
+import os
+import subprocess
+
+from .policy import Scheduling
+
+
+class GPUQueryError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class GPU:
+    index: int
+    uuid: str
+    free_mib: int
+    busy: bool = False
+    process_memory_mib: dict[int, int] = field(default_factory=dict)
+
+
+def _query(fields: str, kind: str) -> str:
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", f"--query-{kind}={fields}", "--format=csv,noheader,nounits"],
+            capture_output=True, text=True, timeout=5, check=False,
+        )
+    except (OSError, subprocess.TimeoutExpired) as error:
+        raise GPUQueryError(f"GPU query unavailable: {error}") from error
+    if result.returncode and (kind != "gpu" or not result.stdout.strip()):
+        raise GPUQueryError(result.stderr.strip() or "nvidia-smi query failed")
+    return result.stdout
+
+
+def discover() -> list[GPU]:
+    devices = []
+    for row in csv.reader(io.StringIO(_query("index,uuid,memory.free", "gpu"))):
+        try:
+            index, uuid, free = [part.strip() for part in row]
+            if uuid.startswith("GPU-"):
+                devices.append(GPU(int(index), uuid, int(free)))
+        except (ValueError, TypeError):
+            continue
+    if not devices:
+        raise GPUQueryError("nvidia-smi returned no healthy GPU records")
+    processes = {}
+    for row in csv.reader(io.StringIO(_query("gpu_uuid,pid,used_gpu_memory", "compute-apps"))):
+        if len(row) == 3 and row[0].strip().startswith("GPU-"):
+            try:
+                uuid, pid, memory = [part.strip() for part in row]
+                used = 0 if memory in {"N/A", "[N/A]", "[Not Supported]"} else int(memory)
+                if int(pid) <= 0 or used < 0:
+                    raise ValueError
+                processes.setdefault(uuid, {})[int(pid)] = used
+            except ValueError:
+                raise GPUQueryError("could not determine GPU process memory") from None
+        elif row and any(part.strip() for part in row) and row[0].strip() != "No running processes found":
+            raise GPUQueryError("could not parse GPU compute-process occupancy")
+    return [GPU(g.index, g.uuid, g.free_mib, g.uuid in processes, processes.get(g.uuid, {})) for g in devices]
+
+
+def _matches(gpu: GPU, token: int | str) -> bool:
+    return str(gpu.index) == str(token) or (
+        isinstance(token, str) and token.startswith("GPU-") and gpu.uuid.startswith(token)
+    )
+
+
+def visible_devices(devices: list[GPU], policy: Scheduling) -> list[GPU]:
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
+    visible_tokens = None if visible is None else [part.strip() for part in visible.split(",")]
+    result = []
+    for gpu in devices:
+        if visible_tokens is not None and not any(_matches(gpu, token) for token in visible_tokens):
+            continue
+        if policy.allowed_gpus is not None and not any(_matches(gpu, token) for token in policy.allowed_gpus):
+            continue
+        result.append(gpu)
+    return result
+
+
+def eligible(devices: list[GPU], policy: Scheduling, reserved: set[str],
+             cooldown: dict[str, float], quarantined: set[str], now: float,
+             memory_mib: int | None = None) -> list[GPU]:
+    minimum = policy.min_free_memory_mib if memory_mib is None else memory_mib
+    result = [gpu for gpu in visible_devices(devices, policy)
+              if gpu.uuid not in quarantined
+              and cooldown.get(gpu.uuid, 0) <= now and gpu.free_mib > minimum
+              and not (policy.require_idle and (gpu.busy or gpu.uuid in reserved))]
+    return sorted(result, key=lambda gpu: (-gpu.free_mib, gpu.index))

task_scheduling/memory.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass, replace
+
+from .gpu import GPU, eligible, visible_devices
+from .policy import Scheduling
+
+
+@dataclass(frozen=True)
+class Reservation:
+    gpu_uuid: str
+    memory_mib: int
+    pid: int | None = None
+
+
+class MemoryPool:
+    def __init__(self, policy: Scheduling):
+        self.policy = policy
+        self.devices: list[GPU] = []
+        self.reservations: dict[int, Reservation] = {}
+
+    def refresh(self, devices: list[GPU]) -> None:
+        self.devices = visible_devices(devices, self.policy)
+
+    def available(self, gpu: GPU) -> int:
+        reserved = sum(max(0, reservation.memory_mib - gpu.process_memory_mib.get(reservation.pid, 0))
+                       for reservation in self.reservations.values() if reservation.gpu_uuid == gpu.uuid)
+        return max(0, gpu.free_mib - self.policy.memory_margin_mib - reserved)
+
+    def candidates(self, memory_mib: int, cooldown: dict[str, float], quarantined: set[str],
+                   now: float, last_gpu: str | None = None) -> list[GPU]:
+        reserved = {reservation.gpu_uuid for reservation in self.reservations.values()}
+        devices = eligible(self.devices, self.policy, reserved, cooldown, quarantined, now, memory_mib=0)
+        return sorted((gpu for gpu in devices if self.available(gpu) >= memory_mib),
+                      key=lambda gpu: (gpu.uuid == last_gpu, self.available(gpu), gpu.index))
+
+    def reserve(self, identifier: int, gpu: GPU, memory_mib: int) -> None:
+        self.reservations[identifier] = Reservation(gpu.uuid, memory_mib)
+
+    def started(self, identifier: int, pid: int) -> None:
+        self.reservations[identifier] = replace(self.reservations[identifier], pid=pid)
+
+    def release(self, identifier: int) -> None:
+        self.reservations.pop(identifier, None)
+
+    def summary(self) -> str:
+        return ", ".join(f"GPU{gpu.index} available {self.available(gpu) / 1024:g}GiB" for gpu in self.devices)

task_scheduling/policy.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+
+
+def integer(value: int, name: str, minimum: int = 0) -> None:
+    if type(value) is not int or value < minimum:
+        raise ValueError(f"{name} must be an integer >= {minimum}")
+
+
+@dataclass(frozen=True)
+class Scheduling:
+    allowed_gpus: tuple[int | str, ...] | None = None
+    min_free_memory_mib: int = 4096
+    memory_margin_mib: int = 0
+    max_parallel_tasks: int = 3
+    require_idle: bool = False
+    max_retries: int = 3
+    poll_interval_seconds: int = 10
+    gpu_cooldown_seconds: int = 60
+    task_timeout_seconds: int = 0
+    wait_seconds: int = 3600
+
+    def __post_init__(self) -> None:
+        for name in ("min_free_memory_mib", "memory_margin_mib", "max_retries",
+                     "gpu_cooldown_seconds", "task_timeout_seconds"):
+            integer(getattr(self, name), name)
+        for name in ("max_parallel_tasks", "poll_interval_seconds", "wait_seconds"):
+            integer(getattr(self, name), name, 1)
+        if type(self.require_idle) is not bool:
+            raise ValueError("require_idle must be a boolean")
+        if self.allowed_gpus is not None:
+            allowed = self.allowed_gpus
+            if not isinstance(allowed, (list, tuple)) or not allowed:
+                raise ValueError("allowed_gpus must be null or a nonempty list of GPU indices/UUIDs")
+            if any(not ((type(gpu) is int and gpu >= 0) or
+                        (isinstance(gpu, str) and gpu.startswith("GPU-") and len(gpu) > 4))
+                   for gpu in allowed):
+                raise ValueError("allowed_gpus entries must be nonnegative indices or GPU- UUIDs")
+            if len(set(allowed)) != len(allowed):
+                raise ValueError("allowed_gpus must not contain duplicates")
+            object.__setattr__(self, "allowed_gpus", tuple(allowed))
+
+    @classmethod
+    def from_settings(cls, settings: dict) -> "Scheduling":
+        if not isinstance(settings, dict) or set(settings) - cls.__dataclass_fields__.keys():
+            raise ValueError("Unknown scheduling settings")
+        return cls(**settings)

task_scheduling/progress.py

@@ -0,0 +1,131 @@
+from contextlib import ExitStack
+from dataclasses import dataclass
+import shutil
+import sys
+import time
+from typing import TextIO
+
+
+@dataclass(frozen=True)
+class WorkerStatus:
+    identifier: int
+    pid: int
+    device: str
+    attempt: int
+    current: int
+    total: int
+    phase: str
+    started: float
+    unit: str = "epoch"
+    memory_mib: int = 0
+
+
+def duration(seconds: float) -> str:
+    seconds = max(0, int(seconds))
+    minutes, seconds = divmod(seconds, 60)
+    hours, minutes = divmod(minutes, 60)
+    days, hours = divmod(hours, 24)
+    if days:
+        return f"{days}d{hours:02d}h"
+    if hours:
+        return f"{hours}h{minutes:02d}m"
+    if minutes:
+        return f"{minutes}m{seconds:02d}s"
+    return f"{seconds}s"
+
+
+class TaskProgress:
+
+    def __init__(self, records: list[dict], slots: int, max_retries: int,
+                 *, label: str = "Tasks", stream: TextIO | None = None):
+        self.label = label
+        self.records, self.slots, self.max_attempts = records, slots, max_retries + 1
+        self.stream = sys.stderr if stream is None else stream
+        self.interactive = self.stream.isatty()
+        self.started = time.monotonic()
+        self.initial_completed = sum(record["status"] == "complete" for record in records)
+        self.slot_ids: list[int | None] = [None] * slots
+        self.lines: list = []
+        self.last_signature = None
+        self.last_printed = float("-inf")
+        self.stack = ExitStack()
+
+    def __enter__(self):
+        try:
+            from tqdm.auto import tqdm
+        except ImportError:
+            self.interactive = False
+        if self.interactive:
+            try:
+                for position in range(self.slots + 2):
+                    self.lines.append(self.stack.enter_context(tqdm(
+                        total=0, bar_format="{desc}", position=position,
+                        leave=position == 0, dynamic_ncols=True, file=self.stream,
+                    )))
+            except BaseException:
+                self.stack.close()
+                raise
+        return self
+
+    def __exit__(self, *exception):
+        return self.stack.__exit__(*exception)
+
+    def snapshot(self, workers: list[WorkerStatus], event: str, *, now: float,
+                 width: int) -> list[str]:
+        by_id = {worker.identifier: worker for worker in workers}
+        self.slot_ids = [identifier if identifier in by_id else None for identifier in self.slot_ids]
+        for worker in workers:
+            if worker.identifier not in self.slot_ids:
+                self.slot_ids[self.slot_ids.index(None)] = worker.identifier
+
+        completed = sum(record["status"] == "complete" for record in self.records)
+        failed = sum(record["status"] in {"failed", "blocked"} for record in self.records)
+        pending = sum(record["status"] == "pending" for record in self.records)
+        retries = sum(max(0, len(record.get("attempts", [])) - 1) for record in self.records)
+        total = len(self.records)
+        elapsed = max(0, now - self.started)
+        fresh_completed = completed - self.initial_completed
+        remaining = total - completed - failed
+        if remaining == 0:
+            eta = "0s"
+        elif fresh_completed > 0 and elapsed > 0:
+            eta = "~" + duration(remaining * elapsed / fresh_completed)
+        else:
+            eta = "estimating"
+        count = f"{self.label} {completed}/{total}"
+        fraction = completed / total if total else 1
+        tail = (f" {fraction:.0%} | run{len(workers)}/{self.slots} wait{pending} fail{failed} retry{retries}"
+                f" | t{duration(elapsed)} ETA {eta}")
+        filled = int(20 * fraction)
+        bar = " [" + "#" * filled + "-" * (20 - filled) + "]"
+        header = count + (bar if len(count + bar + tail) <= width else "") + tail
+        result = [header]
+        for index, identifier in enumerate(self.slot_ids, start=1):
+            if identifier is None:
+                result.append(f"Slot {index} | idle")
+                continue
+            worker = by_id[identifier]
+            budget = f" budget{worker.memory_mib / 1024:g}GiB" if worker.memory_mib else ""
+            result.append(
+                f"Slot {index} | {worker.device} PID{worker.pid}{budget} | {self.label} {identifier:03d} | "
+                f"try{worker.attempt}/{self.max_attempts} | {worker.unit}{worker.current}/{worker.total} | "
+                f"{worker.phase} | elapsed {duration(now - worker.started)}"
+            )
+        result.append("Latest: " + " ".join(event.split()))
+        return [line if len(line) <= width else line[:max(0, width - 3)] + "..." for line in result]
+
+    def render(self, workers: list[WorkerStatus], event: str) -> None:
+        now = time.monotonic()
+        width = max(20, shutil.get_terminal_size(fallback=(120, 24)).columns) if self.interactive else 160
+        frame = self.snapshot(workers, event, now=now, width=width)
+        if self.interactive:
+            for line, text in zip(self.lines, frame):
+                line.set_description_str(text, refresh=False)
+                line.refresh()
+        else:
+            signature = (tuple((r["status"], len(r.get("attempts", []))) for r in self.records),
+                         tuple((w.identifier, w.pid, w.current, w.phase) for w in workers), event)
+            if signature != self.last_signature or now - self.last_printed >= 60:
+                self.stream.write("\n".join(frame) + "\n")
+                self.stream.flush()
+                self.last_printed, self.last_signature = now, signature

task_scheduling/runtime.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+import os
+from pathlib import Path
+import signal
+import subprocess
+from typing import Callable
+
+from .gpu import GPU, discover
+from .task import Command
+
+
+def stop_process_tree(process: subprocess.Popen) -> None:
+    for sig in (signal.SIGTERM, signal.SIGKILL):
+        try:
+            os.killpg(process.pid, sig)
+        except ProcessLookupError:
+            pass
+        if sig == signal.SIGTERM:
+            try:
+                process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                pass
+    process.wait()
+
+
+@dataclass
+class LocalRuntime:
+    discover: Callable[[], list[GPU]] = discover
+    stop: Callable[[subprocess.Popen], None] = stop_process_tree
+
+    def launch(self, command: Command, directory: Path, gpu: GPU | None, inherited_fds: tuple[int, ...]):
+        environment = {**os.environ, **command.env, "PYTHONUNBUFFERED": "1",
+                       "CUDA_VISIBLE_DEVICES": "" if gpu is None else gpu.uuid,
+                       "TASK_PROGRESS_FILE": str(directory / "progress.json")}
+        with (directory / "worker.log").open("w") as log:
+            return subprocess.Popen(command.args, cwd=command.cwd, env=environment, stdout=log,
+                                    stderr=subprocess.STDOUT, start_new_session=True, pass_fds=inherited_fds)

task_scheduling/scheduler.py

@@ -0,0 +1,461 @@
+from collections import deque
+from contextlib import ExitStack
+from copy import deepcopy
+from dataclasses import asdict, dataclass
+import fcntl
+import json
+import os
+from pathlib import Path
+import re
+import subprocess
+import tempfile
+import time
+from typing import Callable, Iterable, TextIO
+
+from .gpu import GPU, GPUQueryError, visible_devices
+from .memory import MemoryPool
+from .policy import Scheduling
+from .progress import TaskProgress, WorkerStatus
+from .runtime import LocalRuntime
+from .storage import StateStore
+from .task import (
+    Assessment, Attempt, CommandHandler, TaskContext, TaskEvent, TaskHandler,
+    TaskOutcome, TaskRecord, TaskSpec,
+)
+
+
+@dataclass
+class RunningTask:
+    process: subprocess.Popen
+    gpu: GPU | None
+    lock: TextIO | None
+    directory: Path
+    started: float
+    missing: int = 0
+
+
+class Scheduler:
+    def __init__(self, root: str | Path, policy: Scheduling | None = None, *,
+                 runtime: LocalRuntime | None = None, inherited_fds: tuple[int, ...] = (),
+                 lock_namespace: str = "task-scheduling", label: str = "Tasks"):
+        if not re.fullmatch(r"[A-Za-z0-9_-]+", lock_namespace):
+            raise ValueError("Lock namespace must contain only letters, digits, underscores and hyphens")
+        self.root = Path(root).resolve()
+        self.policy = policy or Scheduling()
+        self.runtime = runtime or LocalRuntime()
+        self.label = label
+        self.lock_namespace = lock_namespace
+        self.inherited_fds = inherited_fds
+        self.store = StateStore(self.root)
+        self.memory = MemoryPool(self.policy)
+        self._records: dict[str, TaskRecord] = {}
+        self._handlers: dict[str, TaskHandler] = {"command": CommandHandler()}
+        self._listeners: dict[str, Callable[[TaskEvent], None]] = {}
+        self._observers: list[Callable[[tuple[TaskRecord, ...]], None]] = []
+        self._running: dict[str, RunningTask] = {}
+        self._cooldown: dict[str, float] = {}
+        self._quarantined: set[str] = set()
+        self._known: set[str] = set()
+        self._reconciled: set[str] = set()
+        self._session: ExitStack | None = None
+        self._lock_fd: int | None = None
+        self._event = "Preparing tasks"
+        self._query_error = ""
+        self._next_scan = 0.0
+
+    @property
+    def tasks(self) -> tuple[TaskRecord, ...]:
+        return tuple(deepcopy(record) for record in self._records.values())
+
+    @property
+    def has_saved_state(self) -> bool:
+        return self.store.path.exists()
+
+    def __enter__(self) -> "Scheduler":
+        if self._session is not None:
+            raise RuntimeError("Scheduler session is already open")
+        session = ExitStack()
+        try:
+            self._lock_fd = session.enter_context(self.store.lock())
+            submitted = [record.spec for record in self._records.values()]
+            self._records = self.store.load()
+            self.memory = MemoryPool(self.policy)
+            self._cooldown.clear()
+            self._quarantined.clear()
+            self._known.clear()
+            self._next_scan = 0.0
+            self._session = session
+            self.submit_many(submitted)
+        except BaseException:
+            self._session = None
+            session.close()
+            raise
+        return self
+
+    def __exit__(self, *exception) -> None:
+        try:
+            if self._running:
+                self._interrupt_running()
+                self._save()
+        finally:
+            session, self._session = self._session, None
+            self._lock_fd = None
+            self._reconciled.clear()
+            if session is not None:
+                session.__exit__(*exception)
+
+    def register_handler(self, name: str, handler: TaskHandler) -> None:
+        if not isinstance(name, str) or not name.strip():
+            raise ValueError("Handler name must be a nonempty string")
+        if name in self._handlers:
+            raise ValueError(f"Handler already registered: {name}")
+        if any(not callable(getattr(handler, method, None)) for method in ("prepare", "assess", "commit", "recover")):
+            raise TypeError("Handlers must implement prepare, assess, commit and recover")
+        self._handlers[name] = handler
+
+    def subscribe(self, name: str, listener: Callable[[TaskEvent], None]) -> None:
+        if not isinstance(name, str) or not name.strip() or name in self._listeners:
+            raise ValueError(f"Listener name is empty or already registered: {name}")
+        if not callable(listener):
+            raise TypeError("Event listeners must be callable")
+        self._listeners[name] = listener
+
+    def observe(self, observer: Callable[[tuple[TaskRecord, ...]], None]) -> None:
+        if not callable(observer):
+            raise TypeError("State observers must be callable")
+        self._observers.append(observer)
+
+    def submit(self, task: TaskSpec) -> None:
+        self.submit_many([task])
+
+    def submit_many(self, tasks: Iterable[TaskSpec]) -> None:
+        submitted = [TaskSpec.from_dict(json.loads(json.dumps(asdict(task), allow_nan=False))) for task in tasks]
+        specs = {identifier: record.spec for identifier, record in self._records.items()}
+        for task in submitted:
+            if task.id in specs and specs[task.id].identity() != task.identity():
+                raise ValueError(f"Task definition changed: {task.id}")
+            existing = self._records.get(task.id)
+            if existing is not None and task.id in self._running and existing.spec != task:
+                raise ValueError(f"Cannot update an executing task: {task.id}")
+            specs[task.id] = task
+        self._validate_graph(specs)
+        for task in submitted:
+            if task.id in self._records:
+                self._records[task.id].spec = task
+            else:
+                self._records[task.id] = TaskRecord(task, len(self._records))
+        if submitted:
+            self._save()
+
+    def import_records(self, records: Iterable[TaskRecord]) -> None:
+        self._require_session()
+        if self._records or self.store.path.exists():
+            raise ValueError("Recovery records can only be imported into a new scheduler")
+        imported = deepcopy(list(records))
+        if len({record.spec.id for record in imported}) != len(imported):
+            raise ValueError("Duplicate task ids in recovery records")
+        if [record.sequence for record in imported] != list(range(len(imported))):
+            raise ValueError("Recovery records must have contiguous sequence numbers")
+        if any(record.status not in {"pending", "running", "committing", "complete", "failed", "blocked"}
+               for record in imported):
+            raise ValueError("Unknown task status in recovery records")
+        self._validate_graph({record.spec.id: record.spec for record in imported})
+        self._records = {record.spec.id: record for record in imported}
+        self._save()
+
+    def reconcile(self) -> None:
+        self._require_session()
+        self._validate_graph({identifier: record.spec for identifier, record in self._records.items()})
+        for record in list(self._records.values()):
+            if record.spec.id in self._reconciled:
+                continue
+            handler = self._handler(record)
+            recovery = handler.recover(self._context(record))
+            if recovery is not None:
+                record.status = recovery.status
+                record.result = deepcopy(recovery.result)
+                if record.status == "complete" and record.attempts:
+                    record.attempts[-1].status = "complete"
+            elif record.status == "complete":
+                raise ValueError(f"Completed task could not be recovered: {record.spec.id}")
+            elif record.status == "committing":
+                raise ValueError(f"Task publication could not be recovered: {record.spec.id}")
+            else:
+                for attempt in record.attempts:
+                    if attempt.status == "running":
+                        attempt.status = "interrupted"
+                record.status = "pending" if record.can_retry(self.policy.max_retries) else "failed"
+            self._reconciled.add(record.spec.id)
+            self._save()
+            if record.status == "committing":
+                self._commit(record)
+
+    def run(self) -> tuple[TaskRecord, ...]:
+        if self._session is None:
+            with self:
+                return self.run()
+        self.reconcile()
+        self._deliver_events()
+        deadline = time.monotonic() + self.policy.wait_seconds
+        display_records = [{"status": record.status, "attempts": record.attempts}
+                           for record in self._records.values()]
+        with TaskProgress(display_records, self.policy.max_parallel_tasks, self.policy.max_retries,
+                          label=self.label) as progress:
+            try:
+                while self._has_work():
+                    now = time.monotonic()
+                    devices = self._scan(now)
+                    if self._finish_tasks(devices, now):
+                        deadline = now + self.policy.wait_seconds
+                    self._block_dependents()
+                    self._deliver_events()
+                    self._start_tasks(devices, now)
+                    if self._pending() and not self._running:
+                        self._check_wait(now, deadline)
+                    display_records[:] = [{"status": record.status,
+                                           "attempts": record.attempts} for record in self._records.values()]
+                    self._render(progress)
+                    if self._has_work():
+                        time.sleep(1)
+            finally:
+                self._interrupt_running()
+                self._save()
+        self._deliver_events()
+        return self.tasks
+
+    @staticmethod
+    def _validate_graph(specs: dict[str, TaskSpec]) -> None:
+        remaining = {identifier: len(spec.dependencies) for identifier, spec in specs.items()}
+        dependents = {identifier: [] for identifier in specs}
+        for identifier, spec in specs.items():
+            for dependency in spec.dependencies:
+                if dependency not in specs:
+                    raise ValueError(f"Unknown dependency: {dependency}")
+                dependents[dependency].append(identifier)
+        ready = deque(identifier for identifier, count in remaining.items() if count == 0)
+        visited = 0
+        while ready:
+            identifier = ready.popleft()
+            visited += 1
+            for dependent in dependents[identifier]:
+                remaining[dependent] -= 1
+                if remaining[dependent] == 0:
+                    ready.append(dependent)
+        if visited != len(specs):
+            raise ValueError("Task dependency cycle")
+
+    def _require_session(self) -> None:
+        if self._session is None:
+            raise RuntimeError("Open a scheduler session before recovering tasks")
+
+    def _save(self) -> None:
+        if self._session is None:
+            return
+        self.store.save(self._records)
+        for observer in self._observers:
+            observer(self.tasks)
+
+    def _handler(self, record: TaskRecord) -> TaskHandler:
+        if record.spec.handler not in self._handlers:
+            raise ValueError(f"Handler is not registered: {record.spec.handler}")
+        return self._handlers[record.spec.handler]
+
+    def _context(self, record: TaskRecord) -> TaskContext:
+        return TaskContext(deepcopy(record.spec), self.root, record.status, tuple(deepcopy(record.attempts)))
+
+    def _has_work(self) -> bool:
+        return bool(self._running) or any(record.status in {"pending", "committing"} for record in self._records.values())
+
+    def _pending(self) -> list[TaskRecord]:
+        return [record for record in self._records.values() if record.status == "pending"]
+
+    def _ready(self) -> list[TaskRecord]:
+        return sorted((record for record in self._pending()
+                       if all(self._records[dependency].status == "complete" for dependency in record.spec.dependencies)),
+                      key=lambda record: (-record.spec.priority, -record.spec.resources.memory_mib, record.sequence))
+
+    def _scan(self, now: float) -> list[GPU] | None:
+        if now < self._next_scan or not any(record.spec.resources.device == "cuda" for record in self._records.values()):
+            return None
+        self._next_scan = now + self.policy.poll_interval_seconds
+        try:
+            devices = self.runtime.discover()
+            self.memory.refresh(devices)
+            self._known.update(gpu.uuid for gpu in visible_devices(devices, self.policy))
+            self._query_error = ""
+            return devices
+        except GPUQueryError as error:
+            self._query_error = str(error)
+            return None
+
+    def _finish_tasks(self, devices: list[GPU] | None, now: float) -> bool:
+        finished = False
+        for identifier, running in list(self._running.items()):
+            if devices is not None and running.gpu is not None:
+                running.missing = 0 if any(gpu.uuid == running.gpu.uuid for gpu in devices) else running.missing + 1
+            lost = running.missing >= 2
+            timeout = self.policy.task_timeout_seconds > 0 and now - running.started > self.policy.task_timeout_seconds
+            code = running.process.poll()
+            if code is None and not lost and not timeout:
+                continue
+            self.runtime.stop(running.process)
+            self._release(identifier, running)
+            record = self._records[identifier]
+            if lost or timeout:
+                assessment = Assessment(False, True, "device" if lost else "timeout",
+                                        "GPU disappeared" if lost else "task timed out")
+            else:
+                assessment = self._handler(record).assess(self._context(record), TaskOutcome(code))
+            if assessment.success:
+                record.status = "committing"
+                self._save()
+                self._commit(record)
+            else:
+                self._fail(record, assessment, running.gpu, now)
+            finished = True
+        return finished
+
+    def _commit(self, record: TaskRecord) -> None:
+        result = self._handler(record).commit(self._context(record))
+        json.dumps(result, allow_nan=False)
+        if not isinstance(result, dict):
+            raise TypeError("Task commit must return a JSON object")
+        record.result = deepcopy(result)
+        record.status = "complete"
+        if record.attempts:
+            record.attempts[-1].status = "complete"
+        self._event = f"{record.spec.name} completed"
+        self._save()
+
+    def _fail(self, record: TaskRecord, assessment: Assessment, gpu: GPU | None, now: float) -> None:
+        attempt = record.attempts[-1]
+        attempt.status, attempt.kind, attempt.message = "failed", assessment.kind, assessment.message
+        attempt.retryable = assessment.retryable
+        retry = record.can_retry(self.policy.max_retries)
+        record.status = "pending" if retry else "failed"
+        if gpu is not None:
+            if assessment.kind == "device":
+                self._quarantined.add(gpu.uuid)
+            else:
+                self._cooldown[gpu.uuid] = now + self.policy.gpu_cooldown_seconds
+        self._event = f"{record.spec.name}: {assessment.message}; {'retry queued' if retry else 'failed'}"
+        self._save()
+
+    def _acquire_gpu(self, record: TaskRecord, now: float) -> tuple[GPU | None, TextIO | None]:
+        last_gpu = record.attempts[-1].gpu if record.attempts else None
+        for gpu in self.memory.candidates(record.spec.resources.memory_mib, self._cooldown,
+                                          self._quarantined, now, last_gpu):
+            path = Path(tempfile.gettempdir()) / f"{self.lock_namespace}-{os.getuid()}-{gpu.uuid}.lock"
+            handle = path.open("a")
+            try:
+                fcntl.flock(handle, (fcntl.LOCK_EX if self.policy.require_idle else fcntl.LOCK_SH) | fcntl.LOCK_NB)
+            except BlockingIOError:
+                handle.close()
+                continue
+            self.memory.reserve(record.sequence, gpu, record.spec.resources.memory_mib)
+            return gpu, handle
+        return None, None
+
+    def _start_tasks(self, devices: list[GPU] | None, now: float) -> None:
+        for record in self._ready():
+            if len(self._running) >= self.policy.max_parallel_tasks:
+                break
+            gpu, lock = None, None
+            if record.spec.resources.device == "cuda":
+                if devices is None:
+                    continue
+                gpu, lock = self._acquire_gpu(record, now)
+                if gpu is None:
+                    continue
+            self._launch(record, gpu, lock, now)
+
+    def _launch(self, record: TaskRecord, gpu: GPU | None, lock: TextIO | None, now: float) -> None:
+        name = re.sub(r"[^A-Za-z0-9_.-]", "_", record.spec.name)[:80]
+        directory = self.root / f"{record.sequence:03d}-{name}" / f"{len(record.attempts) + 1:04d}"
+        process = None
+        try:
+            record.status = "running"
+            record.attempts.append(Attempt(str(directory), gpu=None if gpu is None else gpu.uuid,
+                                           memory_mib=record.spec.resources.memory_mib))
+            self._save()
+            directory.mkdir(parents=True)
+            command = self._handler(record).prepare(self._context(record))
+            inherited = (*self.inherited_fds, self._lock_fd)
+            if lock is not None:
+                inherited = (*inherited, lock.fileno())
+            process = self.runtime.launch(command, directory, gpu, inherited)
+            record.attempts[-1].pid = process.pid
+            if gpu is not None:
+                self.memory.started(record.sequence, process.pid)
+            self._running[record.spec.id] = RunningTask(process, gpu, lock, directory, now)
+        except BaseException:
+            if process is not None:
+                self.runtime.stop(process)
+            self.memory.release(record.sequence)
+            if lock is not None:
+                lock.close()
+            record.status = "pending"
+            record.attempts[-1].status = "interrupted"
+            raise
+        self._event = (f"{record.spec.name} started on {'CPU' if gpu is None else 'GPU' + str(gpu.index)} "
+                       f"PID{process.pid} budget {record.spec.resources.memory_mib / 1024:g}GiB")
+        self._save()
+
+    def _release(self, identifier: str, running: RunningTask) -> None:
+        self.memory.release(self._records[identifier].sequence)
+        if running.lock is not None:
+            running.lock.close()
+        del self._running[identifier]
+
+    def _interrupt_running(self) -> None:
+        for identifier, running in list(self._running.items()):
+            self.runtime.stop(running.process)
+            self._release(identifier, running)
+            record = self._records[identifier]
+            record.status = "pending"
+            record.attempts[-1].status = "interrupted"
+
+    def _block_dependents(self) -> None:
+        changed = True
+        while changed:
+            changed = False
+            for record in self._pending():
+                if any(self._records[dependency].status in {"failed", "blocked"} for dependency in record.spec.dependencies):
+                    record.status = "blocked"
+                    changed = True
+                    self._save()
+
+    def _deliver_events(self) -> None:
+        for record in list(self._records.values()):
+            if record.status not in {"complete", "failed", "blocked"}:
+                continue
+            event_id = f"{record.spec.id}:{len(record.attempts)}:{record.status}"
+            for name, listener in self._listeners.items():
+                if record.acknowledged.get(name) == event_id:
+                    continue
+                listener(TaskEvent(event_id, deepcopy(record.spec), record.status,
+                                   tuple(deepcopy(record.attempts)), deepcopy(record.result)))
+                record.acknowledged[name] = event_id
+                self._save()
+
+    def _check_wait(self, now: float, deadline: float) -> None:
+        self._event = "Waiting for eligible GPUs" + (f": {self._query_error}" if self._query_error else "")
+        if self._known and self._known <= self._quarantined:
+            raise RuntimeError("All allowed GPUs quarantined; repair GPUs before rerunning")
+        if now > deadline:
+            raise TimeoutError("No eligible GPU became available before timeout")
+
+    def _render(self, progress: TaskProgress) -> None:
+        workers = []
+        for identifier, running in self._running.items():
+            state = {"current": 0, "total": 0, "phase": "initializing", "unit": "step"}
+            try:
+                state.update(json.loads((running.directory / "progress.json").read_text()))
+            except (OSError, ValueError):
+                pass
+            record = self._records[identifier]
+            workers.append(WorkerStatus(record.sequence, running.process.pid,
+                "CPU" if running.gpu is None else f"GPU{running.gpu.index}", len(record.attempts),
+                state["current"], state["total"], f"{record.spec.name} {state['phase']}", running.started,
+                state["unit"], record.spec.resources.memory_mib))
+        progress.render(workers, f"{self.memory.summary()}; {self._event}" if self.memory.devices else self._event)

task_scheduling/storage.py

@@ -0,0 +1,49 @@
+from contextlib import contextmanager
+from dataclasses import asdict
+import fcntl
+import json
+from pathlib import Path
+
+from .task import TaskRecord
+
+
+def atomic_json(path: str | Path, value: object) -> None:
+    path = Path(path)
+    temporary = path.with_suffix(".tmp")
+    temporary.write_text(json.dumps(value, indent=2, allow_nan=False))
+    temporary.replace(path)
+
+
+class StateStore:
+    def __init__(self, root: Path):
+        self.root = root
+        self.path = root / "scheduler.json"
+
+    @contextmanager
+    def lock(self):
+        self.root.mkdir(parents=True, exist_ok=True)
+        with (self.root / "scheduler.lock").open("a") as handle:
+            try:
+                fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            except BlockingIOError:
+                raise RuntimeError(f"Scheduler is still running: {self.root}") from None
+            yield handle.fileno()
+
+    def load(self) -> dict[str, TaskRecord]:
+        if not self.path.exists():
+            return {}
+        document = json.loads(self.path.read_text())
+        if document["version"] != 1:
+            raise ValueError("Unsupported scheduler state version")
+        records = [TaskRecord.from_dict(record) for record in document["tasks"]]
+        if len({record.spec.id for record in records}) != len(records):
+            raise ValueError("Duplicate task ids in scheduler state")
+        if [record.sequence for record in records] != list(range(len(records))):
+            raise ValueError("Invalid task sequence numbers in scheduler state")
+        if any(record.status not in {"pending", "running", "committing", "complete", "failed", "blocked"}
+               for record in records):
+            raise ValueError("Unknown task status in scheduler state")
+        return {record.spec.id: record for record in records}
+
+    def save(self, records: dict[str, TaskRecord]) -> None:
+        atomic_json(self.path, {"version": 1, "tasks": [asdict(record) for record in records.values()]})

task_scheduling/task.py

@@ -0,0 +1,197 @@
+from dataclasses import dataclass, field
+import json
+from pathlib import Path
+from typing import Literal, Protocol
+
+from .policy import integer
+
+
+TaskStatus = Literal["pending", "running", "committing", "complete", "failed", "blocked"]
+
+
+@dataclass(frozen=True)
+class ResourceRequest:
+    device: Literal["cpu", "cuda"] = "cpu"
+    memory_mib: int = 0
+
+    def __post_init__(self) -> None:
+        if self.device not in {"cpu", "cuda"}:
+            raise ValueError("Task devices must be cpu or cuda")
+        integer(self.memory_mib, "memory_mib")
+        if self.device == "cpu" and self.memory_mib:
+            raise ValueError("CPU tasks cannot reserve GPU memory")
+
+
+@dataclass(frozen=True)
+class TaskSpec:
+    id: str
+    handler: str
+    name: str = "task"
+    payload: dict = field(default_factory=dict)
+    dependencies: tuple[str, ...] = ()
+    priority: int = 0
+    resources: ResourceRequest = field(default_factory=ResourceRequest)
+
+    def __post_init__(self) -> None:
+        if any(not isinstance(value, str) or not value.strip() for value in (self.id, self.handler, self.name)):
+            raise ValueError("Task id, handler and name must be nonempty strings")
+        if type(self.priority) is not int:
+            raise ValueError("Task priority must be an integer")
+        if not isinstance(self.dependencies, (list, tuple)) or any(
+            not isinstance(value, str) or not value for value in self.dependencies
+        ):
+            raise ValueError("Task dependencies must contain task ids")
+        if self.id in self.dependencies or len(set(self.dependencies)) != len(self.dependencies):
+            raise ValueError("Task dependencies must be unique and cannot reference the task itself")
+        object.__setattr__(self, "dependencies", tuple(self.dependencies))
+        if not isinstance(self.payload, dict):
+            raise ValueError("Task payload must be a JSON object")
+        if not isinstance(self.resources, ResourceRequest):
+            raise TypeError("Task resources must be a ResourceRequest")
+        json.dumps(self.payload, allow_nan=False)
+
+    def identity(self) -> dict:
+        return {"id": self.id, "handler": self.handler, "payload": self.payload,
+                "dependencies": list(self.dependencies)}
+
+    @classmethod
+    def from_dict(cls, values: dict) -> "TaskSpec":
+        return cls(**{**values, "resources": ResourceRequest(**values["resources"])})
+
+
+@dataclass(frozen=True)
+class Command:
+    args: tuple[str, ...]
+    cwd: str | Path | None = None
+    env: dict[str, str] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.args, (list, tuple)) or not self.args or any(
+            not isinstance(value, str) for value in self.args
+        ):
+            raise ValueError("Command args must be a nonempty list of strings")
+        if any(not isinstance(key, str) or not isinstance(value, str) for key, value in self.env.items()):
+            raise ValueError("Command environment must map strings to strings")
+        object.__setattr__(self, "args", tuple(self.args))
+
+
+@dataclass
+class Attempt:
+    directory: str
+    status: str = "running"
+    gpu: str | None = None
+    pid: int | None = None
+    memory_mib: int = 0
+    kind: str = ""
+    message: str = ""
+    retryable: bool = False
+
+
+@dataclass
+class TaskRecord:
+    spec: TaskSpec
+    sequence: int
+    status: TaskStatus = "pending"
+    attempts: list[Attempt] = field(default_factory=list)
+    result: dict = field(default_factory=dict)
+    acknowledged: dict[str, str] = field(default_factory=dict)
+
+    def can_retry(self, max_retries: int) -> bool:
+        failures = [attempt for attempt in self.attempts if attempt.status == "failed"]
+        return all(attempt.retryable for attempt in failures) and len(failures) <= max_retries
+
+    @classmethod
+    def from_dict(cls, values: dict) -> "TaskRecord":
+        return cls(**{**values, "spec": TaskSpec.from_dict(values["spec"]),
+                      "attempts": [Attempt(**attempt) for attempt in values["attempts"]]})
+
+
+@dataclass(frozen=True)
+class TaskContext:
+    task: TaskSpec
+    root: Path
+    status: TaskStatus
+    attempts: tuple[Attempt, ...]
+
+    @property
+    def directory(self) -> Path:
+        return Path(self.attempts[-1].directory) if self.attempts else self.root
+
+    @property
+    def gpu(self) -> str | None:
+        return self.attempts[-1].gpu if self.attempts else None
+
+
+@dataclass(frozen=True)
+class TaskOutcome:
+    returncode: int | None
+    timed_out: bool = False
+    device_lost: bool = False
+
+
+@dataclass(frozen=True)
+class Assessment:
+    success: bool
+    retryable: bool = False
+    kind: str = "process"
+    message: str = ""
+
+    def __post_init__(self) -> None:
+        if type(self.success) is not bool or type(self.retryable) is not bool:
+            raise TypeError("Assessment success and retryable must be booleans")
+        if not isinstance(self.kind, str) or not isinstance(self.message, str):
+            raise TypeError("Assessment kind and message must be strings")
+
+
+@dataclass(frozen=True)
+class Recovery:
+    status: Literal["complete", "failed", "committing"]
+    result: dict = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.status not in {"complete", "failed", "committing"}:
+            raise ValueError("Recovery status must be complete, failed or committing")
+        if not isinstance(self.result, dict):
+            raise TypeError("Recovery result must be a JSON object")
+        json.dumps(self.result, allow_nan=False)
+
+
+@dataclass(frozen=True)
+class TaskEvent:
+    id: str
+    task: TaskSpec
+    status: TaskStatus
+    attempts: tuple[Attempt, ...]
+    result: dict
+
+
+class TaskHandler(Protocol):
+    def prepare(self, context: TaskContext) -> Command: ...
+
+    def assess(self, context: TaskContext, outcome: TaskOutcome) -> Assessment: ...
+
+    def commit(self, context: TaskContext) -> dict: ...
+
+    def recover(self, context: TaskContext) -> Recovery | None: ...
+
+
+class CommandHandler:
+    def prepare(self, context: TaskContext) -> Command:
+        return Command(**context.task.payload)
+
+    def assess(self, context: TaskContext, outcome: TaskOutcome) -> Assessment:
+        if outcome.device_lost:
+            return Assessment(False, True, "device", "GPU disappeared")
+        if outcome.timed_out:
+            return Assessment(False, True, "timeout", "task timed out")
+        return Assessment(outcome.returncode == 0, True, "process", f"worker exited {outcome.returncode}")
+
+    def commit(self, context: TaskContext) -> dict:
+        return {}
+
+    def recover(self, context: TaskContext) -> Recovery | None:
+        if context.status == "committing":
+            return Recovery("committing")
+        if context.status == "complete":
+            return Recovery("complete")
+        return None