From aad8f9a45e4a0c2f2bf3343fa3b0075d0bac6a0b Mon Sep 17 00:00:00 2001 From: Lars Schneidenbach Date: Wed, 26 Nov 2025 16:26:59 -0500 Subject: [PATCH] TID-based overlap detection with dynamic range Signed-off-by: Lars Schneidenbach --- src/aiu_trace_analyzer/core/acelyzer.py | 6 +- src/aiu_trace_analyzer/pipeline/__init__.py | 1 + src/aiu_trace_analyzer/pipeline/overlap.py | 98 ++++++++++++++++--- .../profiles/everything.json | 2 + .../profiles/torch_minimal.json | 2 + 5 files changed, 92 insertions(+), 17 deletions(-) diff --git a/src/aiu_trace_analyzer/core/acelyzer.py b/src/aiu_trace_analyzer/core/acelyzer.py index eed6555..ff49058 100644 --- a/src/aiu_trace_analyzer/core/acelyzer.py +++ b/src/aiu_trace_analyzer/core/acelyzer.py @@ -485,8 +485,12 @@ def register_processing_functions(self, process.register_stage(callback=event_pipe.assert_ts_sequence, context=monotonic_ts_ctx_a) # register pre-processing: resolve overlap conflicts caused by partially overlapping slices - overlap_ctx = event_pipe.OverlapDetectionContext(overlap_resolve=self._overlap_option_from_arg(args.overlap), + overlap_arg = self._overlap_option_from_arg(args.overlap) + overlap_ctx = event_pipe.OverlapDetectionContext(overlap_resolve=overlap_arg, ts_shift_threshold=self.defaults["ts_shift_threshold"]) + if overlap_arg == event_pipe.OverlapDetectionContext.OVERLAP_RESOLVE_TID: + process.register_stage(callback=event_pipe.detect_partial_overlap_tids, context=overlap_ctx) + process.register_stage(callback=event_pipe.pipeline_barrier, context=event_pipe._main_barrier_context) process.register_stage(callback=event_pipe.detect_partial_overlap_events, context=overlap_ctx) # validate that the overlap has not messed up the event stream ordering diff --git a/src/aiu_trace_analyzer/pipeline/__init__.py b/src/aiu_trace_analyzer/pipeline/__init__.py index 46002a1..09d9bb4 100644 --- a/src/aiu_trace_analyzer/pipeline/__init__.py +++ b/src/aiu_trace_analyzer/pipeline/__init__.py @@ -59,6 +59,7 @@ from aiu_trace_analyzer.pipeline.normalize import normalize_phase1, normalize_phase2 from aiu_trace_analyzer.pipeline.correctness import event_sanity_checks from aiu_trace_analyzer.pipeline.overlap import ( + detect_partial_overlap_tids, detect_partial_overlap_events, assert_ts_sequence, assert_global_ts_sequence, diff --git a/src/aiu_trace_analyzer/pipeline/overlap.py b/src/aiu_trace_analyzer/pipeline/overlap.py index 1f537a9..ba52335 100644 --- a/src/aiu_trace_analyzer/pipeline/overlap.py +++ b/src/aiu_trace_analyzer/pipeline/overlap.py @@ -3,16 +3,22 @@ import copy import aiu_trace_analyzer.logger as aiulog -from aiu_trace_analyzer.pipeline import AbstractContext, EventPairDetectionContext +from aiu_trace_analyzer.pipeline import AbstractContext, AbstractHashQueueContext, TwoPhaseWithBarrierContext from aiu_trace_analyzer.types import TraceEvent, GlobalIngestData from aiu_trace_analyzer.pipeline.tools import PipelineContextTool class OverlapTracking(tuple[float, bool, list[float]]): + """ + tuple of: + float: the ts of the currently 'active' event + bool: whether this event stream/queue is blocked or not (has active event) + list[float]: list of end-ts for the stack of active events + """ pass -class OverlapDetectionContext(EventPairDetectionContext): +class OverlapDetectionContext(TwoPhaseWithBarrierContext): ''' Management structures and functions to deal with overlapping events. Solves without storing the events themselves, just keeps track of @@ -35,7 +41,8 @@ class OverlapDetectionContext(EventPairDetectionContext): def __init__(self, overlap_resolve=OVERLAP_RESOLVE_DROP, - ts_shift_threshold=0.0 + ts_shift_threshold=0.0, + max_tid_streams=5, ) -> None: super().__init__() self.overlap_resolve = overlap_resolve @@ -43,6 +50,8 @@ def __init__(self, self.async_id = 0 self.async_queues = {} self.ts_shift_threshold = ts_shift_threshold + self.tid_space = {} + self.max_tid_streams = max_tid_streams def __del__(self) -> None: level = aiulog.WARN if self.resolved else aiulog.INFO @@ -54,7 +63,7 @@ def __del__(self) -> None: def overlap_detection(self, event: TraceEvent) -> list[TraceEvent]: tid = event["tid"] if "tid" in event else 0 - queue_id = self.queue_hash(event["pid"], tid) + queue_id = self.event_data_hash(event, ["pid", "tid"], ignore_missing=True) if queue_id not in self.queues: self.queues[queue_id] = (0.0, False, []) @@ -113,6 +122,14 @@ def get_overlap_time(self, ts: float, end: float, qstate: OverlapTracking) -> fl overlap_time = max(overlap_time, e - ts) return overlap_time + def find_next_tid(self, event: TraceEvent) -> int: + if event["tid"] not in self.tid_space[event["pid"]]: + aiulog.log( + aiulog.ERROR, + "POD: insufficient dynamic range for tid-based overlap resolution. Increase max_tid_space.") + new_tid = self.tid_space[event["pid"]][event["tid"]] + return new_tid + # solve a detected overlap between a pair of pairs def handle_overlap(self, oevent: TraceEvent, @@ -151,7 +168,7 @@ def handle_overlap(self, self.resolved += 1 return rlist elif self.overlap_resolve == self.OVERLAP_RESOLVE_TID: - oevent["tid"] += 1 + oevent["tid"] = self.find_next_tid(oevent) # feed offending event back into the detector with the new TID to make sure # there are no collisions there either rlist = self.overlap_detection(oevent) @@ -194,16 +211,65 @@ def update_async_event_queue(self, queueID, async_event, current) -> list[TraceE self.async_queues[queueID] = remain return rlist + def collect_tid_space(self, event: TraceEvent) -> None: + pid, tid = event["pid"], event["tid"] + if pid not in self.tid_space: + self.tid_space[pid] = {-1: set()} + + # collect all detected tids under key `-1` + self.tid_space[pid][-1].add(tid) + + if tid not in self.tid_space[pid]: + self.tid_space[pid][tid] = [] + + def _create_tid_space(self, tid: int, exclude: list[int]) -> list[int]: + tlist = [] + next_tid = tid + while len(tlist) < self.max_tid_streams: + next_tid += 1 + if next_tid not in exclude: + tlist.append(next_tid) + return tlist + def drain(self): - revents = [] - # make sure to drain the queue of async 'e' events that might have been hold - # back past the end of the last main event of a stream - while len(self.async_queues) > 0: - _, aq = self.async_queues.popitem() - # make sure to keep everything sorted - aq.sort(key=lambda e: e['ts']) - revents += aq - return revents + if self.overlap_resolve == self.OVERLAP_RESOLVE_TID and self.phase == self._COLLECTION_PHASE: + new_tspace = {} + for pid, tspace in self.tid_space.items(): + new_tspace = {} + # collect candidate lists for each known tid from input + exclude: set = tspace[-1] + for tid in tspace.keys(): + if tid == -1: + continue + tcandidates = self._create_tid_space(tid, exclude) + self.tid_space[pid][tid] = tcandidates + exclude.update(tcandidates) + new_tspace[tid] = tcandidates[0] + for src_tid, next_tid in zip(tcandidates[:-1], tcandidates[1:]): + new_tspace[src_tid] = next_tid + + aiulog.log(aiulog.TRACE, "POD: total tid_space:", self.tid_space[pid]) + self.tid_space[pid] = copy.deepcopy(new_tspace) + aiulog.log(aiulog.TRACE, "POD: tid neighbors:", new_tspace) + return super().drain() + else: + revents = [] + # make sure to drain the queue of async 'e' events that might have been hold + # back past the end of the last main event of a stream + while len(self.async_queues) > 0: + _, aq = self.async_queues.popitem() + # make sure to keep everything sorted + aq.sort(key=lambda e: e['ts']) + revents += aq + return revents + + +def detect_partial_overlap_tids(event: TraceEvent, context: AbstractContext) -> list[TraceEvent]: + assert isinstance(context, OverlapDetectionContext) + + if event["ph"] in "X": + context.collect_tid_space(event) + return [event] # mapping function callback @@ -218,7 +284,7 @@ def detect_partial_overlap_events(event: TraceEvent, context: AbstractContext) - ################################################################### # Timestamps sequence checking to make sure time stamps stay sorted -class TSSequenceContext(EventPairDetectionContext): +class TSSequenceContext(AbstractHashQueueContext): def __init__(self, ts3check: bool = False): super().__init__() self.TS_cmpt_end = {} @@ -236,7 +302,7 @@ def __del__(self): def insert(self, event: TraceEvent, queue_id=None): if not queue_id: - queue_id = self.queue_hash(event["pid"], event["tid"]) + queue_id = self.event_data_hash(event, ["pid", "tid"], ignore_missing=True) if queue_id not in self.queues: self.queues[queue_id] = (0.0, 1e99) diff --git a/src/aiu_trace_analyzer/profiles/everything.json b/src/aiu_trace_analyzer/profiles/everything.json index 2de74c5..3a3f709 100644 --- a/src/aiu_trace_analyzer/profiles/everything.json +++ b/src/aiu_trace_analyzer/profiles/everything.json @@ -23,6 +23,8 @@ {"recombine_cpu_events": true}, {"sort_events": true}, {"assert_ts_sequence": true}, + {"detect_partial_overlap_tids": true}, + {"pipeline_barrier": true}, {"detect_partial_overlap_events": true}, {"assert_ts_sequence": true}, {"collect_iteration_stats": true}, diff --git a/src/aiu_trace_analyzer/profiles/torch_minimal.json b/src/aiu_trace_analyzer/profiles/torch_minimal.json index cb63361..1f49f5f 100644 --- a/src/aiu_trace_analyzer/profiles/torch_minimal.json +++ b/src/aiu_trace_analyzer/profiles/torch_minimal.json @@ -23,6 +23,8 @@ {"recombine_cpu_events": false}, {"sort_events": true}, {"assert_ts_sequence": false}, + {"detect_partial_overlap_tids": true}, + {"pipeline_barrier": true}, {"detect_partial_overlap_events": true}, {"assert_ts_sequence": true}, {"collect_iteration_stats": true},