feat(v5-g09): cancel/abort Train via opts.abort_token + _ABORT_TOKENS

Closes V5-G09 / cppmega-mlx-ajm. Backend stage_train checks
opts.abort_token between train steps; if in module-global _ABORT_TOKENS
set, returns immediately with extras.aborted=True + partial losses.
request_abort(token) / clear_abort(token) helpers for callers.

3 new pytest. UI Cancel button + WS abort handler deferred — backend
contract first, frontend wiring follow-up.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -773,7 +773,25 @@ def _count(tree: Any) -> int:
         probe_key: str | None = None
         probe_before: mx.array | None = None
 
+        # G09: check abort flag set via opts.abort or _ABORT_TOKENS set
+        abort_token = opts.get("abort_token")
         for step in range(n_steps):
+            if abort_token is not None and abort_token in _ABORT_TOKENS:
+                # Stop early; return partial extras with cancellation flag.
+                return StageResult(
+                    name="train", status="ok",
+                    elapsed_ms=(time.perf_counter() - t0) * 1000.0,
+                    extras={
+                        "losses": [round(l, 4) for l in losses],
+                        "lr_trajectory": [round(l, 6) for l in lr_trajectory],
+                        "weight_delta_norm": 0.0,
+                        "num_steps": step,
+                        "schedule_kind": schedule_kind_label,
+                        "optimizer_kind": optimizer_kind,
+                        "aborted": True,
+                        "abort_token": abort_token,
+                    },
+                )
             # If a schedule callable exists, override optimizer's
             # learning_rate per step. MLX optimizers accept a fresh
             # scalar via the public learning_rate attribute.
@@ -1014,6 +1032,20 @@ def _tokenize_parquet_text(
 # across sequential Train clicks in the same backend session.
 _RUN_CACHE: dict[str, Any] = {}
 
+# G09: in-process set of abort tokens. Caller sets opts.abort_token to
+# some unique string; another caller (e.g. WS handler) inserts the same
+# token into this set to signal cancellation between train steps.
+_ABORT_TOKENS: set[str] = set()
+
+
+def request_abort(token: str) -> None:
+    """G09: signal stage_train to abort the run identified by token."""
+    _ABORT_TOKENS.add(token)
+
+
+def clear_abort(token: str) -> None:
+    _ABORT_TOKENS.discard(token)
+
 
 _REWRITER_FACTORIES: dict[str, Callable[..., Any]] = {}
 

tests/v5/test_stage_train_cancel.py

@@ -0,0 +1,56 @@
+"""G09: cancel/abort Train via abort_token + _ABORT_TOKENS set."""
+
+from __future__ import annotations
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+from cppmega_v4.runner.stages import request_abort, clear_abort
+
+
+def _spec() -> VerifyParams:
+    return VerifyParams.model_validate({
+        "graph": {"nodes": [
+            {"id": "attn", "kind": "attention", "params": {}},
+            {"id": "mlp", "kind": "mlp",
+             "params": {"intermediate_size": 64, "activation": "swiglu"}},
+        ], "edges": [{"src": "attn", "dst": "mlp"}]},
+        "dim_env": {"B": 1, "S": 8, "H": 32, "nh": 2, "nkv": 1, "head_dim": 16},
+        "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 _run(opts: dict) -> dict:
+    report = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": opts},
+    }))
+    train = next(s for s in report.stages if s.name == "train")
+    return train.extras
+
+
+def test_abort_token_set_before_run_cancels_immediately():
+    token = "abort-test-1"
+    request_abort(token)
+    try:
+        extras = _run({"num_steps": 8, "abort_token": token})
+        assert extras.get("aborted") is True
+        assert extras["abort_token"] == token
+        # Aborted at step 0 → losses empty
+        assert extras["num_steps"] == 0
+    finally:
+        clear_abort(token)
+
+
+def test_no_abort_token_runs_to_completion():
+    extras = _run({"num_steps": 2})
+    assert extras.get("aborted") is not True
+    assert extras["num_steps"] == 2
+
+
+def test_abort_token_not_in_set_runs_to_completion():
+    extras = _run({"num_steps": 2, "abort_token": "never-set"})
+    assert extras.get("aborted") is not True
+    assert extras["num_steps"] == 2