Feat/split refactor generic - #968
Draft
devalshahamd wants to merge 5 commits into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces significant improvements and refactoring to the trace iteration root detection and annotation utilities, as well as a new, consolidated
split_inferencemodule for inference trace splitting. The main changes enhance the accuracy, extensibility, and maintainability of iteration detection and annotation parsing, and provide a unified interface for trace splitting functionality.Summary of most important changes:
1. Iteration Root Detection Refactor and Enhancement
Trace2Tree/inference_iteration_roots.pyhas been rewritten for greater accuracy and extensibility. The new logic introduces dataclasses and utilities for robust period detection, supports multiple label tiers (e.g., Python frames, CPU ops), and exposes detailed diagnostics. The detection now uses a more principled approach to candidate period selection, block extraction, and candidate ranking. [1] [2].importlinteris updated to reflect the newsplit_inferencesubmodule structure, ensuring correct dependency management.2. Annotation Utilities and Family Keying
TraceUtils/annotation_utils.pynow provide functions for canonicalizing event names into skeletons (to group related operations regardless of numeric suffixes), clustering by skeleton, and parsing annotations with memoization for efficiency. This supports more robust event family detection and grouping.parse_annotation, and theiteration_detailsaggregation uses the cached parser for efficiency and consistency.IterationAnnotationclass now always has aresolutionattribute, improving consistency across instances.3. New Split Inference Module
TraceUtils/split_inference/__init__.pymodule is introduced, consolidating and re-exporting all major trace splitting, window selection, and extraction utilities under a single namespace for easier and more organized usage.References: