torchmonarch-nightly 2025.7.1__cp311-cp311-manylinux2014_x86_64.whl → 2025.7.25__cp311-cp311-manylinux2014_x86_64.whl

monarch/{allocator.py → _src/actor/allocator.py}

@@ -10,18 +10,17 @@ import abc
 import logging
 from typing import final, Optional
 
-from monarch import ActorFuture as Future
-from monarch._rust_bindings.hyperactor_extension.alloc import (  # @manual=//monarch/monarch_extension:monarch_extension
+from monarch._rust_bindings.monarch_hyperactor.alloc import (  # @manual=//monarch/monarch_extension:monarch_extension
     Alloc,
     AllocSpec,
-)
-
-from monarch._rust_bindings.monarch_hyperactor.alloc import (  # @manual=//monarch/monarch_extension:monarch_extension
     LocalAllocatorBase,
     ProcessAllocatorBase,
     RemoteAllocatorBase,
+    SimAllocatorBase,
 )
 
+from monarch._src.actor.future import Future
+
 ALLOC_LABEL_PROC_MESH_NAME = "procmesh.monarch.meta.com/name"
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -43,10 +42,7 @@ class ProcessAllocator(ProcessAllocatorBase):
         Returns:
         - A future that will be fulfilled when the requested allocation is fulfilled.
         """
-        return Future(
-            lambda: self.allocate_nonblocking(spec),
-            lambda: self.allocate_blocking(spec),
-        )
+        return Future(impl=lambda: self.allocate_nonblocking(spec), requires_loop=False)
 
 
 @final
@@ -65,10 +61,26 @@ class LocalAllocator(LocalAllocatorBase):
         Returns:
         - A future that will be fulfilled when the requested allocation is fulfilled.
         """
-        return Future(
-            lambda: self.allocate_nonblocking(spec),
-            lambda: self.allocate_blocking(spec),
-        )
+        return Future(impl=lambda: self.allocate_nonblocking(spec), requires_loop=False)
+
+
+@final
+class SimAllocator(SimAllocatorBase):
+    """
+    An allocator that allocates by spawning actors into the current process using simulated channels for transport
+    """
+
+    def allocate(self, spec: AllocSpec) -> Future[Alloc]:
+        """
+        Allocate a process according to the provided spec.
+
+        Arguments:
+        - `spec`: The spec to allocate according to.
+
+        Returns:
+        - A future that will be fulfilled when the requested allocation is fulfilled.
+        """
+        return Future(impl=lambda: self.allocate_nonblocking(spec), requires_loop=False)
 
 
 class RemoteAllocInitializer(abc.ABC):
@@ -214,7 +226,4 @@ class RemoteAllocator(RemoteAllocatorBase):
         Returns:
         - A future that will be fulfilled when the requested allocation is fulfilled.
         """
-        return Future(
-            lambda: self.allocate_nonblocking(spec),
-            lambda: self.allocate_blocking(spec),
-        )
+        return Future(impl=lambda: self.allocate_nonblocking(spec), requires_loop=False)

monarch/_src/actor/bootstrap_main.py

@@ -0,0 +1,73 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+This is the main function for the boostrapping a new process using a ProcessAllocator.
+"""
+
+import asyncio
+import importlib.resources
+import logging
+import multiprocessing
+import os
+import sys
+
+# Import torch to avoid import-time races if a spawned actor tries to import torch.
+try:
+    import torch  # @manual
+except ImportError:
+    pass
+
+
+async def main():
+    from monarch._rust_bindings.monarch_hyperactor.bootstrap import bootstrap_main
+
+    await bootstrap_main()
+
+
+def invoke_main():
+    # if this is invoked with the stdout piped somewhere, then print
+    # changes its buffering behavior. So we default to the standard
+    # behavior of std out as if it were a terminal.
+    sys.stdout.reconfigure(line_buffering=True)
+    global bootstrap_main
+
+    # TODO: figure out what from worker_main.py we should reproduce here.
+    from monarch._src.actor.telemetry import TracingForwarder
+
+    if os.environ.get("MONARCH_ERROR_DURING_BOOTSTRAP_FOR_TESTING") == "1":
+        raise RuntimeError("Error during bootstrap for testing")
+
+    # forward logs to rust tracing. Defaults to on.
+    if os.environ.get("MONARCH_PYTHON_LOG_TRACING", "1") == "1":
+        logging.root.addHandler(TracingForwarder(level=logging.DEBUG))
+        # install opentelemetry tracing
+
+    try:
+        with (
+            importlib.resources.as_file(
+                importlib.resources.files("monarch") / "py-spy"
+            ) as pyspy,
+        ):
+            if pyspy.exists():
+                os.environ["PYSPY_BIN"] = str(pyspy)
+            # fallback to using local py-spy
+    except Exception as e:
+        logging.warning(f"Failed to set up py-spy: {e}")
+
+    from monarch._src.actor.debugger import remote_breakpointhook
+
+    sys.breakpointhook = remote_breakpointhook
+
+    # Start an event loop for PythonActors to use.
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    # Ensure that processes started via `multiprocessing` are spawned, not forked.
+    # forking is a terrible default, see: https://github.com/python/cpython/issues/84559
+    multiprocessing.set_start_method("spawn", force=True)
+    invoke_main()  # pragma: no cover

monarch/{code_sync.py → _src/actor/code_sync/__init__.py}

@@ -5,6 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 from monarch._rust_bindings.monarch_extension.code_sync import (  # noqa: F401
+    CodeSyncMeshClient,
     RemoteWorkspace,
-    RsyncMeshClient,
+    WorkspaceLocation,
+    WorkspaceShape,
 )

monarch/_src/actor/code_sync/auto_reload.py

@@ -0,0 +1,223 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import contextlib
+import dataclasses
+import importlib
+import importlib.abc
+import importlib.util
+import itertools
+import site
+import sys
+import threading
+from pathlib import Path
+from types import ModuleType
+from typing import Dict, List, Optional, Tuple
+
+from monarch._src.actor.actor_mesh import Actor
+from monarch._src.actor.endpoint import endpoint
+
+
+class SysAuditHookGuard(contextlib.AbstractContextManager):
+    """
+    A guard (and context manager), which will unregister an import hook when
+    closed or deleted.
+    """
+
+    def __init__(self, hooks, idx):
+        self._hooks = hooks
+        self._idx = idx
+
+    def close(self):
+        self._hooks.pop(self._idx, None)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+    def __del__(self):
+        self.close()
+
+
+class SysAuditHookMultiplexer:
+    """
+    Multiplexes import hooks to multiple hooks.
+
+    `sys.addaudithook`s can only be added and not removed, so this class provides
+    a global singleton that can be used to multiplex multiple hooks which support
+    removal.
+    """
+
+    def __init__(self):
+        self._idx = itertools.count()
+        self._hooks = {}
+
+    def _callback(self, event, args):
+        for hook in self._hooks.values():
+            hook(event, args)
+
+    def add(self, hook) -> SysAuditHookGuard:
+        idx = next(self._idx)
+        self._hooks[idx] = hook
+        return SysAuditHookGuard(self._hooks, idx)
+
+    _instance_lock = threading.Lock()
+    _instance = None
+
+    @classmethod
+    def singleton(cls):
+        if cls._instance is None:
+            with cls._instance_lock:
+                if cls._instance is None:
+                    cls._instance = SysAuditHookMultiplexer()
+                    sys.addaudithook(cls._instance._callback)
+        return cls._instance
+
+
+@dataclasses.dataclass
+class ThreadLocalState(threading.local):
+    last_import: Optional[str] = None
+
+
+class SysAuditImportHook:
+    """
+    An audit hook that processes and coalesces import/exec events and calls a
+    user-defined callback with the module name and module object which was
+    imported.
+    """
+
+    def __init__(self, callback):
+        self._callback = callback
+        self._state = ThreadLocalState()
+
+    @classmethod
+    def install(cls, callback) -> SysAuditHookGuard:
+        return SysAuditHookMultiplexer.singleton().add(SysAuditImportHook(callback))
+
+    def _py_filename(self, filename: Path) -> Path:
+        if filename.suffix in (".pyc", ".pyo"):
+            return filename.with_suffix(".py")
+        return filename
+
+    def __call__(self, event, args):
+        if event == "import":
+            # While `filename` is specific as an argument to the import event, it's
+            # almost always `None`, so we need to wait for a subsequent exec event
+            # to get the filename.
+            module, _, _, _, _ = args
+            self._state.last_import = module
+        elif event == "exec":
+            module_name = self._state.last_import
+            if module_name is None:
+                return
+            # We always expect an exec right after an import, so we can clear the
+            # last import module name we store.
+            self._state.last_import = None
+            module = sys.modules.get(module_name)
+            if module is None:
+                return
+            if getattr(module, "__file__", None) is None:
+                return
+            (code_obj,) = args
+            if code_obj.co_filename is None:
+                return
+            # code objects store the original source name, not the pyc
+            if self._py_filename(Path(module.__file__)) != Path(code_obj.co_filename):
+                return
+            self._callback(module_name, module)
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Fingerprint:
+    mtime: float
+    size: int
+
+    @classmethod
+    def for_path(cls, path: Path) -> "Fingerprint":
+        stat = path.stat()
+        return Fingerprint(mtime=stat.st_mtime, size=stat.st_size)
+
+
+class AutoReloader:
+    """
+    Track changes to modules and reloads them when they change.
+    """
+
+    def __init__(self, reload=importlib.reload):
+        self._reload = reload
+        self._tracked_modules: Dict[str, Tuple[Path, Fingerprint]] = {}
+        self._track_all_imported()
+
+    def _maybe_track_module(self, name: str, module: ModuleType):
+        filename = getattr(module, "__file__", None)
+        if filename is None:
+            return
+        if filename == "static-extension":
+            return
+        filename = Path(filename)
+
+        # It's rare for modules to have relative path names, but can happen in
+        # weird special situations (e.g. `_ops.py` from `torch.ops`).
+        if not filename.is_absolute():
+            return
+
+        # Ignore builtin modules.
+        if filename.is_relative_to(sys.prefix):
+            for dirpath in site.getsitepackages():
+                if filename.is_relative_to(dirpath):
+                    break
+            else:
+                return
+
+        self._tracked_modules[name] = (
+            filename,
+            Fingerprint.for_path(filename),
+        )
+
+    def _track_all_imported(self):
+        for name, module in sys.modules.items():
+            if module is None:
+                continue
+            self._maybe_track_module(name, module)
+
+    def import_callback(self, name: str, module: ModuleType):
+        """
+        Callback for when a module has been imported.
+        """
+
+        self._maybe_track_module(name, module)
+
+    def reload_changes(self) -> List[str]:
+        """
+        Reload all modules that have changed since they were last imported.
+        """
+
+        reloaded = []
+
+        for module_name, (filename, stored_fingerprint) in list(
+            self._tracked_modules.items()
+        ):
+            fingerprint = Fingerprint.for_path(filename)
+            if fingerprint == stored_fingerprint:
+                continue
+            reloaded.append(module_name)
+            self._reload(sys.modules[module_name])
+            self._tracked_modules[module_name] = (filename, fingerprint)
+
+        return reloaded
+
+
+class AutoReloadActor(Actor):
+    def __init__(self):
+        self._reloader = AutoReloader()
+        self._hook_guard = SysAuditImportHook.install(self._reloader.import_callback)
+
+    @endpoint
+    async def reload(self) -> None:
+        changed = self._reloader.reload_changes()
+        print(f"reloaded modules: {changed}")