From 7beaf7ca43bb2c2703f0f32976e364100f5b0e78 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Thu, 10 Sep 2026 22:18:03 +0800 Subject: [PATCH] fix: allow restart-analysis replay in assert_cache_hit mode When assert_cache_hit=True and a cache entry exists but load() returns None due to restart_analysis_count > 0, fall through to the normal compile path instead of raising RuntimeError. standalone_compile will trigger the same TensorifyScalarRestartAnalysis that occurred during bake, causing dynamo to re-trace with the correct graph shape. On retry, load() succeeds. Without this fix, verify_compile_cache.sh fails with a false cache miss for any subgraph that encountered RestartAnalysis during bake. Add regression test: test_assert_cache_hit_restart_analysis.py --- magi_compiler/magi_backend/magi_backend.py | 20 ++++- .../test_assert_cache_hit_restart_analysis.py | 80 +++++++++++++++++++ 2 files changed, 96 insertions(+), 4 deletions(-) create mode 100644 tests/feature_tests/cache/test_assert_cache_hit_restart_analysis.py diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 8a487d2..8d6379e 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -215,10 +215,22 @@ def compile( return compiled_graph if self.compile_config.assert_cache_hit: - raise RuntimeError( - f"MAGI_COMPILE_ASSERT_CACHE_HIT: cache miss for runtime_shape={runtime_shape} " - f"graph_index={graph_index}. The pre-baked compile cache does not cover this subgraph." - ) + if cache_entry not in self.cache: + raise RuntimeError( + f"MAGI_COMPILE_ASSERT_CACHE_HIT: cache miss for runtime_shape={runtime_shape} " + f"graph_index={graph_index}. The pre-baked compile cache does not cover this subgraph." + ) + cache_handle = self.cache[cache_entry] + if cache_handle.restart_analysis_count == 0: + raise RuntimeError( + f"MAGI_COMPILE_ASSERT_CACHE_HIT: cache load failed for runtime_shape={runtime_shape} " + f"graph_index={graph_index}. Cache entry exists but artifact could not be loaded " + f"(restart_analysis_count=0)." + ) + # restart_analysis_count > 0 — restart-analysis replay in progress. + # Fall through to normal compile path: standalone_compile will trigger + # TensorifyScalarRestartAnalysis (same as bake), dynamo re-traces, and + # on retry load() succeeds (graph shape matches cached artifact). # Step2: Compile the graph key = f"artifact_shape_{runtime_shape}_subgraph_{graph_index}" diff --git a/tests/feature_tests/cache/test_assert_cache_hit_restart_analysis.py b/tests/feature_tests/cache/test_assert_cache_hit_restart_analysis.py new file mode 100644 index 0000000..57aec1c --- /dev/null +++ b/tests/feature_tests/cache/test_assert_cache_hit_restart_analysis.py @@ -0,0 +1,80 @@ +# 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. + +from __future__ import annotations + +import ast +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_assert_cache_hit_succeeds_with_restart_analysis(tmp_path: Path): + """Regression test: assert_cache_hit must not false-alarm on restart_analysis_count. + + When a subgraph triggers RestartAnalysis during bake, the cache records + restart_analysis_count > 0. In normal (non-verify) mode the loader + replays those skips before returning the real artifact. But under + assert_cache_hit=True (verify mode), that replay returns None which the + caller interprets as a cache miss → RuntimeError. + + This test: + process-1 — compile (bake): warms the cache, encounters RestartAnalysis + process-2 — load with MAGI_COMPILE_ASSERT_CACHE_HIT=1 (verify): + must succeed after the fix (skip restart replay in verify mode) + """ + helper_path = Path(__file__).parent / "cache_reuse_helper" / "restart_analysis_cache_helper.py" + cache_root = tmp_path / "cache" + out_bake = tmp_path / "bake.json" + out_verify = tmp_path / "verify.json" + + env = os.environ.copy() + env["MAGI_LOGGING_LEVEL"] = "info" + env["MAGI_COMPILE_CACHE_ROOT_DIR"] = str(cache_root) + + # ── Process 1: bake (warm cache) ────────────────────────────────── + cmd_bake = [sys.executable, str(helper_path), "--output", str(out_bake)] + p_bake = subprocess.run(cmd_bake, env=env, capture_output=True, text=True) + assert p_bake.returncode == 0, f"bake process failed\nstdout:\n{p_bake.stdout}\nstderr:\n{p_bake.stderr}" + assert ( + "standalone_compile raised RestartAnalysis" in p_bake.stderr + ), "bake process did not encounter RestartAnalysis — test precondition violated" + + # Verify at least one cache handle has restart_analysis_count > 0 + cache_files = list(cache_root.rglob("subgraph_indices.py")) + assert cache_files, "no cache file generated during bake" + any_marked = False + for cache_file in cache_files: + raw = ast.literal_eval(cache_file.read_text()) + for _, handle in raw.items(): + if len(handle) >= 3 and int(handle[2]) > 0: + any_marked = True + break + if any_marked: + break + assert any_marked, "expected at least one cache handle with restart_analysis_count > 0" + + # ── Process 2: verify (assert_cache_hit) ────────────────────────── + verify_env = {**env, "MAGI_COMPILE_ASSERT_CACHE_HIT": "1"} + cmd_verify = [sys.executable, str(helper_path), "--output", str(out_verify)] + p_verify = subprocess.run(cmd_verify, env=verify_env, capture_output=True, text=True) + + assert p_verify.returncode == 0, ( + f"verify (assert_cache_hit) failed with restart_analysis_count > 0.\n" + f"This is the bug: restart-replay returns None → false cache miss.\n" + f"stderr:\n{p_verify.stderr}" + ) + + # Sanity-check: outputs should be numerically close + payload_bake = json.loads(out_bake.read_text()) + payload_verify = json.loads(out_verify.read_text()) + assert payload_verify["shape"] == payload_bake["shape"] + assert abs(payload_bake["sum"] - payload_verify["sum"]) < 1e-2