verifiers 0.1.15.dev421__tar.gz → 0.1.15.dev423__tar.gz

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: verifiers
-Version: 0.1.15.dev421
+Version: 0.1.15.dev423
 Summary: Verifiers: Environments for LLM Reinforcement Learning
 Project-URL: Homepage, https://github.com/primeintellect-ai/verifiers
 Project-URL: Documentation, https://github.com/primeintellect-ai/verifiers

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/pyproject.toml

@@ -166,6 +166,7 @@ flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE" }
 vf-eval = "verifiers.scripts.eval:main"
 eval = "verifiers.v1.cli.eval.main:main"
 validate = "verifiers.v1.cli.validate:main"
+debug = "verifiers.v1.cli.debug:main"
 serve = "verifiers.v1.cli.serve:main"
 init = "verifiers.v1.cli.init:main"
 vf-gepa = "verifiers.scripts.gepa:main"

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/tests/test_sandbox_debug_env.py

@@ -33,12 +33,13 @@ def test_swe_debug_env_is_deprecated_subclass_not_alias():
     assert issubclass(SWEDebugEnv, SandboxDebugEnv)
     assert SandboxDebugRubric.__name__ == "SandboxDebugRubric"
 
-    with pytest.warns(DeprecationWarning, match="use SandboxDebugEnv"):
+    with pytest.warns(DeprecationWarning, match="native v1 `debug` CLI") as warnings:
         env = SWEDebugEnv(
             DummySandboxTaskSet(),
             debug_step="command",
             debug_command="true",
         )
 
+    assert len(warnings) == 1
     assert isinstance(env, SandboxDebugEnv)
     assert env.labels == ["sandbox-debug"]

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/verifiers/envs/experimental/composable/sandbox_debug_env.py

@@ -2,7 +2,8 @@
 
 import shlex
 import time
-from typing import Any, Literal
+from typing import Any, ClassVar, Literal
+from warnings import warn
 
 import verifiers as vf
 from prime_sandboxes import CreateSandboxRequest
@@ -38,6 +39,8 @@ class SandboxDebugEnv(SandboxMixin, vf.MultiTurnEnv):
     - exit: optionally run task tests and score them
     """
 
+    emit_deprecation_warning: ClassVar[bool] = True
+
     def __init__(
         self,
         taskset: SandboxTaskSet,
@@ -59,6 +62,12 @@ class SandboxDebugEnv(SandboxMixin, vf.MultiTurnEnv):
         output_tail_chars: int = 2000,
         **sandbox_kwargs: Any,
     ):
+        if self.emit_deprecation_warning:
+            warn(
+                "SandboxDebugEnv is deprecated; use the native v1 `debug` CLI for v1 tasksets.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         if debug_step not in ("none", "gold_patch", "command", "script"):
             raise ValueError(f"Unsupported debug_step: {debug_step!r}")
         if debug_step == "command" and not debug_command:

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/verifiers/envs/experimental/composable/swe_debug_env.py

@@ -11,7 +11,8 @@ class SWEDebugRubric(SandboxDebugRubric):
 
     def __init__(self, **kwargs: Any):
         warn(
-            "SWEDebugRubric is deprecated; use SandboxDebugRubric.",
+            "SWEDebugRubric is deprecated with the composable debug envs; use the "
+            "native v1 `debug` CLI for v1 tasksets.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -21,9 +22,11 @@ class SWEDebugRubric(SandboxDebugRubric):
 class SWEDebugEnv(SandboxDebugEnv):
     """Deprecated wrapper for SandboxDebugEnv."""
 
+    emit_deprecation_warning = False
+
     def __init__(self, *args: Any, **kwargs: Any):
         warn(
-            "SWEDebugEnv is deprecated; use SandboxDebugEnv.",
+            "SWEDebugEnv is deprecated; use the native v1 `debug` CLI for v1 tasksets.",
             DeprecationWarning,
             stacklevel=2,
         )

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/verifiers/v1/GUIDE.md

@@ -767,8 +767,9 @@ pass `docker` / `prime` / `modal`.
 
 # CLI reference
 
-Three commands, all `uv run <cmd>`: **`eval`** (run + score a model), **`validate`** (model-free
-gold check), **`init`** (scaffold). They share a resolution layer:
+Four commands, all `uv run <cmd>`: **`eval`** (run + score a model), **`validate`**
+(model-free setup/gold checks), **`debug`** (setup + one shell action), **`init`**
+(scaffold). They share a resolution layer:
 
 - **Positional id** — a leading bare token is the taskset id: `eval gsm8k-v1` == `eval --taskset.id
   gsm8k-v1` (for `init` it's the env name).
@@ -823,19 +824,24 @@ trace per line, appended as each rollout finishes — durable mid-run), and `eva
 
 ## `validate`
 
-Run each task's `validate` hook — a model-free check that the ground truth holds (the gold patch
-makes the tests pass, the verifier accepts the gold answer) — in a runtime with the taskset's
-`setup` applied. No model, no harness.
+Model-free checks that a task is sound. By default each task gets two independent
+runtimes — a **gold** check (the taskset's `setup` then its `validate` hook: the gold patch
+makes the tests pass, the verifier accepts the gold answer) and a **setup-only** check —
+reported as one aggregate row. Restrict to one check with `--only-gold` / `--only-setup`.
+No model, no harness.
 
 ```bash
 uv run validate gsm8k-v1 -n 20 --runtime.type subprocess
+uv run validate swebench-v1 -n 1 --runtime.type prime --only-setup
+uv run validate swebench-v1 -n 1 --runtime.type prime --only-gold
 ```
 
 | flag | default | meaning |
 | --- | --- | --- |
 | `<taskset-id>` / `--taskset.id` | — | taskset to validate |
 | `--runtime.type` | `docker` | runtime for `setup` + `validate` (a gold check often needs the task's container) |
-| `--setup-timeout` / `--validate-timeout` | None | per-hook wall-clock caps |
+| `--only-setup` / `--only-gold` | off | run just the setup-only or just the gold check (default: both) |
+| `--timeout.setup` / `--timeout.total` | None | per-task wall-clock caps for `setup` and the `validate` hook |
 | `-n`/`--num-tasks`, `-s`/`--shuffle`, `-c`/`--max-concurrent` (128) | | task selection + concurrency |
 | `-v`/`--verbose`, `--no-rich` | | logging / disable the dashboard |
 
@@ -845,6 +851,29 @@ captured as a result row (one bad task is data, not a crash) with reason `valid`
 default runtime is **docker** (unlike eval's subprocess), and a subprocess runtime against a
 `NEEDS_CONTAINER` / image-bearing taskset aborts with a clear error.
 
+## `debug`
+
+Set up each selected task, run one explicit shell action, and save the trace. This is for
+inspecting task state without a model: it does not apply gold patches, call `validate`, run
+`finalize`, or score rewards.
+
+```bash
+uv run debug swebench-v1 -n 1 --runtime.type prime --command 'pwd; git status --short | head'
+uv run debug swebench-v1 -n 1 --runtime.type prime --script-path ./inspect.sh
+```
+
+| flag | default | meaning |
+| --- | --- | --- |
+| `<taskset-id>` / `--taskset.id` | — | taskset to debug |
+| `--command` / `--script-path` | — | exactly one inline command or host script to upload and execute |
+| `--runtime.type` | `docker` | runtime for setup + the debug action |
+| `--timeout.setup` / `--timeout.total` | None | per-task wall-clock caps for `setup` and the debug action |
+| `-n`/`--num-tasks`, `-s`/`--shuffle`, `-c`/`--max-concurrent` (128) | | task selection + concurrency |
+| `-o`/`--output-dir` | fresh debug run dir | where `config.toml` and `results.jsonl` are written |
+
+Each saved trace has command/script metadata, exit status, elapsed time, timeout/error fields,
+and the full stdout/stderr under `trace.info["debug"]`.
+
 ## `init`
 
 Scaffold a new v1 environment package under `--path` (default `./environments`), following the

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/verifiers/v1/README.md

@@ -41,7 +41,8 @@ uv sync --python 3.12 --extra harbor
 ```bash
 uv run init my-task-v1           # scaffold a new environment (--add-tool/--add-user/--add-harness)
 uv run eval gsm8k-v1 -n 5 -r 3   # single-turn math; default harness; subprocess runtime
-uv run validate gsm8k-v1 -n 5    # model-free: run each task's gold check (the `validate` hook)
+uv run validate gsm8k-v1 -n 5    # model-free gold + setup-only checks (--only-gold / --only-setup)
+uv run debug gsm8k-v1 -n 1 --command 'pwd'  # setup, run one shell action, save traces
 uv run eval -h                   # typed help (+ the local example tasksets/harnesses)
 ```
 

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/verifiers/v1/cli/dashboard/eval.py

@@ -497,16 +497,22 @@ def Rows(groups: list[list[Rollout]], now: float, runtime_type: str) -> Table:
 
 class Pager:
     """Which page of overflowing rollout rows is on screen. Auto-advances on a timer until the
-    user takes over with the left/right arrows, after which it stays where they leave it. `count`
-    (the page count, set each render by `_paginate`) gates the arrows: they're inert while a single
-    page fits, so a stray press before rollouts overflow can't switch off auto-advance or offset the
-    starting page once paging begins. The chosen page is clamped to `count` (it can shrink on a
-    resize; it otherwise only grows, as rollouts are never removed)."""
+    user takes over with the left/right arrows, after which it stays where they leave it. Paging
+    opens on the first page: the timer is anchored to `origin` (the clock at the first paged frame),
+    so `int((now - origin) / _PAGE_SECONDS)` is 0 when paging begins and rotates from there, rather
+    than off the raw wall clock (which would open on an arbitrary page). `_paginate` clears `origin`
+    whenever everything fits on one page, so paging re-anchors and re-opens on page 1 if rollouts
+    overflow again after a resize. `count` (the page count, set each render by `_paginate`) gates the
+    arrows: they're inert while a single page fits, so a stray press before rollouts overflow can't
+    switch off auto-advance or offset the starting page once paging begins. The chosen page is
+    clamped to `count` (it can shrink on a resize; it otherwise only grows, as rollouts are never
+    removed)."""
 
     def __init__(self) -> None:
         self.page = 0
         self.manual = False
         self.count = 1
+        self.origin: float | None = None
 
     def on_key(self, key: str) -> None:
         if key in ("left", "right") and self.count > 1:
@@ -517,7 +523,10 @@ class Pager:
         # Track the auto page while it drives, so the first arrow continues from what's on screen
         # rather than jumping back to page 1. Clamp in manual mode (count can shrink on resize).
         if not self.manual:
-            self.page = int(now / _PAGE_SECONDS) % self.count
+            # Anchor the timer to the first paged frame, so paging opens on page 1 then rotates.
+            if self.origin is None:
+                self.origin = now
+            self.page = int((now - self.origin) / _PAGE_SECONDS) % self.count
         else:
             self.page = max(0, min(self.page, self.count - 1))
         return self.page
