Feat/devalshah/generic splitter

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

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

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,32 +250,29 @@ 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:
             ts, dur, cat = e.get("ts"), e.get("dur"), e.get("cat")
             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.