fix test_simple_inductor_graph_partition (#26522)

Signed-off-by: Boyuan Feng <boyuan@meta.com>

Author: BoyuanFeng

tests/compile/piecewise/test_simple.py

@@ -143,10 +143,14 @@ def test_simple_piecewise_compile(use_inductor):
 
 @torch.inference_mode()
 @pytest.mark.parametrize("splitting_ops", [["silly.attention"], []])
-def test_simple_inductor_graph_partition(splitting_ops):
+def test_simple_inductor_graph_partition(splitting_ops, monkeypatch):
     if not is_torch_equal_or_newer("2.9.0.dev"):
         pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
 
+    # disable compile cache so that we run separately for different splitting_ops
+    # and get the expected number of cudagraphs captured.
+    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
+
     _run_simple_model(
         # Inductor graph partition automatically resets splitting_ops to an empty list
         splitting_ops=splitting_ops,

vllm/compilation/compiler_interface.py

@@ -332,7 +332,10 @@ def hijack_load(*args, **kwargs):
                 nonlocal file_path
                 compiled_fn = inductor_compiled_graph.current_callable
                 file_path = compiled_fn.__code__.co_filename  # noqa
-                if not file_path.startswith(self.base_cache_dir):
+                if (
+                    not file_path.startswith(self.base_cache_dir)
+                    and compiled_fn.__closure__ is not None
+                ):
                     # hooked in the align_inputs_from_check_idxs function
                     # in torch/_inductor/utils.py
                     for cell in compiled_fn.__closure__:
@@ -359,7 +362,10 @@ def hijacked_compile_fx_inner(*args, **kwargs):
                     nonlocal file_path
                     compiled_fn = inductor_compiled_graph.current_callable
                     file_path = compiled_fn.__code__.co_filename  # noqa
-                    if not file_path.startswith(self.base_cache_dir):
+                    if (
+                        not file_path.startswith(self.base_cache_dir)
+                        and compiled_fn.__closure__ is not None
+                    ):
                         # hooked in the align_inputs_from_check_idxs function
                         # in torch/_inductor/utils.py
                         for cell in compiled_fn.__closure__: