Add opt-in inspection instrumentation for stratum-gui

.github/workflows/python_tests.yml

@@ -95,7 +95,7 @@ jobs:
       - name: Install project with test dependencies
         # `cpu` pins the CPU torch build: the runners have no GPU, and the
         # cu126 build would pull ~2.8 GB of nvidia wheels on every Linux job.
-        run: uv sync --extra test --extra cpu
+        run: uv sync --extra test --extra cpu --extra hardware
 
       - name: Run Python tests with coverage
         if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12'

codecov.yml

@@ -9,4 +9,8 @@ coverage:
         threshold: 1%
     patch:
       default:
-        target: 85% 
+        target: 85%
+
+# Inspection-only instrumentation for stratum-gui; off in production runs.
+ignore:
+  - "stratum/_instrumentation/**"

pyproject.toml

@@ -60,6 +60,9 @@ cu126 = [
 benchmark = [
     "psutil>=7.2"
 ]
+hardware = [
+    "psutil>=7.2"
+]
 dev = [
     "ipykernel>=7.3.0",
 ]
@@ -125,6 +128,7 @@ omit = [
     "**/tests/**",
     "*/__pycache__/*",
     "*/_rust_backend_native.*",
+    "*/stratum/_instrumentation/*",  # inspection-only, see codecov.yml
 ]
 
 [tool.coverage.report]

stratum/__init__.py

@@ -2,7 +2,7 @@
 from importlib.metadata import PackageNotFoundError, version as _dist_version
 from typing import Any, List
 
-from ._config import set_config, get_config, config
+from ._config import set_config, get_config, config, dev_config
 
 def _first_version(*names):
     for n in names:
@@ -32,6 +32,7 @@ def _first_version(*names):
     "StringEncoder",
     "OneHotEncoder",
     "config",
+    "dev_config",
     "__version__",
     "__skrub_version__",
     "versions",

stratum/_api.py

@@ -3,6 +3,7 @@
 from skrub import DataOp
 from skrub._data_ops._evaluation import find_scoring_node
 
+from stratum import _instrumentation as instrumentation
 from stratum._config import FLAGS
 from stratum.optimizer._optimize import SearchConfig, optimize
 from stratum.optimizer.logical._scoring import resolve_declared_scoring, resolve_scoring
@@ -19,7 +20,9 @@ class RunTimings:
     optimization: float
     execution: float
 
+
 #TODO: Rename this file
+@instrumentation.record_run("grid_search")
 def grid_search(dag: DataOp, cv=None, scoring=None, return_predictions=False, env=None):
     """Perform grid search with cross-validation on a DataOp DAG.
 
@@ -51,43 +54,41 @@ def grid_search(dag: DataOp, cv=None, scoring=None, return_predictions=False, en
                           return_predictions=return_predictions, declared_kwargs=declared)
     # Resolve variables to constants at compile time, so the scheduler runs
     # without an environment.
-    optimize_start = perf_counter()
-    linearized_dag, split_pos, flagged_ops = optimize(dag, env=env, search=search)
-    optimize_end = perf_counter()
-    sched = SequentialScheduler(linearized_dag, split_pos, flagged_ops, FLAGS.stats, t0=t0)
-
     # Without an explicit `cv`, the folds come from the splitter the plan declares with
     # `mark_as_X(cv=...)`, which the plan computes itself (see BuildSplitterOp).
-    preds = sched.grid_search(cv)
-    execution_end = perf_counter()
-
-    stats_printer(sched, RunTimings(
-        total=execution_end - optimize_start,
-        optimization=optimize_end - optimize_start,
-        execution=execution_end - optimize_end,
-    ))
+    sched, preds = _execute(dag, env, t0, lambda sched: sched.grid_search(cv), search=search)
 
     return (sched,preds) if return_predictions else sched
 
 
+@instrumentation.record_run("evaluate")
 def evaluate(dag: DataOp, seed: int = 42, test_size = 0.2):
     """Evaluate a DataOp DAG with train/test split."""
     t0 = perf_counter()
     # Resolve variables to constants at compile time, so the scheduler runs
     # without an environment.
     env = get_data(dag)
+    _, out = _execute(dag, env, t0, lambda sched: sched.evaluate(seed, test_size))
+    return out
+
+
+def _execute(dag, env, t0, execute, search=None):
+    instrumentation.optimization_started()
     optimize_start = perf_counter()
-    linearized_dag, split_pos, flagged_ops = optimize(dag, env=env)
+    linearized_dag, split_pos, flagged_ops = optimize(dag, env=env, search=search)
     optimize_end = perf_counter()
     sched = SequentialScheduler(linearized_dag, split_pos, flagged_ops, FLAGS.stats, t0=t0)
-    out = sched.evaluate(seed, test_size)
+    instrumentation.execution_started(sched)
+    out = execute(sched)
     execution_end = perf_counter()
-    stats_printer(sched, RunTimings(
+    run = RunTimings(
         total=execution_end - optimize_start,
         optimization=optimize_end - optimize_start,
         execution=execution_end - optimize_end,
-    ))
-    return out
+    )
+    instrumentation.record_timings(run)
+    stats_printer(sched, run)
+    return sched, out
 
 
 def stats_printer(sched: SequentialScheduler, run: RunTimings):

stratum/_config.py

@@ -101,6 +101,37 @@ class _Flags:
 
 FLAGS = _Flags()
 
+
+@dataclass
+class _DevFlags:
+    """Developer-only switches, not part of set_config/get_config.
+
+    Set them directly (``DEV_FLAGS.diagnostics_dir = ...``) or scoped with
+    ``with stratum.dev_config(diagnostics_dir=...):``.
+    """
+    # Export plan snapshots and per-node timings for stratum-gui after each run.
+    diagnostics_dir: str | os.PathLike | None = None
+    # Sample CPU/RSS in a separate process (needs the hardware extra).
+    hardware_utilization: bool = False
+    # Seconds between hardware samples.
+    hardware_sampling_interval: float = 0.1
+
+DEV_FLAGS = _DevFlags()
+
+
+@contextmanager
+def dev_config(**kwargs):
+    """Temporarily override DEV_FLAGS inside a context."""
+    unknown = set(kwargs) - set(vars(DEV_FLAGS))
+    if unknown:
+        raise TypeError(f"Unknown dev flag(s): {sorted(unknown)}")
+    original = vars(DEV_FLAGS).copy()
+    vars(DEV_FLAGS).update(kwargs)
+    try:
+        yield DEV_FLAGS
+    finally:
+        vars(DEV_FLAGS).update(original)
+
 def set_config(rust_backend: bool | None = None,
     num_threads: int | None = None,
     debug_timing: bool | None = None,

stratum/_instrumentation/__init__.py

@@ -0,0 +1,10 @@
+"""Inspection-only instrumentation consumed by Stratum Inspector (stratum-gui).
+
+Not part of the production runtime. Everything here is off by default and only
+activates when ``DEV_FLAGS.diagnostics_dir`` or ``DEV_FLAGS.hardware_utilization`` is set
+(directly or via ``stratum.dev_config``)
+(or a caller opens ``capture_plans``); otherwise every hook is a no-op.
+"""
+from stratum._instrumentation._run import (
+    execution_started, optimization_started, record_run, record_timings,
+)

stratum/_instrumentation/_diagnostics.py

@@ -0,0 +1,205 @@
+"""Optional, versioned run artifacts. No web server or GUI dependency lives here."""
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+import logging
+import os
+from pathlib import Path
+import json
+from time import perf_counter
+from uuid import uuid4
+
+logger = logging.getLogger(__name__)
+_active_run = ContextVar("stratum_diagnostics", default=None)
+_plan_capture = ContextVar("stratum_plan_capture", default=None)
+_capture_svg = ContextVar("stratum_capture_svg", default=False)
+
+
+@dataclass
+class NodeTiming:
+    """Operator processing only; input/output buffer work is measured separately."""
+
+    count: int = 0
+    seconds: float = 0.0
+    buffer_seconds: float = 0.0
+    failures: int = 0
+    phases: dict = field(default_factory=dict)
+
+    def record(self, mode, fold, seconds, buffer_seconds, failed=False):
+        self.count += 1
+        self.seconds += seconds
+        self.buffer_seconds += buffer_seconds
+        self.failures += int(failed)
+        key = (mode, fold)
+        phase = self.phases.get(key)
+        if phase is None:
+            phase = self.phases[key] = {
+                "mode": mode, "fold": fold, "count": 0, "seconds": 0.0,
+                "buffer_seconds": 0.0, "failures": 0,
+            }
+        phase["count"] += 1
+        phase["seconds"] += seconds
+        phase["buffer_seconds"] += buffer_seconds
+        phase["failures"] += int(failed)
+
+    def to_dict(self):
+        return {
+            "count": self.count, "seconds": self.seconds,
+            "buffer_seconds": self.buffer_seconds, "failures": self.failures,
+            "phases": list(self.phases.values()),
+        }
+
+
+def plan_snapshot(name, ops, split_pos=None, recompute_ops=()):
+    """Use plan-local positions as IDs; freeze labels before further IR mutations."""
+    ids = {op: i for i, op in enumerate(ops)}
+    recompute = set(recompute_ops)
+    return {
+        "name": name, "split_pos": split_pos,
+        "nodes": [{
+            "id": i, "label": str(op), "implementation": type(op).__name__,
+            "inputs": [ids[inp] for inp in op.inputs],
+            "output_type": op.output_type.name,
+            "dead_in_fit": op.dead_in_fit, "recompute": op in recompute,
+        } for i, op in enumerate(ops)],
+    }
+
+
+def capture_plan(name, ops, split_pos=None, recompute_ops=()):
+    """Only traverse when runtime diagnostics or in-memory capture is enabled."""
+    plans = _plan_capture.get()
+    if plans is not None:
+        ops = list(ops)
+        snapshot = plan_snapshot(name, ops, split_pos, recompute_ops)
+        plans.append(snapshot)
+        if _capture_svg.get():
+            from stratum._instrumentation._graph import graph_svg, op_graph
+            snapshot["svg"] = graph_svg(op_graph(ops, name))
+        return
+    run = _active_run.get()
+    if run is not None:
+        try:
+            run.plans.append(plan_snapshot(name, list(ops), split_pos, recompute_ops))
+        except Exception:
+            logger.warning("Could not capture diagnostics plan %s", name, exc_info=True)
+
+
+def diagnostics_enabled():
+    return _active_run.get() is not None or _plan_capture.get() is not None
+
+
+@contextmanager
+def capture_plans(*, include_svg=False):
+    """Capture optimizer checkpoints in memory without a scheduler or disk export.
+
+    Unlike best-effort runtime export, capture errors propagate to the caller.
+    Nested captures restore the previous sink when they finish.
+    """
+    plans = []
+    token = _plan_capture.set(plans)
+    svg_token = _capture_svg.set(include_svg)
+    try:
+        yield plans
+    finally:
+        _capture_svg.reset(svg_token)
+        _plan_capture.reset(token)
+
+
+def current_diagnostics():
+    return _active_run.get()
+
+
+class RunDiagnostics:
+    def __init__(self, directory, kind, config):
+        self.directory = Path(directory).expanduser()
+        self.run_id = uuid4().hex
+        self.kind = kind
+        self.config = config
+        self.started_at = datetime.now(timezone.utc).isoformat()
+        self.plans = []
+        self.scheduler = None
+        self.timings = None
+        self.stage = "preparation"
+        self.optimization_start = None
+        self.execution_start = None
+        self.hardware = None
+
+    def partial_timings(self, start):
+        """Timings of a run that did not reach ``record_timings``, split by stage marks."""
+        end = perf_counter()
+        if self.optimization_start is None:
+            return {"total": end - start, "optimization": None, "execution": None}
+        execution_start = self.execution_start or end
+        return {
+            "total": end - self.optimization_start,
+            "optimization": execution_start - self.optimization_start,
+            "execution": end - self.execution_start if self.execution_start else None,
+        }
+
+    def publish(self, error=None):
+        sched = self.scheduler
+        document = {
+            "schema_version": 1, "run_id": self.run_id, "kind": self.kind,
+            "started_at": self.started_at,
+            "finished_at": datetime.now(timezone.utc).isoformat(),
+            "status": "failed" if error is not None else "completed",
+            "error": {"type": type(error).__name__, "message": str(error),
+                      "stage": self.stage} if error is not None else None,
+            "config": self.config, "plans": self.plans,
+            "timings": self.timings,
+            "node_timings": {}, "heavy_hitters": [], "buffer_pool": None,
+            "hardware": self.hardware,
+        }
+        if sched is not None:
+            node_timings = (sched.observer.node_timings if sched.observer is not None else None) or {}
+            document["node_timings"] = {
+                str(i): node_timings[op].to_dict()
+                for i, op in enumerate(sched.linearized_dag)
+                if op in node_timings
+            }
+            groups = {}
+            for op, timing in node_timings.items():
+                label = str(op)
+                group = groups.setdefault(label, {"label": label, "count": 0, "seconds": 0.0})
+                group["count"] += timing.count
+                group["seconds"] += timing.seconds
+            document["heavy_hitters"] = sorted(groups.values(), key=lambda g: g["seconds"], reverse=True)
+            document["buffer_pool"] = asdict(sched.pool.stats)
+
+        self.directory.mkdir(parents=True, exist_ok=True)
+        destination = self.directory / f"{self.run_id}.json"
+        temporary = self.directory / f".{self.run_id}.tmp"
+        try:
+            with temporary.open("x", encoding="utf-8") as stream:
+                json.dump(document, stream, separators=(",", ":"), allow_nan=False)
+            os.replace(temporary, destination)
+        finally:
+            temporary.unlink(missing_ok=True)
+
+
+@contextmanager
+def diagnostic_run(directory, kind, config):
+    """Export failures too; diagnostics must never change the computation outcome."""
+    if directory is None:
+        yield None
+        return
+    run = RunDiagnostics(directory, kind, config)
+    token = _active_run.set(run)
+    start = perf_counter()
+    error = None
+    try:
+        yield run
+    except BaseException as exc:
+        error = exc
+        raise
+    finally:
+        _active_run.reset(token)
+        if run.timings is None:
+            run.timings = run.partial_timings(start)
+        try:
+            run.publish(error)
+        except Exception:
+            logger.warning("Could not export Stratum diagnostics to %s", directory, exc_info=True)