Latency metrics, collective signal, and verification export

[kwang3170]

• Adds per-request TTFT/TPOT/goodput from vLLM's RequestOutput.metrics, plus a wall-clock anchor on the scheduler
  sampler so the GPU trace, scheduler samples, and request timestamps share one clock; both land in scheduler_stats.json.

• Runs node_rollup's existing comm math on the live trace and converts it into ranked CollectiveCause entries in
  residuals.json, mirroring how scheduler causes already work.

• Customer report - Serializes each EngineABResult (until now read for two numbers and discarded) into verification.json with the full engine config on both sides, rep count, scatter, and environment.
  (My bad on last standup, I misinterpreted the task earlier and realised this part of the report wasn't actually done, so its here now in this PR)

[github-actions]

Code Review by Gemini

The changes introduce significant new features: latency metrics, collective communication signal attribution, and a detailed verification export. The implementation is well-structured, with new modules for distinct functionalities and careful integration into the existing `run_loop`. The accompanying tests are comprehensive and cover various scenarios, including edge cases and error handling.

No critical bugs were found. The code demonstrates good practices like handling `None` values, defensive programming, and clear documentation.

### Minor Improvements

**1. `gitm/optimizer/verification_export.py` - `build_record` function**

The fallback logic for `candidate_config` is a best-effort attempt to reconstruct the configuration if `gitm_llm_kwargs` is not available from the `live_engine`. While robust, for a "customer report" aiming for full reproducibility, it's ideal to always capture the *actual* full configuration. If `gitm_llm_kwargs` is the authoritative source, relying on `c.spec.knobs` as a fallback might introduce subtle discrepancies if the intervention involves more complex changes not fully captured by `c.spec.knobs`.

Consider adding a comment or assertion to clarify when this fallback is expected to be hit (e.g., only for `DryRunApplicator` or specific engine versions) and to emphasize that for `LiveEngineApplicator`, `gitm_llm_kwargs` should ideally always be present and complete.

```diff
--- a/gitm/optimizer/verification_export.py
+++ b/gitm/optimizer/verification_export.py
@@ -109,6 +109,10 @@
             candidate_cfg = dict(getattr(live_engine, "gitm_llm_kwargs", None) or {})
             if not candidate_cfg and baseline_cfg:
                 candidate_cfg = {**baseline_cfg, **(c.spec.knobs or {c.spec.knob: c.spec.value})}
+            # NOTE: The fallback above (if candidate_cfg is empty) reconstructs the config
+            # from baseline + the intervention spec. This is a best-effort for cases
+            # where the live_engine does not expose its full kwargs (e.g., DryRunApplicator).
+            # For LiveEngineApplicator, gitm_llm_kwargs should ideally be complete.
             verification.append(
                 build_record(
                     c.spec, ab, result,

2. gitm/scheduler/loop.py - candidate_cfg fallback logic

This is the call site for the build_record function, and the same consideration applies. The fallback logic for candidate_cfg is present here.

--- a/gitm/scheduler/loop.py
+++ b/gitm/scheduler/loop.py
@@ -560,6 +560,10 @@
             live_engine = getattr(applicator, "engine", cfg.engine)
             candidate_cfg = dict(getattr(live_engine, "gitm_llm_kwargs", None) or {})
             if not candidate_cfg and baseline_cfg:
                 candidate_cfg = {**baseline_cfg, **(c.spec.knobs or {c.spec.knob: c.spec.value})}
+            # NOTE: The fallback above (if candidate_cfg is empty) reconstructs the config
+            # from baseline + the intervention spec. This is a best-effort for cases
+            # where the live_engine does not expose its full kwargs (e.g., DryRunApplicator).
+            # For LiveEngineApplicator, gitm_llm_kwargs should ideally be complete.
             verification.append(
                 build_record(
                     c.spec, ab, result,

[github-actions]

🤖 Claude Code Review

Review

🐛 Bugs

worst_device_comm passes wrong object to device_comm_stats
device_comm_stats receives a Trace whose events field is replaced with a raw list[KernelEvent], but the original trace.events is likely typed as list[Event] (a union including non-kernel events). The model_copy(update={"events": evs}) replacement silently passes only kernel events, which may break device_comm_stats if it relies on the full event list for wall-time or other fields. More critically, by_dev groups on k.device_id but then trace.model_copy(update={"events": evs}) doesn't update device_count, so multi-device sub-traces report wrong device counts.

goodput_rps is None when window_s == 0 but not when it's very small

goodput = (len(met) / window_s) if window_s else None

window_s is max(..., 0.0), so a near-zero window (clock granularity) yields a huge but non-None rate. The comment says to guard against division by ~0, but only an exact zero is caught. Use if window_s and window_s > 1e-6 or similar.

_percentile uses round() which rounds half-to-even
round(0.5 * (n-1)) for even n-1 will round down when the expected nearest-rank behaviour rounds up. For small n (e.g. n=2, p95 → round(0.95)=1 which is fine, but n=3, p50 → round(1.0)=1, ok) this is mostly fine but deviates from standard nearest-rank at boundary cases. Document the choice or use math.floor(q * len + 0.5).

test_real_scatter_widens_the_band_beyond_the_floor asserts a specific rel_std value
The test asserts agreement_band ≈ 0.10 (10%) from baseline_std=5, candidate_std=5 on a 100 tok/s baseline. This implicitly assumes EngineABResult.rel_std is computed as (std_a + std_b) / baseline_tps. That formula is never shown in this diff — if rel_std is computed differently (e.g. pooled, or only baseline), the test will silently pass against the wrong value or break on an internal change.

DeviceCommStats missing device_id field in test assertion
test_cross_device_compute_does_not_mask_exposed_comm asserts stats.device_id == 0, but DeviceCommStats isn't shown to have a device_id field. If worst_device_comm returns the DeviceCommStats from the per-device call (which operates on a sub-trace with only that device's kernels), there's no guarantee DeviceCommStats carries the originating device ID.

---

🔒 Security

write_verification creates parent directories without restriction

out_path.parent.mkdir(parents=True, exist_ok=True)

If out_path comes from user/config input and is an absolute path, this will silently create arbitrary directory trees. Validate out_path is within an expected run_dir before writing.

---

⚡ Performance

worst_device_comm iterates trace.kernels() twice
trace.kernels() is called once for the early-exit check and again implicitly inside each device_comm_stats call. If kernels() is an O(n) filter over all events, this is 1 + D passes for D devices. Pass the pre-grouped lists directly if the API allows, or cache the result.

Per-iteration getattr(applicator, "engine", ...) in the hot path
In the optimization loop, live_engine = getattr(applicator, "engine", cfg.engine) and the subsequent getattr(live_engine, "gitm_llm_kwargs", None) are repeated every candidate. These are cheap but if gitm_llm_kwargs is a property that copies a large config dict, this is unnecessary allocation. Consider making config snapshots explicit and cheap.

---

📊 Reproducibility

t0_wall_ns and t0_ns captured sequentially, not atomically

self._t0_ns = time.perf_counter_ns()
self._t0_wall_ns = time.time_ns()

These two calls have an unmeasured gap (typically microseconds). Any code using t0_wall_ns + sample.t_ns to align with time.time_ns()-based request timestamps inherits this offset. The gap is documented as "exact" in the comment, which overstates the precision. Document the expected error bound (~1–10 µs) so consumers know the join tolerance.

SchedulerStatsSummary.t0_wall_ns = 0 as sentinel is ambiguous
Zero is a valid time_ns() value on Unix epoch. A sentinel of None with t0_wall_ns: int | None = 0 would be unambiguous, or at minimum the field docstring should note "0 means unset, not epoch."

---

💡 Suggestions

collective_causes floor constants aren't validated against each other
If a caller passes exposed_floor > dominant_floor the two causes fire with inverted semantics. Add an assertion or swap silently:

assert exposed_floor <= dominant_floor, "exposed_floor should be <= dominant_floor"

build_record silently falls back to a synthesized candidate_cfg

if not candidate_cfg and baseline_cfg:
    candidate_cfg = {**baseline_cfg, **(c.spec.knobs or {c.spec.knob: c.spec.value})}

This synthesized config may be wrong if the apply path changed other knobs as side-effects (e.g. a restart that resets defaults). The fallback should at minimum log a warning so operators know the exported config is inferred, not observed.

req_records list cast is unnecessarily defensive

req_records = list(run_out.get("requests") or []) if isinstance(run_out, dict) else []

request_records_from_outputs already handles None/empty; the or [] guard is fine but list(...) on an already-list result is a gratuitous copy. Minor, but in a hot loop it adds up.

No test for summarize_requests with all-None-TTFT records
The case where every record has arrival_wall_s but first_token_wall_s is None (prefill never reported) isn't tested — n_ttft == 0, goodput_rps should be None because no SLO can be met without TTFT, but the current meets_slo correctly returns False in that case. Worth an explicit test to lock in this behaviour.

write_verification return type annotation missing
The function is annotated -> str in the docstring but returns str(out_path). If out_path was already a str, this double-converts. Trivial, but inconsistent with the Path-first philosophy elsewhere.

[github-actions]

🤖 Claude Code Review

Code Review

🐛 Bugs

worst_device_comm passes wrong object to device_comm_stats
trace.model_copy(update={"events": evs}) passes a list of KernelEvent objects as events, but the original trace likely holds list[TraceEvent] (a union including MemcpyEvent, SyncEvent, etc.). More critically, device_comm_stats calls trace.kernels() — if the split filters only KernelEvent objects, other event types are silently dropped on each per-device copy. Verify trace.kernels() vs trace.events semantics here.

_as_dims redundant guard

if isinstance(val, list | tuple) and len(val) >= 1:
    x = as_int(val[0], 1) if len(val) > 0 else 1  # `len(val) > 0` always True here

The inner if len(val) > 0 is dead code given the outer >= 1 guard. Not harmful but misleading.

DeviceCommStats missing device_id field in tests
test_cross_device_compute_does_not_mask_exposed_comm asserts stats.device_id == 0, but DeviceCommStats is not shown to carry a device_id field — worst_device_comm returns whatever device_comm_stats returns, which may not record which device it came from. This test would fail or require device_comm_stats to be extended.

goodput_rps when window_s == 0

goodput = (len(met) / window_s) if window_s else None

window_s can be exactly 0.0 (two requests at identical timestamps), which is falsy — so goodput correctly becomes None. But the test test_arrival_without_first_token_meets_no_slo asserts goodput_rps == 0.0 with a 2s window and zero SLO-meeting requests. That path (0 / 2.0) returns 0.0, which is correct — but if window_s were 0, it would return None instead of 0.0, which could surprise callers expecting a rate of zero. The inconsistency (None vs 0.0) warrants a comment or type-narrowing.

_percentile rounding edge case
int(round(q * (len(ordered) - 1))) for q=0.99 on a 2-element list gives round(0.99) = 1 (index 1). This is the max, not P99 of two samples. Minor, but the docstring says "nearest-rank" without acknowledging this degeneracy for tiny samples.

agreement_band computation uses rel_std not derived from baseline_std/candidate_std

agreement_band=max(ab.rel_std, MIN_NOISE_BAND)

EngineABResult.rel_std is not shown in the diff. If it's defined as, say, candidate_std / candidate_tps, the test test_real_scatter_widens_the_band_beyond_the_floor hardcodes an expectation of 0.10 (5/100 + 5/100?), but the actual formula is opaque. The test may be fragile.

---

🔒 Security

write_verification uses caller-supplied path without restriction

out_path.parent.mkdir(parents=True, exist_ok=True)
out_path.write_text(...)

If out_path comes from config or user input, an attacker could write to arbitrary locations. At minimum, validate it stays within the expected run_dir.

---

⚡ Performance

Per-iteration getattr on hot path in apply loop

live_engine = getattr(applicator, "engine", cfg.engine)
candidate_cfg = dict(getattr(live_engine, "gitm_llm_kwargs", None) or {})

This is inside the intervention loop. If gitm_llm_kwargs is a large dict (full engine config), copying it on every iteration is wasteful. Consider a shallow copy only when the key set actually changes, or snapshot only the known-varying keys.

worst_device_comm iterates kernels twice
trace.kernels() is called once to build by_dev, then device_comm_stats calls trace.kernels() again (via the model copy) per device. If kernels() is O(n), this is O(devices × n). For the common single-device case the fast path works, but the multi-device path reconstructs per-device traces only to have them re-scanned.

---

📊 Reproducibility

t0_wall_ns captured sequentially, not atomically

self._t0_ns = time.perf_counter_ns()
self._t0_wall_ns = time.time_ns()

These two calls are not atomic. On a loaded system, context switches between them introduce an unknown offset (potentially milliseconds). The comment claims "sub-microsecond" but that's only a best-case. This offset is then used to align scheduler samples against request timestamps at second-scale latencies — acceptable in practice, but the "sub-microsecond" claim in the comment is misleading and should be softened.

---

💡 Suggestions

run_out type unsafety

req_records = list(run_out.get("requests") or []) if isinstance(run_out, dict) else []

run_out is typed Any and initialized to None. If a future runner returns a non-dict non-None value (e.g. a float TPS), this silently drops requests. A more explicit isinstance check or a typed return dataclass from WorkloadRunner would catch this at definition time.

baseline_cfg snapshot may be empty even when engine exists

baseline_cfg = dict(getattr(cfg.engine, "gitm_llm_kwargs", None) or {})

If gitm_llm_kwargs is absent (engine doesn't implement it), baseline_cfg is {} and the fallback:

candidate_cfg = {**baseline_cfg, **(c.spec.knobs or {c.spec.knob: c.spec.value})}

produces a candidate config that only contains the changed knob — not genuinely useful for reproduction. A warning log when the attribute is missing would surface this silently degraded case.

collective_causes floor thresholds not configurable at call site from loop
worst_device_comm is called with default exposed_floor=0.05/dominant_floor=0.20 from loop.py. These are significant operational thresholds but aren't exposed in LoopConfig, making them hard to tune per-workload without code changes.

Missing __all__ in new public modules
collective_signal.py and verification_export.py export several symbols but define no __all__. This makes wildcard imports and API surface unclear.

test_nccl_fixture_produces_a_collective_cause depends on a fixture file
synthetic_4xA100_nccl.json is referenced but not in the diff. If this file is missing from CI, the test fails with a confusing FileNotFoundError rather than a skip. Consider pytest.importorskip or a fixture existence check.

[github-actions]

Code Review by Gemini

The changes introduce significant new features for latency metrics, collective communication analysis, and verification export. The code is generally well-structured, documented, and accompanied by comprehensive tests.

However, there is a potential compatibility issue related to Python version syntax.

Issue: Python Version Compatibility for Type Union Syntax

The new X | Y syntax for type unions in isinstance() checks is a feature introduced in Python 3.10. If the project's minimum supported Python version is older than 3.10, these changes will cause a SyntaxError.

Files and Lines to Change:

1. gitm/importers/torch_trace.py

--- a/gitm/importers/torch_trace.py
+++ b/gitm/importers/torch_trace.py
@@ -84,7 +84,7 @@ def _arg_get(args: dict[str, Any] | None, keys: tuple[str, ...]) -> Any:
     lower = {str(k).lower(): v for k, v in args.items()}
     for k in keys:
         if k.lower() in lower:
             return lower[k.lower()]
     return None
 
 
 def _as_dims(val: Any) -> tuple[int, int, int]:
     if val is None:
         return (1, 1, 1)
-    if isinstance(val, (list, tuple)) and len(val) >= 1:
+    if isinstance(val, list | tuple) and len(val) >= 1:
         x = as_int(val[0], 1) if len(val) > 0 else 1
         y = as_int(val[1], 1) if len(val) > 1 else 1
         z = as_int(val[2], 1) if len(val) > 2 else 1
         return (max(x, 1), max(y, 1), max(z, 1))
     if isinstance(val, dict):
         return (
             max(as_int(val.get("x", val.get("X", 1)), 1), 1),
             max(as_int(val.get("y", val.get("Y", 1)), 1), 1),
             max(as_int(val.get("z", val.get("Z", 1)), 1), 1),
         )

Suggested Change:
Revert to the tuple syntax for isinstance() for broader compatibility with Python versions older than 3.10.

--- a/gitm/importers/torch_trace.py
+++ b/gitm/importers/torch_trace.py
@@ -84,7 +84,7 @@ def _arg_get(args: dict[str, Any] | None, keys: tuple[str, ...]) -> Any:
     lower = {str(k).lower(): v for k, v in args.items()}
     for k in keys:
         if k.lower() in lower:
             return lower[k.lower()]
     return None
 
 
 def _as_dims(val: Any) -> tuple[int, int, int]:
     if val is None:
         return (1, 1, 1)
-    if isinstance(val, list | tuple) and len(val) >= 1:
+    if isinstance(val, (list, tuple)) and len(val) >= 1:
         x = as_int(val[0], 1) if len(val) > 0 else 1
         y = as_int(val[1], 1) if len(val) > 1 else 1
         z = as_int(val[2], 1) if len(val) > 2 else 1
         return (max(x, 1), max(y, 1), max(z, 1))
     if isinstance(val, dict):
         return (
             max(as_int(val.get("x", val.get("X", 1)), 1), 1),
             max(as_int(val.get("y", val.get("Y", 1)), 1), 1),
             max(as_int(val.get("z", val.get("Z", 1)), 1), 1),
         )

2. tests/test_importers_hypothesis.py

--- a/tests/test_importers_hypothesis.py
+++ b/tests/test_importers_hypothesis.py
@@ -34,7 +34,7 @@ def test_prop_us_to_ns_monotonic(ts: float, dur: float):
         }
     )
     # Default import path uses LightKernel for scale; --strict uses KernelEvent.
-    assert isinstance(ev, (KernelEvent, LightKernel))
+    assert isinstance(ev, KernelEvent | LightKernel)
     assert ev.start_ns == int(ts * 1000.0)
     assert ev.end_ns == int((ts + dur) * 1000.0)
     assert ev.end_ns >= ev.start_ns or dur == 0.0
@@ -67,7 +67,7 @@ def test_prop_missing_args_never_crash(cat, name, has_grid, has_stream, has_devi
     # memset names must skip for Memcpy-ish cats
     assume("memset" not in name.lower())
     ev = event_from_chrome(
         {"ph": "X", "cat": cat, "name": name, "ts": 1.0, "dur": 2.0, "args": args}
     )
     # May be None for weird name/cat combos; must not raise.
     if ev is not None:
-        assert isinstance(ev, (KernelEvent, MemcpyEvent, LightKernel, LightMemcpy))
+        assert isinstance(ev, KernelEvent | MemcpyEvent | LightKernel | LightMemcpy)
         assert getattr(ev, "kind", None) in ("kernel", "memcpy")
         assert ev.end_ns >= ev.start_ns
         # Light events are not pydantic; only validate the full Trace path
         # when the object supports model_dump (strict/full KernelEvent path).
         if hasattr(ev, "model_dump"):
             Trace.model_validate(
                 {
                     "workload_id": "w",
                     "fingerprint": "f",
                     "run_id": "r",

Suggested Change:
Revert to the tuple syntax for isinstance() for broader compatibility with Python versions older than 3.10.

--- a/tests/test_importers_hypothesis.py
+++ b/tests/test_importers_hypothesis.py
@@ -34,7 +34,7 @@ def test_prop_us_to_ns_monotonic(ts: float, dur: float):
         }
     )
     # Default import path uses LightKernel for scale; --strict uses KernelEvent.
-    assert isinstance(ev, KernelEvent | LightKernel)
+    assert isinstance(ev, (KernelEvent, LightKernel))
     assert ev.start_ns == int(ts * 1000.0)
     assert ev.end_ns == int((ts + dur) * 1000.0)
     assert ev.end_ns >= ev.start_ns or dur == 0.0
@@ -67,7 +67,7 @@ def test_prop_missing_args_never_crash(cat, name, has_grid, has_stream, has_devi
     # memset names must skip for Memcpy-ish cats
     assume("memset" not in name.lower())
     ev = event_from_chrome(
         {"ph": "X", "cat": cat, "name": name, "ts": 1.0, "dur": 2.0, "args": args}
     )
     # May be None for weird name/cat combos; must not raise.
     if ev is not None:
-        assert isinstance(ev, KernelEvent | MemcpyEvent | LightKernel | LightMemcpy)
+        assert isinstance(ev, (KernelEvent, MemcpyEvent, LightKernel, LightMemcpy))
         assert getattr(ev, "kind", None) in ("kernel", "memcpy")
         assert ev.end_ns >= ev.start_ns
         # Light events are not pydantic; only validate the full Trace path
         # when the object supports model_dump (strict/full KernelEvent path).
         if hasattr(ev, "model_dump"):
             Trace.model_validate(
                 {
                     "workload_id": "w",
                     "fingerprint": "f",
                     "run_id": "r",

[github-actions]

🤖 Claude Code Review

Review

🐛 Bugs

worst_device_comm passes raw event lists, not Trace objects

device_comm_stats(trace.model_copy(update={"events": evs}))

evs here is a list[KernelEvent] (built via by_dev.setdefault(k.device_id, []).append(k)), but Trace.events almost certainly expects list[TraceEvent] (kernels + memcpy + sync). If device_comm_stats iterates trace.kernels() (a filtered view), this silently works; if it accesses trace.events directly and the schema validates on assignment, it either crashes or silently loses non-kernel events from the per-device split. Verify model_copy accepts a bare kernel list.

DeviceCommStats has no device_id field but the test asserts stats.device_id == 0

# test_collective_signal.py:108
assert stats.device_id == 0

If DeviceCommStats doesn't carry device_id, this test will fail with AttributeError on first run. Either the field needs adding to the dataclass or the test assertion is wrong.

goodput_rps is 0.0 but typed as float | None

# test_serving_latency.py:97
assert s.goodput_rps == 0.0  # window is known (2.0s), zero requests met SLO

In summarize_requests: goodput = (len(met) / window_s) if window_s else None. When window_s > 0 and len(met) == 0, this correctly returns 0.0. But the test comment in test_arrival_without_first_token_meets_no_slo says "window is known" — however, window_s requires both arrivals and finishes. Those records have arrival_wall_s=0.0 and finished_wall_s set, so the window is computable and this works. But it's fragile: if either list is empty (e.g. a build that only populates finished_wall_s), window_s is None and the test breaks silently.

agreement_band formula for rel_std is not shown

agreement_band=max(ab.rel_std, MIN_NOISE_BAND),

EngineABResult.rel_std appears to be a relative std of one side (baseline or candidate), but the test asserts abs(rec.agreement_band - 0.10) < 1e-9 for baseline_std=5.0, candidate_std=5.0 on baseline_tps=100.0. That implies rel_std = (5+5)/100 = 0.10. If rel_std is defined as only one side's std divided by baseline, it would be 0.05, failing the test. The definition of rel_std on EngineABResult needs auditing against this expectation.

_as_dims redundant guard

if isinstance(val, list | tuple) and len(val) >= 1:
    x = as_int(val[0], 1) if len(val) > 0 else 1  # `> 0` is always True here

The len(val) > 0 inner check is dead code — the outer guard already ensures len >= 1. Minor but misleading.

---

🔒 Security

out_path from caller-controlled input written without sanitization

out_path = Path(out_path)
out_path.parent.mkdir(parents=True, exist_ok=True)
out_path.write_text(...)

If out_path comes from config/API input (e.g. run_dir / "verification.json" where run_dir is user-supplied), a path like ../../etc/cron.d/gitm would traverse out of the intended directory. Consider asserting the resolved path stays within an expected root.

---

⚡ Performance

Per-iteration getattr in hot apply loop

live_engine = getattr(applicator, "engine", cfg.engine)
candidate_cfg = dict(getattr(live_engine, "gitm_llm_kwargs", None) or {})

This is inside the intervention loop. If gitm_llm_kwargs is a large dict (full vLLM config), copying it every iteration is wasteful. Snapshot only the keys that differ, or copy lazily only on kept interventions.

sorted(by_dev.items()) in worst_device_comm materializes full event lists

per_device = [
    device_comm_stats(trace.model_copy(update={"events": evs}))
    for _, evs in sorted(by_dev.items())
]

For a 4-GPU trace with millions of kernels, this creates 4 full Trace copies via model_copy. Consider passing device-filtered iterators or computing stats without copying the trace object.

---

📊 Reproducibility

t0_wall_ns and t0_ns captured sequentially, not atomically

self._t0_ns = time.perf_counter_ns()
self._t0_wall_ns = time.time_ns()

These are two separate syscalls. On a loaded host the OS can schedule the process out between them, introducing a skew of milliseconds between the monotonic and wall anchors — potentially larger than the "sub-microsecond" the comment claims. The offset is used to join scheduler samples to request timestamps; a multi-ms skew on a fast prefill could misattribute latency. Consider taking both clocks multiple times and using the pair with the smallest gap, or document the expected error bound more conservatively.

Non-deterministic max() tiebreak in worst_device_comm

return max(per_device, key=lambda c: c.exposed_comm_share_of_wall)

When two devices have equal exposed_comm_share_of_wall (common in synthetic tests), max returns whichever appears first in the list — which depends on dict insertion order. This is deterministic in Python 3.7+ but worth documenting, especially since the test asserts stats.device_id == 0 and would silently pass for the wrong reason on a tie.

---

💡 Suggestions

req_records extraction is fragile to non-dict runner output

req_records = list(run_out.get("requests") or []) if isinstance(run_out, dict) else []

If a custom runner accidentally returns a list instead of a dict, run_out is silently treated as having no requests. Consider logging a warning when run_out is not None and not isinstance(run_out, dict).

build_record docstring warns about mutation, but the caller doesn't enforce timing

baseline_cfg = dict(getattr(cfg.engine, "gitm_llm_kwargs", None) or {})
result = apply_intervention(...)  # mutates engine

The snapshot is taken before apply_intervention, which is correct. But candidate_cfg falls back to constructing from baseline_cfg + spec.knobs when live_engine reports nothing — this silent fallback could produce a misleading config diff if the engine applies additional side-effects beyond the named knob. At minimum, log when the fallback is used.

No test for write_verification with an empty records list writing a valid document — build_export handles it, but write_verification is only tested with one record. An empty verification file is a valid outcome (all interventions rolled back before A/B) and worth a test.

_percentile uses nearest-rank (round), but p99 on small samples (e.g. 3 requests) will alias to p100. For a benchmarking tool reporting percentiles to customers, this should be documented or a note added that <~100 requests makes p99 unreliable.

[github-actions]

Code Review by Gemini

The changes introduce latency metrics, collective communication analysis, and a detailed verification export. The overall structure and new features are well-designed and tested.

However, there is a potential bug in how multi-device traces are handled when calculating collective communication causes.

Bug: Incorrect multi-device trace splitting in worst_device_comm

The worst_device_comm function in gitm/optimizer/collective_signal.py aims to analyze communication per device. Its docstring correctly states that device_comm_stats (from node_rollup) assumes a single-device trace. However, the current implementation only filters KernelEvents by device_id when creating per-device traces. If the original Trace object contains MemcpyEvents or SyncEvents for a specific device, these events will be missed in the per-device analysis, leading to an underestimation of communication time for that device. node_rollup's _COMM_PATTERNS includes "memcpy", indicating that MemcpyEvents are considered communication.

File: gitm/optimizer/collective_signal.py

--- a/gitm/optimizer/collective_signal.py
+++ b/gitm/optimizer/collective_signal.py
@@ -40,19 +40,24 @@
     The worst device is the one reported because a collective is a barrier: the
     step waits on the GPU that spent longest exposed.
     """
-    kernels = trace.kernels()
-    if not kernels:
+    if not trace.events:
         return None
-    by_dev: dict[int, list] = {}
-    for k in kernels:
-        by_dev.setdefault(k.device_id, []).append(k)
-    if len(by_dev) == 1:
-        return device_comm_stats(trace)
-    per_device = [
-        device_comm_stats(trace.model_copy(update={"events": evs}))
-        for _, evs in sorted(by_dev.items())
-    ]
-    return max(per_device, key=lambda c: c.exposed_comm_share_of_wall)
+
+    # Group all events by device_id.
+    by_dev_all_events: dict[int, list] = {}
+    for event in trace.events:
+        # Only consider events that have a device_id for per-device analysis.
+        # KernelEvent, MemcpyEvent, and SyncEvent all have a device_id.
+        if hasattr(event, 'device_id'):
+            by_dev_all_events.setdefault(event.device_id, []).append(event)
+
+    if not by_dev_all_events:
+        return None
+
+    per_device_stats = []
+    for dev_id, dev_events in sorted(by_dev_all_events.items()):
+        single_device_trace = trace.model_copy(update={"events": dev_events})
+        per_device_stats.append(device_comm_stats(single_device_trace))
+
+    return max(per_device_stats, key=lambda c: c.exposed_comm_share_of_wall)