From 8f49df50f580beb5398e3a6b74263a4396eee94b Mon Sep 17 00:00:00 2001 From: Elias Strauss Date: Wed, 14 Oct 2026 20:07:39 +0200 Subject: [PATCH] Add opt-in inspection instrumentation for stratum-gui MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Plan snapshots, per-node timings and hardware tracking, consumed by stratum-gui. Inspection only and off by default: enabled via set_config(diagnostics_dir=..., hardware_utilization=True). All code lives in stratum/_instrumentation; the scheduler exposes one optional observer hook and the API calls hooks that no-op when disabled. Disabled, the scheduler is indistinguishable from main (3.81 vs 3.87 µs per no-op operator); API hooks add ~0.2 µs per grid_search/evaluate call. --- .github/workflows/python_tests.yml | 2 +- codecov.yml | 6 +- pyproject.toml | 4 + stratum/__init__.py | 3 +- stratum/_api.py | 37 +-- stratum/_config.py | 31 +++ stratum/_instrumentation/__init__.py | 10 + stratum/_instrumentation/_diagnostics.py | 205 +++++++++++++++ stratum/_instrumentation/_graph.py | 45 ++++ stratum/_instrumentation/_hardware.py | 235 ++++++++++++++++++ stratum/_instrumentation/_hardware_sampler.py | 122 +++++++++ stratum/_instrumentation/_observer.py | 66 +++++ stratum/_instrumentation/_run.py | 80 ++++++ stratum/optimizer/_optimize.py | 8 + stratum/runtime/_scheduler.py | 9 + tests/instrumentation/__init__.py | 0 tests/instrumentation/test_diagnostics.py | 139 +++++++++++ tests/instrumentation/test_hardware.py | 218 ++++++++++++++++ uv.lock | 6 +- 19 files changed, 1204 insertions(+), 22 deletions(-) create mode 100644 stratum/_instrumentation/__init__.py create mode 100644 stratum/_instrumentation/_diagnostics.py create mode 100644 stratum/_instrumentation/_graph.py create mode 100644 stratum/_instrumentation/_hardware.py create mode 100644 stratum/_instrumentation/_hardware_sampler.py create mode 100644 stratum/_instrumentation/_observer.py create mode 100644 stratum/_instrumentation/_run.py create mode 100644 tests/instrumentation/__init__.py create mode 100644 tests/instrumentation/test_diagnostics.py create mode 100644 tests/instrumentation/test_hardware.py diff --git a/.github/workflows/python_tests.yml b/.github/workflows/python_tests.yml index 14ec0dd7..42ffecf1 100644 --- a/.github/workflows/python_tests.yml +++ b/.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' diff --git a/codecov.yml b/codecov.yml index 548da3e7..68d88fcf 100644 --- a/codecov.yml +++ b/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/**" diff --git a/pyproject.toml b/pyproject.toml index b9f06a4f..19be3044 100644 --- a/pyproject.toml +++ b/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] diff --git a/stratum/__init__.py b/stratum/__init__.py index a3edd7b9..01d158e6 100644 --- a/stratum/__init__.py +++ b/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", diff --git a/stratum/_api.py b/stratum/_api.py index 4264455b..4ed6eda7 100644 --- a/stratum/_api.py +++ b/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): diff --git a/stratum/_config.py b/stratum/_config.py index 67bb2ba6..34a4793e 100644 --- a/stratum/_config.py +++ b/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, diff --git a/stratum/_instrumentation/__init__.py b/stratum/_instrumentation/__init__.py new file mode 100644 index 00000000..141375d9 --- /dev/null +++ b/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, +) diff --git a/stratum/_instrumentation/_diagnostics.py b/stratum/_instrumentation/_diagnostics.py new file mode 100644 index 00000000..e27ca005 --- /dev/null +++ b/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) diff --git a/stratum/_instrumentation/_graph.py b/stratum/_instrumentation/_graph.py new file mode 100644 index 00000000..760418bf --- /dev/null +++ b/stratum/_instrumentation/_graph.py @@ -0,0 +1,45 @@ +"""Graphviz plan rendering for stratum-gui snapshots (same layout as show_graph).""" +from __future__ import annotations + +from graphviz import Digraph + +from stratum._config import get_config +from stratum.optimizer._op_utils import validate_operands +from stratum.optimizer.logical._ops import ChoiceOp + + +def plan_graph(nodes, filename: str = 'plan'): + """Graph of snapshot nodes ({"id", "label", "inputs"}), with stable SVG node IDs.""" + dot = Digraph( + comment=filename, + format=get_config()["graph_format"], + graph_attr={'rankdir': 'BT', 'ranksep': '0.3', 'nodesep': '0.2'}, + node_attr={'fontname': 'Helvetica', 'fontsize': '11'}, + ) + for node in nodes: + label = node["label"] + label = label.replace("<", "'").replace(">", "'") if label is not None else "None" + dot.node(str(node["id"]), label, id=f"node-{node['id']}") + for input_id in node["inputs"]: + dot.edge(str(input_id), str(node["id"])) + return dot + + +def op_graph(ops, filename: str = 'plan'): + """Plan graph with stage-local IDs for node inspection.""" + ops = list(ops) + ids = {op: i for i, op in enumerate(ops)} + nodes = [] + for op in ops: + validate_operands(op) + op.update_name() + nodes.append({"id": ids[op], + "label": op.name if isinstance(op, ChoiceOp) else str(op), + "inputs": [ids[inp] for inp in op.inputs]}) + return plan_graph(nodes, filename) + + +def graph_svg(dot): + """Render in memory, excluding the XML declaration and external DTD.""" + svg = dot.pipe(format="svg").decode("utf-8") + return svg[svg.index("= {"fit_transform", "predict"} for t in fitted) + assert run["buffer_pool"]["hits"] > 0 + + +def test_evaluate_export_and_config_restore(tmp_path): + original = st.get_config() + with st.dev_config(diagnostics_dir=tmp_path): + result = evaluate(pipeline()) + assert len(result) == 4 + assert artifact(tmp_path)["kind"] == "evaluate" + assert st.get_config() == original + + +def test_dev_config_restores_and_rejects_unknown_flags(tmp_path): + from stratum._config import DEV_FLAGS + with st.dev_config(diagnostics_dir=tmp_path, hardware_sampling_interval=.5): + assert DEV_FLAGS.diagnostics_dir == tmp_path + assert DEV_FLAGS.diagnostics_dir is None + assert DEV_FLAGS.hardware_sampling_interval == .1 + with pytest.raises(TypeError, match="diagnostic_dir"): + with st.dev_config(diagnostic_dir=tmp_path): + pass + assert "diagnostics_dir" not in st.get_config() + + +def fail(value, mode): + if mode != "preview": + raise ValueError("deliberate failure") + return value + + +def test_failed_run_exports_partial_measurements_and_preserves_exception(tmp_path): + dag = st.as_data_op(1).skb.apply_func(fail, st.eval_mode()) + with st.config(stats=False), st.dev_config(diagnostics_dir=tmp_path): + with pytest.raises(RuntimeError, match="deliberate failure"): + evaluate(dag) + run = artifact(tmp_path) + assert run["status"] == "failed" + assert run["error"]["stage"] == "execution" + assert "deliberate failure" in run["error"]["message"] + assert sum(t["failures"] for t in run["node_timings"].values()) == 1 + + +def test_optimization_failure_exports_stage(tmp_path, monkeypatch): + def fail_optimize(*args, **kwargs): + raise ValueError("invalid plan") + monkeypatch.setattr(st._api, "optimize", fail_optimize) + with st.dev_config(diagnostics_dir=tmp_path): + with pytest.raises(ValueError, match="invalid plan"): + evaluate(st.as_data_op(1)) + run = artifact(tmp_path) + assert run["error"]["stage"] == "optimization" + assert run["plans"] == [] + assert run["timings"]["execution"] is None + + +def test_export_failure_does_not_fail_computation(tmp_path, caplog): + blocked = tmp_path / "file" + blocked.write_text("not a directory") + with st.dev_config(diagnostics_dir=blocked): + result = evaluate(pipeline()) + assert len(result) == 4 + assert "Could not export Stratum diagnostics" in caplog.text diff --git a/tests/instrumentation/test_hardware.py b/tests/instrumentation/test_hardware.py new file mode 100644 index 00000000..8a23860d --- /dev/null +++ b/tests/instrumentation/test_hardware.py @@ -0,0 +1,218 @@ +import json +from time import perf_counter, sleep +from types import SimpleNamespace + +import pytest + +import stratum as st +from stratum._api import evaluate, grid_search +from stratum._instrumentation import _hardware +from stratum._instrumentation._hardware import HardwareTracker, hardware_run +from stratum._instrumentation._hardware_sampler import sample_tree +from stratum.runtime._scheduler import SequentialScheduler +from stratum._instrumentation._observer import observe +from tests.instrumentation.test_diagnostics import SameLabelOp, pipeline + + +def test_disabled_tracking_does_not_spawn_or_read_cpu_counters(monkeypatch): + def forbidden(*args, **kwargs): + pytest.fail("Disabled hardware profiling performed work") + monkeypatch.setattr(_hardware.subprocess, "Popen", forbidden) + monkeypatch.setattr(_hardware, "process_time", forbidden) + with st.dev_config(hardware_utilization=False): + result = grid_search(pipeline(), cv=2, scoring="neg_mean_squared_error") + assert not hasattr(result, "hardware_utilization") + + +def test_sampler_start_failure_preserves_result_and_exports_error(tmp_path, monkeypatch): + def unavailable(*args, **kwargs): + raise OSError("sampler deliberately unavailable") + monkeypatch.setattr(_hardware.subprocess, "Popen", unavailable) + with st.dev_config(hardware_utilization=True, diagnostics_dir=tmp_path): + result = evaluate(pipeline()) + assert len(result) == 4 + run = json.loads(next(tmp_path.glob("*.json")).read_text()) + assert run["status"] == "completed" + assert run["hardware"]["status"] == "unavailable" + assert "deliberately unavailable" in run["hardware"]["error"] + + +def test_sampler_report_failure_preserves_result_and_reaps_sampler(tmp_path, monkeypatch): + class BrokenSampler: + def __init__(self, *args, **kwargs): + from io import BytesIO + self.stdout = BytesIO(b'{"ready":true}\n') + self.returncode = None + def communicate(self, input=None, timeout=None): + self.returncode = 0 + return b"invalid report", b"" + def poll(self): + return self.returncode + def kill(self): + self.returncode = -9 + monkeypatch.setattr(_hardware.subprocess, "Popen", BrokenSampler) + with st.dev_config(hardware_utilization=True, diagnostics_dir=tmp_path): + result = grid_search(pipeline(), cv=2, scoring="neg_mean_squared_error") + assert len(result.results_) == 1 + run = json.loads(next(tmp_path.glob("*.json")).read_text()) + assert run["status"] == "completed" + assert run["hardware"]["status"] == "unavailable" + assert run["hardware"]["nodes"] # Engine CPU measurements survive a sampler failure. + + +def test_hardware_starts_before_input_preparation(tmp_path, monkeypatch): + from stratum import _api + original = _api.get_data + def get_data(dag): + assert _hardware.get_hardware_tracker() is not None + sleep(.03) + return original(dag) + monkeypatch.setattr(_api, "get_data", get_data) + with st.dev_config(hardware_utilization=True, diagnostics_dir=tmp_path): + evaluate(pipeline()) + run = json.loads(next(tmp_path.glob("*.json")).read_text()) + assert run["hardware"]["wall_seconds"] >= run["timings"]["total"] + .03 + + +def test_memory_join_uses_half_open_windows_and_reports_missing_samples(monkeypatch): + tracker = HardwareTracker() + a, b, c = object(), object(), object() + tracker.bind_scheduler(SimpleNamespace(linearized_dag=[a, b, c])) + cpu = iter([0.2, 0.1, 0.001]) + monkeypatch.setattr(_hardware, "process_time", lambda: next(cpu)) + for op, start, end in [(a, 0, 1), (b, 1, 2), (c, 2, 2.05)]: + phase = tracker.end(op, "predict", 1, 0, start, end) + phase.tracking_seconds += .00001 + tracker.boundary_seconds += .00001 + tracker.boundary_calls += 2 + tracker.started, tracker.finished, tracker.cpu_seconds = 0, 2.05, .301 + report = tracker._assemble({"samples": [ + {"timestamp": .25, "tree_rss_bytes": 100}, + {"timestamp": .75, "tree_rss_bytes": 300}, + {"timestamp": 1.0, "tree_rss_bytes": 500}, + {"timestamp": 1.5, "tree_rss_bytes": 700}, + {"timestamp": 2.05, "tree_rss_bytes": 999}, + {"timestamp": 1.25, "collection_start": .99, "tree_rss_bytes": 999}, + ]}) + nodes = report["nodes"] + assert nodes["0"]["rss_mean_bytes"] == 200 + assert nodes["0"]["calls_with_memory_samples"] == 1 + assert nodes["1"]["rss_peak_bytes"] == 700 + assert nodes["2"]["rss_mean_bytes"] is None + assert nodes["2"]["rss_peak_bytes"] is None + assert nodes["2"]["memory_sample_count"] == 0 + assert nodes["0"]["process_cpu_percent"] == pytest.approx(20) + assert sum(n["tracking_seconds"] for n in nodes.values()) == pytest.approx(report["overhead"]["boundary_seconds"]) + + +def test_window_retention_is_bounded_but_cpu_totals_remain_exact(monkeypatch): + monkeypatch.setattr(_hardware, "MAX_WINDOWS", 1) + monkeypatch.setattr(_hardware, "process_time", lambda: .2) + tracker = HardwareTracker() + op = object() + tracker.bind_scheduler(SimpleNamespace(linearized_dag=[op])) + tracker.end(op, "predict", 0, 0, 1, 2) + tracker.end(op, "predict", 1, 0, 2, 3) + assert len(tracker.windows) == 1 + assert tracker.dropped_windows == 1 + assert sum(p.process_cpu_seconds for p in tracker.nodes[0].values()) == pytest.approx(.4) + assert sum(p.calls for p in tracker.nodes[0].values()) == 2 + + +def test_sampler_excludes_its_own_process_and_handles_pid_reuse(): + class Proc: + def __init__(self, pid, born, cpu, rss): + self.pid, self.born, self.cpu, self.rss = pid, born, cpu, rss + def oneshot(self): + from contextlib import nullcontext + return nullcontext() + def memory_info(self): + return SimpleNamespace(rss=self.rss) + def cpu_times(self): + return SimpleNamespace(user=self.cpu, system=0) + def create_time(self): + return self.born + root = Proc(1, 0, 5, 100) + worker = Proc(2, 1, 2, 200) + sampler = Proc(3, 1, 50, 1000) + root.children = lambda **kwargs: [worker, sampler] + errors = SimpleNamespace(NoSuchProcess=ProcessLookupError, AccessDenied=PermissionError) + previous = {} + assert sample_tree(root, errors, previous, 3) == (100, 300, 7) + root.cpu += 1 + worker.born, worker.cpu = 2, .5 # Same PID, new process. + assert sample_tree(root, errors, previous, 3) == (100, 300, 1.5) + assert len(previous) == 2 + + +class BusyOp(SameLabelOp): + def process(self, mode, inputs): + end = perf_counter() + .12 + while perf_counter() < end: + sum(range(1000)) + return 1 + + +class IdleOp(SameLabelOp): + def process(self, mode, inputs): + sleep(.12) + return 1 + + +def test_real_service_tracks_cpu_memory_and_reaps_process(): + busy, idle = BusyOp(), IdleOp() + with hardware_run(True, .02) as tracker: + assert tracker.error is None, tracker.error + sched = SequentialScheduler([busy, idle], None, []) + observe(sched, tracker, collect_node_timings=False) + sched.process_op(busy) + sched.process_op(idle) + sched.pool.remove_all() + assert tracker.process.poll() == 0 + report = sched.hardware_utilization + assert report["status"] == "completed" + assert report["sample_count"] >= 3 + assert report["overhead"]["boundary_calls"] == 4 + assert report["overhead"]["service_cpu_seconds"] > 0 + assert report["overhead"]["startup_seconds"] > 0 + nodes = report["nodes"] + assert nodes["0"]["process_cpu_seconds"] > nodes["1"]["process_cpu_seconds"] + assert nodes["0"]["process_cpu_percent"] > nodes["1"]["process_cpu_percent"] + assert nodes["0"]["memory_sample_count"] > 0 + assert nodes["1"]["rss_peak_bytes"] > 0 + + +def memory_workload(value, mode): + if mode != "preview": + data = bytearray(16 * 2**20) + data[::4096] = b"x" * (len(data) // 4096) + sleep(.08) + return value + + +def test_api_exports_hardware_and_overhead_without_stats(tmp_path, capsys): + dag = st.as_data_op(1).skb.apply_func(memory_workload, st.eval_mode()) + with st.config(stats=False), st.dev_config( + hardware_utilization=True, hardware_sampling_interval=.02, diagnostics_dir=tmp_path): + evaluate(dag) + run = json.loads(next(tmp_path.glob("*.json")).read_text()) + hw = run["hardware"] + assert hw["status"] == "completed", hw["error"] + assert hw["nodes"] + assert any(n["rss_peak_bytes"] is not None for n in hw["nodes"].values()) + assert hw["wall_seconds"] >= run["timings"]["total"] + assert "Tracking overhead" not in capsys.readouterr().out + + +def test_failed_operation_still_records_hardware_and_stops_service(tmp_path, capsys): + from tests.instrumentation.test_diagnostics import fail + dag = st.as_data_op(1).skb.apply_func(fail, st.eval_mode()) + with st.config(stats=True), st.dev_config( + hardware_utilization=True, hardware_sampling_interval=.02, diagnostics_dir=tmp_path): + with pytest.raises(RuntimeError, match="deliberate failure"): + evaluate(dag) + run = json.loads(next(tmp_path.glob("*.json")).read_text()) + assert run["status"] == "failed" + assert run["hardware"]["status"] == "completed" + assert run["hardware"]["nodes"] + assert "Tracking overhead" in capsys.readouterr().out diff --git a/uv.lock b/uv.lock index 957b7bb8..a85286c1 100644 --- a/uv.lock +++ b/uv.lock @@ -2137,6 +2137,9 @@ cu126 = [ dev = [ { name = "ipykernel" }, ] +hardware = [ + { name = "psutil" }, +] test = [ { name = "catboost" }, { name = "coverage", extra = ["toml"] }, @@ -2165,6 +2168,7 @@ requires-dist = [ { name = "pillow", specifier = ">=12.3.0" }, { name = "polars" }, { name = "psutil", marker = "extra == 'benchmark'", specifier = ">=7.2" }, + { name = "psutil", marker = "extra == 'hardware'", specifier = ">=7.2" }, { name = "pyarrow", specifier = ">=23.0.1" }, { name = "pygments", specifier = ">=2.20.0" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=9.0.3" }, @@ -2178,7 +2182,7 @@ requires-dist = [ { name = "urllib3", specifier = ">=2.7.0" }, { name = "xgboost", marker = "extra == 'test'", specifier = ">=3.1.1" }, ] -provides-extras = ["test", "cpu", "cu126", "benchmark", "dev"] +provides-extras = ["test", "cpu", "cu126", "benchmark", "hardware", "dev"] [package.metadata.requires-dev] dev = [{ name = "maturin", specifier = ">=1.10.1" }]