Avoid writing temporary modules to disk

apmorton · pytorchmergebot · commit a633a85901aa · 2025-10-27T02:41:53.000Z
diff --git a/torch/distributed/nn/jit/instantiator.py b/torch/distributed/nn/jit/instantiator.py
@@ -1,11 +1,8 @@
 #!/usr/bin/python3
 # mypy: allow-untyped-defs
-import atexit
 import importlib
-import logging
-import os
+import importlib.util
 import sys
-import tempfile
 from typing import Optional
 
 import torch
@@ -14,15 +11,7 @@
 )
 
 
-logger = logging.getLogger(__name__)
-
-
 _FILE_PREFIX = "_remote_module_"
-_TEMP_DIR = tempfile.TemporaryDirectory()
-INSTANTIATED_TEMPLATE_DIR_PATH = _TEMP_DIR.name
-atexit.register(_TEMP_DIR.cleanup)
-logger.info("Created a temporary directory at %s", INSTANTIATED_TEMPLATE_DIR_PATH)
-sys.path.append(INSTANTIATED_TEMPLATE_DIR_PATH)
 
 
 def get_arg_return_types_from_interface(module_interface):
@@ -63,40 +52,35 @@ def get_arg_return_types_from_interface(module_interface):
     return args_str, arg_types_str, return_type_str
 
 
-def _write(out_path, text):
-    old_text: Optional[str]
-    try:
-        with open(out_path) as f:
-            old_text = f.read()
-    except OSError:
-        old_text = None
-    if old_text != text:
-        with open(out_path, "w") as f:
-            logger.info("Writing %s", out_path)
-            f.write(text)
-    else:
-        logger.info("Skipped writing %s", out_path)
+class StringLoader(importlib.abc.SourceLoader):
+    def __init__(self, data):
+        self.data = data
+
+    def get_source(self, fullname):
+        return self.data
+    
+    def get_data(self, path):
+        return self.data.encode("utf-8")
+    
+    def get_filename(self, fullname):
+        return fullname
 
 
 def _do_instantiate_remote_module_template(
     generated_module_name, str_dict, enable_moving_cpu_tensors_to_cuda
 ):
-    generated_code_text = get_remote_module_template(
-        enable_moving_cpu_tensors_to_cuda
-    ).format(**str_dict)
-    out_path = os.path.join(
-        INSTANTIATED_TEMPLATE_DIR_PATH, f"{generated_module_name}.py"
+    if generated_module_name in sys.modules:
+        return sys.modules[generated_module_name]
+
+    spec = importlib.util.spec_from_loader(
+        generated_module_name,
+        StringLoader(get_remote_module_template(enable_moving_cpu_tensors_to_cuda).format(**str_dict)),
+        origin='torch-jit',
     )
-    _write(out_path, generated_code_text)
-
-    # From importlib doc,
-    # > If you are dynamically importing a module that was created since
-    # the interpreter began execution (e.g., created a Python source file),
-    # you may need to call invalidate_caches() in order for the new module
-    # to be noticed by the import system.
-    importlib.invalidate_caches()
-    generated_module = importlib.import_module(f"{generated_module_name}")
-    return generated_module
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[generated_module_name] = module
+    spec.loader.exec_module(module)
+    return module
 
 
 def instantiate_scriptable_remote_module_template(
