[inductor] Test megacache loading during kernel finalization (#196763) (#196763)

Summary:

S686093 occurred when asynchronous megacache population loaded live Triton kernels while compilation finalized entries in `CompiledTritonKernels`. The shipped backout defers bundle loading until compilation, but that concurrency boundary had no regression coverage.

Adapt the finalizer mechanism from D114901960 into deterministic CPU tests. Events force `torch.compiler.load_cache_artifacts()` to run during cache clearing, entry removal, and replacement. Assert that population preserves the compiler's cache state and that the persisted bundle remains loadable afterward on the compile thread. Mirror the tests into `xplat/caffe2`.

Authored with Codex.

Test Plan:
```
buck2 test fbcode//caffe2/test/inductor:codecache -- test_cache_hot_load_during_kernel_finalization
```
All three cases pass: https://www.internalfb.com/intern/testinfra/testrun/21110623299276337

Negative control: temporarily restored eager `TritonBundler.read_and_emit()` in `InductorCacheArtifact.populate_cache()` and ran the same command. All three cases failed at the cache-identity assertion, with no timeout or crash: https://www.internalfb.com/intern/testinfra/testrun/19703248415725483
Removed the temporary mutation afterward; no production-code changes remain.

```
arc f fbcode/caffe2/test/inductor/test_codecache.py xplat/caffe2/test/inductor/test_codecache.py
arc lint -a fbcode/caffe2/test/inductor/test_codecache.py xplat/caffe2/test/inductor/test_codecache.py
arc lint --engine extra --never-apply-patches fbcode/caffe2/test/inductor/test_codecache.py xplat/caffe2/test/inductor/test_codecache.py
```
Fast and comprehensive lint pass. Formatting reports no applicable linters.

Formatter verification after rebasing:
```
tools/third-party/ruff/ruff format --check --no-cache --config fbcode/caffe2/pyproject.toml fbcode/caffe2/test/inductor/test_codecache.py xplat/caffe2/test/inductor/test_codecache.py
tools/third-party/ruff/ruff check --no-cache --config fbcode/caffe2/pyproject.toml fbcode/caffe2/test/inductor/test_codecache.py xplat/caffe2/test/inductor/test_codecache.py
arc lint -a fbcode/caffe2/test/inductor/test_codecache.py xplat/caffe2/test/inductor/test_codecache.py
arc lint --engine extra --never-apply-patches fbcode/caffe2/test/inductor/test_codecache.py xplat/caffe2/test/inductor/test_codecache.py
```
Both mirrored files pass formatting and lint with packaged Ruff 0.14.0; PyTorch CI pins Ruff 0.14.4. Formatting preserves both Python syntax trees.

Re-ran the focused Buck regression command on the rebased revision. The run remained in remote test discovery with zero results after about eight minutes and was canceled before applying the formatting-only changes: https://www.internalfb.com/intern/testinfra/testrun/33776997240090459
This rerun is inconclusive; the earlier passing run and negative control remain the regression evidence.

V3 CI verification (2026-09-15): GitHub lint passed, including `Lint / lintrunner-noclang-partial / lint`. All six `inductor_core` CPU jobs passed (both shards on Python 3.11, 3.12, and 3.13); the core-suite selection includes `inductor/test_codecache`. All seven `inductor_amx` CPU shards also passed.

The remaining two CPU failures are the `inductor_aoti_fallback` shards in https://github.com/pytorch/pytorch/actions/runs/34959021263. These jobs are explicitly documented as known red and marked `unstable` in `.github/workflows/inductor-unittest.yml`; their explicit test selection in `.ci/pytorch/test.sh` excludes `test_codecache`. They remain reported as failures in Phabricator. The `github-export-checks` and `pytorch-oss-checks` gates report missing GitHub approval for https://github.com/pytorch/pytorch/pull/196763; the synced PR has zero GitHub reviews. These require an approving GitHub review before retrying the gates.

Reviewed By: aorenste, warrendeng

Differential Revision: D119397497

Pull Request resolved: https://github.com/pytorch/pytorch/pull/196763
Approved by: https://github.com/aorenste

test/inductor/test_codecache.py

@@ -13,9 +13,11 @@ import subprocess
 import sys
 import tempfile
 import textwrap
+import threading
 import types
 import unittest
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from typing import Any, cast
 from typing_extensions import override
@@ -30,6 +32,7 @@ from torch._dynamo.utils import counters
 from torch._functorch import config as functorch_config
 from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache
 from torch._inductor import config, config_comms, metrics
+from torch._inductor.async_compile import CompiledTritonKernels
 from torch._inductor.cache_key import (
     AUTOTUNE_CACHE_KEY_STRATEGY,
     CacheKeyStrategy,
@@ -49,6 +52,7 @@ from torch._inductor.codecache import (
     FxGraphCachePickler,
     FxGraphHashDetails,
     PyCodeCache,
+    StaticAutotunerFuture,
     TensorMetadata,
     TensorMetadataAndValues,
     triton_key,
@@ -66,6 +70,11 @@ from torch._inductor.graph import GraphLowering
 from torch._inductor.mock_cache import global_stats, PatchCaches, Stats
 from torch._inductor.runtime.runtime_utils import cache_dir
 from torch._inductor.test_case import run_tests, TestCase
+from torch._inductor.triton_bundler import (
+    StaticallyLaunchedAutotuner,
+    TritonBundle,
+    TritonBundler,
+)
 from torch._inductor.utils import clear_caches, fresh_cache
 from torch._library import capture_triton
 from torch._subclasses import FakeTensorMode
@@ -1603,6 +1612,87 @@ class TestFxGraphCache(TestCase):
     def test_cache_hot_load_empty(self):
         self.assertIsNone(torch.compiler.save_cache_artifacts())
 
+    @config.patch(
+        bundle_triton_into_fx_graph_cache=True,
+        use_static_triton_launcher=True,
+    )
+    @parametrize("operation", ("clear", "remove", "replace"))
+    def test_cache_hot_load_during_kernel_finalization(self, operation):
+        kernel_src = "megacache_concurrent_kernel"
+        bundle = TritonBundle(
+            [],
+            [
+                StaticallyLaunchedAutotuner(
+                    CompiledTritonKernels.key(kernel_src),
+                    "concurrent_kernel",
+                    types.SimpleNamespace(compile_results=[]),
+                )
+            ],
+        )
+        content = pickle.dumps(types.SimpleNamespace(_triton_bundle=bundle))
+        graph_key = "f" + "0" * 51
+        CacheArtifactManager.record_artifact("inductor", graph_key, content)
+        artifacts = torch.compiler.save_cache_artifacts()
+        self.assertIsNotNone(artifacts)
+        artifact_bytes, _ = artifacts
+        CacheArtifactManager.clear()
+
+        finalizing = threading.Event()
+        loaded = threading.Event()
+        finalized = []
+
+        class FinalizedKernel:
+            def __del__(self):
+                # Kernel finalizers can release the GIL while the cache drops
+                # its last reference. Force megacache loading into that window.
+                finalizing.set()
+                finalized.append(loaded.wait(timeout=30))
+
+        def load_artifacts():
+            try:
+                self.assertTrue(finalizing.wait(timeout=30))
+                return torch.compiler.load_cache_artifacts(artifact_bytes)
+            finally:
+                loaded.set()
+
+        self.addCleanup(CompiledTritonKernels.cache_clear)
+        CompiledTritonKernels.save(kernel_src, FinalizedKernel())
+        self.assertIsInstance(CompiledTritonKernels.get(kernel_src), FinalizedKernel)
+        replacement = mock.sentinel.compiled_kernel
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            loading = executor.submit(load_artifacts)
+            if operation == "clear":
+                CompiledTritonKernels.cache_clear()
+            elif operation == "remove":
+                CompiledTritonKernels.remove_future(kernel_src)
+            else:
+                CompiledTritonKernels.save(kernel_src, replacement)
+            cache_info = loading.result(timeout=30)
+
+        self.assertEqual(finalized, [True])
+        self.assertIsNotNone(cache_info)
+        self.assertEqual(cache_info.inductor_artifacts, [graph_key])
+        expected = replacement if operation == "replace" else None
+        self.assertIs(CompiledTritonKernels.get(kernel_src), expected)
+        self.assertEqual(
+            counters["inductor"]["triton_bundler_load_static_autotuner"], 0
+        )
+
+        directory = FxGraphCache._get_tmp_dir_for_key(graph_key)
+        filenames = os.listdir(directory)
+        self.assertEqual(len(filenames), 1)
+        with open(os.path.join(directory, filenames[0]), "rb") as f:
+            restored = pickle.load(f)
+
+        # The bundle must still be loadable on the compile thread.
+        metadata = TritonBundler.read_and_emit(restored._triton_bundle)
+        self.assertEqual(
+            metadata.statically_launched_kernel_names, ["concurrent_kernel"]
+        )
+        self.assertIsInstance(
+            CompiledTritonKernels.get(kernel_src), StaticAutotunerFuture
+        )
+
     def test_cache_hot_load_generic(self):
         class CacheStub:
             def __init__(self):