diff --git a/magi_compiler/config.py b/magi_compiler/config.py index 44fbd03..f1a92e1 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -312,11 +312,19 @@ def __repr__(self, indent: int = 4): def model_rank_dir_name(model_idx: int, model_tag: str | None) -> str: - """Directory name for a model instance: ``model_{idx}[_{tag}]_rank_{rank}``.""" + """Directory name for a model instance: ``model_{idx}[_{tag}]_rank_{rank}_ws{world_size}``. + + world_size is included because runtime model behavior (e.g. MoE EP head + padding, tensor strides in custom-op meta functions) can differ across + world sizes. Without it, switching between e.g. EP=6 and EP=8 would + silently reuse stale compiled artifacts and trigger stride-mismatch + assertions at runtime. + """ rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + world_size = torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1 if model_tag: - return f"model_{model_idx}_{model_tag}_rank_{rank}" - return f"model_{model_idx}_rank_{rank}" + return f"model_{model_idx}_{model_tag}_rank_{rank}_ws{world_size}" + return f"model_{model_idx}_rank_{rank}_ws{world_size}" def debug_dump_path(cache_root_dir: str, model_idx: int, model_tag: str | None = None) -> Path: diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 9569bb2..ddf9ed2 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -128,12 +128,12 @@ def initialize_cache(self, cache_dir: Path): self.cache_dir.mkdir(parents=True, exist_ok=True) magi_logger.info("Using cache directory: %s for MagiCompiler", cache_dir) + self.cache = {} + self._remaining_restart_skips = {} + if self.cache_file_path.exists(): - # load the cache from the file with self.cache_file_path.open() as f: - # Parse Python literals using ast.literal_eval, which is a safe alternative to eval(). raw = ast.literal_eval(f.read()) - self.cache = {} for entry, handle in raw.items(): cache_entry = CacheEntry(*entry) cache_handle = CacheHandle(*handle) @@ -160,6 +160,13 @@ def load(self, graph: fx.GraphModule, example_inputs: list[Any], cache_entry: Ca return None cache_handle = self.cache[cache_entry] + + artifact_dir = Path(cache_handle.path) + if not artifact_dir.exists(): + magi_logger.warning("Stale cache entry removed (artifact dir missing): %s", cache_handle.path) + del self.cache[cache_entry] + return None + if cache_entry.graph_index not in self._remaining_restart_skips: self._remaining_restart_skips[cache_entry.graph_index] = cache_handle.restart_analysis_count remaining = self._remaining_restart_skips[cache_entry.graph_index] @@ -224,6 +231,7 @@ def _maybe_store_cache_entry( if self.disable_cache: return False if cache_handle is None: + self.cache.pop(cache_entry, None) return False prev_handle = self.cache.get(cache_entry) diff --git a/tests/feature_tests/test_cache_invariant.py b/tests/feature_tests/test_cache_invariant.py new file mode 100644 index 0000000..ad794fb --- /dev/null +++ b/tests/feature_tests/test_cache_invariant.py @@ -0,0 +1,146 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Cache invariant tests: every entry in ``CompilerManager.cache`` must have a +valid on-disk artifact directory. These tests reproduce three classes of stale- +cache bugs present on ``origin/main`` (before fix). + +Each test constructs a ``CompilerManager`` with known-bad state and asserts the +invariant that the code *should* maintain. On unfixed code they FAIL. +""" + +from __future__ import annotations + +import pprint +from pathlib import Path +from unittest.mock import patch + +import torch +import torch.fx as fx + +from magi_compiler.config import CompileConfig +from magi_compiler.magi_backend._cache_data_cls import CacheEntry, CacheHandle +from magi_compiler.magi_backend.magi_backend import CompilerManager + + +def _make_cm(tmp_path: Path) -> CompilerManager: + cfg = CompileConfig(cache_root_dir=str(tmp_path), backend="eager", _cli_parse_args=False) + return CompilerManager(cfg) + + +def _write_index(cache_dir: Path, data: dict) -> None: + (cache_dir / "subgraph_indices.py").write_text(pprint.pformat(data)) + + +def _dummy_graph() -> fx.GraphModule: + return fx.GraphModule(torch.nn.Module(), fx.Graph()) + + +# --------------------------------------------------------------------------- +# Bug 1: load() does not check whether artifact dir exists on disk. +# +# On origin/main, load() blindly passes the handle to compiler.load() which +# calls CompiledArtifact.load() on a non-existent path → crash. +# +# Expected after fix: load() returns None and removes the stale entry. +# --------------------------------------------------------------------------- + + +def test_load_rejects_missing_artifact_dir(tmp_path: Path) -> None: + cache_dir = tmp_path / "model_bug1" + cache_dir.mkdir(parents=True) + bogus = cache_dir / "artifact_shape_None_subgraph_0" + assert not bogus.exists() + + data = {(None, 0, "inductor_standalone"): ("artifact_shape_None_subgraph_0", str(bogus), 0)} + _write_index(cache_dir, data) + + cm = _make_cm(tmp_path) + cm.initialize_cache(cache_dir) + + entry = CacheEntry(None, 0, "inductor_standalone") + assert entry in cm.cache, "precondition: entry loaded from index" + + # Patch compiler.load so that if CompilerManager.load reaches it we know + # the guard is missing (it should never be called for a missing dir). + with patch.object(cm.compiler, "load", side_effect=AssertionError("compiler.load should not be called")): + result = cm.load(_dummy_graph(), [], entry) + + assert result is None, "load() must return None for missing artifact dir" + assert entry not in cm.cache, "stale entry must be removed from cache" + + +# --------------------------------------------------------------------------- +# Bug 2: _maybe_store_cache_entry(entry, None, ...) does not remove a stale +# handle that was rehydrated from subgraph_indices.py. +# +# On origin/main, it just returns False; the ghost handle survives in memory +# and save_to_file() will re-serialize it. +# +# Expected after fix: the stale entry is deleted from self.cache. +# --------------------------------------------------------------------------- + + +def test_store_none_handle_removes_stale_entry(tmp_path: Path) -> None: + cache_dir = tmp_path / "model_bug2" + cache_dir.mkdir(parents=True) + bogus = cache_dir / "artifact_shape_None_subgraph_0" + + data = {(None, 0, "inductor_standalone"): ("artifact_shape_None_subgraph_0", str(bogus), 1)} + _write_index(cache_dir, data) + + cm = _make_cm(tmp_path) + cm.initialize_cache(cache_dir) + + entry = CacheEntry(None, 0, "inductor_standalone") + assert entry in cm.cache, "precondition: entry loaded from index" + + cm._maybe_store_cache_entry(entry, None, None, "artifact_shape_None_subgraph_0") + + assert entry not in cm.cache, "stale entry must be removed when handle is None" + + +# --------------------------------------------------------------------------- +# Bug 3: initialize_cache() does not reset self.cache when +# subgraph_indices.py is absent. +# +# On origin/main, `self.cache = {}` only runs inside the +# `if self.cache_file_path.exists()` branch. So a second call to +# initialize_cache (same dir, no index file) preserves ghost handles +# injected between the two calls. +# +# Expected after fix: self.cache is always reset to {}. +# --------------------------------------------------------------------------- + + +def test_reinit_without_indices_clears_memory(tmp_path: Path) -> None: + cache_dir = tmp_path / "model_bug3" + cache_dir.mkdir(parents=True) + + cm = _make_cm(tmp_path) + cm.initialize_cache(cache_dir) + + # Simulate a ghost handle left from a partial first compilation. + bogus = cache_dir / "artifact_shape_None_subgraph_0" + entry = CacheEntry(None, 0, "inductor_standalone") + cm.cache[entry] = CacheHandle("artifact_shape_None_subgraph_0", str(bogus), 1) + cm._remaining_restart_skips[0] = 1 + + assert not (cache_dir / "subgraph_indices.py").exists(), "precondition: no index file" + + # Re-initialize same directory — should clear everything. + cm.initialize_cache(cache_dir) + + assert entry not in cm.cache, "ghost handle must be cleared on re-init" + assert 0 not in cm._remaining_restart_skips, "restart skips must be cleared when no index"