Improve async workers to handle forking for async compile (pytorch#142072)

Pull Request resolved: pytorch#142072
Approved by: https://github.com/masnesral

Author: aorenste

torch/_inductor/async_compile.py

@@ -1,6 +1,7 @@
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
+import atexit
 import functools
 import logging
 import os
@@ -327,3 +328,10 @@ def wait(self, scope: Dict[str, Any]) -> None:
     pass
 else:
     AsyncCompile.warm_pool()
+
+# On exit give the workers a chance to clean themselves up. Without this the
+# resource_tracker can complain about leaked semaphores coming from the
+# ProcessPoolExecutor:
+#   UserWarning: resource_tracker: There appear to be 5 leaked semaphore objects
+#   to clean up at shutdown
+atexit.register(shutdown_compile_workers)

torch/_inductor/compile_worker/__main__.py

@@ -1,15 +1,25 @@
 # mypy: allow-untyped-defs
 import argparse
+import functools
+import importlib
 import logging
 import os
 import sys
+from typing import Type, TypeVar
 
 from torch._inductor.async_compile import pre_fork_setup
-from torch._inductor.compile_worker.subproc_pool import SubprocMain
+from torch._inductor.compile_worker.subproc_pool import (
+    SubprocKind,
+    SubprocMain,
+    SubprocPickler,
+)
 from torch._inductor.compile_worker.watchdog import _async_compile_initializer
 from torch._inductor.runtime.compile_tasks import _set_triton_ptxas_path
 
 
+_T = TypeVar("_T")
+
+
 log = logging.getLogger(__name__)
 
 _set_triton_ptxas_path()
@@ -22,9 +32,26 @@
     pass
 
 
+def _lookup_and_create_type(base: Type[_T], qname: str) -> _T:
+    """
+    Given a base type and qualified name: import & lookup that name, check
+    that it's of the given type and then instantiate it.
+    """
+    pkg, name = qname.rsplit(".", 1)
+    mod = importlib.import_module(pkg)
+    ty = getattr(mod, name)
+    if not issubclass(ty, base):
+        raise TypeError(f"Type {ty} is not a subtype of {base}")
+    return ty()
+
+
 def main():
     try:
         parser = argparse.ArgumentParser()
+        parser.add_argument(
+            "--pickler", type=functools.partial(_lookup_and_create_type, SubprocPickler)
+        )
+        parser.add_argument("--kind", type=SubprocKind)
         parser.add_argument("--workers", type=int)
         parser.add_argument("--parent", type=int)
         parser.add_argument("--read-fd", type=int)
@@ -38,7 +65,8 @@ def main():
         pre_fork_setup()
 
         _async_compile_initializer(args.parent)
-        SubprocMain(args.workers, read_fd, write_fd).main()
+
+        SubprocMain(args.pickler, args.kind, args.workers, read_fd, write_fd).main()
     except Exception:
         log.exception("Uncaught exception in compile_worker subprocess")
 

torch/_inductor/compile_worker/subproc_pool.py

@@ -9,9 +9,11 @@
 import sys
 import threading
 import traceback
+import typing
 from concurrent.futures import Future, ProcessPoolExecutor
 from concurrent.futures.process import BrokenProcessPool
-from typing import Any, BinaryIO, Callable, Dict, Tuple, TypeVar
+from enum import Enum
+from typing import Any, BinaryIO, Callable, Dict, Optional, Tuple, TypeVar
 from typing_extensions import Never, ParamSpec
 
 # _thread_safe_fork is needed because the subprocesses in the pool can read
@@ -88,14 +90,39 @@ def __init__(self, details: str) -> None:
         super().__init__(f"An exception occurred in a subprocess:\n\n{details}")
 
 
+class SubprocPickler:
+    """
+    Allows a caller to provide a custom pickler for passing data with the
+    subprocess.
+    """
+
+    def dumps(self, obj: object) -> bytes:
+        return pickle.dumps(obj, pickle.HIGHEST_PROTOCOL)
+
+    def loads(self, data: bytes) -> object:
+        return pickle.loads(data)
+
+
+class SubprocKind(Enum):
+    FORK = "fork"
+    SPAWN = "spawn"
+
+
 class SubprocPool:
     """
     Mimic a concurrent.futures.ProcessPoolExecutor, but wrap it in
     a subprocess.Popen() to try to avoid issues with forking/spawning
     """
 
-    def __init__(self, nprocs: int) -> None:
+    def __init__(
+        self,
+        nprocs: int,
+        pickler: Optional[SubprocPickler] = None,
+        kind: SubprocKind = SubprocKind.FORK,
+    ) -> None:
         entry = os.path.join(os.path.dirname(__file__), "__main__.py")
+        self.pickler = pickler or SubprocPickler()
+        self.kind = kind
 
         subproc_read_fd, write_fd = os.pipe()
         read_fd, subproc_write_fd = os.pipe()
@@ -105,6 +132,8 @@ def __init__(self, nprocs: int) -> None:
         cmd = [
             sys.executable,
             entry,
+            f"--pickler={self.pickler.__class__.__module__}.{self.pickler.__class__.__name__}",
+            f"--kind={self.kind.value}",
             f"--workers={nprocs}",
             f"--parent={os.getpid()}",
             f"--read-fd={str(subproc_read_fd)}",
@@ -143,7 +172,7 @@ def submit(
     ) -> Future[_T]:
         if args or kwargs:
             job_fn = functools.partial(job_fn, *args, **kwargs)
-        job_data = pickle.dumps(job_fn, pickle.HIGHEST_PROTOCOL)
+        job_data = self.pickler.dumps(job_fn)
         future: Future[_T]
         with self.futures_lock:
             job_id = next(self.job_id_count)
@@ -156,31 +185,48 @@ def submit(
         return future
 
     def _read_thread(self) -> None:
-        try:
-            while True:
+        while True:
+            data = b""
+            try:
                 job_id, data = _recv_msg(self.read_pipe)
-                if job_id < 0:
-                    if self.running:
-                        log.warning("SubprocPool unclean exit")
-                    self.read_pipe.close()
+            except Exception as e:
+                # Something went wrong during the read. There's no way we have a
+                # valid job_id.
+                log.exception("failure in subproc_pool._recv_msg")
+                job_id = -1
+
+            if job_id < 0:
+                # read_pipe returned None or got exception
+                if self.running:
+                    log.warning("SubprocPool unclean exit")
+                    self.running = False
+                self.read_pipe.close()
+                # Cancel all the pending futures.
+                self.shutdown()
+                return
+
+            try:
+                result = self.pickler.loads(data)
+            except Exception as e:
+                # Something went wrong unpickling. We have a job_id so just
+                # notify that particular future and continue on.
+                log.exception("unpickle failure in SubprocPool._read_thread")
+                result = e
+
+            with self.futures_lock:
+                if not self.running:
                     return
-                result = pickle.loads(data)
-                with self.futures_lock:
-                    if not self.running:
-                        return
-                    if isinstance(result, _SubprocExceptionInfo):
-                        # An exception occurred in the submitted job
-                        self.pending_futures[job_id].set_exception(
-                            SubprocException(result.details)
-                        )
-                    elif isinstance(result, Exception):
-                        # An exception occurred in some of our subprocess machinery.
-                        self.pending_futures[job_id].set_exception(result)
-                    else:
-                        self.pending_futures[job_id].set_result(result)
-                    del self.pending_futures[job_id]
-        except Exception:
-            log.exception("failure in SubprocPool._read_thread")
+                if isinstance(result, _SubprocExceptionInfo):
+                    # An exception occurred in the submitted job
+                    self.pending_futures[job_id].set_exception(
+                        SubprocException(result.details)
+                    )
+                elif isinstance(result, Exception):
+                    # An exception occurred in some of our subprocess machinery.
+                    self.pending_futures[job_id].set_exception(result)
+                else:
+                    self.pending_futures[job_id].set_result(result)
+                del self.pending_futures[job_id]
 
     def shutdown(self) -> None:
         try:
@@ -204,7 +250,16 @@ def shutdown(self) -> None:
 class SubprocMain:
     """Communicates with a SubprocPool in the parent process, called by __main__.py"""
 
-    def __init__(self, nprocs: int, read_pipe: BinaryIO, write_pipe: BinaryIO) -> None:
+    def __init__(
+        self,
+        pickler: SubprocPickler,
+        kind: SubprocKind,
+        nprocs: int,
+        read_pipe: BinaryIO,
+        write_pipe: BinaryIO,
+    ) -> None:
+        self.pickler = pickler
+        self.kind = kind
         self.read_pipe = read_pipe
         self.write_pipe = write_pipe
         self.write_lock = threading.Lock()
@@ -215,7 +270,7 @@ def __init__(self, nprocs: int, read_pipe: BinaryIO, write_pipe: BinaryIO) -> No
     def _new_pool(self, nprocs: int, warm: bool) -> ProcessPoolExecutor:
         pool = ProcessPoolExecutor(
             nprocs,
-            mp_context=multiprocessing.get_context("fork"),
+            mp_context=multiprocessing.get_context(self.kind.value),
             initializer=functools.partial(_async_compile_initializer, os.getpid()),
         )
         multiprocessing.util.Finalize(None, pool.shutdown, exitpriority=sys.maxsize)
@@ -253,7 +308,9 @@ def submit(self, job_id: int, data: bytes) -> None:
                 self.pool = self._new_pool(self.nprocs, False)
 
     def _submit_inner(self, job_id: int, data: bytes) -> None:
-        future = self.pool.submit(functools.partial(SubprocMain.do_job, data))
+        future = self.pool.submit(
+            functools.partial(SubprocMain.do_job, self.pickler, data)
+        )
 
         def callback(_: Future[Any]) -> None:
             if not self.running:
@@ -262,7 +319,7 @@ def callback(_: Future[Any]) -> None:
                 result = future.result()
             except Exception as e:
                 log.exception("Error in subprocess")
-                result = pickle.dumps(e, pickle.HIGHEST_PROTOCOL)
+                result = self.pickler.dumps(e)
             assert isinstance(result, bytes)
             with self.write_lock:
                 if self.running:
@@ -272,14 +329,15 @@ def callback(_: Future[Any]) -> None:
         future.add_done_callback(callback)
 
     @staticmethod
-    def do_job(data: bytes) -> bytes:
+    def do_job(pickler: SubprocPickler, data: bytes) -> bytes:
         # do the pickle/unpickle in the sub-subproc
-        job = pickle.loads(data)
+        job = typing.cast(Callable[[], object], pickler.loads(data))
+
         try:
             result = job()
         except Exception:
             result = _SubprocExceptionInfo(traceback.format_exc())
-        return pickle.dumps(result, pickle.HIGHEST_PROTOCOL)
+        return pickler.dumps(result)
 
 
 def _warm_process_pool(pool: ProcessPoolExecutor, n: int) -> None: