torchmonarch-nightly 2025.8.1__cp311-cp311-manylinux2014_x86_64.whl → 2025.9.3__cp311-cp311-manylinux2014_x86_64.whl

monarch/_src/actor/allocator.py

@@ -8,7 +8,8 @@
 
 import abc
 import logging
-from typing import Awaitable, final, Optional, TYPE_CHECKING
+from dataclasses import dataclass
+from typing import Dict, final, Literal, Optional
 
 from monarch._rust_bindings.monarch_hyperactor.alloc import (  # @manual=//monarch/monarch_extension:monarch_extension
     Alloc,
@@ -18,21 +19,59 @@ from monarch._rust_bindings.monarch_hyperactor.alloc import (  # @manual=//monar
     RemoteAllocatorBase,
     SimAllocatorBase,
 )
-from monarch._src.actor.future import Future
 
-if TYPE_CHECKING:
-    from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
+from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask, Shared
+from monarch._src.actor.future import DeprecatedNotAFuture, Future
+
 
 ALLOC_LABEL_PROC_MESH_NAME = "procmesh.monarch.meta.com/name"
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+@dataclass
+class AllocHandle(DeprecatedNotAFuture):
+    _hy_alloc: "Shared[Alloc]"
+    _extent: Dict[str, int]
+    _stream_logs: bool
+
+    def reshape(self, extent: Dict[str, int]) -> "AllocHandle":
+        async def task() -> Alloc:
+            alloc = await self._hy_alloc
+            return alloc.reshape(extent)
+
+        return AllocHandle(
+            PythonTask.from_coroutine(task()).spawn(), extent, self._stream_logs
+        )
+
+    @property
+    def initialized(self) -> Future[Literal[True]]:
+        """
+        Future completes with 'True' when the alloc has initialized.
+        Because alloc are remote objects, there is no guarentee that the alloc is
+        still usable after this completes, only that at some point in the past it was usable.
+        """
+
+        async def task() -> Literal[True]:
+            await self._hy_alloc
+            return True
+
+        return Future(coro=task())
+
+    @property
+    def stream_logs(self) -> bool:
+        """
+        Whether to stream stdout/stderr logs from the allocated processes back to the client.
+        The default behavior is determined by the underlying allocator.
+        """
+        return self._stream_logs
+
+
 class AllocateMixin(abc.ABC):
     @abc.abstractmethod
     def allocate_nonblocking(self, spec: AllocSpec) -> "PythonTask[Alloc]": ...
 
-    def allocate(self, spec: AllocSpec) -> "Future[Alloc]":
+    def allocate(self, spec: AllocSpec) -> "AllocHandle":
         """
         Allocate a process according to the provided spec.
 
@@ -42,7 +81,25 @@ class AllocateMixin(abc.ABC):
         Returns:
         - A future that will be fulfilled when the requested allocation is fulfilled.
         """
-        return Future(coro=self.allocate_nonblocking(spec))
+        return AllocHandle(
+            self.allocate_nonblocking(spec).spawn(),
+            spec.extent,
+            self._stream_logs(),
+        )
+
+    @abc.abstractmethod
+    def _stream_logs(self) -> bool:
+        """
+        Whether to stream stdout/stderr logs from the allocated processes back to the client.
+        A common pattern is if the processes are allocated on the same host as the client,
+        then it is not necessary to stream logs back. But if the processes are remotely allocated,
+        it is recommended to stream logs back. It is up to each allocator to decide the default behavior.
+
+        Returns:
+        - A boolean indicating whether to stream logs back to the client.
+        """
+
+        ...
 
 
 @final
@@ -51,6 +108,9 @@ class ProcessAllocator(ProcessAllocatorBase, AllocateMixin):
     An allocator that allocates by spawning local processes.
     """
 
+    def _stream_logs(self) -> bool:
+        return False
+
 
 @final
 class LocalAllocator(LocalAllocatorBase, AllocateMixin):
@@ -58,6 +118,9 @@ class LocalAllocator(LocalAllocatorBase, AllocateMixin):
     An allocator that allocates by spawning actors into the current process.
     """
 
+    def _stream_logs(self) -> bool:
+        return False
+
 
 @final
 class SimAllocator(SimAllocatorBase, AllocateMixin):
@@ -65,6 +128,9 @@ class SimAllocator(SimAllocatorBase, AllocateMixin):
     An allocator that allocates by spawning actors into the current process using simulated channels for transport
     """
 
+    def _stream_logs(self) -> bool:
+        return False
+
 
 class RemoteAllocInitializer(abc.ABC):
     """Subclass-able Python interface for `hyperactor_mesh::alloc::remoteprocess:RemoteProcessAllocInitializer`.
@@ -198,3 +264,6 @@ class RemoteAllocator(RemoteAllocatorBase, AllocateMixin):
     An allocator that allocates by spawning actors on a remote host.
     The remote host must be running hyperactor's remote-process-allocator.
     """
+
+    def _stream_logs(self) -> bool:
+        return True

monarch/_src/actor/bootstrap_main.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 """
 This is the main function for the boostrapping a new process using a ProcessAllocator.
 """
@@ -17,7 +19,7 @@ import sys
 
 # Import torch to avoid import-time races if a spawned actor tries to import torch.
 try:
-    import torch  # @manual
+    import torch  # @manual  # noqa: F401
 except ImportError:
     pass
 
@@ -36,14 +38,15 @@ def invoke_main():
     global bootstrap_main
 
     # TODO: figure out what from worker_main.py we should reproduce here.
-    from monarch._src.actor.telemetry import TracingForwarder
+    from monarch._src.actor.telemetry import TracingForwarder  # noqa
 
     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))
+        # we can stream python logs now; no need to forward them to rust processes
+        pass
         # install opentelemetry tracing
 
     try:
@@ -58,7 +61,7 @@ def invoke_main():
     except Exception as e:
         logging.warning(f"Failed to set up py-spy: {e}")
 
-    from monarch._src.actor.debugger import remote_breakpointhook
+    from monarch._src.actor.debugger.debugger import remote_breakpointhook
 
     sys.breakpointhook = remote_breakpointhook
 

monarch/_src/actor/code_sync/__init__.py

@@ -6,7 +6,9 @@
 
 from monarch._rust_bindings.monarch_extension.code_sync import (  # noqa: F401
     CodeSyncMeshClient,
+    CodeSyncMethod,
     RemoteWorkspace,
+    WorkspaceConfig,
     WorkspaceLocation,
     WorkspaceShape,
 )

monarch/_src/actor/debugger/__init__.py

@@ -0,0 +1,7 @@
+# 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.
+
+# pyre-unsafe