@@ -530,7 +539,10 @@ def _paginate(
     selecting the one `pager` points at. Returns (this page's groups, 0-based index, page count) —
     a single page when everything already fits."""
     if sum(len(g) for g in groups) <= rows_per_page:
-        pager.count = 1  # everything fits on one page — arrows stay inert
+        # Everything fits: arrows stay inert, and clear the anchor so paging re-opens on page 1 if
+        # rollouts overflow again later (e.g. a terminal resize) rather than off the stale origin.
+        pager.count = 1
+        pager.origin = None
         return groups, 0, 1
     pages: list[list[list[Rollout]]] = []
     current: list[list[Rollout]] = []

verifiers-0.1.15.dev423/verifiers/v1/cli/debug.py

@@ -0,0 +1,327 @@
+"""The debug entrypoint: setup tasks, run one shell action, and persist traces."""
+
+import asyncio
+import contextlib
+import logging
+import random
+import shlex
+import signal
+import sys
+import time
+import traceback
+from collections.abc import Awaitable
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from pydantic_config import cli
+
+import verifiers.v1 as vf
+from verifiers.v1.cli.output import append_trace, save_config
+from verifiers.v1.cli.resolve import (
+    extract_id,
+    references_config_file,
+    with_positional_taskset,
+)
+from verifiers.v1.configs.debug import DebugConfig
+from verifiers.v1.decorators import invoke
+from verifiers.v1.env import resolve_runtime_config
+from verifiers.v1.runtimes import ProgramResult, Runtime, make_runtime
+from verifiers.v1.state import state_cls
+from verifiers.v1.taskset import Taskset
+from verifiers.v1.trace import Error, Trace
+from verifiers.v1.utils.logging import setup_logging
+
+logger = logging.getLogger(__name__)
+
+USAGE = (
+    "usage: uv run debug [<taskset-id>] (--command <cmd> | --script-path <path>) "
+    "[--runtime.type subprocess] [options] [@ file.toml]\n"
+    "       runs setup, then one command or uploaded host script, and saves traces"
+)
+
+
+def _narrow(argv: list[str]) -> type[DebugConfig]:
+    """`DebugConfig` with `taskset` narrowed to the config type of the id on the CLI."""
+    taskset_id = extract_id(argv, "taskset")
+    if not taskset_id:
+        return DebugConfig
+    ftype = vf.taskset_config_type(taskset_id)
+    return type(
+        DebugConfig.__name__,
+        (DebugConfig,),
+        {"__annotations__": {"taskset": ftype}, "taskset": ftype(id=taskset_id)},
+    )
+
+
+def output_path(config: DebugConfig) -> Path:
+    if config.output_dir is not None:
+        return config.output_dir
+    return Path("outputs") / f"{config.taskset.name}--debug" / config.uuid
+
+
+def task_info(task) -> dict[str, Any]:
+    return {"idx": task.idx, "name": task.name, "workdir": task.workdir}
+
+
+def runtime_info(runtime: Runtime) -> dict[str, Any]:
+    return {
+        "type": runtime.type,
+        "name": runtime.name,
+        "descriptor": runtime.descriptor,
+    }
+
+
+def result_info(result: ProgramResult, start: float) -> dict[str, Any]:
+    ok = result.exit_code == 0
+    return {
+        "ok": ok,
+        "reason": "pass" if ok else "nonzero_exit",
+        "exit_code": result.exit_code,
+        "elapsed": round(time.time() - start, 2),
+        "stdout": result.stdout or "",
+        "stderr": result.stderr or "",
+    }
+
+
+def error_info(
+    error: BaseException,
+    start: float,
+    timeout: float | None,
+    stage: str,
+) -> dict[str, Any]:
+    if isinstance(error, asyncio.TimeoutError):
+        message = (
+            f"{stage} timed out after {timeout}s"
+            if timeout is not None
+            else f"{stage} timed out"
+        )
+        reason = "timeout"
+    elif isinstance(error, asyncio.CancelledError):
+        message = f"{stage} cancelled"
+        reason = "cancelled"
+    else:
+        message = str(error)
+        reason = "error"
+    return {
+        "ok": False,
+        "reason": reason,
+        "exit_code": None,
+        "elapsed": round(time.time() - start, 2),
+        "error": message,
+        "error_type": type(error).__name__,
+        "stdout": "",
+        "stderr": "",
+    }
+
+
+def capture_trace_error(trace: Trace, error: BaseException) -> None:
+    if isinstance(error, Exception):
+        trace.capture_error(error)
+        return
+    trace.errors.append(
+        Error(
+            type=type(error).__name__,
+            message=str(error),
+            traceback=traceback.format_exc(),
+        )
+    )
+    trace.stop("error")
+
+
+def record_debug_error(
+    trace: Trace,
+    debug: dict[str, Any],
+    runtime: Runtime,
+    error: BaseException,
+    setup_timeout: float | None,
+    action_timeout: float | None,
+) -> None:
+    if not trace.timing.setup.end:
+        trace.timing.setup.end = time.time()
+    if trace.timing.generation.start and not trace.timing.generation.end:
+        trace.timing.generation.end = time.time()
+    in_action = bool(trace.timing.generation.start)
+    stage = "debug action" if in_action else "setup"
+    timeout = action_timeout if in_action else setup_timeout
+    error_start = (
+        trace.timing.generation.start if in_action else trace.timing.setup.start
+    )
+    debug.update(error_info(error, error_start, timeout, stage))
+    debug.setdefault("runtime", runtime_info(runtime))
+    capture_trace_error(trace, error)
+
+
+def record_action_failure(trace: Trace, debug: dict[str, Any]) -> None:
+    message = debug.get("error") or f"debug action failed: {debug['reason']}"
+    trace.errors.append(
+        Error(
+            type=debug.get("error_type") or "DebugActionError",
+            message=str(message),
+            traceback=None,
+        )
+    )
+
+
+async def run_timed(
+    action: Awaitable[ProgramResult], config: DebugConfig
+) -> dict[str, Any]:
+    """Run the whole debug action under one `--timeout.total` budget."""
+    start = time.time()
+    try:
+        result = await asyncio.wait_for(action, config.timeout.total)
+    except Exception as e:
+        return error_info(e, start, config.timeout.total, "debug action")
+    return result_info(result, start)
+
+
+async def run_action(runtime: Runtime, config: DebugConfig) -> dict[str, Any]:
+    if config.command is not None:
+        return {
+            "action": "command",
+            "command": config.command,
+            **(await run_timed(runtime.run(["sh", "-lc", config.command], {}), config)),
+        }
+    assert config.script_path is not None
+    # /tmp keyed by the (run-unique) runtime name: outside the task workdir so the script
+    # doesn't show up in the repo state being inspected, and never shared between runs on
+    # the subprocess runtime, where an absolute path is a host path.
+    remote_path = f"/tmp/{runtime.name}-script.sh"
+    quoted = shlex.quote(remote_path)
+    command = f"chmod +x {quoted} && {quoted}"
+
+    async def upload_and_run() -> ProgramResult:
+        await runtime.write(remote_path, config.script_path.read_bytes())
+        return await runtime.run(["sh", "-lc", command], {})
+
+    return {
+        "action": "script",
+        "script_path": str(config.script_path),
+        "remote_script_path": remote_path,
+        "command": command,
+        **(await run_timed(upload_and_run(), config)),
+    }
+
+
+async def debug_task(taskset: Taskset, task, config: DebugConfig) -> tuple[Trace, bool]:
+    trace = Trace(task=task, state=state_cls(type(taskset))())
+    debug = {
+        "task": task_info(task),
+        "action": "command" if config.command is not None else "script",
+        "ok": False,
+        "reason": "not_run",
+    }
+    cancelled = False
+    runtime = make_runtime(
+        resolve_runtime_config(config.runtime, task),
+        name=f"debug-{task.idx}-{uuid4().hex[:8]}",
+    )
+    setup_timeout = (
+        config.timeout.setup if config.timeout.setup is not None else task.timeout.setup
+    )
+    try:
+        trace.timing.setup.start = time.time()
+        await runtime.start()
+        debug["runtime"] = runtime_info(runtime)
+        await asyncio.wait_for(
+            invoke(taskset.setup, {"task": task, "trace": trace, "runtime": runtime}),
+            setup_timeout,
+        )
+        trace.timing.setup.end = time.time()
+
+        trace.timing.generation.start = time.time()
+        debug.update(await run_action(runtime, config))
+        trace.timing.generation.end = time.time()
+        if not debug.get("ok"):
+            record_action_failure(trace, debug)
+        trace.stop(str(debug["reason"]))
+    except asyncio.CancelledError as e:
+        cancelled = True
+        record_debug_error(
+            trace, debug, runtime, e, setup_timeout, config.timeout.total
+        )
+    except Exception as e:
+        record_debug_error(
+            trace, debug, runtime, e, setup_timeout, config.timeout.total
+        )
+    finally:
+        trace.info["debug"] = debug
+        try:
+            await runtime.stop()
+        except asyncio.CancelledError:
+            # a task cancellation delivered mid-stop would abort before the caller can
+            # persist the trace — absorb it here; the caller re-raises after appending
+            cancelled = True
+        except Exception:
+            logger.warning("runtime teardown failed (task %s)", task.idx, exc_info=True)
+    return trace, cancelled
+
+
+async def run_debug(config: DebugConfig) -> list[Trace]:
+    taskset = vf.load_taskset(config.taskset)
+    tasks = taskset.load_tasks()
+    if config.shuffle:
+        random.Random(0).shuffle(tasks)
+    if config.num_tasks is not None:
+        tasks = tasks[: config.num_tasks]
+    if isinstance(config.runtime, vf.SubprocessConfig) and (
+        taskset.NEEDS_CONTAINER or any(t.image for t in tasks)
+    ):
+        raise SystemExit(
+            "taskset needs a container runtime to debug - pass --runtime.type docker (or prime)"
+        )
+
+    out = output_path(config)
+    save_config(config, out)
+    logger.info(
+        "debugging %d task(s) from %s on the %s runtime",
+        len(tasks),
+        config.name,
+        config.runtime.type,
+    )
+    logger.info("results: %s", out)
+
+    sem = asyncio.Semaphore(config.max_concurrent) if config.max_concurrent else None
+    write_lock = asyncio.Lock()
+
+    async def one(task) -> Trace:
+        async with sem or contextlib.nullcontext():
+            trace, cancelled = await debug_task(taskset, task, config)
+        await append_trace(out, trace, write_lock)
+        info = trace.info["debug"]
+        detail = f" - {info['error']}" if info.get("error") else ""
+        logger.info(
+            "idx=%s ok=%s reason=%s%s",
+            task.idx,
+            info["ok"],
+            info["reason"],
+            detail,
+        )
+        if cancelled:
+            raise asyncio.CancelledError
+        return trace
+
+    return await asyncio.gather(*(one(task) for task in tasks))
+
+
+def main(argv: list[str] | None = None) -> None:
+    argv = with_positional_taskset(list(sys.argv[1:]) if argv is None else list(argv))
+
+    if not argv or any(arg in ("-h", "--help") for arg in argv):
+        print(USAGE)
+        sys.argv = [sys.argv[0], "--help"]
+        cli(_narrow(argv))
+        return
+    if not extract_id(argv, "taskset") and not references_config_file(argv):
+        raise SystemExit(USAGE)
+
+    config_type = _narrow(argv)
+    sys.argv = [sys.argv[0], *argv]
+    config = cli(config_type)
+    setup_logging("DEBUG" if config.verbose else "INFO")
+    signal.signal(signal.SIGTERM, lambda *_: (_ for _ in ()).throw(KeyboardInterrupt()))
+    asyncio.run(run_debug(config))
+
+
+if __name__ == "__main__":
+    main()

{verifiers-0.1.15.dev421 → verifiers-0.1.15.dev423}/verifiers/v1/cli/output.py

@@ -15,10 +15,11 @@ import asyncio
 from pathlib import Path
 
 import tomli_w
-from pydantic import TypeAdapter
+from pydantic import BaseModel, TypeAdapter
 
 from verifiers.v1.configs.eval import EvalConfig
 from verifiers.v1.trace import Trace
+from verifiers.v1.utils.aio import run_shielded
 
 
 def output_path(config: EvalConfig) -> Path:
@@ -30,7 +31,7 @@ def output_path(config: EvalConfig) -> Path:
     return Path("outputs") / name / config.uuid
 
 
-def write_config(config: EvalConfig, results_dir: Path) -> Path:
+def write_config(config: BaseModel, results_dir: Path) -> Path:
     """Write the run's resolved `config.toml` (re-readable via `@ config.toml`); return its
     path. mode="json" makes values TOML-friendly (Path -> str, etc.); exclude_none drops the
     nulls TOML can't represent."""
@@ -41,7 +42,7 @@ def write_config(config: EvalConfig, results_dir: Path) -> Path:
     return config_path
 
 
-def save_config(config: EvalConfig, results_dir: Path) -> None:
+def save_config(config: BaseModel, results_dir: Path) -> None:
     """Set up the run's output dir: write `config.toml` and start a fresh (empty)
     `results.jsonl`. Call once up front, before traces start landing."""
     write_config(config, results_dir)
@@ -65,10 +66,6 @@ async def append_trace(results_dir: Path, trace: Trace, lock: asyncio.Lock) -> N
         async with lock:
             await asyncio.to_thread(write_trace, results_dir, trace)
 
-    # Shield lock acquisition and the worker so finalized traces survive cancellation.
-    persist_task = asyncio.create_task(persist())
-    try:
-        await asyncio.shield(persist_task)
-    except asyncio.CancelledError:
-        await persist_task
-        raise
+    # Run lock acquisition and the worker to completion even under cancellation, so
+    # finalized traces are never lost mid-write (`run_shielded` re-raises the cancellation).
+    await run_shielded(persist())