Refactor from_frame_stack into a rule-table architecture (no behavior change)

mosaic/libmosaic/utils/data_utils.py

@@ -9,9 +9,10 @@
 
 import enum
 import logging
+import re
 from collections import defaultdict
 from dataclasses import dataclass, field, fields as dataclass_fields
-from typing import Any, List, Optional, Union
+from typing import Any, Callable, List, Optional, Union
 
 NUM_GPUS_PER_HOST = 8
 
@@ -42,6 +43,61 @@ class Frame:
     line: int
 
 
+# ---------------------------------------------------------------------------
+# Frame-stack predicate helpers
+#
+# These small helpers form the building blocks for the categorization rule
+# table. Each helper scans the whole frame stack and returns whether the
+# stack contains a frame matching the supplied criterion. They make the rule
+# table declarative and easy to extend.
+# ---------------------------------------------------------------------------
+
+
+def _any_filename_contains(stack: List[Frame], substring: str) -> bool:
+    """Return True if any frame's filename contains the given substring."""
+    return any(substring in frame.filename for frame in stack)
+
+
+def _any_filename_matches_regex(
+    stack: List[Frame], compiled_regex: "re.Pattern[str]"
+) -> bool:
+    """Return True if any frame's filename matches the given compiled regex."""
+    return any(compiled_regex.search(frame.filename) for frame in stack)
+
+
+def _any_frame_name_in(stack: List[Frame], names: frozenset) -> bool:
+    """Return True if any frame's name is in the given set."""
+    return any(frame.name in names for frame in stack)
+
+
+def _any_frame_name_contains(stack: List[Frame], substring: str) -> bool:
+    """Return True if any frame's name contains the given substring."""
+    return any(substring in frame.name for frame in stack)
+
+
+def _any_frame_name_startswith(stack: List[Frame], prefix: str) -> bool:
+    """Return True if any frame's name starts with the given prefix."""
+    return any(frame.name.startswith(prefix) for frame in stack)
+
+
+def _no_frame_name_in(stack: List[Frame], names: frozenset) -> bool:
+    """Return True if NO frame's name is in the given set."""
+    return not _any_frame_name_in(stack, names)
+
+
+# Frozensets used by the rule table. Defined at module scope so the
+# (predicate, category) lambdas can close over them without re-allocating
+# per call.
+_BACKWARD_GRAD_HELPER_NAMES: frozenset = frozenset(
+    {"clip_grad_norm_", "calc_grad_norm"}
+)
+_OPTIMIZER_VANILLA_NAMES: frozenset = frozenset({"custom_adamw", "_init_group"})
+_NET_AGGREGATE_NAMES: frozenset = frozenset(
+    {"dist_max", "dist_mean", "dist_sum", "_aggregate"}
+)
+_STATS_NAMES: frozenset = frozenset({"compute_grad_stats"})
+
+
 class AllocationType(enum.Enum):
     PARAMETER = 0
     ACTIVATION = 1
@@ -56,36 +112,14 @@ class AllocationType(enum.Enum):
 
     @classmethod
     def from_frame_stack(cls, frame_stack: List[Frame]) -> "AllocationType":
-        has_fsdp = any(
-            "fully_sharded_data_parallel.py" in frame.filename for frame in frame_stack
-        )
-        if has_fsdp:
-            return cls.FSDP
-
-        for frame in frame_stack:
-            if "forward" in frame.name:
-                return cls.ACTIVATION
-            if "param" in frame.name:
-                return cls.PARAMETER
-            if "flatten_params_wrapper" in frame.filename:
-                return cls.PARAMETER
-            if "build_norm_fn" in frame.name:
-                return cls.PARAMETER
-            if "__init__" in frame.name and "parallel/layers.py" in frame.filename:
-                return cls.PARAMETER
-            if "backward" in frame.name:
-                return cls.BACKWARD
-            if frame.name in ["clip_grad_norm_", "calc_grad_norm"]:
-                return cls.BACKWARD
-            if frame.name in ["custom_adamw", "_init_group"]:
-                return cls.OPTIMIZER
-            if frame.name in ["dist_max", "dist_mean", "dist_sum", "_aggregate"]:
-                return cls.NET
-            if "validate_process_group" in frame.name:
-                return cls.NET_SETUP
-            if frame.name == "compute_grad_stats":
-                return cls.STATS
+        """Classify a frame stack into an AllocationType.
 
+        Iterates the module-level rule table; the first matching rule wins.
+        Falls back to UNKNOWN if no rule matches.
+        """
+        for predicate, category in _RULE_TABLE:
+            if predicate(frame_stack):
+                return category
         return cls.UNKNOWN
 
     @classmethod
