feat(v7-h10): abort_token honoured by verify_build_spec + dry_forward

Closes item 19 of the UI<->API wiring honesty list. Previously only
stage_train respected _ABORT_TOKENS, so a UI cancel during a long
verify on a big graph had zero effect until train started.

- run_pipeline picks up any abort_token declared in stage_options
  (typically train.abort_token) and gates every stage entry.
- New _stage_abort_token helper resolves the same token from inside
  individual stages.
- stage_verify_build_spec and stage_dry_forward gate at entry and
  return status='cancelled' instead of running to completion.

Honest scope: cancellation is at stage boundaries — MLX itself can't
be interrupted mid-call, so verify_build_spec on a partially-finished
graph still completes the current stage's work before yielding. UI now
has at-most-one-stage of latency per cancel click.

bd: cppmega-mlx-lmor
4/4 pytest + 14/14 lifecycle regression.

Author: apstenku123

cppmega_v4/runner/pipeline.py

@@ -82,11 +82,41 @@ def run_pipeline(spec: VerifyParams, pipeline: Pipeline) -> PipelineReport:
     """Run ``pipeline`` over ``spec`` and return the report.
 
     Stops on the first failed stage unless ``continue_on_failure`` is set.
+
+    V7-H10: between stages, check if the pipeline-wide abort_token (any
+    stage that declares one in its stage_options) has been signalled via
+    pipeline.abort RPC. If so, mark all remaining stages as cancelled —
+    this makes verify_build_spec / dry_forward etc. cooperatively
+    cancellable at stage boundaries even though they can't be
+    interrupted mid-call.
     """
+    from cppmega_v4.runner.stages import _ABORT_TOKENS, clear_abort
+
+    def _pipeline_abort_token() -> str | None:
+        """Pick the first abort_token declared in any stage's options."""
+        for _name, _opts in (pipeline.stage_options or {}).items():
+            tok = (_opts or {}).get("abort_token")
+            if tok:
+                return str(tok)
+        return None
+
     t0 = time.perf_counter()
     ctx = StageContext(spec=spec, options=pipeline.stage_options)
     results: list[StageResult] = []
+    abort_token = _pipeline_abort_token()
     for name in pipeline.stages:
+        # V7-H10: cancellation gate before each stage.
+        if abort_token is not None and abort_token in _ABORT_TOKENS:
+            results.append(StageResult(
+                name=name, status="cancelled", elapsed_ms=0.0,
+                error={"type": "Aborted", "abort_token": abort_token},
+            ))
+            for remaining in pipeline.stages[len(results):]:
+                results.append(StageResult(
+                    name=remaining, status="skipped", elapsed_ms=0.0,
+                ))
+            clear_abort(abort_token)
+            break
         stage = STAGE_REGISTRY[name]
         result = stage(ctx)
         results.append(result)

cppmega_v4/runner/stages.py

@@ -159,6 +159,22 @@ def _cancelled_train_result(
     )
 
 
+def _stage_abort_token(ctx: StageContext, stage_name: str) -> str | None:
+    """V7-H10: resolve the pipeline-wide abort token for a stage entry.
+
+    Prefer the per-stage opts override, fall back to any abort_token
+    declared by any other stage in the pipeline (typically train) so a
+    single UI cancel propagates to verify / dry_forward / etc."""
+    own = ctx.opts(stage_name).get("abort_token")
+    if own:
+        return str(own)
+    for _opts in (ctx.options or {}).values():
+        tok = (_opts or {}).get("abort_token")
+        if tok:
+            return str(tok)
+    return None
+
+
 def stage_parse(ctx: StageContext) -> StageResult:
     """Materialise loss/optim/graph from the wire-form spec."""
     t0 = time.perf_counter()
@@ -184,6 +200,17 @@ def stage_parse(ctx: StageContext) -> StageResult:
 
 def stage_verify_build_spec(ctx: StageContext) -> StageResult:
     t0 = time.perf_counter()
+    # V7-H10: respect pipeline-level abort_token at stage entry so a
+    # cancel issued during a long graph verify still produces a
+    # cancelled result instead of forcing the user to wait it out.
+    _tok = _stage_abort_token(ctx, "verify_build_spec")
+    if _tok and _tok in _ABORT_TOKENS:
+        clear_abort(_tok)
+        return StageResult(
+            name="verify_build_spec", status="cancelled",
+            elapsed_ms=(time.perf_counter() - t0) * 1000.0,
+            error={"type": "Aborted", "abort_token": _tok},
+        )
     try:
         if ctx.build_spec is None:
             return _fail("verify_build_spec", t0,
@@ -331,6 +358,15 @@ def stage_build_model(ctx: StageContext) -> StageResult:
 
 def stage_dry_forward(ctx: StageContext) -> StageResult:
     t0 = time.perf_counter()
+    # V7-H10: same cancel gate as stage_verify_build_spec.
+    _tok = _stage_abort_token(ctx, "dry_forward")
+    if _tok and _tok in _ABORT_TOKENS:
+        clear_abort(_tok)
+        return StageResult(
+            name="dry_forward", status="cancelled",
+            elapsed_ms=(time.perf_counter() - t0) * 1000.0,
+            error={"type": "Aborted", "abort_token": _tok},
+        )
     try:
         opts = ctx.opts("dry_forward")
         seq = int(opts.get("S", 8))

tests/v4/test_pipeline_abort_verify_dry.py

@@ -0,0 +1,116 @@
+"""V7-H10: abort_token honoured by verify_build_spec + dry_forward.
+
+Previously only stage_train respected _ABORT_TOKENS. UI clicking cancel
+during a long verify on a big graph had no effect — the user had to
+wait it out. This pins:
+
+  - stage_verify_build_spec returns 'cancelled' at entry if aborted.
+  - stage_dry_forward returns 'cancelled' at entry if aborted.
+  - run_pipeline driver, between stages, halts and marks remaining
+    stages 'skipped' if abort_token is set.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from cppmega_v4.jsonrpc.dispatcher import _pipeline_abort
+from cppmega_v4.jsonrpc.schema import PipelineAbortParams, VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+from cppmega_v4.runner.stages import _ABORT_TOKENS, request_abort
+from cppmega_v4.runtime import run_registry as rr
+
+
+@pytest.fixture(autouse=True)
+def _reset():
+    _ABORT_TOKENS.clear()
+    rr.reset()
+    yield
+    _ABORT_TOKENS.clear()
+    rr.reset()
+
+
+def _spec() -> VerifyParams:
+    return VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "attn", "kind": "attention",
+                 "params": {"num_heads": 4, "head_dim": 64}},
+                {"id": "mlp", "kind": "mlp", "params": {}},
+            ],
+            "edges": [{"src": "attn", "dst": "mlp"}],
+        },
+        "dim_env": {"B": 1, "S": 8, "H": 128,
+                    "nh": 2, "nkv": 1, "head_dim": 64},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["mlp"]},
+        "optim": {"kind": "adamw",
+                  "groups": [{"matcher": "all", "lr": 1e-3,
+                              "weight_decay": 0.01,
+                              "betas": [0.9, 0.95]}]},
+    })
+
+
+def test_v7_h10_verify_build_spec_stage_function_cancels_at_entry():
+    """Direct stage call: the entry-gate inside stage_verify_build_spec
+    fires when its abort_token is in _ABORT_TOKENS, independent of the
+    pipeline driver."""
+    from cppmega_v4.runner.stages import (
+        StageContext, stage_verify_build_spec, stage_parse,
+    )
+    ctx = StageContext(spec=_spec(),
+                       options={"verify_build_spec":
+                                {"abort_token": "ABRT-V"}})
+    parse_result = stage_parse(ctx)
+    assert parse_result.status == "ok"
+
+    request_abort("ABRT-V")
+    v = stage_verify_build_spec(ctx)
+    assert v.status == "cancelled"
+    assert v.error == {"type": "Aborted", "abort_token": "ABRT-V"}
+
+
+def test_v7_h10_dry_forward_stage_function_cancels_at_entry():
+    """Direct stage call: stage_dry_forward's entry-gate fires."""
+    from cppmega_v4.runner.stages import (
+        StageContext, stage_dry_forward, stage_parse, stage_build_model,
+        stage_verify_build_spec, stage_resolve_shapes,
+        stage_estimate_memory,
+    )
+    ctx = StageContext(spec=_spec(),
+                       options={"dry_forward":
+                                {"abort_token": "ABRT-DF"}})
+    assert stage_parse(ctx).status == "ok"
+    assert stage_verify_build_spec(ctx).status == "ok"
+    assert stage_resolve_shapes(ctx).status == "ok"
+    assert stage_estimate_memory(ctx).status == "ok"
+    assert stage_build_model(ctx).status == "ok"
+
+    request_abort("ABRT-DF")
+    df = stage_dry_forward(ctx)
+    assert df.status == "cancelled"
+    assert df.error == {"type": "Aborted", "abort_token": "ABRT-DF"}
+
+
+def test_v7_h10_pipeline_driver_between_stages_halts_on_abort_rpc():
+    """Run a pipeline whose train stage carries an abort_token; trigger
+    pipeline.abort RPC up-front so the very first between-stage gate
+    fires before parse even runs."""
+    _pipeline_abort(PipelineAbortParams(run_id="ABRT-PD"))
+    rep = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": {"num_steps": 2,
+                                    "abort_token": "ABRT-PD"}},
+    }))
+    parse = next(s for s in rep.stages if s.name == "parse")
+    assert parse.status == "cancelled"
+    assert rep.overall_status == "cancelled"
+
+
+def test_v7_h10_no_abort_token_still_runs_to_completion():
+    rep = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": {"num_steps": 2}},
+    }))
+    assert rep.overall_status == "ok"
+    for s in rep.stages:
+        assert s.status == "ok", f"{s.name}: {s.status} {s.error}"