Revert "Cache the value of torch_key in subproc (#151057)"

This reverts commit 5f5805a6ac.

Reverted https://github.com/pytorch/pytorch/pull/151057 on behalf of https://github.com/facebook-github-bot due to Diff reverted internally ([comment](https://github.com/pytorch/pytorch/pull/151057#issuecomment-2816614510))

torch/_inductor/codecache.py

@@ -662,37 +662,12 @@ def build_code_hash(
             build_code_hash(spec.submodule_search_locations, f"{spec.name}.", hasher)
 
 
-def torch_key_cache(func: Callable[[], bytes]) -> Callable[[], bytes]:
-    """
-    This function is a reimplementation of functools.lru_cache with a
-    set function that allows prepopulating the cache.
-    """
-    # Use list for reference semantics
-    _cache: list[bytes] = []
-
-    def wrapper() -> bytes:
-        if len(_cache) == 0:
-            _cache.append(func())
-        return _cache[0]
-
-    def set_val(val: bytes) -> None:
-        assert len(_cache) == 0
-        _cache.append(val)
-
-    def clear() -> None:
-        _cache.clear()
-
-    wrapper.set = set_val  # type: ignore[attr-defined]
-    wrapper.clear = clear  # type: ignore[attr-defined]
-    return wrapper
-
-
-@torch_key_cache
+@functools.lru_cache(None)
 def torch_key() -> bytes:
     """
     Compute a key that contains relevant information about torch source files
     """
-    with dynamo_timed("inductor_codecache_torch_key", log_pt2_compile_event=False):
+    with dynamo_timed("inductor_codecache_torch_key", log_pt2_compile_event=True):
         if not config.is_fbcode():
 
             def get_code_hash(root: str) -> bytes:

torch/_inductor/compile_worker/__main__.py

@@ -1,6 +1,5 @@
 # mypy: allow-untyped-defs
 import argparse
-import base64
 import functools
 import importlib
 import logging
@@ -9,7 +8,6 @@ import sys
 from typing import TypeVar
 
 from torch._inductor.async_compile import pre_fork_setup
-from torch._inductor.codecache import torch_key
 from torch._inductor.compile_worker.subproc_pool import (
     SubprocKind,
     SubprocMain,
@@ -58,7 +56,6 @@ def main():
         parser.add_argument("--parent", type=int)
         parser.add_argument("--read-fd", type=int)
         parser.add_argument("--write-fd", type=int)
-        parser.add_argument("--torch-key", type=str)
         args = parser.parse_args()
         if os.getppid() != args.parent:
             sys.exit(0)
@@ -67,8 +64,6 @@ def main():
 
         pre_fork_setup()
 
-        torch_key.set(base64.b64decode(args.torch_key.encode("utf-8")))  # type: ignore[attr-defined]
-
         _async_compile_initializer(args.parent)
 
         SubprocMain(args.pickler, args.kind, args.workers, read_fd, write_fd).main()

torch/_inductor/compile_worker/subproc_pool.py

@@ -1,4 +1,3 @@
-import base64
 import functools
 import itertools
 import logging
@@ -21,7 +20,6 @@ from typing_extensions import Never, ParamSpec
 # justknobs, e.g., in the Triton compiler. For internal, the import installs
 # functionality to destroy singletons before forking and re-enable them after.
 import torch._thread_safe_fork  # noqa: F401
-from torch._inductor.codecache import torch_key
 from torch._inductor.compile_worker.utils import _async_compile_initializer
 from torch._inductor.utils import get_ld_library_path
 
@@ -117,7 +115,6 @@ class SubprocPool:
         read_fd, subproc_write_fd = os.pipe()
         self.write_pipe = os.fdopen(write_fd, "wb")
         self.read_pipe = os.fdopen(read_fd, "rb")
-        torch_key_str = base64.b64encode(torch_key()).decode("utf-8")
 
         cmd = [
             sys.executable,
@@ -128,7 +125,6 @@ class SubprocPool:
             f"--parent={os.getpid()}",
             f"--read-fd={str(subproc_read_fd)}",
             f"--write-fd={str(subproc_write_fd)}",
-            f"--torch-key={torch_key_str}",
         ]
         self.process = subprocess.Popen(
             cmd,