diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index a3fed2e..fcdbbbe 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -239,6 +239,26 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: if cfg.engine is None and runner is not None: cfg.engine = getattr(runner, "engine", None) + # A factory-built runner may know its own workload id better than the + # guessed/default one above (e.g. a caller passed ``workload_runner`` + # directly with no ``cfg.workload``, so ``workload`` fell through to the + # "vllm-decode" default regardless of what the runner actually is). It + # must be re-checked here, after the runner is resolved, not folded into + # the initial guess a few lines up — at that point neither the runner nor + # ``cfg.engine`` (populated from it just above) exist yet. + # + # ``cfg.workload`` (the field) is never reassigned anywhere in this + # function — it stays exactly the caller's original input for the whole + # call, unlike the local ``workload`` var this line progressively + # resolves. So this is an unambiguous "did the caller pin one explicitly" + # check, not a proxy for it. Deliberately not falling back to + # ``cfg.engine.workload_id`` here too: no current runner sets both an + # engine and a top-level workload_id, so there's no real precedence + # question yet — the runner's own attribute is preferred as the most + # specific source when it exists. + if cfg.workload is None and runner is not None: + workload = getattr(runner, "workload_id", None) or workload + # Sample the engine scheduler (queue depth, batch occupancy, preemptions) # over the same window as the CUPTI capture — engine-level telemetry the GPU # trace can't see. A no-op when no engine is attached (empty series). diff --git a/tests/test_run_loop_workload.py b/tests/test_run_loop_workload.py index 5affde5..bc1609f 100644 --- a/tests/test_run_loop_workload.py +++ b/tests/test_run_loop_workload.py @@ -10,6 +10,8 @@ from contextlib import contextmanager from pathlib import Path +import pytest + from .conftest import make_kernel, make_trace # Digest of an empty kernel list — sha256(repr([]))[:16]. A real trace must not @@ -32,6 +34,127 @@ def test_no_data_guard_does_not_fabricate_claims(tmp_path: Path): assert "NO DATA" in result["report_md"] +_UNSET = object() # identity sentinel — a real workload_id could legitimately be any string + + +class _Runner: + """A minimal workload_runner double — explicit about whether/what + ``workload_id`` it carries, instead of monkey-patching a plain function + (which silently loses the attribute if ever wrapped, e.g. functools.partial).""" + + def __init__(self, workload_id: object = _UNSET) -> None: + # _UNSET means: don't set the attribute at all, so + # getattr(..., "workload_id", None) exercises its own default path + # rather than reading an attribute that happens to be None. + if workload_id is not _UNSET: + self.workload_id = workload_id + + def __call__(self) -> dict: + return {"events": 1} + + +@pytest.fixture +def captured_workload_id(monkeypatch): + """Patches ``capture``/``sync_device`` so ``run_loop`` runs without a real + GPU, and returns the dict that the ``workload_id`` it was called with + lands in. Shared by every workload_id-relabeling test below to avoid + repeating the same three-line monkeypatch setup six times.""" + import gitm.scheduler.loop as loop + + captured: dict = {} + + @contextmanager + def fake_capture(out_path, *, workload_id="w", fingerprint="f", run_id=None): + captured["workload_id"] = workload_id + kernels = [make_kernel("k", start_ns=i * 100, end_ns=i * 100 + 90) for i in range(10)] + yield make_trace(events=kernels, vendor="nvidia", run_id=run_id or "r") + + monkeypatch.setattr(loop, "capture", fake_capture) + monkeypatch.setattr(loop, "sync_device", lambda: None) + return captured + + +def test_runner_workload_id_relabels_trace_when_workload_unset(tmp_path, captured_workload_id): + """A factory-built runner's own ``workload_id`` (e.g. set by + ``_vllm_decode_factory``, see #60) must relabel the trace/capture call + instead of being silently discarded. + + Regression guard: ``run_loop`` computed its ``workload`` label from + ``cfg.engine.workload_id`` before the runner (and ``cfg.engine``) were even + resolved — the runner's own attribute was set on the wrong object at the + wrong time and was never actually read, so a runner passed via + ``workload_runner`` with no explicit ``cfg.workload`` always fell back to + the hardcoded "vllm-decode" default regardless of what it actually was. + """ + from gitm import optimize + + optimize(budget="1s", scratch=str(tmp_path), workload_runner=_Runner("custom-workload")) + + assert captured_workload_id["workload_id"] == "custom-workload" + + +def test_runner_workload_id_does_not_override_explicit_workload(tmp_path, captured_workload_id): + """An explicit ``cfg.workload`` always wins over the runner's self-reported id.""" + from gitm import optimize + + optimize( + workload="hft", budget="1s", scratch=str(tmp_path), + workload_runner=_Runner("custom-workload"), + ) + + assert captured_workload_id["workload_id"] == "hft" + + +def test_runner_workload_id_none_falls_back_to_default(tmp_path, captured_workload_id): + """A factory that didn't populate ``workload_id`` (explicit ``None``, a + plausible real-world state) falls back to the guessed/default label + rather than relabeling to "None".""" + from gitm import optimize + + optimize(budget="1s", scratch=str(tmp_path), workload_runner=_Runner(None)) + + assert captured_workload_id["workload_id"] == "vllm-decode" + + +def test_runner_without_workload_id_attribute_falls_back_to_default(tmp_path, captured_workload_id): + """A plain runner with no ``workload_id`` attribute at all (the common case + for a directly-passed callable, not a registry factory) exercises the + ``getattr(..., "workload_id", None)`` default path, not just the explicit + ``None`` case above.""" + from gitm import optimize + + optimize(budget="1s", scratch=str(tmp_path), workload_runner=_Runner()) + + assert captured_workload_id["workload_id"] == "vllm-decode" + + +def test_runner_workload_id_empty_string_falls_back_to_default(tmp_path, captured_workload_id): + """An empty-string ``workload_id`` is treated the same as unset/None, not + as a real (if unusual) label — ``getattr(...) or workload`` falls back on + any falsy value, by design, not just ``None``. Documented explicitly here + since it's a plausible footgun otherwise.""" + from gitm import optimize + + optimize(budget="1s", scratch=str(tmp_path), workload_runner=_Runner("")) + + assert captured_workload_id["workload_id"] == "vllm-decode" + + +def test_no_runner_and_no_explicit_workload_uses_default(tmp_path, captured_workload_id): + """With nothing to read from at all (no runner, no cfg.workload, no + cfg.engine), the simplest path still resolves to the hardcoded default — + the relabeling guard must not require a runner to be present. Relies on + the same "no vLLM installed" test-environment fact + ``test_no_data_guard_does_not_fabricate_claims`` already depends on + (rather than monkeypatching ``get_factory``, which would silently stop + exercising this path if the lookup ever moved or gained a second site).""" + from gitm import optimize + + optimize(budget="1s", scratch=str(tmp_path)) + + assert captured_workload_id["workload_id"] == "vllm-decode" + + def test_runner_runs_inside_capture_and_produces_real_trace(tmp_path: Path, monkeypatch): """An injected runner is invoked inside the capture window; with kernels in the trace the loop proceeds to real claims with a non-empty fingerprint."""