[Fix] Key compile state by topology and resolve it per call

magi_compiler/_api.py

@@ -76,8 +76,34 @@ def get_attr_name_for_wrapper_installed_flag() -> str:
     return "_magi_wrapper_installed"
 
 
+def get_attr_name_for_bound_wrapper_flag(entry_name: str) -> str:
+    """Name the marker saying *entry_name* has already been patched on this instance.
+
+    Deliberately not the class-level flag: the class decorator sets that one and then
+    patches each new instance from __init__, and an instance inherits it, so reading it
+    here would skip the patching it is supposed to guard. Per method, so an instance can
+    carry more than one entry point.
+    """
+    return f"_magi_wrapper_installed_for_{entry_name}"
+
+
 def get_attr_name_for_state(entry_name: str) -> str:
-    return f"_magi_state_for_{entry_name}"
+    """Name the attribute holding an entry point's compile state, one per topology.
+
+    The state owns the captured bytecode and AOT artifacts that _run_orchestration's fast
+    paths replay directly, without going through dynamo's guards. Those artifacts have the
+    ProcessGroup they traced with baked in, because a ProcessGroup is not something dynamo
+    can turn into a graph input. So a runtime that changes CP/DP between calls -- what
+    adaptive DP does between requests -- needs one state per topology; a shared one would
+    replay the first topology's graph under every later one. Keyed this way each topology
+    compiles once and is reused whenever the runtime returns to it.
+
+    MAGI_COMPILE_TOPOLOGY_KEY is set by the runtime on every mesh change; empty means a
+    single fixed topology, i.e. the plain name.
+    """
+    topology = os.environ.get("MAGI_COMPILE_TOPOLOGY_KEY", "")
+    suffix = f"__{topology}" if topology else ""
+    return f"_magi_state_for_{entry_name}{suffix}"
 
 
 def _run_orchestration(state: MagiCompileState, args, kwargs):
