Introduce pipeline configuration files

profiles/default.json

@@ -0,0 +1 @@
+{}

profiles/everything.json

@@ -0,0 +1,55 @@
+{
+    "stages": [
+        {"drop_timestamp_reversed_events": true},
+        {"create_slice_from_BE": true},
+        {"normalize_phase1": true},
+        {"frequency_align_collect": true},
+        {"pipeline_barrier": true},
+        {"frequency_align_apply": true},
+        {"normalize_phase2": true},
+        {"event_sanity_checks": true},
+        {"time_align_collect" : true},
+        {"pipeline_barrier": true},
+        {"time_align_apply": true},
+        {"remove_ids_from_name": true},
+        {"map_tid_to_range": true},
+        {"cycle_count_to_wallclock": true},
+        {"tighten_hts_by_instr_type": true},
+        {"tripple_phased_events": true},
+        {"mp_sync_tight_v1": true},
+        {"mp_ts_calibration_v2": true},
+        {"queueing_counter": true},
+        {"drop_global_events": true},
+        {"recombine_cpu_events": true},
+        {"sort_events": true},
+        {"assert_ts_sequence": true},
+        {"detect_partial_overlap_events": true},
+        {"assert_ts_sequence": true},
+        {"collect_iteration_stats": true},
+        {"extract_power_event": true},
+        {"sort_events": true},
+        {"compute_power": true},
+        {"extract_data_transfer_event": true},
+        {"compute_bandwidth": true},
+        {"compute_utilization_fingerprints": true},
+        {"assert_ts_sequence": true},
+        {"communication_event_collection": true},
+        {"pipeline_barrier": true},
+        {"communication_event_apply": true},
+        {"compute_utilization": true},
+        {"sort_events": true},
+        {"assert_global_ts_sequence": true},
+        {"flow_prepare_event_data": true},
+        {"flow_extraction": true},
+        {"mp_calc_bw_v2": true},
+        {"mp_calc_bw": true},
+        {"calculate_stats": true},
+        {"processing_filter": true},
+        {"flow_data_cleanup": true},
+        {"cycle_count_conversion_cleanup": true},
+        {"cleanup_copy_of_device_ts": true},
+        {"tb_refinement_intrusive": true},
+        {"tb_refinement_lightweight": true},
+        {"calculate_stats_v2": true}
+    ]
+}

src/aiu_trace_analyzer/core/engine.py

@@ -4,6 +4,7 @@
 import aiu_trace_analyzer.core.processing as processing
 import aiu_trace_analyzer.export.exporter as output
 
+
 class Engine:
     # main engine that drives the processing
     # high-level:
@@ -13,9 +14,9 @@ class Engine:
     #   * when iteration is finished: flush any potential exporter buffers
 
     def __init__(self,
-                  importer: ingest.AbstractTraceIngest,
-                  processor: processing.EventProcessor,
-                  exporter: output.AbstractTraceExporter) -> None:
+                 importer: ingest.AbstractTraceIngest,
+                 processor: processing.EventProcessor,
+                 exporter: output.AbstractTraceExporter) -> None:
         self.importer = importer
         self.processor = processor
         self.exporter = exporter
@@ -24,14 +25,14 @@ def run(self) -> int:
         # pull from ingest various sources as iterator
         for next_item in self.importer:
             # convert/process
-            events = self.processor.process( next_item )
+            events = self.processor.process(next_item)
             # push to export
-            self.exporter.export( events )
+            self.exporter.export(events)
 
         # drain the context buffers (if any)
         drain = self.processor.drain()
         # export any events emitted during drain
-        self.exporter.export( drain )
+        self.exporter.export(drain)
         # flush the export buffers
         self.exporter.flush()
         return 0

src/aiu_trace_analyzer/core/processing.py

@@ -8,6 +8,7 @@
 from aiu_trace_analyzer.types import TraceEvent
 from aiu_trace_analyzer.core.duplicate_hold import IntermediateDuplicateAndHoldContext, duplicate_and_hold
 from aiu_trace_analyzer.export.exporter import JsonFileTraceExporter
+from aiu_trace_analyzer.core.stage_profile import StageProfile, StageProfileChecker
 
 _MINREQKEYS = ["ph", "ts", "pid", "name"]
 
@@ -20,14 +21,15 @@ class EventProcessor:
     High-level 3 stage:
       1. pass event(s) through registered pre-processing functions
       2. convert from python dict no AbstractEventType object
-      3. pass those through registered post-processing functions (none exist yet)
     '''
-    def __init__(self, intermediate: str = None) -> None:
+    def __init__(self, profile: StageProfile = None, intermediate: str = None) -> None:
         self.stages = []
         self.stages.append((EventProcessor.sanity_check, None, {}))
+        self.profile = profile
         self.event_count = 0
         self.intermediate = intermediate
         self.stage_count = 0
+        self.stage_check = StageProfileChecker(self.profile)
 
     def __del__(self) -> None:
         aiulog.log(aiulog.INFO, "Exported events: ", self.event_count)
@@ -39,11 +41,19 @@ def __del__(self) -> None:
        * a dictionary for k/v config arguments
     '''
     def register_stage(self, callback, context: procCTX.AbstractContext = None, **kwargs):
+        if not self.stage_check.fwd_find_stage(callback.__name__):
+            aiulog.log(aiulog.DEBUG, "DAH: Skipping registration of", callback.__name__, ": disabled in profile.")
+            return
+        else:
+            aiulog.log(aiulog.DEBUG, "DAH: registering: ", callback.__name__)
+
         self.stages.append((callback, context, kwargs))
 
         # if intermediate results are requested, register an additional special function+context
         if self.intermediate:
-            next_intermediate = IntermediateDuplicateAndHoldContext(JsonFileTraceExporter(target_uri=f'{self.intermediate}_{callback.__name__}_{self.stage_count}'))
+            next_intermediate = IntermediateDuplicateAndHoldContext(
+                JsonFileTraceExporter(target_uri=f'{self.intermediate}_{callback.__name__}_{self.stage_count}')
+            )
             aiulog.log(aiulog.TRACE, "DAH: registering preprocessing stage export:", next_intermediate.exporter.target_uri)
             self.stages.append((duplicate_and_hold, next_intermediate, None))
             self.stage_count += 1
@@ -54,12 +64,12 @@ def register_stage(self, callback, context: procCTX.AbstractContext = None, **kw
     Returns empty list if input is invalid (which drops the event and ends processing of this event)
     '''
     @staticmethod
-    def sanity_check( event: TraceEvent, _: procCTX.AbstractContext) -> list[TraceEvent]:
+    def sanity_check(event: TraceEvent, _: procCTX.AbstractContext) -> list[TraceEvent]:
         for check in _MINREQKEYS:
             if check not in event:
                 aiulog.log(aiulog.ERROR, "Event failed sanityCheck: ", check, "is not in", event)
-                return [] # TODO: should be exception
-        return [ event ]
+                return []    # TODO: should be exception
+        return [event]
 
     def process(self, event: TraceEvent) -> list[aiuev.AbstractEventType]:
         # turn into a list, pre/post have do be able to expand single events into lists
@@ -76,7 +86,7 @@ def process(self, event: TraceEvent) -> list[aiuev.AbstractEventType]:
     # walk through the registered pre-processing hooks for the event
     # split any returned list of events into single events for each next stage pre-processor
     def pre_process(self, event: TraceEvent) -> list[TraceEvent]:
-        event_list = [ event ]
+        event_list = [event]
         for pre_process, context, keyword_dictionary in self.stages:
             next_event_list = []
             for event in event_list:
@@ -102,16 +112,15 @@ def convert_events(self, event_list: list[TraceEvent]) -> list[aiuev.AbstractEve
             if "args" not in event:
                 event["args"] = {}
             # any key that's not listed is to be moved into args to be preserved
-            for key,val in event.items():
-                if key not in ["ph","ts","pid","tid","name","cat","args","id","bp","dur"]:
+            for key, val in event.items():
+                if key not in ["ph", "ts", "pid", "tid", "name", "cat", "args", "id", "bp", "dur"]:
                     event["args"][key] = val
 
             new_event = aiuev.AbstractEventType.from_dict(event)
 
             output_event_list.append(new_event)
         return output_event_list
 
-
     def drain(self) -> list[aiuev.AbstractEventType]:
         # walk through the registered pre-processing hooks for the event
         # split any returned list of events into single events for each next stage pre-processor

src/aiu_trace_analyzer/core/stage_profile.py

@@ -0,0 +1,48 @@
+# Copyright 2024-2025 IBM Corporation
+
+import json
+from pathlib import Path
+
+
+class StageProfile:
+    _everything_profile = 'profiles/everything.json'
+
+    def __init__(self, profile_data: dict):
+        self.profile = self._ingest_profile_data(profile_data)
+
+    @classmethod
+    def from_json(cls, file: Path):
+        with open(file, 'r') as config_fd:
+            profile_data = json.load(config_fd)
+
+        # if a profile is empty, then assume all stages to be enabled
+        if len(profile_data) == 0:
+            with open(StageProfile._everything_profile, 'r') as config_fd:
+                profile_data = json.load(config_fd)
+
+        profile = StageProfile(profile_data)
+        return profile
+
+    def _ingest_profile_data(self, profile_data: dict) -> list[str]:
+        if 'stages' not in profile_data:
+            raise KeyError("Profile data is missing 'stages' key.")
+
+        profile = []
+        for stage_data in profile_data['stages']:
+            stage, enabled = stage_data.popitem()
+            if enabled:
+                profile.append(stage)
+        return profile
+
+
+class StageProfileChecker:
+    def __init__(self, profile: StageProfile):
+        self.stages = profile
+        self.reg_idx = 0
+
+    def fwd_find_stage(self, stage: str) -> bool:
+        for incr, st in enumerate(self.stages.profile[self.reg_idx:]):
+            if stage == st:
+                self.reg_idx += incr
+                return True
+        return False

src/aiu_trace_analyzer/pipeline/__init__.py

@@ -9,6 +9,7 @@
   2)  implement your processing function and corresponding context (derive from existing contexts or AbstractContext)
   3)  add import lines below for your new context and function(s)
   4)  integrate to acelyzer.py as needed (watch the order of registered functions, because they're almost never independent)
+  5)  add the new callback function names to any personalities/profiles in matching order
 
   Simple example to get started could be TIDMappingContext and map_tid_to_range in tid_mapping.py
 '''
@@ -49,23 +50,21 @@
 # for reference of the template, you'd do here:
 #       from aiu_trace_analyzer.pipeline.template import MyStructsAndFunctionsForCrossEventContext
 
-
-
 # import the separated processing functions
-from aiu_trace_analyzer.pipeline.mappings import *
+from aiu_trace_analyzer.pipeline.mappings import map_complete_to_duration, remove_ids_from_name
 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_events, assert_ts_sequence, assert_global_ts_sequence, recombine_cpu_events
 from aiu_trace_analyzer.pipeline.inverse_ts import drop_timestamp_reversed_events
 from aiu_trace_analyzer.pipeline.sort import sort_events
 from aiu_trace_analyzer.pipeline.make_slice import create_slice_from_BE
-from aiu_trace_analyzer.pipeline.power import extract_power_event,check_power_ts_sequence,compute_power
+from aiu_trace_analyzer.pipeline.power import extract_power_event, check_power_ts_sequence, compute_power
 from aiu_trace_analyzer.pipeline.filter import processing_filter
 from aiu_trace_analyzer.pipeline.tid_mapping import map_tid_to_range
 from aiu_trace_analyzer.pipeline.tripple_event import tripple_phased_events
 from aiu_trace_analyzer.pipeline.timesync import cycle_count_to_wallclock, cycle_count_conversion_cleanup, realign_dts_to_hts, \
                                           tighten_hts_by_instr_type, get_opIds_from_event, cleanup_copy_of_device_ts
-from aiu_trace_analyzer.pipeline.dma import extract_data_transfer_event,compute_bandwidth
+from aiu_trace_analyzer.pipeline.dma import extract_data_transfer_event, compute_bandwidth
 from aiu_trace_analyzer.pipeline.stats import calculate_stats
 from aiu_trace_analyzer.pipeline.stats_v2 import calculate_stats_v2
 from aiu_trace_analyzer.pipeline.mp_calc_bw import mp_calc_bw

tests/aiu_trace_analyzer/core/test_stage_profile.py

@@ -0,0 +1,43 @@
+# Copyright 2024-2025 IBM Corporation
+
+import pytest
+
+from aiu_trace_analyzer.core.stage_profile import StageProfile, StageProfileChecker
+
+
+@pytest.fixture
+def default_profile_config() -> str:
+    return "profiles/default.json"
+
+
+@pytest.fixture
+def everything_profile() -> str:
+    return "profiles/everything.json"
+
+
+@pytest.fixture
+def default_profile(default_profile_config) -> StageProfile:
+    return StageProfile.from_json(default_profile_config)
+
+
+@pytest.fixture
+def default_stage_checker(default_profile) -> StageProfileChecker:
+    return StageProfileChecker(default_profile)
+
+
+def test_default_is_everything(default_profile_config, everything_profile):
+    def_stage_profile = StageProfile.from_json(default_profile_config)
+    all_stage_profile = StageProfile.from_json(everything_profile)
+
+    for d, a in zip(def_stage_profile.profile, all_stage_profile.profile):
+        assert d == a
+
+
+def test_fwd_find(default_stage_checker):
+    test_stages = [('create_slice_from_BE', True, 1), ('do_not_move_index', False, 1), ('pipeline_barrier', True, 4)]
+
+    assert isinstance(default_stage_checker, StageProfileChecker)
+
+    for (stage, found, idx) in test_stages:
+        assert default_stage_checker.fwd_find_stage(stage) == found
+        assert default_stage_checker.reg_idx == idx