manifold-cli 0.0.39__py3-none-any.whl

manifold_cli/runner/gpu.py

@@ -0,0 +1,182 @@
+"""GPU and driver detection for runner placement, and GPU usage sampling."""
+
+import csv
+import subprocess
+
+from manifold_cli.models.observability import GpuUsage
+from manifold_cli.models.runner import RunnerGpuDevice, RunnerGpuFacts
+
+MEBIBYTES_PER_GIBIBYTE = 1024
+BYTES_PER_MEBIBYTE = 1024**2
+# `nvidia-smi` prints a value it cannot read in brackets, such as `[N/A]` or
+# `[Not Supported]`.
+UNREADABLE_VALUE_PREFIX = "["
+# A driver module still loading blocks `nvidia-smi` rather than failing it.
+NVIDIA_SMI_TIMEOUT_SECONDS = 10.0
+# The `nvidia-smi --query-gpu` fields of `RunnerGpuDevice` past its index,
+# name and memory, in the order that `_parse_gpu_device` unpacks them. A
+# driver that does not know one of them refuses the whole query, so
+# `gpu_facts` asks for them in a second query.
+GPU_DETAIL_FIELDS = (
+    "index",
+    "compute_cap",
+    "mig.mode.current",
+    "ecc.mode.current",
+    "compute_mode",
+    "power.limit",
+    "clocks.max.sm",
+    "pcie.link.gen.max",
+    "pcie.link.width.max",
+)
+
+
+def gpu_facts() -> RunnerGpuFacts:
+    """Return the visible GPU facts, or empty facts when the probe fails.
+
+    A GPU's details are None when the driver cannot report them.
+    """
+    try:
+        rows = _query_gpus(("index", "name", "memory.total", "driver_version"))
+        details = _gpu_details()
+        devices = [
+            _parse_gpu_device(index, name, memory_mib, details.get(index))
+            for index, name, memory_mib, _driver_version in rows
+        ]
+        driver_version = next((driver for *_, driver in rows), None)
+        return RunnerGpuFacts(devices=devices, driver_version=driver_version)
+    except (
+        FileNotFoundError,
+        subprocess.CalledProcessError,
+        subprocess.TimeoutExpired,
+        ValueError,
+    ):
+        return RunnerGpuFacts(devices=[], driver_version=None)
+
+
+def _gpu_details() -> dict[str, list[str | None]]:
+    """Return each GPU's details by index, or none when the driver refuses."""
+    try:
+        rows = _query_gpus(GPU_DETAIL_FIELDS)
+    except subprocess.CalledProcessError:
+        return {}
+    return {index: [_readable_value(value) for value in rest] for index, *rest in rows}
+
+
+def _query_gpus(fields: tuple[str, ...]) -> list[list[str]]:
+    result = subprocess.run(
+        [
+            "nvidia-smi",
+            "--query-gpu=" + ",".join(fields),
+            "--format=csv,noheader,nounits",
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=NVIDIA_SMI_TIMEOUT_SECONDS,
+    )
+    return [
+        [value.strip() for value in row]
+        for row in csv.reader(result.stdout.splitlines())
+    ]
+
+
+def _parse_gpu_device(
+    index: str, name: str, memory_mib: str, details: list[str | None] | None
+) -> RunnerGpuDevice:
+    (
+        compute_cap,
+        mig_mode,
+        ecc_mode,
+        compute_mode,
+        power_limit_w,
+        max_sm_clock_mhz,
+        pcie_generation,
+        pcie_width,
+    ) = details or [None] * (len(GPU_DETAIL_FIELDS) - 1)
+    return RunnerGpuDevice(
+        index=int(index),
+        name=name,
+        memory_gb=round(int(memory_mib) / MEBIBYTES_PER_GIBIBYTE),
+        compute_cap=compute_cap,
+        is_mig_enabled=_is_enabled(mig_mode),
+        is_ecc_enabled=_is_enabled(ecc_mode),
+        compute_mode=compute_mode,
+        power_limit_w=None if power_limit_w is None else float(power_limit_w),
+        max_sm_clock_mhz=None if max_sm_clock_mhz is None else int(max_sm_clock_mhz),
+        pcie_generation=None if pcie_generation is None else int(pcie_generation),
+        pcie_width=None if pcie_width is None else int(pcie_width),
+    )
+
+
+def _readable_value(value: str) -> str | None:
+    if value.startswith(UNREADABLE_VALUE_PREFIX):
+        return None
+    return value
+
+
+def _is_enabled(mode: str | None) -> bool | None:
+    if mode is None:
+        return None
+    return mode == "Enabled"
+
+
+def gpu_usage() -> list[GpuUsage]:
+    """Return the utilisation and memory use of each visible GPU.
+
+    Leave out a GPU for which `nvidia-smi` prints any of the values in
+    brackets, and return an empty list when the probe fails.
+    """
+    argv = [
+        "nvidia-smi",
+        "--query-gpu=index,name,utilization.gpu,memory.used,memory.total",
+        "--format=csv,noheader,nounits",
+    ]
+    try:
+        result = subprocess.run(
+            argv,
+            check=True,
+            capture_output=True,
+            text=True,
+            timeout=NVIDIA_SMI_TIMEOUT_SECONDS,
+        )
+        rows = [
+            [value.strip() for value in row]
+            for row in csv.reader(result.stdout.splitlines())
+        ]
+        return [
+            GpuUsage(
+                index=int(index),
+                name=name,
+                utilization_percent=float(utilization),
+                memory_used_bytes=int(memory_used_mib) * BYTES_PER_MEBIBYTE,
+                memory_total_bytes=int(memory_total_mib) * BYTES_PER_MEBIBYTE,
+            )
+            for index, name, utilization, memory_used_mib, memory_total_mib in rows
+            if not any(
+                value.startswith(UNREADABLE_VALUE_PREFIX)
+                for value in (utilization, memory_used_mib, memory_total_mib)
+            )
+        ]
+    except (
+        FileNotFoundError,
+        subprocess.CalledProcessError,
+        subprocess.TimeoutExpired,
+        ValueError,
+    ):
+        return []
+
+
+def gpu_memory_gb(devices: list[RunnerGpuDevice]) -> int:
+    """Return the smallest visible capacity, the floor a worker can rely on."""
+    return min((device.memory_gb for device in devices), default=0)
+
+
+def round_robin_gpu_device(devices: list[int], index: int) -> int | None:
+    """The device the worker at `index` lands on, cycling through `devices`.
+
+    None where the box exposes no GPU, so a worker's launch args fall back to
+    every card rather than a device index that names none of them.
+    """
+    if len(devices) == 0:
+        return None
+    return devices[index % len(devices)]

manifold_cli/runner/heartbeat.py

@@ -0,0 +1,202 @@
+"""Beating a side for as long as the server leaves it here.
+
+The server reclaims a side that goes quiet for the window it is swept
+against, so a side that is serving has to say so. The beat is also the only
+thing that says the side is gone: it is a conditional write, and one that
+matches no row is the instruction to stop the worker. Nothing else reports that
+a peer failed, that the run was cancelled, or that a sweep reclaimed this side,
+and nothing else needs to. The refusal includes the run's status, and
+`_beat_until_lost` handles a terminal status as a run that ended rather than a
+side the sweep reclaimed.
+
+A peer that completed is the same instruction arriving as an answer rather than
+a refusal, because the run is over and the side is still assigned to this
+runner. A policy container serves until it is stopped, so its own worker ending
+is not what says the run is done, and a side that only stopped when its peer
+failed would serve a finished run until something reclaimed it.
+
+A side is beaten for as long as its work is driven, not from the moment the
+policy accepts the benchmark's signature. Until then the heartbeat is sent with
+`serving=false`: the server does not change the side, so the startup window
+still covers loading an image, and it marks the runner seen and refuses the
+heartbeat once the run is released. The first serving heartbeat is sent once
+the policy has accepted the benchmark's signature, and the server moves the side
+to `running` on that heartbeat.
+"""
+
+import asyncio
+import contextlib
+from collections.abc import Coroutine
+from typing import Any, NoReturn
+
+from rich.markup import escape
+
+from manifold_cli.clients.manifold import ManifoldApiClient
+from manifold_cli.models.run import (
+    TERMINAL_STATUSES,
+    Assignment,
+    AssignmentNotAssignedToRunnerError,
+    RunSide,
+    RunStatus,
+    SideOutcome,
+)
+from manifold_cli.runner.plan import RunnerError
+from manifold_cli.runner.workload import Log
+
+
+class PeerCompletedError(Exception):
+    """The run's other side completed, so this one has nothing left to serve.
+
+    Raised out of the beat rather than returned, so the work being beaten for is
+    cancelled where it stands and the worker's teardown runs, as it does when
+    the side is refused. The side that catches this reports completed: it served
+    the whole run, and the peer's outcome is the run's.
+    """
+
+
+class RunEndedError(Exception):
+    """The run is terminal, so the server ended this side with the run.
+
+    Raised out of the heartbeat for the same reason as `PeerCompletedError`.
+    The side that catches this does not report an outcome, because the server
+    has already ended the side.
+    """
+
+    def __init__(self, status: RunStatus):
+        super().__init__(f"the run was {status.value}")
+        self.status = status
+
+    def stopped_line(self, side: RunSide) -> str:
+        """The log line for `side` stopping, in red unless the stop is expected."""
+        line = f"The run was {self.status.value}, so the {side.value} stopped."
+        if self.status in _EXPECTED_ENDINGS:
+            return line
+        return f"[red]{line}[/red]"
+
+
+_EXPECTED_ENDINGS = frozenset({RunStatus.CANCELED, RunStatus.COMPLETED})
+
+
+async def while_beating[T](
+    work: Coroutine[Any, Any, T],
+    *,
+    client: ManifoldApiClient,
+    assignment: Assignment,
+    period_s: float,
+    ready: asyncio.Event,
+    log: Log,
+    side_lost_event: asyncio.Event,
+) -> T:
+    """Drive `work`, beating the side for as long as `work` runs.
+
+    `assignment` is the attempt returned by the dispatch. Each heartbeat
+    includes its identifiers, and errors use its side.
+
+    `ready` is set once the policy has accepted the benchmark's signature. Until
+    it is set the side is beaten with `serving=false`: the server keeps the side
+    in the startup window rather than measuring it against the heartbeat window
+    for the whole of its image pull, and refuses the next beat once the run is
+    released, rather than the first serving one. Setting `ready` ends the wait
+    between beats, so the first serving beat is sent at once.
+
+    `side_lost_event` is set once the server has reclaimed the side, and before
+    the work is cancelled.
+
+    Returns what `work` returned. Raises RunEndedError where the server refused
+    the heartbeat because the run is terminal, RunnerError where it took the side
+    back otherwise, and PeerCompletedError where it answered that the run is
+    over, having cancelled the work in every case: the worker's teardown runs on
+    that cancellation, the same way it does when its bridge dies.
+    """
+    working = asyncio.ensure_future(work)
+    beating = asyncio.ensure_future(
+        _beat_until_lost(
+            client,
+            assignment=assignment,
+            period_s=period_s,
+            ready=ready,
+            log=log,
+            side_lost_event=side_lost_event,
+        )
+    )
+    try:
+        await asyncio.wait({working, beating}, return_when=asyncio.FIRST_COMPLETED)
+        if working.done():
+            # Raises whatever the work raised, which is this side's outcome.
+            return working.result()
+        # Never returns: the beat only finishes by raising.
+        return await beating
+    finally:
+        for task in (working, beating):
+            await _cancel_and_wait(task)
+
+
+async def _beat_until_lost(
+    client: ManifoldApiClient,
+    *,
+    assignment: Assignment,
+    period_s: float,
+    ready: asyncio.Event,
+    log: Log,
+    side_lost_event: asyncio.Event,
+) -> NoReturn:
+    """Beat every `period_s` until the server refuses to take one.
+
+    Replay uploads use their own executor, and all heartbeats for a side use
+    one asynchronous client. Replay uploads therefore cannot delay a
+    heartbeat. Ending the side cancels any in-flight request and closes the
+    client before the side's event loop ends.
+
+    A beat that does not reach Manifold is logged and beaten again. The window
+    holds several beats, so one lost to a deployment that is restarting does not
+    cost the side. A revoked credential is refused the same way here, and the
+    daemon's next poll is what ends the process.
+
+    Raises RunEndedError once the server refuses a heartbeat for a terminal run,
+    RunnerError once it has reclaimed the side otherwise, and
+    PeerCompletedError once it says the peer has completed.
+    """
+    side = assignment.side
+    async with client.assignment_heartbeats() as heartbeats:
+        while True:
+            serving = ready.is_set()
+            try:
+                beat = await heartbeats.send(assignment, serving=serving)
+            except AssignmentNotAssignedToRunnerError as exc:
+                if exc.run_status is not None and exc.run_status in TERMINAL_STATUSES:
+                    raise RunEndedError(exc.run_status) from exc
+                side_lost_event.set()
+                raise RunnerError(
+                    f"The server reclaimed the {side.value}: {exc}"
+                ) from exc
+            except Exception as exc:
+                log(f"A heartbeat did not reach Manifold: {escape(str(exc))}")
+            else:
+                if beat.peer_outcome is SideOutcome.COMPLETED:
+                    raise PeerCompletedError(f"the {_peer_of(side).value} completed")
+            if serving:
+                await asyncio.sleep(period_s)
+            else:
+                # Wait on `ready` as well as the period, so the first serving
+                # beat is sent when the worker comes up rather than at the next
+                # tick.
+                with contextlib.suppress(TimeoutError):
+                    await asyncio.wait_for(ready.wait(), period_s)
+
+
+async def _cancel_and_wait(task: asyncio.Task[Any]) -> None:
+    """Cancel `task` and finish its cleanup despite repeated cancellation."""
+    if not task.done():
+        task.cancel()
+        while not task.done():
+            try:
+                await asyncio.shield(task)
+            except asyncio.CancelledError:
+                continue
+    with contextlib.suppress(asyncio.CancelledError):
+        task.exception()
+
+
+def _peer_of(side: RunSide) -> RunSide:
+    """The side of the pairing assigned to the other runner."""
+    return RunSide.POLICY if side is RunSide.BENCHMARK else RunSide.BENCHMARK

manifold_cli/runner/loopback.py

@@ -0,0 +1,102 @@
+"""Worker plans with no container, no GPU, and nothing to pull.
+
+The loopback pair is the double for the adapter contract, and what the first
+end-to-end run exercises: the policy worker serves the built-in echo policy on
+the address it is given, and the benchmark worker round-trips payloads to it
+through the bridge and writes the results rollup a benchmark image writes. Both
+are modules of this CLI, spawned with the interpreter running the runner, so a
+plan they produce needs no operator setup.
+
+They ignore a version's `config`, the cache directory and the GPU arguments;
+the plan they build is fixed. They do not ignore the shard: a double that ran
+the whole episode set however it was partitioned would not stand in for the
+thing being doubled.
+"""
+
+import sys
+from pathlib import Path
+from typing import Any
+from uuid import UUID
+
+from manifold_cli.runner.plan import (
+    RESULTS_ROOT,
+    ProcessOnly,
+    ShardAssignment,
+    SidePreparation,
+    TcpProbe,
+    WireBinding,
+    WorkerPlan,
+)
+
+# An echo policy binds its socket as soon as the interpreter is up.
+READY_TIMEOUT_SECONDS = 5.0
+
+
+class LoopbackAdapter:
+    """Plans the echo policy and the ping benchmark that talk to each other."""
+
+    def plan_policy(
+        self,
+        config: dict[str, Any],
+        *,
+        run_id: UUID,
+        assignment_id: UUID,
+        wire: WireBinding,
+        cache_dir: Path | None = None,
+        gpu_device: int | None = None,
+    ) -> WorkerPlan:
+        """Plan an echo policy serving on `wire`."""
+        return WorkerPlan(
+            argv=(
+                sys.executable,
+                # Unbuffered: the runner pipes a worker's output, and a pipe
+                # makes the interpreter block-buffer, which would hold these
+                # lines back until the worker exited.
+                "-u",
+                "-m",
+                "manifold_cli.runner.loopback_policy",
+                "--listen",
+                wire.address,
+            ),
+            readiness=TcpProbe(wire=wire, timeout_s=READY_TIMEOUT_SECONDS),
+        )
+
+    def plan_side_preparation(self, config: dict[str, Any]) -> SidePreparation | None:
+        """Do not plan a preparation for the built-in workers."""
+        return None
+
+    def plan_benchmark(
+        self,
+        config: dict[str, Any],
+        *,
+        run_id: UUID,
+        assignment_id: UUID,
+        wire: WireBinding,
+        shard: ShardAssignment,
+        live_view_wire: WireBinding | None = None,
+        results_root: Path = RESULTS_ROOT,
+        cache_dir: Path | None = None,
+        gpu_device: int | None = None,
+    ) -> WorkerPlan:
+        """Plan a ping benchmark driving `shard`, writing under `results_root`."""
+        output_dir = results_root / run_id.hex / shard.label
+        return WorkerPlan(
+            argv=(
+                sys.executable,
+                "-u",
+                "-m",
+                "manifold_cli.runner.loopback_bench",
+                "--server",
+                wire.address,
+                "--output-dir",
+                str(output_dir),
+                "--shard-index",
+                str(shard.index),
+                "--num-shards",
+                str(shard.count),
+                "--total-episodes",
+                str(shard.total_episodes),
+            ),
+            readiness=ProcessOnly(),
+            results_dir=output_dir / "results",
+        )

manifold_cli/runner/loopback_bench.py

@@ -0,0 +1,183 @@
+"""The stand-in benchmark worker: round-trips payloads through the bridge.
+
+`LoopbackAdapter` spawns this module. It stands in for a benchmark container: it
+dials the address the benchmark bridge accepts on, drives its shard of the run's
+episodes as echo round-trips against the policy on the far side, and writes a
+results rollup under its output directory, where a benchmark image writes its
+results.
+
+The episode count is the run's, taken from `--total-episodes`, as a real image
+takes it from the run rather than deciding it. This module only says how a
+single episode is driven.
+
+The rollup is what the runner reads back from a benchmark image, because the
+bridge relays the wire verbatim and an image's own protocol never reaches the
+runner. So this one carries the layout and the shape a real image writes:
+`<output-dir>/results/<benchmark>.json`, a flat document keyed `records` with one
+entry per episode -- what manifold-sdk's `write_rollup` leaves behind and
+`parse_episode_rollup` reads back. The runner scans the results directory in a
+later slice.
+"""
+
+import argparse
+import asyncio
+import json
+import os
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+
+from manifold_cli.transport.diagnostics import ping
+from manifold_cli.transport.interfaces import Connector
+from manifold_cli.transport.tcp import TcpConnector, split_host_port
+
+BENCHMARK = "loopback"
+
+# One round-trip stands in for one step: the echo policy answers every payload,
+# as a policy answers an observation with an action.
+STEPS_PER_EPISODE = 4
+
+# One entry per task, with the payload every round-trip of it carries. The
+# global episode set cycles them, so how many episodes there are is the run's
+# to state and not this module's. Two of them, because a record's episode index
+# is global across tasks and a rollup of a single task would not show it.
+TASKS: tuple[tuple[str, int], ...] = (
+    ("echo-4k", 4 * 1024),
+    ("echo-64k", 64 * 1024),
+)
+
+
+@dataclass(frozen=True)
+class EpisodeRecord:
+    """One episode as the rollup reports it."""
+
+    episode_idx: int
+    task_name: str
+    success: bool
+    steps: int
+    initialization_sec: float
+    started_at: datetime
+    ended_at: datetime
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Ping a policy through the bridge.")
+    parser.add_argument(
+        "--server", required=True, help="Address to dial, as HOST:PORT."
+    )
+    parser.add_argument(
+        "--output-dir", required=True, help="Directory to write the rollup under."
+    )
+    parser.add_argument(
+        "--total-episodes",
+        required=True,
+        type=int,
+        help="Episodes the whole run drives, which the shards divide.",
+    )
+    parser.add_argument(
+        "--num-shards", required=True, type=int, help="Shards the run is split into."
+    )
+    parser.add_argument(
+        "--shard-index", required=True, type=int, help="Which of them this one is."
+    )
+    args = parser.parse_args()
+    raise SystemExit(
+        asyncio.run(
+            _run(
+                args.server,
+                Path(args.output_dir),
+                total_episodes=args.total_episodes,
+                num_shards=args.num_shards,
+                shard_index=args.shard_index,
+            )
+        )
+    )
+
+
+def _episode_ids(total_episodes: int, num_shards: int, shard_index: int) -> list[int]:
+    """The global episode ids this shard owns, as manifold-sdk partitions them.
+
+    Shard `i` owns every index where `idx % num_shards == i`, which covers the
+    set exactly once. The ids stay global, so a record names the episode the run
+    planned rather than the position it happened to hold in one shard.
+    """
+    return [idx for idx in range(total_episodes) if idx % num_shards == shard_index]
+
+
+async def _run(
+    server: str,
+    output_dir: Path,
+    *,
+    total_episodes: int,
+    num_shards: int,
+    shard_index: int,
+) -> int:
+    host, port = split_host_port(server)
+    connector = TcpConnector(host, port)
+    records: list[EpisodeRecord] = []
+    for episode_idx in _episode_ids(total_episodes, num_shards, shard_index):
+        task, payload_bytes = TASKS[episode_idx % len(TASKS)]
+        record = await _run_episode(
+            connector, task, payload_bytes, episode_idx=episode_idx
+        )
+        elapsed_ms = (record.ended_at - record.started_at).total_seconds() * 1000
+        _say(
+            f"{task} episode {record.episode_idx}: {record.steps} round-trips "
+            f"of {payload_bytes} bytes in {elapsed_ms:.2f} ms"
+        )
+        records.append(record)
+    path = _write_rollup(output_dir, records)
+    _say(f"wrote the rollup to {path}")
+    return 0 if all(record.success for record in records) else 1
+
+
+async def _run_episode(
+    connector: Connector, task: str, payload_bytes: int, *, episode_idx: int
+) -> EpisodeRecord:
+    """Round-trip one payload per step, and record what the steps came to."""
+    started_at = datetime.now(UTC)
+    matched = 0
+    for _ in range(STEPS_PER_EPISODE):
+        result = await ping(connector, payload=os.urandom(payload_bytes))
+        matched += 1 if result.matched else 0
+    return EpisodeRecord(
+        episode_idx=episode_idx,
+        task_name=task,
+        success=matched == STEPS_PER_EPISODE,
+        steps=STEPS_PER_EPISODE,
+        initialization_sec=0.0,
+        started_at=started_at,
+        ended_at=datetime.now(UTC),
+    )
+
+
+def _write_rollup(output_dir: Path, records: Sequence[EpisodeRecord]) -> Path:
+    """Write the rollup where a supervisor scans for one."""
+    rollup = {
+        "records": [
+            {
+                "episode_idx": record.episode_idx,
+                "task_name": record.task_name,
+                "success": record.success,
+                "steps": record.steps,
+                "initialization_sec": record.initialization_sec,
+                "started_at": record.started_at.isoformat(),
+                "ended_at": record.ended_at.isoformat(),
+            }
+            for record in records
+        ]
+    }
+    results_dir = output_dir / "results"
+    results_dir.mkdir(parents=True, exist_ok=True)
+    path = results_dir / f"{BENCHMARK}.json"
+    path.write_text(json.dumps(rollup))
+    return path
+
+
+def _say(message: str) -> None:
+    print(message, flush=True)
+
+
+if __name__ == "__main__":
+    main()