[Fix] Include world_size in cache directory name to prevent EP cross-contamination

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:

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)

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"