fix: forkserver w/ preloaded modules/registry

Author: nkemnitz

.coveragerc

@@ -2,6 +2,7 @@
 omit =
   *_tmp.py
   zetta_utils/log.py
+  zetta_utils/builder/preload/try_load.py
   zetta_utils/cli/task_mgmt.py
   zetta_utils/task_management/subtask_structure.py
   zetta_utils/task_management/segment.py

pyproject.toml

@@ -313,4 +313,7 @@ warn_unused_ignores = true
 known_third_party = "wandb"
 profile = "black"
 skip = ["specs"]
-skip_glob = ["**/__init__.py"]
+skip_glob = [
+  "**/__init__.py",
+  "zetta_utils/builder/preload/*.py"
+]

zetta_utils/__init__.py

@@ -36,88 +36,21 @@
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
 
-def _load_core_modules():
-    """Load core modules that were previously imported at package level."""
-    from . import log, typing, parsing, builder, common, constants
-    from . import geometry, distributions, layer, ng
-
-    # Add builder module suppression now that it's loaded
-    log.add_supress_traceback_module(builder)
-
-
-def load_all_modules():
-    _load_core_modules()
-    load_inference_modules()
-    load_training_modules()
-    from . import task_management
+def load_all_modules():  # pragma: no cover
+    import zetta_utils.builder.preload.all
 
 
 def try_load_train_inference():  # pragma: no cover
-    try:
-        _load_core_modules()
-    except Exception as e:  # pylint: disable=broad-exception-caught
-        logger.exception(e)
-
-    try:
-        load_inference_modules()
-
-    except Exception as e:  # pylint: disable=broad-exception-caught
-        logger.exception(e)
-
-    try:
-        load_training_modules()
-    except Exception as e:  # pylint: disable=broad-exception-caught
-        logger.exception(e)
-
-    try:
-        from . import mazepa_addons
-    except Exception as e:  # pylint: disable=broad-exception-caught
-        logger.exception(e)
+    import zetta_utils.builder.preload.try_load
 
 
 def load_submodules():  # pragma: no cover
     from . import internal
 
 
-def load_inference_modules():
-    _load_core_modules()
-    from . import (
-        augmentations,
-        convnet,
-        mazepa,
-        mazepa_layer_processing,
-        tensor_ops,
-        tensor_typing,
-        tensor_mapping,
-    )
-    from .layer import volumetric
-    from .layer.volumetric import cloudvol
-    from .message_queues import sqs
-
-    from . import mazepa_addons
-    from . import message_queues
-    from . import cloud_management
-
-    load_submodules()
-
-
-def load_training_modules():
-    _load_core_modules()
-    from . import (
-        augmentations,
-        convnet,
-        mazepa,
-        tensor_ops,
-        tensor_typing,
-        training,
-        tensor_mapping,
-    )
-    from .layer import volumetric, db_layer
-    from .layer.db_layer import datastore, firestore
-    from .layer.volumetric import cloudvol
-
-    from . import mazepa_addons
-    from . import message_queues
-    from . import cloud_management
-
-    load_submodules()
+def load_inference_modules():  # pragma: no cover
+    import zetta_utils.builder.preload.inference
+
+
+def load_training_modules():  # pragma: no cover
+    import zetta_utils.builder.preload.training

zetta_utils/builder/preload/__init__.py

@@ -0,0 +1,15 @@
+# pylint: disable=unused-import, wrong-import-position
+"""All module imports."""
+
+import time
+
+from zetta_utils import log
+
+_start = time.perf_counter()
+
+from zetta_utils.builder.preload import inference
+from zetta_utils.builder.preload import training
+from zetta_utils import task_management
+
+_elapsed = time.perf_counter() - _start
+log.get_logger("zetta_utils").debug(f"Preload all modules: {_elapsed:.2f}s")

zetta_utils/builder/preload/all.py

@@ -0,0 +1,15 @@
+# pylint: disable=unused-import, wrong-import-position
+"""Core module imports - shared by all load modes."""
+
+import time
+
+_start = time.perf_counter()
+
+from zetta_utils import log, typing, parsing, builder, common, constants
+from zetta_utils import geometry, distributions, layer, ng
+
+# Add builder module suppression now that it's loaded
+log.add_supress_traceback_module(builder)
+
+_elapsed = time.perf_counter() - _start
+log.get_logger("zetta_utils").debug(f"Preload core modules: {_elapsed:.2f}s")

zetta_utils/builder/preload/core.py

@@ -0,0 +1,33 @@
+# pylint: disable=unused-import, wrong-import-position
+"""Inference module imports."""
+
+import time
+
+from zetta_utils import log
+
+_start = time.perf_counter()
+
+# Import core first
+from zetta_utils.builder.preload import core
+
+from zetta_utils import (
+    augmentations,
+    convnet,
+    mazepa,
+    mazepa_layer_processing,
+    tensor_ops,
+    tensor_typing,
+    tensor_mapping,
+)
+from zetta_utils.layer import volumetric
+from zetta_utils.layer.volumetric import cloudvol
+from zetta_utils.message_queues import sqs
+
+from zetta_utils import mazepa_addons
+from zetta_utils import message_queues
+from zetta_utils import cloud_management
+
+from zetta_utils import internal
+
+_elapsed = time.perf_counter() - _start
+log.get_logger("zetta_utils").debug(f"Preload inference modules: {_elapsed:.2f}s")