@@ -104,8 +138,6 @@ def from_frame_stack_with_custom(
     @classmethod
     def _matches_custom_pattern(cls, frame_stack: List[Frame], pattern: str) -> bool:
         """Check if any frame matches the custom pattern using regex"""
-        import re
-
         try:
             compiled_pattern = re.compile(pattern)
             for frame in frame_stack:
@@ -120,6 +152,85 @@ def _matches_custom_pattern(cls, frame_stack: List[Frame], pattern: str) -> bool
         return False
 
 
+# ---------------------------------------------------------------------------
+# Categorization rule table
+#
+# Each entry is (predicate, category). `from_frame_stack` walks the table
+# top-to-bottom and returns the category of the first rule whose predicate
+# returns True. The table is defined at module scope (after AllocationType)
+# so the predicates can close over the enum values directly.
+#
+# Rules are migrated 1-to-1 from the original if-cascade in their original
+# priority order: this preserves behavior for every existing test stack.
+# ---------------------------------------------------------------------------
+_RULE_TABLE: List[tuple[Callable[[List[Frame]], bool], AllocationType]] = [
+    # FSDP — any frame in FSDP source dominates everything else.
+    (
+        lambda s: _any_filename_contains(s, "fully_sharded_data_parallel.py"),
+        AllocationType.FSDP,
+    ),
+    # ACTIVATION — anything mentioning "forward" in its function name.
+    (
+        lambda s: _any_frame_name_contains(s, "forward"),
+        AllocationType.ACTIVATION,
+    ),
+    # PARAMETER — anything mentioning "param" in its function name.
+    (
+        lambda s: _any_frame_name_contains(s, "param"),
+        AllocationType.PARAMETER,
+    ),
+    # PARAMETER — flat-params wrapper.
+    (
+        lambda s: _any_filename_contains(s, "flatten_params_wrapper"),
+        AllocationType.PARAMETER,
+    ),
+    # PARAMETER — norm-fn builder.
+    (
+        lambda s: _any_frame_name_contains(s, "build_norm_fn"),
+        AllocationType.PARAMETER,
+    ),
+    # PARAMETER — module __init__ inside parallel/layers.py.
+    # This is a per-frame compound condition: a single frame must satisfy
+    # both name and filename predicates.
+    (
+        lambda s: any(
+            "__init__" in f.name and "parallel/layers.py" in f.filename for f in s
+        ),
+        AllocationType.PARAMETER,
+    ),
+    # BACKWARD — anything mentioning "backward" in its function name.
+    (
+        lambda s: _any_frame_name_contains(s, "backward"),
+        AllocationType.BACKWARD,
+    ),
+    # BACKWARD — known gradient-norm helpers.
+    (
+        lambda s: _any_frame_name_in(s, _BACKWARD_GRAD_HELPER_NAMES),
+        AllocationType.BACKWARD,
+    ),
+    # OPTIMIZER — vanilla optimizer entry points.
+    (
+        lambda s: _any_frame_name_in(s, _OPTIMIZER_VANILLA_NAMES),
+        AllocationType.OPTIMIZER,
+    ),
+    # NET — distributed aggregation primitives.
+    (
+        lambda s: _any_frame_name_in(s, _NET_AGGREGATE_NAMES),
+        AllocationType.NET,
+    ),
+    # NET_SETUP — process-group validation.
+    (
+        lambda s: _any_frame_name_contains(s, "validate_process_group"),
+        AllocationType.NET_SETUP,
+    ),
+    # STATS — gradient-stat computation.
+    (
+        lambda s: _any_frame_name_in(s, _STATS_NAMES),
+        AllocationType.STATS,
+    ),
+]
+
+
 @dataclass
 class TraceEvent:
     action: str

test/test_custom_profiling.py

@@ -10,6 +10,12 @@
 from mosaic.cmd.entry_point import get_memory_profile
 from mosaic.cmd.get_memory_profile import main
 from mosaic.libmosaic.utils.data_utils import (
+    _any_filename_contains,
+    _any_filename_matches_regex,
+    _any_frame_name_contains,
+    _any_frame_name_in,
+    _any_frame_name_startswith,
+    _no_frame_name_in,
     AllocationType,
     Frame,
     MemoryUsage,
@@ -351,6 +357,67 @@ def test_special_regex_characters_in_patterns(self) -> None:
         self.assertEqual(category, "parens")
 
 
+class TestFrameStackPredicateHelpers(TestCase):
+    """Unit tests for the small frame-stack predicate helpers used by the
+    AllocationType rule table."""
+
+    def setUp(self) -> None:
+        super().setUp()
+        # A small synthetic stack used by every assertion below.
+        self.stack: list[Frame] = [
+            Frame(name="forward", filename="model.py", line=10),
+            Frame(name="custom_adamw", filename="optim.py", line=20),
+            Frame(name="my_kernel_v2", filename="kernel.cu", line=30),
+        ]
+
+    def test_any_filename_contains(self) -> None:
+        self.assertTrue(_any_filename_contains(self.stack, "model.py"))
+        self.assertTrue(_any_filename_contains(self.stack, "optim"))
+        self.assertFalse(_any_filename_contains(self.stack, "missing.py"))
+        # Empty stack always returns False.
+        self.assertFalse(_any_filename_contains([], "anything"))
+
+    def test_any_filename_matches_regex(self) -> None:
+        import re
+
+        cu_pattern = re.compile(r"\.cu$")
+        py_pattern = re.compile(r"^model\.py$")
+        miss_pattern = re.compile(r"^never$")
+        self.assertTrue(_any_filename_matches_regex(self.stack, cu_pattern))
+        self.assertTrue(_any_filename_matches_regex(self.stack, py_pattern))
+        self.assertFalse(_any_filename_matches_regex(self.stack, miss_pattern))
+        self.assertFalse(_any_filename_matches_regex([], cu_pattern))
+
+    def test_any_frame_name_in(self) -> None:
+        self.assertTrue(
+            _any_frame_name_in(self.stack, frozenset({"custom_adamw", "other"}))
+        )
+        self.assertFalse(
+            _any_frame_name_in(self.stack, frozenset({"never", "missing"}))
+        )
+        self.assertFalse(_any_frame_name_in([], frozenset({"forward"})))
+
+    def test_any_frame_name_contains(self) -> None:
+        self.assertTrue(_any_frame_name_contains(self.stack, "forward"))
+        self.assertTrue(_any_frame_name_contains(self.stack, "kernel"))
+        self.assertFalse(_any_frame_name_contains(self.stack, "missing"))
+        self.assertFalse(_any_frame_name_contains([], "forward"))
+
+    def test_any_frame_name_startswith(self) -> None:
+        self.assertTrue(_any_frame_name_startswith(self.stack, "forward"))
+        self.assertTrue(_any_frame_name_startswith(self.stack, "my_kernel"))
+        # "kernel" appears mid-name, so startswith should not match.
+        self.assertFalse(_any_frame_name_startswith(self.stack, "kernel"))
+        self.assertFalse(_any_frame_name_startswith([], "forward"))
+
+    def test_no_frame_name_in(self) -> None:
+        # Inverse of _any_frame_name_in.
+        self.assertFalse(_no_frame_name_in(self.stack, frozenset({"custom_adamw"})))
+        self.assertTrue(_no_frame_name_in(self.stack, frozenset({"never"})))
+        # Empty stack means no frame is in the set, so True.
+        self.assertTrue(_no_frame_name_in([], frozenset({"forward"})))
+
+
 class TestOmegaConfIntegration(TestCase):
     """Tests for OmegaConf integration with custom profiling"""