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

monarch/bootstrap_main.py

@@ -4,56 +4,17 @@
 # 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 warnings
 
-import asyncio
-import importlib.resources
-import logging
-import os
-import sys
+warnings.warn(
+    "monarch.bootstrap_main is deprecated, please use from monarch._src.actor.bootstrap_main instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
 
-# Import torch to avoid import-time races if a spawned actor tries to import torch.
-import torch  # noqa[F401]
-
-
-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.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))
-
-    try:
-        with (
-            importlib.resources.path("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}")
-
-    # Start an event loop for PythonActors to use.
-    asyncio.run(main())
+from monarch._src.actor.bootstrap_main import *  # noqa
 
 
 if __name__ == "__main__":
+    # noqa
     invoke_main()  # pragma: no cover

monarch/common/client.py

@@ -37,6 +37,7 @@ from monarch._rust_bindings.monarch_extension.client import (  # @manual=//monar
     LogLevel,
     WorldState,
 )
+from monarch._src.actor.shape import NDSlice
 from monarch.common import messages
 from monarch.common.borrows import Borrow, StorageAliases
 from monarch.common.controller_api import LogMessage, MessageResult, TController
@@ -47,7 +48,6 @@ from monarch.common.invocation import DeviceException, RemoteException, Seq
 from monarch.common.recording import flatten_messages, Recording
 
 from monarch.common.reference import Ref, Referenceable
-from monarch.common.shape import NDSlice
 from monarch.common.stream import StreamRef
 from monarch.common.tensor import Tensor
 from monarch.common.tree import tree_map

monarch/common/controller_api.py

@@ -13,9 +13,10 @@ from monarch._rust_bindings.monarch_extension.client import (  # @manual=//monar
     WorldState,
 )
 
+from monarch._src.actor.shape import NDSlice
+
 from monarch.common.invocation import DeviceException, RemoteException, Seq
 from monarch.common.reference import Ref
-from monarch.common.shape import NDSlice
 from monarch.common.tensor import Tensor
 
 

monarch/common/device_mesh.py

@@ -28,16 +28,16 @@ from typing import (
 
 import monarch.common.messages as messages
 import torch
-from monarch.common.shape import MeshTrait
+from monarch._src.actor.shape import MeshTrait, NDSlice, Shape
 
 from torch.utils._python_dispatch import TorchDispatchMode
 from torch.utils._pytree import tree_map
+from torch.utils.weak import weakref
 
 from ._tensor_to_table import tensor_to_table
 from .context_manager import activate_first_context_manager
 from .messages import Dims
 from .reference import Referenceable
-from .shape import NDSlice, Shape
 from .stream import Stream
 from .tensor import MeshSliceTensor, Tensor
 
@@ -171,6 +171,7 @@ class DeviceMesh(Referenceable, MeshTrait):
         self.exit = lambda: None
         self.ref = None
         self._active_mesh_context = None
+        self._subset_of: Optional[weakref.ReferenceType["DeviceMesh"]] = None
 
     def define_remotely(self):
         if self.ref is None:
@@ -228,8 +229,17 @@ class DeviceMesh(Referenceable, MeshTrait):
     def _new_with_shape(self, shape: Shape) -> "DeviceMesh":
         mesh = DeviceMesh(self.client, shape.ndslice, tuple(shape.labels))
         mesh.exit = self.exit
+        mesh._subset_of = weakref.ref(self)
         return mesh
 
+    def _is_subset_of(self, other: "DeviceMesh") -> bool:
+        p = self
+        while p is not None:
+            if p is other:
+                return True
+            p = None if p._subset_of is None else p._subset_of()
+        return False
+
     def __call__(self, **kwargs) -> "DeviceMesh":
         """
         device_mesh(batch=3) or device_mesh(batch=slice(3, None))

monarch/common/messages.py

@@ -17,18 +17,21 @@ from typing import (
     NamedTuple,
     Optional,
     Protocol,
+    Sequence,
     Tuple,
     TYPE_CHECKING,
 )
 
 from monarch._rust_bindings.monarch_extension import tensor_worker
+from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox
+
+from monarch._src.actor.shape import NDSlice
 from monarch.common.function import ResolvableFromCloudpickle, ResolvableFunction
 from monarch.common.invocation import DeviceException, RemoteException
 from monarch.common.reference import Referenceable
 from monarch.common.tree import flattener
 from pyre_extensions import none_throws
 
-from .shape import NDSlice
 from .tensor_factory import TensorFactory
 
 if TYPE_CHECKING:
@@ -424,6 +427,23 @@ class SendTensor(NamedTuple):
         )
 
 
+class SendResultOfActorCall(NamedTuple):
+    seq: int
+    broker_id: Tuple[str, int]
+    local_state: Sequence[Tensor | tensor_worker.Ref]
+    mutates: List[tensor_worker.Ref]
+    stream: tensor_worker.StreamRef
+
+
+class CallActorMethod(NamedTuple):
+    seq: int
+    result: object
+    broker_id: Tuple[str, int]
+    local_state: Sequence[Tensor | tensor_worker.Ref]
+    mutates: List[tensor_worker.Ref]
+    stream: tensor_worker.StreamRef
+
+
 class SplitComm(NamedTuple):
     dims: Dims
     device_mesh: DeviceMesh

monarch/common/recording.py

@@ -10,9 +10,9 @@ import traceback
 from collections import defaultdict
 from typing import cast, Dict, Generator, List, NamedTuple, Tuple, TYPE_CHECKING, Union
 
-from monarch.common.reference import Ref
+from monarch._src.actor.shape import iter_ranks
 
-from monarch.common.shape import iter_ranks
+from monarch.common.reference import Ref
 
 from monarch.common.tensor import InputChecker
 
@@ -21,8 +21,9 @@ from . import messages
 if TYPE_CHECKING:
     from monarch.common.client import Client
 
+from monarch._src.actor.shape import NDSlice
+
 from .reference import Referenceable
-from .shape import NDSlice
 from .tensor import Tensor
 
 logger = logging.getLogger(__name__)

monarch/common/remote.py

@@ -8,12 +8,12 @@
 
 import functools
 import logging
-import warnings
 
 from logging import Logger
 from typing import (
     Any,
     Callable,
+    cast,
     Dict,
     Generic,
     Literal,
@@ -28,12 +28,18 @@ from typing import (
 import monarch.common.messages as messages
 
 import torch
+from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox
+from monarch._rust_bindings.monarch_hyperactor.shape import Shape
+from monarch._src.actor.actor_mesh import Port, PortTuple
+from monarch._src.actor.endpoint import Extent, Selection
 
-from monarch.common import _coalescing, device_mesh, messages, stream
+from monarch.common import _coalescing, device_mesh, stream
+from monarch.common.future import Future as OldFuture
 
 if TYPE_CHECKING:
     from monarch.common.client import Client
 
+from monarch._src.actor.endpoint import Endpoint
 from monarch.common.device_mesh import RemoteProcessGroup
 from monarch.common.fake import fake_call
 
@@ -49,9 +55,9 @@ from monarch.common.function_caching import (
     TensorGroup,
     TensorPlaceholder,
 )
-from monarch.common.future import Future
 from monarch.common.messages import Dims
-from monarch.common.tensor import dtensor_check, dtensor_dispatch
+
+from monarch.common.tensor import dtensor_check, dtensor_dispatch, InputChecker
 from monarch.common.tree import flatten, tree_map
 from torch import autograd, distributed as dist
 from typing_extensions import ParamSpec
@@ -62,42 +68,96 @@ P = ParamSpec("P")
 R = TypeVar("R")
 T = TypeVar("T")
 
-Propagator = Callable | Literal["mocked", "cached", "inspect"] | None
-
 
-class Remote(Generic[P, R]):
+class Remote(Generic[P, R], Endpoint[P, R]):
     def __init__(self, impl: Any, propagator_arg: Propagator):
+        super().__init__(propagator_arg)
         self._remote_impl = impl
-        self._propagator_arg = propagator_arg
-        self._cache: Optional[dict] = None
+
+    def _call_name(self) -> Any:
+        return self._remote_impl
+
+    def _send(
+        self,
+        args: Tuple[Any, ...],
+        kwargs: Dict[str, Any],
+        port: "Optional[Port]" = None,
+        selection: Selection = "all",
+    ) -> Extent:
+        ambient_mesh = device_mesh._active
+        propagator = self._fetch_propagate
+        rfunction = self._maybe_resolvable
+        # a None rfunction is an optimization for the identity function (lambda x: x)
+        if rfunction is None:
+            preprocess_message = None
+            rfunction = ResolvableFunctionFromPath("ident")
+        else:
+            preprocess_message = rfunction
+        _, dtensors, mutates, tensor_mesh = dtensor_check(
+            propagator, rfunction, args, kwargs, ambient_mesh, stream._active
+        )
+
+        if ambient_mesh is None:
+            raise ValueError(
+                "Calling a 'remote' monarch function requires an active proc_mesh (`with proc_mesh.activate():`)"
+            )
+
+        if not ambient_mesh._is_subset_of(tensor_mesh):
+            raise ValueError(
+                f"The current mesh {ambient_mesh} is not a subset of the mesh on which the tensors being used are defined {tensor_mesh}"
+            )
+
+        client: "Client" = ambient_mesh.client
+        if _coalescing.is_active(client):
+            raise NotImplementedError("NYI: fetching results during a coalescing block")
+        stream_ref = stream._active._to_ref(client)
+
+        fut = (port, ambient_mesh._ndslice)
+
+        ident = client.new_node(mutates, dtensors, cast("OldFuture", fut))
+
+        client.send(
+            ambient_mesh._ndslice,
+            messages.SendValue(
+                ident,
+                None,
+                mutates,
+                preprocess_message,
+                args,
+                kwargs,
+                stream_ref,
+            ),
+        )
+        # we have to ask for status updates
+        # from workers to be sure they have finished
+        # enough work to count this future as finished,
+        # and all potential errors have been reported
+        client._request_status()
+        return Extent(ambient_mesh._labels, ambient_mesh._ndslice.sizes)
+
+    def _port(self, once: bool = False) -> "PortTuple[R]":
+        ambient_mesh = device_mesh._active
+        if ambient_mesh is None:
+            raise ValueError(
+                "FIXME - cannot create a port without an active proc_mesh, because there is not way to create a port without a mailbox"
+            )
+        mesh_controller = getattr(ambient_mesh.client, "_mesh_controller", None)
+        if mesh_controller is None:
+            raise ValueError(
+                "Cannot create raw port objects with an old-style tensor engine controller."
+            )
+        mailbox: Mailbox = mesh_controller._mailbox
+        return PortTuple.create(mailbox, once)
 
     @property
     def _resolvable(self):
         return resolvable_function(self._remote_impl)
 
-    def _propagate(self, args, kwargs, fake_args, fake_kwargs):
-        if self._propagator_arg is None or self._propagator_arg == "cached":
-            if self._cache is None:
-                self._cache = {}
-            return _cached_propagation(self._cache, self._resolvable, args, kwargs)
-        elif self._propagator_arg == "inspect":
-            return None
-        elif self._propagator_arg == "mocked":
-            raise NotImplementedError("mocked propagation")
-        else:
-            return fake_call(self._propagator_arg, *fake_args, **fake_kwargs)
-
-    def _fetch_propagate(self, args, kwargs, fake_args, fake_kwargs):
-        if self._propagator_arg is None:
-            return  # no propgator provided, so we just assume no mutations
-        return self._propagate(args, kwargs, fake_args, fake_kwargs)
-
-    def _pipe_propagate(self, args, kwargs, fake_args, fake_kwargs):
-        if not callable(self._propagator_arg):
-            raise ValueError("Must specify explicit callable for pipe")
-        return self._propagate(args, kwargs, fake_args, fake_kwargs)
+    @property
+    def _maybe_resolvable(self):
+        return None if self._remote_impl is None else self._resolvable
 
-    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+    def _rref(self, args, kwargs):
         return dtensor_dispatch(
             self._resolvable,
             self._propagate,
@@ -107,12 +167,8 @@ class Remote(Generic[P, R]):
             stream._active,
         )
 
-    def call_on_shard_and_fetch(
-        self, *args, shard: Dict[str, int] | None = None, **kwargs
-    ) -> Future[R]:
-        return _call_on_shard_and_fetch(
-            self._resolvable, self._fetch_propagate, *args, shard=shard, **kwargs
-        )
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        return self.rref(*args, **kwargs)
 
 
 # This can't just be Callable because otherwise we are not
@@ -151,14 +207,43 @@ def remote(function: Any = None, *, propagate: Propagator = None) -> Any:
     return Remote(function, propagate)
 
 
-def _call_on_shard_and_fetch(
-    rfunction: ResolvableFunction | None,
-    propagator: Any,
+remote_identity = Remote(None, lambda x: x)
+
+
+def call_on_shard_and_fetch(
+    remote: Endpoint[P, R], *args, shard: Dict[str, int] | None = None, **kwargs
+) -> OldFuture[R]:
+    # We have to flatten the tensors twice: first to discover
+    # which mesh we are working on to shard it, and then again when doing the
+    # dtensor_check in send. This complexity is a consequence of doing
+    # implicit inference of the mesh from the tensors.
+    dtensors, unflatten = flatten((args, kwargs), lambda x: isinstance(x, torch.Tensor))
+    with InputChecker.from_flat_args(
+        remote._call_name(), dtensors, unflatten
+    ) as checker:
+        checker.check_mesh_stream_local(device_mesh._active, stream._active)
+
+        if not hasattr(checker.mesh.client, "_mesh_controller"):
+            return _old_call_on_shard_and_fetch(
+                cast("Remote[P, R]", remote),
+                *args,
+                shard=shard,
+                **kwargs,
+            )
+
+        selected_slice = checker.mesh._process(shard)
+        shard_mesh = checker.mesh._new_with_shape(Shape(["_"], selected_slice))
+        with shard_mesh.activate():
+            return cast("OldFuture[R]", remote.call_one(*args, **kwargs))
+
+
+def _old_call_on_shard_and_fetch(
+    remote_obj: Remote[P, R],
     /,
     *args: object,
     shard: dict[str, int] | None = None,
     **kwargs: object,
-) -> Future:
+) -> OldFuture[R]:
     """
     Call `function` at the coordinates `shard` of the current device mesh, and retrieve the result as a Future.
         function - the remote function to call
@@ -166,6 +251,9 @@ def _call_on_shard_and_fetch(
         shard - a dictionary from mesh dimension name to coordinate of the shard
                 If None, this will fetch from coordinate 0 for all dimensions (useful after all_reduce/all_gather)
     """
+
+    rfunction = remote_obj._maybe_resolvable
+    propagator = remote_obj._fetch_propagate
     ambient_mesh = device_mesh._active
 
     if rfunction is None:
@@ -180,15 +268,9 @@ def _call_on_shard_and_fetch(
     client: "Client" = mesh.client
     if _coalescing.is_active(client):
         raise NotImplementedError("NYI: fetching results during a coalescing block")
+    stream_ref = stream._active._to_ref(client)
     return client.fetch(
-        mesh,
-        stream._active._to_ref(client),
-        shard,
-        preprocess_message,
-        args,
-        kwargs,
-        mutates,
-        dtensors,
+        mesh, stream_ref, shard, preprocess_message, args, kwargs, mutates, dtensors
     )
 
 
@@ -270,8 +352,9 @@ _miss = 0
 _hit = 0
 
 
-def _cached_propagation(_cache, rfunction, args, kwargs):
+def _cached_propagation(_cache, rfunction: ResolvableFunction, args, kwargs):
     tensors, shape_key = hashable_tensor_flatten(args, kwargs)
+    # pyre-ignore
     inputs_group = TensorGroup([t._fake for t in tensors])
     requires_grads = tuple(t.requires_grad for t in tensors)
     key = (shape_key, inputs_group.pattern, requires_grads)
@@ -280,8 +363,8 @@ def _cached_propagation(_cache, rfunction, args, kwargs):
     if key not in _cache:
         _miss += 1
         args_no_pg, kwargs_no_pg = tree_map(_mock_pgs, (args, kwargs))
-        result_with_placeholders, output_pattern = _propagate.call_on_shard_and_fetch(
-            function=rfunction, args=args_no_pg, kwargs=kwargs_no_pg
+        result_with_placeholders, output_pattern = call_on_shard_and_fetch(
+            _propagate, function=rfunction, args=args_no_pg, kwargs=kwargs_no_pg
         ).result()
 
         _, unflatten_result = flatten(

monarch/common/tensor.py

@@ -40,12 +40,13 @@ from .borrows import StorageAliases
 if TYPE_CHECKING:
     from monarch.common.device_mesh import DeviceMesh
 
+from monarch._src.actor.shape import NDSlice
+
 from .fake import fake_call
 from .function import Propagator, ResolvableFunction
 from .invocation import Invocation
 from .messages import Dims
 from .reference import Referenceable
-from .shape import NDSlice
 from .stream import Stream
 from .tree import flatten
 

monarch/controller/backend.py

@@ -13,9 +13,9 @@ import socket
 from abc import ABC, abstractmethod
 from typing import List, NamedTuple, Optional, Sequence, Tuple
 
-from monarch.common import messages
+from monarch._src.actor.shape import iter_ranks, Slices as Ranks
 
-from monarch.common.shape import iter_ranks, Slices as Ranks
+from monarch.common import messages
 from monarch_supervisor import (
     Context,
     FunctionCall,

monarch/controller/controller.py

@@ -19,11 +19,12 @@ from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarc
     ActorId,
 )
 
+from monarch._src.actor.shape import NDSlice
+
 from monarch.common import messages
 from monarch.common.controller_api import LogMessage, MessageResult
 from monarch.common.invocation import DeviceException, Seq
 from monarch.common.reference import Ref
-from monarch.common.shape import NDSlice
 from monarch.common.tensor import Tensor
 from monarch.controller import debugger
 

monarch/controller/rust_backend/controller.py

@@ -29,11 +29,12 @@ from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarc
 )
 
 from monarch._rust_bindings.monarch_messages.debugger import DebuggerAction
+
+from monarch._src.actor.shape import NDSlice
 from monarch.common.controller_api import LogMessage, MessageResult
 from monarch.common.device_mesh import no_mesh
 from monarch.common.invocation import DeviceException, RemoteException
 from monarch.common.messages import SupportsToRustMessage
-from monarch.common.shape import NDSlice
 from monarch.common.tensor import Tensor
 from monarch.controller.debugger import read as debugger_read, write as debugger_write
 from pyre_extensions import none_throws

monarch/fetch.py

@@ -9,13 +9,13 @@
 This is a utility file for fetching a shard of a tensor from remote.
 """
 
-from typing import TypeVar
+from typing import cast, TypeVar
 
 from monarch.common.device_mesh import no_mesh
 
 from monarch.common.future import Future
 
-from monarch.common.remote import _call_on_shard_and_fetch
+from monarch.common.remote import call_on_shard_and_fetch, remote_identity
 
 T = TypeVar("T")
 
@@ -37,9 +37,7 @@ def fetch_shard(
             shard = {}
         shard.update(kwargs)
 
-    return _call_on_shard_and_fetch(
-        None, lambda *args, **kwargs: None, obj, shard=shard
-    )
+    return cast("Future[T]", call_on_shard_and_fetch(remote_identity, obj, shard=shard))
 
 
 def show(obj: T, shard: dict[str, int] | None = None, **kwargs: int) -> object: