Skip to content

Feat/split refactor generic - #968

Draft
devalshahamd wants to merge 5 commits into
mainfrom
feat/split_refactor_generic
Draft

Feat/split refactor generic#968
devalshahamd wants to merge 5 commits into
mainfrom
feat/split_refactor_generic

Conversation

@devalshahamd

Copy link
Copy Markdown
Contributor

This pull request introduces significant improvements and refactoring to the trace iteration root detection and annotation utilities, as well as a new, consolidated split_inference module 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

  • The generic iteration root detection in Trace2Tree/inference_iteration_roots.py has 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]
  • The import linter configuration .importlinter is updated to reflect the new split_inference submodule structure, ensuring correct dependency management.

2. Annotation Utilities and Family Keying

  • The annotation utilities in TraceUtils/annotation_utils.py now 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.
  • The annotation parsing is now cached and exposed via parse_annotation, and the iteration_details aggregation uses the cached parser for efficiency and consistency.
  • The IterationAnnotation class now always has a resolution attribute, improving consistency across instances.

3. New Split Inference Module

  • A new TraceUtils/split_inference/__init__.py module 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:

  • Iteration root detection refactor and diagnostics: [1] [2]
  • Import linter update for new submodule structure:
  • Annotation utilities: skeletonization, clustering, and caching: [1] [2] [3]
  • New split_inference module:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants