diff --git a/TraceLens/Trace2Tree/inference_iteration_roots.py b/TraceLens/Trace2Tree/inference_iteration_roots.py index 1e50ff95..9ea7c95d 100644 --- a/TraceLens/Trace2Tree/inference_iteration_roots.py +++ b/TraceLens/Trace2Tree/inference_iteration_roots.py @@ -256,21 +256,17 @@ def deepest_container(lo: float, hi: float) -> Optional[dict]: continue root["parent"] = host_node["UID"] host_node.setdefault("children", []).append(root["UID"]) - # Flag the new ancestry GPU-bearing so the descent will follow it: the - # reattached subtree carries kernels the host frames previously lacked. + gpu_uids = root.get("gpu_events", []) ancestor: Optional[dict] = host_node - while ancestor is not None and not ancestor.get("_kernel_bearing"): - ancestor["_kernel_bearing"] = True + while ancestor is not None and ancestor.get("non_gpu_path", False): + if gpu_uids: + ancestor.setdefault("gpu_events", []).extend(gpu_uids) + ancestor.pop("non_gpu_path", None) ancestor = tree.get_parent_event(ancestor) reattached += 1 return tree -def _gpu_bearing(event: dict) -> bool: - """Whether ``event`` has any GPU work under it (native or reattached).""" - return bool(event.get("gpu_events") or event.get("_kernel_bearing")) - - def _descendant_gpu_time(tree: TraceToTree, nodes: Sequence[dict]) -> float: """Total GPU time under ``nodes`` in the tree, each kernel counted once.""" seen: set = set() @@ -319,10 +315,10 @@ def _blocks_by_pattern( block.append(child) pos += 1 j += 1 - elif not _gpu_bearing(child): + elif child.get("non_gpu_path", False): j += 1 # skip a kernel-less intruder, keep matching this position else: - break # kernel-bearing deviation: a real break, stop matching + break # GPU-bearing deviation: a real break, stop matching if pos == period: blocks.append(block) i = j diff --git a/TraceLens/TraceUtils/annotation_utils.py b/TraceLens/TraceUtils/annotation_utils.py index c58ab6ac..37e87bb2 100644 --- a/TraceLens/TraceUtils/annotation_utils.py +++ b/TraceLens/TraceUtils/annotation_utils.py @@ -454,50 +454,20 @@ def dominant_cluster(groups: Dict[str, List[str]]) -> Tuple[Optional[str], float return skeleton, len(groups[skeleton]) / total -# --- cached identity and inheritance ---------------------------------------- -# Stamped onto roots whose window and identity come from different annotations. -PROVENANCE_KEY = "split_provenance" - - +# --- cached identity --------------------------------------------------------- @lru_cache(maxsize=None) def parse_annotation(name: str) -> IterationAnnotation: """Memoized parse. Treat the result as read-only; instances are shared. - - Construction runs up to seven regex matches, and detection parses the same - few hundred distinct names across hundreds of thousands of instances. """ return IterationAnnotation(name) def is_parseable(name: str) -> bool: """True when a parser recognized the name, so its metadata is real. - - The distinction matters because unparseable names still yield a full detail - dict, just a fabricated one (one decode-equivalent request), so callers that - test the numbers instead of this predicate cannot tell the two apart. """ return parse_annotation(name).matched -def inherit_identity(target: dict, source: dict) -> dict: - """Copy of ``target`` that parses as ``source``. - - Timestamps, process and thread stay with ``target``: the outer span decides - the extraction window while the inner annotation supplies phase and batch - size. Both names are recorded so the two-level choice stays auditable. - """ - prior = target.get(PROVENANCE_KEY) or {} - out = dict(target) - out["name"] = source.get("name", "") - out[PROVENANCE_KEY] = { - # Keep the original window owner when a root is relabelled more than - # once, otherwise the second pass erases where the span came from. - "window_from": prior.get("window_from") or target.get("name", ""), - "identity_from": source.get("name", ""), - } - return out - - def has_context(detail: dict) -> bool: """Step runs at least one prefill (context) request.""" return detail.get("context_requests", 0) > 0 diff --git a/TraceLens/TraceUtils/split_inference/detect_utils.py b/TraceLens/TraceUtils/split_inference/detect_utils.py index 02ef8d31..cb37dd30 100644 --- a/TraceLens/TraceUtils/split_inference/detect_utils.py +++ b/TraceLens/TraceUtils/split_inference/detect_utils.py @@ -22,12 +22,10 @@ from statistics import median from typing import Dict, Iterable, List, Optional, Sequence, Tuple -from ..annotation_utils import PROVENANCE_KEY, name_skeleton - -# Projections *enclose* the kernels they describe, so summing GPU time over both -# double-counts. Kept apart here and recombined by consumers that want both. +# A GPU annotation span *encloses* the kernels it describes, so summing GPU time +# over both double-counts. Kept apart here, recombined by consumers wanting both. GPU_KERNEL_CATEGORIES = ("kernel", "gpu_memcpy", "gpu_memset") -PROJECTION_CATEGORY = "gpu_user_annotation" +GPU_USER_ANNOTATION = "gpu_user_annotation" # Coverage to accept roots outright, and the floor below which a trace is # unsplittable rather than degraded. @@ -47,8 +45,8 @@ class DetectStatus(IntEnum): """Whether the trace can be split. Deliberately separate from phase trust.""" SPLITTABLE = 0 - NOT_SPLITTABLE = 1 - DEGRADED = 2 + NOT_SPLITTABLE = 2 + DEGRADED = 1 class PhaseConfidence(str, Enum): @@ -63,18 +61,13 @@ class PhaseConfidence(str, Enum): class CoverageReport: """Result of a GPU-time coverage audit. - ``covered_selected`` measures the roots' extraction windows, gaps included, - since that is what the output contains; ``covered_spans`` measures the bare - annotation spans. Judging on bare spans alone hunts for extra roots whenever - an iteration's sampling step runs just outside its annotation. + ``covered_selected`` measures the roots' extraction windows, gaps included; ``covered_spans`` measures the bare annotation spans. """ strategy: str - covered_any: float covered_selected: float covered_spans: float gpu_busy: float - window: Optional[Tuple[float, float]] = None @property def span_share(self) -> float: @@ -89,13 +82,7 @@ def span_share(self) -> float: @property def passes(self) -> bool: - """Whether the roots explain enough GPU work, without being stretched. - - Gating on the roots rather than on all annotations is what makes this a - real check: a run whose annotations blanket the timeline while the roots - cover fifteen of five hundred iterations scores near-perfectly on the - permissive measure. - """ + """Whether the roots explain enough GPU work, without being stretched.""" return self.covered_selected >= COVERAGE_GATE and ( self.span_share >= MIN_SPAN_SHARE ) @@ -123,12 +110,15 @@ def to_manifest(self) -> dict: "phase_confidence": self.phase_confidence.value, "n_roots": len(self.roots), "attribution_strategy": cov.strategy if cov else None, - "coverage_any_annotation": round(cov.covered_any, 4) if cov else None, "coverage_selected_roots": round(cov.covered_selected, 4) if cov else None, "coverage_root_spans_only": round(cov.covered_spans, 4) if cov else None, "root_span_share": round(cov.span_share, 4) if cov else None, } - manifest.update(self.diagnostics) + # Underscore keys are objects passed between stages -- the event map, for + # one -- not findings. Serializing them would swamp the manifest. + manifest.update( + {k: v for k, v in self.diagnostics.items() if not k.startswith("_")} + ) return manifest @@ -219,8 +209,8 @@ def group_by_thread(events: Iterable[dict]) -> Dict[Tuple, List[dict]]: class SpanSet: """A disjoint, sorted union of ``(start, end)`` time spans. - Overlapping input is the norm -- annotations nest, projections repeat per - stream -- so merging on construction is what makes membership a bisect. + Overlapping input is the norm -- annotations nest, GPU annotation spans + repeat per stream -- so merging on construction makes membership a bisect. """ def __init__(self, spans: Iterable[Tuple[float, float]] = ()): @@ -260,23 +250,20 @@ def bounds(self) -> Optional[Tuple[float, float]]: class GpuAttribution: """Attributes GPU kernels to annotations and measures coverage. - Two strategies, chosen per trace. ``projection`` (the kernel starts inside a - ``gpu_user_annotation`` span) is preferred: cheaper, and immune to graph - capture, where correlations cannot be walked at all. ``correlation`` (the - launch traces back to a CPU op inside an annotation) is a mandatory - fallback, since some traces have no projections -- notably any trace that is - itself a previous split output. - - Attribution to *any* annotation counts, not just the selected root: the - question is "are we missing whole regions of work", not fine accounting. + Two strategies, chosen per set of instances + ``gpu_span`` (the kernel starts inside a ``gpu_user_annotation`` span) is + preferred because it needs no launch link, and applies when *every* instance + has a GPU counterpart: one instance missing its span silently undercounts the + whole set, so a single miss sends the set to ``correlation`` (the launch + traces back to a CPU op inside the instance). """ - STRATEGY_PROJECTION = "projection" + STRATEGY_GPU_SPAN = "gpu_span" STRATEGY_CORRELATION = "correlation" def __init__(self, events: Iterable[dict]): self.kernels: List[dict] = [] - self.projections: List[dict] = [] + self.gpu_annotation_spans: List[dict] = [] corr_cpu: List[dict] = [] self._corr_kernels: Dict[int, List[dict]] = {} for e in events: @@ -284,8 +271,8 @@ def __init__(self, events: Iterable[dict]): if ts is None or dur is None: continue corr = (e.get("args") or {}).get("correlation") - if cat == PROJECTION_CATEGORY: - self.projections.append(e) + if cat == GPU_USER_ANNOTATION: + self.gpu_annotation_spans.append(e) elif cat in GPU_KERNEL_CATEGORIES: self.kernels.append(e) if corr is not None: @@ -295,13 +282,15 @@ def __init__(self, events: Iterable[dict]): self.kernels.sort(key=lambda x: x["ts"]) self._kernel_starts = [k["ts"] for k in self.kernels] - self.strategy = ( - self.STRATEGY_PROJECTION if self.projections else self.STRATEGY_CORRELATION - ) - # Built on demand: the projection path normally never needs it, but it - # remains reachable as a cross-check when projections look untrustworthy. + self.gpu_busy = sum(k["dur"] for k in self.kernels) + self._corr_cpu = corr_cpu self._cpu_index_cache: Optional[IntervalIndex] = None + self._spans_by_external_id: Dict[object, List[dict]] = {} + for span in self.gpu_annotation_spans: + ext = (span.get("args") or {}).get("External id") + if ext is not None: + self._spans_by_external_id.setdefault(ext, []).append(span) @property def _cpu_index(self) -> IntervalIndex: @@ -337,93 +326,47 @@ def kernels_for(self, spans: Sequence[dict]) -> List[dict]: out.append(k) return out - def gpu_time_by_correlation(self, spans: Sequence[dict]) -> float: - """GPU time launched from CPU ops inside ``spans``.""" - return sum(k["dur"] for k in self.kernels_for(spans)) - - def gpu_time_for_family(self, skeleton: str, instances: Sequence[dict]) -> float: + def attributed_kernels(self, instances: Sequence[dict]) -> Tuple[List[dict], str]: + """Kernels belonging to ``instances``, and which strategy found them.""" + matched = [ + self._spans_by_external_id.get((e.get("args") or {}).get("External id")) + for e in instances + ] + if instances and all(matched): + spans = SpanSet.of_events([s for group in matched for s in group]) + return [ + k for k in self._kernels_in(spans.bounds) if spans.covers(k["ts"]) + ], self.STRATEGY_GPU_SPAN + return self.kernels_for(instances), self.STRATEGY_CORRELATION + + def gpu_time_for_family(self, instances: Sequence[dict]) -> float: """GPU time attributable to one annotation family.""" - if self.strategy != self.STRATEGY_PROJECTION: - return self.gpu_time_by_correlation(instances) - spans = self._projection_union({skeleton}) - if not spans: - return 0.0 - return sum( - k["dur"] for k in self._kernels_in(spans.bounds) if spans.covers(k["ts"]) - ) - - # -- coverage ------------------------------------------------------------ - def audit( - self, annotations: Sequence[dict], roots: Sequence[dict] - ) -> CoverageReport: - """Measure GPU coverage by all annotations, and by the roots alone. - - High ``covered_any`` next to low ``covered_selected`` means the roots - sit at the wrong nesting level, and widening should fix it. - """ - # Credit the roots two ways: projections survive graph capture, and - # launch correlations work for roots no annotation produced. Matching by - # name alone scores every synthetic root zero and stalls escalation. - if self.strategy == self.STRATEGY_PROJECTION: - by_name = self._projection_union(_window_names(roots)) - any_spans = self._projection_union(None) - else: - by_name = SpanSet() - any_spans = SpanSet.of_events(self.kernels_for(annotations)) - - all_kernels = self.kernels - busy = sum(k["dur"] for k in all_kernels) - if busy <= 0: - return CoverageReport(self.strategy, 0.0, 0.0, 0.0, 0.0, None) - - tiles, _ = build_root_tiles(roots) - tile_spans = [ - {"pid": pid, "tid": tid, "ts": start, "dur": end - start} - for (pid, tid, _), (start, end) in tiles.items() + return sum(k["dur"] for k in self.attributed_kernels(instances)[0]) + + def audit(self, roots: Sequence[dict]) -> CoverageReport: + """Measure what share of the trace's GPU time the roots account for.""" + root_kernels, strategy = self.attributed_kernels(roots) + if self.gpu_busy <= 0: + return CoverageReport(strategy, 0.0, 0.0, self.gpu_busy) + + # Extraction hands out whole windows, so judge the window: first root's + # start to the last one's end, per thread + windows = [ + { + "pid": pid, + "tid": tid, + "ts": group[0].get("ts", 0), + "dur": group[-1].get("ts", 0) + + group[-1].get("dur", 0) + - group[0].get("ts", 0), + } + for (pid, tid), group in group_by_thread(roots).items() ] - covered = {} - for label, spans in ( - ("any", any_spans), - ("spans", by_name | SpanSet.of_events(self.kernels_for(roots))), - ("tiles", by_name | SpanSet.of_events(self.kernels_for(tile_spans))), - ): - covered[label] = sum(k["dur"] for k in all_kernels if spans.covers(k["ts"])) + selected = {id(k): k for k in root_kernels} + selected.update({id(k): k for k in self.kernels_for(windows)}) return CoverageReport( - self.strategy, - covered["any"] / busy, - covered["tiles"] / busy, - covered["spans"] / busy, - busy, - None, - ) - - def _projection_union(self, names: Optional[set]) -> SpanSet: - """Union of projection spans, optionally restricted by annotation name. - - Taken across every GPU thread rather than per stream: a kernel on a side - stream inside an annotated region is genuinely covered, and per-stream - matching would call it uncovered and escalate for nothing. - """ - return SpanSet.of_events( - p - for p in self.projections - if names is None or name_skeleton(p.get("name", "")) in names + strategy, + sum(k["dur"] for k in selected.values()) / self.gpu_busy, + sum(k["dur"] for k in root_kernels) / self.gpu_busy, + self.gpu_busy, ) - - -def _window_names(roots: Sequence[dict]) -> set: - """Skeletons to match projections against for the selected roots. - - A root enriched in step 1.5 carries the inner annotation's name, so its - projections are recorded under the outer span it actually came from. - """ - names = set() - for r in roots: - prov = r.get(PROVENANCE_KEY) or {} - names.add(name_skeleton(prov.get("window_from") or r.get("name", ""))) - return names - - -# Escalation probes were removed: on the fallback corpus they never once improved -# a split, and the coverage gate now grades roots directly (splittable / degraded -# / not-splittable) with no intermediate probe ladder. diff --git a/TraceLens/TraceUtils/split_inference/execution_roots.py b/TraceLens/TraceUtils/split_inference/execution_roots.py index 5db87a3a..35d6f888 100644 --- a/TraceLens/TraceUtils/split_inference/execution_roots.py +++ b/TraceLens/TraceUtils/split_inference/execution_roots.py @@ -12,12 +12,12 @@ _descendant_gpu_time, _entry_roots, _reattach_worker_threads, - GPU_KERNEL_CATS, ) from ...Trace2Tree.trace_to_tree import TraceToTree from ..annotation_utils import ( find_known_annotations, is_parseable, + name_skeleton, ) from .detect_utils import ( COVERAGE_FLOOR, @@ -37,6 +37,9 @@ detect_from_sibling_roots, ) +# A candidate must already explain this much GPU work before bookends are worth +# trying: below it the shortfall is the pattern being wrong, not a missing warmup. +BOOKEND_FLOOR = 0.50 __all__ = [ "COVERAGE_FLOOR", @@ -59,7 +62,7 @@ def _detect_from_known_annotations( known = find_known_annotations(annotations) if not known: return None - + # ToDo: We do not need this check, clean up in a follow up PR. labelled = sum(1 for r in known if is_parseable(r.get("name", ""))) if labelled == len(known): confidence = PhaseConfidence.HIGH @@ -92,6 +95,8 @@ def _detect_from_unknown_family( regular = [f for f in families if f.regular] if not regular: return None + # Find the family with the lowest rank, which is the family with the most GPU work. + # Tie-break by the family with lowest interarrival CV family = min(regular, key=lambda f: f.rank) return RootSet( roots=sorted(family.instances, key=lambda e: e.get("ts", 0)), @@ -105,71 +110,104 @@ def _detect_from_unknown_family( ) +# Enough to show a mixed root set without one pathological case flooding the log. +_LOGGED_SKELETONS = 3 + + +def _log_attempt(step: str, root_set: Optional[RootSet]) -> None: + """Report what one cascade step found, so the escalation is traceable.""" + if root_set is None: + print(f"[roots] {step}: no candidate") + return + cov = root_set.coverage + + skeletons = sorted({name_skeleton(r.get("name", "")) for r in root_set.roots}) + shown = ", ".join( + s if len(s) <= 60 else s[:57] + "..." for s in skeletons[:_LOGGED_SKELETONS] + ) + if len(skeletons) > _LOGGED_SKELETONS: + shown += f", +{len(skeletons) - _LOGGED_SKELETONS} more" + head = ( + f"[roots] {step}: {len(root_set.roots)} roots via {root_set.method} " + f"[{shown}]" + ) + if cov is None: + print(f"{head}, coverage not measured, status={root_set.status.name}") + return + print( + f"{head}, covered_selected={cov.covered_selected:.1%}, " + f"span_share={cov.span_share:.1%} " + f"({cov.strategy}), status={root_set.status.name}" + ) + + def _try_bookend_enhancement( candidate: RootSet, tree: TraceToTree, total_gpu: float, ) -> Optional[RootSet]: - """Add warmup and/or wrapup bookend roots to improve coverage. - - Uses the before_uids / after_uids lists stored in diagnostics by the - branch_descent and sibling_roots detectors. These are the UIDs of - GPU-bearing siblings that fall outside the repeating pattern. - Only adds a bookend if it contributes GPU time. + """Add warmup and wrapup roots for the work outside the repeating pattern. + + A period anchors on a repeating run, so a warmup pass that does an + iteration's work with extra setup, and a wrapup that trails it, sit outside + every block and their GPU time is reported as unaccounted. The detectors + leave those leftovers in ``before_uids``/``after_uids``; this turns each side + into one root, but only if it carries GPU time -- a bookend of pure CPU + teardown would dilute the split without explaining anything. """ if not total_gpu or not candidate.roots: return None - before_uids = candidate.diagnostics.get("before_uids", []) - after_uids = candidate.diagnostics.get("after_uids", []) - if not before_uids and not after_uids: + uid_map = tree.events_by_uid + before = [ + uid_map[uid] + for uid in candidate.diagnostics.get("before_uids", ()) + if uid in uid_map + ] + after = [ + uid_map[uid] + for uid in candidate.diagnostics.get("after_uids", ()) + if uid in uid_map + ] + if not before and not after: return None - before = [tree.events_by_uid[uid] for uid in before_uids if uid in tree.events_by_uid] - after = [tree.events_by_uid[uid] for uid in after_uids if uid in tree.events_by_uid] - - before_gpu = _descendant_gpu_time(tree, before) if before else 0 - after_gpu = _descendant_gpu_time(tree, after) if after else 0 - - iter_gpu = candidate.diagnostics.get("iter_gpu_time", 0) - new_cov = (before_gpu + iter_gpu + after_gpu) / total_gpu + before_gpu = _descendant_gpu_time(tree, before) if before else 0.0 + after_gpu = _descendant_gpu_time(tree, after) if after else 0.0 + if before_gpu <= 0 and after_gpu <= 0: + return None - new_roots = list(candidate.roots) - if before and before_gpu > 0: - before_sorted = sorted(before, key=lambda e: e["ts"]) - warmup = dict(before_sorted[0]) - warmup["name"] = "warmup" - warmup["dur"] = ( - before_sorted[-1]["ts"] + before_sorted[-1].get("dur", 0) - before_sorted[0]["ts"] - ) - new_roots.insert(0, warmup) - if after and after_gpu > 0: - after_sorted = sorted(after, key=lambda e: e["ts"]) - wrapup = dict(after_sorted[0]) - wrapup["name"] = "wrapup" - wrapup["dur"] = ( - after_sorted[-1]["ts"] + after_sorted[-1].get("dur", 0) - after_sorted[0]["ts"] - ) - new_roots.append(wrapup) + def _span(events: Sequence[dict], name: str) -> dict: + ordered = sorted(events, key=lambda e: e["ts"]) + last = ordered[-1] + root = dict(ordered[0]) + root["name"] = name + root["dur"] = last["ts"] + last.get("dur", 0) - ordered[0]["ts"] + return root - diag = dict(candidate.diagnostics) - diag["bookend_enhancement"] = True - diag["branch_coverage"] = round(new_cov, 4) + roots = list(candidate.roots) + diagnostics = dict(candidate.diagnostics) if before_gpu > 0: - diag["warmup_gpu_pct"] = round(100 * before_gpu / total_gpu, 1) + roots.insert(0, _span(before, "warmup")) + diagnostics["warmup_gpu_pct"] = round(100 * before_gpu / total_gpu, 1) if after_gpu > 0: - diag["wrapup_gpu_pct"] = round(100 * after_gpu / total_gpu, 1) - + roots.append(_span(after, "wrapup")) + diagnostics["wrapup_gpu_pct"] = round(100 * after_gpu / total_gpu, 1) + + coverage = ( + before_gpu + diagnostics.get("iter_gpu_time", 0.0) + after_gpu + ) / total_gpu + diagnostics["bookend_enhancement"] = True + diagnostics["branch_coverage"] = round(coverage, 4) return RootSet( - roots=new_roots, + roots=roots, method=candidate.method, phase_confidence=candidate.phase_confidence, - status=_grade(new_cov), - diagnostics=diag, + status=_grade(coverage), + diagnostics=diagnostics, ) - def find_iteration_roots(events: Sequence[dict]) -> RootSet: """Find iteration roots and report how much GPU work they account for. @@ -185,12 +223,13 @@ def find_iteration_roots(events: Sequence[dict]) -> RootSet: annotations = collect_annotations(events) best_fallback: Optional[RootSet] = None - def _try(root_set: Optional[RootSet]) -> Optional[RootSet]: + def _try(step: str, root_set: Optional[RootSet]) -> Optional[RootSet]: """Audit coverage; return the root_set if it passes, else save as fallback.""" nonlocal best_fallback if root_set is None: + _log_attempt(step, None) return None - coverage = attribution.audit(annotations, root_set.roots) + coverage = attribution.audit(root_set.roots) root_set.coverage = coverage known_labels = root_set.phase_confidence is PhaseConfidence.HIGH @@ -198,12 +237,14 @@ def _try(root_set: Optional[RootSet]) -> Optional[RootSet]: known_labels or not root_set.diagnostics.get("suspiciously_few_roots") ): root_set.status = DetectStatus.SPLITTABLE + _log_attempt(step, root_set) return root_set if coverage.covered_selected >= COVERAGE_FLOOR: root_set.status = DetectStatus.DEGRADED else: root_set.status = DetectStatus.NOT_SPLITTABLE + _log_attempt(step, root_set) if best_fallback is None or ( root_set.status.value < best_fallback.status.value @@ -212,18 +253,20 @@ def _try(root_set: Optional[RootSet]) -> Optional[RootSet]: return None # --- 1. Known annotation patterns ----------------------------------------- - result = _try(_detect_from_known_annotations(annotations)) + result = _try("1 known annotations", _detect_from_known_annotations(annotations)) if result is not None: return result # --- 2. Unknown annotation families --------------------------------------- - result = _try(_detect_from_unknown_family(annotations, attribution)) + result = _try( + "2 unknown families", _detect_from_unknown_family(annotations, attribution) + ) if result is not None: return result # --- 3 & 4. Tree-based detectors (built once) ---------------------------- try: - tree = TraceToTree(list(events), prune_nongpu_paths=False) + tree = TraceToTree(list(events), prune_nongpu_paths=True) tree.build_tree(add_python_func=True) except Exception as exc: print(f"TraceToTree build failed ({exc}), skipping tree detectors.") @@ -239,37 +282,53 @@ def _try(root_set: Optional[RootSet]) -> Optional[RootSet]: tree = _reattach_worker_threads(tree) entry_roots = _entry_roots(tree) total_gpu = _total_gpu_time(tree) - uid_map = tree.events_by_uid def _attach_uid_map(root_set: RootSet) -> RootSet: - root_set.diagnostics["_events_by_uid"] = uid_map + """Hand extraction the tree it needs to collect each root's ancestors. + + Underscore-prefixed so ``to_manifest`` leaves it out: the map is the whole + trace, and serializing it into a manifest would dwarf the manifest. + """ + root_set.diagnostics["_events_by_uid"] = tree.events_by_uid return root_set + if total_gpu == 0: + return RootSet( + roots=[], + method="none", + status=DetectStatus.NOT_SPLITTABLE, + diagnostics={"reason": "no GPU work"}, + ) # --- 3. Branch descent ---------------------------------------------------- branch_set = detect_from_branch_descent(tree, entry_roots, total_gpu) + if branch_set is not None: + branch_set.coverage = attribution.audit(branch_set.roots) + _log_attempt("3 branch descent", branch_set) if branch_set is not None and branch_set.status is DetectStatus.SPLITTABLE: return _attach_uid_map(branch_set) # --- 4. Sibling roots ---------------------------------------------------- sibling_set = detect_from_sibling_roots(tree, entry_roots, total_gpu) + if sibling_set is not None: + sibling_set.coverage = attribution.audit(sibling_set.roots) + _log_attempt("4 sibling roots", sibling_set) if sibling_set is not None and sibling_set.status is DetectStatus.SPLITTABLE: return _attach_uid_map(sibling_set) - # --- 5. Bookend enhancement ------------------------------------------------ - # If a generic detector found iterations covering >=50% of GPU time but - # not enough to pass, check whether adding a warmup block (before first - # iteration) and/or wrapup block (after last iteration) improves coverage. - BOOKEND_FLOOR = 0.50 + # --- 5. Bookend enhancement ---------------------------------------------- + # Only for a candidate that already found the pattern: the shortfall it can + # fix is a warmup or wrapup left outside the blocks, not a wrong period. bookend_set = None for candidate in (branch_set, sibling_set): if candidate is None or not candidate.roots: continue - cov = candidate.diagnostics.get("branch_coverage", 0) - if cov < BOOKEND_FLOOR: + if candidate.diagnostics.get("branch_coverage", 0) < BOOKEND_FLOOR: continue bookend_set = _try_bookend_enhancement(candidate, tree, total_gpu) if bookend_set is not None: + bookend_set.coverage = attribution.audit(bookend_set.roots) break + _log_attempt("5 bookend enhancement", bookend_set) if bookend_set is not None and bookend_set.status is DetectStatus.SPLITTABLE: return _attach_uid_map(bookend_set) @@ -279,9 +338,11 @@ def _attach_uid_map(root_set: RootSet) -> RootSet: candidate is not None and candidate.status is not DetectStatus.NOT_SPLITTABLE ): + _log_attempt("fallback (best usable)", candidate) return _attach_uid_map(candidate) for candidate in (bookend_set, branch_set, sibling_set, best_fallback): if candidate is not None: + _log_attempt("fallback (last resort)", candidate) return _attach_uid_map(candidate) return RootSet( roots=[], diff --git a/TraceLens/TraceUtils/split_inference/root_detection.py b/TraceLens/TraceUtils/split_inference/root_detection.py index afa1f919..5a59a892 100644 --- a/TraceLens/TraceUtils/split_inference/root_detection.py +++ b/TraceLens/TraceUtils/split_inference/root_detection.py @@ -6,10 +6,9 @@ """Annotation families and tree-based detection helpers. -Step 1 groups every annotation into families, known and unknown alike. Step 1.5 -chooses which nesting level is the iteration and where its metadata comes from. +Steps 1 and 2 group every annotation into families, known and unknown alike. ``detect_from_branch_descent`` and ``detect_from_sibling_roots`` handle the -tree-based detection path. +tree-based detection path for traces no annotation family explains. """ from dataclasses import dataclass @@ -28,7 +27,6 @@ _blocks_by_pattern, _descendant_gpu_time, _find_repeating_period, - _gpu_bearing, ) from ...Trace2Tree.trace_to_tree import TraceToTree from ..annotation_utils import ( @@ -46,19 +44,13 @@ RootSet, ) -# A family must enclose more than this share of another family's instances for -# the nesting relation to hold. Majority rather than unanimity, because real -# traces have ragged edges: warmup instances outside the loop, a truncated final -# iteration. Requiring every instance rejects the correct relation nearly always. -NESTING_MAJORITY = 0.5 - @dataclass class AnnotationFamily: """All instances of one logical annotation, keyed by its skeleton. - CPU-side only: projections duplicate one annotation across streams, so any - count taken from them is inflated. They give the "has GPU work" signal only. + CPU-side only: GPU annotation spans duplicate one annotation across streams, + so any count from them is inflated. They give "has GPU work" signal only. """ skeleton: str @@ -83,11 +75,7 @@ def rank(self) -> tuple: def _interarrival_cv(instances: Sequence[dict]) -> float: - """Variation in the spacing between consecutive instances. - - A once-per-iteration annotation arrives on a steady cadence. Used to rank - families, not reject them, so it needs no threshold. - """ + """Variation in the spacing between consecutive instances.""" stamps = sorted(e.get("ts", 0) for e in instances) gaps = [b - a for a, b in zip(stamps, stamps[1:]) if b > a] if len(gaps) < 2: @@ -99,12 +87,7 @@ def _interarrival_cv(instances: Sequence[dict]) -> float: def collect_annotations(events: Sequence[dict]) -> List[dict]: """CPU-side *marker* annotations, in time order. - A split point must be a semantic iteration marker, not an operation. Real - tensor ops -- collectives especially (``nccl:*`` / ``gloo:*``) -- are emitted - as annotations too, but they carry ``Input Dims`` because they act on - tensors. Iteration markers (``step[DECODE]``, ``execute_...``, ``DataLoader``) - never do. Excluding anything with ``Input Dims`` keeps a collective from being - chosen as the iteration root, which is how an all-gather was winning before. + Input Dims are checked for to avoid selecting operation annotations as iteration roots. """ annotations = [ e @@ -121,14 +104,7 @@ def collect_annotations(events: Sequence[dict]) -> List[dict]: def build_families( annotations: Sequence[dict], attribution: GpuAttribution ) -> List[AnnotationFamily]: - """Group annotations into families and drop the ones with no GPU work. - - Pruning first is the point of the ordering: roughly half the families are - pure CPU scheduling chatter that nesting analysis would otherwise carry. - - A family survives if *any* instance has GPU work, not all -- a scheduler - family may have thousands of instances where only hundreds wrap real work. - """ + """Group annotations and drop the ones with no GPU work.""" grouped: Dict[str, List[dict]] = {} for event in annotations: grouped.setdefault(name_skeleton(event.get("name", "")), []).append(event) @@ -137,117 +113,214 @@ def build_families( AnnotationFamily( skeleton=skeleton, instances=instances, - gpu_time=attribution.gpu_time_for_family(skeleton, instances), + gpu_time=attribution.gpu_time_for_family(instances), parseable=any(is_parseable(e.get("name", "")) for e in instances), interarrival_cv=_interarrival_cv(instances), ) for skeleton, instances in grouped.items() ] - # "No projection means no GPU work" is an inference, not a guarantee. If it - # would delete every family on a trace that plainly has kernels, it is wrong - # here -- fall back to following launch correlations instead. - if families and attribution.kernels and not any(f.gpu_time for f in families): - for family in families: - family.gpu_time = attribution.gpu_time_by_correlation(family.instances) - return [f for f in families if f.gpu_time > 0] # --- steps ------------------------------------------------------------------ +def _total_gpu_time(tree: TraceToTree) -> float: + return sum( + e.get("dur", 0) + for e in tree.events_by_uid.values() + if e.get("cat") in GPU_KERNEL_CATS + ) -def _compute_gpu_signature( - tree: TraceToTree, block: Sequence[dict], -) -> List[str]: - """Normalized names contributing >50% of GPU time in a representative block. +def _grade(coverage: float) -> DetectStatus: + if coverage >= COVERAGE_GATE: + return DetectStatus.SPLITTABLE + if coverage >= COVERAGE_FLOOR: + return DetectStatus.DEGRADED + return DetectStatus.NOT_SPLITTABLE - Returns the smallest prefix of names (sorted by descending GPU contribution) - whose combined GPU time exceeds half the block total. Used by the bookend - promotion step to decide whether a prefix/suffix region is a real iteration. + +def _child_groups(tree: TraceToTree, ordered: Sequence[dict]) -> tuple: + """Children grouped by name, with the GPU time under each group. + + Computed once per node and shared by both candidate paths, since walking the + subtree is the expensive part of visiting a node. """ - name_gpu: Dict[str, float] = {} - total = 0.0 - for event in block: - gpu = _descendant_gpu_time(tree, [event]) - norm = normalize_name_for_comparison(event.get("name", "")) - name_gpu[norm] = name_gpu.get(norm, 0) + gpu - total += gpu - if not total: - return [] - sorted_names = sorted(name_gpu.items(), key=lambda x: -x[1]) - sig: List[str] = [] - accumulated = 0.0 - for name, gpu in sorted_names: - sig.append(name) - accumulated += gpu - if accumulated > total * 0.5: - break - return sig + groups: Dict[str, List[dict]] = {} + for child in ordered: + groups.setdefault(child.get("name", ""), []).append(child) + gpu = {name: _descendant_gpu_time(tree, inst) for name, inst in groups.items()} + return groups, gpu -def _matches_gpu_signature( - events: Sequence[dict], signature: List[str], -) -> bool: - """True when *signature* names appear as a subsequence in *events*.""" - names = [normalize_name_for_comparison(e.get("name", "")) for e in events] - sig_idx = 0 - for name in names: - if sig_idx < len(signature) and name == signature[sig_idx]: - sig_idx += 1 - return sig_idx == len(signature) +def _branch_candidate( + tree: TraceToTree, + roots: List[dict], + blocked: List[dict], + total_gpu: float, + depth: int, + source: str, + period: Optional[int], + gpu_time: Optional[float] = None, + extra: Optional[Dict] = None, +) -> RootSet: + if gpu_time is None: + gpu_time = _descendant_gpu_time(tree, blocked) + cov = gpu_time / total_gpu + diagnostics = { + "period_label_tier": BRANCH_DESCENT_TIER, + "period": period, + "period_depth": depth, + "branch_source": source, + "branch_coverage": round(cov, 4), + "iter_gpu_time": gpu_time, + } + diagnostics.update(extra or {}) + return RootSet( + roots=roots, + method=f"generic:{BRANCH_DESCENT_TIER}", + phase_confidence=PhaseConfidence.UNKNOWN, + status=_grade(cov), + diagnostics=diagnostics, + ) + + + +def _bookend_diagnostics( + blocked: Sequence[dict], + prefix: Sequence[dict], + suffix: Sequence[dict], +) -> Dict: + """The unpromoted remainder, for the cascade to offer as bookend roots. + + Only UIDs: resolving them to events and pricing their GPU time is the + cascade's job, and doing it here would walk the subtree a second time on the + descent's hot path. + """ + blocked_uids = {e.get("UID") for e in blocked} + return { + "before_uids": [ + e.get("UID") for e in prefix if e.get("UID") not in blocked_uids + ], + "after_uids": [ + e.get("UID") for e in suffix if e.get("UID") not in blocked_uids + ], + } -def _promote_bookend_iterations( +def _periodic_candidate( tree: TraceToTree, + node: dict, ordered: Sequence[dict], - start: int, - unit_blocks: List[List[dict]], -) -> tuple: - """Promote prefix/suffix regions to iteration blocks when they match - the GPU signature of the detected iterations. - - Returns ``(unit_blocks, prefix, suffix)`` where *unit_blocks* may have - gained entries at the front/back and *prefix*/*suffix* contain only the - events that were NOT promoted. - """ - rep_block = unit_blocks[len(unit_blocks) // 2] - iter_sig = _compute_gpu_signature(tree, rep_block) + gputime_by_name: Dict[str, float], + total_gpu: float, + depth: int, +) -> Optional[RootSet]: + """One candidate per repetition of a contiguous repeating child-name run. - prefix = list(ordered[:start]) + Frames whose *name* carries no GPU work anywhere are dropped before the + period search. + """ + live = [e for e in ordered if gputime_by_name.get(e.get("name", ""), 0.0) > 0] + if len(live) < MIN_LABEL_CHILDREN: + return None + # Normalized names, because ``_blocks_by_pattern`` matches on them: the two + # have to agree, and normalizing is what lets a python frame repeat at all + # when its line number shifts between iterations. + period, pattern, start = _find_repeating_period( + [normalize_name_for_comparison(e.get("name", "")) for e in live] + ) + if period is None or period == 1: + return None + unit_blocks = _blocks_by_pattern(live, pattern, start) + if len(unit_blocks) < MIN_LABEL_CHILDREN: + return None + prefix = list(live[:start]) blocked_uids = {e.get("UID") for b in unit_blocks for e in b} last_block_end = ( unit_blocks[-1][-1]["ts"] + unit_blocks[-1][-1].get("dur", 0) ) suffix = [ - e for e in ordered + e for e in live if e["ts"] >= last_block_end and e.get("UID") not in blocked_uids ] + iteration_roots: List[dict] = [] + blocked: List[dict] = [] + for block in unit_blocks: + first, last = block[0], block[-1] + event = dict(first) + event["name"] = node.get("name", event.get("name", "")) + event["dur"] = (last["ts"] + last.get("dur", 0)) - first["ts"] + iteration_roots.append(event) + blocked.extend(block) + return _branch_candidate( + tree, + iteration_roots, + blocked, + total_gpu, + depth, + "period", + period, + extra=_bookend_diagnostics(blocked, prefix, suffix), + ) - if iter_sig and prefix and _matches_gpu_signature(prefix, iter_sig): - unit_blocks.insert(0, prefix) - prefix = [] - if iter_sig and suffix and _matches_gpu_signature(suffix, iter_sig): - unit_blocks.append(suffix) - suffix = [] - - return unit_blocks, prefix, suffix +def _grouped_candidate( + tree: TraceToTree, + groups: Dict[str, List[dict]], + gputime_by_name: Dict[str, float], + total_gpu: float, + depth: int, +) -> Optional[RootSet]: + """One candidate from the recurring child frame that carries the GPU work. -def _total_gpu_time(tree: TraceToTree) -> float: - return sum( - e.get("dur", 0) - for e in tree.events_by_uid.values() - if e.get("cat") in GPU_KERNEL_CATS - ) + A *conditional* loop body has no contiguous period. Grouping by name ignores + the gaps, exactly as ``build_families`` does one level up. + Ranked by GPU time, then cadence, then count. Only the winning family becomes roots. + """ + ranked = [ + ( + gputime_by_name[name], + -_interarrival_cv(instances), + len(instances), + name, + instances, + ) + for name, instances in groups.items() + if len(instances) >= MIN_LABEL_CHILDREN + ] + if not ranked: + return None + ranked.sort(key=lambda r: (r[0], r[1], r[2]), reverse=True) + gpu_time, _, _, _, instances = ranked[0] + if gpu_time <= 0: + return None -def _grade(coverage: float) -> DetectStatus: - if coverage >= COVERAGE_GATE: - return DetectStatus.SPLITTABLE - if coverage >= COVERAGE_FLOOR: - return DetectStatus.DEGRADED - return DetectStatus.NOT_SPLITTABLE + # Siblings own disjoint subtrees, so their GPU times add without overlap. + with_gpu = [r for r in ranked if r[0] > 0] + extra = { + "branch_families_with_gpu": len(with_gpu), + "branch_families_combined_coverage": round( + sum(r[0] for r in with_gpu) / total_gpu, 4 + ), + "branch_runner_up_families": [ + {"name": name, "count": count, "gpu_share": round(gpu / total_gpu, 4)} + for gpu, _, count, name, _ in with_gpu[1:4] + ], + } + # The group's own frames, so there is no synthetic event to fabricate. + return _branch_candidate( + tree, + list(instances), + list(instances), + total_gpu, + depth, + "frame_family", + None, + gpu_time=gpu_time, + extra=extra, + ) def detect_from_branch_descent( @@ -257,14 +330,12 @@ def detect_from_branch_descent( ) -> Optional[RootSet]: """Walk the call tree to find the frame whose children repeat and cover the GPU. - The BFS keeps descending past nodes whose repeating pattern explains too - little GPU work (sub-loops). Returns a :class:`RootSet` for the best - candidate found, or ``None`` when no repeating pattern exists at all. - The caller decides whether coverage is acceptable. + Two candidates compete at every node -- a contiguous repeating name run, and + the best name-grouped child frame -- and the one covering more GPU work wins. + The BFS keeps descending past nodes whose candidates explain too little GPU + work (sub-loops). Returns the best :class:`RootSet` found, or ``None``. The + caller decides whether coverage is acceptable. """ - if not total_gpu: - return None - best: Optional[RootSet] = None queue = deque((r, 0) for r in entry_roots) visited = 0 @@ -274,60 +345,43 @@ def detect_from_branch_descent( if visited > BRANCH_MAX_NODES: break children = tree.get_children_events(node) - gpu_children = [c for c in children if _gpu_bearing(c)] - if len(gpu_children) >= MIN_LABEL_CHILDREN: - ordered = sorted(gpu_children, key=lambda e: e.get("ts", 0)) - period, pattern, start = _find_repeating_period( - [normalize_name_for_comparison(e.get("name", "")) for e in ordered] - ) - if period is not None: - unit_blocks = _blocks_by_pattern(ordered, pattern, start) - if len(unit_blocks) >= MIN_LABEL_CHILDREN: - unit_blocks, prefix, suffix = _promote_bookend_iterations( - tree, ordered, start, unit_blocks, - ) - iteration_roots = [] - blocked = [] - for block in unit_blocks: - first, last = block[0], block[-1] - event = dict(first) - event["name"] = node.get("name", event.get("name", "")) - event["dur"] = (last["ts"] + last.get("dur", 0)) - first["ts"] - iteration_roots.append(event) - blocked.extend(block) - iter_gpu_time = _descendant_gpu_time(tree, blocked) - cov = iter_gpu_time / total_gpu - blocked_uids = {e.get("UID") for e in blocked} - before_uids = [ - e.get("UID") for e in prefix - if e.get("UID") not in blocked_uids - ] - after_uids = [ - e.get("UID") for e in suffix - if e.get("UID") not in blocked_uids - ] - candidate = RootSet( - roots=iteration_roots, - method=f"generic:{BRANCH_DESCENT_TIER}", - phase_confidence=PhaseConfidence.UNKNOWN, - status=_grade(cov), - diagnostics={ - "period_label_tier": BRANCH_DESCENT_TIER, - "period": period, - "period_depth": depth, - "branch_coverage": round(cov, 4), - "iter_gpu_time": iter_gpu_time, - "before_uids": before_uids, - "after_uids": after_uids, - }, - ) - if cov >= BRANCH_COVERAGE_GATE: - return candidate - if best is None or cov > best.diagnostics.get("branch_coverage", 0): - best = candidate + if len(children) >= MIN_LABEL_CHILDREN: + ordered = sorted(children, key=lambda e: e.get("ts", 0)) + groups, gputime_by_name = _child_groups(tree, ordered) + for candidate in ( + _periodic_candidate( + tree, node, ordered, gputime_by_name, total_gpu, depth + ), + _grouped_candidate(tree, groups, gputime_by_name, total_gpu, depth), + ): + if candidate is None: + continue + cov = candidate.diagnostics["branch_coverage"] + # A frame explaining no GPU work is never an answer, however + # early it is found: pure CPU output processing was winning. + if cov > 0 and ( + best is None or cov > best.diagnostics["branch_coverage"] + ): + best = candidate + if ( + best is not None + and best.diagnostics["branch_coverage"] >= BRANCH_COVERAGE_GATE + ): + break for child in children: - if _gpu_bearing(child): + if not child.get("non_gpu_path", False): queue.append((child, depth + 1)) + + if best is not None: + print( + f"[roots] branch best: {len(best.roots)} roots via " + f"{best.diagnostics['branch_source']} under " + f"'{best.roots[0].get('name', '')[:70]}', " + f"period={best.diagnostics['period']}, " + f"depth={best.diagnostics['period_depth']}, " + f"coverage={best.diagnostics['branch_coverage']:.1%} " + f"-> {best.status.name}" + ) return best @@ -352,17 +406,22 @@ def detect_from_sibling_roots( if period is None: return None - n_blocks = (len(ordered) - start) // period + blocks = (len(ordered) - start) // period unit_blocks = [ - list(ordered[start + i * period : start + (i + 1) * period]) - for i in range(n_blocks) + list(ordered[start + index * period : start + (index + 1) * period]) + for index in range(blocks) ] if not unit_blocks: return None - - unit_blocks, prefix, suffix = _promote_bookend_iterations( - tree, ordered, start, unit_blocks, + prefix = list(ordered[:start]) + blocked_uids_tmp = {e.get("UID") for b in unit_blocks for e in b} + last_block_end = ( + unit_blocks[-1][-1]["ts"] + unit_blocks[-1][-1].get("dur", 0) ) + suffix = [ + e for e in ordered + if e["ts"] >= last_block_end and e.get("UID") not in blocked_uids_tmp + ] sibling_roots = [] blocked = [] @@ -373,30 +432,21 @@ def detect_from_sibling_roots( sibling_roots.append(event) blocked.extend(block) - iter_gpu_time = _descendant_gpu_time(tree, blocked) if total_gpu else 0.0 + iter_gpu_time = _descendant_gpu_time(tree, blocked) cov = iter_gpu_time / total_gpu if total_gpu else 0.0 - blocked_uids = {e.get("UID") for e in blocked} - before_uids = [ - e.get("UID") for e in prefix - if e.get("UID") not in blocked_uids - ] - after_uids = [ - e.get("UID") for e in suffix - if e.get("UID") not in blocked_uids - ] + diagnostics = _bookend_diagnostics(blocked, prefix, suffix) + diagnostics.update( + { + "period_label_tier": "sibling_roots", + "period": period, + "branch_coverage": round(cov, 4), + "iter_gpu_time": iter_gpu_time, + } + ) return RootSet( roots=sibling_roots, method="generic:sibling_roots", phase_confidence=PhaseConfidence.UNKNOWN, status=_grade(cov), - diagnostics={ - "period_label_tier": "sibling_roots", - "period": period, - "branch_coverage": round(cov, 4), - "iter_gpu_time": iter_gpu_time, - "before_uids": before_uids, - "after_uids": after_uids, - }, + diagnostics=diagnostics, ) - - diff --git a/TraceLens/TraceUtils/split_inference/steady_state_window.py b/TraceLens/TraceUtils/split_inference/steady_state_window.py index 86bf158d..77d4cf28 100644 --- a/TraceLens/TraceUtils/split_inference/steady_state_window.py +++ b/TraceLens/TraceUtils/split_inference/steady_state_window.py @@ -42,6 +42,7 @@ def _identify_regions_by_peak( values: list[int], num_steps: int, label: str = "Steady state", + min_run: int = 5, ) -> tuple[list[tuple[int, int]], int]: """Find contiguous regions where ``values`` are near the global peak. @@ -68,8 +69,8 @@ def _identify_regions_by_peak( if steady_state_started: prev_events_in_steady -= 1 - if prev_events_in_steady > 5 and not steady_state_started: - print(f"{label} started at index {i - 5}") + if prev_events_in_steady > min_run and not steady_state_started: + print(f"{label} started at index {i - min_run}") steady_state_started = True start_index = i - prev_events_in_steady + 1 @@ -86,7 +87,7 @@ def _identify_regions_by_peak( prev_events_in_steady = 0 if steady_state_started and not steady_state_ended: - regions.append((start_index, i)) + regions.append((start_index, i + 1)) print(f"{label} regions: {regions}") @@ -418,13 +419,21 @@ def _identify_regions_by_decode_baseline( decode_regions, global_max = _identify_regions_by_peak( decode_bs, num_steps, label="Steady state (decode baseline)", + min_run=1, ) - # Map decode-space indices back to full-iteration indices + # Map decode-space indices back to full-iteration indices and extend + # into adjacent prefill_bearing iterations so the region captures the + # full serving workload (the annotation path includes these naturally). + n = len(phase_labels) full_regions: list[tuple[int, int]] = [] for ds, de in decode_regions: full_start = decode_indices[ds] full_end = decode_indices[min(de - 1, len(decode_indices) - 1)] + 1 + while full_start > 0 and phase_labels[full_start - 1] == "prefill_bearing": + full_start -= 1 + while full_end < n and phase_labels[full_end] == "prefill_bearing": + full_end += 1 full_regions.append((full_start, full_end)) return full_regions, global_max diff --git a/TraceLens/TraceUtils/split_inference/trace_extraction.py b/TraceLens/TraceUtils/split_inference/trace_extraction.py index d060a5c0..2a390c04 100644 --- a/TraceLens/TraceUtils/split_inference/trace_extraction.py +++ b/TraceLens/TraceUtils/split_inference/trace_extraction.py @@ -28,14 +28,14 @@ from .detect_utils import ( GPU_KERNEL_CATEGORIES, - PROJECTION_CATEGORY, + GPU_USER_ANNOTATION, build_root_tiles, ) -# Kernels plus the annotation projections that describe them. Anything summing +# Kernels plus the GPU annotation spans that describe them. Anything summing # GPU *time* must use GPU_KERNEL_CATEGORIES instead, since a projection encloses # the kernels it describes and counting both double-counts. -GPU_EVENT_CATEGORIES = [*GPU_KERNEL_CATEGORIES, PROJECTION_CATEGORY] +GPU_EVENT_CATEGORIES = [*GPU_KERNEL_CATEGORIES, GPU_USER_ANNOTATION] def get_filename(filepath: str) -> dict: @@ -134,7 +134,7 @@ def infer_batch_sizes_from_shapes( e for e in cpu_events[lo:hi] if win_ts <= e["ts"] < win_end and e["dur"] <= win_dur ] - batch_sizes.append(most_common_first_dim(window_events)) + batch_sizes.append(most_common_first_dim(window_events, exclude_mem_ops=True)) return batch_sizes diff --git a/TraceLens/TraceUtils/split_inference_trace_annotation.py b/TraceLens/TraceUtils/split_inference_trace_annotation.py index 7af4ca6b..68f17b58 100644 --- a/TraceLens/TraceUtils/split_inference_trace_annotation.py +++ b/TraceLens/TraceUtils/split_inference_trace_annotation.py @@ -367,8 +367,8 @@ def main(): if detection.coverage: print( f"GPU coverage ({detection.coverage.strategy}): " - f"{detection.coverage.covered_any:.1%} by any annotation, " - f"{detection.coverage.covered_selected:.1%} by the selected roots" + f"{detection.coverage.covered_selected:.1%} by the selected roots, " + f"{detection.coverage.span_share:.1%} of that inside their spans" ) # Create output directory diff --git a/TraceLens/util.py b/TraceLens/util.py index 487e63b4..08cde286 100644 --- a/TraceLens/util.py +++ b/TraceLens/util.py @@ -1114,17 +1114,32 @@ def get_events(pftrace_data: dict) -> List[dict]: return pftrace_data.get("traceEvents", []) -def most_common_first_dim(events: list[dict]) -> int | None: +_MEMORY_VIEW_OPS = frozenset({ + "aten::select", "aten::slice", "aten::as_strided", "aten::narrow", + "aten::copy_", "aten::_to_copy", "aten::to", + "aten::index_put_", "aten::_index_put_impl_", + "aten::resize_", "aten::resolve_conj", "aten::resolve_neg", + "aten::expand", "aten::permute", "aten::transpose", "aten::contiguous", + "aten::view", "aten::reshape", "aten::unsqueeze", "aten::squeeze", + "aten::flatten", "aten::unflatten", +}) + + +def most_common_first_dim( + events: list[dict], exclude_mem_ops: bool = False, +) -> int | None: """Return the most common first dimension across all ``Input Dims`` of cpu_op events. - Scans every ``cpu_op`` event's ``Input Dims`` argument, collects the first - element of each dimension list, and returns the most frequent value. - Returns ``None`` when no cpu_op carries ``Input Dims``. + When *exclude_mem_ops* is True, skips memory/view ops whose tensor + dimensions reflect cache or layout sizes rather than the batch dimension. + Returns ``None`` when no eligible cpu_op carries ``Input Dims``. """ first_dims: list[int] = [] for e in events: if e.get("cat") != "cpu_op": continue + if exclude_mem_ops and e.get("name", "") in _MEMORY_VIEW_OPS: + continue input_dims = e.get("args", {}).get("Input Dims") if not input_dims: continue diff --git a/tests/test_split_inference_trace_annotation.py b/tests/test_split_inference_trace_annotation.py index cb1200e6..8c17623a 100644 --- a/tests/test_split_inference_trace_annotation.py +++ b/tests/test_split_inference_trace_annotation.py @@ -43,10 +43,15 @@ from tests.fixtures.traces import INFERENCE_ROOT from TraceLens.Trace2Tree.inference_iteration_roots import ( _entry_roots, + _find_repeating_period, _reattach_worker_threads, ) from TraceLens.Trace2Tree.trace_to_tree import TraceToTree +from TraceLens.util import normalize_name_for_comparison +from TraceLens.TraceUtils.split_inference.detect_utils import DetectStatus from TraceLens.TraceUtils.split_inference.root_detection import ( + _child_groups, + _periodic_candidate, _total_gpu_time, detect_from_branch_descent, ) @@ -247,18 +252,14 @@ def _details(num_requests, context_requests=0): def test_identify_steady_state_regions_clear_region(): iter_details = [_details(2) for _ in range(4)] + [_details(20) for _ in range(30)] - regions, global_max = _identify_regions_inference( - iter_details, num_steps=32 - ) + regions, global_max = _identify_regions_inference(iter_details, num_steps=32) assert global_max == 20 assert regions == [(4, 33)] def test_identify_steady_state_regions_fallback(): iter_details = [_details(20 if i % 2 == 0 else 2) for i in range(10)] - regions, global_max = _identify_regions_inference( - iter_details, num_steps=12 - ) + regions, global_max = _identify_regions_inference(iter_details, num_steps=12) assert global_max == 20 assert len(regions) == 1 assert regions == [(4, 6)] @@ -743,9 +744,7 @@ def test_find_steady_state_inference_invalid_mode(): trace = _make_trace(names) roots = split.find_iteration_roots(trace["traceEvents"]).roots with pytest.raises(ValueError, match="Unknown mode"): - split.find_steady_state_inference( - roots, num_steps=4, mode="invalid" - ) + split.find_steady_state_inference(roots, num_steps=4, mode="invalid") def test_find_steady_state_inference_conc_mismatch_warning(capsys): @@ -1248,3 +1247,319 @@ def test_capture_tree_cache(self, tmp_path): p.write_text(json.dumps(events)) _get_cached_capture_tree((f"k{i}", str(p)), str(p)) assert len(tcm._capture_tree_cache) <= tcm._CAPTURE_TREE_CACHE_MAX_SIZE + + +def _conditional_loop_events( + turns: int, worked: set, noisy: set = frozenset(), post: bool = False +) -> List[Dict]: + """An event loop where only the turns in ``worked`` run the GPU frame. + + Every turn also emits a kernel-free bookkeeping frame, so the child-name + sequence has no contiguous period -- the shape that made real sglang overlap + scheduling undetectable. + + ``noisy`` turns emit one extra kernel-free frame, which is enough to break + the stride even when every turn does work. ``post`` adds a second, much + cheaper GPU frame per working turn, so the loop body is two frames wide. + """ + events: List[Dict] = [ + { + "ph": "X", + "cat": "cpu_op", + "name": "event_loop", + "pid": 1, + "tid": 1, + "ts": 0, + "dur": turns * 2000, + "args": {"Sequence number": 0}, + } + ] + corr = 900 + for turn in range(turns): + base = 100 + turn * 2000 + idle = (("idle", 200),) if turn in noisy else () + for name, offset in (("get_next_batch", 0), ("bookkeeping", 100)) + idle: + events.append( + { + "ph": "X", + "cat": "cpu_op", + "name": name, + "pid": 1, + "tid": 1, + "ts": base + offset, + "dur": 50, + "args": {"Sequence number": turn}, + } + ) + if turn not in worked: + continue + frames = [("run_batch", 300, 400, 100)] + if post: + frames.append(("post_process", 800, 100, 5)) + for name, offset, dur, kernel_dur in frames: + events.append( + { + "ph": "X", + "cat": "cpu_op", + "name": name, + "pid": 1, + "tid": 1, + "ts": base + offset, + "dur": dur, + "args": {"Sequence number": turn, "correlation": corr}, + } + ) + events.extend( + [ + { + "ph": "X", + "cat": "cuda_runtime", + "name": "hipLaunchKernel", + "pid": 1, + "tid": 1, + "ts": base + offset + 10, + "dur": 5, + "args": {"correlation": corr}, + }, + { + "ph": "X", + "cat": "kernel", + "name": f"{name}_kernel", + "pid": 0, + "tid": 7, + "ts": base + offset + 50, + "dur": kernel_dur, + "args": {"correlation": corr, "stream": 7}, + }, + { + "ph": "s", + "id": corr, + "pid": 0, + "tid": 7, + "ts": base + offset + 50, + "cat": "ac2g", + "name": "ac2g", + }, + { + "ph": "f", + "id": corr, + "pid": 0, + "tid": 7, + "ts": base + offset + 50 + kernel_dur, + "cat": "ac2g", + "name": "ac2g", + "bp": "e", + }, + ] + ) + corr += 1 + return events + + +def _descend(events: List[Dict]): + tree = TraceToTree(events, prune_nongpu_paths=True) + tree.build_tree(add_python_func=True) + _reattach_worker_threads(tree) + return detect_from_branch_descent(tree, _entry_roots(tree), _total_gpu_time(tree)) + + +def test_branch_descent_finds_conditional_loop_body(): + """The working frame is found even though it runs on only some turns.""" + worked = {0, 1, 3, 4, 6, 8, 10, 11} + result = _descend(_conditional_loop_events(14, worked)) + + assert result is not None + assert result.diagnostics["branch_source"] == "frame_family" + assert len(result.roots) == len(worked) + assert {r["name"] for r in result.roots} == {"run_batch"} + assert result.status is DetectStatus.SPLITTABLE + + +def test_branch_descent_periodicity_ignores_kernel_free_frames(): + """Chatter is dropped before the period search, not only after it. + + ``_blocks_by_pattern`` already tolerates kernel-free intruders, but that + tolerance never runs when their irregular spacing is the very thing stopping + a period from being found. + """ + events = _conditional_loop_events( + 12, worked=set(range(12)), noisy={2, 5, 9}, post=True + ) + tree = TraceToTree(events, prune_nongpu_paths=True) + tree.build_tree(add_python_func=True) + _reattach_worker_threads(tree) + node = next(e for e in tree.events_by_uid.values() if e.get("name") == "event_loop") + ordered = sorted(tree.get_children_events(node), key=lambda e: e.get("ts", 0)) + _, gputime_by_name = _child_groups(tree, ordered) + + # The idle frames land on 3 of 12 turns, leaving no contiguous repeating run. + assert _find_repeating_period([e.get("name", "") for e in ordered])[0] is None + + candidate = _periodic_candidate( + tree, node, ordered, gputime_by_name, _total_gpu_time(tree), 0 + ) + assert candidate is not None + assert candidate.diagnostics["branch_coverage"] > 0.9 + + +def test_branch_descent_rejects_frames_with_no_gpu_work(): + """A frame explaining no GPU work is never an answer, however often it runs. + + The loop below is perfectly periodic, so the pattern path does produce a + candidate -- it just accounts for none of the trace's GPU time, which sits on + an unrelated frame. Reporting it was how 63 CPU-only roots got emitted. + """ + events = _conditional_loop_events(14, worked=set()) + events.extend( + [ + { + "ph": "X", + "cat": "cpu_op", + "name": "unrelated_work", + "pid": 1, + "tid": 1, + "ts": 90000, + "dur": 500, + "args": {"correlation": 42}, + }, + { + "ph": "X", + "cat": "kernel", + "name": "gemm_kernel", + "pid": 0, + "tid": 7, + "ts": 90100, + "dur": 200, + "args": {"correlation": 42, "stream": 7}, + }, + ] + ) + # Guard the guard: the detector must reach the coverage check, not bail out + # early on a trace with no GPU work at all. + assert sum(e["dur"] for e in events if e.get("cat") == "kernel") > 0 + assert _descend(events) is None + + +VLLM_PRIMARY = ( + "execute_{i}_context_3(sq128sk256sqsq1sqsk1)_generation_2(sq1sk300sqsq1sqsk1)" +) +SGLANG_DECODE = "step[DECODE bs={i}]" +SGLANG_EXTEND = "step[EXTEND bs=2 toks={t}]" +VLLM_BACKUP = "execute_context_3({i})_generation_2(50)" + + +def _warmup_loop_events(turns: int = 6) -> List[Dict]: + """A loop whose first turn does the iteration's work plus extra setup. + + The setup frame breaks the stride, so the period search anchors after the + first turn and leaves it outside every block. That is the shape bookend + promotion exists for: the turn is a real iteration, and dropping it reports + its GPU time as work no root explains. + """ + events: List[Dict] = [ + { + "ph": "X", + "cat": "cpu_op", + "name": "event_loop", + "pid": 1, + "tid": 1, + "ts": 0, + "dur": turns * 2000, + "args": {"Sequence number": 0}, + } + ] + corr = 500 + for turn in range(turns): + base = 100 + turn * 2000 + second = ("setup", 10) if turn == 0 else ("mlp", 50) + for name, offset, kernel_dur in ( + ("attention", 300, 100), + (second[0], 800, second[1]), + ): + events.append( + { + "ph": "X", + "cat": "cpu_op", + "name": name, + "pid": 1, + "tid": 1, + "ts": base + offset, + "dur": 200, + "args": {"Sequence number": turn, "correlation": corr}, + } + ) + events.extend( + [ + { + "ph": "X", + "cat": "cuda_runtime", + "name": "hipLaunchKernel", + "pid": 1, + "tid": 1, + "ts": base + offset + 10, + "dur": 5, + "args": {"correlation": corr}, + }, + { + "ph": "X", + "cat": "kernel", + "name": f"{name}_kernel", + "pid": 0, + "tid": 7, + "ts": base + offset + 50, + "dur": kernel_dur, + "args": {"correlation": corr, "stream": 7}, + }, + { + "ph": "s", + "id": corr, + "pid": 0, + "tid": 7, + "ts": base + offset + 50, + "cat": "ac2g", + "name": "ac2g", + }, + { + "ph": "f", + "id": corr, + "pid": 0, + "tid": 7, + "ts": base + offset + 50 + kernel_dur, + "cat": "ac2g", + "name": "ac2g", + "bp": "e", + }, + ] + ) + corr += 1 + return events + + +def test_bookend_promotion_adopts_a_warmup_turn_that_does_the_same_work(): + """The off-stride first turn is kept, because its GPU signature matches.""" + turns = 6 + events = _warmup_loop_events(turns) + tree = TraceToTree(events, prune_nongpu_paths=True) + tree.build_tree(add_python_func=True) + _reattach_worker_threads(tree) + node = next(e for e in tree.events_by_uid.values() if e.get("name") == "event_loop") + ordered = sorted(tree.get_children_events(node), key=lambda e: e.get("ts", 0)) + _, gputime_by_name = _child_groups(tree, ordered) + + # The setup frame means the repeating run starts only at the second turn. + live = [e for e in ordered if gputime_by_name.get(e.get("name", ""), 0.0) > 0] + _, _, start = _find_repeating_period( + [normalize_name_for_comparison(e.get("name", "")) for e in live] + ) + assert start > 0 + + candidate = _periodic_candidate( + tree, node, ordered, gputime_by_name, _total_gpu_time(tree), 0 + ) + assert candidate is not None + # Every turn is a root, warmup included, and nothing is left for the cascade + # to bolt on afterwards. + assert len(candidate.roots) == turns + assert candidate.roots[0]["ts"] < live[start]["ts"] + assert candidate.diagnostics["before_uids"] == [] + assert candidate.diagnostics["branch_coverage"] > 0.99 diff --git a/tests/test_split_root_detection.py b/tests/test_split_root_detection.py index 5ec306b9..6babe729 100644 --- a/tests/test_split_root_detection.py +++ b/tests/test_split_root_detection.py @@ -6,14 +6,14 @@ """Tests for the coverage-gated splitter: components and the flow built on them.""" +import pytest + from TraceLens.Trace2Tree.inference_iteration_roots import ( find_period_candidates, ) from TraceLens.TraceUtils.annotation_utils import ( - PROVENANCE_KEY, cluster_by_skeleton, dominant_cluster, - inherit_identity, is_parseable, name_skeleton, parse_annotation, @@ -43,7 +43,7 @@ # --------------------------------------------------------------------------- # # Event builders # --------------------------------------------------------------------------- # -def annotation(name, ts, dur, pid=1, tid=10): +def annotation(name, ts, dur, pid=1, tid=10, ext=None): return { "name": name, "cat": "user_annotation", @@ -52,7 +52,7 @@ def annotation(name, ts, dur, pid=1, tid=10): "dur": dur, "pid": pid, "tid": tid, - "args": {}, + "args": {} if ext is None else {"External id": ext}, } @@ -82,7 +82,9 @@ def kernel(ts, dur, corr, name="gemm", pid=1, tid=99): } -def projection(name, ts, dur, pid=1, tid=99): +def gpu_annotation_span(name, ts, dur, pid=1, tid=99, ext=None): + # Carries the External id its CPU annotation carries: that shared id is how + # a GPU span is known to describe one specific instance. return { "name": name, "cat": "gpu_user_annotation", @@ -91,21 +93,21 @@ def projection(name, ts, dur, pid=1, tid=99): "dur": dur, "pid": pid, "tid": tid, - "args": {}, + "args": {} if ext is None else {"External id": ext}, } -def serving_trace(count=16, name_template=VLLM, period=1000, with_projection=False): +def serving_trace(count=16, name_template=VLLM, period=1000, with_gpu_annotation=False): """One annotation per iteration, each launching one kernel.""" events, corr = [], 500 for i in range(count): base = 1000 + i * period name = name_template.format(i=i) - events.append(annotation(name, base, 100)) + events.append(annotation(name, base, 100, ext=i)) events.append(launch(base + 10, corr)) events.append(kernel(base + 200, 40, corr)) - if with_projection: - events.append(projection(name, base + 200, 40)) + if with_gpu_annotation: + events.append(gpu_annotation_span(name, base + 200, 40, ext=i)) corr += 1 return events @@ -169,27 +171,6 @@ def test_resolution_exists_on_every_annotation(self): """Stage 2 keys on resolution; a missing attribute would raise.""" assert parse_annotation("step[DECODE bs=4]").resolution is None - def test_inherit_keeps_window_and_takes_identity(self): - outer = annotation("scheduler.process_batch_result", 1000, 500) - inner = annotation("step[DECODE bs=7]", 1050, 100) - merged = inherit_identity(outer, inner) - - assert merged["ts"] == 1000 and merged["dur"] == 500 - assert merged["name"] == "step[DECODE bs=7]" - assert parse_annotation(merged["name"]).generation_requests == 7 - assert merged[PROVENANCE_KEY] == { - "window_from": "scheduler.process_batch_result", - "identity_from": "step[DECODE bs=7]", - } - assert outer["name"] == "scheduler.process_batch_result" - - def test_second_inheritance_keeps_the_original_window_owner(self): - outer = annotation("scheduler.process_batch_result", 1000, 500) - once = inherit_identity(outer, annotation("step[DECODE bs=1]", 1050, 100)) - twice = inherit_identity(once, annotation("step[EXTEND bs=2 toks=8]", 1060, 10)) - assert twice[PROVENANCE_KEY]["window_from"] == "scheduler.process_batch_result" - assert twice[PROVENANCE_KEY]["identity_from"] == "step[EXTEND bs=2 toks=8]" - # --------------------------------------------------------------------------- # # C2: periodicity @@ -270,25 +251,41 @@ def test_group_by_thread_sorts_within_group(self): # C4: GPU attribution # --------------------------------------------------------------------------- # class TestGpuAttribution: - def test_prefers_projections_when_present(self): - attribution = GpuAttribution(serving_trace(4, with_projection=True)) - assert attribution.strategy == GpuAttribution.STRATEGY_PROJECTION + def test_prefers_gpu_annotation_spans_when_present(self): + events = serving_trace(4, with_gpu_annotation=True) + attribution = GpuAttribution(events) + _, strategy = attribution.attributed_kernels(collect_annotations(events)) + assert strategy == GpuAttribution.STRATEGY_GPU_SPAN + + def test_falls_back_to_correlation_without_gpu_annotation_spans(self): + events = serving_trace(4) + attribution = GpuAttribution(events) + _, strategy = attribution.attributed_kernels(collect_annotations(events)) + assert strategy == GpuAttribution.STRATEGY_CORRELATION - def test_falls_back_to_correlation_without_projections(self): - attribution = GpuAttribution(serving_trace(4)) - assert attribution.strategy == GpuAttribution.STRATEGY_CORRELATION + def test_strategy_is_chosen_per_root_set_not_per_trace(self): + """One instance without a GPU counterpart must not measure by spans. - def test_projections_are_excluded_from_gpu_busy_time(self): - """Counting a projection as GPU time double-counts the kernels inside it.""" - events = serving_trace(4, with_projection=True) + A name join would credit this set with the annotated iterations' spans; + the External id join sees the gap and falls back to correlation. + """ + events = serving_trace(4, with_gpu_annotation=True) + annotations = collect_annotations(events) + unannotated = annotation("step[DECODE bs=9]", 90_000, 400) + attribution = GpuAttribution(events + [unannotated]) + _, strategy = attribution.attributed_kernels(annotations + [unannotated]) + assert strategy == GpuAttribution.STRATEGY_CORRELATION + + def test_gpu_annotation_spans_are_excluded_from_gpu_busy_time(self): + """Counting an annotation span as GPU time double-counts the kernels inside.""" + events = serving_trace(4, with_gpu_annotation=True) annotations = collect_annotations(events) - assert GpuAttribution(events).audit(annotations, annotations).gpu_busy == 4 * 40 + assert GpuAttribution(events).audit(annotations).gpu_busy == 4 * 40 def test_full_coverage_when_every_kernel_is_annotated(self): events = serving_trace(8) annotations = collect_annotations(events) - report = GpuAttribution(events).audit(annotations, annotations) - assert report.covered_any == 1.0 + report = GpuAttribution(events).audit(annotations) assert report.covered_selected == 1.0 assert report.passes @@ -312,17 +309,19 @@ def test_work_just_outside_a_root_counts_once_windows_extend(self): corr += 1 roots = collect_annotations(events) - report = GpuAttribution(events).audit(roots, roots) - assert report.covered_spans < 1.0 - assert report.covered_selected == 1.0 - assert report.span_share >= 0.9 + report = GpuAttribution(events).audit(roots) + # A tenth of GPU time is launched after the annotations close. + assert report.covered_spans == pytest.approx(0.9) + # The windows reclaim it, bar the final iteration's tail, which falls + # outside the last root and so outside the audited window. + assert report.covered_selected > report.covered_spans assert report.passes def test_sparse_roots_stretched_over_many_iterations_do_not_pass(self): """Coverage from window extension alone is not root coverage.""" events = serving_trace(40) annotations = collect_annotations(events) - report = GpuAttribution(events).audit(annotations, annotations[::20]) + report = GpuAttribution(events).audit(annotations[::20]) assert report.covered_selected > report.covered_spans assert report.span_share < 0.5 assert not report.passes @@ -335,33 +334,33 @@ def test_gate_measures_the_roots_not_every_annotation(self): """ events = serving_trace(40) annotations = collect_annotations(events) - report = GpuAttribution(events).audit(annotations, annotations[:2]) - assert report.covered_any == 1.0 - assert not report.passes + attribution = GpuAttribution(events) + assert attribution.audit(annotations).covered_selected == 1.0 + assert not attribution.audit(annotations[:2]).passes def test_unannotated_work_lowers_coverage(self): events = serving_trace(8) # A kernel with no launch site inside any annotation. events.append(kernel(1500, 4000, 99999, name="orphan")) attribution = GpuAttribution(events) - report = attribution.audit(collect_annotations(events), []) - assert report.covered_any < COVERAGE_GATE + report = attribution.audit(collect_annotations(events)) + assert report.covered_selected < COVERAGE_GATE def test_selected_roots_can_cover_less_than_all_annotations(self): """The signature of roots sitting at the wrong nesting level.""" events = serving_trace(8) annotations = collect_annotations(events) attribution = GpuAttribution(events) - report = attribution.audit(annotations, annotations[:2]) - assert report.covered_any == 1.0 - assert report.covered_selected < report.covered_any + assert ( + attribution.audit(annotations[:2]).covered_selected + < attribution.audit(annotations).covered_selected + ) def test_family_gpu_time(self): events = serving_trace(4) attribution = GpuAttribution(events) annotations = collect_annotations(events) - skeleton = name_skeleton(annotations[0]["name"]) - assert attribution.gpu_time_for_family(skeleton, annotations) == 4 * 40 + assert attribution.gpu_time_for_family(annotations) == 4 * 40 # --------------------------------------------------------------------------- # @@ -412,7 +411,7 @@ def test_healthy_trace_resolves_without_probes(self): assert result.phase_confidence is PhaseConfidence.HIGH assert result.method == "annotation:tier" assert len(result) == 16 - assert result.coverage.covered_any == 1.0 + assert result.coverage.covered_selected == 1.0 def test_partial_known_falls_through_to_unknown_family(self): """Only 3 of 20 iterations match a known pattern. Known annotations @@ -521,7 +520,7 @@ def test_uncovered_work_grades_directly_without_probes(self): events = serving_trace(16) events.append(kernel(1500, 500_000, 99999, name="unaccounted")) result = find_iteration_roots(events) - assert result.coverage.covered_any < COVERAGE_GATE + assert result.coverage.covered_selected < COVERAGE_GATE assert result.status in (DetectStatus.DEGRADED, DetectStatus.NOT_SPLITTABLE) def test_manifest_reports_quality(self): @@ -529,7 +528,7 @@ def test_manifest_reports_quality(self): assert manifest["status"] == 0 assert manifest["phase_confidence"] == "high" assert manifest["n_roots"] == 16 - assert manifest["coverage_any_annotation"] == 1.0 + assert manifest["coverage_selected_roots"] == 1.0 assert manifest["attribution_strategy"] == "correlation" diff --git a/tests/test_trace_split.py b/tests/test_trace_split.py index 81426670..e700980c 100644 --- a/tests/test_trace_split.py +++ b/tests/test_trace_split.py @@ -355,8 +355,8 @@ def test_trace_split_no_annotations(dirpath, trace_gz, tmp_path): # Run all modes in a single invocation ss_ref = os.path.join(dirpath, "steady_state_traces") dp_ref = os.path.join(dirpath, "phase_split_traces") - run_ss = os.path.isdir(ss_ref) and not is_llm - run_dp = os.path.isdir(dp_ref) and not is_llm + run_ss = os.path.isdir(ss_ref) + run_dp = os.path.isdir(dp_ref) out = str(tmp_path / "output") os.makedirs(out, exist_ok=True) @@ -390,26 +390,67 @@ def test_trace_split_no_annotations(dirpath, trace_gz, tmp_path): f"tolerance={tolerance}" ) - # --- find-steady-state + divide-phases: compare kernel counts --- + # --- find-steady-state: compare window types and kernel counts --- if run_ss: - gen_kernels = _kernel_events(_collect_events(out)) - ref_kernels = _kernel_events( - _strip_annotations(_collect_events(ss_ref)) - ) - assert len(gen_kernels) >= len(ref_kernels), ( - f"find-steady-state: stripped has fewer kernels ({len(gen_kernels)}) " - f"than annotated ({len(ref_kernels)})" + def _ss_window_type(filename): + return filename.split("_steady_state")[0] + + gen_ss_files = sorted(f for f in _list_gz(out) if "steady_state" in f) + ref_ss_files = _list_gz(ss_ref) + gen_types = sorted(_ss_window_type(f) for f in gen_ss_files) + ref_types = sorted(_ss_window_type(f) for f in ref_ss_files) + assert gen_types == ref_types, ( + f"find-steady-state: window type mismatch — " + f"generated {gen_types}, reference {ref_types}" ) + for ref_file in ref_ss_files: + wtype = _ss_window_type(ref_file) + gen_file = next(f for f in gen_ss_files if _ss_window_type(f) == wtype) + gen_kernels = _kernel_events( + DataLoader.load_data(os.path.join(out, gen_file))["traceEvents"] + ) + ref_kernels = _kernel_events( + _strip_annotations( + DataLoader.load_data( + os.path.join(ss_ref, ref_file) + )["traceEvents"] + ) + ) + ss_tolerance = max(1, int(len(ref_kernels) * 0.05)) + assert len(gen_kernels) >= len(ref_kernels) - ss_tolerance, ( + f"find-steady-state '{wtype}': stripped has fewer kernels " + f"({len(gen_kernels)}) than annotated ({len(ref_kernels)}), " + f"tolerance={ss_tolerance}" + ) if run_dp: - gen_kernels = _kernel_events(_collect_events(out, recursive=True)) - ref_kernels = _kernel_events( - _strip_annotations(_collect_events(dp_ref, recursive=True)) - ) - assert len(gen_kernels) >= len(ref_kernels), ( - f"divide-phases: stripped has fewer kernels ({len(gen_kernels)}) " - f"than annotated ({len(ref_kernels)})" + ref_phases = sorted( + d for d in os.listdir(dp_ref) + if os.path.isdir(os.path.join(dp_ref, d)) ) + assert ref_phases, f"No phase subdirectories in {dp_ref}" + for phase_dir in ref_phases: + gen_phase_path = os.path.join(out, phase_dir) + ref_phase_path = os.path.join(dp_ref, phase_dir) + assert os.path.isdir(gen_phase_path), ( + f"divide-phases: missing phase directory '{phase_dir}' in output" + ) + gen_files = _list_gz(gen_phase_path) + ref_files = _list_gz(ref_phase_path) + assert len(gen_files) == len(ref_files), ( + f"divide-phases '{phase_dir}': chunk count mismatch — " + f"generated {len(gen_files)}, reference {len(ref_files)}" + ) + gen_kernels = _kernel_events(_collect_events(gen_phase_path)) + ref_kernels = _kernel_events( + _strip_annotations(_collect_events(ref_phase_path)) + ) + dp_tolerance = max(1, int(len(ref_kernels) * 0.05)) + assert len(gen_kernels) >= len(ref_kernels) - dp_tolerance, ( + f"divide-phases '{phase_dir}': stripped has fewer kernels " + f"({len(gen_kernels)}) than annotated ({len(ref_kernels)}), " + f"tolerance={dp_tolerance}" + ) # --------------------------------------------------------------------------- diff --git a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_1.json.gz b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_1.json.gz index d513d66f..ea3aa0dc 100644 Binary files a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_1.json.gz and b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_1.json.gz differ diff --git a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_3.json.gz b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_3.json.gz index 8e4bcac9..b767b05e 100644 Binary files a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_3.json.gz and b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_3.json.gz differ diff --git a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_4.json.gz b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_4.json.gz index e0890b84..6d931ff9 100644 Binary files a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_4.json.gz and b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_iteration_4.json.gz differ diff --git a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_warmup.json.gz b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_warmup.json.gz index 5fb55cf3..ec3df635 100644 Binary files a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_warmup.json.gz and b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_warmup.json.gz differ diff --git a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_wrapup.json.gz b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_wrapup.json.gz index 77f0383b..01ec5944 100644 Binary files a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_wrapup.json.gz and b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/split_traces/profile_trace_rank_0_wrapup.json.gz differ diff --git a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/steady_state_traces/steady_state_profile_trace_rank_0.json.gz b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/steady_state_traces/steady_state_profile_trace_rank_0.json.gz index 98babd27..85de7ba9 100644 Binary files a/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/steady_state_traces/steady_state_profile_trace_rank_0.json.gz and b/tests/traces/trace_splitter_traces/xdit_hunyuanvideo/steady_state_traces/steady_state_profile_trace_rank_0.json.gz differ