Isaac-GR00T/tests/test_support/runtime.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Shared subprocess/runtime helpers for tests."""

from __future__ import annotations

import contextlib
import json
import os
import pathlib
import re
import shutil
import socket
import subprocess
import tempfile
import time


DEFAULT_SERVER_STARTUP_SECONDS = 600.0


def _default_cache_path() -> pathlib.Path:
    """Return the cache root directory."""
    if "TEST_CACHE_PATH" in os.environ:
        return pathlib.Path(os.environ["TEST_CACHE_PATH"])

    local_fallback = pathlib.Path.home() / ".cache" / "g00t"
    local_fallback.mkdir(parents=True, exist_ok=True)
    return local_fallback


TEST_CACHE_PATH = _default_cache_path()


def get_root() -> pathlib.Path:
    """Return the root directory of the repository."""
    return pathlib.Path(__file__).resolve().parents[2]


def fast_copy_tree(
    src: pathlib.Path,
    dst: pathlib.Path,
    *,
    symlinks: bool = False,
) -> None:
    """Recursively copy *src* → *dst*, optimised for slow-NFS sources.

    Many-small-files trees over a contended NFS mount stall
    ``shutil.copytree`` badly: each file pays an open + stat + read +
    close NFS round-trip, so metadata-bound trees (e.g. a Python venv
    with pandas/sympy/torch test fixtures, ~50k tiny files) collapse
    from thousands of files/s on a healthy mount to single-digit
    files/s on a contended one. Job 313566325 observed 3.4 files/s,
    blowing the 1200 s ``pytest-timeout`` on a single venv copy that
    finishes in <60 s on a quiet runner.

    This helper streams the tree through a single ``tar`` pipe — one
    sequential read on the source side, one sequential write on the
    destination, no per-file network round-trips. Falls back to
    ``shutil.copytree(dirs_exist_ok=True)`` when ``tar`` is missing
    (e.g. some minimal containers) so the helper is always usable.

    Args:
        src: source directory (must exist).
        dst: destination directory (created when absent; existing
            entries are merged, matching
            ``shutil.copytree(dirs_exist_ok=True)``).
        symlinks: when ``False`` (default) symlinks are dereferenced
            and real bytes are copied — load-bearing for ``/shared``
            sources, where preserving symlinks would route runtime
            reads back over NFS and defeat the purpose of staging.
            When ``True`` symlinks are preserved verbatim (used for
            ``libero_uv`` venvs whose internal ``bin/python`` symlink
            points outside the tree and must stay a symlink).
    """
    src = pathlib.Path(src)
    dst = pathlib.Path(dst)
    if not src.is_dir():
        raise NotADirectoryError(f"fast_copy_tree source is not a directory: {src}")
    dst.mkdir(parents=True, exist_ok=True)

    tar_bin = shutil.which("tar")
    if tar_bin is None:
        shutil.copytree(src, dst, dirs_exist_ok=True, symlinks=symlinks)
        return

    # ``tar -h`` (--dereference) follows symlinks at archive-creation
    # time, matching ``shutil.copytree(symlinks=False)``.
    create_flags = "-cpf" if symlinks else "-chpf"
    src_cmd = [tar_bin, "-C", str(src), create_flags, "-", "."]
    dst_cmd = [tar_bin, "-C", str(dst), "-xpf", "-"]

    src_proc = subprocess.Popen(src_cmd, stdout=subprocess.PIPE)
    try:
        try:
            dst_proc = subprocess.Popen(dst_cmd, stdin=src_proc.stdout)
        except BaseException:
            # dst_proc never started, so nobody will drain src_proc's stdout.
            # Reap src_proc explicitly to avoid leaving a zombie.
            src_proc.kill()
            src_proc.wait()
            raise
    finally:
        # Close our copy of the pipe so dst_proc gets SIGPIPE if it dies
        # before src_proc finishes writing — otherwise we'd deadlock.
        if src_proc.stdout is not None:
            src_proc.stdout.close()

    dst_rc = dst_proc.wait()
    src_rc = src_proc.wait()
    if src_rc != 0 or dst_rc != 0:
        raise subprocess.CalledProcessError(
            dst_rc or src_rc,
            f"tar-pipe {src} → {dst} (src_rc={src_rc} dst_rc={dst_rc})",
        )


def checkpoint_tree_ready(path: pathlib.Path) -> bool:
    """Return True if *path* looks like a HuggingFace ``transformers`` checkpoint dir."""
    if not (path / "config.json").is_file():
        return False
    index_file = path / "model.safetensors.index.json"
    if not index_file.is_file():
        return True
    shards = set(json.loads(index_file.read_text()).get("weight_map", {}).values())
    return all((path / shard).is_file() for shard in shards)


# Backward-compatible alias used by other test files.
libero_checkpoint_tree_ready = checkpoint_tree_ready


def demo_dataset_tree_ready(path: pathlib.Path) -> bool:
    """Return True if *path* looks like a bundled LeRobot demo dataset.

    Checks for ``meta/modality.json``, at least one ``.parquet`` under ``data/``,
    and at least one ``.mp4`` under ``videos/``.
    """
    if not (path / "meta" / "modality.json").is_file():
        return False
    data_dir = path / "data"
    if not data_dir.is_dir() or not any(data_dir.rglob("*.parquet")):
        return False
    videos_dir = path / "videos"
    if not videos_dir.is_dir() or not any(videos_dir.rglob("*.mp4")):
        return False
    return True


# Backward-compatible alias used by other test files.
libero_demo_tree_ready = demo_dataset_tree_ready


# ---------------------------------------------------------------------------
# Generic model checkpoint resolver
# ---------------------------------------------------------------------------


def resolve_model_checkpoint_path(
    *,
    hf_repo_id: str,
    hf_subdir: str | None = None,
    path_override_env: str | None = None,
    repo_root: pathlib.Path | None = None,
) -> pathlib.Path:
    """Resolve a GR00T model checkpoint, downloading from HuggingFace if needed.

    Resolution order:

    1. Environment variable *path_override_env* (must be a complete checkpoint).
    2. ``<repo_root>/checkpoints/<model_name>/<subdir>``.
    3. Git worktree toplevel + same relative ``checkpoints/...`` path.
    4. ``TEST_CACHE_PATH/models/<model_name>/<subdir>``, downloading from
       HuggingFace when missing (requires ``HF_TOKEN``).

    Args:
        hf_repo_id: HuggingFace repo id, e.g. ``"nvidia/GR00T-N1.7-LIBERO"``.
        hf_subdir: Optional subdirectory within the repo (e.g. ``"libero_10"``).
        path_override_env: Name of an env var that, when set, overrides all other
            resolution.
        repo_root: Repository root (auto-detected if ``None``).
    """
    root = repo_root if repo_root is not None else get_root()
    model_name = hf_repo_id.split("/")[-1]
    rel_path = f"checkpoints/{model_name}/{hf_subdir}" if hf_subdir else f"checkpoints/{model_name}"

    if path_override_env:
        override = os.environ.get(path_override_env, "").strip()
        if override:
            p = pathlib.Path(override).expanduser().resolve()
            assert checkpoint_tree_ready(p), (
                f"{path_override_env} does not point to a complete checkpoint directory: {p}"
            )
            return p

    local = root / rel_path
    if checkpoint_tree_ready(local):
        return local

    try:
        toplevel = subprocess.check_output(
            ["git", "rev-parse", "--show-toplevel"],
            cwd=str(root),
            text=True,
            timeout=30,
        ).strip()
        git_cp = pathlib.Path(toplevel) / rel_path
        if checkpoint_tree_ready(git_cp):
            return git_cp
    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
        pass

    shared_root = TEST_CACHE_PATH / "models" / model_name
    shared = shared_root / hf_subdir if hf_subdir else shared_root
    if checkpoint_tree_ready(shared):
        return shared

    allow = [f"{hf_subdir}/*"] if hf_subdir else None
    token = os.environ.get("HF_TOKEN", "")
    assert token, "HF_TOKEN is required to download gated models. Set via: export HF_TOKEN=hf_..."

    from huggingface_hub import snapshot_download

    snapshot_download(
        repo_id=hf_repo_id,
        local_dir=str(shared_root),
        token=token,
        **({"allow_patterns": allow} if allow else {}),
    )
    assert checkpoint_tree_ready(shared), (
        f"Checkpoint at {shared} is incomplete after resolve/download "
        "(check HF_TOKEN, network, and Hugging Face repo layout)."
    )
    return shared


# ---------------------------------------------------------------------------
# Generic demo dataset resolver
# ---------------------------------------------------------------------------


def resolve_demo_dataset(
    *,
    dataset_name: str,
    path_override_env: str | None = None,
    global_env_var: str | None = None,
    hf_download_env_var: str | None = None,
    repo_root: pathlib.Path | None = None,
) -> pathlib.Path:
    """Resolve a bundled LeRobot demo dataset by *dataset_name*.

    Resolution order (first match wins):

    1. *path_override_env* environment variable (per-test override).
    2. *global_env_var* environment variable (CI / local override).
    3. ``<repo_root>/demo_data/<dataset_name>`` — normal clone with Git LFS.
    4. ``TEST_CACHE_PATH/datasets/<dataset_name>`` — shared PVC / local cache.
    5. If *hf_download_env_var* is set in the environment, download from HuggingFace
       into the shared path (requires ``HF_TOKEN``).

    Args:
        dataset_name: Directory name under ``demo_data/`` (e.g. ``"libero_demo"``
            or ``"droid_sample"``).
        path_override_env: Name of an env var for per-test path override.
        global_env_var: Name of an env var for a global path override
            (e.g. ``"LIBERO_DEMO_DATASET_PATH"``).
        hf_download_env_var: Name of an env var whose value is a HuggingFace
            *dataset* repo id; when set, the dataset is downloaded into shared
            storage.
        repo_root: Repository root (auto-detected if ``None``).
    """
    root = repo_root if repo_root is not None else get_root()

    if path_override_env:
        alt = os.environ.get(path_override_env, "").strip()
        if alt:
            resolved = pathlib.Path(alt).expanduser().resolve()
            assert demo_dataset_tree_ready(resolved), (
                f"{path_override_env} does not point to a complete demo dataset: {resolved}"
            )
            return resolved

    if global_env_var:
        env_path = os.environ.get(global_env_var, "").strip()
        if env_path:
            resolved = pathlib.Path(env_path).expanduser().resolve()
            assert demo_dataset_tree_ready(resolved), (
                f"{global_env_var} does not point to a complete demo dataset tree: {resolved}"
            )
            return resolved

    in_repo = root / "demo_data" / dataset_name
    if demo_dataset_tree_ready(in_repo):
        return in_repo

    shared = TEST_CACHE_PATH / "datasets" / dataset_name
    if demo_dataset_tree_ready(shared):
        return shared

    if hf_download_env_var:
        hf_dataset = os.environ.get(hf_download_env_var, "").strip()
        if hf_dataset:
            token = os.environ.get("HF_TOKEN", "")
            assert token, (
                f"HF_TOKEN is required to download {hf_download_env_var} into shared storage"
            )
            shared.parent.mkdir(parents=True, exist_ok=True)
            from huggingface_hub import snapshot_download

            snapshot_download(
                repo_id=hf_dataset,
                repo_type="dataset",
                local_dir=str(shared),
                token=token,
            )
            assert demo_dataset_tree_ready(shared), (
                f"Downloaded HuggingFace dataset into {shared} but it does not match "
                "expected LeRobot demo dataset layout"
            )
            return shared

    raise AssertionError(
        f"{dataset_name} dataset not found. "
        f"It ships in-repo under demo_data/{dataset_name} (requires Git LFS). "
        f"Alternatives: set {global_env_var or path_override_env or 'a path env var'} "
        f"to an existing checkout; populate {shared} on the shared drive "
        "(CI_SHARED_DRIVE_PATH / ~/.cache/g00t)."
    )


# ---------------------------------------------------------------------------
# LIBERO-specific convenience wrappers (backward compatibility)
# ---------------------------------------------------------------------------

_LIBERO_N17_LIBERO_REPO = "nvidia/GR00T-N1.7-LIBERO"
_LIBERO_N17_LIBERO_SUBDIR = "libero_10"


def resolve_libero_n17_libero10_checkpoint_path(
    repo_root: pathlib.Path | None = None,
    *,
    path_override_env: str,
) -> pathlib.Path:
    """Resolve the LIBERO-finetuned GR00T-N1.7 checkpoint (``libero_10`` subfolder).

    Thin wrapper around :func:`resolve_model_checkpoint_path` kept for backward
    compatibility with existing callers.
    """
    return resolve_model_checkpoint_path(
        hf_repo_id=_LIBERO_N17_LIBERO_REPO,
        hf_subdir=_LIBERO_N17_LIBERO_SUBDIR,
        path_override_env=path_override_env,
        repo_root=repo_root,
    )


def resolve_libero_demo_dataset_path(
    repo_root: pathlib.Path | None = None,
    *,
    path_override_env: str | None = None,
) -> pathlib.Path:
    """Return the path to the LIBERO ``libero_demo`` dataset.

    Thin wrapper around :func:`resolve_demo_dataset` kept for backward
    compatibility with existing callers.
    """
    return resolve_demo_dataset(
        dataset_name="libero_demo",
        path_override_env=path_override_env,
        global_env_var="LIBERO_DEMO_DATASET_PATH",
        hf_download_env_var="GR00T_LIBERO_DEMO_HF_DATASET",
        repo_root=repo_root,
    )


_DROID_N17_REPO = "nvidia/GR00T-N1.7-DROID"


def resolve_droid_n17_checkpoint_path(
    repo_root: pathlib.Path | None = None,
    *,
    path_override_env: str,
) -> pathlib.Path:
    """Resolve the DROID-finetuned GR00T-N1.7 checkpoint.

    Resolution order:

    1. Environment variable named by *path_override_env* (must be a complete checkpoint).
    2. ``<repo_root>/checkpoints/GR00T-N1.7-DROID``.
    3. Git worktree toplevel + same relative ``checkpoints/...`` path.
    4. ``TEST_CACHE_PATH/models/GR00T-N1.7-DROID``, downloading
       from Hugging Face when missing (requires ``HF_TOKEN``).

    Raises:
        AssertionError: if overrides are incomplete or download leaves a broken tree.
    """
    root = repo_root if repo_root is not None else get_root()

    override = os.environ.get(path_override_env, "").strip()
    if override:
        p = pathlib.Path(override).expanduser().resolve()
        assert libero_checkpoint_tree_ready(p), (
            f"{path_override_env} does not point to a complete checkpoint directory: {p}"
        )
        return p

    local = root / "checkpoints/GR00T-N1.7-DROID"
    if libero_checkpoint_tree_ready(local):
        return local

    try:
        toplevel = subprocess.check_output(
            ["git", "rev-parse", "--show-toplevel"],
            cwd=str(root),
            text=True,
            timeout=30,
        ).strip()
        git_cp = pathlib.Path(toplevel) / "checkpoints/GR00T-N1.7-DROID"
        if libero_checkpoint_tree_ready(git_cp):
            return git_cp
    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
        pass

    return resolve_model_checkpoint_path(
        hf_repo_id=_DROID_N17_REPO,
        path_override_env=path_override_env,
        repo_root=repo_root,
    )


def resolve_droid_demo_dataset_path(
    repo_root: pathlib.Path | None = None,
    *,
    path_override_env: str | None = None,
) -> pathlib.Path:
    """Return the path to the DROID ``droid_sample`` dataset (small LeRobot bundle).

    This is the 3-episode DROID demo described in the README under
    ``demo_data/droid_sample`` (Git LFS in the Isaac-GR00T repo).

    Resolution order (first match wins):

    0. If *path_override_env* is set and that variable is non-empty in the
       environment, its path is used (must satisfy :func:`libero_demo_tree_ready`).
    1. ``DROID_DEMO_DATASET_PATH`` — explicit directory (CI or local override).
    2. ``<repo_root>/demo_data/droid_sample`` — normal clone with Git LFS.
    3. ``TEST_CACHE_PATH/datasets/droid_sample`` — shared PVC / local cache.

    Raises:
        AssertionError: if no usable tree is found.
    """
    root = repo_root if repo_root is not None else get_root()

    if path_override_env:
        alt = os.environ.get(path_override_env, "").strip()
        if alt:
            resolved = pathlib.Path(alt).expanduser().resolve()
            assert libero_demo_tree_ready(resolved), (
                f"{path_override_env} does not point to a complete droid_sample-style dataset: "
                f"{resolved}"
            )
            return resolved

    env_path = os.environ.get("DROID_DEMO_DATASET_PATH", "").strip()
    if env_path:
        resolved = pathlib.Path(env_path).expanduser().resolve()
        assert libero_demo_tree_ready(resolved), (
            f"DROID_DEMO_DATASET_PATH does not point to a complete droid_sample tree: {resolved}"
        )
        return resolved

    in_repo = root / "demo_data" / "droid_sample"
    if libero_demo_tree_ready(in_repo):
        return in_repo

    shared = TEST_CACHE_PATH / "datasets" / "droid_sample"
    if libero_demo_tree_ready(shared):
        return shared

    raise AssertionError(
        "droid_sample dataset not found. It ships in-repo under demo_data/droid_sample "
        "(requires Git LFS). Alternatives: set DROID_DEMO_DATASET_PATH to an existing checkout; "
        f"populate {shared} on the shared drive (CI_SHARED_DRIVE_PATH / ~/.cache/g00t)."
    )


EGL_VENDOR_DIRS = [
    pathlib.Path("/usr/share/glvnd/egl_vendor.d"),
    pathlib.Path("/etc/glvnd/egl_vendor.d"),
    pathlib.Path("/usr/local/share/glvnd/egl_vendor.d"),
]


def hf_hub_download_cmd(repo_id: str, filename: str, local_dir: str) -> list[str]:
    """Build a ``uv run python -c`` command that downloads a file from HuggingFace.

    Reads HF_TOKEN from the environment and passes it explicitly so gated repos
    work without requiring ``huggingface-cli login``.  Raises AssertionError if
    HF_TOKEN is not set.
    """
    token = os.environ.get("HF_TOKEN", "")
    assert token, (
        "HF_TOKEN environment variable is not set. "
        "A HuggingFace token with access to gated repos is required. "
        "Set it via: export HF_TOKEN=hf_..."
    )
    return [
        "uv",
        "run",
        "python",
        "-c",
        f"from huggingface_hub import hf_hub_download; "
        f"hf_hub_download(repo_id={repo_id!r}, filename={filename!r}, "
        f"local_dir={local_dir!r}, token={token!r})",
    ]


# GPU names that contain these tokens are known to have RT cores.
# Compute-only data-center GPUs (A100, H100, H200, B200, V100, etc.) do not.
_RT_CORE_GPU_PATTERNS = (
    r"\brtx\b",  # RTX 20xx/30xx/40xx/50xx, Quadro RTX, RTX Ax000
    r"\bl40s?\b",  # L40 / L40S
    r"\bl4\b",  # L4
)


def has_rt_core_gpu() -> bool:
    """Return True if any available GPU has RT cores (required for Vulkan ray tracing).

    Checks ``nvidia-smi`` GPU names against known RT-capable product lines.
    Returns False if nvidia-smi is unavailable or no matching GPU is found.
    """
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode != 0:
            return False
        for name in result.stdout.strip().splitlines():
            if any(re.search(pat, name.strip().lower()) for pat in _RT_CORE_GPU_PATTERNS):
                return True
    except Exception:
        pass
    return False


def find_nvidia_egl_vendor_file() -> pathlib.Path:
    """Return the first NVIDIA EGL vendor JSON file found, or raise FileNotFoundError."""
    for vendor_dir in EGL_VENDOR_DIRS:
        for candidate in vendor_dir.glob("*nvidia*.json") if vendor_dir.is_dir() else []:
            return candidate
    searched = ", ".join(str(d) for d in EGL_VENDOR_DIRS)
    raise FileNotFoundError(
        f"NVIDIA EGL vendor file not found (searched: {searched}). "
        "robosuite requires EGL_PLATFORM_DEVICE_EXT which is only provided by the "
        "NVIDIA EGL implementation. Install the NVIDIA GL/EGL packages or run on a "
        "host with the full NVIDIA driver stack."
    )


def resolve_shared_uv_cache_dir() -> pathlib.Path | None:
    """Return a writable uv cache path, or None.

    Only redirects the uv cache when TEST_CACHE_PATH is set — on dev
    machines uv's default cache (~/.cache/uv) is already local and fast, so
    there is no benefit to overriding it.
    """
    if "TEST_CACHE_PATH" not in os.environ:
        return None
    cache_dir = TEST_CACHE_PATH / "uv-cache"
    try:
        cache_dir.mkdir(parents=True, exist_ok=True)
        return cache_dir
    except OSError:
        print(
            f"[cache] warning: uv cache unavailable at {cache_dir}; "
            "falling back to uv default cache dir"
        )
        return None


def build_shared_hf_cache_env(cache_key: str) -> dict[str, str]:
    """Build HF cache environment variables for a cache key."""
    hf_cache_dir = TEST_CACHE_PATH / f"hf-cache/{cache_key}"
    try:
        hub_cache_dir = hf_cache_dir / "hub"
        transformers_cache_dir = hf_cache_dir / "transformers"
        datasets_cache_dir = hf_cache_dir / "datasets"
        hub_cache_dir.mkdir(parents=True, exist_ok=True)
        transformers_cache_dir.mkdir(parents=True, exist_ok=True)
        datasets_cache_dir.mkdir(parents=True, exist_ok=True)
    except OSError:
        print(
            f"[cache] warning: Hugging Face cache unavailable at {hf_cache_dir}; "
            "falling back to defaults"
        )
        return {}

    return {
        "HF_HOME": str(hf_cache_dir),
        "HF_HUB_CACHE": str(hub_cache_dir),
        "HUGGINGFACE_HUB_CACHE": str(hub_cache_dir),
        "TRANSFORMERS_CACHE": str(transformers_cache_dir),
        "HF_DATASETS_CACHE": str(datasets_cache_dir),
    }


def assert_port_available(host: str, port: int) -> None:
    """Raise AssertionError if the port is already bound.

    Call this before starting a model server subprocess to catch port conflicts
    early (e.g. a leftover process from a previous test run or two tests
    inadvertently assigned the same port).
    """
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            s.bind((host, port))
        except OSError as exc:
            raise AssertionError(
                f"Port {port} on {host} is already in use. "
                "Each test file uses a unique port — check for a conflicting "
                "process or a previous test run that did not shut down cleanly."
            ) from exc


def start_server_process(
    server_code: str,
    *,
    cwd: pathlib.Path,
    env: dict[str, str],
) -> tuple[subprocess.Popen, pathlib.Path]:
    """Start a model server subprocess with stderr captured to a temp file.

    Returns the Popen object and the path to the stderr log file.  On failure
    the caller should read and print the log so CI output includes the error.
    """
    stderr_log = pathlib.Path(tempfile.mktemp(prefix="server_stderr_", suffix=".log"))
    stderr_fh = open(stderr_log, "w")  # noqa: SIM115
    proc = subprocess.Popen(
        ["bash", "-c", server_code],
        cwd=cwd,
        env=env,
        stdout=stderr_fh,
        stderr=stderr_fh,
    )
    return proc, stderr_log


def _dump_server_log(log_path: pathlib.Path, tail_chars: int = 8000) -> str:
    """Read the tail of a server log file and return it as a string."""
    try:
        text = log_path.read_text()
        return text[-tail_chars:] if len(text) > tail_chars else text
    except OSError:
        return "<server log not available>"


def wait_for_server_ready(
    proc: subprocess.Popen,
    host: str,
    port: int,
    timeout_s: float,
    server_log: pathlib.Path | None = None,
) -> None:
    """Wait until the server accepts TCP connections, or raise if it dies/times out."""
    deadline = time.monotonic() + timeout_s
    while True:
        if proc.poll() is not None:
            log_info = ""
            if server_log is not None:
                log_info = f"\nServer output:\n{_dump_server_log(server_log)}"
            raise AssertionError(
                f"Model server failed to start.\nreturncode={proc.returncode}{log_info}"
            )
        try:
            with socket.create_connection((host, port), timeout=1.0):
                elapsed = time.monotonic() - deadline + timeout_s
                print(f"Model server ready after {elapsed:.1f}s.")
                return
        except OSError:
            if time.monotonic() >= deadline:
                if proc.poll() is None:
                    proc.terminate()
                    try:
                        proc.wait(timeout=15)
                    except subprocess.TimeoutExpired:
                        proc.kill()
                        proc.wait(timeout=15)
                log_info = ""
                if server_log is not None:
                    log_info = f"\nServer output:\n{_dump_server_log(server_log)}"
                raise AssertionError(
                    "Model server did not become ready before timeout.\n"
                    f"timeout_seconds={timeout_s}\n"
                    f"Set the corresponding env var to override.{log_info}"
                )
            time.sleep(0.5)


def run_subprocess_step(
    cmd: list[str],
    *,
    step: str,
    cwd: pathlib.Path,
    env: dict[str, str],
    timeout_s: int | float | None = None,
    stream_output: bool = False,
    log_prefix: str = "examples",
    failure_prefix: str = "Subprocess step failed",
    output_tail_chars: int = 8000,
) -> tuple[subprocess.CompletedProcess, float]:
    """Run a subprocess step with consistent timing/logging/failure formatting."""
    print(f"[{log_prefix}] step={step} command={' '.join(cmd)}", flush=True)
    start = time.perf_counter()
    run_kwargs = {
        "cwd": cwd,
        "env": env,
        "check": False,
    }
    if timeout_s is not None:
        run_kwargs["timeout"] = timeout_s
    if not stream_output:
        run_kwargs["capture_output"] = True
        run_kwargs["text"] = True
    result = subprocess.run(cmd, **run_kwargs)
    elapsed_s = time.perf_counter() - start
    print(f"[{log_prefix}] step={step} elapsed_s={elapsed_s:.2f}", flush=True)

    if result.returncode != 0:
        if stream_output:
            output_info = "See streamed test logs above for subprocess output."
        else:
            output = (result.stdout or "") + (result.stderr or "")
            output_info = f"output_tail=\n{output[-output_tail_chars:]}"
        raise AssertionError(
            f"{failure_prefix}: {step}\n"
            f"elapsed_s={elapsed_s:.2f}\n"
            f"returncode={result.returncode}\n"
            f"command={' '.join(cmd)}\n"
            f"{output_info}"
        )
    return result, elapsed_s


@contextlib.contextmanager
def timed(label: str):
    """Context manager that prints the wall-clock duration of a labelled phase.

    Usage::

        with timed("model load"):
            model = load_model(...)
    """
    print(f"[timing] {label} — starting", flush=True)
    t0 = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - t0
        print(f"[timing] {label} — done in {elapsed:.1f}s", flush=True)