@@ -209,14 +235,18 @@ def _magi_compile_bound_method(
     if not callable(getattr(instance, method_name, None)):
         raise AttributeError(f"{instance.__class__.__name__} instance has no callable method '{method_name}'")
 
-    state_attr = get_attr_name_for_state(method_name)
-    if getattr(instance, state_attr, None) is not None:
+    installed_attr = get_attr_name_for_bound_wrapper_flag(method_name)
+    if getattr(instance, installed_attr, False):
         return instance
 
     old_method = getattr(instance, method_name)
 
     @torch.compiler.disable()
     def new_call(*args, **kwargs):
+        # Per call, not per wrap: the name carries the topology, and adaptive DP changes
+        # topology between calls. Binding it at construction pinned every later call to
+        # the topology that happened to be live back then.
+        state_attr = get_attr_name_for_state(method_name)
         state = getattr(instance, state_attr, None)
         if state is None:
             _lazy_init_magi_state(instance, instance, dynamic_arg_dims, conf, model_tag, method_name, state_attr)
@@ -232,19 +262,19 @@ def new_call(*args, **kwargs):
         return _run_orchestration(state, args, kwargs)
 
     setattr(instance, method_name, new_call)
-    setattr(instance, get_attr_name_for_wrapper_installed_flag(), True)
+    setattr(instance, installed_attr, True)
     return instance
 
 
 def _magi_compile_function(func: Callable, dynamic_arg_dims: dict[str, int | list[int]], conf: CompileConfig, model_tag: str):
     """Wrap a function entry with compiled routing."""
-    state_attr = get_attr_name_for_state("function")
-    if getattr(func, state_attr, None) is not None:
+    if getattr(func, get_attr_name_for_wrapper_installed_flag(), False):
         return func
 
     @torch.compiler.disable()
     @functools.wraps(func)  # for the original function name and docstring
     def wrapper(*args, **kwargs):
+        state_attr = get_attr_name_for_state("function")  # per call: see _magi_compile_bound_method
         state = getattr(wrapper, state_attr, None)
         if state is None:
             _lazy_init_magi_state(wrapper, func, dynamic_arg_dims, conf, model_tag, None, state_attr)
@@ -255,6 +285,7 @@ def wrapper(*args, **kwargs):
 
         return _run_orchestration(state, args, kwargs)
 
+    setattr(wrapper, get_attr_name_for_wrapper_installed_flag(), True)
     return wrapper
 
 

tests/feature_tests/test_cache_topology_isolation.py

@@ -3,16 +3,24 @@
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 
-"""Verify that different parallel topologies produce isolated cache directories,
-preventing EP/CP cross-contamination of compiled artifacts."""
+"""Verify that different parallel topologies stay isolated, on disk and in memory.
+
+Compiled artifacts have the ProcessGroup they traced with baked in, so a runtime that
+changes CP/DP between calls -- adaptive DP reacting to queue depth -- must not reuse one
+topology's artifacts under another. Two things have to be keyed by topology for that:
+the cache directory, and the compile state held on the instance.
+"""
 
 from __future__ import annotations
 
+import os
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import torch
 
+from magi_compiler._api import get_attr_name_for_state, get_attr_name_for_wrapper_installed_flag
 from magi_compiler.config import _get_parallel_topology, magi_cache_dump_path, model_rank_dir_name
 
 
@@ -94,3 +102,90 @@ def test_same_topology_same_path(self):
                 name1 = model_rank_dir_name(0, None)
                 name2 = model_rank_dir_name(0, None)
         assert name1 == name2
+
+
+class TestCompileStateIsolation:
+    """The compile state on an instance is the in-memory counterpart of the cache dir.
+
+    It owns the bytecode and AOT artifacts that the fast paths in _run_orchestration replay
+    directly, without going through dynamo's guards -- so if two topologies share one state,
+    the second silently runs the first one's graph, on the first one's process group.
+    """
+
+    def test_a_shared_state_is_what_lets_one_graph_serve_two_topologies(self):
+        """The failure mode: with the key pinned, both topologies land on one attribute."""
+        with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "pinned"}):
+            under_cp8 = get_attr_name_for_state("forward")
+            under_cp4 = get_attr_name_for_state("forward")
+        assert under_cp8 == under_cp4
+
+    def test_each_topology_gets_its_own_state(self):
+        with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp8_dp1"}):
+            under_cp8 = get_attr_name_for_state("forward")
+        with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp4_dp2"}):
+            under_cp4 = get_attr_name_for_state("forward")
+        assert under_cp8 != under_cp4
+        assert "cp8_dp1" in under_cp8 and "cp4_dp2" in under_cp4
+
+    def test_a_single_fixed_topology_keeps_the_plain_name(self):
+        """No key means no adaptive runtime, and the attribute must not change shape."""
+        with patch.dict(os.environ, {}, clear=False):
+            os.environ.pop("MAGI_COMPILE_TOPOLOGY_KEY", None)
+            assert get_attr_name_for_state("forward") == "_magi_state_for_forward"
+
+    def test_the_topology_at_call_time_decides_not_the_one_at_wrap_time(self):
+        """The regression this guards.
+
+        The attribute used to be resolved once while wrapping the instance, so a model built
+        under cp=8 kept reaching for the cp=8 state no matter what the runtime switched to
+        afterwards -- which made keying by topology inert.
+        """
+        from magi_compiler._api import _magi_compile_bound_method
+
+        reached = []
+
+        def fake_init(holder, target, dims, conf, tag, method, state_attr):
+            reached.append(state_attr)
+            setattr(
+                holder,
+                state_attr,
+                SimpleNamespace(
+                    compile_config=SimpleNamespace(offload_config=SimpleNamespace(model_cpu_offload=False)),
+                    jit_compiled_code=None,
+                ),
+            )
+
+        class Probe(torch.nn.Module):
+            def forward(self, x):
+                return x
+
+        probe = Probe()
+        with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp8_dp1"}):
+            _magi_compile_bound_method(probe, {"x": 0}, SimpleNamespace(), "probe", method_name="forward")
+
+        with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp4_dp2"}):
+            with patch("magi_compiler._api._lazy_init_magi_state", side_effect=fake_init):
+                with patch("magi_compiler._api._run_orchestration", return_value=None):
+                    probe.forward(torch.zeros(2))
+
+        assert reached == ["_magi_state_for_forward__cp4_dp2"], reached
+
+    def test_the_class_decorator_still_patches_every_instance(self):
+        """A mark on the class must not answer for an instance.
+
+        magi_compile on a class marks the class and patches each instance from __init__.
+        Reading that inherited mark to decide whether an instance was patched made every
+        instance skip it: forward stayed the class's own and ran eagerly, and no state was
+        ever attached -- which is invisible except as timings that match eager.
+        """
+        from magi_compiler import magi_compile
+
+        @magi_compile(dynamic_arg_dims={"x": 0})
+        class Block(torch.nn.Module):
+            def forward(self, x):
+                return x
+
+        block = Block()
+
+        assert getattr(Block, get_attr_name_for_wrapper_installed_flag(), False), "the class went unmarked"
+        assert "forward" in vars(block), "the instance kept the class's forward, so nothing was compiled"