feat: add train cancel rpc and ui

Author: apstenku123

cppmega_v4/jsonrpc/__init__.py

@@ -28,6 +28,8 @@
     JsonRpcError,
     JsonRpcRequest,
     JsonRpcResponse,
+    PipelineAbortParams,
+    PipelineAbortResult,
     ProbeRunParams,
     ProbeRunResult,
     SuggestAdaptersParams,
@@ -53,6 +55,8 @@
     "JsonRpcResponse",
     "LRUCache",
     "METHOD_REGISTRY",
+    "PipelineAbortParams",
+    "PipelineAbortResult",
     "ProbeRunParams",
     "ProbeRunResult",
     "SCHEMA_VERSION",

cppmega_v4/jsonrpc/dispatcher.py

@@ -50,6 +50,8 @@
     JsonRpcError,
     JsonRpcRequest,
     JsonRpcResponse,
+    PipelineAbortParams,
+    PipelineAbortResult,
     PipelineRunParams,
     PipelineRunResult,
     ProbeRunParams,
@@ -90,6 +92,10 @@
         PipelineRunParams,
         lambda p, c: _pipeline_run(p),
     ),
+    "pipeline.abort": (
+        PipelineAbortParams,
+        lambda p, c: _pipeline_abort(p),
+    ),
     "tokenizer.encode_visualize": (
         EncodeVisualizeParams,
         lambda p, c: encode_visualize(p, cache=c),
@@ -134,6 +140,12 @@ def _pipeline_run(params: PipelineRunParams) -> PipelineRunResult:
     return PipelineRunResult.model_validate(report.to_dict())
 
 
+def _pipeline_abort(params: PipelineAbortParams) -> PipelineAbortResult:
+    from cppmega_v4.runner.stages import request_abort
+    request_abort(params.run_id)
+    return PipelineAbortResult(run_id=params.run_id)
+
+
 def dispatch(
     payload: Mapping[str, Any] | JsonRpcRequest,
     *,

cppmega_v4/jsonrpc/schema.py

@@ -663,13 +663,30 @@ class PipelineRunParams(BaseModel):
     pipeline: PipelinePayload
 
 
+class PipelineAbortParams(BaseModel):
+    """Request cancellation of an in-flight pipeline train stage."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    run_id: str = Field(min_length=1)
+
+
+class PipelineAbortResult(BaseModel):
+    """Acknowledgement that the train abort token has been set."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    status: Literal["abort_requested"] = "abort_requested"
+    run_id: str
+
+
 class StageResult(BaseModel):
     """One stage execution result."""
 
     model_config = ConfigDict(extra="allow")
 
     name: str
-    status: Literal["ok", "skipped", "fail"]
+    status: Literal["ok", "skipped", "fail", "cancelled"]
     elapsed_ms: float
     warnings: int = 0
     errors: int = 0
@@ -682,7 +699,7 @@ class PipelineRunResult(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
     stages: list[StageResult]
-    overall_status: Literal["ok", "fail"]
+    overall_status: Literal["ok", "fail", "cancelled"]
     total_elapsed_ms: float
 
 
@@ -707,6 +724,7 @@ class PipelineRunResult(BaseModel):
     "backend.status",
     "probe.run",
     "pipeline.run",
+    "pipeline.abort",
 })
 
 
@@ -717,6 +735,7 @@ class PipelineRunResult(BaseModel):
     "build_preset_specs",
     "probe.run",
     "pipeline.run",
+    "pipeline.abort",
     "backend.status",
     "tokenizer.encode_visualize",
     "tokenizer.list_presets",

cppmega_v4/jsonrpc/server.py

@@ -73,7 +73,7 @@ def cache_clear() -> dict[str, str]:
 
     @app.post("/rpc")
     async def rpc(payload: dict) -> dict:
-        response = dispatch(payload, cache=cache)
+        response = await _dispatch(payload, cache)
         return response.model_dump(mode="json", exclude_none=True)
 
     @app.websocket("/ws")
@@ -92,7 +92,7 @@ async def ws(socket: WebSocket) -> None:
                                   "data": {"detail": str(exc)}},
                     })
                     continue
-                response = dispatch(payload, cache=cache)
+                response = await _dispatch(payload, cache)
                 await socket.send_json(
                     response.model_dump(mode="json", exclude_none=True)
                 )
@@ -108,6 +108,12 @@ async def ws(socket: WebSocket) -> None:
     return app
 
 
+async def _dispatch(payload: dict, cache: LRUCache):
+    if payload.get("method") == "pipeline.run":
+        return await asyncio.to_thread(dispatch, payload, cache=cache)
+    return dispatch(payload, cache=cache)
+
+
 async def _heartbeat(socket: WebSocket) -> None:
     while True:
         await asyncio.sleep(_HEARTBEAT_INTERVAL_S)

cppmega_v4/runner/pipeline.py

@@ -90,7 +90,10 @@ def run_pipeline(spec: VerifyParams, pipeline: Pipeline) -> PipelineReport:
         stage = STAGE_REGISTRY[name]
         result = stage(ctx)
         results.append(result)
-        if result.status == "fail" and not pipeline.continue_on_failure:
+        if (
+            result.status in {"fail", "cancelled"}
+            and not pipeline.continue_on_failure
+        ):
             # Mark every subsequent stage as skipped for visibility.
             for remaining in pipeline.stages[len(results):]:
                 results.append(StageResult(
@@ -100,6 +103,8 @@ def run_pipeline(spec: VerifyParams, pipeline: Pipeline) -> PipelineReport:
     overall = "ok"
     if any(r.status == "fail" for r in results):
         overall = "fail"
+    elif any(r.status == "cancelled" for r in results):
+        overall = "cancelled"
     elapsed = (time.perf_counter() - t0) * 1000.0
     return PipelineReport(
         stages=results, overall_status=overall, total_elapsed_ms=elapsed,

cppmega_v4/runner/stages.py

@@ -42,7 +42,7 @@
 from cppmega_v4.jsonrpc.schema import VerifyParams
 
 
-StageStatus = Literal["ok", "skipped", "fail"]
+StageStatus = Literal["ok", "skipped", "fail", "cancelled"]
 
 
 @dataclass
@@ -119,6 +119,36 @@ def _fail(name: str, t0: float, exc: BaseException) -> StageResult:
     )
 
 
+def _cancelled_train_result(
+    t0: float,
+    *,
+    abort_token: str,
+    step: int,
+    losses: list[float],
+    lr_trajectory: list[float],
+    schedule_kind_label: str,
+    optimizer_kind: str,
+    weight_delta_norm: float = 0.0,
+) -> StageResult:
+    clear_abort(abort_token)
+    return StageResult(
+        name="train", status="cancelled",
+        elapsed_ms=(time.perf_counter() - t0) * 1000.0,
+        extras={
+            "losses": [round(loss_item, 4) for loss_item in losses],
+            "lr_trajectory": [
+                round(lr_item, 6) for lr_item in lr_trajectory
+            ],
+            "weight_delta_norm": round(weight_delta_norm, 6),
+            "num_steps": step,
+            "schedule_kind": schedule_kind_label,
+            "optimizer_kind": optimizer_kind,
+            "aborted": True,
+            "abort_token": abort_token,
+        },
+    )
+
+
 def stage_parse(ctx: StageContext) -> StageResult:
     """Materialise loss/optim/graph from the wire-form spec."""
     t0 = time.perf_counter()
@@ -726,7 +756,14 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
                         token_count = len(tokens)
                 except Exception:
                     pass
+        train_token_embedding = None
+        train_input_source = "random"
+        if data_source != "synthetic":
+            train_token_embedding = nn.Embedding(vocab_size, hidden)
+            train_input_source = "token_embedding"
         all_modules = nn.Sequential(*modules, *lm_heads)
+        if train_token_embedding is not None:
+            all_modules.train_token_embedding = train_token_embedding
         if side_channel_token_embedding is not None:
             all_modules.side_channel_token_embedding = side_channel_token_embedding
         opt, optimizer_kind = _build_optimizer(spec_optim, lr)
@@ -939,21 +976,14 @@ def _count(tree: Any) -> int:
         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(loss_item, 4)
-                                   for loss_item in losses],
-                        "lr_trajectory": [round(lr_item, 6)
-                                          for lr_item 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,
-                    },
+                return _cancelled_train_result(
+                    t0,
+                    abort_token=str(abort_token),
+                    step=step,
+                    losses=losses,
+                    lr_trajectory=lr_trajectory,
+                    schedule_kind_label=schedule_kind_label,
+                    optimizer_kind=optimizer_kind,
                 )
             # If a schedule callable exists, override optimizer's
             # learning_rate per step. MLX optimizers accept a fresh
@@ -965,8 +995,11 @@ def _count(tree: Any) -> int:
             else:
                 lr_trajectory.append(lr)
 
-            emb = mx.random.normal(shape=(batch, seq, hidden), key=rng_key)
-            rng_key, _ = mx.random.split(rng_key)
+            if train_token_embedding is None:
+                emb = mx.random.normal(shape=(batch, seq, hidden), key=rng_key)
+                rng_key, _ = mx.random.split(rng_key)
+            else:
+                emb = train_token_embedding(targets)
             loss, grads = loss_and_grad(all_modules, emb, targets)
             mx.eval(loss, grads)
             if probe_key is None:
@@ -1002,6 +1035,25 @@ def _count(tree: Any) -> int:
             opt.update(all_modules, grads)
             mx.eval(all_modules.parameters(), opt.state)
             losses.append(float(loss.item()))
+            if abort_token is not None and abort_token in _ABORT_TOKENS:
+                partial_delta = 0.0
+                if probe_key is not None and probe_before is not None:
+                    after_flat = dict(
+                        nn.utils.tree_flatten(all_modules.parameters()))
+                    partial_delta = float(
+                        mx.linalg.norm(
+                            after_flat[probe_key] - probe_before).item()
+                    )
+                return _cancelled_train_result(
+                    t0,
+                    abort_token=str(abort_token),
+                    step=step + 1,
+                    losses=losses,
+                    lr_trajectory=lr_trajectory,
+                    schedule_kind_label=schedule_kind_label,
+                    optimizer_kind=optimizer_kind,
+                    weight_delta_norm=partial_delta,
+                )
 
         try:
             if hasattr(mx, "metal"):
@@ -1110,6 +1162,7 @@ def _count(tree: Any) -> int:
                 "schedule_kind": schedule_kind_label,
                 "optimizer_kind": optimizer_kind,
                 "data_source": data_source,
+                "train_input_source": train_input_source,
                 "token_count": token_count,
                 "tokenizer_used": tokenizer_used,
                 "loss_kind": (

tests/v4/test_jsonrpc_dispatcher.py

@@ -8,6 +8,7 @@
     LRUCache,
     dispatch,
 )
+from cppmega_v4.runner.stages import clear_abort
 
 
 _GRAPH = {
@@ -105,3 +106,41 @@ def test_dispatch_pipeline_run_round_trip():
     assert resp.error is None
     assert resp.result["overall_status"] == "ok"
     assert len(resp.result["stages"]) == 2
+
+
+def test_dispatch_pipeline_abort_requests_stage_abort():
+    run_id = "dispatcher-abort-1"
+    clear_abort(run_id)
+    try:
+        resp = dispatch({
+            "jsonrpc": "2.0", "id": "abort", "method": "pipeline.abort",
+            "params": {"run_id": run_id},
+        })
+        assert resp.error is None
+        assert resp.result == {
+            "status": "abort_requested",
+            "run_id": run_id,
+        }
+
+        envelope = {
+            "jsonrpc": "2.0", "id": "train", "method": "pipeline.run",
+            "params": {
+                "spec": {"graph": _GRAPH, "dim_env": _DIM_ENV,
+                         "loss": _LOSS, "optim": _OPTIM},
+                "pipeline": {
+                    "stages": ["parse", "verify_build_spec", "build_model",
+                               "train"],
+                    "stage_options": {
+                        "train": {"num_steps": 8, "abort_token": run_id},
+                    },
+                },
+            },
+        }
+        run_resp = dispatch(envelope)
+        assert run_resp.error is None
+        assert run_resp.result["overall_status"] == "cancelled"
+        train = run_resp.result["stages"][-1]
+        assert train["status"] == "cancelled"
+        assert train["aborted"] is True
+    finally:
+        clear_abort(run_id)