From 9fe7a70a9754efafe708a0a6040a8fd85b9ebfc8 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 16:59:34 +0800 Subject: [PATCH 01/34] docs: add run store design spec (rev 2, post-Codex-review) --- .../specs/2026-08-15-run-store-design.md | 155 ++++++++++++++++++ 1 file changed, 155 insertions(+) create mode 100644 docs/superpowers/specs/2026-08-15-run-store-design.md diff --git a/docs/superpowers/specs/2026-08-15-run-store-design.md b/docs/superpowers/specs/2026-08-15-run-store-design.md new file mode 100644 index 0000000..63789c3 --- /dev/null +++ b/docs/superpowers/specs/2026-08-15-run-store-design.md @@ -0,0 +1,155 @@ +# Run Store Design (rev 2, post-Codex-review) + +Status: DRAFT rev 2 — revised per Codex review 2026-08-15; pending user approval. + +## Problem + +ComfyResearch persists no training results server-side. `POST /api/train` streams NDJSON and keeps nothing; the browser stashes `lossHistory` / `memoryCheckpoint_b64` back into node `data`, which lands in `workspace.json` or graph-library templates (hence the 21 MB library bloat). Results are lost if the stream drops, there is no run history, no cross-run comparison, and no way for an agent to submit a run and query results later. + +## Goals (v1) + +1. Every training run is persisted server-side with a `run_id`: metadata, hyperparameters, scalar metric histories, terminal status. +2. Agent-first async API: submit returns `run_id` immediately; execution is detached from any HTTP stream; poll/list/filter/bulk-delete. +3. A runs panel in the canvas UI: run list with grouping, curve-overlay comparison. +4. Supports agent mass-rollouts: thousands of runs, concurrent creation, O(1) write path, grouping, batch GC, crash recovery. + +## Non-goals (v1) + +- Checkpoint/blob externalization (phase 2; streams pass `checkpoint_b64` through unchanged; the store never persists it). +- Backfill of existing embedded results. +- wandb/MLflow export. Multi-user/auth. +- CRL runs (`post_train`'s CRL generator variant): explicitly excluded from capture in v1; classic trainer runs only. + +## Architecture + +**Files are the source of truth; SQLite is a rebuildable read index. Persistence happens exclusively in the API-host server process.** + +### Capture point: engine layer, not the HTTP generator + +`RunWriter` hooks where trainer events are actually produced/consumed — wrapping `iter_trainer_events` / `iter_trainer_events_from_context` consumption sites: + +- Single run (`/api/train` local path): the stream generator drives a `RunWriter` alongside yielding. +- **Sweep / coordinate descent**: `train_sweep.py` and `train_coordinate_descent.py` consume inner-run events internally and only emit wrapper events upward; each **inner run gets its own `RunWriter`** (own `run_id`, `group_id` = sweep session id) at that consumption site. An API-level tee cannot see these — this is why capture lives in the engine. +- **Remote runs**: the API-host side consumes the SSH event stream (`remote/ssh.py` iterator); `RunWriter` attaches there. The remote CLI process never writes the store (it runs on the remote host's checkout — writing there would create a split-brain store). +- **Async submit** (below): the worker drives `RunWriter` directly. + +### Execution detached from the response stream (agent path) + +Starlette cancels streaming work on client disconnect, and the remote iterator kills the SSH process in `finally` — so a tee alone cannot make runs durable. Therefore: + +- **`POST /api/runs` (async submit)**: validates via the existing prepare pipeline (400 on bad graph), registers the run, enqueues onto a server-owned worker pool (thread pool, configurable cap, default = 2; excess runs queue as `queued`), returns **202 + `run_id`** synchronously. Optional `Idempotency-Key` header: same key returns the existing run instead of double-submitting. Execution and persistence are independent of any client connection. +- **`POST /api/train` (existing, browser path)**: unchanged semantics — stream drives execution, disconnect aborts (matches "closing the tab cancels"). A `RunWriter` tees what is emitted, and abort-on-disconnect finalizes the run as `aborted` with partial metrics. First event added: `{"type": "run_registered", "run_id"}`. + +Agents use submit+poll; the browser keeps its live stream. Both paths converge on the same engine-level capture. + +### On-disk layout (source of truth) + +``` +data/runs/ + {run_id}/ # "run-" + uuid4().hex[:12] + run.json # RunRecord (atomic write, rewritten on status change) + metrics.ndjson # append-only per-step scalar DELTA rows + results.json # terminal snapshot: full histories (atomic write, once) +``` + +### Metrics scheme (delta-based; cumulative-history trap addressed) + +The trainer emits **cumulative** histories on every `metrics` event; appending raw payloads would be O(n²) in storage. Instead `RunWriter` keeps the last-seen length per series and appends **only new entries** as scalar rows to `metrics.ndjson`. It also retains the latest cumulative snapshot in memory so that: + +- baseline step-0 values recorded before the first `metrics` event are captured from the first emission's full history (deltas start from index 0); +- on `aborted` / `error` — which emit no histories — `results.json` is still written from the last-seen snapshot; +- on `complete`, `results.json` is written from the terminal event's authoritative histories. + +**Authority rule**: if `results.json` exists it is the authoritative series; `metrics.ndjson` is authoritative only for runs that died without one (crash). `GET /api/runs/{id}/metrics` serves exactly one source and reports which (`"source": "results" | "ndjson"`), never a merge. + +Stripped from all persisted data in v1: `checkpoint_b64`, `plot_png_base64`, embedding/attention histories. + +### RunRecord (`run.json`) — new `comfy_research/schemas/run_record.py` + +``` +run_id, schema_version: 1 +group_id: str | None # sweep session / rollout batch +parent_id: str | None # originating run (e.g. resume-continuation) or agent session +origin: "human" | "agent" | "sweep" +status: "queued" | "running" | "completed" | "failed" | "aborted" | "paused" | "crashed" | "unreadable" +created_at, started_at, finished_at: float | None (unix ms) +trainer_node_id, device, error_detail +graph: GraphDocument # config-only snapshot (results/UI blobs stripped via existing tier logic) +hyperparams: dict[str, scalar] # flattened via generated param models +``` + +`TrainRequest` gains optional `run_origin` / `run_group_id` / `run_parent_id` (default `origin="human"`) — existing clients unaffected. + +**`paused` semantics (v1)**: terminal in the store. Resume still happens through the existing client-held `checkpoint_b64` flow; the continuation is a **new run** with `parent_id` = the paused run. The store does not promise resume-by-run_id (that requires phase-2 checkpoint persistence). + +### Crash recovery + +- The index carries `last_heartbeat_at`, updated (index-only, coalesced to ≥1 s intervals) on each event batch. +- On server startup — and lazily whenever a `running` row's heartbeat is older than 60 s — the run is reconciled to `crashed` (file rewritten, index updated). No permanently stuck `running` rows. +- Interrupted `metrics.ndjson` tail lines are tolerated on read: a non-parsing final line is dropped with a logged warning. +- A `run.json` that fails to parse surfaces as `status="unreadable"` in listings (id + path) — a deliberate break from the graph-library silent-skip pattern. + +### SQLite index (`data/runs/index.db`) — write discipline + +- WAL mode, **per-process connection**, `busy_timeout=5000`, transactions kept to single upserts. WAL serializes writers; it does not make many-writer contention free — so index writes are limited to: registration, coalesced heartbeat, terminal upsert (with summary columns). Per-step metrics never touch the index. +- Columns: RunRecord scalars + `hyperparams_json` + summaries materialized at terminal state: `final_loss`, `final_test_loss`, `best_test_loss`, `steps_completed`, `duration_seconds`. +- Not a truth source: validated on startup; missing/corrupt → rebuild by scanning `data/runs/*/run.json` (also `scripts/rebuild_run_index.py`). Write path is file-first; index failure logs a warning, never fails the run. + +## HTTP API (new router `api/runs.py`) + +| Method | Path | Notes | +| --- | --- | --- | +| POST | `/api/runs` | async submit: body = TrainRequest (+ origin/group/parent), 202 + `{run_id}`; optional `Idempotency-Key` | +| GET | `/api/runs` | filters `status`, `origin`, `group_id`, `since`, `hyperparam.`, `ids=` (bulk fetch); `order_by` (e.g. `-final_loss`); **cursor pagination** (`cursor`/`limit`, default 100) | +| GET | `/api/runs/{run_id}` | RunRecord + summary | +| GET | `/api/runs/{run_id}/metrics` | single-source series (see authority rule), optional `downsample=N` | +| GET | `/api/runs/groups` | per-group status counts + best summaries | +| POST | `/api/runs/{run_id}/abort` | for submitted runs (delegates to existing train control) | +| DELETE | `/api/runs/{run_id}` | refuses `running`/`queued` | +| DELETE | `/api/runs` | bulk; requires ≥1 filter; terminal-state runs only | + +Errors use a structured envelope `{"code", "detail"}` (new pattern; existing endpoints migrate later). + +## Runs panel UI (frontend) + +- Left-rail "Runs" panel: table (status dot, group/trainer title, created, duration, final loss), grouped by `group_id`, `origin=agent` groups collapsed by default. +- Multi-select → overlay loss/test-loss curves (reuse existing curve components), legend annotated with differing hyperparams. +- Row actions: open stored config graph on canvas; delete. +- Polls `GET /api/runs` while any run is non-terminal. + +## GC / retention + +- **Single-owner**: only the API server process runs GC (startup + periodic), never the CLI. +- Eligible: terminal-state runs only, older than a 10-minute grace period. +- Default policy: keep all `origin=human`; for `origin=agent`, prune oldest beyond `max_runs_agent` (default 2000). Every prune logs what was dropped — no silent caps. +- Agent-facing cleanup: bulk `DELETE /api/runs?group_id=...`. +- Optional `data/runs/config.json`: `max_runs_agent`, `max_age_days_agent`, worker-pool size. + +## Testing + +- Unit: delta extraction from cumulative emissions (incl. step-0 baseline, abort-without-histories); authority rule; hyperparam flattening; index rebuild equivalence (write N, delete index, rebuild, compare query results). +- API: submit→poll→terminal RunRecord correctness; idempotency key; filters + cursor pagination; bulk-delete guards; unreadable surfacing; abort of queued/running submitted runs. +- Failure: kill worker mid-run → heartbeat reconciliation to `crashed`; truncated ndjson tail tolerated. +- Concurrency: worker pool at cap + interleaved queries against WAL index. +- Sweep: inner runs each captured with `group_id`; wrapper events unchanged. +- Frontend: grouped list from fixture; overlay selection. + +## Phasing + +- v1: everything above. +- Phase 2: checkpoint/blob persistence in a local content-addressed artifact store (reuse `remote/execution_artifacts.py` design) → real resume-by-run_id; backfill; structured-error migration of existing endpoints; CRL run capture; optional wandb exporter. + +## Codex review disposition (2026-08-15) + +| # | Finding | Disposition | +| --- | --- | --- | +| 1 | Critical: disconnect cancels training; tee ≠ durability | Added detached execution: `POST /api/runs` worker pool; `/api/train` keeps browser semantics (disconnect = abort, finalized as `aborted`) | +| 2 | High: CLI RunWriter = split-brain on remote host | Persistence is API-host-only; capture attaches to the SSH stream consumer | +| 3 | High: API-level tee misses sweep/CD inner runs | Capture moved to engine-level consumption sites; inner runs get own RunWriter | +| 4 | High: cumulative histories → O(n²) or loss | Delta-based writer + snapshot; explicit single-source authority rule | +| 5 | High: no crash recovery | Heartbeat + `crashed` status + startup/lazy reconciliation | +| 6 | Medium: WAL over-claimed | Write discipline: per-process conn, busy_timeout, start/heartbeat/terminal writes only | +| 7 | Medium: GC races live writers | Single-owner GC, terminal-only, grace period | +| 8 | Medium: stream-first API | 202 submit, idempotency key, cursor pagination, `ids=` bulk fetch, `queued` state | +| 9 | Medium: `paused` misleading without checkpoints | Documented terminal-in-store; resume = new run with `parent_id`; real resume in phase 2 | +| 10 | Medium: CRL variant unaddressed | Explicitly excluded from v1, listed in phase 2 | From 94c51663fb9a3f59ef71735af37d46bf7a3958e0 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 17:28:49 +0800 Subject: [PATCH 02/34] docs: add run store implementation plan (post-Codex-review) Codex review dispositions: run_id-scoped abort via per-trainer FIFO (no cross-abort, no head-of-line blocking); load_series() authority rule shared by rebuild/reconcile/API; startup wiring moved into _lifespan; SERIES_COLORS export noted; run_registered parser branch; validate with resume/policy; unreadable rows listable+deletable; bulk delete collects before deleting; removed reference to nonexistent contract test; venv/torch prerequisite. --- .../superpowers/plans/2026-08-15-run-store.md | 2657 +++++++++++++++++ 1 file changed, 2657 insertions(+) create mode 100644 docs/superpowers/plans/2026-08-15-run-store.md diff --git a/docs/superpowers/plans/2026-08-15-run-store.md b/docs/superpowers/plans/2026-08-15-run-store.md new file mode 100644 index 0000000..e0c936c --- /dev/null +++ b/docs/superpowers/plans/2026-08-15-run-store.md @@ -0,0 +1,2657 @@ +# Run Store Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Persist every training run server-side (files as truth + rebuildable SQLite index) with an agent-first async submit/query API and a canvas Runs panel. + +**Architecture:** A `RunWriter` captures trainer NDJSON events at the engine layer (delta-extracting cumulative histories into `data/runs/{run_id}/metrics.ndjson`, terminal snapshot into `results.json`, metadata in `run.json`), mirrored into a rebuildable SQLite index (`data/runs/index.db`, WAL). A worker pool executes async-submitted runs detached from any HTTP stream. Spec: `docs/superpowers/specs/2026-08-15-run-store-design.md`. + +**Tech Stack:** Python 3.10+, FastAPI, Pydantic v2, stdlib `sqlite3`; React + TypeScript + Vitest frontend (hand-rolled SVG charts, plain fetch). + +## Global Constraints + +- Files are the source of truth; the SQLite index must be rebuildable from `data/runs/*/run.json` at any time. Index write failure logs a warning, never fails a run. +- All JSON writes atomic: `.tmp` + `Path.replace` (mirror `graph_library.py:_atomic_write_json`). +- Never persist `checkpoint_b64`, `memoryCheckpoint_b64`, `plot_png_base64`, `observable_embedding_histories`, `observable_attention_slice_histories` in the run store. +- All persisted event data passes through `sanitize_train_ndjson_value` (NaN/Inf → None). +- SQLite: WAL mode, `timeout=5` connections, single-upsert transactions; index writes only at registration / coalesced heartbeat (≥1 s apart) / terminal state. +- CRL runs (`crl_trainer` node type) are NOT captured in v1; classic trainer runs only. +- Statuses: `queued | running | completed | failed | aborted | paused | crashed | unreadable`. Terminal = all except `queued`/`running`. +- New API error responses use `{"code": "", "detail": ""}` via HTTPException detail dict. +- Frontend: no new chart libs; new CSS uses `var(--cr-*)` tokens only (a raw hex fails `npm run verify:css-tokens`); backend URLs are relative `/api/...`. +- Backend tests: pytest, files in `comfy_research/tests/`; run store tests must isolate via the `COMFYRESEARCH_RUNS_DIR` env var + `tmp_path`. Frontend tests: Vitest in `frontend/src/graph/__tests__/`, `// @vitest-environment jsdom` where DOM is needed. +- Run backend tests with `python -m pytest -v` from the repo root, INSIDE the project's environment (torch and requirements.txt installed — a bare interpreter fails on `import torch`). Verify with `python -c "import torch"` before starting; if it fails, ask the user which environment to use rather than pip-installing torch. + +## File Structure + +| File | Responsibility | +| --- | --- | +| `comfy_research/schemas/run_record.py` (new) | `RunRecord` model, run-id minting, result-key stripping, hyperparam flattening | +| `comfy_research/engine/runs/run_store.py` (new) | Paths, atomic file persistence, `MetricsDeltaTracker`, results/metrics read-back | +| `comfy_research/engine/runs/run_index.py` (new) | SQLite index: schema, upsert, query, rebuild, stale-run reconciliation | +| `comfy_research/engine/runs/run_writer.py` (new) | `RunWriter` facade (files + index + heartbeat + finalization), `capture_events` | +| `comfy_research/engine/runs/run_worker.py` (new) | Async-submit worker pool, per-trainer serialization, abort, idempotency | +| `comfy_research/engine/runs/run_gc.py` (new) | Retention config + single-owner GC | +| `comfy_research/schemas/train_request.py` (modify) | Optional `run_origin` / `run_group_id` / `run_parent_id` | +| `comfy_research/api/train.py` (modify) | Tee local + remote stream generators through `capture_events` | +| `comfy_research/engine/runs/train_sweep.py` (modify) | Per-inner-run capture with `group_id` = sweep session | +| `comfy_research/api/runs.py` (new) | `/api/runs` router | +| `comfy_research/main.py` (modify) | Register router; startup reconciliation + GC | +| `frontend/src/graph/runsApi.ts` (new) | Typed fetch client for `/api/runs` | +| `frontend/src/components/RunsPanel.tsx` (new) | Rail panel: grouped run list, selection, overlay chart | +| `frontend/src/components/railTypes.ts`, `LeftNavRail.tsx`, `ResearchCanvas.tsx` (modify) | Rail wiring + open-run-graph action | +| `docs/en/reference/runs-api.md` (new), `docs/en/reference/data-contracts.md` (modify) | API + on-disk contract docs | + +--- + +### Task 1: RunRecord schema, stripping, hyperparam flattening + +**Files:** +- Create: `comfy_research/schemas/run_record.py` +- Test: `comfy_research/tests/test_run_record.py` + +**Interfaces:** +- Consumes: `GraphDocument`, `Node` from `comfy_research.schemas.graph`; `load_node_manifest` from `comfy_research.generated.node_manifest`. +- Produces: `RunRecord` (pydantic, fields below), `new_run_id() -> str`, `strip_result_data(nodes: list[Node]) -> list[Node]`, `flatten_hyperparams(nodes: list[Node]) -> dict[str, float | int | str | bool]`, `RUN_RESULT_DATA_KEYS: frozenset[str]`, `TERMINAL_STATUSES: frozenset[str]`, `now_ms() -> float`. + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_run_record.py +from __future__ import annotations + +from comfy_research.schemas.graph import Edge, GraphDocument, Node +from comfy_research.schemas.run_record import ( + RUN_RESULT_DATA_KEYS, + RunRecord, + flatten_hyperparams, + new_run_id, + now_ms, + strip_result_data, +) + + +def _trainer_node() -> Node: + return Node( + id="t1", + type="trainer", + data={ + "trainingSteps": 4, + "computeDevice": "cpu", + "instanceTitle": "Trainer", + "lossHistory": [1.0, 0.5], + "memoryCheckpoint_b64": "QUJD", + "plotPngBase64": "aW1n", + }, + ) + + +def test_new_run_id_prefix_and_uniqueness() -> None: + a, b = new_run_id(), new_run_id() + assert a.startswith("run-") and len(a) == 16 + assert a != b + + +def test_strip_result_data_removes_blobs_keeps_config() -> None: + stripped = strip_result_data([_trainer_node()]) + data = stripped[0].data + assert data["trainingSteps"] == 4 + assert data["instanceTitle"] == "Trainer" + assert "lossHistory" not in data + assert "memoryCheckpoint_b64" not in data + assert "plotPngBase64" not in data + # original untouched + assert "lossHistory" in _trainer_node().data + + +def test_result_keys_cover_known_blobs() -> None: + for key in ("checkpoint_b64", "memoryCheckpoint_b64", "plotPngBase64", + "lossHistory", "testLossHistory", "regLossHistory", "stepTicks", + "observableMetricHistories", "embeddingHistory", + "attentionMapFrames", "valueHistory", "runSummary", "lastError"): + assert key in RUN_RESULT_DATA_KEYS + + +def test_flatten_hyperparams_declared_scalars_only() -> None: + flat = flatten_hyperparams([_trainer_node()]) + assert flat["t1.trainingSteps"] == 4 + assert flat["t1.computeDevice"] == "cpu" + assert "t1.lossHistory" not in flat # not a declared field + assert "t1.instanceTitle" not in flat # declared but excluded as label + + +def test_run_record_roundtrip() -> None: + rec = RunRecord( + run_id=new_run_id(), + origin="agent", + status="queued", + created_at=now_ms(), + trainer_node_id="t1", + graph=GraphDocument(version=1, nodes=strip_result_data([_trainer_node()]), edges=[]), + hyperparams=flatten_hyperparams([_trainer_node()]), + ) + again = RunRecord.model_validate(rec.model_dump(mode="json")) + assert again.run_id == rec.run_id + assert again.schema_version == 1 + assert again.group_id is None and again.finished_at is None +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_run_record.py -v` +Expected: FAIL with `ModuleNotFoundError: comfy_research.schemas.run_record` + +- [ ] **Step 3: Write the implementation** + +```python +# comfy_research/schemas/run_record.py +"""RunRecord: persisted metadata for one training run (``data/runs/{run_id}/run.json``).""" +from __future__ import annotations + +import time +import uuid +from typing import Any, Literal + +from pydantic import BaseModel, Field + +from comfy_research.generated.node_manifest import load_node_manifest +from comfy_research.schemas.graph import GraphDocument, Node + +RunStatus = Literal[ + "queued", "running", "completed", "failed", "aborted", "paused", "crashed", "unreadable" +] +TERMINAL_STATUSES: frozenset[str] = frozenset( + {"completed", "failed", "aborted", "paused", "crashed"} +) + +# Result/UI payloads the browser stashes into node data; never persisted in run configs. +# Superset mirror of frontend/src/graph/graphFileExportTier.ts strip lists. +RUN_RESULT_DATA_KEYS: frozenset[str] = frozenset( + { + "checkpoint_b64", "memoryCheckpoint_b64", + "lossHistory", "testLossHistory", "regLossHistory", + "stepTicks", "epochTicks", "observableMetricHistories", + "lastTrainLoopSeconds", "plotPngBase64", "valueHistory", + "embeddingHistory", "attentionMapFrames", "previewGrid", + "histogramPng", "imageGrid", "runSummary", "lastError", + "lastSweepSummary", "observableEmbeddingHistories", + "observableAttentionSliceHistories", + } +) + +_HYPERPARAM_EXCLUDED_FIELDS = frozenset({"instanceTitle"}) + + +class RunRecord(BaseModel): + run_id: str + schema_version: int = 1 + group_id: str | None = None + parent_id: str | None = None + origin: Literal["human", "agent", "sweep"] = "human" + status: RunStatus = "queued" + created_at: float + started_at: float | None = None + finished_at: float | None = None + trainer_node_id: str + device: str = "" + error_detail: str = "" + graph: GraphDocument + hyperparams: dict[str, Any] = Field(default_factory=dict) + + +def now_ms() -> float: + return time.time() * 1000.0 + + +def new_run_id() -> str: + return "run-" + uuid.uuid4().hex[:12] + + +def strip_result_data(nodes: list[Node]) -> list[Node]: + """Copy nodes with result/UI blobs removed from ``data`` (config-only snapshot).""" + out: list[Node] = [] + for n in nodes: + data = {k: v for k, v in (n.data or {}).items() if k not in RUN_RESULT_DATA_KEYS} + out.append(n.model_copy(update={"data": data}, deep=True)) + return out + + +def _declared_field_keys(node_type: str) -> list[str]: + for entry in load_node_manifest(): + if entry.get("type") == node_type: + return [f["key"] for f in entry.get("fields", [])] + return [] + + +def flatten_hyperparams(nodes: list[Node]) -> dict[str, Any]: + """``{node_id}.{field}`` -> scalar, for manifest-declared fields present in node data.""" + flat: dict[str, Any] = {} + for n in nodes: + data = n.data or {} + for key in _declared_field_keys(str(n.type)): + if key in _HYPERPARAM_EXCLUDED_FIELDS or key not in data: + continue + v = data[key] + if isinstance(v, (bool, int, float, str)): + flat[f"{n.id}.{key}"] = v + return flat +``` + +Before finalizing, check the actual shape returned by `load_node_manifest()` (`comfy_research/generated/node_manifest.py`) — if it returns an object keyed by type or a list under a `"nodes"` key, adapt `_declared_field_keys` accordingly (keep the returned-keys contract identical). + +- [ ] **Step 4: Run test to verify it passes** + +Run: `python -m pytest comfy_research/tests/test_run_record.py -v` +Expected: PASS (5 tests) + +- [ ] **Step 5: Commit** + +```bash +git add comfy_research/schemas/run_record.py comfy_research/tests/test_run_record.py +git commit -m "feat: add RunRecord schema with result stripping and hyperparam flattening" +``` + +--- + +### Task 2: File store + metrics delta tracker + +**Files:** +- Create: `comfy_research/engine/runs/run_store.py` +- Test: `comfy_research/tests/test_run_store.py` + +**Interfaces:** +- Consumes: `RunRecord`, `TERMINAL_STATUSES` (Task 1); `sanitize_train_ndjson_value`. +- Produces: `runs_root() -> Path` (env `COMFYRESEARCH_RUNS_DIR` override, else `/data/runs`), `run_dir(run_id) -> Path`, `write_run_record(rec) -> None` (atomic), `read_run_record(run_id) -> RunRecord | None`, `append_metric_rows(run_id, rows: list[dict]) -> None`, `read_metric_rows(run_id) -> list[dict]` (tolerates truncated tail), `write_results(run_id, payload: dict) -> None` (atomic, sanitized), `read_results(run_id) -> dict | None`, `load_series(run_id) -> tuple[str, dict]` (single implementation of the authority rule: `("results", ...)` when `results.json` exists, else `("ndjson", ...)` rebuilt from metric rows), `class MetricsDeltaTracker` with `extract(event: dict) -> list[dict]` and `snapshot() -> dict`, `SERIES_KEYS`, `summarize(snapshot: dict) -> dict` (final_loss, final_test_loss, best_test_loss, steps_completed). + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_run_store.py +from __future__ import annotations + +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms +from comfy_research.engine.runs import run_store + + +def _record(run_id: str) -> RunRecord: + return RunRecord( + run_id=run_id, origin="agent", status="running", created_at=now_ms(), + trainer_node_id="t1", graph=GraphDocument(version=1, nodes=[], edges=[]), + ) + + +def _metrics_event(n: int) -> dict: + return { + "type": "metrics", "step": n, + "loss_history": [1.0 / (i + 1) for i in range(n)], + "test_loss_history": [2.0 / (i + 1) for i in range(n)], + "reg_loss_history": [], + "step_ticks": list(range(n)), + "epoch_ticks": [], + "observable_metric_histories": {"obs1:acc": [float(i) for i in range(n)]}, + "observable_warnings": {}, + } + + +def test_runs_root_env_override(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path / "r")) + assert run_store.runs_root() == tmp_path / "r" + + +def test_record_roundtrip_and_atomicity(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + rid = new_run_id() + run_store.write_run_record(_record(rid)) + rec = run_store.read_run_record(rid) + assert rec is not None and rec.run_id == rid + assert not list((tmp_path / rid).glob("*.tmp")) + assert run_store.read_run_record("run-missing00000") is None + + +def test_delta_tracker_appends_only_new_rows() -> None: + tracker = run_store.MetricsDeltaTracker() + rows1 = tracker.extract(_metrics_event(2)) + rows2 = tracker.extract(_metrics_event(5)) + assert len(rows1) == 2 and len(rows2) == 3 + assert rows2[0]["idx"] == 2 and rows2[0]["step"] == 2 + assert rows2[-1]["loss"] == 1.0 / 5 + assert rows2[-1]["obs"]["obs1:acc"] == 4.0 + snap = tracker.snapshot() + assert len(snap["loss_history"]) == 5 + + +def test_delta_tracker_ignores_non_metrics_events() -> None: + tracker = run_store.MetricsDeltaTracker() + assert tracker.extract({"type": "progress", "step": 1, "total": 4}) == [] + + +def test_metrics_ndjson_roundtrip_tolerates_truncated_tail(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + rid = new_run_id() + run_store.write_run_record(_record(rid)) + run_store.append_metric_rows(rid, [{"idx": 0, "loss": 1.0}, {"idx": 1, "loss": 0.5}]) + path = tmp_path / rid / "metrics.ndjson" + path.write_text(path.read_text() + '{"idx": 2, "lo', encoding="utf-8") + rows = run_store.read_metric_rows(rid) + assert [r["idx"] for r in rows] == [0, 1] + + +def test_results_sanitized_and_stripped(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + rid = new_run_id() + run_store.write_run_record(_record(rid)) + run_store.write_results(rid, { + "loss_history": [1.0, float("nan")], + "checkpoint_b64": "QUJD", + "plot_png_base64": "aW1n", + "observable_embedding_histories": {"a": [1]}, + }) + res = run_store.read_results(rid) + assert res == {"loss_history": [1.0, None]} + + +def test_summarize() -> None: + tracker = run_store.MetricsDeltaTracker() + tracker.extract(_metrics_event(4)) + s = run_store.summarize(tracker.snapshot()) + assert s["final_loss"] == 0.25 + assert s["best_test_loss"] == 0.5 + assert s["steps_completed"] == 4 +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_run_store.py -v` +Expected: FAIL with `ImportError` / module not found + +- [ ] **Step 3: Write the implementation** + +```python +# comfy_research/engine/runs/run_store.py +"""File layer of the run store: ``data/runs/{run_id}/`` is the source of truth.""" +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +from comfy_research.schemas.run_record import RunRecord +from comfy_research.schemas.train_request import sanitize_train_ndjson_value + +_REPO_ROOT = Path(__file__).resolve().parents[3] + +SERIES_KEYS = ("loss_history", "test_loss_history", "reg_loss_history", + "step_ticks", "epoch_ticks") +_RESULT_STRIP_KEYS = frozenset( + {"checkpoint_b64", "plot_png_base64", "visualization_node_ids", + "observable_viz_updates", "observable_embedding_histories", + "observable_attention_slice_histories", "type"} +) + + +def runs_root() -> Path: + env = os.environ.get("COMFYRESEARCH_RUNS_DIR", "").strip() + return Path(env) if env else _REPO_ROOT / "data" / "runs" + + +def run_dir(run_id: str) -> Path: + if not run_id or "/" in run_id or "\\" in run_id or run_id.startswith("."): + raise ValueError(f"invalid run_id: {run_id!r}") + return runs_root() / run_id + + +def _atomic_write_json(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + tmp.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") + tmp.replace(path) + + +def write_run_record(rec: RunRecord) -> None: + _atomic_write_json(run_dir(rec.run_id) / "run.json", rec.model_dump(mode="json")) + + +def read_run_record(run_id: str) -> RunRecord | None: + path = run_dir(run_id) / "run.json" + try: + return RunRecord.model_validate_json(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + + +def append_metric_rows(run_id: str, rows: list[dict]) -> None: + if not rows: + return + path = run_dir(run_id) / "metrics.ndjson" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as f: + for row in rows: + f.write(json.dumps(sanitize_train_ndjson_value(row), separators=(",", ":")) + "\n") + + +def read_metric_rows(run_id: str) -> list[dict]: + path = run_dir(run_id) / "metrics.ndjson" + if not path.is_file(): + return [] + rows: list[dict] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + try: + rows.append(json.loads(line)) + except json.JSONDecodeError: + break # truncated tail from an interrupted append; drop it + return rows + + +def write_results(run_id: str, payload: dict) -> None: + kept = {k: v for k, v in payload.items() if k not in _RESULT_STRIP_KEYS} + _atomic_write_json(run_dir(run_id) / "results.json", + sanitize_train_ndjson_value(kept)) + + +def read_results(run_id: str) -> dict | None: + path = run_dir(run_id) / "results.json" + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + + +def load_series(run_id: str) -> tuple[str, dict]: + """Authority rule: ``results.json`` if present, else series rebuilt from ``metrics.ndjson``. + + Returns ("results" | "ndjson", series-dict). The single place that decides which + source wins — the API, index rebuild, and reconciliation must all use it. + """ + results = read_results(run_id) + if results is not None: + return "results", results + rows = read_metric_rows(run_id) + return "ndjson", { + "loss_history": [r.get("loss") for r in rows], + "test_loss_history": [r["test_loss"] for r in rows if "test_loss" in r], + "reg_loss_history": [r["reg_loss"] for r in rows if "reg_loss" in r], + "step_ticks": [r["step"] for r in rows if "step" in r], + "epoch_ticks": [r["epoch"] for r in rows if "epoch" in r], + } + + +class MetricsDeltaTracker: + """Turn cumulative-history ``metrics`` events into append-only delta rows. + + The trainer re-sends full histories on every emission; appending raw payloads + would be O(n^2) in storage. Track last-seen length, emit only new indices, + keep the latest cumulative snapshot for terminal fallback. + """ + + def __init__(self) -> None: + self._latest: dict[str, Any] = {k: [] for k in SERIES_KEYS} + self._latest["observable_metric_histories"] = {} + self._seen = 0 + + def extract(self, event: dict) -> list[dict]: + if event.get("type") != "metrics": + return [] + for key in SERIES_KEYS: + v = event.get(key) + if isinstance(v, list): + self._latest[key] = list(v) + obs = event.get("observable_metric_histories") + if isinstance(obs, dict): + self._latest["observable_metric_histories"] = { + str(k): list(v) for k, v in obs.items() if isinstance(v, list) + } + loss = self._latest["loss_history"] + rows: list[dict] = [] + for i in range(self._seen, len(loss)): + row: dict[str, Any] = {"idx": i, "loss": loss[i]} + for name, key in (("step", "step_ticks"), ("test_loss", "test_loss_history"), + ("reg_loss", "reg_loss_history"), ("epoch", "epoch_ticks")): + series = self._latest[key] + if i < len(series): + row[name] = series[i] + obs_row = { + k: v[i] + for k, v in self._latest["observable_metric_histories"].items() + if i < len(v) + } + if obs_row: + row["obs"] = obs_row + rows.append(row) + self._seen = len(loss) + return rows + + def snapshot(self) -> dict: + return {k: list(v) if isinstance(v, list) else dict(v) for k, v in self._latest.items()} + + +def summarize(snapshot: dict) -> dict: + loss = snapshot.get("loss_history") or [] + test = [x for x in (snapshot.get("test_loss_history") or []) if isinstance(x, (int, float))] + return { + "final_loss": loss[-1] if loss else None, + "final_test_loss": test[-1] if test else None, + "best_test_loss": min(test) if test else None, + "steps_completed": len(loss), + } +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `python -m pytest comfy_research/tests/test_run_store.py -v` +Expected: PASS (7 tests) + +- [ ] **Step 5: Commit** + +```bash +git add comfy_research/engine/runs/run_store.py comfy_research/tests/test_run_store.py +git commit -m "feat: add run store file layer with metrics delta tracker" +``` + +--- + +### Task 3: SQLite index — upsert, query, rebuild, reconcile + +**Files:** +- Create: `comfy_research/engine/runs/run_index.py` +- Test: `comfy_research/tests/test_run_index.py` + +**Interfaces:** +- Consumes: `run_store.runs_root/read_run_record/write_run_record`, `RunRecord`, `TERMINAL_STATUSES`, `now_ms`. +- Produces: `index_path() -> Path`, `upsert_run(rec: RunRecord, summary: dict | None = None) -> None`, `touch_heartbeat(run_id: str, at_ms: float) -> None`, `class RunQuery` (pydantic-free dataclass: `status`, `origin`, `group_id`, `since_ms`, `ids`, `hyperparams: dict[str, str]`, `order_by="-created_at"`, `limit=100`, `cursor=None`), `query_runs(q: RunQuery) -> tuple[list[dict], str | None]` (rows + next cursor), `group_summary() -> list[dict]`, `delete_rows(run_ids: list[str]) -> None`, `rebuild_index() -> int`, `reconcile_stale_running(timeout_ms: float = 60_000) -> list[str]`. +- Row dict keys: `run_id, group_id, parent_id, origin, status, created_at, started_at, finished_at, last_heartbeat_at, trainer_node_id, device, error_detail, hyperparams (dict), final_loss, final_test_loss, best_test_loss, steps_completed, duration_seconds`. + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_run_index.py +from __future__ import annotations + +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery + + +def _record(status: str = "running", origin: str = "agent", + group: str | None = None, lr: float = 0.01) -> RunRecord: + return RunRecord( + run_id=new_run_id(), origin=origin, status=status, created_at=now_ms(), + group_id=group, trainer_node_id="t1", + graph=GraphDocument(version=1, nodes=[], edges=[]), + hyperparams={"opt1.lr": lr}, + ) + + +def _seed(monkeypatch, tmp_path): + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + + +def test_upsert_and_query_filters(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + a = _record(status="completed", lr=0.01) + b = _record(status="failed", lr=0.1) + for r in (a, b): + run_store.write_run_record(r) + run_index.upsert_run(r) + rows, cursor = run_index.query_runs(RunQuery(status="completed")) + assert [r["run_id"] for r in rows] == [a.run_id] + assert cursor is None + rows, _ = run_index.query_runs(RunQuery(hyperparams={"opt1.lr": "0.1"})) + assert [r["run_id"] for r in rows] == [b.run_id] + rows, _ = run_index.query_runs(RunQuery(ids=[a.run_id])) + assert rows[0]["hyperparams"] == {"opt1.lr": 0.01} + + +def test_cursor_pagination(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + made = [] + for _ in range(5): + r = _record(status="completed") + run_store.write_run_record(r) + run_index.upsert_run(r) + made.append(r.run_id) + page1, cur1 = run_index.query_runs(RunQuery(limit=2)) + page2, cur2 = run_index.query_runs(RunQuery(limit=2, cursor=cur1)) + page3, cur3 = run_index.query_runs(RunQuery(limit=2, cursor=cur2)) + ids = [r["run_id"] for r in page1 + page2 + page3] + assert sorted(ids) == sorted(made) and len(ids) == 5 + assert cur3 is None + + +def test_summary_columns_and_order_by(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + a, b = _record(status="completed"), _record(status="completed") + for r, loss in ((a, 0.5), (b, 0.1)): + run_store.write_run_record(r) + run_index.upsert_run(r, summary={"final_loss": loss, "final_test_loss": None, + "best_test_loss": None, "steps_completed": 3}) + rows, _ = run_index.query_runs(RunQuery(order_by="final_loss")) + assert [r["run_id"] for r in rows] == [b.run_id, a.run_id] + assert rows[0]["final_loss"] == 0.1 and rows[0]["steps_completed"] == 3 + + +def test_rebuild_surfaces_unreadable_run(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + good = _record(status="completed") + run_store.write_run_record(good) + run_index.upsert_run(good) + bad_dir = tmp_path / "run-corrupted0001" + bad_dir.mkdir() + (bad_dir / "run.json").write_text("{not json", encoding="utf-8") + assert run_index.rebuild_index() == 2 + rows, _ = run_index.query_runs(RunQuery(status="unreadable")) + assert [r["run_id"] for r in rows] == ["run-corrupted0001"] + + +def test_rebuild_from_files(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + recs = [_record(status="completed") for _ in range(3)] + for r in recs: + run_store.write_run_record(r) + run_index.upsert_run(r) + before, _ = run_index.query_runs(RunQuery()) + run_index.index_path().unlink() + assert run_index.rebuild_index() == 3 + after, _ = run_index.query_runs(RunQuery()) + assert {r["run_id"] for r in after} == {r["run_id"] for r in before} + + +def test_reconcile_stale_running(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + stale = _record(status="running") + run_store.write_run_record(stale) + run_index.upsert_run(stale) + run_index.touch_heartbeat(stale.run_id, now_ms() - 120_000) + fresh = _record(status="running") + run_store.write_run_record(fresh) + run_index.upsert_run(fresh) + run_index.touch_heartbeat(fresh.run_id, now_ms()) + + crashed = run_index.reconcile_stale_running(timeout_ms=60_000) + assert crashed == [stale.run_id] + assert run_store.read_run_record(stale.run_id).status == "crashed" + rows, _ = run_index.query_runs(RunQuery(ids=[fresh.run_id])) + assert rows[0]["status"] == "running" + + +def test_group_summary_and_delete(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + a = _record(status="completed", group="g1") + b = _record(status="failed", group="g1") + for r in (a, b): + run_store.write_run_record(r) + run_index.upsert_run(r) + groups = run_index.group_summary() + g1 = next(g for g in groups if g["group_id"] == "g1") + assert g1["counts"] == {"completed": 1, "failed": 1} + run_index.delete_rows([a.run_id, b.run_id]) + rows, _ = run_index.query_runs(RunQuery(group_id="g1")) + assert rows == [] +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_run_index.py -v` +Expected: FAIL with module not found + +- [ ] **Step 3: Write the implementation** + +```python +# comfy_research/engine/runs/run_index.py +"""Rebuildable SQLite read-index over ``data/runs/*/run.json`` (never a truth source).""" +from __future__ import annotations + +import json +import logging +import sqlite3 +from dataclasses import dataclass, field +from pathlib import Path + +from comfy_research.engine.runs import run_store +from comfy_research.schemas.run_record import RunRecord, now_ms + +logger = logging.getLogger(__name__) + +_COLUMNS = ( + "run_id", "group_id", "parent_id", "origin", "status", "created_at", + "started_at", "finished_at", "last_heartbeat_at", "trainer_node_id", + "device", "error_detail", "hyperparams_json", "final_loss", + "final_test_loss", "best_test_loss", "steps_completed", "duration_seconds", +) +_ORDERABLE = {"created_at", "finished_at", "final_loss", "final_test_loss", + "best_test_loss", "steps_completed", "duration_seconds"} + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS runs ( + run_id TEXT PRIMARY KEY, group_id TEXT, parent_id TEXT, + origin TEXT NOT NULL, status TEXT NOT NULL, + created_at REAL NOT NULL, started_at REAL, finished_at REAL, + last_heartbeat_at REAL, trainer_node_id TEXT NOT NULL, + device TEXT DEFAULT '', error_detail TEXT DEFAULT '', + hyperparams_json TEXT NOT NULL DEFAULT '{}', + final_loss REAL, final_test_loss REAL, best_test_loss REAL, + steps_completed INTEGER, duration_seconds REAL +); +CREATE INDEX IF NOT EXISTS idx_runs_status ON runs(status); +CREATE INDEX IF NOT EXISTS idx_runs_group ON runs(group_id); +CREATE INDEX IF NOT EXISTS idx_runs_created ON runs(created_at); +""" + + +def index_path() -> Path: + return run_store.runs_root() / "index.db" + + +def _connect() -> sqlite3.Connection: + index_path().parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(index_path(), timeout=5) + conn.execute("PRAGMA journal_mode=WAL") + conn.executescript(_SCHEMA) + return conn + + +def upsert_run(rec: RunRecord, summary: dict | None = None) -> None: + s = summary or {} + duration = ( + (rec.finished_at - rec.started_at) / 1000.0 + if rec.finished_at is not None and rec.started_at is not None else None + ) + try: + with _connect() as conn: + conn.execute( + f"INSERT OR REPLACE INTO runs ({','.join(_COLUMNS)}) " + f"VALUES ({','.join('?' * len(_COLUMNS))})", + (rec.run_id, rec.group_id, rec.parent_id, rec.origin, rec.status, + rec.created_at, rec.started_at, rec.finished_at, now_ms(), + rec.trainer_node_id, rec.device, rec.error_detail, + json.dumps(rec.hyperparams), s.get("final_loss"), + s.get("final_test_loss"), s.get("best_test_loss"), + s.get("steps_completed"), duration), + ) + except sqlite3.Error: + logger.warning("run index upsert failed for %s", rec.run_id, exc_info=True) + + +def touch_heartbeat(run_id: str, at_ms: float) -> None: + try: + with _connect() as conn: + conn.execute("UPDATE runs SET last_heartbeat_at=? WHERE run_id=?", (at_ms, run_id)) + except sqlite3.Error: + logger.warning("run index heartbeat failed for %s", run_id, exc_info=True) + + +@dataclass +class RunQuery: + status: str | None = None + origin: str | None = None + group_id: str | None = None + since_ms: float | None = None + ids: list[str] | None = None + hyperparams: dict[str, str] = field(default_factory=dict) + order_by: str = "-created_at" + limit: int = 100 + cursor: str | None = None + + +def _row_to_dict(row: tuple) -> dict: + d = dict(zip(_COLUMNS, row)) + d["hyperparams"] = json.loads(d.pop("hyperparams_json") or "{}") + return d + + +def query_runs(q: RunQuery) -> tuple[list[dict], str | None]: + key = q.order_by.lstrip("-") + if key not in _ORDERABLE: + key, q = "created_at", RunQuery(**{**q.__dict__, "order_by": "-created_at"}) + direction = "DESC" if q.order_by.startswith("-") else "ASC" + where, params = ["1=1"], [] + for col, val in (("status", q.status), ("origin", q.origin), ("group_id", q.group_id)): + if val is not None: + where.append(f"{col}=?") + params.append(val) + if q.since_ms is not None: + where.append("created_at>=?") + params.append(q.since_ms) + if q.ids: + where.append(f"run_id IN ({','.join('?' * len(q.ids))})") + params.extend(q.ids) + for hk, hv in q.hyperparams.items(): + where.append("CAST(json_extract(hyperparams_json, ?) AS TEXT)=?") + params.extend([f'$."{hk}"', hv]) + if q.cursor: + cv, cid = q.cursor.rsplit(":", 1) + op = "<" if direction == "DESC" else ">" + where.append(f"({key} {op} ? OR ({key} = ? AND run_id {op} ?))") + params.extend([float(cv), float(cv), cid]) + sql = (f"SELECT {','.join(_COLUMNS)} FROM runs WHERE {' AND '.join(where)} " + f"ORDER BY {key} {direction} NULLS LAST, run_id {direction} LIMIT ?") + limit = max(1, min(int(q.limit), 500)) + with _connect() as conn: + rows = [_row_to_dict(r) for r in conn.execute(sql, [*params, limit + 1])] + next_cursor = None + if len(rows) > limit: + rows = rows[:limit] + last = rows[-1] + next_cursor = f"{last[key]}:{last['run_id']}" + return rows, next_cursor + + +def group_summary() -> list[dict]: + with _connect() as conn: + raw = conn.execute( + "SELECT group_id, status, COUNT(*), MIN(best_test_loss), MIN(final_loss) " + "FROM runs WHERE group_id IS NOT NULL GROUP BY group_id, status" + ).fetchall() + groups: dict[str, dict] = {} + for gid, status, count, best_test, best_final in raw: + g = groups.setdefault(gid, {"group_id": gid, "counts": {}, + "best_test_loss": None, "best_final_loss": None}) + g["counts"][status] = count + for k, v in (("best_test_loss", best_test), ("best_final_loss", best_final)): + if v is not None and (g[k] is None or v < g[k]): + g[k] = v + return sorted(groups.values(), key=lambda g: g["group_id"]) + + +def delete_rows(run_ids: list[str]) -> None: + if not run_ids: + return + with _connect() as conn: + conn.execute(f"DELETE FROM runs WHERE run_id IN ({','.join('?' * len(run_ids))})", run_ids) + + +def rebuild_index() -> int: + root = run_store.runs_root() + count = 0 + if index_path().exists(): + index_path().unlink() + for entry in sorted(root.glob("run-*/run.json")): + run_id = entry.parent.name + rec = run_store.read_run_record(run_id) + if rec is None: + # Corrupted run.json: surface it, don't silently skip (spec: unreadable rows). + try: + with _connect() as conn: + conn.execute( + "INSERT OR REPLACE INTO runs (run_id, origin, status, created_at, " + "trainer_node_id) VALUES (?, 'human', 'unreadable', 0, '')", + (run_id,), + ) + count += 1 + except sqlite3.Error: + logger.warning("could not index unreadable run %s", run_id, exc_info=True) + continue + _, series = run_store.load_series(rec.run_id) + upsert_run(rec, run_store.summarize(series)) + count += 1 + return count + + +def reconcile_stale_running(timeout_ms: float = 60_000) -> list[str]: + cutoff = now_ms() - timeout_ms + with _connect() as conn: + stale = [r[0] for r in conn.execute( + "SELECT run_id FROM runs WHERE status IN ('running','queued') " + "AND (last_heartbeat_at IS NULL OR last_heartbeat_at < ?)", (cutoff,) + )] + crashed = [] + for run_id in stale: + rec = run_store.read_run_record(run_id) + if rec is None or rec.status not in ("running", "queued"): + continue + rec = rec.model_copy(update={"status": "crashed", "finished_at": now_ms()}) + run_store.write_run_record(rec) + _, series = run_store.load_series(run_id) + upsert_run(rec, run_store.summarize(series)) + crashed.append(run_id) + return crashed +``` + +Note: SQLite < 3.30 lacks `NULLS LAST`; if CI's Python bundles an older SQLite the `ORDER BY` should become `ORDER BY ({key} IS NULL), {key} {direction}, run_id {direction}` — check `sqlite3.sqlite_version` and use the portable form directly if simpler. + +- [ ] **Step 4: Run test to verify it passes** + +Run: `python -m pytest comfy_research/tests/test_run_index.py -v` +Expected: PASS (7 tests) + +- [ ] **Step 5: Commit** + +```bash +git add comfy_research/engine/runs/run_index.py comfy_research/tests/test_run_index.py +git commit -m "feat: add rebuildable SQLite run index with cursor queries and reconciliation" +``` + +--- + +### Task 4: RunWriter facade + capture_events + +**Files:** +- Create: `comfy_research/engine/runs/run_writer.py` +- Test: `comfy_research/tests/test_run_writer.py` + +**Interfaces:** +- Consumes: Tasks 1–3. +- Produces: + - `class RunWriter`: `__init__(self, record: RunRecord)` (persists record + index row), `mark_running() -> None`, `on_event(event: dict) -> None` (delta append, coalesced heartbeat ≥1 s, terminal handling), `finalize(status: str, error_detail: str = "") -> None` (idempotent), `finalize_disconnect() -> None` (→ `aborted` if not yet terminal), property `is_terminal: bool`, property `record: RunRecord`. + - `capture_events(events: Iterator[dict], writer: RunWriter) -> Iterator[dict]`. +- Terminal mapping inside `on_event`: `complete → completed` (results.json from event payload), `aborted → aborted`, `paused → paused` (results.json from event payload), `error → failed` (results.json from tracker snapshot, `error_detail` from event). `aborted` uses tracker snapshot for results.json. + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_run_writer.py +from __future__ import annotations + +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.engine.runs.run_writer import RunWriter, capture_events + + +def _writer(monkeypatch, tmp_path) -> RunWriter: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + return RunWriter(RunRecord( + run_id=new_run_id(), origin="agent", status="queued", created_at=now_ms(), + trainer_node_id="t1", graph=GraphDocument(version=1, nodes=[], edges=[]), + )) + + +def _events() -> list[dict]: + return [ + {"type": "progress", "step": 0, "total": 2}, + {"type": "metrics", "step": 2, "loss_history": [1.0, 0.5], + "test_loss_history": [], "reg_loss_history": [], "step_ticks": [0, 1], + "epoch_ticks": [], "observable_metric_histories": {}, "observable_warnings": {}}, + {"type": "complete", "checkpoint_b64": "QUJD", "plot_png_base64": "aW1n", + "loss_history": [1.0, 0.5, 0.25], "test_loss_history": [], + "reg_loss_history": [], "step_ticks": [0, 1, 2], "epoch_ticks": [], + "observable_viz_updates": [], "observable_metric_histories": {}, + "observable_embedding_histories": {}, "observable_attention_slice_histories": {}, + "observable_warnings": {}, "train_loop_seconds": 0.1, + "visualization_node_ids": []}, + ] + + +def test_complete_flow(monkeypatch, tmp_path) -> None: + w = _writer(monkeypatch, tmp_path) + seen = list(capture_events(iter(_events()), w)) + assert [e["type"] for e in seen] == ["progress", "metrics", "complete"] + rec = run_store.read_run_record(w.record.run_id) + assert rec.status == "completed" and rec.started_at and rec.finished_at + res = run_store.read_results(rec.run_id) + assert res["loss_history"] == [1.0, 0.5, 0.25] + assert "checkpoint_b64" not in res and "plot_png_base64" not in res + rows, _ = run_index.query_runs(RunQuery(ids=[rec.run_id])) + assert rows[0]["status"] == "completed" and rows[0]["final_loss"] == 0.25 + assert len(run_store.read_metric_rows(rec.run_id)) == 2 + + +def test_disconnect_mid_stream_finalizes_aborted(monkeypatch, tmp_path) -> None: + w = _writer(monkeypatch, tmp_path) + gen = capture_events(iter(_events()), w) + next(gen) + next(gen) # consumed progress + metrics, then client goes away + gen.close() + rec = run_store.read_run_record(w.record.run_id) + assert rec.status == "aborted" + assert run_store.read_results(rec.run_id)["loss_history"] == [1.0, 0.5] + + +def test_error_event_finalizes_failed(monkeypatch, tmp_path) -> None: + w = _writer(monkeypatch, tmp_path) + list(capture_events(iter([{"type": "error", "detail": "boom"}]), w)) + rec = run_store.read_run_record(w.record.run_id) + assert rec.status == "failed" and rec.error_detail == "boom" + + +def test_finalize_idempotent(monkeypatch, tmp_path) -> None: + w = _writer(monkeypatch, tmp_path) + list(capture_events(iter(_events()), w)) + w.finalize_disconnect() # no-op after completed + assert run_store.read_run_record(w.record.run_id).status == "completed" +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_run_writer.py -v` +Expected: FAIL with module not found + +- [ ] **Step 3: Write the implementation** + +```python +# comfy_research/engine/runs/run_writer.py +"""RunWriter: persist one training run's lifecycle from its NDJSON event stream.""" +from __future__ import annotations + +import logging +from typing import Any, Iterator + +from comfy_research.engine.runs import run_index, run_store +from comfy_research.schemas.run_record import RunRecord, TERMINAL_STATUSES, now_ms + +logger = logging.getLogger(__name__) + +_HEARTBEAT_MIN_INTERVAL_MS = 1000.0 +_TERMINAL_EVENT_STATUS = {"complete": "completed", "aborted": "aborted", + "paused": "paused", "error": "failed"} + + +class RunWriter: + def __init__(self, record: RunRecord) -> None: + self._record = record + self._tracker = run_store.MetricsDeltaTracker() + self._last_heartbeat = 0.0 + run_store.write_run_record(record) + run_index.upsert_run(record) + + @property + def record(self) -> RunRecord: + return self._record + + @property + def is_terminal(self) -> bool: + return self._record.status in TERMINAL_STATUSES + + def _update(self, **changes: Any) -> None: + self._record = self._record.model_copy(update=changes) + run_store.write_run_record(self._record) + + def mark_running(self) -> None: + if self._record.status == "queued": + self._update(status="running", started_at=now_ms()) + run_index.upsert_run(self._record) + + def on_event(self, event: dict) -> None: + if self.is_terminal: + return + etype = str(event.get("type", "")) + rows = self._tracker.extract(event) + if rows: + try: + run_store.append_metric_rows(self._record.run_id, rows) + except OSError: + logger.warning("metric append failed for %s", self._record.run_id, exc_info=True) + now = now_ms() + if now - self._last_heartbeat >= _HEARTBEAT_MIN_INTERVAL_MS: + run_index.touch_heartbeat(self._record.run_id, now) + self._last_heartbeat = now + status = _TERMINAL_EVENT_STATUS.get(etype) + if status is None: + return + if etype in ("complete", "paused"): + run_store.write_results(self._record.run_id, dict(event)) + else: + run_store.write_results(self._record.run_id, self._tracker.snapshot()) + detail = str(event.get("detail", "")) if etype == "error" else "" + self.finalize(status, error_detail=detail) + + def finalize(self, status: str, error_detail: str = "") -> None: + if self.is_terminal: + return + self._update(status=status, finished_at=now_ms(), error_detail=error_detail) + snap = run_store.read_results(self._record.run_id) or self._tracker.snapshot() + run_index.upsert_run(self._record, run_store.summarize(snap)) + + def finalize_disconnect(self) -> None: + if not self.is_terminal: + run_store.write_results(self._record.run_id, self._tracker.snapshot()) + self.finalize("aborted") + + +def capture_events(events: Iterator[dict], writer: RunWriter) -> Iterator[dict]: + """Tee events into the writer; disconnect (generator close) finalizes as aborted.""" + writer.mark_running() + try: + for event in events: + try: + writer.on_event(event) + except Exception: + logger.warning("run capture failed for %s", writer.record.run_id, exc_info=True) + yield event + finally: + writer.finalize_disconnect() +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `python -m pytest comfy_research/tests/test_run_writer.py -v` +Expected: PASS (4 tests) + +- [ ] **Step 5: Commit** + +```bash +git add comfy_research/engine/runs/run_writer.py comfy_research/tests/test_run_writer.py +git commit -m "feat: add RunWriter facade with capture_events tee" +``` + +--- + +### Task 5: Capture in `/api/train` (local + remote) and TrainRequest run fields + +**Files:** +- Modify: `comfy_research/schemas/train_request.py` (add fields to `TrainRequest`) +- Modify: `comfy_research/api/train.py:359-463` (`post_train`) +- Test: `comfy_research/tests/test_train_run_capture.py` + +**Interfaces:** +- Consumes: `RunWriter`, `capture_events`, `RunRecord`, `new_run_id`, `now_ms`, `strip_result_data`, `flatten_hyperparams`, `GraphDocument`. +- Produces: `TrainRequest.run_origin: Literal["human","agent","sweep"] = "human"`, `run_group_id: str | None = None`, `run_parent_id: str | None = None`; helper `build_run_record(body: TrainRequest, *, origin=None, group_id=None, status="queued") -> RunRecord` in `run_writer.py`; `/api/train` emits `{"type": "run_registered", "run_id": ...}` before any training event — on the local classic path it is the FIRST event; on the remote path it comes after the bootstrap/`phase` events (which are emitted before validation completes); the CRL path emits none (v1 exclusion). +- Frontend stream parser: `frontend/src/graph/readNdjsonTrainStream.ts` routes unknown event types through its progress-callback fallback — add an explicit `run_registered` branch that ignores the event (or records the id for later use), so it never reaches the fallback. + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_train_run_capture.py +from __future__ import annotations + +import json + +from fastapi.testclient import TestClient + +from comfy_research.main import app +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.tests.train_test_fixtures import minimal_cpu_train_request + + +def _events(text: str) -> list[dict]: + return [json.loads(line) for line in text.splitlines() if line.strip()] + + +def test_post_train_registers_and_persists_run(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + body = {**minimal_cpu_train_request(), "run_origin": "agent", "run_group_id": "g1"} + response = TestClient(app).post("/api/train", json=body) + assert response.status_code == 200 + events = _events(response.text) + assert events[0]["type"] == "run_registered" + run_id = events[0]["run_id"] + assert any(e["type"] == "complete" for e in events) + + rec = run_store.read_run_record(run_id) + assert rec.status == "completed" + assert rec.origin == "agent" and rec.group_id == "g1" + assert rec.graph.nodes # config snapshot present + for node in rec.graph.nodes: + assert "memoryCheckpoint_b64" not in (node.data or {}) + res = run_store.read_results(run_id) + assert len(res["loss_history"]) == 4 and "checkpoint_b64" not in res + rows, _ = run_index.query_runs(RunQuery(ids=[run_id])) + assert rows[0]["status"] == "completed" and rows[0]["hyperparams"] + + +def test_post_train_default_origin_human(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + response = TestClient(app).post("/api/train", json=minimal_cpu_train_request()) + run_id = _events(response.text)[0]["run_id"] + assert run_store.read_run_record(run_id).origin == "human" +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_train_run_capture.py -v` +Expected: FAIL — first event is `progress`, not `run_registered` + +- [ ] **Step 3: Implement** + +In `comfy_research/schemas/train_request.py`, extend `TrainRequest`: + +```python +class TrainRequest(BaseModel): + trainer_node_id: str + nodes: list[Node] = Field(default_factory=list) + edges: list[Edge] = Field(default_factory=list) + resume: dict[str, Any] | None = None + hessian_oversized_policy: Literal["skip", "force"] | None = None + run_origin: Literal["human", "agent", "sweep"] = "human" + run_group_id: str | None = None + run_parent_id: str | None = None +``` + +Add to `comfy_research/engine/runs/run_writer.py` (imports: `TrainRequest`, `GraphDocument`, `new_run_id`, `strip_result_data`, `flatten_hyperparams` from `comfy_research.schemas.run_record`): + +```python +def build_run_record( + body: TrainRequest, + *, + origin: str | None = None, + group_id: str | None = None, + status: str = "queued", +) -> RunRecord: + nodes = strip_result_data(body.nodes) + trainer = next((n for n in body.nodes if n.id == body.trainer_node_id), None) + device = str((trainer.data or {}).get("computeDevice", "")) if trainer else "" + return RunRecord( + run_id=new_run_id(), + origin=origin or body.run_origin, + group_id=group_id if group_id is not None else body.run_group_id, + parent_id=body.run_parent_id, + status=status, + created_at=now_ms(), + trainer_node_id=body.trainer_node_id, + device=device, + graph=GraphDocument(version=1, nodes=nodes, edges=body.edges), + hyperparams=flatten_hyperparams(body.nodes), + ) +``` + +In `comfy_research/api/train.py`, add imports: + +```python +from comfy_research.engine.runs.run_writer import RunWriter, build_run_record, capture_events +``` + +Replace the local classic-path generator (currently lines 455-457): + +```python + writer = RunWriter(build_run_record(body)) + + def generate(): + yield _ndjson_encode({"type": "run_registered", "run_id": writer.record.run_id}) + for event in capture_events(iter_trainer_events_from_context(ctx), writer): + yield _ndjson_encode(event) +``` + +For the remote path, `iter_remote_train_stdout_lines` yields already-encoded NDJSON lines; wrap `generate_remote`'s final `yield from` (line 415): + +```python + writer = RunWriter(build_run_record(body)) + yield _ndjson_encode({"type": "run_registered", "run_id": writer.record.run_id}) + writer.mark_running() + try: + for raw in iter_remote_train_stdout_lines(body, config=cfg_holder["cfg"]): + try: + writer.on_event(json.loads(bytes(raw).decode("utf-8"))) + except Exception: + pass # unparseable remote line: forward but don't capture + yield raw + finally: + writer.finalize_disconnect() +``` + +Confirm `iter_remote_train_stdout_lines` yields `bytes` (check `comfy_research/remote/ssh.py`); if it yields `str`, drop the `.decode`. The CRL path (`generate_crl`) is left untouched (v1 exclusion). The `run_registered` event goes AFTER `prepare_trainer_run` so invalid graphs still 400 without creating a run record. + +Also update `frontend/src/graph/readNdjsonTrainStream.ts`: add a `run_registered` case that is a no-op (before the fallback branch that currently funnels unknown types into the progress callback), plus modify Task 5's files list to include it in the commit. + +- [ ] **Step 4: Run new and existing tests** + +Run: `python -m pytest comfy_research/tests/test_train_run_capture.py comfy_research/tests/test_train_api_integration.py -v` +Expected: the two new tests PASS. `test_post_train_streams_real_cpu_training_result` will FAIL on `assert [event["step"] for event in progress] == [0, 1, 2, 3]` only if `run_registered` broke its event parsing — it filters by type, so it should PASS; if any existing assertion trips over the new first event, update that assertion to skip `run_registered`. + +- [ ] **Step 5: Commit** + +```bash +git add comfy_research/schemas/train_request.py comfy_research/api/train.py \ + comfy_research/engine/runs/run_writer.py comfy_research/tests/test_train_run_capture.py +git commit -m "feat: capture /api/train runs (local and remote) into the run store" +``` + +--- + +### Task 6: Sweep inner-run capture + +**Files:** +- Modify: `comfy_research/engine/runs/train_sweep.py:234-323` (`iter_sweep_events`) +- Test: `comfy_research/tests/test_sweep_run_capture.py` + +**Interfaces:** +- Consumes: `RunWriter`, `build_run_record`, `capture_events`, `TrainRequest`. +- Produces: each non-CRL sweep combo persisted as its own run with `origin="sweep"`, `group_id=`; wrapper events unchanged. + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_sweep_run_capture.py +from __future__ import annotations + +import json + +from fastapi.testclient import TestClient + +from comfy_research.main import app +from comfy_research.engine.runs import run_index +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.tests.train_test_fixtures import minimal_cpu_train_request + + +def test_sweep_inner_runs_captured_with_group(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + base = minimal_cpu_train_request() + body = { + "sweep_session_id": "sweep-test-1", + "trainer_node_id": base["trainer_node_id"], + "nodes": base["nodes"], + "edges": base["edges"], + "axes": [{"node_id": base["trainer_node_id"], "data_path": "trainingSteps", + "values": [2, 3]}], + "metric": {"kind": "final_train_loss"}, + } + response = TestClient(app).post("/api/train/sweep", json=body) + assert response.status_code == 200 + events = [json.loads(l) for l in response.text.splitlines() if l.strip()] + assert any(e["type"] == "sweep_complete" for e in events) + + rows, _ = run_index.query_runs(RunQuery(group_id="sweep-test-1")) + assert len(rows) == 2 + assert all(r["origin"] == "sweep" and r["status"] == "completed" for r in rows) + steps = sorted(r["hyperparams"][f"{base['trainer_node_id']}.trainingSteps"] for r in rows) + assert steps == [2, 3] +``` + +Before finalizing the test, read `SweepAxis` (`train_sweep.py:26-45`) and copy its exact field names for the axes payload (`node_id` / `data_path` / `values` — adjust to the real names). + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_sweep_run_capture.py -v` +Expected: FAIL — `query_runs` returns 0 rows + +- [ ] **Step 3: Implement** + +In `train_sweep.py`, inside `iter_sweep_events`'s combo loop, wrap the classic (non-CRL) branch. Current code: + +```python + ctx_sup = prepare_trainer_run(...) + event_iter = iter_trainer_events_from_context(ctx_sup) +``` + +becomes: + +```python + ctx_sup = prepare_trainer_run( + nodes_p, + body.edges, + body.trainer_node_id, + resume=None, + hessian_oversized_policy="skip", + ) + inner_req = TrainRequest( + trainer_node_id=body.trainer_node_id, + nodes=nodes_p, + edges=body.edges, + ) + inner_writer = RunWriter( + build_run_record(inner_req, origin="sweep", group_id=session_id) + ) + event_iter = capture_events( + iter_trainer_events_from_context(ctx_sup), inner_writer + ) +``` + +Imports at top of `train_sweep.py`: + +```python +from comfy_research.engine.runs.run_writer import RunWriter, build_run_record, capture_events +from comfy_research.schemas.train_request import TrainRequest +``` + +The sweep loop `break`s out of `event_iter` on `complete`/`aborted`/`paused` without exhausting it; `capture_events` still finalizes because `on_event` marks terminal before the break, and generator GC triggers the `finally` for the error path. To make finalization deterministic (not GC-dependent), close explicitly: after the `for ev in event_iter:` loop add `event_iter.close()` (guard with `if hasattr(event_iter, "close")` since the CRL branch stays unwrapped). + +Apply the identical wrapping to `comfy_research/engine/runs/train_coordinate_descent.py` (~line 208), where `iter_trainer_events_from_context` is consumed the same way inside the tuning loop: build an `inner_req = TrainRequest(trainer_node_id=..., nodes=, edges=...)` with the actual local variable names at that site, use `origin="sweep"` and `group_id=`, and add the same `event_iter.close()` guard. If that file's inner consumption differs structurally (e.g. no reusable event iterator), capture only the sweep path in this task and record coordinate descent as a follow-up item in the plan's final commit message — do not leave it silently uncaptured. + +- [ ] **Step 4: Run tests** + +Run: `python -m pytest comfy_research/tests/test_sweep_run_capture.py comfy_research/tests/test_sweep_api.py -v` (if `test_sweep_api.py` doesn't exist, run `python -m pytest comfy_research/tests -k sweep -v`) +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add comfy_research/engine/runs/train_sweep.py comfy_research/tests/test_sweep_run_capture.py +git commit -m "feat: capture sweep inner runs into the run store with sweep session grouping" +``` + +--- + +### Task 7: Async-submit worker pool + +**Files:** +- Create: `comfy_research/engine/runs/run_worker.py` +- Test: `comfy_research/tests/test_run_worker.py` + +**Interfaces:** +- Consumes: `prepare_trainer_run`, `iter_trainer_events_from_context`, `request_abort` (`train_control.py`), `RunWriter`, `build_run_record`, `capture_events`, `TrainRequest`. +- Produces: `class RunWorkerPool`: `submit(body: TrainRequest, idempotency_key: str | None = None) -> RunRecord` (validates via `prepare_trainer_run(..., validate_only=True)` with the request's `resume` and `hessian_oversized_policy` so validation matches execution — raises `HTTPException` upward on bad graph; returns existing record for a known idempotency key), `abort(run_id: str) -> bool`, `shutdown(wait: bool = False) -> None`; module singleton `get_worker_pool() -> RunWorkerPool` and `reset_worker_pool_for_tests() -> None`. +- Same-trainer serialization WITHOUT occupying executor slots: a run whose `trainer_node_id` is already executing waits in a per-trainer FIFO (plain data, no thread); when the active run finishes, the pool submits the next queued one. A waiting run never blocks a worker thread, so different-trainer runs always get slots (no head-of-line blocking). +- Abort is `run_id`-scoped: a run waiting in the FIFO is removed and finalized `aborted` directly (no `request_abort`); `request_abort(trainer_node_id)` is only issued when the pool's currently-executing run for that trainer IS this `run_id` — never for a queued run, so it cannot cross-abort another run. (Caveat, documented: the underlying `train_control` registry is a single slot per trainer id, so a concurrent browser `/api/train` on the SAME trainer id could still receive the signal — pre-existing product semantics; and an abort landing before the training loop registers the trainer is lost, cooperative-abort semantics as today.) +- Submit rejects remote-GPU graphs (`HTTPException(400, {"code": "remote_not_supported", ...})`) — remote runs go through streaming `/api/train`. + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_run_worker.py +from __future__ import annotations + +import time + +import pytest +from fastapi import HTTPException + +from comfy_research.engine.runs import run_store +from comfy_research.engine.runs.run_worker import get_worker_pool, reset_worker_pool_for_tests +from comfy_research.schemas.train_request import TrainRequest +from comfy_research.tests.train_test_fixtures import minimal_cpu_train_request + + +def _wait_terminal(run_id: str, timeout_s: float = 30.0) -> str: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + rec = run_store.read_run_record(run_id) + if rec and rec.status not in ("queued", "running"): + return rec.status + time.sleep(0.05) + raise AssertionError(f"run {run_id} never reached terminal state") + + +@pytest.fixture(autouse=True) +def _isolated(tmp_path, monkeypatch): + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + reset_worker_pool_for_tests() + yield + reset_worker_pool_for_tests() + + +def test_submit_runs_to_completion() -> None: + body = TrainRequest.model_validate({**minimal_cpu_train_request(), "run_origin": "agent"}) + rec = get_worker_pool().submit(body) + assert rec.status == "queued" and rec.origin == "agent" + assert _wait_terminal(rec.run_id) == "completed" + assert len(run_store.read_results(rec.run_id)["loss_history"]) == 4 + + +def test_submit_invalid_graph_raises_400_and_persists_nothing(tmp_path) -> None: + body = TrainRequest.model_validate({ + "trainer_node_id": "trainer", + "nodes": [{"id": "trainer", "type": "trainer", + "data": {"trainingSteps": 1, "computeDevice": "cpu"}}], + "edges": [], + }) + with pytest.raises(HTTPException) as exc: + get_worker_pool().submit(body) + assert exc.value.status_code == 400 + assert not list(tmp_path.glob("run-*")) + + +def test_idempotency_key_returns_same_run() -> None: + body = TrainRequest.model_validate(minimal_cpu_train_request()) + a = get_worker_pool().submit(body, idempotency_key="k1") + b = get_worker_pool().submit(body, idempotency_key="k1") + assert a.run_id == b.run_id + _wait_terminal(a.run_id) + + +def test_same_trainer_runs_serialize() -> None: + pool = get_worker_pool() + body = TrainRequest.model_validate(minimal_cpu_train_request()) + a = pool.submit(body) + b = pool.submit(body) + assert _wait_terminal(a.run_id) == "completed" + assert _wait_terminal(b.run_id) == "completed" + + +def test_abort_waiting_run_is_scoped_to_that_run() -> None: + pool = get_worker_pool() + slow = minimal_cpu_train_request() + for n in slow["nodes"]: + if n["id"] == slow["trainer_node_id"]: + n["data"] = {**n["data"], "trainingSteps": 2000} + a = pool.submit(TrainRequest.model_validate(slow)) + b = pool.submit(TrainRequest.model_validate(slow)) # waits in FIFO behind a (same trainer) + assert pool.abort(b.run_id) is True + # b dies immediately, without any train_control signal that could hit a + assert run_store.read_run_record(b.run_id).status == "aborted" + assert run_store.read_run_record(a.run_id).status in ("queued", "running") + # abort the running run; cooperative signal can land before the training loop + # registers the trainer, so retry until it takes effect + deadline = time.monotonic() + 30.0 + while time.monotonic() < deadline: + pool.abort(a.run_id) + rec = run_store.read_run_record(a.run_id) + if rec.status not in ("queued", "running"): + break + time.sleep(0.1) + assert run_store.read_run_record(a.run_id).status == "aborted" + assert pool.abort("run-nonexistent0") is False + + +def _renamed_fixture(suffix: str) -> dict: + """Same minimal graph under fresh ids, so it counts as a different trainer.""" + req = minimal_cpu_train_request() + for n in req["nodes"]: + n["id"] = n["id"] + suffix + for e in req["edges"]: + e["id"] = e["id"] + suffix + e["source"] = e["source"] + suffix + e["target"] = e["target"] + suffix + req["trainer_node_id"] = req["trainer_node_id"] + suffix + return req + + +def test_different_trainer_not_blocked_by_same_trainer_queue() -> None: + pool = get_worker_pool() + slow = minimal_cpu_train_request() + for n in slow["nodes"]: + if n["id"] == slow["trainer_node_id"]: + n["data"] = {**n["data"], "trainingSteps": 2000} + a = pool.submit(TrainRequest.model_validate(slow)) + b = pool.submit(TrainRequest.model_validate(slow)) # waits in FIFO, occupies no slot + fast = pool.submit(TrainRequest.model_validate(_renamed_fixture("-x"))) + # with head-of-line blocking, `fast` would be stuck behind b in the 2-slot pool + assert _wait_terminal(fast.run_id) == "completed" + for rid in (a.run_id, b.run_id): + pool.abort(rid) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_run_worker.py -v` +Expected: FAIL with module not found + +- [ ] **Step 3: Write the implementation** + +```python +# comfy_research/engine/runs/run_worker.py +"""Server-owned worker pool: async-submitted runs execute detached from any HTTP stream.""" +from __future__ import annotations + +import logging +import threading +from collections import deque +from concurrent.futures import ThreadPoolExecutor + +from fastapi import HTTPException + +from comfy_research.engine.runs.run_writer import RunWriter, build_run_record, capture_events +from comfy_research.engine.runs.trainer_run import ( + iter_trainer_events_from_context, + prepare_trainer_run, +) +from comfy_research.engine.runs.train_control import request_abort +from comfy_research.engine.runs.ai4science_alias import remap_ai4science_node_types +from comfy_research.schemas.train_request import TrainRequest + +logger = logging.getLogger(__name__) + +_DEFAULT_SLOTS = 2 + + +def _prefers_remote_gpu(body: TrainRequest) -> bool: + node = next((n for n in body.nodes if n.id == body.trainer_node_id), None) + data = getattr(node, "data", None) or {} + spec = str(data.get("computeDevice", "")).strip().lower() + return (spec == "cuda" or spec.startswith("cuda:")) and data.get("remoteGpu") is True + + +class RunWorkerPool: + def __init__(self, slots: int = _DEFAULT_SLOTS) -> None: + self._executor = ThreadPoolExecutor(max_workers=slots, thread_name_prefix="run-worker") + self._lock = threading.Lock() + self._active: dict[str, str] = {} # trainer_node_id -> run_id currently executing + self._waiting: dict[str, deque[str]] = {} # trainer_node_id -> queued run_ids (FIFO) + self._writers: dict[str, RunWriter] = {} + self._idempotency: dict[str, str] = {} + self._records: dict[str, TrainRequest] = {} + + def submit(self, body: TrainRequest, idempotency_key: str | None = None): + if _prefers_remote_gpu(body): + raise HTTPException(status_code=400, detail={ + "code": "remote_not_supported", + "detail": "Async submit runs locally only; use streaming POST /api/train for remote GPU runs.", + }) + with self._lock: + if idempotency_key and idempotency_key in self._idempotency: + run_id = self._idempotency[idempotency_key] + return self._writers[run_id].record + mapped = remap_ai4science_node_types(body.nodes) + prepare_trainer_run( + mapped, body.edges, body.trainer_node_id, + resume=body.resume, + hessian_oversized_policy=body.hessian_oversized_policy, + validate_only=True, + ) + writer = RunWriter(build_run_record(body)) + run_id = writer.record.run_id + trainer_id = body.trainer_node_id + with self._lock: + self._writers[run_id] = writer + self._records[run_id] = body + if idempotency_key: + self._idempotency[idempotency_key] = run_id + if trainer_id in self._active: + # Same trainer already executing: wait in FIFO (no thread blocked), + # because the train_control registry is a single slot per trainer id. + self._waiting.setdefault(trainer_id, deque()).append(run_id) + return writer.record + self._active[trainer_id] = run_id + self._executor.submit(self._execute, run_id) + return writer.record + + def _execute(self, run_id: str) -> None: + with self._lock: + writer = self._writers.get(run_id) + body = self._records.get(run_id) + if writer is None or body is None: + return + try: + if not writer.is_terminal: + mapped = remap_ai4science_node_types(body.nodes) + ctx = prepare_trainer_run( + mapped, body.edges, body.trainer_node_id, + resume=body.resume, + hessian_oversized_policy=body.hessian_oversized_policy, + ) + for _ in capture_events(iter_trainer_events_from_context(ctx), writer): + pass + except HTTPException as exc: + writer.finalize("failed", error_detail=str(exc.detail)) + except Exception as exc: + logger.warning("submitted run %s crashed", run_id, exc_info=True) + writer.finalize("failed", error_detail=f"{type(exc).__name__}: {exc}") + finally: + self._dispatch_next(body.trainer_node_id) + + def _dispatch_next(self, trainer_id: str) -> None: + with self._lock: + queue = self._waiting.get(trainer_id) + next_id = None + while queue: + candidate = queue.popleft() + w = self._writers.get(candidate) + if w is not None and not w.is_terminal: + next_id = candidate + break + if next_id is None: + self._active.pop(trainer_id, None) + if queue is not None and not queue: + self._waiting.pop(trainer_id, None) + return + self._active[trainer_id] = next_id + self._executor.submit(self._execute, next_id) + + def abort(self, run_id: str) -> bool: + with self._lock: + writer = self._writers.get(run_id) + body = self._records.get(run_id) + if writer is None or body is None or writer.is_terminal: + return False + trainer_id = body.trainer_node_id + queue = self._waiting.get(trainer_id) + if queue is not None and run_id in queue: + queue.remove(run_id) # waiting run: finalize directly, no train_control signal + writer.finalize("aborted") + return True + is_running_here = self._active.get(trainer_id) == run_id + if not is_running_here: + return False + request_abort(trainer_id) # only when THIS run holds the trainer; cooperative + return True + + def shutdown(self, wait: bool = False) -> None: + self._executor.shutdown(wait=wait, cancel_futures=True) + + +_pool: RunWorkerPool | None = None +_pool_lock = threading.Lock() + + +def get_worker_pool() -> RunWorkerPool: + global _pool + with _pool_lock: + if _pool is None: + _pool = RunWorkerPool() + return _pool + + +def reset_worker_pool_for_tests() -> None: + global _pool + with _pool_lock: + if _pool is not None: + _pool.shutdown(wait=True) + _pool = None +``` + +Known v1 limits (documented, accepted): idempotency map is in-memory (lost on restart); an abort signal landing before the training loop registers the trainer is lost (cooperative abort — same `/api/train/control` semantics; the API test retries); a concurrent browser `/api/train` on the same trainer id shares the single `train_control` slot with the pool's active run (pre-existing product behavior). + +- [ ] **Step 4: Run test to verify it passes** + +Run: `python -m pytest comfy_research/tests/test_run_worker.py -v` +Expected: PASS (6 tests; the abort test may take ~10 s) + +- [ ] **Step 5: Commit** + +```bash +git add comfy_research/engine/runs/run_worker.py comfy_research/tests/test_run_worker.py +git commit -m "feat: add async-submit run worker pool with per-trainer serialization" +``` + +--- + +### Task 8: `/api/runs` router + startup wiring + +**Files:** +- Create: `comfy_research/api/runs.py` +- Modify: `comfy_research/main.py` (router registration ~line 104-125; startup reconciliation) +- Test: `comfy_research/tests/test_runs_api.py` + +**Interfaces:** +- Consumes: Tasks 1–4, 7 (`run_store`, `run_index`, `RunQuery`, `get_worker_pool`). +- Produces routes (all JSON): + - `POST /api/runs` — body `TrainRequest`; 202 `{"run_id", "status"}`; honors `Idempotency-Key` header. + - `GET /api/runs` — query params `status, origin, group_id, since, ids (comma-sep), order_by, limit, cursor`, plus any `hyperparam.=`; returns `{"runs": [...], "next_cursor": str | null}`. + - `GET /api/runs/groups` — `{"groups": [...]}` (registered before the `/{run_id}` route). + - `GET /api/runs/{run_id}` — RunRecord JSON + `"summary"`; 404 `{"code": "run_not_found", ...}`. + - `GET /api/runs/{run_id}/metrics` — `{"source": "results"|"ndjson", "data": {...}, "downsampled": bool}`; `?downsample=N` keeps every ceil(len/N)-th point of each series. + - `POST /api/runs/{run_id}/abort` — `{"ok": bool}`; 404 if unknown to the pool. + - `DELETE /api/runs/{run_id}` — 409 `{"code": "run_active", ...}` for `queued`/`running`; else removes dir + index row. + - `DELETE /api/runs` — same query filters; 400 `{"code": "filter_required", ...}` if no filter given; deletes terminal-state matches only; returns `{"deleted": int}`. + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_runs_api.py +from __future__ import annotations + +import time + +import pytest +from fastapi.testclient import TestClient + +from comfy_research.main import app +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_worker import reset_worker_pool_for_tests +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms +from comfy_research.tests.train_test_fixtures import minimal_cpu_train_request + + +@pytest.fixture(autouse=True) +def _isolated(tmp_path, monkeypatch): + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + reset_worker_pool_for_tests() + yield + reset_worker_pool_for_tests() + + +def _seed_terminal(status: str = "completed", group: str | None = None) -> str: + rec = RunRecord( + run_id=new_run_id(), origin="agent", status=status, created_at=now_ms(), + finished_at=now_ms(), group_id=group, trainer_node_id="t1", + graph=GraphDocument(version=1, nodes=[], edges=[]), + hyperparams={"t1.trainingSteps": 4}, + ) + run_store.write_run_record(rec) + run_store.write_results(rec.run_id, {"loss_history": [1.0, 0.5], + "test_loss_history": [], "step_ticks": [0, 1]}) + run_index.upsert_run(rec, {"final_loss": 0.5, "final_test_loss": None, + "best_test_loss": None, "steps_completed": 2}) + return rec.run_id + + +def _wait_terminal(client: TestClient, run_id: str, timeout_s: float = 30.0) -> str: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + status = client.get(f"/api/runs/{run_id}").json()["status"] + if status not in ("queued", "running"): + return status + time.sleep(0.05) + raise AssertionError("never terminal") + + +def test_submit_then_poll_lifecycle() -> None: + client = TestClient(app) + resp = client.post("/api/runs", json={**minimal_cpu_train_request(), "run_origin": "agent"}) + assert resp.status_code == 202 + run_id = resp.json()["run_id"] + assert _wait_terminal(client, run_id) == "completed" + metrics = client.get(f"/api/runs/{run_id}/metrics").json() + assert metrics["source"] == "results" + assert len(metrics["data"]["loss_history"]) == 4 + + +def test_submit_invalid_graph_400() -> None: + resp = TestClient(app).post("/api/runs", json={ + "trainer_node_id": "trainer", + "nodes": [{"id": "trainer", "type": "trainer", + "data": {"trainingSteps": 1, "computeDevice": "cpu"}}], + "edges": [], + }) + assert resp.status_code == 400 + + +def test_list_filters_and_hyperparam() -> None: + client = TestClient(app) + a = _seed_terminal("completed", group="g1") + _seed_terminal("failed", group="g1") + body = client.get("/api/runs", params={"status": "completed"}).json() + assert [r["run_id"] for r in body["runs"]] == [a] + body = client.get("/api/runs", params={"hyperparam.t1.trainingSteps": "4"}).json() + assert len(body["runs"]) == 2 + groups = client.get("/api/runs/groups").json()["groups"] + assert groups[0]["group_id"] == "g1" and groups[0]["counts"]["failed"] == 1 + + +def test_get_missing_run_404_with_code() -> None: + resp = TestClient(app).get("/api/runs/run-doesnotexist") + assert resp.status_code == 404 + assert resp.json()["detail"]["code"] == "run_not_found" + + +def test_delete_guards() -> None: + client = TestClient(app) + rid = _seed_terminal("completed") + assert client.delete("/api/runs").status_code == 400 + running = RunRecord( + run_id=new_run_id(), origin="agent", status="running", created_at=now_ms(), + trainer_node_id="t1", graph=GraphDocument(version=1, nodes=[], edges=[]), + ) + run_store.write_run_record(running) + run_index.upsert_run(running) + run_index.touch_heartbeat(running.run_id, now_ms()) + assert client.delete(f"/api/runs/{running.run_id}").status_code == 409 + assert client.delete(f"/api/runs/{rid}").status_code == 200 + assert run_store.read_run_record(rid) is None + resp = client.delete("/api/runs", params={"status": "failed"}) + assert resp.status_code == 200 and resp.json()["deleted"] == 0 + + +def test_bulk_delete_by_group() -> None: + client = TestClient(app) + _seed_terminal("completed", group="g2") + _seed_terminal("aborted", group="g2") + resp = client.delete("/api/runs", params={"group_id": "g2"}) + assert resp.json()["deleted"] == 2 +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_runs_api.py -v` +Expected: FAIL — 404s on every `/api/runs` route + +- [ ] **Step 3: Write the implementation** + +```python +# comfy_research/api/runs.py +"""Run store API: async submit, query, metrics, abort, delete. Files are truth.""" +from __future__ import annotations + +import shutil + +from fastapi import APIRouter, Header, HTTPException, Request + +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.engine.runs.run_worker import get_worker_pool +from comfy_research.schemas.run_record import TERMINAL_STATUSES +from comfy_research.schemas.train_request import TrainRequest + +router = APIRouter(prefix="/api/runs", tags=["runs"]) + +_HYPERPARAM_PREFIX = "hyperparam." + + +def _err(status: int, code: str, detail: str) -> HTTPException: + return HTTPException(status_code=status, detail={"code": code, "detail": detail}) + + +def _query_from_request(request: Request) -> RunQuery: + p = request.query_params + hyper = {k[len(_HYPERPARAM_PREFIX):]: v for k, v in p.items() + if k.startswith(_HYPERPARAM_PREFIX)} + ids = [s for s in (p.get("ids") or "").split(",") if s] or None + return RunQuery( + status=p.get("status"), origin=p.get("origin"), group_id=p.get("group_id"), + since_ms=float(p["since"]) if p.get("since") else None, + ids=ids, hyperparams=hyper, + order_by=p.get("order_by") or "-created_at", + limit=int(p.get("limit") or 100), cursor=p.get("cursor"), + ) + + +@router.post("", status_code=202) +def submit_run(body: TrainRequest, + idempotency_key: str | None = Header(default=None)) -> dict: + rec = get_worker_pool().submit(body, idempotency_key=idempotency_key) + return {"run_id": rec.run_id, "status": rec.status} + + +@router.get("") +def list_runs(request: Request) -> dict: + rows, next_cursor = run_index.query_runs(_query_from_request(request)) + return {"runs": rows, "next_cursor": next_cursor} + + +@router.get("/groups") +def list_groups() -> dict: + return {"groups": run_index.group_summary()} + + +@router.get("/{run_id}") +def get_run(run_id: str) -> dict: + rec = run_store.read_run_record(run_id) + if rec is None: + if (run_store.run_dir(run_id) / "run.json").exists(): + # file present but unparseable: surface it, don't 404 (spec: unreadable) + return {"run_id": run_id, "status": "unreadable", "summary": None} + raise _err(404, "run_not_found", f"No run {run_id!r}.") + _, series = run_store.load_series(run_id) + return {**rec.model_dump(mode="json"), "summary": run_store.summarize(series)} + + +@router.get("/{run_id}/metrics") +def get_run_metrics(run_id: str, downsample: int | None = None) -> dict: + if run_store.read_run_record(run_id) is None: + raise _err(404, "run_not_found", f"No run {run_id!r}.") + source, data = run_store.load_series(run_id) + downsampled = False + if downsample and downsample > 0: + out = {} + for key, series in data.items(): + if isinstance(series, list) and len(series) > downsample: + stride = -(-len(series) // downsample) + out[key] = series[::stride] + downsampled = True + else: + out[key] = series + data = out + return {"source": source, "data": data, "downsampled": downsampled} + + +@router.post("/{run_id}/abort") +def abort_run(run_id: str) -> dict: + ok = get_worker_pool().abort(run_id) + if not ok: + rec = run_store.read_run_record(run_id) + if rec is None: + raise _err(404, "run_not_found", f"No run {run_id!r}.") + return {"ok": False} + return {"ok": True} + + +def _delete_run_files(run_id: str) -> None: + shutil.rmtree(run_store.run_dir(run_id), ignore_errors=True) + run_index.delete_rows([run_id]) + + +_DELETABLE_STATUSES = TERMINAL_STATUSES | {"unreadable"} + + +@router.delete("/{run_id}") +def delete_run(run_id: str) -> dict: + rec = run_store.read_run_record(run_id) + if rec is None: + if (run_store.run_dir(run_id) / "run.json").exists(): + _delete_run_files(run_id) # unreadable: deletable, that's the point + return {"ok": True} + raise _err(404, "run_not_found", f"No run {run_id!r}.") + if rec.status not in _DELETABLE_STATUSES: + raise _err(409, "run_active", "Abort the run before deleting it.") + _delete_run_files(run_id) + return {"ok": True} + + +@router.delete("") +def bulk_delete(request: Request) -> dict: + q = _query_from_request(request) + if not any([q.status, q.origin, q.group_id, q.since_ms, q.ids, q.hyperparams]): + raise _err(400, "filter_required", + "Bulk delete requires at least one filter (status/origin/group_id/since/ids).") + # Collect ALL matching ids first (paginate to the end), then delete the deletable + # ones — deleting while paginating would shift the cursor and skip matches. + q.limit = 500 + to_delete: list[str] = [] + while True: + rows, cursor = run_index.query_runs(q) + to_delete.extend(r["run_id"] for r in rows if r["status"] in _DELETABLE_STATUSES) + if cursor is None: + break + q.cursor = cursor + for run_id in to_delete: + _delete_run_files(run_id) + return {"deleted": len(to_delete)} +``` + +In `comfy_research/main.py`: `from comfy_research.api.runs import router as runs_router` and `app.include_router(runs_router)` alongside the existing includes. Startup work goes into the EXISTING `_lifespan` function (`main.py:70` — read it first and add to it; do NOT run this at `create_app()` time, which executes on module import before tests can monkeypatch `COMFYRESEARCH_RUNS_DIR`): + +```python + # inside _lifespan, before the yield, alongside the existing startup work + try: + from comfy_research.engine.runs.run_index import reconcile_stale_running + reconcile_stale_running() + except Exception: # never block startup on store recovery + logging.getLogger(__name__).warning("run store reconciliation failed", exc_info=True) +``` + +Ensure `logging` is imported in `main.py` (add the import if absent). `TestClient(app)` runs the lifespan when used as a context manager; plain `TestClient(app).get(...)` calls skip it, which is fine — reconciliation is separately unit-tested in Task 3, so API tests here don't depend on lifespan execution. + +(Route-order note: `/groups` is declared before `/{run_id}` — FastAPI matches in declaration order. Also verify the SPA catch-all `GET /{full_path:path}` in `main.py` stays registered AFTER the new router — routers are added in `create_app()` before the catch-all, and the catch-all already 404s paths starting with `api`.) + +- [ ] **Step 4: Run test to verify it passes** + +Run: `python -m pytest comfy_research/tests/test_runs_api.py -v` +Expected: PASS (6 tests) + +- [ ] **Step 5: Check for a route-contract test** + +Run: `grep -rl "openapi()" comfy_research/tests/ tests/ 2>/dev/null` +If any test asserts on the OpenAPI path list, run it; if it pins exact routes, add the new `/api/runs*` paths to its expectation. If the grep finds nothing, skip this step. + +- [ ] **Step 6: Commit** + +```bash +git add comfy_research/api/runs.py comfy_research/main.py comfy_research/tests/test_runs_api.py +git commit -m "feat: add /api/runs router with async submit, query, metrics, and deletes" +``` + +--- + +### Task 9: GC / retention + +**Files:** +- Create: `comfy_research/engine/runs/run_gc.py` +- Modify: `comfy_research/main.py` (call after reconciliation), `comfy_research/engine/runs/run_worker.py` (`_execute` tail) +- Test: `comfy_research/tests/test_run_gc.py` + +**Interfaces:** +- Consumes: `run_store`, `run_index`, `RunQuery`, `TERMINAL_STATUSES`. +- Produces: `load_gc_config() -> dict` (reads optional `data/runs/config.json`; defaults `{"max_runs_agent": 2000, "max_age_days_agent": None, "worker_slots": 2}`), `run_gc_once() -> list[str]` (returns pruned run ids; only `origin="agent"`, terminal status, `finished_at` older than a 10-minute grace period; oldest-first beyond `max_runs_agent`; logs every pruned id). + +- [ ] **Step 1: Write the failing test** + +```python +# comfy_research/tests/test_run_gc.py +from __future__ import annotations + +from comfy_research.engine.runs import run_gc, run_index, run_store +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms + + +def _seed(status: str, origin: str, finished_ms_ago: float) -> str: + rec = RunRecord( + run_id=new_run_id(), origin=origin, status=status, + created_at=now_ms() - finished_ms_ago - 1000, + finished_at=(now_ms() - finished_ms_ago) if status != "running" else None, + trainer_node_id="t1", graph=GraphDocument(version=1, nodes=[], edges=[]), + ) + run_store.write_run_record(rec) + run_index.upsert_run(rec) + return rec.run_id + + +def test_gc_prunes_only_old_terminal_agent_runs(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + (tmp_path / "config.json").write_text('{"max_runs_agent": 2}', encoding="utf-8") + hour = 3_600_000.0 + keep_human = _seed("completed", "human", 10 * hour) + keep_recent = _seed("completed", "agent", 0.0) # inside grace period + keep_running = _seed("running", "agent", 5 * hour) + newest = _seed("completed", "agent", 1 * hour) + older = _seed("completed", "agent", 2 * hour) + oldest = _seed("completed", "agent", 3 * hour) + + pruned = run_gc.run_gc_once() + # agent rows newest-first by created_at: [keep_recent, newest, older, oldest, keep_running]; + # cap 2 keeps [keep_recent, newest]; over-cap = [older, oldest, keep_running], of which + # only terminal runs past the 10-minute grace period are prunable. + assert set(pruned) == {older, oldest} + for rid in (keep_human, keep_recent, keep_running, newest): + assert run_store.read_run_record(rid) is not None + for rid in pruned: + assert run_store.read_run_record(rid) is None + + +def test_gc_default_config_noop_under_cap(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + _seed("completed", "agent", 3_600_000.0) + assert run_gc.run_gc_once() == [] +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `python -m pytest comfy_research/tests/test_run_gc.py -v` +Expected: FAIL with module not found + +- [ ] **Step 3: Write the implementation** + +```python +# comfy_research/engine/runs/run_gc.py +"""Single-owner retention GC for agent-origin runs. Runs only in the API server process.""" +from __future__ import annotations + +import json +import logging +import shutil + +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.schemas.run_record import TERMINAL_STATUSES, now_ms + +logger = logging.getLogger(__name__) + +_GRACE_MS = 10 * 60 * 1000.0 +_DEFAULTS = {"max_runs_agent": 2000, "max_age_days_agent": None, "worker_slots": 2} + + +def load_gc_config() -> dict: + path = run_store.runs_root() / "config.json" + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + raw = {} + return {**_DEFAULTS, **{k: raw[k] for k in _DEFAULTS if k in raw}} + + +def run_gc_once() -> list[str]: + cfg = load_gc_config() + rows: list[dict] = [] + q = RunQuery(origin="agent", order_by="-created_at", limit=500) + while True: + page, cursor = run_index.query_runs(q) + rows.extend(page) + if cursor is None: + break + q.cursor = cursor + cutoff_ms = now_ms() - _GRACE_MS + max_age = cfg["max_age_days_agent"] + age_cutoff = now_ms() - max_age * 86_400_000.0 if max_age else None + + def prunable(r: dict) -> bool: + return (r["status"] in TERMINAL_STATUSES + and r["finished_at"] is not None and r["finished_at"] < cutoff_ms) + + pruned: list[str] = [] + over_cap = rows[cfg["max_runs_agent"]:] # rows are newest-first + for r in over_cap: + if prunable(r): + pruned.append(r["run_id"]) + if age_cutoff is not None: + for r in rows[: cfg["max_runs_agent"]]: + if prunable(r) and r["finished_at"] < age_cutoff: + pruned.append(r["run_id"]) + for run_id in pruned: + shutil.rmtree(run_store.run_dir(run_id), ignore_errors=True) + run_index.delete_rows(pruned) + for run_id in pruned: + logger.info("run GC pruned %s", run_id) + return pruned +``` + +Wire-up: in `main.py`, call `run_gc_once()` inside the same guarded `try` as reconciliation. In `run_worker.py`, at the very end of `_execute` (after the trainer lock releases), add: + +```python + try: + from comfy_research.engine.runs.run_gc import run_gc_once + run_gc_once() + except Exception: + logger.warning("post-run GC failed", exc_info=True) +``` + +Also in `run_worker.py`, size the pool from config: in `get_worker_pool()`, `RunWorkerPool(slots=load_gc_config()["worker_slots"])`. + +- [ ] **Step 4: Run tests** + +Run: `python -m pytest comfy_research/tests/test_run_gc.py comfy_research/tests/test_run_worker.py -v` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add comfy_research/engine/runs/run_gc.py comfy_research/engine/runs/run_worker.py \ + comfy_research/main.py comfy_research/tests/test_run_gc.py +git commit -m "feat: add single-owner run GC with agent retention policy" +``` + +--- + +### Task 10: Frontend runs API client + rail panel (list) + +**Files:** +- Create: `frontend/src/graph/runsApi.ts`, `frontend/src/components/RunsPanel.tsx` +- Modify: `frontend/src/components/railTypes.ts`, `frontend/src/components/LeftNavRail.tsx`, `frontend/src/components/ResearchCanvas.tsx` (~line 5996-6048 slot block), `frontend/src/index.css` (panel styles), `frontend/src/graph/__tests__/leftNavRail.v1.seam.test.tsx` (labels list) +- Test: `frontend/src/graph/__tests__/runsPanel.test.tsx` + +**Interfaces:** +- Consumes: `GET /api/runs`, `GET /api/runs/{id}/metrics` (Task 8 shapes); `SweepVizLinePlot` + `PlotSeries` (Task 11 uses them; this task only lists). +- Produces: + - `runsApi.ts`: `type RunRow = { run_id: string; group_id: string | null; origin: string; status: string; created_at: number; finished_at: number | null; trainer_node_id: string; final_loss: number | null; final_test_loss: number | null; steps_completed: number | null; duration_seconds: number | null; hyperparams: Record }`, `fetchRuns(params?: Record): Promise<{ runs: RunRow[]; next_cursor: string | null }>`, `fetchRunMetrics(runId: string): Promise<{ source: string; data: Record }>`, `fetchRunRecord(runId: string): Promise<{ graph: GraphDocument } & Record>`, `deleteRun(runId: string): Promise`. + - `RunsPanel.tsx`: `export function RunsPanel(props: { onOpenRunGraph: (runId: string) => void })` — fetches on mount, polls every 3 s while any run has status `queued`/`running`, groups rows by `group_id` (ungrouped first), collapses `origin === "sweep" || origin === "agent"` groups by default, checkbox per row feeding selection state (used by Task 11), delete button per terminal row. + - Rail: `RailPrimarySection` gains `"runs"`; `LeftNavRail` `primaryItems` gains `{ id: "runs", label: "Runs" }` + a `case "runs"` icon; `ResearchCanvas` renders `{railSection === "runs" ?
: null}` (conditional-mount form, matching Templates). `openRunGraphInNewProject` is a stub in this task (`console.warn`); Task 11 implements it. + +- [ ] **Step 1: Write the failing test** + +```tsx +// frontend/src/graph/__tests__/runsPanel.test.tsx +// @vitest-environment jsdom +import { act } from "react"; +import { createRoot } from "react-dom/client"; +import { afterEach, beforeEach, expect, test, vi } from "vitest"; + +import { RunsPanel } from "../../components/RunsPanel"; + +(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true; + +const RUNS = { + runs: [ + { run_id: "run-aaa", group_id: null, origin: "human", status: "completed", + created_at: 1, finished_at: 2, trainer_node_id: "t1", final_loss: 0.25, + final_test_loss: null, steps_completed: 4, duration_seconds: 1.5, hyperparams: {} }, + { run_id: "run-bbb", group_id: "sweep-1", origin: "sweep", status: "failed", + created_at: 3, finished_at: 4, trainer_node_id: "t1", final_loss: null, + final_test_loss: null, steps_completed: 0, duration_seconds: null, hyperparams: {} }, + ], + next_cursor: null, +}; + +let host: HTMLDivElement; + +beforeEach(() => { + host = document.createElement("div"); + document.body.appendChild(host); + vi.stubGlobal("fetch", vi.fn(async () => ({ + ok: true, + json: async () => RUNS, + })) as unknown as typeof fetch); +}); + +afterEach(() => { + vi.unstubAllGlobals(); + host.remove(); +}); + +test("renders run rows grouped, sweep group collapsed by default", async () => { + const root = createRoot(host); + await act(async () => { + root.render( {}} />); + }); + await act(async () => { await Promise.resolve(); }); + expect(host.textContent).toContain("run-aaa"); + expect(host.textContent).toContain("sweep-1"); + // collapsed group hides its member row until expanded + expect(host.textContent).not.toContain("run-bbb"); + const toggle = host.querySelector('[data-testid="run-group-toggle-sweep-1"]'); + expect(toggle).not.toBeNull(); + await act(async () => { toggle!.click(); }); + expect(host.textContent).toContain("run-bbb"); + await act(async () => { root.unmount(); }); +}); +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd frontend && npx vitest run src/graph/__tests__/runsPanel.test.tsx` +Expected: FAIL — module `../../components/RunsPanel` not found + +- [ ] **Step 3: Implement** + +`frontend/src/graph/runsApi.ts`: + +```ts +import type { GraphDocument } from "../types/graph"; + +export type RunRow = { + run_id: string; + group_id: string | null; + origin: string; + status: string; + created_at: number; + finished_at: number | null; + trainer_node_id: string; + final_loss: number | null; + final_test_loss: number | null; + steps_completed: number | null; + duration_seconds: number | null; + hyperparams: Record; +}; + +async function readJson(res: Response): Promise { + if (!res.ok) { + const text = await res.text().catch(() => ""); + throw new Error(text || res.statusText); + } + return (await res.json()) as T; +} + +export async function fetchRuns( + params?: Record, +): Promise<{ runs: RunRow[]; next_cursor: string | null }> { + const qs = params ? `?${new URLSearchParams(params)}` : ""; + return readJson(await fetch(`/api/runs${qs}`, { cache: "no-store" })); +} + +export async function fetchRunMetrics( + runId: string, +): Promise<{ source: string; data: Record }> { + return readJson(await fetch(`/api/runs/${runId}/metrics`, { cache: "no-store" })); +} + +export async function fetchRunRecord( + runId: string, +): Promise<{ graph: GraphDocument } & Record> { + return readJson(await fetch(`/api/runs/${runId}`, { cache: "no-store" })); +} + +export async function deleteRun(runId: string): Promise { + await readJson(await fetch(`/api/runs/${runId}`, { method: "DELETE" })); +} +``` + +`frontend/src/components/RunsPanel.tsx` (structure; reuse `cr-nodes-panel*` chrome): + +```tsx +import { useCallback, useEffect, useMemo, useState } from "react"; + +import { deleteRun, fetchRuns, type RunRow } from "../graph/runsApi"; + +const ACTIVE_STATUSES = new Set(["queued", "running"]); +const POLL_MS = 3000; +const TERMINAL_DELETABLE = new Set(["completed", "failed", "aborted", "paused", "crashed"]); + +type RunGroup = { key: string; label: string; rows: RunRow[]; collapsedByDefault: boolean }; + +function groupRuns(rows: RunRow[]): RunGroup[] { + const ungrouped: RunRow[] = []; + const byGroup = new Map(); + for (const row of rows) { + if (row.group_id) { + const list = byGroup.get(row.group_id) ?? []; + list.push(row); + byGroup.set(row.group_id, list); + } else { + ungrouped.push(row); + } + } + const groups: RunGroup[] = []; + if (ungrouped.length) { + groups.push({ key: "", label: "Runs", rows: ungrouped, collapsedByDefault: false }); + } + for (const [key, groupRows] of byGroup) { + const collapsed = groupRows.every((r) => r.origin === "sweep" || r.origin === "agent"); + groups.push({ key, label: key, rows: groupRows, collapsedByDefault: collapsed }); + } + return groups; +} + +export function RunsPanel({ + onOpenRunGraph, + selectedRunIds, + onToggleSelect, +}: { + onOpenRunGraph: (runId: string) => void; + selectedRunIds?: ReadonlySet; + onToggleSelect?: (runId: string) => void; +}) { + const [rows, setRows] = useState([]); + const [error, setError] = useState(null); + const [expanded, setExpanded] = useState>({}); + + const load = useCallback(() => { + fetchRuns({ limit: "200" }) + .then((body) => { setRows(body.runs); setError(null); }) + .catch((e: Error) => setError(e.message)); + }, []); + + useEffect(() => { + load(); + }, [load]); + + const anyActive = useMemo(() => rows.some((r) => ACTIVE_STATUSES.has(r.status)), [rows]); + useEffect(() => { + if (!anyActive) return; + const id = window.setInterval(load, POLL_MS); + return () => window.clearInterval(id); + }, [anyActive, load]); + + const groups = useMemo(() => groupRuns(rows), [rows]); + + return ( + + ); +} +``` + +Rail wiring: +- `railTypes.ts`: `export type RailPrimarySection = "nodes" | "observables" | "templates" | "runs";` +- `LeftNavRail.tsx`: append `{ id: "runs", label: "Runs" }` to `primaryItems`; add `case "runs":` to `RailIcon` returning a simple inline SVG (e.g. three horizontal bars of differing length, `stroke="currentColor"`). +- `ResearchCanvas.tsx` (slot block ~line 5996): add after the observables slot: + +```tsx + {railSection === "runs" ? ( +
+ { console.warn("open run", runId); }} /> +
+ ) : null} +``` + +- `leftNavRail.v1.seam.test.tsx`: update `V1_RAIL_LABELS` to `["Nodes", "Observables", "Templates", "Runs"]` (this list is an explicit-decision seam; this plan is that decision). +- `index.css`: add `.cr-runs-panel__row { display: flex; gap: 6px; align-items: center; }`, `.cr-runs-panel__status-dot { width: 8px; height: 8px; border-radius: 50%; background: var(--cr-text-4); }`, `[data-status="completed"].cr-runs-panel__status-dot { background: var(--cr-chart-2); }`, `[data-status="failed"].cr-runs-panel__status-dot, [data-status="crashed"].cr-runs-panel__status-dot { background: var(--cr-chart-1); }`, `.cr-runs-panel__error { color: var(--cr-chart-1); }` — token vars only, no raw hex. + +- [ ] **Step 4: Run tests + hex ratchet** + +Run: `cd frontend && npx vitest run src/graph/__tests__/runsPanel.test.tsx src/graph/__tests__/leftNavRail.v1.seam.test.tsx && npm run verify:css-tokens` +Expected: PASS, ratchet clean + +- [ ] **Step 5: Commit** + +```bash +git add frontend/src/graph/runsApi.ts frontend/src/components/RunsPanel.tsx \ + frontend/src/components/railTypes.ts frontend/src/components/LeftNavRail.tsx \ + frontend/src/components/ResearchCanvas.tsx frontend/src/index.css \ + frontend/src/graph/__tests__/runsPanel.test.tsx \ + frontend/src/graph/__tests__/leftNavRail.v1.seam.test.tsx +git commit -m "feat: add Runs rail panel with grouped run list" +``` + +--- + +### Task 11: Overlay comparison chart + open-run-graph + +**Files:** +- Create: `frontend/src/graph/runCompareOverlay.ts` +- Modify: `frontend/src/components/RunsPanel.tsx` (selection state + chart), `frontend/src/components/ResearchCanvas.tsx` (`openRunGraphInNewProject`) +- Test: `frontend/src/graph/__tests__/runCompareOverlay.test.ts` + +**Interfaces:** +- Consumes: `PlotSeries`/`PlotPoint` + `SERIES_COLORS` from `frontend/src/graph/sweepVizPlot.ts`; `SweepVizLinePlot` from `frontend/src/components/nodes/SweepVizLinePlot.tsx`; `fetchRunMetrics`, `fetchRunRecord` (Task 10); `sanitizeLoadedGraph`, `newProjectId`-style project creation as done by `openSavedGraphInNewProject` (`ResearchCanvas.tsx:5761`). +- Produces: + - `buildRunCompareSeries(inputs: Array<{ runId: string; label: string; data: Record }>): PlotSeries[]` — one series per run from `loss_history` vs `step_ticks` (falls back to index when `step_ticks` shorter), plus a dashed series per run from `test_loss_history` when non-empty; colors cycle `SERIES_COLORS`; test series reuse their run's color with `strokeDasharray: "4 3"`. + - `openRunGraphInNewProject(runId: string)` in `ResearchCanvas` — fetches the record, `sanitizeLoadedGraph(record.graph)`, creates a new project tab titled `Run {runId}` (same shape as `openSavedGraphInNewProject` minus `librarySource`), activates it. + - `RunsPanel` moves selection to internal state: drop the Task-10 `selectedRunIds`/`onToggleSelect` props and hold `const [selectedRunIds, setSelectedRunIds] = useState>(new Set())` inside the component (update the Task-10 test if its props usage breaks). When ≥1 selected, fetch metrics per selected run (cache in a `useRef(Map)`), render ``. + +- [ ] **Step 1: Write the failing test** + +```ts +// frontend/src/graph/__tests__/runCompareOverlay.test.ts +import { expect, test } from "vitest"; + +import { buildRunCompareSeries } from "../runCompareOverlay"; + +test("one solid series per run, dashed test series, distinct colors", () => { + const series = buildRunCompareSeries([ + { runId: "run-a", label: "run-a", + data: { loss_history: [1, 0.5], test_loss_history: [2, 1], step_ticks: [0, 1] } }, + { runId: "run-b", label: "run-b", + data: { loss_history: [3, 2, 1], test_loss_history: [], step_ticks: [0, 1, 2] } }, + ]); + expect(series.map((s) => s.id)).toEqual(["run-a", "run-a:test", "run-b"]); + expect(series[0].points.map((p) => p.y)).toEqual([1, 0.5]); + expect(series[0].points.map((p) => p.x)).toEqual([0, 1]); + expect(series[1].strokeDasharray).toBeTruthy(); + expect(series[1].color).toBe(series[0].color); + expect(series[2].color).not.toBe(series[0].color); +}); + +test("falls back to index when step_ticks missing", () => { + const [s] = buildRunCompareSeries([ + { runId: "run-c", label: "run-c", data: { loss_history: [5, 4], step_ticks: [] } }, + ]); + expect(s.points.map((p) => p.x)).toEqual([0, 1]); +}); +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd frontend && npx vitest run src/graph/__tests__/runCompareOverlay.test.ts` +Expected: FAIL — module not found + +- [ ] **Step 3: Implement** + +```ts +// frontend/src/graph/runCompareOverlay.ts +import { SERIES_COLORS, type PlotPoint, type PlotSeries } from "./sweepVizPlot"; + +type RunSeriesInput = { runId: string; label: string; data: Record }; + +function toPoints(ys: number[], xs: number[], rowId: string): PlotPoint[] { + return ys + .map((y, i) => ({ x: xs[i] ?? i, xDisplay: String(xs[i] ?? i), y, rowId })) + .filter((p) => Number.isFinite(p.y)); +} + +export function buildRunCompareSeries(inputs: RunSeriesInput[]): PlotSeries[] { + const series: PlotSeries[] = []; + inputs.forEach((input, i) => { + const color = SERIES_COLORS[i % SERIES_COLORS.length]; + const steps = input.data.step_ticks ?? []; + const loss = input.data.loss_history ?? []; + series.push({ id: input.runId, label: input.label, color, + points: toPoints(loss, steps, input.runId) }); + const test = input.data.test_loss_history ?? []; + if (test.length) { + series.push({ id: `${input.runId}:test`, label: `${input.label} (test)`, color, + strokeDasharray: "4 3", points: toPoints(test, steps, input.runId) }); + } + }); + return series; +} +``` + +`SERIES_COLORS` in `sweepVizPlot.ts` (line ~109) is currently a NON-exported local constant — add `export` to it as part of this task (verify no name collision first; it's an array of `var(--cr-chart-N)` strings). In `RunsPanel.tsx`, add the selection + chart described in Interfaces. In `ResearchCanvas.tsx`, implement `openRunGraphInNewProject` as a `useCallback` next to `openSavedGraphInNewProject` (line ~5761), reusing its project-creation shape: + +```tsx + const openRunGraphInNewProject = useCallback((runId: string) => { + void fetchRunRecord(runId).then((record) => { + const { nodes, edges } = sanitizeLoadedGraph(record.graph); + const id = newProjectId(); + setProjects((list) => [ + ...list, + { + id, + title: formatProjectTabTitle(`Run ${runId}`), + canvas: { + id: newCanvasId(), + title: `Run ${runId}`, + nodes, + edges, + savedViewport: record.graph.viewport ?? null, + viewportApplyNonce: 1, + dirty: false, + }, + }, + ]); + setActiveProjectId(id); + }); + }, []); +``` + +Copy the exact project/canvas object shape from `openSavedGraphInNewProject` at implementation time (helper names `newProjectId` / `newCanvasId` / `formatProjectTabTitle` must match what that function actually uses; omit `librarySource`). Replace the Task 10 `console.warn` stub with this callback. + +- [ ] **Step 4: Run all frontend tests + build gate** + +Run: `cd frontend && npx vitest run && npm run verify:css-tokens` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add frontend/src/graph/runCompareOverlay.ts frontend/src/components/RunsPanel.tsx \ + frontend/src/components/ResearchCanvas.tsx frontend/src/graph/__tests__/runCompareOverlay.test.ts +git commit -m "feat: add run comparison overlay chart and open-run-graph action" +``` + +--- + +### Task 12: Docs + +**Files:** +- Create: `docs/en/reference/runs-api.md` +- Modify: `docs/en/reference/data-contracts.md` (add "Run store" section), `docs/en/reference/training-api.md` (mention `run_registered` event + `run_origin`/`run_group_id`/`run_parent_id` fields) + +**Interfaces:** documentation of Task 8's routes and Task 2's on-disk layout, exactly as implemented. + +- [ ] **Step 1: Write `docs/en/reference/runs-api.md`** + +Cover, with the same table style as `training-api.md`: the async submit contract (202, `Idempotency-Key`, remote-GPU 400), every `/api/runs*` route with query params, the status lifecycle (`queued → running → completed | failed | aborted | paused | crashed`, plus `crashed` reconciliation and `unreadable`), the metrics authority rule (`source` field), the error envelope `{"code", "detail"}`, cursor pagination, bulk-delete filter requirement, and the GC defaults (`data/runs/config.json`). Include one curl example: submit `minimal_cpu_train_request`-shaped JSON, poll `GET /api/runs/{id}`, fetch `GET /api/runs/{id}/metrics`. + +- [ ] **Step 2: Update `data-contracts.md`** + +Add a "Run store (`data/runs/`)" section: directory layout (`run.json` / `metrics.ndjson` / `results.json`), files-are-truth + rebuildable `index.db` (and `scripts` note: rebuild happens automatically on corruption; never hand-edit `index.db`), what is never persisted (checkpoints, PNGs, embedding/attention histories), `paused` = terminal-in-store with resume-as-new-run via `parent_id`. + +- [ ] **Step 3: Build docs** + +Run: `python scripts/build_docs.py` +Expected: build succeeds; fix any Sphinx warnings introduced by the new page (add it to the reference toctree — mirror how `training-api.md` is listed). + +- [ ] **Step 4: Full test sweep** + +Run: `python -m pytest comfy_research/tests -x -q -m "not repro and not slow" && cd frontend && npx vitest run` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add docs/en/reference/runs-api.md docs/en/reference/data-contracts.md docs/en/reference/training-api.md +git commit -m "docs: document the run store API and on-disk contracts" +``` From c94d8d4a3787df57d921f035020f3d23067c934d Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 17:34:38 +0800 Subject: [PATCH 03/34] feat: add RunRecord schema with result stripping and hyperparam flattening --- comfy_research/schemas/run_record.py | 90 +++++++++++++++++++++++++ comfy_research/tests/test_run_record.py | 76 +++++++++++++++++++++ 2 files changed, 166 insertions(+) create mode 100644 comfy_research/schemas/run_record.py create mode 100644 comfy_research/tests/test_run_record.py diff --git a/comfy_research/schemas/run_record.py b/comfy_research/schemas/run_record.py new file mode 100644 index 0000000..1715b8e --- /dev/null +++ b/comfy_research/schemas/run_record.py @@ -0,0 +1,90 @@ +"""RunRecord: persisted metadata for one training run (``data/runs/{run_id}/run.json``).""" +from __future__ import annotations + +import time +import uuid +from typing import Any, Literal + +from pydantic import BaseModel, Field + +from comfy_research.generated.node_manifest import load_node_manifest +from comfy_research.schemas.graph import GraphDocument, Node + +RunStatus = Literal[ + "queued", "running", "completed", "failed", "aborted", "paused", "crashed", "unreadable" +] +TERMINAL_STATUSES: frozenset[str] = frozenset( + {"completed", "failed", "aborted", "paused", "crashed"} +) + +# Result/UI payloads the browser stashes into node data; never persisted in run configs. +# Superset mirror of frontend/src/graph/graphFileExportTier.ts strip lists. +RUN_RESULT_DATA_KEYS: frozenset[str] = frozenset( + { + "checkpoint_b64", "memoryCheckpoint_b64", + "lossHistory", "testLossHistory", "regLossHistory", + "stepTicks", "epochTicks", "observableMetricHistories", + "lastTrainLoopSeconds", "plotPngBase64", "valueHistory", + "embeddingHistory", "attentionMapFrames", "previewGrid", + "histogramPng", "imageGrid", "runSummary", "lastError", + "lastSweepSummary", "observableEmbeddingHistories", + "observableAttentionSliceHistories", + } +) + +_HYPERPARAM_EXCLUDED_FIELDS = frozenset({"instanceTitle"}) + + +class RunRecord(BaseModel): + run_id: str + schema_version: int = 1 + group_id: str | None = None + parent_id: str | None = None + origin: Literal["human", "agent", "sweep"] = "human" + status: RunStatus = "queued" + created_at: float + started_at: float | None = None + finished_at: float | None = None + trainer_node_id: str + device: str = "" + error_detail: str = "" + graph: GraphDocument + hyperparams: dict[str, Any] = Field(default_factory=dict) + + +def now_ms() -> float: + return time.time() * 1000.0 + + +def new_run_id() -> str: + return "run-" + uuid.uuid4().hex[:12] + + +def strip_result_data(nodes: list[Node]) -> list[Node]: + """Copy nodes with result/UI blobs removed from ``data`` (config-only snapshot).""" + out: list[Node] = [] + for n in nodes: + data = {k: v for k, v in (n.data or {}).items() if k not in RUN_RESULT_DATA_KEYS} + out.append(n.model_copy(update={"data": data}, deep=True)) + return out + + +def _declared_field_keys(node_type: str) -> list[str]: + for entry in load_node_manifest(): + if entry.get("type") == node_type: + return [f["key"] for f in entry.get("fields", [])] + return [] + + +def flatten_hyperparams(nodes: list[Node]) -> dict[str, Any]: + """``{node_id}.{field}`` -> scalar, for manifest-declared fields present in node data.""" + flat: dict[str, Any] = {} + for n in nodes: + data = n.data or {} + for key in _declared_field_keys(n.type.value): + if key in _HYPERPARAM_EXCLUDED_FIELDS or key not in data: + continue + v = data[key] + if isinstance(v, (bool, int, float, str)): + flat[f"{n.id}.{key}"] = v + return flat diff --git a/comfy_research/tests/test_run_record.py b/comfy_research/tests/test_run_record.py new file mode 100644 index 0000000..e59c82c --- /dev/null +++ b/comfy_research/tests/test_run_record.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from comfy_research.schemas.graph import Edge, GraphDocument, Node +from comfy_research.schemas.run_record import ( + RUN_RESULT_DATA_KEYS, + RunRecord, + flatten_hyperparams, + new_run_id, + now_ms, + strip_result_data, +) + + +def _trainer_node() -> Node: + return Node( + id="t1", + type="trainer", + data={ + "trainingSteps": 4, + "computeDevice": "cpu", + "instanceTitle": "Trainer", + "lossHistory": [1.0, 0.5], + "memoryCheckpoint_b64": "QUJD", + "plotPngBase64": "aW1n", + }, + ) + + +def test_new_run_id_prefix_and_uniqueness() -> None: + a, b = new_run_id(), new_run_id() + assert a.startswith("run-") and len(a) == 16 + assert a != b + + +def test_strip_result_data_removes_blobs_keeps_config() -> None: + stripped = strip_result_data([_trainer_node()]) + data = stripped[0].data + assert data["trainingSteps"] == 4 + assert data["instanceTitle"] == "Trainer" + assert "lossHistory" not in data + assert "memoryCheckpoint_b64" not in data + assert "plotPngBase64" not in data + # original untouched + assert "lossHistory" in _trainer_node().data + + +def test_result_keys_cover_known_blobs() -> None: + for key in ("checkpoint_b64", "memoryCheckpoint_b64", "plotPngBase64", + "lossHistory", "testLossHistory", "regLossHistory", "stepTicks", + "observableMetricHistories", "embeddingHistory", + "attentionMapFrames", "valueHistory", "runSummary", "lastError"): + assert key in RUN_RESULT_DATA_KEYS + + +def test_flatten_hyperparams_declared_scalars_only() -> None: + flat = flatten_hyperparams([_trainer_node()]) + assert flat["t1.trainingSteps"] == 4 + assert flat["t1.computeDevice"] == "cpu" + assert "t1.lossHistory" not in flat # not a declared field + assert "t1.instanceTitle" not in flat # declared but excluded as label + + +def test_run_record_roundtrip() -> None: + rec = RunRecord( + run_id=new_run_id(), + origin="agent", + status="queued", + created_at=now_ms(), + trainer_node_id="t1", + graph=GraphDocument(version=1, nodes=strip_result_data([_trainer_node()]), edges=[]), + hyperparams=flatten_hyperparams([_trainer_node()]), + ) + again = RunRecord.model_validate(rec.model_dump(mode="json")) + assert again.run_id == rec.run_id + assert again.schema_version == 1 + assert again.group_id is None and again.finished_at is None From 6796fbf6ca7c13517c0a948f87a44461acb3f17c Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 17:39:51 +0800 Subject: [PATCH 04/34] feat: add run store file layer with metrics delta tracker --- comfy_research/engine/runs/run_store.py | 169 ++++++++++++++++++++++++ comfy_research/tests/test_run_store.py | 91 +++++++++++++ 2 files changed, 260 insertions(+) create mode 100644 comfy_research/engine/runs/run_store.py create mode 100644 comfy_research/tests/test_run_store.py diff --git a/comfy_research/engine/runs/run_store.py b/comfy_research/engine/runs/run_store.py new file mode 100644 index 0000000..cef6331 --- /dev/null +++ b/comfy_research/engine/runs/run_store.py @@ -0,0 +1,169 @@ +"""File layer of the run store: ``data/runs/{run_id}/`` is the source of truth.""" +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +from comfy_research.schemas.run_record import RunRecord +from comfy_research.schemas.train_request import sanitize_train_ndjson_value + +_REPO_ROOT = Path(__file__).resolve().parents[3] + +SERIES_KEYS = ("loss_history", "test_loss_history", "reg_loss_history", + "step_ticks", "epoch_ticks") +_RESULT_STRIP_KEYS = frozenset( + {"checkpoint_b64", "plot_png_base64", "visualization_node_ids", + "observable_viz_updates", "observable_embedding_histories", + "observable_attention_slice_histories", "type"} +) + + +def runs_root() -> Path: + env = os.environ.get("COMFYRESEARCH_RUNS_DIR", "").strip() + return Path(env) if env else _REPO_ROOT / "data" / "runs" + + +def run_dir(run_id: str) -> Path: + if not run_id or "/" in run_id or "\\" in run_id or run_id.startswith("."): + raise ValueError(f"invalid run_id: {run_id!r}") + return runs_root() / run_id + + +def _atomic_write_json(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + tmp.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") + tmp.replace(path) + + +def write_run_record(rec: RunRecord) -> None: + _atomic_write_json(run_dir(rec.run_id) / "run.json", rec.model_dump(mode="json")) + + +def read_run_record(run_id: str) -> RunRecord | None: + path = run_dir(run_id) / "run.json" + try: + return RunRecord.model_validate_json(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + + +def append_metric_rows(run_id: str, rows: list[dict]) -> None: + if not rows: + return + path = run_dir(run_id) / "metrics.ndjson" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as f: + for row in rows: + f.write(json.dumps(sanitize_train_ndjson_value(row), separators=(",", ":")) + "\n") + + +def read_metric_rows(run_id: str) -> list[dict]: + path = run_dir(run_id) / "metrics.ndjson" + if not path.is_file(): + return [] + rows: list[dict] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + try: + rows.append(json.loads(line)) + except json.JSONDecodeError: + break # truncated tail from an interrupted append; drop it + return rows + + +def write_results(run_id: str, payload: dict) -> None: + kept = {k: v for k, v in payload.items() if k not in _RESULT_STRIP_KEYS} + _atomic_write_json(run_dir(run_id) / "results.json", + sanitize_train_ndjson_value(kept)) + + +def read_results(run_id: str) -> dict | None: + path = run_dir(run_id) / "results.json" + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + + +def load_series(run_id: str) -> tuple[str, dict]: + """Authority rule: ``results.json`` if present, else series rebuilt from ``metrics.ndjson``. + + Returns ("results" | "ndjson", series-dict). The single place that decides which + source wins — the API, index rebuild, and reconciliation must all use it. + """ + results = read_results(run_id) + if results is not None: + return "results", results + rows = read_metric_rows(run_id) + return "ndjson", { + "loss_history": [r.get("loss") for r in rows], + "test_loss_history": [r["test_loss"] for r in rows if "test_loss" in r], + "reg_loss_history": [r["reg_loss"] for r in rows if "reg_loss" in r], + "step_ticks": [r["step"] for r in rows if "step" in r], + "epoch_ticks": [r["epoch"] for r in rows if "epoch" in r], + } + + +class MetricsDeltaTracker: + """Turn cumulative-history ``metrics`` events into append-only delta rows. + + The trainer re-sends full histories on every emission; appending raw payloads + would be O(n^2) in storage. Track last-seen length, emit only new indices, + keep the latest cumulative snapshot for terminal fallback. + """ + + def __init__(self) -> None: + self._latest: dict[str, Any] = {k: [] for k in SERIES_KEYS} + self._latest["observable_metric_histories"] = {} + self._seen = 0 + + def extract(self, event: dict) -> list[dict]: + if event.get("type") != "metrics": + return [] + for key in SERIES_KEYS: + v = event.get(key) + if isinstance(v, list): + self._latest[key] = list(v) + obs = event.get("observable_metric_histories") + if isinstance(obs, dict): + self._latest["observable_metric_histories"] = { + str(k): list(v) for k, v in obs.items() if isinstance(v, list) + } + loss = self._latest["loss_history"] + rows: list[dict] = [] + for i in range(self._seen, len(loss)): + row: dict[str, Any] = {"idx": i, "loss": loss[i]} + for name, key in (("step", "step_ticks"), ("test_loss", "test_loss_history"), + ("reg_loss", "reg_loss_history"), ("epoch", "epoch_ticks")): + series = self._latest[key] + if i < len(series): + row[name] = series[i] + obs_row = { + k: v[i] + for k, v in self._latest["observable_metric_histories"].items() + if i < len(v) + } + if obs_row: + row["obs"] = obs_row + rows.append(row) + self._seen = len(loss) + return rows + + def snapshot(self) -> dict: + return {k: list(v) if isinstance(v, list) else dict(v) for k, v in self._latest.items()} + + +def summarize(snapshot: dict) -> dict: + loss = snapshot.get("loss_history") or [] + test = [x for x in (snapshot.get("test_loss_history") or []) if isinstance(x, (int, float))] + return { + "final_loss": loss[-1] if loss else None, + "final_test_loss": test[-1] if test else None, + "best_test_loss": min(test) if test else None, + "steps_completed": len(loss), + } diff --git a/comfy_research/tests/test_run_store.py b/comfy_research/tests/test_run_store.py new file mode 100644 index 0000000..f36abf7 --- /dev/null +++ b/comfy_research/tests/test_run_store.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms +from comfy_research.engine.runs import run_store + + +def _record(run_id: str) -> RunRecord: + return RunRecord( + run_id=run_id, origin="agent", status="running", created_at=now_ms(), + trainer_node_id="t1", graph=GraphDocument(version=1, nodes=[], edges=[]), + ) + + +def _metrics_event(n: int) -> dict: + return { + "type": "metrics", "step": n, + "loss_history": [1.0 / (i + 1) for i in range(n)], + "test_loss_history": [2.0 / (i + 1) for i in range(n)], + "reg_loss_history": [], + "step_ticks": list(range(n)), + "epoch_ticks": [], + "observable_metric_histories": {"obs1:acc": [float(i) for i in range(n)]}, + "observable_warnings": {}, + } + + +def test_runs_root_env_override(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path / "r")) + assert run_store.runs_root() == tmp_path / "r" + + +def test_record_roundtrip_and_atomicity(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + rid = new_run_id() + run_store.write_run_record(_record(rid)) + rec = run_store.read_run_record(rid) + assert rec is not None and rec.run_id == rid + assert not list((tmp_path / rid).glob("*.tmp")) + assert run_store.read_run_record("run-missing00000") is None + + +def test_delta_tracker_appends_only_new_rows() -> None: + tracker = run_store.MetricsDeltaTracker() + rows1 = tracker.extract(_metrics_event(2)) + rows2 = tracker.extract(_metrics_event(5)) + assert len(rows1) == 2 and len(rows2) == 3 + assert rows2[0]["idx"] == 2 and rows2[0]["step"] == 2 + assert rows2[-1]["loss"] == 1.0 / 5 + assert rows2[-1]["obs"]["obs1:acc"] == 4.0 + snap = tracker.snapshot() + assert len(snap["loss_history"]) == 5 + + +def test_delta_tracker_ignores_non_metrics_events() -> None: + tracker = run_store.MetricsDeltaTracker() + assert tracker.extract({"type": "progress", "step": 1, "total": 4}) == [] + + +def test_metrics_ndjson_roundtrip_tolerates_truncated_tail(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + rid = new_run_id() + run_store.write_run_record(_record(rid)) + run_store.append_metric_rows(rid, [{"idx": 0, "loss": 1.0}, {"idx": 1, "loss": 0.5}]) + path = tmp_path / rid / "metrics.ndjson" + path.write_text(path.read_text() + '{"idx": 2, "lo', encoding="utf-8") + rows = run_store.read_metric_rows(rid) + assert [r["idx"] for r in rows] == [0, 1] + + +def test_results_sanitized_and_stripped(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + rid = new_run_id() + run_store.write_run_record(_record(rid)) + run_store.write_results(rid, { + "loss_history": [1.0, float("nan")], + "checkpoint_b64": "QUJD", + "plot_png_base64": "aW1n", + "observable_embedding_histories": {"a": [1]}, + }) + res = run_store.read_results(rid) + assert res == {"loss_history": [1.0, None]} + + +def test_summarize() -> None: + tracker = run_store.MetricsDeltaTracker() + tracker.extract(_metrics_event(4)) + s = run_store.summarize(tracker.snapshot()) + assert s["final_loss"] == 0.25 + assert s["best_test_loss"] == 0.5 + assert s["steps_completed"] == 4 From 5d0691fda872a625e039464e719502fe08955fb9 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 17:44:46 +0800 Subject: [PATCH 05/34] feat: add rebuildable SQLite run index with cursor queries and reconciliation --- comfy_research/engine/runs/run_index.py | 207 ++++++++++++++++++++++++ comfy_research/tests/test_run_index.py | 123 ++++++++++++++ 2 files changed, 330 insertions(+) create mode 100644 comfy_research/engine/runs/run_index.py create mode 100644 comfy_research/tests/test_run_index.py diff --git a/comfy_research/engine/runs/run_index.py b/comfy_research/engine/runs/run_index.py new file mode 100644 index 0000000..34996db --- /dev/null +++ b/comfy_research/engine/runs/run_index.py @@ -0,0 +1,207 @@ +"""Rebuildable SQLite read-index over ``data/runs/*/run.json`` (never a truth source).""" +from __future__ import annotations + +import json +import logging +import sqlite3 +from dataclasses import dataclass, field +from pathlib import Path + +from comfy_research.engine.runs import run_store +from comfy_research.schemas.run_record import RunRecord, now_ms + +logger = logging.getLogger(__name__) + +_COLUMNS = ( + "run_id", "group_id", "parent_id", "origin", "status", "created_at", + "started_at", "finished_at", "last_heartbeat_at", "trainer_node_id", + "device", "error_detail", "hyperparams_json", "final_loss", + "final_test_loss", "best_test_loss", "steps_completed", "duration_seconds", +) +_ORDERABLE = {"created_at", "finished_at", "final_loss", "final_test_loss", + "best_test_loss", "steps_completed", "duration_seconds"} + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS runs ( + run_id TEXT PRIMARY KEY, group_id TEXT, parent_id TEXT, + origin TEXT NOT NULL, status TEXT NOT NULL, + created_at REAL NOT NULL, started_at REAL, finished_at REAL, + last_heartbeat_at REAL, trainer_node_id TEXT NOT NULL, + device TEXT DEFAULT '', error_detail TEXT DEFAULT '', + hyperparams_json TEXT NOT NULL DEFAULT '{}', + final_loss REAL, final_test_loss REAL, best_test_loss REAL, + steps_completed INTEGER, duration_seconds REAL +); +CREATE INDEX IF NOT EXISTS idx_runs_status ON runs(status); +CREATE INDEX IF NOT EXISTS idx_runs_group ON runs(group_id); +CREATE INDEX IF NOT EXISTS idx_runs_created ON runs(created_at); +""" + + +def index_path() -> Path: + return run_store.runs_root() / "index.db" + + +def _connect() -> sqlite3.Connection: + index_path().parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(index_path(), timeout=5) + conn.execute("PRAGMA journal_mode=WAL") + conn.executescript(_SCHEMA) + return conn + + +def upsert_run(rec: RunRecord, summary: dict | None = None) -> None: + s = summary or {} + duration = ( + (rec.finished_at - rec.started_at) / 1000.0 + if rec.finished_at is not None and rec.started_at is not None else None + ) + try: + with _connect() as conn: + conn.execute( + f"INSERT OR REPLACE INTO runs ({','.join(_COLUMNS)}) " + f"VALUES ({','.join('?' * len(_COLUMNS))})", + (rec.run_id, rec.group_id, rec.parent_id, rec.origin, rec.status, + rec.created_at, rec.started_at, rec.finished_at, now_ms(), + rec.trainer_node_id, rec.device, rec.error_detail, + json.dumps(rec.hyperparams), s.get("final_loss"), + s.get("final_test_loss"), s.get("best_test_loss"), + s.get("steps_completed"), duration), + ) + except sqlite3.Error: + logger.warning("run index upsert failed for %s", rec.run_id, exc_info=True) + + +def touch_heartbeat(run_id: str, at_ms: float) -> None: + try: + with _connect() as conn: + conn.execute("UPDATE runs SET last_heartbeat_at=? WHERE run_id=?", (at_ms, run_id)) + except sqlite3.Error: + logger.warning("run index heartbeat failed for %s", run_id, exc_info=True) + + +@dataclass +class RunQuery: + status: str | None = None + origin: str | None = None + group_id: str | None = None + since_ms: float | None = None + ids: list[str] | None = None + hyperparams: dict[str, str] = field(default_factory=dict) + order_by: str = "-created_at" + limit: int = 100 + cursor: str | None = None + + +def _row_to_dict(row: tuple) -> dict: + d = dict(zip(_COLUMNS, row)) + d["hyperparams"] = json.loads(d.pop("hyperparams_json") or "{}") + return d + + +def query_runs(q: RunQuery) -> tuple[list[dict], str | None]: + key = q.order_by.lstrip("-") + if key not in _ORDERABLE: + key, q = "created_at", RunQuery(**{**q.__dict__, "order_by": "-created_at"}) + direction = "DESC" if q.order_by.startswith("-") else "ASC" + where, params = ["1=1"], [] + for col, val in (("status", q.status), ("origin", q.origin), ("group_id", q.group_id)): + if val is not None: + where.append(f"{col}=?") + params.append(val) + if q.since_ms is not None: + where.append("created_at>=?") + params.append(q.since_ms) + if q.ids: + where.append(f"run_id IN ({','.join('?' * len(q.ids))})") + params.extend(q.ids) + for hk, hv in q.hyperparams.items(): + where.append("CAST(json_extract(hyperparams_json, ?) AS TEXT)=?") + params.extend([f'$."{hk}"', hv]) + if q.cursor: + cv, cid = q.cursor.rsplit(":", 1) + op = "<" if direction == "DESC" else ">" + where.append(f"({key} {op} ? OR ({key} = ? AND run_id {op} ?))") + params.extend([float(cv), float(cv), cid]) + sql = (f"SELECT {','.join(_COLUMNS)} FROM runs WHERE {' AND '.join(where)} " + f"ORDER BY {key} {direction} NULLS LAST, run_id {direction} LIMIT ?") + limit = max(1, min(int(q.limit), 500)) + with _connect() as conn: + rows = [_row_to_dict(r) for r in conn.execute(sql, [*params, limit + 1])] + next_cursor = None + if len(rows) > limit: + rows = rows[:limit] + last = rows[-1] + next_cursor = f"{last[key]}:{last['run_id']}" + return rows, next_cursor + + +def group_summary() -> list[dict]: + with _connect() as conn: + raw = conn.execute( + "SELECT group_id, status, COUNT(*), MIN(best_test_loss), MIN(final_loss) " + "FROM runs WHERE group_id IS NOT NULL GROUP BY group_id, status" + ).fetchall() + groups: dict[str, dict] = {} + for gid, status, count, best_test, best_final in raw: + g = groups.setdefault(gid, {"group_id": gid, "counts": {}, + "best_test_loss": None, "best_final_loss": None}) + g["counts"][status] = count + for k, v in (("best_test_loss", best_test), ("best_final_loss", best_final)): + if v is not None and (g[k] is None or v < g[k]): + g[k] = v + return sorted(groups.values(), key=lambda g: g["group_id"]) + + +def delete_rows(run_ids: list[str]) -> None: + if not run_ids: + return + with _connect() as conn: + conn.execute(f"DELETE FROM runs WHERE run_id IN ({','.join('?' * len(run_ids))})", run_ids) + + +def rebuild_index() -> int: + root = run_store.runs_root() + count = 0 + if index_path().exists(): + index_path().unlink() + for entry in sorted(root.glob("run-*/run.json")): + run_id = entry.parent.name + rec = run_store.read_run_record(run_id) + if rec is None: + # Corrupted run.json: surface it, don't silently skip (spec: unreadable rows). + try: + with _connect() as conn: + conn.execute( + "INSERT OR REPLACE INTO runs (run_id, origin, status, created_at, " + "trainer_node_id) VALUES (?, 'human', 'unreadable', 0, '')", + (run_id,), + ) + count += 1 + except sqlite3.Error: + logger.warning("could not index unreadable run %s", run_id, exc_info=True) + continue + _, series = run_store.load_series(rec.run_id) + upsert_run(rec, run_store.summarize(series)) + count += 1 + return count + + +def reconcile_stale_running(timeout_ms: float = 60_000) -> list[str]: + cutoff = now_ms() - timeout_ms + with _connect() as conn: + stale = [r[0] for r in conn.execute( + "SELECT run_id FROM runs WHERE status IN ('running','queued') " + "AND (last_heartbeat_at IS NULL OR last_heartbeat_at < ?)", (cutoff,) + )] + crashed = [] + for run_id in stale: + rec = run_store.read_run_record(run_id) + if rec is None or rec.status not in ("running", "queued"): + continue + rec = rec.model_copy(update={"status": "crashed", "finished_at": now_ms()}) + run_store.write_run_record(rec) + _, series = run_store.load_series(run_id) + upsert_run(rec, run_store.summarize(series)) + crashed.append(run_id) + return crashed diff --git a/comfy_research/tests/test_run_index.py b/comfy_research/tests/test_run_index.py new file mode 100644 index 0000000..eb67c3a --- /dev/null +++ b/comfy_research/tests/test_run_index.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery + + +def _record(status: str = "running", origin: str = "agent", + group: str | None = None, lr: float = 0.01) -> RunRecord: + return RunRecord( + run_id=new_run_id(), origin=origin, status=status, created_at=now_ms(), + group_id=group, trainer_node_id="t1", + graph=GraphDocument(version=1, nodes=[], edges=[]), + hyperparams={"opt1.lr": lr}, + ) + + +def _seed(monkeypatch, tmp_path): + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + + +def test_upsert_and_query_filters(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + a = _record(status="completed", lr=0.01) + b = _record(status="failed", lr=0.1) + for r in (a, b): + run_store.write_run_record(r) + run_index.upsert_run(r) + rows, cursor = run_index.query_runs(RunQuery(status="completed")) + assert [r["run_id"] for r in rows] == [a.run_id] + assert cursor is None + rows, _ = run_index.query_runs(RunQuery(hyperparams={"opt1.lr": "0.1"})) + assert [r["run_id"] for r in rows] == [b.run_id] + rows, _ = run_index.query_runs(RunQuery(ids=[a.run_id])) + assert rows[0]["hyperparams"] == {"opt1.lr": 0.01} + + +def test_cursor_pagination(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + made = [] + for _ in range(5): + r = _record(status="completed") + run_store.write_run_record(r) + run_index.upsert_run(r) + made.append(r.run_id) + page1, cur1 = run_index.query_runs(RunQuery(limit=2)) + page2, cur2 = run_index.query_runs(RunQuery(limit=2, cursor=cur1)) + page3, cur3 = run_index.query_runs(RunQuery(limit=2, cursor=cur2)) + ids = [r["run_id"] for r in page1 + page2 + page3] + assert sorted(ids) == sorted(made) and len(ids) == 5 + assert cur3 is None + + +def test_summary_columns_and_order_by(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + a, b = _record(status="completed"), _record(status="completed") + for r, loss in ((a, 0.5), (b, 0.1)): + run_store.write_run_record(r) + run_index.upsert_run(r, summary={"final_loss": loss, "final_test_loss": None, + "best_test_loss": None, "steps_completed": 3}) + rows, _ = run_index.query_runs(RunQuery(order_by="final_loss")) + assert [r["run_id"] for r in rows] == [b.run_id, a.run_id] + assert rows[0]["final_loss"] == 0.1 and rows[0]["steps_completed"] == 3 + + +def test_rebuild_surfaces_unreadable_run(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + good = _record(status="completed") + run_store.write_run_record(good) + run_index.upsert_run(good) + bad_dir = tmp_path / "run-corrupted0001" + bad_dir.mkdir() + (bad_dir / "run.json").write_text("{not json", encoding="utf-8") + assert run_index.rebuild_index() == 2 + rows, _ = run_index.query_runs(RunQuery(status="unreadable")) + assert [r["run_id"] for r in rows] == ["run-corrupted0001"] + + +def test_rebuild_from_files(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + recs = [_record(status="completed") for _ in range(3)] + for r in recs: + run_store.write_run_record(r) + run_index.upsert_run(r) + before, _ = run_index.query_runs(RunQuery()) + run_index.index_path().unlink() + assert run_index.rebuild_index() == 3 + after, _ = run_index.query_runs(RunQuery()) + assert {r["run_id"] for r in after} == {r["run_id"] for r in before} + + +def test_reconcile_stale_running(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + stale = _record(status="running") + run_store.write_run_record(stale) + run_index.upsert_run(stale) + run_index.touch_heartbeat(stale.run_id, now_ms() - 120_000) + fresh = _record(status="running") + run_store.write_run_record(fresh) + run_index.upsert_run(fresh) + run_index.touch_heartbeat(fresh.run_id, now_ms()) + + crashed = run_index.reconcile_stale_running(timeout_ms=60_000) + assert crashed == [stale.run_id] + assert run_store.read_run_record(stale.run_id).status == "crashed" + rows, _ = run_index.query_runs(RunQuery(ids=[fresh.run_id])) + assert rows[0]["status"] == "running" + + +def test_group_summary_and_delete(tmp_path, monkeypatch) -> None: + _seed(monkeypatch, tmp_path) + a = _record(status="completed", group="g1") + b = _record(status="failed", group="g1") + for r in (a, b): + run_store.write_run_record(r) + run_index.upsert_run(r) + groups = run_index.group_summary() + g1 = next(g for g in groups if g["group_id"] == "g1") + assert g1["counts"] == {"completed": 1, "failed": 1} + run_index.delete_rows([a.run_id, b.run_id]) + rows, _ = run_index.query_runs(RunQuery(group_id="g1")) + assert rows == [] From 4fc9199f3617524b42ecfa2da81a9fa9d6a63635 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 17:50:47 +0800 Subject: [PATCH 06/34] fix: NULL-safe cursor pagination and portable NULL ordering in run index --- comfy_research/engine/runs/run_index.py | 30 +++++++++++--- comfy_research/tests/test_run_index.py | 52 +++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 6 deletions(-) diff --git a/comfy_research/engine/runs/run_index.py b/comfy_research/engine/runs/run_index.py index 34996db..69e0830 100644 --- a/comfy_research/engine/runs/run_index.py +++ b/comfy_research/engine/runs/run_index.py @@ -118,21 +118,39 @@ def query_runs(q: RunQuery) -> tuple[list[dict], str | None]: for hk, hv in q.hyperparams.items(): where.append("CAST(json_extract(hyperparams_json, ?) AS TEXT)=?") params.extend([f'$."{hk}"', hv]) + + # Handle cursor pagination with NULL-safe keyset ordering if q.cursor: - cv, cid = q.cursor.rsplit(":", 1) - op = "<" if direction == "DESC" else ">" - where.append(f"({key} {op} ? OR ({key} = ? AND run_id {op} ?))") - params.extend([float(cv), float(cv), cid]) + if q.cursor.startswith("n:"): # NULL key cursor + cid = q.cursor[2:] + op = "<" if direction == "DESC" else ">" + where.append(f"{key} IS NULL AND run_id {op} ?") + params.append(cid) + else: # Non-NULL key cursor (format "v::") + parts = q.cursor.rsplit(":", 1) + cv_str = parts[0][2:] # Remove "v:" prefix + cid = parts[1] + cv = float(cv_str) + op = "<" if direction == "DESC" else ">" + where.append(f"({key} {op} ? OR ({key} = ? AND run_id {op} ?) OR {key} IS NULL)") + params.extend([cv, cv, cid]) + + # Portable NULL ordering: ({key} IS NULL) ASC sorts 0 for non-NULL, 1 for NULL (NULLs last) sql = (f"SELECT {','.join(_COLUMNS)} FROM runs WHERE {' AND '.join(where)} " - f"ORDER BY {key} {direction} NULLS LAST, run_id {direction} LIMIT ?") + f"ORDER BY ({key} IS NULL) ASC, {key} {direction}, run_id {direction} LIMIT ?") limit = max(1, min(int(q.limit), 500)) with _connect() as conn: rows = [_row_to_dict(r) for r in conn.execute(sql, [*params, limit + 1])] + next_cursor = None if len(rows) > limit: rows = rows[:limit] last = rows[-1] - next_cursor = f"{last[key]}:{last['run_id']}" + # Encode cursor: "n:" for NULL key, "v::" for non-NULL + if last[key] is None: + next_cursor = f"n:{last['run_id']}" + else: + next_cursor = f"v:{last[key]}:{last['run_id']}" return rows, next_cursor diff --git a/comfy_research/tests/test_run_index.py b/comfy_research/tests/test_run_index.py index eb67c3a..6461664 100644 --- a/comfy_research/tests/test_run_index.py +++ b/comfy_research/tests/test_run_index.py @@ -121,3 +121,55 @@ def test_group_summary_and_delete(tmp_path, monkeypatch) -> None: run_index.delete_rows([a.run_id, b.run_id]) rows, _ = run_index.query_runs(RunQuery(group_id="g1")) assert rows == [] + + +def test_cursor_pagination_with_nulls(tmp_path, monkeypatch) -> None: + """Test cursor pagination with NULL values in order-by column. + + Regression test for crash when cursor ends on row with NULL order key. + Ensures no dropped or duplicated rows across pages. + """ + _seed(monkeypatch, tmp_path) + # Create 3 runs with NULL final_loss and 2 with values + null_runs = [] + for _ in range(3): + r = _record(status="completed") + run_store.write_run_record(r) + run_index.upsert_run(r, summary={"final_loss": None, "final_test_loss": None, + "best_test_loss": None, "steps_completed": 0}) + null_runs.append(r.run_id) + + value_runs = [] + for i, loss in enumerate([0.5, 0.1]): + r = _record(status="completed") + run_store.write_run_record(r) + run_index.upsert_run(r, summary={"final_loss": loss, "final_test_loss": None, + "best_test_loss": None, "steps_completed": 1}) + value_runs.append(r.run_id) + + all_runs = value_runs + null_runs + + # Query with order_by final_loss (descending), paginate with limit=2 + page1, cur1 = run_index.query_runs(RunQuery(order_by="-final_loss", limit=2)) + page2, cur2 = run_index.query_runs(RunQuery(order_by="-final_loss", limit=2, cursor=cur1)) + page3, cur3 = run_index.query_runs(RunQuery(order_by="-final_loss", limit=2, cursor=cur2)) + + # Collect all pages + pages = [page1, page2, page3] + all_pages = [] + for p in pages: + all_pages.extend([r["run_id"] for r in p]) + + # Assertions + assert len(all_pages) == 5, f"Expected 5 rows, got {len(all_pages)}" + assert sorted(all_pages) == sorted(all_runs), f"Mismatch in collected rows" + assert len(set(all_pages)) == len(all_pages), "Duplicate rows found across pages" + + # Verify NULL rows come after non-NULL ones (DESC sort means high to low, NULLs last) + flat_rows = page1 + page2 + page3 + non_null_indices = [i for i, r in enumerate(flat_rows) if r["final_loss"] is not None] + null_indices = [i for i, r in enumerate(flat_rows) if r["final_loss"] is None] + + if non_null_indices and null_indices: + assert max(non_null_indices) < min(null_indices), \ + "NULL rows should come after non-NULL rows (NULLS LAST)" From 6b116aa57ce815fb2406955bb88d972a4b6c1b45 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 17:55:01 +0800 Subject: [PATCH 07/34] feat: add RunWriter facade with capture_events tee --- comfy_research/engine/runs/run_writer.py | 90 ++++++++++++++++++++++++ comfy_research/tests/test_run_writer.py | 70 ++++++++++++++++++ 2 files changed, 160 insertions(+) create mode 100644 comfy_research/engine/runs/run_writer.py create mode 100644 comfy_research/tests/test_run_writer.py diff --git a/comfy_research/engine/runs/run_writer.py b/comfy_research/engine/runs/run_writer.py new file mode 100644 index 0000000..7939b99 --- /dev/null +++ b/comfy_research/engine/runs/run_writer.py @@ -0,0 +1,90 @@ +"""RunWriter: persist one training run's lifecycle from its NDJSON event stream.""" +from __future__ import annotations + +import logging +from typing import Any, Iterator + +from comfy_research.engine.runs import run_index, run_store +from comfy_research.schemas.run_record import RunRecord, TERMINAL_STATUSES, now_ms + +logger = logging.getLogger(__name__) + +_HEARTBEAT_MIN_INTERVAL_MS = 1000.0 +_TERMINAL_EVENT_STATUS = {"complete": "completed", "aborted": "aborted", + "paused": "paused", "error": "failed"} + + +class RunWriter: + def __init__(self, record: RunRecord) -> None: + self._record = record + self._tracker = run_store.MetricsDeltaTracker() + self._last_heartbeat = 0.0 + run_store.write_run_record(record) + run_index.upsert_run(record) + + @property + def record(self) -> RunRecord: + return self._record + + @property + def is_terminal(self) -> bool: + return self._record.status in TERMINAL_STATUSES + + def _update(self, **changes: Any) -> None: + self._record = self._record.model_copy(update=changes) + run_store.write_run_record(self._record) + + def mark_running(self) -> None: + if self._record.status == "queued": + self._update(status="running", started_at=now_ms()) + run_index.upsert_run(self._record) + + def on_event(self, event: dict) -> None: + if self.is_terminal: + return + etype = str(event.get("type", "")) + rows = self._tracker.extract(event) + if rows: + try: + run_store.append_metric_rows(self._record.run_id, rows) + except OSError: + logger.warning("metric append failed for %s", self._record.run_id, exc_info=True) + now = now_ms() + if now - self._last_heartbeat >= _HEARTBEAT_MIN_INTERVAL_MS: + run_index.touch_heartbeat(self._record.run_id, now) + self._last_heartbeat = now + status = _TERMINAL_EVENT_STATUS.get(etype) + if status is None: + return + if etype in ("complete", "paused"): + run_store.write_results(self._record.run_id, dict(event)) + else: + run_store.write_results(self._record.run_id, self._tracker.snapshot()) + detail = str(event.get("detail", "")) if etype == "error" else "" + self.finalize(status, error_detail=detail) + + def finalize(self, status: str, error_detail: str = "") -> None: + if self.is_terminal: + return + self._update(status=status, finished_at=now_ms(), error_detail=error_detail) + snap = run_store.read_results(self._record.run_id) or self._tracker.snapshot() + run_index.upsert_run(self._record, run_store.summarize(snap)) + + def finalize_disconnect(self) -> None: + if not self.is_terminal: + run_store.write_results(self._record.run_id, self._tracker.snapshot()) + self.finalize("aborted") + + +def capture_events(events: Iterator[dict], writer: RunWriter) -> Iterator[dict]: + """Tee events into the writer; disconnect (generator close) finalizes as aborted.""" + writer.mark_running() + try: + for event in events: + try: + writer.on_event(event) + except Exception: + logger.warning("run capture failed for %s", writer.record.run_id, exc_info=True) + yield event + finally: + writer.finalize_disconnect() diff --git a/comfy_research/tests/test_run_writer.py b/comfy_research/tests/test_run_writer.py new file mode 100644 index 0000000..aa7415b --- /dev/null +++ b/comfy_research/tests/test_run_writer.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.engine.runs.run_writer import RunWriter, capture_events + + +def _writer(monkeypatch, tmp_path) -> RunWriter: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + return RunWriter(RunRecord( + run_id=new_run_id(), origin="agent", status="queued", created_at=now_ms(), + trainer_node_id="t1", graph=GraphDocument(version=1, nodes=[], edges=[]), + )) + + +def _events() -> list[dict]: + return [ + {"type": "progress", "step": 0, "total": 2}, + {"type": "metrics", "step": 2, "loss_history": [1.0, 0.5], + "test_loss_history": [], "reg_loss_history": [], "step_ticks": [0, 1], + "epoch_ticks": [], "observable_metric_histories": {}, "observable_warnings": {}}, + {"type": "complete", "checkpoint_b64": "QUJD", "plot_png_base64": "aW1n", + "loss_history": [1.0, 0.5, 0.25], "test_loss_history": [], + "reg_loss_history": [], "step_ticks": [0, 1, 2], "epoch_ticks": [], + "observable_viz_updates": [], "observable_metric_histories": {}, + "observable_embedding_histories": {}, "observable_attention_slice_histories": {}, + "observable_warnings": {}, "train_loop_seconds": 0.1, + "visualization_node_ids": []}, + ] + + +def test_complete_flow(monkeypatch, tmp_path) -> None: + w = _writer(monkeypatch, tmp_path) + seen = list(capture_events(iter(_events()), w)) + assert [e["type"] for e in seen] == ["progress", "metrics", "complete"] + rec = run_store.read_run_record(w.record.run_id) + assert rec.status == "completed" and rec.started_at and rec.finished_at + res = run_store.read_results(rec.run_id) + assert res["loss_history"] == [1.0, 0.5, 0.25] + assert "checkpoint_b64" not in res and "plot_png_base64" not in res + rows, _ = run_index.query_runs(RunQuery(ids=[rec.run_id])) + assert rows[0]["status"] == "completed" and rows[0]["final_loss"] == 0.25 + assert len(run_store.read_metric_rows(rec.run_id)) == 2 + + +def test_disconnect_mid_stream_finalizes_aborted(monkeypatch, tmp_path) -> None: + w = _writer(monkeypatch, tmp_path) + gen = capture_events(iter(_events()), w) + next(gen) + next(gen) # consumed progress + metrics, then client goes away + gen.close() + rec = run_store.read_run_record(w.record.run_id) + assert rec.status == "aborted" + assert run_store.read_results(rec.run_id)["loss_history"] == [1.0, 0.5] + + +def test_error_event_finalizes_failed(monkeypatch, tmp_path) -> None: + w = _writer(monkeypatch, tmp_path) + list(capture_events(iter([{"type": "error", "detail": "boom"}]), w)) + rec = run_store.read_run_record(w.record.run_id) + assert rec.status == "failed" and rec.error_detail == "boom" + + +def test_finalize_idempotent(monkeypatch, tmp_path) -> None: + w = _writer(monkeypatch, tmp_path) + list(capture_events(iter(_events()), w)) + w.finalize_disconnect() # no-op after completed + assert run_store.read_run_record(w.record.run_id).status == "completed" From 5cebcc9b49cc34c87b214ba39166d758994f3890 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:06:53 +0800 Subject: [PATCH 08/34] feat: capture /api/train runs (local and remote) into the run store --- comfy_research/api/train.py | 19 +++++++- comfy_research/engine/runs/run_writer.py | 35 ++++++++++++++- comfy_research/schemas/train_request.py | 3 ++ .../tests/test_train_run_capture.py | 44 +++++++++++++++++++ frontend/src/graph/readNdjsonTrainStream.ts | 5 ++- 5 files changed, 102 insertions(+), 4 deletions(-) create mode 100644 comfy_research/tests/test_train_run_capture.py diff --git a/comfy_research/api/train.py b/comfy_research/api/train.py index edc556d..af91319 100644 --- a/comfy_research/api/train.py +++ b/comfy_research/api/train.py @@ -19,6 +19,7 @@ from comfy_research.engine.crl.crl_run import iter_crl_events_from_context, prepare_crl_run from comfy_research.engine.runs.cuda_devices import list_local_cuda_devices from comfy_research.engine.runs.trainer_run import iter_trainer_events_from_context, prepare_trainer_run +from comfy_research.engine.runs.run_writer import RunWriter, build_run_record, capture_events from comfy_research.schemas.graph import NodeKind from comfy_research.config.remote_train_config import ( RemoteTrainConfig, @@ -412,7 +413,18 @@ def generate_remote(): set_last_validation_result(False, str(exc)) yield _ndjson_encode({"type": "error", "detail": str(exc)}) return - yield from iter_remote_train_stdout_lines(body, config=cfg_holder["cfg"]) + writer = RunWriter(build_run_record(body)) + yield _ndjson_encode({"type": "run_registered", "run_id": writer.record.run_id}) + writer.mark_running() + try: + for raw in iter_remote_train_stdout_lines(body, config=cfg_holder["cfg"]): + try: + writer.on_event(json.loads(raw.decode("utf-8"))) + except Exception: + pass # unparseable remote line: forward but don't capture + yield raw + finally: + writer.finalize_disconnect() return StreamingResponse( generate_remote(), @@ -452,8 +464,11 @@ def generate_crl(): hessian_oversized_policy=body.hessian_oversized_policy, ) + writer = RunWriter(build_run_record(body)) + def generate(): - for event in iter_trainer_events_from_context(ctx): + yield _ndjson_encode({"type": "run_registered", "run_id": writer.record.run_id}) + for event in capture_events(iter_trainer_events_from_context(ctx), writer): yield _ndjson_encode(event) return StreamingResponse( diff --git a/comfy_research/engine/runs/run_writer.py b/comfy_research/engine/runs/run_writer.py index 7939b99..83e29a1 100644 --- a/comfy_research/engine/runs/run_writer.py +++ b/comfy_research/engine/runs/run_writer.py @@ -5,7 +5,16 @@ from typing import Any, Iterator from comfy_research.engine.runs import run_index, run_store -from comfy_research.schemas.run_record import RunRecord, TERMINAL_STATUSES, now_ms +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import ( + RunRecord, + TERMINAL_STATUSES, + flatten_hyperparams, + new_run_id, + now_ms, + strip_result_data, +) +from comfy_research.schemas.train_request import TrainRequest logger = logging.getLogger(__name__) @@ -76,6 +85,30 @@ def finalize_disconnect(self) -> None: self.finalize("aborted") +def build_run_record( + body: TrainRequest, + *, + origin: str | None = None, + group_id: str | None = None, + status: str = "queued", +) -> RunRecord: + nodes = strip_result_data(body.nodes) + trainer = next((n for n in body.nodes if n.id == body.trainer_node_id), None) + device = str((trainer.data or {}).get("computeDevice", "")) if trainer else "" + return RunRecord( + run_id=new_run_id(), + origin=origin or body.run_origin, + group_id=group_id if group_id is not None else body.run_group_id, + parent_id=body.run_parent_id, + status=status, + created_at=now_ms(), + trainer_node_id=body.trainer_node_id, + device=device, + graph=GraphDocument(version=1, nodes=nodes, edges=body.edges), + hyperparams=flatten_hyperparams(body.nodes), + ) + + def capture_events(events: Iterator[dict], writer: RunWriter) -> Iterator[dict]: """Tee events into the writer; disconnect (generator close) finalizes as aborted.""" writer.mark_running() diff --git a/comfy_research/schemas/train_request.py b/comfy_research/schemas/train_request.py index 69afa14..d1b7a80 100644 --- a/comfy_research/schemas/train_request.py +++ b/comfy_research/schemas/train_request.py @@ -17,6 +17,9 @@ class TrainRequest(BaseModel): edges: list[Edge] = Field(default_factory=list) resume: dict[str, Any] | None = None hessian_oversized_policy: Literal["skip", "force"] | None = None + run_origin: Literal["human", "agent", "sweep"] = "human" + run_group_id: str | None = None + run_parent_id: str | None = None def sanitize_train_ndjson_value(obj: Any) -> Any: diff --git a/comfy_research/tests/test_train_run_capture.py b/comfy_research/tests/test_train_run_capture.py new file mode 100644 index 0000000..67bbbbb --- /dev/null +++ b/comfy_research/tests/test_train_run_capture.py @@ -0,0 +1,44 @@ +# comfy_research/tests/test_train_run_capture.py +from __future__ import annotations + +import json + +from fastapi.testclient import TestClient + +from comfy_research.main import app +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.tests.train_test_fixtures import minimal_cpu_train_request + + +def _events(text: str) -> list[dict]: + return [json.loads(line) for line in text.splitlines() if line.strip()] + + +def test_post_train_registers_and_persists_run(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + body = {**minimal_cpu_train_request(), "run_origin": "agent", "run_group_id": "g1"} + response = TestClient(app).post("/api/train", json=body) + assert response.status_code == 200 + events = _events(response.text) + assert events[0]["type"] == "run_registered" + run_id = events[0]["run_id"] + assert any(e["type"] == "complete" for e in events) + + rec = run_store.read_run_record(run_id) + assert rec.status == "completed" + assert rec.origin == "agent" and rec.group_id == "g1" + assert rec.graph.nodes # config snapshot present + for node in rec.graph.nodes: + assert "memoryCheckpoint_b64" not in (node.data or {}) + res = run_store.read_results(run_id) + assert len(res["loss_history"]) == 4 and "checkpoint_b64" not in res + rows, _ = run_index.query_runs(RunQuery(ids=[run_id])) + assert rows[0]["status"] == "completed" and rows[0]["hyperparams"] + + +def test_post_train_default_origin_human(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + response = TestClient(app).post("/api/train", json=minimal_cpu_train_request()) + run_id = _events(response.text)[0]["run_id"] + assert run_store.read_run_record(run_id).origin == "human" diff --git a/frontend/src/graph/readNdjsonTrainStream.ts b/frontend/src/graph/readNdjsonTrainStream.ts index acfe8d4..ec6c36d 100644 --- a/frontend/src/graph/readNdjsonTrainStream.ts +++ b/frontend/src/graph/readNdjsonTrainStream.ts @@ -99,7 +99,8 @@ export async function readNdjsonTrainStream( | TrainStreamPaused | TrainStreamAborted | TrainStreamError - | { type: "remote_session" }; + | { type: "remote_session" } + | { type: "run_registered"; run_id: string }; if (ev.type === "complete") complete = ev; else if (ev.type === "paused") paused = ev; else if (ev.type === "aborted") aborted = true; @@ -109,6 +110,8 @@ export async function readNdjsonTrainStream( else if (ev.type === "metrics") options?.onMetrics?.(ev); else if (ev.type === "remote_session") { options?.onRemoteSession?.(); + } else if (ev.type === "run_registered") { + // No-op: run id is not currently surfaced to the UI. } else onProgress(ev as TrainStreamProgress); }; From c40f55d62c8003cab450b50671eab1b67edf85db Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:12:09 +0800 Subject: [PATCH 09/34] fix: gitignore data/runs and isolate train API tests from the run store --- .gitignore | 1 + comfy_research/tests/test_train_api_integration.py | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 2dc4b73..8ba5cd4 100644 --- a/.gitignore +++ b/.gitignore @@ -68,3 +68,4 @@ data/graph_library/assets.json /docs/locales/**/*.mo .diff-worktrees/ data/runtime/ +data/runs/ diff --git a/comfy_research/tests/test_train_api_integration.py b/comfy_research/tests/test_train_api_integration.py index 9e7b086..c200351 100644 --- a/comfy_research/tests/test_train_api_integration.py +++ b/comfy_research/tests/test_train_api_integration.py @@ -13,7 +13,8 @@ def _ndjson_events(response_text: str) -> list[dict[str, object]]: return [json.loads(line) for line in response_text.splitlines() if line.strip()] -def test_post_train_streams_real_cpu_training_result() -> None: +def test_post_train_streams_real_cpu_training_result(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) response = TestClient(app).post("/api/train", json=minimal_cpu_train_request()) assert response.status_code == 200 @@ -35,7 +36,8 @@ def test_post_train_streams_real_cpu_training_result() -> None: ) -def test_post_train_rejects_invalid_graph_before_streaming() -> None: +def test_post_train_rejects_invalid_graph_before_streaming(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) response = TestClient(app).post( "/api/train", json={ From d293d1c70f20377a50e3335218af07a1b60d9a21 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:20:53 +0800 Subject: [PATCH 10/34] feat: capture sweep inner runs into the run store with sweep session grouping Wraps the non-CRL branch of iter_sweep_events with RunWriter/capture_events so each sweep combo persists as its own run (origin="sweep", group_id=). The wrapped iterator is closed explicitly after the consumption loop (guarded by hasattr) so finalization on early break (complete/aborted/paused) is deterministic rather than GC-dependent; the CRL branch stays unwrapped but gets the same close() guard for consistency. Applied the identical wrapping to train_coordinate_descent.py: both the baseline evaluation (_run_train_once) and the per-candidate evaluation inside the round/axis loop now persist inner runs under the coordinate- descent session id, since both consume iter_trainer_events_from_context the same way as the sweep path. Also fixes a pre-existing bug in validate_sweep_request that stringified the NodeKind enum before calling has_capability (str(NodeKind.trainer) == "NodeKind.trainer", never matching), which made POST /api/train/sweep 400 on every request. Every other call site in the codebase passes the enum directly; this brings train_sweep.py in line and was required for the new test (and the endpoint itself) to work at all. --- .../engine/runs/train_coordinate_descent.py | 38 ++++++++++++++++--- comfy_research/engine/runs/train_sweep.py | 18 ++++++++- .../tests/test_sweep_run_capture.py | 34 +++++++++++++++++ 3 files changed, 82 insertions(+), 8 deletions(-) create mode 100644 comfy_research/tests/test_sweep_run_capture.py diff --git a/comfy_research/engine/runs/train_coordinate_descent.py b/comfy_research/engine/runs/train_coordinate_descent.py index d6ce7d4..e0a33f1 100644 --- a/comfy_research/engine/runs/train_coordinate_descent.py +++ b/comfy_research/engine/runs/train_coordinate_descent.py @@ -21,8 +21,10 @@ _parse_data_path, _trainer_training_steps_override, ) +from comfy_research.engine.runs.run_writer import RunWriter, build_run_record, capture_events from comfy_research.engine.runs.trainer_run import iter_trainer_events_from_context, prepare_trainer_run from comfy_research.schemas.graph import Edge, Node +from comfy_research.schemas.train_request import TrainRequest MAX_COORDINATE_DESCENT_ROUNDS = 24 MAX_AXIS_VALUES = 128 @@ -210,6 +212,7 @@ def _run_train_once( edges: list[Edge], trainer_node_id: str, training_steps_override: int | None, + group_id: str, ) -> tuple[dict[str, Any] | None, str | None]: nodes_p = _trainer_training_steps_override(nodes, trainer_node_id, training_steps_override) try: @@ -225,14 +228,23 @@ def _run_train_once( return None, d if isinstance(d, str) else str(d) except Exception as e: return None, str(e) - for ev in iter_trainer_events_from_context(ctx): + inner_req = TrainRequest(trainer_node_id=trainer_node_id, nodes=nodes_p, edges=edges) + inner_writer = RunWriter(build_run_record(inner_req, origin="sweep", group_id=group_id)) + event_iter = capture_events(iter_trainer_events_from_context(ctx), inner_writer) + result: tuple[dict[str, Any] | None, str | None] = (None, "no_complete_event") + for ev in event_iter: if ev.get("type") == "complete": - return ev, None + result = (ev, None) + break if ev.get("type") == "aborted": - return None, "training_aborted" + result = (None, "training_aborted") + break if ev.get("type") == "paused": - return None, "training_paused_not_supported_in_tuning" - return None, "no_complete_event" + result = (None, "training_paused_not_supported_in_tuning") + break + if hasattr(event_iter, "close"): + event_iter.close() + return result def _is_better(a: tuple[float, float, float], b: tuple[float, float, float] | None) -> bool: @@ -303,6 +315,7 @@ def iter_coordinate_descent_events(body: TrainCoordinateDescentRequest) -> Itera body.edges, body.trainer_node_id, body.training_steps_override, + session_id, ) eval_index += 1 yield { @@ -358,7 +371,18 @@ def iter_coordinate_descent_events(body: TrainCoordinateDescentRequest) -> Itera resume=None, hessian_oversized_policy="skip", ) - for ev in iter_trainer_events_from_context(ctx): + inner_req = TrainRequest( + trainer_node_id=body.trainer_node_id, + nodes=nodes_p, + edges=body.edges, + ) + inner_writer = RunWriter( + build_run_record(inner_req, origin="sweep", group_id=session_id) + ) + event_iter = capture_events( + iter_trainer_events_from_context(ctx), inner_writer + ) + for ev in event_iter: if ev.get("type") == "complete": complete = ev break @@ -368,6 +392,8 @@ def iter_coordinate_descent_events(body: TrainCoordinateDescentRequest) -> Itera if ev.get("type") == "paused": err = "training_paused_not_supported_in_tuning" break + if hasattr(event_iter, "close"): + event_iter.close() if err is None and complete is None: err = "no_complete_event" if err is None and complete is not None: diff --git a/comfy_research/engine/runs/train_sweep.py b/comfy_research/engine/runs/train_sweep.py index 52749c3..4592ab7 100644 --- a/comfy_research/engine/runs/train_sweep.py +++ b/comfy_research/engine/runs/train_sweep.py @@ -17,8 +17,10 @@ unregister_sweep_session, ) from comfy_research.engine.crl.crl_run import iter_crl_events_from_context, prepare_crl_run +from comfy_research.engine.runs.run_writer import RunWriter, build_run_record, capture_events from comfy_research.engine.runs.trainer_run import iter_trainer_events_from_context, prepare_trainer_run from comfy_research.schemas.graph import Edge, Node, NodeKind +from comfy_research.schemas.train_request import TrainRequest MAX_SWEEP_POINTS = 256 @@ -191,7 +193,7 @@ def validate_sweep_request(body: TrainSweepRequest) -> None: nmap = {n.id: n for n in body.nodes} if body.trainer_node_id not in nmap: raise ValueError("trainer_node_id not found in nodes.") - tk = str(nmap[body.trainer_node_id].type) + tk = nmap[body.trainer_node_id].type if not has_capability(tk, "trainer_runner"): raise ValueError("trainer_node_id must refer to a trainer or crl_trainer node.") @@ -276,7 +278,17 @@ def iter_sweep_events(body: TrainSweepRequest) -> Iterator[dict[str, Any]]: resume=None, hessian_oversized_policy="skip", ) - event_iter = iter_trainer_events_from_context(ctx_sup) + inner_req = TrainRequest( + trainer_node_id=body.trainer_node_id, + nodes=nodes_p, + edges=body.edges, + ) + inner_writer = RunWriter( + build_run_record(inner_req, origin="sweep", group_id=session_id) + ) + event_iter = capture_events( + iter_trainer_events_from_context(ctx_sup), inner_writer + ) complete: dict[str, Any] | None = None for ev in event_iter: if ev.get("type") == "complete": @@ -288,6 +300,8 @@ def iter_sweep_events(body: TrainSweepRequest) -> Iterator[dict[str, Any]]: if ev.get("type") == "paused": err = "training_paused_not_supported_in_sweep" break + if hasattr(event_iter, "close"): + event_iter.close() if err: pass elif complete is None: diff --git a/comfy_research/tests/test_sweep_run_capture.py b/comfy_research/tests/test_sweep_run_capture.py new file mode 100644 index 0000000..bb44c2a --- /dev/null +++ b/comfy_research/tests/test_sweep_run_capture.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import json + +from fastapi.testclient import TestClient + +from comfy_research.main import app +from comfy_research.engine.runs import run_index +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.tests.train_test_fixtures import minimal_cpu_train_request + + +def test_sweep_inner_runs_captured_with_group(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + base = minimal_cpu_train_request() + body = { + "sweep_session_id": "sweep-test-1", + "trainer_node_id": base["trainer_node_id"], + "nodes": base["nodes"], + "edges": base["edges"], + "axes": [{"node_id": base["trainer_node_id"], "path": "trainingSteps", + "values": [2, 3]}], + "metric": {"kind": "final_train_loss"}, + } + response = TestClient(app).post("/api/train/sweep", json=body) + assert response.status_code == 200 + events = [json.loads(l) for l in response.text.splitlines() if l.strip()] + assert any(e["type"] == "sweep_complete" for e in events) + + rows, _ = run_index.query_runs(RunQuery(group_id="sweep-test-1")) + assert len(rows) == 2 + assert all(r["origin"] == "sweep" and r["status"] == "completed" for r in rows) + steps = sorted(r["hyperparams"][f"{base['trainer_node_id']}.trainingSteps"] for r in rows) + assert steps == [2, 3] From 9f58d2a8f691f183c8c7d455833ab18daa94562e Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:29:48 +0800 Subject: [PATCH 11/34] feat: add async-submit run worker pool with per-trainer serialization --- comfy_research/engine/runs/run_worker.py | 158 +++++++++++++++++++++++ comfy_research/tests/test_run_worker.py | 121 +++++++++++++++++ 2 files changed, 279 insertions(+) create mode 100644 comfy_research/engine/runs/run_worker.py create mode 100644 comfy_research/tests/test_run_worker.py diff --git a/comfy_research/engine/runs/run_worker.py b/comfy_research/engine/runs/run_worker.py new file mode 100644 index 0000000..4150c78 --- /dev/null +++ b/comfy_research/engine/runs/run_worker.py @@ -0,0 +1,158 @@ +# comfy_research/engine/runs/run_worker.py +"""Server-owned worker pool: async-submitted runs execute detached from any HTTP stream.""" +from __future__ import annotations + +import logging +import threading +from collections import deque +from concurrent.futures import ThreadPoolExecutor + +from fastapi import HTTPException + +from comfy_research.engine.runs.run_writer import RunWriter, build_run_record, capture_events +from comfy_research.engine.runs.trainer_run import ( + iter_trainer_events_from_context, + prepare_trainer_run, +) +from comfy_research.engine.runs.train_control import request_abort +from comfy_research.engine.runs.ai4science_alias import remap_ai4science_node_types +from comfy_research.schemas.train_request import TrainRequest + +logger = logging.getLogger(__name__) + +_DEFAULT_SLOTS = 2 + + +def _prefers_remote_gpu(body: TrainRequest) -> bool: + node = next((n for n in body.nodes if n.id == body.trainer_node_id), None) + data = getattr(node, "data", None) or {} + spec = str(data.get("computeDevice", "")).strip().lower() + return (spec == "cuda" or spec.startswith("cuda:")) and data.get("remoteGpu") is True + + +class RunWorkerPool: + def __init__(self, slots: int = _DEFAULT_SLOTS) -> None: + self._executor = ThreadPoolExecutor(max_workers=slots, thread_name_prefix="run-worker") + self._lock = threading.Lock() + self._active: dict[str, str] = {} # trainer_node_id -> run_id currently executing + self._waiting: dict[str, deque[str]] = {} # trainer_node_id -> queued run_ids (FIFO) + self._writers: dict[str, RunWriter] = {} + self._idempotency: dict[str, str] = {} + self._records: dict[str, TrainRequest] = {} + + def submit(self, body: TrainRequest, idempotency_key: str | None = None): + if _prefers_remote_gpu(body): + raise HTTPException(status_code=400, detail={ + "code": "remote_not_supported", + "detail": "Async submit runs locally only; use streaming POST /api/train for remote GPU runs.", + }) + with self._lock: + if idempotency_key and idempotency_key in self._idempotency: + run_id = self._idempotency[idempotency_key] + return self._writers[run_id].record + mapped = remap_ai4science_node_types(body.nodes) + prepare_trainer_run( + mapped, body.edges, body.trainer_node_id, + resume=body.resume, + hessian_oversized_policy=body.hessian_oversized_policy, + validate_only=True, + ) + writer = RunWriter(build_run_record(body)) + run_id = writer.record.run_id + trainer_id = body.trainer_node_id + with self._lock: + self._writers[run_id] = writer + self._records[run_id] = body + if idempotency_key: + self._idempotency[idempotency_key] = run_id + if trainer_id in self._active: + # Same trainer already executing: wait in FIFO (no thread blocked), + # because the train_control registry is a single slot per trainer id. + self._waiting.setdefault(trainer_id, deque()).append(run_id) + return writer.record + self._active[trainer_id] = run_id + self._executor.submit(self._execute, run_id) + return writer.record + + def _execute(self, run_id: str) -> None: + with self._lock: + writer = self._writers.get(run_id) + body = self._records.get(run_id) + if writer is None or body is None: + return + try: + if not writer.is_terminal: + mapped = remap_ai4science_node_types(body.nodes) + ctx = prepare_trainer_run( + mapped, body.edges, body.trainer_node_id, + resume=body.resume, + hessian_oversized_policy=body.hessian_oversized_policy, + ) + for _ in capture_events(iter_trainer_events_from_context(ctx), writer): + pass + except HTTPException as exc: + writer.finalize("failed", error_detail=str(exc.detail)) + except Exception as exc: + logger.warning("submitted run %s crashed", run_id, exc_info=True) + writer.finalize("failed", error_detail=f"{type(exc).__name__}: {exc}") + finally: + self._dispatch_next(body.trainer_node_id) + + def _dispatch_next(self, trainer_id: str) -> None: + with self._lock: + queue = self._waiting.get(trainer_id) + next_id = None + while queue: + candidate = queue.popleft() + w = self._writers.get(candidate) + if w is not None and not w.is_terminal: + next_id = candidate + break + if next_id is None: + self._active.pop(trainer_id, None) + if queue is not None and not queue: + self._waiting.pop(trainer_id, None) + return + self._active[trainer_id] = next_id + self._executor.submit(self._execute, next_id) + + def abort(self, run_id: str) -> bool: + with self._lock: + writer = self._writers.get(run_id) + body = self._records.get(run_id) + if writer is None or body is None or writer.is_terminal: + return False + trainer_id = body.trainer_node_id + queue = self._waiting.get(trainer_id) + if queue is not None and run_id in queue: + queue.remove(run_id) # waiting run: finalize directly, no train_control signal + writer.finalize("aborted") + return True + is_running_here = self._active.get(trainer_id) == run_id + if not is_running_here: + return False + request_abort(trainer_id) # only when THIS run holds the trainer; cooperative + return True + + def shutdown(self, wait: bool = False) -> None: + self._executor.shutdown(wait=wait, cancel_futures=True) + + +_pool: RunWorkerPool | None = None +_pool_lock = threading.Lock() + + +def get_worker_pool() -> RunWorkerPool: + global _pool + with _pool_lock: + if _pool is None: + _pool = RunWorkerPool() + return _pool + + +def reset_worker_pool_for_tests() -> None: + global _pool + with _pool_lock: + if _pool is not None: + _pool.shutdown(wait=True) + _pool = None diff --git a/comfy_research/tests/test_run_worker.py b/comfy_research/tests/test_run_worker.py new file mode 100644 index 0000000..369e5c3 --- /dev/null +++ b/comfy_research/tests/test_run_worker.py @@ -0,0 +1,121 @@ +# comfy_research/tests/test_run_worker.py +from __future__ import annotations + +import time + +import pytest +from fastapi import HTTPException + +from comfy_research.engine.runs import run_store +from comfy_research.engine.runs.run_worker import get_worker_pool, reset_worker_pool_for_tests +from comfy_research.schemas.train_request import TrainRequest +from comfy_research.tests.train_test_fixtures import minimal_cpu_train_request + + +def _wait_terminal(run_id: str, timeout_s: float = 30.0) -> str: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + rec = run_store.read_run_record(run_id) + if rec and rec.status not in ("queued", "running"): + return rec.status + time.sleep(0.05) + raise AssertionError(f"run {run_id} never reached terminal state") + + +@pytest.fixture(autouse=True) +def _isolated(tmp_path, monkeypatch): + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + reset_worker_pool_for_tests() + yield + reset_worker_pool_for_tests() + + +def test_submit_runs_to_completion() -> None: + body = TrainRequest.model_validate({**minimal_cpu_train_request(), "run_origin": "agent"}) + rec = get_worker_pool().submit(body) + assert rec.status == "queued" and rec.origin == "agent" + assert _wait_terminal(rec.run_id) == "completed" + assert len(run_store.read_results(rec.run_id)["loss_history"]) == 4 + + +def test_submit_invalid_graph_raises_400_and_persists_nothing(tmp_path) -> None: + body = TrainRequest.model_validate({ + "trainer_node_id": "trainer", + "nodes": [{"id": "trainer", "type": "trainer", + "data": {"trainingSteps": 1, "computeDevice": "cpu"}}], + "edges": [], + }) + with pytest.raises(HTTPException) as exc: + get_worker_pool().submit(body) + assert exc.value.status_code == 400 + assert not list(tmp_path.glob("run-*")) + + +def test_idempotency_key_returns_same_run() -> None: + body = TrainRequest.model_validate(minimal_cpu_train_request()) + a = get_worker_pool().submit(body, idempotency_key="k1") + b = get_worker_pool().submit(body, idempotency_key="k1") + assert a.run_id == b.run_id + _wait_terminal(a.run_id) + + +def test_same_trainer_runs_serialize() -> None: + pool = get_worker_pool() + body = TrainRequest.model_validate(minimal_cpu_train_request()) + a = pool.submit(body) + b = pool.submit(body) + assert _wait_terminal(a.run_id) == "completed" + assert _wait_terminal(b.run_id) == "completed" + + +def test_abort_waiting_run_is_scoped_to_that_run() -> None: + pool = get_worker_pool() + slow = minimal_cpu_train_request() + for n in slow["nodes"]: + if n["id"] == slow["trainer_node_id"]: + n["data"] = {**n["data"], "trainingSteps": 2000} + a = pool.submit(TrainRequest.model_validate(slow)) + b = pool.submit(TrainRequest.model_validate(slow)) # waits in FIFO behind a (same trainer) + assert pool.abort(b.run_id) is True + # b dies immediately, without any train_control signal that could hit a + assert run_store.read_run_record(b.run_id).status == "aborted" + assert run_store.read_run_record(a.run_id).status in ("queued", "running") + # abort the running run; cooperative signal can land before the training loop + # registers the trainer, so retry until it takes effect + deadline = time.monotonic() + 30.0 + while time.monotonic() < deadline: + pool.abort(a.run_id) + rec = run_store.read_run_record(a.run_id) + if rec.status not in ("queued", "running"): + break + time.sleep(0.1) + assert run_store.read_run_record(a.run_id).status == "aborted" + assert pool.abort("run-nonexistent0") is False + + +def _renamed_fixture(suffix: str) -> dict: + """Same minimal graph under fresh ids, so it counts as a different trainer.""" + req = minimal_cpu_train_request() + for n in req["nodes"]: + n["id"] = n["id"] + suffix + for e in req["edges"]: + e["id"] = e["id"] + suffix + e["source"] = e["source"] + suffix + e["target"] = e["target"] + suffix + req["trainer_node_id"] = req["trainer_node_id"] + suffix + return req + + +def test_different_trainer_not_blocked_by_same_trainer_queue() -> None: + pool = get_worker_pool() + slow = minimal_cpu_train_request() + for n in slow["nodes"]: + if n["id"] == slow["trainer_node_id"]: + n["data"] = {**n["data"], "trainingSteps": 2000} + a = pool.submit(TrainRequest.model_validate(slow)) + b = pool.submit(TrainRequest.model_validate(slow)) # waits in FIFO, occupies no slot + fast = pool.submit(TrainRequest.model_validate(_renamed_fixture("-x"))) + # with head-of-line blocking, `fast` would be stuck behind b in the 2-slot pool + assert _wait_terminal(fast.run_id) == "completed" + for rid in (a.run_id, b.run_id): + pool.abort(rid) From e0e751dc50a32c5d4f17b0e7427b277df975c834 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:35:48 +0800 Subject: [PATCH 12/34] fix: single-flight idempotency reservation in run worker submit --- comfy_research/engine/runs/run_worker.py | 80 ++++++++++++++++-------- comfy_research/tests/test_run_worker.py | 15 +++++ 2 files changed, 68 insertions(+), 27 deletions(-) diff --git a/comfy_research/engine/runs/run_worker.py b/comfy_research/engine/runs/run_worker.py index 4150c78..cae0f21 100644 --- a/comfy_research/engine/runs/run_worker.py +++ b/comfy_research/engine/runs/run_worker.py @@ -38,6 +38,7 @@ def __init__(self, slots: int = _DEFAULT_SLOTS) -> None: self._waiting: dict[str, deque[str]] = {} # trainer_node_id -> queued run_ids (FIFO) self._writers: dict[str, RunWriter] = {} self._idempotency: dict[str, str] = {} + self._idempotency_inflight: dict[str, threading.Event] = {} self._records: dict[str, TrainRequest] = {} def submit(self, body: TrainRequest, idempotency_key: str | None = None): @@ -46,33 +47,58 @@ def submit(self, body: TrainRequest, idempotency_key: str | None = None): "code": "remote_not_supported", "detail": "Async submit runs locally only; use streaming POST /api/train for remote GPU runs.", }) - with self._lock: - if idempotency_key and idempotency_key in self._idempotency: - run_id = self._idempotency[idempotency_key] - return self._writers[run_id].record - mapped = remap_ai4science_node_types(body.nodes) - prepare_trainer_run( - mapped, body.edges, body.trainer_node_id, - resume=body.resume, - hessian_oversized_policy=body.hessian_oversized_policy, - validate_only=True, - ) - writer = RunWriter(build_run_record(body)) - run_id = writer.record.run_id - trainer_id = body.trainer_node_id - with self._lock: - self._writers[run_id] = writer - self._records[run_id] = body - if idempotency_key: - self._idempotency[idempotency_key] = run_id - if trainer_id in self._active: - # Same trainer already executing: wait in FIFO (no thread blocked), - # because the train_control registry is a single slot per trainer id. - self._waiting.setdefault(trainer_id, deque()).append(run_id) - return writer.record - self._active[trainer_id] = run_id - self._executor.submit(self._execute, run_id) - return writer.record + reserved = False + if idempotency_key: + # Single-flight: only the caller that wins the reservation does the + # prepare/RunWriter work; concurrent same-key callers wait for it and + # then return its record, instead of each racing to persist their own run. + while True: + with self._lock: + if idempotency_key in self._idempotency: + run_id = self._idempotency[idempotency_key] + return self._writers[run_id].record + event = self._idempotency_inflight.get(idempotency_key) + if event is None: + self._idempotency_inflight[idempotency_key] = threading.Event() + reserved = True + break + event.wait() # outside the lock; re-check self._idempotency once released + try: + mapped = remap_ai4science_node_types(body.nodes) + prepare_trainer_run( + mapped, body.edges, body.trainer_node_id, + resume=body.resume, + hessian_oversized_policy=body.hessian_oversized_policy, + validate_only=True, + ) + writer = RunWriter(build_run_record(body)) + run_id = writer.record.run_id + trainer_id = body.trainer_node_id + submit_now = False + with self._lock: + self._writers[run_id] = writer + self._records[run_id] = body + if idempotency_key: + self._idempotency[idempotency_key] = run_id + if trainer_id in self._active: + # Same trainer already executing: wait in FIFO (no thread blocked), + # because the train_control registry is a single slot per trainer id. + self._waiting.setdefault(trainer_id, deque()).append(run_id) + else: + self._active[trainer_id] = run_id + submit_now = True + if submit_now: + self._executor.submit(self._execute, run_id) + return writer.record + finally: + if reserved: + # Release the reservation (success or failure) so a waiting caller + # either sees the persisted record or gets to retry itself — a + # validation failure must not leave a poisoned reservation. + with self._lock: + event = self._idempotency_inflight.pop(idempotency_key, None) + if event is not None: + event.set() def _execute(self, run_id: str) -> None: with self._lock: diff --git a/comfy_research/tests/test_run_worker.py b/comfy_research/tests/test_run_worker.py index 369e5c3..85fedb9 100644 --- a/comfy_research/tests/test_run_worker.py +++ b/comfy_research/tests/test_run_worker.py @@ -2,6 +2,7 @@ from __future__ import annotations import time +from concurrent.futures import ThreadPoolExecutor import pytest from fastapi import HTTPException @@ -59,6 +60,20 @@ def test_idempotency_key_returns_same_run() -> None: _wait_terminal(a.run_id) +def test_concurrent_submits_with_same_idempotency_key_single_flight(tmp_path) -> None: + pool = get_worker_pool() + body = TrainRequest.model_validate(minimal_cpu_train_request()) + n = 4 + with ThreadPoolExecutor(max_workers=n) as tpe: + futures = [tpe.submit(pool.submit, body, "concurrent-key") for _ in range(n)] + records = [f.result(timeout=30.0) for f in futures] + run_ids = {rec.run_id for rec in records} + assert len(run_ids) == 1, f"expected a single run_id, got {run_ids}" + run_id = next(iter(run_ids)) + assert _wait_terminal(run_id) == "completed" + assert list(tmp_path.glob("run-*")) == [tmp_path / run_id] + + def test_same_trainer_runs_serialize() -> None: pool = get_worker_pool() body = TrainRequest.model_validate(minimal_cpu_train_request()) From 984b739b7c7c78e02d0f076c7e90e146dcbed80a Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:40:27 +0800 Subject: [PATCH 13/34] feat: add /api/runs router with async submit, query, metrics, and deletes --- comfy_research/api/runs.py | 138 ++++++++++++++++++++++++++ comfy_research/main.py | 8 ++ comfy_research/tests/test_runs_api.py | 112 +++++++++++++++++++++ 3 files changed, 258 insertions(+) create mode 100644 comfy_research/api/runs.py create mode 100644 comfy_research/tests/test_runs_api.py diff --git a/comfy_research/api/runs.py b/comfy_research/api/runs.py new file mode 100644 index 0000000..1c3f039 --- /dev/null +++ b/comfy_research/api/runs.py @@ -0,0 +1,138 @@ +# comfy_research/api/runs.py +"""Run store API: async submit, query, metrics, abort, delete. Files are truth.""" +from __future__ import annotations + +import shutil + +from fastapi import APIRouter, Header, HTTPException, Request + +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.engine.runs.run_worker import get_worker_pool +from comfy_research.schemas.run_record import TERMINAL_STATUSES +from comfy_research.schemas.train_request import TrainRequest + +router = APIRouter(prefix="/api/runs", tags=["runs"]) + +_HYPERPARAM_PREFIX = "hyperparam." + + +def _err(status: int, code: str, detail: str) -> HTTPException: + return HTTPException(status_code=status, detail={"code": code, "detail": detail}) + + +def _query_from_request(request: Request) -> RunQuery: + p = request.query_params + hyper = {k[len(_HYPERPARAM_PREFIX):]: v for k, v in p.items() + if k.startswith(_HYPERPARAM_PREFIX)} + ids = [s for s in (p.get("ids") or "").split(",") if s] or None + return RunQuery( + status=p.get("status"), origin=p.get("origin"), group_id=p.get("group_id"), + since_ms=float(p["since"]) if p.get("since") else None, + ids=ids, hyperparams=hyper, + order_by=p.get("order_by") or "-created_at", + limit=int(p.get("limit") or 100), cursor=p.get("cursor"), + ) + + +@router.post("", status_code=202) +def submit_run(body: TrainRequest, + idempotency_key: str | None = Header(default=None)) -> dict: + rec = get_worker_pool().submit(body, idempotency_key=idempotency_key) + return {"run_id": rec.run_id, "status": rec.status} + + +@router.get("") +def list_runs(request: Request) -> dict: + rows, next_cursor = run_index.query_runs(_query_from_request(request)) + return {"runs": rows, "next_cursor": next_cursor} + + +@router.get("/groups") +def list_groups() -> dict: + return {"groups": run_index.group_summary()} + + +@router.get("/{run_id}") +def get_run(run_id: str) -> dict: + rec = run_store.read_run_record(run_id) + if rec is None: + if (run_store.run_dir(run_id) / "run.json").exists(): + # file present but unparseable: surface it, don't 404 (spec: unreadable) + return {"run_id": run_id, "status": "unreadable", "summary": None} + raise _err(404, "run_not_found", f"No run {run_id!r}.") + _, series = run_store.load_series(run_id) + return {**rec.model_dump(mode="json"), "summary": run_store.summarize(series)} + + +@router.get("/{run_id}/metrics") +def get_run_metrics(run_id: str, downsample: int | None = None) -> dict: + if run_store.read_run_record(run_id) is None: + raise _err(404, "run_not_found", f"No run {run_id!r}.") + source, data = run_store.load_series(run_id) + downsampled = False + if downsample and downsample > 0: + out = {} + for key, series in data.items(): + if isinstance(series, list) and len(series) > downsample: + stride = -(-len(series) // downsample) + out[key] = series[::stride] + downsampled = True + else: + out[key] = series + data = out + return {"source": source, "data": data, "downsampled": downsampled} + + +@router.post("/{run_id}/abort") +def abort_run(run_id: str) -> dict: + ok = get_worker_pool().abort(run_id) + if not ok: + rec = run_store.read_run_record(run_id) + if rec is None: + raise _err(404, "run_not_found", f"No run {run_id!r}.") + return {"ok": False} + return {"ok": True} + + +def _delete_run_files(run_id: str) -> None: + shutil.rmtree(run_store.run_dir(run_id), ignore_errors=True) + run_index.delete_rows([run_id]) + + +_DELETABLE_STATUSES = TERMINAL_STATUSES | {"unreadable"} + + +@router.delete("/{run_id}") +def delete_run(run_id: str) -> dict: + rec = run_store.read_run_record(run_id) + if rec is None: + if (run_store.run_dir(run_id) / "run.json").exists(): + _delete_run_files(run_id) # unreadable: deletable, that's the point + return {"ok": True} + raise _err(404, "run_not_found", f"No run {run_id!r}.") + if rec.status not in _DELETABLE_STATUSES: + raise _err(409, "run_active", "Abort the run before deleting it.") + _delete_run_files(run_id) + return {"ok": True} + + +@router.delete("") +def bulk_delete(request: Request) -> dict: + q = _query_from_request(request) + if not any([q.status, q.origin, q.group_id, q.since_ms, q.ids, q.hyperparams]): + raise _err(400, "filter_required", + "Bulk delete requires at least one filter (status/origin/group_id/since/ids).") + # Collect ALL matching ids first (paginate to the end), then delete the deletable + # ones — deleting while paginating would shift the cursor and skip matches. + q.limit = 500 + to_delete: list[str] = [] + while True: + rows, cursor = run_index.query_runs(q) + to_delete.extend(r["run_id"] for r in rows if r["status"] in _DELETABLE_STATUSES) + if cursor is None: + break + q.cursor = cursor + for run_id in to_delete: + _delete_run_files(run_id) + return {"deleted": len(to_delete)} diff --git a/comfy_research/main.py b/comfy_research/main.py index b1c2184..b140ab4 100644 --- a/comfy_research/main.py +++ b/comfy_research/main.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from contextlib import asynccontextmanager from pathlib import Path @@ -24,6 +25,7 @@ parametric_path_sampler, pca, predict, + runs, svd, train, user_linear_datasets, @@ -84,6 +86,11 @@ async def _lifespan(_: FastAPI): "After editing frontend/src/, run: cd frontend && npm run build — then reload the tab (Cmd+Shift+R). " "For live dev, run COMFYRESEARCH_PORT= npm run dev in frontend/ and open http://127.0.0.1:5173.", ) + try: + from comfy_research.engine.runs.run_index import reconcile_stale_running + reconcile_stale_running() + except Exception: # never block startup on store recovery + logging.getLogger(__name__).warning("run store reconciliation failed", exc_info=True) yield @@ -112,6 +119,7 @@ def create_app() -> FastAPI: app.include_router(user_linear_datasets.router) app.include_router(user_symbolic_func_datasets.router) app.include_router(train.router) + app.include_router(runs.router) app.include_router(collect.router) app.include_router(parametric_path_sampler.router) app.include_router(activation_tensor.router) diff --git a/comfy_research/tests/test_runs_api.py b/comfy_research/tests/test_runs_api.py new file mode 100644 index 0000000..663896c --- /dev/null +++ b/comfy_research/tests/test_runs_api.py @@ -0,0 +1,112 @@ +# comfy_research/tests/test_runs_api.py +from __future__ import annotations + +import time + +import pytest +from fastapi.testclient import TestClient + +from comfy_research.main import app +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_worker import reset_worker_pool_for_tests +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms +from comfy_research.tests.train_test_fixtures import minimal_cpu_train_request + + +@pytest.fixture(autouse=True) +def _isolated(tmp_path, monkeypatch): + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + reset_worker_pool_for_tests() + yield + reset_worker_pool_for_tests() + + +def _seed_terminal(status: str = "completed", group: str | None = None) -> str: + rec = RunRecord( + run_id=new_run_id(), origin="agent", status=status, created_at=now_ms(), + finished_at=now_ms(), group_id=group, trainer_node_id="t1", + graph=GraphDocument(version=1, nodes=[], edges=[]), + hyperparams={"t1.trainingSteps": 4}, + ) + run_store.write_run_record(rec) + run_store.write_results(rec.run_id, {"loss_history": [1.0, 0.5], + "test_loss_history": [], "step_ticks": [0, 1]}) + run_index.upsert_run(rec, {"final_loss": 0.5, "final_test_loss": None, + "best_test_loss": None, "steps_completed": 2}) + return rec.run_id + + +def _wait_terminal(client: TestClient, run_id: str, timeout_s: float = 30.0) -> str: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + status = client.get(f"/api/runs/{run_id}").json()["status"] + if status not in ("queued", "running"): + return status + time.sleep(0.05) + raise AssertionError("never terminal") + + +def test_submit_then_poll_lifecycle() -> None: + client = TestClient(app) + resp = client.post("/api/runs", json={**minimal_cpu_train_request(), "run_origin": "agent"}) + assert resp.status_code == 202 + run_id = resp.json()["run_id"] + assert _wait_terminal(client, run_id) == "completed" + metrics = client.get(f"/api/runs/{run_id}/metrics").json() + assert metrics["source"] == "results" + assert len(metrics["data"]["loss_history"]) == 4 + + +def test_submit_invalid_graph_400() -> None: + resp = TestClient(app).post("/api/runs", json={ + "trainer_node_id": "trainer", + "nodes": [{"id": "trainer", "type": "trainer", + "data": {"trainingSteps": 1, "computeDevice": "cpu"}}], + "edges": [], + }) + assert resp.status_code == 400 + + +def test_list_filters_and_hyperparam() -> None: + client = TestClient(app) + a = _seed_terminal("completed", group="g1") + _seed_terminal("failed", group="g1") + body = client.get("/api/runs", params={"status": "completed"}).json() + assert [r["run_id"] for r in body["runs"]] == [a] + body = client.get("/api/runs", params={"hyperparam.t1.trainingSteps": "4"}).json() + assert len(body["runs"]) == 2 + groups = client.get("/api/runs/groups").json()["groups"] + assert groups[0]["group_id"] == "g1" and groups[0]["counts"]["failed"] == 1 + + +def test_get_missing_run_404_with_code() -> None: + resp = TestClient(app).get("/api/runs/run-doesnotexist") + assert resp.status_code == 404 + assert resp.json()["detail"]["code"] == "run_not_found" + + +def test_delete_guards() -> None: + client = TestClient(app) + rid = _seed_terminal("completed") + assert client.delete("/api/runs").status_code == 400 + running = RunRecord( + run_id=new_run_id(), origin="agent", status="running", created_at=now_ms(), + trainer_node_id="t1", graph=GraphDocument(version=1, nodes=[], edges=[]), + ) + run_store.write_run_record(running) + run_index.upsert_run(running) + run_index.touch_heartbeat(running.run_id, now_ms()) + assert client.delete(f"/api/runs/{running.run_id}").status_code == 409 + assert client.delete(f"/api/runs/{rid}").status_code == 200 + assert run_store.read_run_record(rid) is None + resp = client.delete("/api/runs", params={"status": "failed"}) + assert resp.status_code == 200 and resp.json()["deleted"] == 0 + + +def test_bulk_delete_by_group() -> None: + client = TestClient(app) + _seed_terminal("completed", group="g2") + _seed_terminal("aborted", group="g2") + resp = client.delete("/api/runs", params={"group_id": "g2"}) + assert resp.json()["deleted"] == 2 From 5174c4a5d3ce4e62e2e8dbec1b4839620be9bc3b Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:44:28 +0800 Subject: [PATCH 14/34] fix: structured 400 for malformed run query params --- comfy_research/api/runs.py | 24 ++++++++++++++++++++++-- comfy_research/tests/test_runs_api.py | 14 ++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/comfy_research/api/runs.py b/comfy_research/api/runs.py index 1c3f039..f4584cc 100644 --- a/comfy_research/api/runs.py +++ b/comfy_research/api/runs.py @@ -21,6 +21,26 @@ def _err(status: int, code: str, detail: str) -> HTTPException: return HTTPException(status_code=status, detail={"code": code, "detail": detail}) +def _parse_float_param(p, name: str) -> float | None: + raw = p.get(name) + if not raw: + return None + try: + return float(raw) + except ValueError: + raise _err(400, "invalid_query_param", f"{name}={raw!r} is not a valid number.") + + +def _parse_int_param(p, name: str, default: int) -> int: + raw = p.get(name) + if not raw: + return default + try: + return int(raw) + except ValueError: + raise _err(400, "invalid_query_param", f"{name}={raw!r} is not a valid integer.") + + def _query_from_request(request: Request) -> RunQuery: p = request.query_params hyper = {k[len(_HYPERPARAM_PREFIX):]: v for k, v in p.items() @@ -28,10 +48,10 @@ def _query_from_request(request: Request) -> RunQuery: ids = [s for s in (p.get("ids") or "").split(",") if s] or None return RunQuery( status=p.get("status"), origin=p.get("origin"), group_id=p.get("group_id"), - since_ms=float(p["since"]) if p.get("since") else None, + since_ms=_parse_float_param(p, "since"), ids=ids, hyperparams=hyper, order_by=p.get("order_by") or "-created_at", - limit=int(p.get("limit") or 100), cursor=p.get("cursor"), + limit=_parse_int_param(p, "limit", 100), cursor=p.get("cursor"), ) diff --git a/comfy_research/tests/test_runs_api.py b/comfy_research/tests/test_runs_api.py index 663896c..84a0caf 100644 --- a/comfy_research/tests/test_runs_api.py +++ b/comfy_research/tests/test_runs_api.py @@ -86,6 +86,20 @@ def test_get_missing_run_404_with_code() -> None: assert resp.json()["detail"]["code"] == "run_not_found" +def test_malformed_query_params_return_structured_400() -> None: + client = TestClient(app) + resp = client.get("/api/runs", params={"limit": "abc"}) + assert resp.status_code == 400 + assert resp.json()["detail"]["code"] == "invalid_query_param" + + resp = client.get("/api/runs", params={"since": "xyz"}) + assert resp.status_code == 400 + assert resp.json()["detail"]["code"] == "invalid_query_param" + + resp = client.delete("/api/runs", params={"since": "xyz"}) + assert resp.status_code == 400 + + def test_delete_guards() -> None: client = TestClient(app) rid = _seed_terminal("completed") From 2ffe9d558b6e1444d43928713ca617c3c93c8aaa Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:46:59 +0800 Subject: [PATCH 15/34] feat: add single-owner run GC with agent retention policy --- comfy_research/engine/runs/run_gc.py | 59 ++++++++++++++++++++++++ comfy_research/engine/runs/run_worker.py | 8 +++- comfy_research/main.py | 2 + comfy_research/tests/test_run_gc.py | 45 ++++++++++++++++++ 4 files changed, 113 insertions(+), 1 deletion(-) create mode 100644 comfy_research/engine/runs/run_gc.py create mode 100644 comfy_research/tests/test_run_gc.py diff --git a/comfy_research/engine/runs/run_gc.py b/comfy_research/engine/runs/run_gc.py new file mode 100644 index 0000000..455deab --- /dev/null +++ b/comfy_research/engine/runs/run_gc.py @@ -0,0 +1,59 @@ +"""Single-owner retention GC for agent-origin runs. Runs only in the API server process.""" +from __future__ import annotations + +import json +import logging +import shutil + +from comfy_research.engine.runs import run_index, run_store +from comfy_research.engine.runs.run_index import RunQuery +from comfy_research.schemas.run_record import TERMINAL_STATUSES, now_ms + +logger = logging.getLogger(__name__) + +_GRACE_MS = 10 * 60 * 1000.0 +_DEFAULTS = {"max_runs_agent": 2000, "max_age_days_agent": None, "worker_slots": 2} + + +def load_gc_config() -> dict: + path = run_store.runs_root() / "config.json" + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + raw = {} + return {**_DEFAULTS, **{k: raw[k] for k in _DEFAULTS if k in raw}} + + +def run_gc_once() -> list[str]: + cfg = load_gc_config() + rows: list[dict] = [] + q = RunQuery(origin="agent", order_by="-created_at", limit=500) + while True: + page, cursor = run_index.query_runs(q) + rows.extend(page) + if cursor is None: + break + q.cursor = cursor + cutoff_ms = now_ms() - _GRACE_MS + max_age = cfg["max_age_days_agent"] + age_cutoff = now_ms() - max_age * 86_400_000.0 if max_age else None + + def prunable(r: dict) -> bool: + return (r["status"] in TERMINAL_STATUSES + and r["finished_at"] is not None and r["finished_at"] < cutoff_ms) + + pruned: list[str] = [] + over_cap = rows[cfg["max_runs_agent"]:] # rows are newest-first + for r in over_cap: + if prunable(r): + pruned.append(r["run_id"]) + if age_cutoff is not None: + for r in rows[: cfg["max_runs_agent"]]: + if prunable(r) and r["finished_at"] < age_cutoff: + pruned.append(r["run_id"]) + for run_id in pruned: + shutil.rmtree(run_store.run_dir(run_id), ignore_errors=True) + run_index.delete_rows(pruned) + for run_id in pruned: + logger.info("run GC pruned %s", run_id) + return pruned diff --git a/comfy_research/engine/runs/run_worker.py b/comfy_research/engine/runs/run_worker.py index cae0f21..8b3fd15 100644 --- a/comfy_research/engine/runs/run_worker.py +++ b/comfy_research/engine/runs/run_worker.py @@ -123,6 +123,11 @@ def _execute(self, run_id: str) -> None: writer.finalize("failed", error_detail=f"{type(exc).__name__}: {exc}") finally: self._dispatch_next(body.trainer_node_id) + try: + from comfy_research.engine.runs.run_gc import run_gc_once + run_gc_once() + except Exception: + logger.warning("post-run GC failed", exc_info=True) def _dispatch_next(self, trainer_id: str) -> None: with self._lock: @@ -172,7 +177,8 @@ def get_worker_pool() -> RunWorkerPool: global _pool with _pool_lock: if _pool is None: - _pool = RunWorkerPool() + from comfy_research.engine.runs.run_gc import load_gc_config + _pool = RunWorkerPool(slots=load_gc_config()["worker_slots"]) return _pool diff --git a/comfy_research/main.py b/comfy_research/main.py index b140ab4..1cc902a 100644 --- a/comfy_research/main.py +++ b/comfy_research/main.py @@ -89,6 +89,8 @@ async def _lifespan(_: FastAPI): try: from comfy_research.engine.runs.run_index import reconcile_stale_running reconcile_stale_running() + from comfy_research.engine.runs.run_gc import run_gc_once + run_gc_once() except Exception: # never block startup on store recovery logging.getLogger(__name__).warning("run store reconciliation failed", exc_info=True) yield diff --git a/comfy_research/tests/test_run_gc.py b/comfy_research/tests/test_run_gc.py new file mode 100644 index 0000000..31220cd --- /dev/null +++ b/comfy_research/tests/test_run_gc.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from comfy_research.engine.runs import run_gc, run_index, run_store +from comfy_research.schemas.graph import GraphDocument +from comfy_research.schemas.run_record import RunRecord, new_run_id, now_ms + + +def _seed(status: str, origin: str, finished_ms_ago: float) -> str: + rec = RunRecord( + run_id=new_run_id(), origin=origin, status=status, + created_at=now_ms() - finished_ms_ago - 1000, + finished_at=(now_ms() - finished_ms_ago) if status != "running" else None, + trainer_node_id="t1", graph=GraphDocument(version=1, nodes=[], edges=[]), + ) + run_store.write_run_record(rec) + run_index.upsert_run(rec) + return rec.run_id + + +def test_gc_prunes_only_old_terminal_agent_runs(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + (tmp_path / "config.json").write_text('{"max_runs_agent": 2}', encoding="utf-8") + hour = 3_600_000.0 + keep_human = _seed("completed", "human", 10 * hour) + keep_recent = _seed("completed", "agent", 0.0) # inside grace period + keep_running = _seed("running", "agent", 5 * hour) + newest = _seed("completed", "agent", 1 * hour) + older = _seed("completed", "agent", 2 * hour) + oldest = _seed("completed", "agent", 3 * hour) + + pruned = run_gc.run_gc_once() + # agent rows newest-first by created_at: [keep_recent, newest, older, oldest, keep_running]; + # cap 2 keeps [keep_recent, newest]; over-cap = [older, oldest, keep_running], of which + # only terminal runs past the 10-minute grace period are prunable. + assert set(pruned) == {older, oldest} + for rid in (keep_human, keep_recent, keep_running, newest): + assert run_store.read_run_record(rid) is not None + for rid in pruned: + assert run_store.read_run_record(rid) is None + + +def test_gc_default_config_noop_under_cap(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("COMFYRESEARCH_RUNS_DIR", str(tmp_path)) + _seed("completed", "agent", 3_600_000.0) + assert run_gc.run_gc_once() == [] From 575e8957017470473516410ed1856e81151beb7a Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:53:11 +0800 Subject: [PATCH 16/34] feat: add Runs rail panel with grouped run list --- frontend/src/components/LeftNavRail.tsx | 7 ++ frontend/src/components/ResearchCanvas.tsx | 11 ++ frontend/src/components/RunsPanel.tsx | 117 ++++++++++++++++++ frontend/src/components/railTypes.ts | 2 +- .../__tests__/leftNavRail.v1.seam.test.tsx | 2 +- .../src/graph/__tests__/runsPanel.test.tsx | 53 ++++++++ frontend/src/graph/runsApi.ts | 47 +++++++ frontend/src/index.css | 31 +++++ 8 files changed, 268 insertions(+), 2 deletions(-) create mode 100644 frontend/src/components/RunsPanel.tsx create mode 100644 frontend/src/graph/__tests__/runsPanel.test.tsx create mode 100644 frontend/src/graph/runsApi.ts diff --git a/frontend/src/components/LeftNavRail.tsx b/frontend/src/components/LeftNavRail.tsx index 40dd615..47d5018 100644 --- a/frontend/src/components/LeftNavRail.tsx +++ b/frontend/src/components/LeftNavRail.tsx @@ -12,6 +12,7 @@ const primaryItems: RailItem[] = [ { id: "nodes", label: "Nodes" }, { id: "observables", label: "Observables" }, { id: "templates", label: "Templates" }, + { id: "runs", label: "Runs" }, ]; const THEME_OPTIONS: { id: CrTheme; label: string }[] = [ @@ -44,6 +45,12 @@ function RailIcon({ name }: { name: string }) { ); + case "runs": + return ( + + + + ); case "settings": return ( diff --git a/frontend/src/components/ResearchCanvas.tsx b/frontend/src/components/ResearchCanvas.tsx index 2aa9758..5e45c8d 100644 --- a/frontend/src/components/ResearchCanvas.tsx +++ b/frontend/src/components/ResearchCanvas.tsx @@ -262,6 +262,7 @@ import { import { ObservablePanel } from "../observables/ObservablePanel"; import { isObservableModelNodeType } from "../observables/modelNodeTypes"; import { researchNodeTypes } from "./nodeTypes"; +import { RunsPanel } from "./RunsPanel"; import { migrateObservableVizNodeTypes } from "../graph/observableVizVariant"; import { beginLibraryNodeDrag, @@ -5799,6 +5800,11 @@ export function ResearchCanvas() { [], ); + const openRunGraphInNewProject = useCallback((runId: string) => { + // TODO(Task 11): fetch the run record and open its graph in a new project tab. + console.warn("open run", runId); + }, []); + useEffect(() => { if (typeof window === "undefined") return; const params = new URLSearchParams(window.location.search); @@ -6041,6 +6047,11 @@ export function ResearchCanvas() { /> ) : null} + {railSection === "runs" ? ( +
+ +
+ ) : null}
(); + for (const row of rows) { + if (row.group_id) { + const list = byGroup.get(row.group_id) ?? []; + list.push(row); + byGroup.set(row.group_id, list); + } else { + ungrouped.push(row); + } + } + const groups: RunGroup[] = []; + if (ungrouped.length) { + groups.push({ key: "", label: "Runs", rows: ungrouped, collapsedByDefault: false }); + } + for (const [key, groupRows] of byGroup) { + const collapsed = groupRows.every((r) => r.origin === "sweep" || r.origin === "agent"); + groups.push({ key, label: key, rows: groupRows, collapsedByDefault: collapsed }); + } + return groups; +} + +export function RunsPanel({ + onOpenRunGraph, + selectedRunIds, + onToggleSelect, +}: { + onOpenRunGraph: (runId: string) => void; + selectedRunIds?: ReadonlySet; + onToggleSelect?: (runId: string) => void; +}) { + const [rows, setRows] = useState([]); + const [error, setError] = useState(null); + const [expanded, setExpanded] = useState>({}); + + const load = useCallback(() => { + fetchRuns({ limit: "200" }) + .then((body) => { setRows(body.runs); setError(null); }) + .catch((e: Error) => setError(e.message)); + }, []); + + useEffect(() => { + load(); + }, [load]); + + const anyActive = useMemo(() => rows.some((r) => ACTIVE_STATUSES.has(r.status)), [rows]); + useEffect(() => { + if (!anyActive) return; + const id = window.setInterval(load, POLL_MS); + return () => window.clearInterval(id); + }, [anyActive, load]); + + const groups = useMemo(() => groupRuns(rows), [rows]); + + return ( + + ); +} diff --git a/frontend/src/components/railTypes.ts b/frontend/src/components/railTypes.ts index 608f5c0..c2ccf44 100644 --- a/frontend/src/components/railTypes.ts +++ b/frontend/src/components/railTypes.ts @@ -1,2 +1,2 @@ /** Primary rail destinations that swap the left sidebar panel. */ -export type RailPrimarySection = "nodes" | "observables" | "templates"; +export type RailPrimarySection = "nodes" | "observables" | "templates" | "runs"; diff --git a/frontend/src/graph/__tests__/leftNavRail.v1.seam.test.tsx b/frontend/src/graph/__tests__/leftNavRail.v1.seam.test.tsx index 6286b03..c62693c 100644 --- a/frontend/src/graph/__tests__/leftNavRail.v1.seam.test.tsx +++ b/frontend/src/graph/__tests__/leftNavRail.v1.seam.test.tsx @@ -52,7 +52,7 @@ beforeEach(() => { delete document.documentElement.dataset.crTheme; }); -const V1_RAIL_LABELS = ["Nodes", "Observables", "Templates"]; +const V1_RAIL_LABELS = ["Nodes", "Observables", "Templates", "Runs"]; function renderRail(host: HTMLElement) { const root = createRoot(host); diff --git a/frontend/src/graph/__tests__/runsPanel.test.tsx b/frontend/src/graph/__tests__/runsPanel.test.tsx new file mode 100644 index 0000000..3a6a743 --- /dev/null +++ b/frontend/src/graph/__tests__/runsPanel.test.tsx @@ -0,0 +1,53 @@ +// @vitest-environment jsdom +import { act } from "react"; +import { createRoot } from "react-dom/client"; +import { afterEach, beforeEach, expect, test, vi } from "vitest"; + +import { RunsPanel } from "../../components/RunsPanel"; + +(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true; + +const RUNS = { + runs: [ + { run_id: "run-aaa", group_id: null, origin: "human", status: "completed", + created_at: 1, finished_at: 2, trainer_node_id: "t1", final_loss: 0.25, + final_test_loss: null, steps_completed: 4, duration_seconds: 1.5, hyperparams: {} }, + { run_id: "run-bbb", group_id: "sweep-1", origin: "sweep", status: "failed", + created_at: 3, finished_at: 4, trainer_node_id: "t1", final_loss: null, + final_test_loss: null, steps_completed: 0, duration_seconds: null, hyperparams: {} }, + ], + next_cursor: null, +}; + +let host: HTMLDivElement; + +beforeEach(() => { + host = document.createElement("div"); + document.body.appendChild(host); + vi.stubGlobal("fetch", vi.fn(async () => ({ + ok: true, + json: async () => RUNS, + })) as unknown as typeof fetch); +}); + +afterEach(() => { + vi.unstubAllGlobals(); + host.remove(); +}); + +test("renders run rows grouped, sweep group collapsed by default", async () => { + const root = createRoot(host); + await act(async () => { + root.render( {}} />); + }); + await act(async () => { await Promise.resolve(); }); + expect(host.textContent).toContain("run-aaa"); + expect(host.textContent).toContain("sweep-1"); + // collapsed group hides its member row until expanded + expect(host.textContent).not.toContain("run-bbb"); + const toggle = host.querySelector('[data-testid="run-group-toggle-sweep-1"]'); + expect(toggle).not.toBeNull(); + await act(async () => { toggle!.click(); }); + expect(host.textContent).toContain("run-bbb"); + await act(async () => { root.unmount(); }); +}); diff --git a/frontend/src/graph/runsApi.ts b/frontend/src/graph/runsApi.ts new file mode 100644 index 0000000..37132cb --- /dev/null +++ b/frontend/src/graph/runsApi.ts @@ -0,0 +1,47 @@ +import type { GraphDocument } from "../types/graph"; + +export type RunRow = { + run_id: string; + group_id: string | null; + origin: string; + status: string; + created_at: number; + finished_at: number | null; + trainer_node_id: string; + final_loss: number | null; + final_test_loss: number | null; + steps_completed: number | null; + duration_seconds: number | null; + hyperparams: Record; +}; + +async function readJson(res: Response): Promise { + if (!res.ok) { + const text = await res.text().catch(() => ""); + throw new Error(text || res.statusText); + } + return (await res.json()) as T; +} + +export async function fetchRuns( + params?: Record, +): Promise<{ runs: RunRow[]; next_cursor: string | null }> { + const qs = params ? `?${new URLSearchParams(params)}` : ""; + return readJson(await fetch(`/api/runs${qs}`, { cache: "no-store" })); +} + +export async function fetchRunMetrics( + runId: string, +): Promise<{ source: string; data: Record }> { + return readJson(await fetch(`/api/runs/${runId}/metrics`, { cache: "no-store" })); +} + +export async function fetchRunRecord( + runId: string, +): Promise<{ graph: GraphDocument } & Record> { + return readJson(await fetch(`/api/runs/${runId}`, { cache: "no-store" })); +} + +export async function deleteRun(runId: string): Promise { + await readJson(await fetch(`/api/runs/${runId}`, { method: "DELETE" })); +} diff --git a/frontend/src/index.css b/frontend/src/index.css index 0b9e4a3..719516c 100644 --- a/frontend/src/index.css +++ b/frontend/src/index.css @@ -1556,6 +1556,37 @@ html.cr-node-over-library-delete .cr-nodes-panel::after { letter-spacing: 0.02em; } +.cr-runs-panel__group { + display: flex; + flex-direction: column; +} + +.cr-runs-panel__row { + display: flex; + gap: 6px; + align-items: center; +} + +.cr-runs-panel__status-dot { + width: 8px; + height: 8px; + border-radius: 50%; + background: var(--cr-text-4); +} + +[data-status="completed"].cr-runs-panel__status-dot { + background: var(--cr-chart-2); +} + +[data-status="failed"].cr-runs-panel__status-dot, +[data-status="crashed"].cr-runs-panel__status-dot { + background: var(--cr-chart-1); +} + +.cr-runs-panel__error { + color: var(--cr-chart-1); +} + .cr-nodes-panel__search-row { display: flex; align-items: center; From ae915af2552f1a8140bde7a02ba8b9b732ab29d4 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 18:57:31 +0800 Subject: [PATCH 17/34] fix: surface run delete failures in the Runs panel --- frontend/src/components/RunsPanel.tsx | 6 ++- .../src/graph/__tests__/runsPanel.test.tsx | 37 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/frontend/src/components/RunsPanel.tsx b/frontend/src/components/RunsPanel.tsx index 7c6abd1..5c605ac 100644 --- a/frontend/src/components/RunsPanel.tsx +++ b/frontend/src/components/RunsPanel.tsx @@ -102,7 +102,11 @@ export function RunsPanel({ {row.final_loss != null ? row.final_loss.toPrecision(3) : "—"} {TERMINAL_DELETABLE.has(row.status) ? ( ) : null} diff --git a/frontend/src/graph/__tests__/runsPanel.test.tsx b/frontend/src/graph/__tests__/runsPanel.test.tsx index 3a6a743..a182beb 100644 --- a/frontend/src/graph/__tests__/runsPanel.test.tsx +++ b/frontend/src/graph/__tests__/runsPanel.test.tsx @@ -51,3 +51,40 @@ test("renders run rows grouped, sweep group collapsed by default", async () => { expect(host.textContent).toContain("run-bbb"); await act(async () => { root.unmount(); }); }); + +test("surfaces an error and keeps the row when deleting a run fails", async () => { + vi.stubGlobal("fetch", vi.fn(async (_input: RequestInfo | URL, init?: RequestInit) => { + if (init?.method === "DELETE") { + return { + ok: false, + status: 500, + statusText: "Internal Server Error", + text: async () => "boom", + json: async () => ({}), + }; + } + return { ok: true, json: async () => RUNS }; + }) as unknown as typeof fetch); + + const root = createRoot(host); + await act(async () => { + root.render( {}} />); + }); + await act(async () => { await Promise.resolve(); }); + + const deleteBtn = host.querySelector('[aria-label="Delete run-aaa"]'); + expect(deleteBtn).not.toBeNull(); + await act(async () => { + deleteBtn!.click(); + await Promise.resolve(); + await Promise.resolve(); + }); + + const errorEl = host.querySelector(".cr-runs-panel__error"); + expect(errorEl).not.toBeNull(); + expect(errorEl!.textContent).toContain("boom"); + // the row must still be present since the delete failed + expect(host.textContent).toContain("run-aaa"); + + await act(async () => { root.unmount(); }); +}); From 90a9f04f42a10e59789157cb83eaef6595cf8cd3 Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 19:02:21 +0800 Subject: [PATCH 18/34] feat: add run comparison overlay chart and open-run-graph action --- frontend/src/components/ResearchCanvas.tsx | 26 +++++- frontend/src/components/RunsPanel.tsx | 87 ++++++++++++++++--- .../graph/__tests__/runCompareOverlay.test.ts | 25 ++++++ frontend/src/graph/runCompareOverlay.ts | 35 ++++++++ frontend/src/graph/sweepVizPlot.ts | 2 +- 5 files changed, 158 insertions(+), 17 deletions(-) create mode 100644 frontend/src/graph/__tests__/runCompareOverlay.test.ts create mode 100644 frontend/src/graph/runCompareOverlay.ts diff --git a/frontend/src/components/ResearchCanvas.tsx b/frontend/src/components/ResearchCanvas.tsx index 5e45c8d..941eee6 100644 --- a/frontend/src/components/ResearchCanvas.tsx +++ b/frontend/src/components/ResearchCanvas.tsx @@ -263,6 +263,7 @@ import { ObservablePanel } from "../observables/ObservablePanel"; import { isObservableModelNodeType } from "../observables/modelNodeTypes"; import { researchNodeTypes } from "./nodeTypes"; import { RunsPanel } from "./RunsPanel"; +import { fetchRunRecord } from "../graph/runsApi"; import { migrateObservableVizNodeTypes } from "../graph/observableVizVariant"; import { beginLibraryNodeDrag, @@ -5801,8 +5802,29 @@ export function ResearchCanvas() { ); const openRunGraphInNewProject = useCallback((runId: string) => { - // TODO(Task 11): fetch the run record and open its graph in a new project tab. - console.warn("open run", runId); + void fetchRunRecord(runId).then((record) => { + const { nodes, edges } = sanitizeLoadedGraph(record.graph); + const id = newProjectId(); + const canvasId = newProjectId(); + setProjects((list) => [ + ...list, + { + id, + title: `Run ${runId}`, + canvas: { + id: canvasId, + title: `Run ${runId}`, + nodes, + edges, + savedViewport: record.graph.viewport ?? null, + viewportApplyNonce: 1, + dirty: false, + }, + }, + ]); + setActiveProjectId(id); + setNotice(null); + }); }, []); useEffect(() => { diff --git a/frontend/src/components/RunsPanel.tsx b/frontend/src/components/RunsPanel.tsx index 5c605ac..3297ff5 100644 --- a/frontend/src/components/RunsPanel.tsx +++ b/frontend/src/components/RunsPanel.tsx @@ -1,6 +1,8 @@ -import { useCallback, useEffect, useMemo, useState } from "react"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { deleteRun, fetchRuns, type RunRow } from "../graph/runsApi"; +import { buildRunCompareSeries } from "../graph/runCompareOverlay"; +import { deleteRun, fetchRunMetrics, fetchRuns, type RunRow } from "../graph/runsApi"; +import { SweepVizLinePlot } from "./nodes/SweepVizLinePlot"; const ACTIVE_STATUSES = new Set(["queued", "running"]); const POLL_MS = 3000; @@ -33,16 +35,27 @@ function groupRuns(rows: RunRow[]): RunGroup[] { export function RunsPanel({ onOpenRunGraph, - selectedRunIds, - onToggleSelect, }: { onOpenRunGraph: (runId: string) => void; - selectedRunIds?: ReadonlySet; - onToggleSelect?: (runId: string) => void; }) { const [rows, setRows] = useState([]); const [error, setError] = useState(null); const [expanded, setExpanded] = useState>({}); + const [selectedRunIds, setSelectedRunIds] = useState>(new Set()); + const [compareSeries, setCompareSeries] = useState>([]); + const metricsCacheRef = useRef(new Map>()); + + const onToggleSelect = useCallback((runId: string) => { + setSelectedRunIds((prev) => { + const next = new Set(prev); + if (next.has(runId)) { + next.delete(runId); + } else { + next.add(runId); + } + return next; + }); + }, []); const load = useCallback(() => { fetchRuns({ limit: "200" }) @@ -63,6 +76,38 @@ export function RunsPanel({ const groups = useMemo(() => groupRuns(rows), [rows]); + useEffect(() => { + if (selectedRunIds.size === 0) { + setCompareSeries([]); + return; + } + let cancelled = false; + const ids = Array.from(selectedRunIds); + void Promise.all( + ids.map(async (runId) => { + const cached = metricsCacheRef.current.get(runId); + if (cached) return { runId, data: cached }; + const res = await fetchRunMetrics(runId); + metricsCacheRef.current.set(runId, res.data); + return { runId, data: res.data }; + }), + ) + .then((results) => { + if (cancelled) return; + setCompareSeries( + buildRunCompareSeries( + results.map((r) => ({ runId: r.runId, label: r.runId, data: r.data })), + ), + ); + }) + .catch((e: Error) => { + if (!cancelled) setError(e.message); + }); + return () => { + cancelled = true; + }; + }, [selectedRunIds]); + return ( ); } diff --git a/frontend/src/graph/__tests__/runCompareOverlay.test.ts b/frontend/src/graph/__tests__/runCompareOverlay.test.ts new file mode 100644 index 0000000..2c4bdeb --- /dev/null +++ b/frontend/src/graph/__tests__/runCompareOverlay.test.ts @@ -0,0 +1,25 @@ +import { expect, test } from "vitest"; + +import { buildRunCompareSeries } from "../runCompareOverlay"; + +test("one solid series per run, dashed test series, distinct colors", () => { + const series = buildRunCompareSeries([ + { runId: "run-a", label: "run-a", + data: { loss_history: [1, 0.5], test_loss_history: [2, 1], step_ticks: [0, 1] } }, + { runId: "run-b", label: "run-b", + data: { loss_history: [3, 2, 1], test_loss_history: [], step_ticks: [0, 1, 2] } }, + ]); + expect(series.map((s) => s.id)).toEqual(["run-a", "run-a:test", "run-b"]); + expect(series[0].points.map((p) => p.y)).toEqual([1, 0.5]); + expect(series[0].points.map((p) => p.x)).toEqual([0, 1]); + expect(series[1].strokeDasharray).toBeTruthy(); + expect(series[1].color).toBe(series[0].color); + expect(series[2].color).not.toBe(series[0].color); +}); + +test("falls back to index when step_ticks missing", () => { + const [s] = buildRunCompareSeries([ + { runId: "run-c", label: "run-c", data: { loss_history: [5, 4], step_ticks: [] } }, + ]); + expect(s.points.map((p) => p.x)).toEqual([0, 1]); +}); diff --git a/frontend/src/graph/runCompareOverlay.ts b/frontend/src/graph/runCompareOverlay.ts new file mode 100644 index 0000000..20ee3ae --- /dev/null +++ b/frontend/src/graph/runCompareOverlay.ts @@ -0,0 +1,35 @@ +import { SERIES_COLORS, type PlotPoint, type PlotSeries } from "./sweepVizPlot"; + +type RunSeriesInput = { runId: string; label: string; data: Record }; + +function toPoints(ys: number[], xs: number[], rowId: string): PlotPoint[] { + return ys + .map((y, i) => ({ x: xs[i] ?? i, xDisplay: String(xs[i] ?? i), y, rowId })) + .filter((p) => Number.isFinite(p.y)); +} + +export function buildRunCompareSeries(inputs: RunSeriesInput[]): PlotSeries[] { + const series: PlotSeries[] = []; + inputs.forEach((input, i) => { + const color = SERIES_COLORS[i % SERIES_COLORS.length]!; + const steps = input.data.step_ticks ?? []; + const loss = input.data.loss_history ?? []; + series.push({ + id: input.runId, + label: input.label, + color, + points: toPoints(loss, steps, input.runId), + }); + const test = input.data.test_loss_history ?? []; + if (test.length) { + series.push({ + id: `${input.runId}:test`, + label: `${input.label} (test)`, + color, + strokeDasharray: "4 3", + points: toPoints(test, steps, input.runId), + }); + } + }); + return series; +} diff --git a/frontend/src/graph/sweepVizPlot.ts b/frontend/src/graph/sweepVizPlot.ts index 3165670..2db548b 100644 --- a/frontend/src/graph/sweepVizPlot.ts +++ b/frontend/src/graph/sweepVizPlot.ts @@ -106,7 +106,7 @@ export function dualAxisWarranted(series: PlotSeries[]): boolean { /* Theme-aware: resolved per active theme by tokens.css (classic values are * the exact legacy hex palette). Consumed as SVG inline styles, which * resolve var() at render time. */ -const SERIES_COLORS = [ +export const SERIES_COLORS = [ "var(--cr-chart-1)", "var(--cr-chart-2)", "var(--cr-chart-3)", From ff78424c34fe418ca6826cbb9b910b05cd0df61a Mon Sep 17 00:00:00 2001 From: abrohamLee Date: Sat, 15 Aug 2026 19:13:27 +0800 Subject: [PATCH 19/34] fix: surface run-open failures and refresh live-run metrics in compare view --- frontend/src/components/ResearchCanvas.tsx | 2 + frontend/src/components/RunsPanel.tsx | 19 ++++- .../researchCanvasOpenRunGraph.wiring.test.ts | 50 +++++++++++++ .../src/graph/__tests__/runsPanel.test.tsx | 73 +++++++++++++++++++ 4 files changed, 140 insertions(+), 4 deletions(-) create mode 100644 frontend/src/graph/__tests__/researchCanvasOpenRunGraph.wiring.test.ts diff --git a/frontend/src/components/ResearchCanvas.tsx b/frontend/src/components/ResearchCanvas.tsx index 941eee6..e28cc42 100644 --- a/frontend/src/components/ResearchCanvas.tsx +++ b/frontend/src/components/ResearchCanvas.tsx @@ -5824,6 +5824,8 @@ export function ResearchCanvas() { ]); setActiveProjectId(id); setNotice(null); + }).catch((e: unknown) => { + setNotice(e instanceof Error ? e.message : "Could not open run graph."); }); }, []); diff --git a/frontend/src/components/RunsPanel.tsx b/frontend/src/components/RunsPanel.tsx index 3297ff5..2c981dc 100644 --- a/frontend/src/components/RunsPanel.tsx +++ b/frontend/src/components/RunsPanel.tsx @@ -85,10 +85,19 @@ export function RunsPanel({ const ids = Array.from(selectedRunIds); void Promise.all( ids.map(async (runId) => { - const cached = metricsCacheRef.current.get(runId); - if (cached) return { runId, data: cached }; + // Live runs (queued/running) must not serve stale cached metrics: their + // loss history keeps growing while the row list polls, so always refetch. + // Terminal runs' metrics are immutable once finished, so caching is safe. + const row = rows.find((r) => r.run_id === runId); + const isLive = row != null && ACTIVE_STATUSES.has(row.status); + if (!isLive) { + const cached = metricsCacheRef.current.get(runId); + if (cached) return { runId, data: cached }; + } const res = await fetchRunMetrics(runId); - metricsCacheRef.current.set(runId, res.data); + if (!isLive) { + metricsCacheRef.current.set(runId, res.data); + } return { runId, data: res.data }; }), ) @@ -106,7 +115,9 @@ export function RunsPanel({ return () => { cancelled = true; }; - }, [selectedRunIds]); + // `rows` is included so a poll-driven reload re-runs this effect: live-run + // metrics bypass the cache above, terminal-run metrics still hit it. + }, [selectedRunIds, rows]); return (