diff --git a/.flake8 b/.flake8 index e69cea0..9484d49 100644 --- a/.flake8 +++ b/.flake8 @@ -5,5 +5,6 @@ max-complexity = 17 per-file-ignores = src/aiu_trace_analyzer/pipeline/__init__.py:F401 + src/aiu_trace_analyzer/verification/__init__.py:F401 src/aiu_trace_analyzer/__init__.py:F401,F403 src/aiu_trace_analyzer/core/acelyzer.py:C901 diff --git a/src/aiu_trace_analyzer/__init__.py b/src/aiu_trace_analyzer/__init__.py index db8629b..9a8e04e 100644 --- a/src/aiu_trace_analyzer/__init__.py +++ b/src/aiu_trace_analyzer/__init__.py @@ -12,3 +12,4 @@ from aiu_trace_analyzer.export import * from aiu_trace_analyzer.core import * from aiu_trace_analyzer.pipeline import * +from aiu_trace_analyzer.verification import * diff --git a/src/aiu_trace_analyzer/core/acelyzer.py b/src/aiu_trace_analyzer/core/acelyzer.py index c754dc1..317fad1 100644 --- a/src/aiu_trace_analyzer/core/acelyzer.py +++ b/src/aiu_trace_analyzer/core/acelyzer.py @@ -15,6 +15,7 @@ import aiu_trace_analyzer.export.exporter as output import aiu_trace_analyzer.logger as aiulog import aiu_trace_analyzer.pipeline as event_pipe +import aiu_trace_analyzer.verification as verify_pipe from aiu_trace_analyzer import __version__ @@ -762,12 +763,22 @@ def register_verification_functions(self, process: processor.EventProcessor, args, exporter: output.AbstractTraceExporter): - verification_ctx = event_pipe.VerificationContext() - kernel_parent_ctx = event_pipe.KernelParentVerificationContext() + # for a verification stage example, see verify.py - process.register_stage(callback=event_pipe.verify, context=verification_ctx) - process.register_stage(callback=event_pipe.kernel_parent_collect, context=kernel_parent_ctx) + # strict: accelerator events have no call/return relationship, so even a fully embedded + # event is a data problem and not a legitimate nesting + overlap_verification_ctx = verify_pipe.OverlapVerificationContext(strict=True) + kernel_parent_ctx = verify_pipe.KernelParentVerificationContext() + + # register anything that verifies the raw/unsorted order of the input before this sort-stage + process.register_stage(callback=event_pipe.sort_events, context=event_pipe.EventSortingContext( + event_types=None, sortkey=self._default_sort_ts_and_rev_dur, global_sort=True)) + + # anything that requires a sorted event stream below this point + process.register_stage(callback=verify_pipe.verify_kernel_overlap, context=overlap_verification_ctx) + process.register_stage(callback=verify_pipe.kernel_parent_collect, context=kernel_parent_ctx) process.register_stage(callback=event_pipe.pipeline_barrier, context=event_pipe._main_barrier_context) - process.register_stage(callback=event_pipe.kernel_parent_verify, context=kernel_parent_ctx) - process.register_stage(callback=event_pipe.verify_cleanup) - process.register_stage(callback=event_pipe.verification_result_filter) + process.register_stage(callback=verify_pipe.kernel_parent_verify, context=kernel_parent_ctx) + + # drop all regular events and only keep report-data after this stage + process.register_stage(callback=verify_pipe.verification_result_filter) diff --git a/src/aiu_trace_analyzer/pipeline/__init__.py b/src/aiu_trace_analyzer/pipeline/__init__.py index 0cacd60..ebd0830 100644 --- a/src/aiu_trace_analyzer/pipeline/__init__.py +++ b/src/aiu_trace_analyzer/pipeline/__init__.py @@ -111,11 +111,4 @@ # for reference of the template, you'd do here: # from aiu_trace_analyzer.pipeline.template import myprocessing -from aiu_trace_analyzer.verification.verify import verify, verify_cleanup -from aiu_trace_analyzer.verification.verify import VerificationContext -from aiu_trace_analyzer.verification.kernel_parent_verify import ( - KernelParentVerificationContext, - kernel_parent_collect, - kernel_parent_verify -) -from aiu_trace_analyzer.verification.report import verification_result_filter +# NOTE: the verification stages live in aiu_trace_analyzer.verification diff --git a/src/aiu_trace_analyzer/pipeline/overlap.py b/src/aiu_trace_analyzer/pipeline/overlap.py index d92ff91..0f212e0 100644 --- a/src/aiu_trace_analyzer/pipeline/overlap.py +++ b/src/aiu_trace_analyzer/pipeline/overlap.py @@ -3,7 +3,10 @@ import copy import aiu_trace_analyzer.logger as aiulog -from aiu_trace_analyzer.pipeline import AbstractContext, AbstractHashQueueContext, TwoPhaseWithBarrierContext +from aiu_trace_analyzer.pipeline import ( + AbstractContext, + AbstractHashQueueContext, + TwoPhaseWithBarrierContext) from aiu_trace_analyzer.types import TraceEvent, GlobalIngestData, TraceWarning from aiu_trace_analyzer.pipeline.tools import PipelineContextTool @@ -32,6 +35,10 @@ class OverlapDetectionContext(TwoPhaseWithBarrierContext): is the time stamp indicating until what time it's active * need to keep a list of active end ts because conflicts might happen within non-critical nested overlapping events + * strict mode: any event that starts while another event of the same stream is still + running is an overlap. The default (non-strict) mode accepts fully embedded events + as a legitimate nesting (e.g. a parent/child function-call relationship) and only + flags partial overlaps. ''' OVERLAP_RESOLVE_DROP = 1 OVERLAP_RESOLVE_TID = 2 @@ -39,10 +46,16 @@ class OverlapDetectionContext(TwoPhaseWithBarrierContext): OVERLAP_RESOLVE_WARN = 4 OVERLAP_RESOLVE_SHIFT = 5 + # default event fields that define the identity of a stream/queue. Derived classes with a static + # partitioning can override this by assignment (never in-place: the list is shared by all + # instances). Hierarchical keys are supported, e.g. "args.stream". + _QUEUE_ID_KEYS = ["pid", "tid"] + def __init__(self, overlap_resolve=OVERLAP_RESOLVE_DROP, ts_shift_threshold=0.0, max_tid_streams=5, + strict=False, ) -> None: super().__init__(warnings=[ TraceWarning( @@ -57,6 +70,14 @@ def __init__(self, self.ts_shift_threshold = ts_shift_threshold self.tid_space = {} self.max_tid_streams = max_tid_streams + self.strict = strict + self.queue_id_keys = None # resolved from the first event, see _select_queue_id_keys() + + # select the event fields that define a stream/queue identity. Called only for the first event, + # so the selection may depend on the input dialect (assumed constant for one program execution). + # Derived classes can override, e.g. `return super()._select_queue_id_keys(event) + ["args.stream"]` + def _select_queue_id_keys(self, _event: TraceEvent) -> list[str]: + return self._QUEUE_ID_KEYS # search for events within the same pid/tid # accumulate a queue of events for each pid/tid @@ -64,7 +85,10 @@ def __init__(self, def overlap_detection(self, event: TraceEvent) -> list[TraceEvent]: tid = event["tid"] if "tid" in event else 0 - queue_id = self.event_data_hash(event, ["pid", "tid"], ignore_missing=True) + if self.queue_id_keys is None: + self.queue_id_keys = self._select_queue_id_keys(event) + aiulog.log(aiulog.DEBUG, "POD: stream identity keys:", self.queue_id_keys) + queue_id = self.event_data_hash(event, self.queue_id_keys, ignore_missing=True) if queue_id not in self.queues: self.queues[queue_id] = (0.0, False, []) @@ -78,13 +102,23 @@ def overlap_detection(self, event: TraceEvent) -> list[TraceEvent]: event_ts = event["ts"] event_end = round(event["ts"] + event["dur"], 4) + # retire the events that have ended before this one starts, so that the blocked status + # refers to the ts of this event (every event leaves its own end-ts behind, so a status + # that's only updated after the check would keep any queue blocked forever) + # in SHIFT-mode, recursion happens within the same queue_id -> the update has to happen later + if self.overlap_resolve != self.OVERLAP_RESOLVE_SHIFT: + self.update_queue_status(event_ts, queue_id) + _, blocked, end_ts = self.queues[queue_id] + aiulog.log(aiulog.TRACE, "POD queue before: ", queue_id, "from", event["pid"], tid, self.queues[queue_id]) assert (blocked and len(end_ts) > 0) or (not blocked and len(end_ts) == 0) if not blocked: self.queues[queue_id] = (event_ts, True, [event_end]) revents = [event] else: - if self.check_overlap_condition(event_ts, event_end, self.queues[queue_id]): + # a blocked queue is all it takes in strict mode: embedded events are overlaps too, + # so there's no additional condition to check + if self.strict or self.check_overlap_condition(event_ts, event_end, self.queues[queue_id]): # actual overlap revents = self.handle_overlap(event, queue_id) else: @@ -95,7 +129,8 @@ def overlap_detection(self, event: TraceEvent) -> list[TraceEvent]: if self.overlap_resolve == self.OVERLAP_RESOLVE_ASYNC: aevents = self.update_async_event_queue(queue_id, None, event_ts) revents = aevents + revents # prepend any async events that need to be injected - self.update_queue_status(event_ts, queue_id) + elif self.overlap_resolve == self.OVERLAP_RESOLVE_SHIFT: + self.update_queue_status(event_ts, queue_id) # shift-mode requires update _after_ resolution aiulog.log(aiulog.TRACE, "POD queue after: ", queue_id, "from", event["pid"], tid, self.queues[queue_id]) return revents @@ -110,10 +145,10 @@ def check_overlap_condition(self, ts, end, qstate: OverlapTracking) -> bool: aiulog.log(aiulog.TRACE, "POD overlap detected", qstate, ts, end) return overlap - # remove only keep entries of end timestamps that are later than the new current head + # remove only entries of end timestamps that are strictly later than the new current head def update_queue_status(self, new_current: float, queue_id: int): end_q = self.queues[queue_id][2] - new_end_q = list(filter(lambda x: x >= new_current, end_q)) + new_end_q = list(filter(lambda x: x > new_current, end_q)) is_blocked = (len(new_end_q) > 0) # unblock the queue if no more end-ts are remaining self.queues[queue_id] = (new_current, is_blocked, new_end_q) @@ -138,17 +173,22 @@ def find_next_tid(self, event: TraceEvent) -> int: new_tid = self.tid_space[event["pid"]][event["tid"]] return new_tid + # single funnel for recording a detected overlap: derived classes can override this to + # attach per-event detail to the warning (the offending event is not available in issue_warning()) + def _record_overlap(self, _oevent: TraceEvent) -> None: + self.issue_warning("overlaps") + # solve a detected overlap between a pair of pairs def handle_overlap(self, oevent: TraceEvent, queue_id: int) -> list[TraceEvent]: if self.overlap_resolve == self.OVERLAP_RESOLVE_DROP: aiulog.log(aiulog.WARN, "Solving overlap conflict by dropping:", oevent) - self.issue_warning("overlaps") + self._record_overlap(oevent) return [] elif self.overlap_resolve == self.OVERLAP_RESOLVE_WARN: aiulog.log(aiulog.WARN, "Detected overlap conflict: ", oevent["name"]) - self.issue_warning("overlaps") + self._record_overlap(oevent) return [oevent] elif self.overlap_resolve == self.OVERLAP_RESOLVE_SHIFT: ts_shift = self.get_overlap_time(oevent["ts"], oevent["ts"]+oevent["dur"], self.queues[queue_id]) @@ -173,21 +213,22 @@ def handle_overlap(self, "us: increase threshold or use different overlap res option.") rlist = [oevent] - self.issue_warning("overlaps") + self._record_overlap(oevent) return rlist elif self.overlap_resolve == self.OVERLAP_RESOLVE_TID: 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) - self.issue_warning("overlaps") + self._record_overlap(oevent) return rlist elif self.overlap_resolve == self.OVERLAP_RESOLVE_ASYNC: + # record before the event is converted to an async b/e pair (which drops its 'dur') + self._record_overlap(oevent) oevent["id"] = self.async_id end_ts = oevent["ts"] + oevent["dur"] oevent.pop("dur") self.async_id += 1 - self.issue_warning("overlaps") e_event = copy.deepcopy(oevent) oevent["ph"] = "b" diff --git a/src/aiu_trace_analyzer/profiles/verification.json b/src/aiu_trace_analyzer/profiles/verification.json index 063d20c..411bd38 100644 --- a/src/aiu_trace_analyzer/profiles/verification.json +++ b/src/aiu_trace_analyzer/profiles/verification.json @@ -1,9 +1,10 @@ { "stages": [ + {"sort_events": true}, + {"verify_kernel_overlap": true}, {"kernel_parent_collect": true}, {"pipeline_barrier": true}, {"kernel_parent_verify": true}, - {"verify_cleanup": true}, {"verification_result_filter": true} ] } diff --git a/src/aiu_trace_analyzer/verification/__init__.py b/src/aiu_trace_analyzer/verification/__init__.py new file mode 100644 index 0000000..1554a10 --- /dev/null +++ b/src/aiu_trace_analyzer/verification/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024-2026 IBM Corporation + +''' +Verification pipeline stages and contexts. + +The stages of the verification pipeline are collected here, separate from the regular processing +pipeline of aiu_trace_analyzer.pipeline. The dependency between the two is one-directional: +verification stages build on the pipeline infrastructure (contexts, detection algorithms), the +regular pipeline never imports anything from here. + +A verification stage consists of a context derived from AbstractVerificationContext (which emits +the accumulated findings as verification meta-data events at drain time) plus one or more callback +functions. For a minimal example of both, see verify.py. +''' + +# import the verification contexts: +from aiu_trace_analyzer.verification.verify import VerificationContext +from aiu_trace_analyzer.verification.kernel_parent_verify import KernelParentVerificationContext +from aiu_trace_analyzer.verification.overlap_verify import OverlapVerificationContext + +# import the verification stage callbacks: +from aiu_trace_analyzer.verification.verify import verify, verify_cleanup +from aiu_trace_analyzer.verification.kernel_parent_verify import ( + kernel_parent_collect, + kernel_parent_verify +) +from aiu_trace_analyzer.verification.overlap_verify import verify_kernel_overlap +from aiu_trace_analyzer.verification.report import verification_result_filter diff --git a/src/aiu_trace_analyzer/verification/kernel_parent_verify.py b/src/aiu_trace_analyzer/verification/kernel_parent_verify.py index 32d4846..8f9c768 100644 --- a/src/aiu_trace_analyzer/verification/kernel_parent_verify.py +++ b/src/aiu_trace_analyzer/verification/kernel_parent_verify.py @@ -44,8 +44,8 @@ import aiu_trace_analyzer.logger as aiulog from aiu_trace_analyzer.types import TraceEvent, TraceWarning -from aiu_trace_analyzer.pipeline import TwoPhaseWithBarrierContext, AbstractContext -from aiu_trace_analyzer.pipeline.context import AbstractVerificationContext +from aiu_trace_analyzer.pipeline.barrier import TwoPhaseWithBarrierContext +from aiu_trace_analyzer.pipeline.context import AbstractContext, AbstractVerificationContext class KernelParentVerificationContext(AbstractVerificationContext, TwoPhaseWithBarrierContext): diff --git a/src/aiu_trace_analyzer/verification/overlap_verify.py b/src/aiu_trace_analyzer/verification/overlap_verify.py new file mode 100644 index 0000000..4d781ff --- /dev/null +++ b/src/aiu_trace_analyzer/verification/overlap_verify.py @@ -0,0 +1,101 @@ +# Copyright 2024-2026 IBM Corporation + +""" +Overlap verification for accelerator compute events. + +Compute events of the same stream cannot overlap in time: the stream processes them one after +another. Any detected overlap therefore indicates a data problem (e.g. inaccurate timestamps) +and is reported as an error-level finding of the verification report. + +This module reuses the overlap detection of the regular processing pipeline +(OverlapDetectionContext) and only replaces the parts that differ in verification mode: + * overlaps are reported, never resolved + * the finding is error-level and records each offending event as an instance + * the streams are identified per input dialect instead of by pid+tid +""" + +from aiu_trace_analyzer.types import TraceEvent, TraceWarning +from aiu_trace_analyzer.pipeline.context import AbstractContext, AbstractVerificationContext +from aiu_trace_analyzer.pipeline.overlap import OverlapDetectionContext +from aiu_trace_analyzer.pipeline.tools import PipelineContextTool + + +class OverlapVerificationContext(OverlapDetectionContext, AbstractVerificationContext): + ''' + Verification-mode variant of the overlap detection: it only reports overlapping compute + events, it never resolves them (always OVERLAP_RESOLVE_WARN). Every detected overlap is + recorded as an instance of an error-level finding and the accumulated findings are emitted + as verification meta-data events by AbstractVerificationContext.drain() (reached via the MRO, + the parent drain chains up to it; no two-phase barrier required for this context). + + Compute streams are identified per dialect: FLEX events have a single stream per pid, while + TORCH events separate the streams within a pid by their 'args.stream' entry. + ''' + test_name = "Compute Overlap Check" + + _OVERLAP_WARNING = "overlaps" + _STREAM_KEY = "stream" + _STREAM_ARG = "args." + _STREAM_KEY + _DEFAULT_STREAM = 0 # convention: actual stream numbers start at 1, 0 means 'no stream entry' + + def __init__(self, strict=False) -> None: + super().__init__(self.OVERLAP_RESOLVE_WARN, strict=strict) + # replace the resolution-oriented warning of the parent: in verification mode nothing is + # resolved, a detected overlap is a finding that has to fail the test + self.add_warning( + TraceWarning( + name=self._OVERLAP_WARNING, + text="Overlapping accelerator events detected: {d[count]}", + data={"count": 0}, + is_error=True, + ) + ) + + def _select_queue_id_keys(self, event: TraceEvent) -> list[str]: + dialect = PipelineContextTool.get_dialect_of_event(event) + assert dialect is not None, \ + "OVL: cannot determine the dialect of the first event." \ + " Register this stage before any stage that removes the jobhash." + if dialect.get("NAME") == "TORCH": + return ["pid", self._STREAM_ARG] + return ["pid"] + + def _record_overlap(self, oevent: TraceEvent) -> None: + super()._record_overlap(oevent) + self.warnings[self._OVERLAP_WARNING].add_instance({ + "name": oevent["name"], + "pid": oevent["pid"], + "tid": oevent["tid"], + "stream": oevent["args"].get(self._STREAM_KEY), + "ts": oevent["ts"], + "dur": oevent["dur"], + }) + + def add_default_stream(self, event: TraceEvent) -> bool: + ''' + the stream-based separation of queues requires the stream entry to exist. Events without + one are all attributed to the same default stream. Returns whether a default was added. + ''' + if self._STREAM_KEY in event["args"]: + return False + event["args"][self._STREAM_KEY] = self._DEFAULT_STREAM + return True + + def remove_default_stream(self, event: TraceEvent) -> None: + del event["args"][self._STREAM_KEY] + + +def verify_kernel_overlap(event: TraceEvent, context: AbstractContext) -> list[TraceEvent]: + assert isinstance(context, OverlapVerificationContext) + + # only compute events are checked; anything else just passes through + if event["ph"] not in "X" or not PipelineContextTool.is_acc_event(event): + return [event] + + # the default is only required to determine the queue/stream of this event: drop it again to + # keep the event unchanged for any downstream stage + default_added = context.add_default_stream(event) + revents = context.overlap_detection(event) + if default_added: + context.remove_default_stream(event) + return revents diff --git a/tests/aiu_trace_analyzer/pipeline/test_overlap.py b/tests/aiu_trace_analyzer/pipeline/test_overlap.py index f5cda27..024fed1 100644 --- a/tests/aiu_trace_analyzer/pipeline/test_overlap.py +++ b/tests/aiu_trace_analyzer/pipeline/test_overlap.py @@ -2,7 +2,7 @@ import pytest -from aiu_trace_analyzer.pipeline.overlap import recombine_cpu_events +from aiu_trace_analyzer.pipeline.overlap import OverlapDetectionContext, recombine_cpu_events from aiu_trace_analyzer.types import TraceEvent @@ -52,3 +52,61 @@ def test_get_cycles(flex_event_with_jobhash: TraceEvent, expected: TraceEvent): # actual result check assert modified[0] == expected + + +########################################################### +# overlap detection: strict vs. non-strict mode + +def _x_event(ts, dur, pid=1, tid=1, name="kernel") -> TraceEvent: + return TraceEvent({"ph": "X", "pid": pid, "tid": tid, "ts": ts, "dur": dur, "name": name, "args": {}}) + + +def _detect(events: list[TraceEvent], strict: bool) -> int: + '''run the events through a warn-only detection and return the number of detected overlaps''' + context = OverlapDetectionContext( + overlap_resolve=OverlapDetectionContext.OVERLAP_RESOLVE_WARN, strict=strict) + for event in events: + context.overlap_detection(event) + return context.warnings["overlaps"].args_list["count"] + + +# (description, events, expected non-strict count, expected strict count) +list_of_overlap_mode_tests = [ + ("no overlap at all", + [_x_event(0.0, 10.0), _x_event(20.0, 10.0)], 0, 0), + ("back-to-back: previous event ends when the next one starts", + [_x_event(0.0, 10.0), _x_event(10.0, 10.0)], 0, 0), + ("partial overlap is flagged in both modes", + [_x_event(0.0, 10.0), _x_event(5.0, 10.0)], 1, 1), + ("fully embedded event is a legitimate nesting only in non-strict mode", + [_x_event(0.0, 100.0), _x_event(10.0, 10.0)], 0, 1), + ("embedded event sharing the start ts", + [_x_event(0.0, 100.0), _x_event(0.0, 10.0)], 0, 1), + ("several independent events on one stream stay unflagged", + [_x_event(0.0, 5.0), _x_event(10.0, 5.0), _x_event(20.0, 5.0), _x_event(30.0, 5.0)], 0, 0), + ("separate streams never overlap each other", + [_x_event(0.0, 10.0, tid=1), _x_event(5.0, 10.0, tid=2)], 0, 0), +] + + +@pytest.mark.parametrize( + "description,events,expected_default,expected_strict", + list_of_overlap_mode_tests) +def test_overlap_modes(description, events, expected_default, expected_strict): + assert _detect(events, strict=False) == expected_default, f"non-strict: {description}" + assert _detect(events, strict=True) == expected_strict, f"strict: {description}" + + +def test_strict_does_not_flag_events_after_the_queue_ran_empty(): + # regression guard: the blocked status of a queue has to be re-evaluated for the ts of the + # incoming event. Every event leaves its own end-ts behind, so a status that is only updated + # after the overlap check would keep the queue blocked and flag every subsequent event. + events = [_x_event(float(i) * 10.0, 5.0) for i in range(10)] + assert _detect(events, strict=True) == 0 + + +def test_strict_flags_every_event_of_an_embedded_chain(): + # one long event with three events nested inside it: all three are overlaps in strict mode + events = [_x_event(0.0, 100.0), _x_event(10.0, 5.0), _x_event(20.0, 5.0), _x_event(30.0, 5.0)] + assert _detect(events, strict=True) == 3 + assert _detect(events, strict=False) == 0