zetta_utils/builder/preload/inference.py

@@ -0,0 +1,33 @@
+# pylint: disable=unused-import, wrong-import-position
+"""Training module imports."""
+
+import time
+
+from zetta_utils import log
+
+_start = time.perf_counter()
+
+# Import core first
+from zetta_utils.builder.preload import core
+
+from zetta_utils import (
+    augmentations,
+    convnet,
+    mazepa,
+    tensor_ops,
+    tensor_typing,
+    training,
+    tensor_mapping,
+)
+from zetta_utils.layer import volumetric, db_layer
+from zetta_utils.layer.db_layer import datastore, firestore
+from zetta_utils.layer.volumetric import cloudvol
+
+from zetta_utils import mazepa_addons
+from zetta_utils import message_queues
+from zetta_utils import cloud_management
+
+from zetta_utils import internal
+
+_elapsed = time.perf_counter() - _start
+log.get_logger("zetta_utils").debug(f"Preload training modules: {_elapsed:.2f}s")

zetta_utils/builder/preload/training.py

@@ -0,0 +1,26 @@
+# pylint: disable=unused-import, broad-exception-caught
+"""Try-load variant with error handling."""
+
+from zetta_utils import log
+
+logger = log.get_logger("zetta_utils")
+
+try:
+    from zetta_utils.builder.preload import core
+except Exception as e:
+    logger.exception(e)
+
+try:
+    from zetta_utils.builder.preload import inference
+except Exception as e:
+    logger.exception(e)
+
+try:
+    from zetta_utils.builder.preload import training
+except Exception as e:
+    logger.exception(e)
+
+try:
+    from zetta_utils import mazepa_addons
+except Exception as e:
+    logger.exception(e)

zetta_utils/builder/preload/try_load.py

@@ -1,9 +1,12 @@
+import multiprocessing
 import os
 import pprint
 import subprocess
 import sys
+import threading
+import time
 from tempfile import NamedTemporaryFile
-from typing import Optional
+from typing import Literal, Optional, cast
 
 import click
 
@@ -15,6 +18,35 @@
 
 logger = log.get_logger("zetta_utils")
 
+LoadMode = Literal["all", "inference", "training", "try"]
+
+_PRELOAD_MODULES: dict[LoadMode, str] = {
+    "all": "zetta_utils.builder.preload.all",
+    "inference": "zetta_utils.builder.preload.inference",
+    "training": "zetta_utils.builder.preload.training",
+    "try": "zetta_utils.builder.preload.try_load",
+}
+
+
+def _noop() -> None:
+    pass
+
+
+def initialize_forkserver(load_mode: LoadMode = "all"):
+    """Initialize forkserver with preloaded modules for the given load mode."""
+    preload_module = _PRELOAD_MODULES[load_mode]
+    logger.info(f"Configuring forkserver with preload module: {preload_module}")
+
+    total_start = time.perf_counter()
+    multiprocessing.set_forkserver_preload([preload_module])
+    ctx = multiprocessing.get_context("forkserver")
+    proc = ctx.Process(target=_noop)
+    proc.start()
+    proc.join()
+
+    total_elapsed = time.perf_counter() - total_start
+    logger.info(f"Forkserver initialized in {total_elapsed:.2f}s (mode: {load_mode})")
+
 
 @click.group()
 @click.option("-v", "--verbose", count=True, default=2)
@@ -102,19 +134,27 @@ def run(
 ):
     """Perform ``zetta_utils.builder.build`` action on file contents."""
     ctx = click.get_current_context()
-    load_mode = ctx.obj.get("load_mode", "all") if ctx and ctx.obj else "all"
+    load_mode = cast(LoadMode, ctx.obj.get("load_mode", "all") if ctx and ctx.obj else "all")
 
-    # Load modules first
+    # Start forkserver init in background while main process loads modules
+    forkserver_thread = threading.Thread(
+        target=initialize_forkserver, args=(load_mode,), name="forkserver_init"
+    )
+    forkserver_thread.start()
+
+    # Load modules in main process (runs in parallel with forkserver init)
     if load_mode == "all":
         zetta_utils.load_all_modules()
     elif load_mode == "inference":  # pragma: no cover
         zetta_utils.load_inference_modules()
     elif load_mode == "try":  # pragma: no cover
         zetta_utils.try_load_train_inference()
     else:  # pragma: no cover
-        assert load_mode == "training"
         zetta_utils.load_training_modules()
 
+    # Wait for forkserver to be ready before proceeding
+    forkserver_thread.join()
+
     from zetta_utils import builder, parsing  # pylint: disable=import-outside-toplevel
     from zetta_utils.run import (  # pylint: disable=import-outside-toplevel
         run_ctx_manager,