diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 43ae1da..c6dc57f 100644 --- a/magi_compiler/_api.py +++ b/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 diff --git a/tests/feature_tests/test_cache_topology_isolation.py b/tests/feature_tests/test_cache_topology_isolation.py index bcc0b21..c7df040 100644 --- a/tests/feature_tests/test_cache_topology_isolation.py +++ b/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"