torchmonarch-nightly 2025.6.10__cp310-cp310-manylinux2014_x86_64.whl → 2025.6.11__cp310-cp310-manylinux2014_x86_64.whl

monarch/actor_mesh.py

@@ -7,6 +7,7 @@
 import asyncio
 import collections
 import contextvars
+import functools
 import inspect
 
 import itertools
@@ -38,6 +39,7 @@ from typing import (
 
 import monarch
 from monarch import ActorFuture as Future
+from monarch._rust_bindings.hyperactor_extension.telemetry import enter_span, exit_span
 
 from monarch._rust_bindings.monarch_hyperactor.actor import PanicFlag, PythonMessage
 from monarch._rust_bindings.monarch_hyperactor.actor_mesh import PythonActorMesh
@@ -49,6 +51,7 @@ from monarch._rust_bindings.monarch_hyperactor.mailbox import (
 )
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
 from monarch._rust_bindings.monarch_hyperactor.shape import Point as HyPoint, Shape
+
 from monarch.common.pickle_flatten import flatten, unflatten
 from monarch.common.shape import MeshTrait, NDSlice
 
@@ -492,13 +495,29 @@ class _Actor:
                 return None
             else:
                 the_method = getattr(self.instance, message.method)._method
-                result = the_method(self.instance, *args, **kwargs)
+
                 if not inspect.iscoroutinefunction(the_method):
+                    enter_span(
+                        the_method.__module__, message.method, str(ctx.mailbox.actor_id)
+                    )
+                    result = the_method(self.instance, *args, **kwargs)
+                    exit_span()
                     if port is not None:
                         port.send("result", result)
                     return None
 
-                return self.run_async(ctx, self.run_task(port, result, panic_flag))
+                async def instrumented():
+                    enter_span(
+                        the_method.__module__, message.method, str(ctx.mailbox.actor_id)
+                    )
+                    result = await the_method(self.instance, *args, **kwargs)
+                    exit_span()
+                    return result
+
+                return self.run_async(
+                    ctx,
+                    self.run_task(port, instrumented(), panic_flag),
+                )
         except Exception as e:
             traceback.print_exc()
             s = ActorError(e)
@@ -510,7 +529,11 @@ class _Actor:
             else:
                 raise s from None
 
-    async def run_async(self, ctx, coroutine):
+    async def run_async(
+        self,
+        ctx: MonarchContext,
+        coroutine: Coroutine[Any, None, Any],
+    ) -> None:
         _context.set(ctx)
         if self.complete_task is None:
             self.complete_task = asyncio.create_task(self._complete())
@@ -564,6 +587,12 @@ def _unpickle(data: bytes, mailbox: Mailbox) -> Any:
 
 
 class Actor(MeshTrait):
+    @functools.cached_property
+    def logger(cls) -> logging.Logger:
+        lgr = logging.getLogger(cls.__class__.__name__)
+        lgr.setLevel(logging.DEBUG)
+        return lgr
+
     @property
     def _ndslice(self) -> NDSlice:
         raise NotImplementedError(

monarch/allocator.py

@@ -4,6 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
+import abc
 from typing import final
 
 from monarch import ActorFuture as Future
@@ -15,6 +18,7 @@ from monarch._rust_bindings.hyperactor_extension.alloc import (  # @manual=//mon
 from monarch._rust_bindings.monarch_hyperactor.alloc import (  # @manual=//monarch/monarch_extension:monarch_extension
     LocalAllocatorBase,
     ProcessAllocatorBase,
+    RemoteAllocatorBase,
 )
 
 
@@ -60,3 +64,66 @@ class LocalAllocator(LocalAllocatorBase):
             lambda: self.allocate_nonblocking(spec),
             lambda: self.allocate_blocking(spec),
         )
+
+
+class RemoteAllocInitializer(abc.ABC):
+    """Subclass-able Python interface for `hyperactor_mesh::alloc::remoteprocess:RemoteProcessAllocInitializer`.
+
+    NOTE: changes to method signatures of this class must be made to the call-site at
+    `PyRemoteProcessAllocInitializer.py_initialize_alloc()` in `monarch/monarch_hyperactor/src/alloc.rs`
+    """
+
+    @abc.abstractmethod
+    async def initialize_alloc(self) -> list[str]:
+        """
+        Return the addresses of the servers that should be used to allocate processes
+        for the proc mesh. The addresses should be running hyperactor's RemoteProcessAllocator.
+
+        Each address is of the form `{transport}!{addr}(:{port})`.
+        This is the string form of `hyperactor::channel::ChannelAddr` (Rust).
+        For example, `tcp!127.0.0.1:1234`.
+
+        NOTE: Currently, all the addresses must have the same transport type and port
+        NOTE: Although this method is currently called once at the initialization of the Allocator,
+            in the future this method can be called multiple times and should return the current set of
+            addresses that are eligible to handle allocation requests.
+
+        """
+        ...
+
+
+class StaticRemoteAllocInitializer(RemoteAllocInitializer):
+    """
+    Returns the static list of server addresses that this initializer
+    was constructed with on each `initialize_alloc()` call.
+    """
+
+    def __init__(self, *addrs: str) -> None:
+        super().__init__()
+        self.addrs: list[str] = list(addrs)
+
+    async def initialize_alloc(self) -> list[str]:
+        return list(self.addrs)
+
+
+@final
+class RemoteAllocator(RemoteAllocatorBase):
+    """
+    An allocator that allocates by spawning actors on a remote host.
+    The remote host must be running hyperactor's remote-process-allocator.
+    """
+
+    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(
+            lambda: self.allocate_nonblocking(spec),
+            lambda: self.allocate_blocking(spec),
+        )

monarch/bootstrap_main.py

@@ -58,7 +58,7 @@ def invoke_main():
 
     # forward logs to rust tracing. Defaults to on.
     if os.environ.get("MONARCH_PYTHON_LOG_TRACING", "1") == "1":
-        logging.root.addHandler(TracingForwarder())
+        logging.root.addHandler(TracingForwarder(level=logging.DEBUG))
 
     try:
         with (

tests/test_allocator.py

@@ -0,0 +1,216 @@
+# 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-strict
+
+import contextlib
+import importlib.resources
+import math
+import os
+import subprocess
+import sys
+import unittest
+from datetime import timedelta
+from typing import Generator
+
+import cloudpickle
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from monarch._rust_bindings.hyperactor_extension.alloc import (
+    AllocConstraints,
+    AllocSpec,
+)
+
+from monarch._rust_bindings.monarch_hyperactor.channel import (
+    ChannelAddr,
+    ChannelTransport,
+)
+from monarch.actor_mesh import Actor, current_rank, current_size, endpoint, ValueMesh
+
+from monarch.allocator import RemoteAllocator, StaticRemoteAllocInitializer
+from monarch.proc_mesh import ProcMesh
+
+from torch.distributed.elastic.utils.distributed import get_free_port
+
+_100_MILLISECONDS = timedelta(milliseconds=100)
+
+
+class TestActor(Actor):
+    """Silly actor that computes the world size by all-reducing rank-hot tensors"""
+
+    def __init__(self) -> None:
+        self.rank: int = current_rank().rank
+        self.world_size: int = math.prod(current_size().values())
+
+    @endpoint
+    async def compute_world_size(self, master_addr: str, master_port: int) -> int:
+        os.environ["MASTER_ADDR"] = master_addr
+        os.environ["MASTER_PORT"] = str(master_port)
+        dist.init_process_group("gloo", rank=self.rank, world_size=self.world_size)
+
+        try:
+            t = F.one_hot(torch.tensor(self.rank), num_classes=dist.get_world_size())
+            dist.all_reduce(t)
+            return int(torch.sum(t).item())
+        finally:
+            dist.destroy_process_group()
+
+
+@contextlib.contextmanager
+def remote_process_allocator() -> Generator[str, None, None]:
+    with importlib.resources.path(__package__, "") as package_path:
+        addr = ChannelAddr.any(ChannelTransport.Unix)
+
+        process_allocator = subprocess.Popen(
+            args=[
+                "process_allocator",
+                f"--addr={addr}",
+            ],
+            env={
+                # prefix PATH with this test module's directory to
+                # give 'process_allocator' and 'monarch_bootstrap' binary resources
+                # in this test module's directory precedence over the installed ones
+                # useful in BUCK where these binaries are added as 'resources' of this test target
+                "PATH": f"{package_path}:{os.getenv('PATH', '')}",
+                "RUST_LOG": "debug",
+            },
+        )
+        try:
+            yield addr
+        finally:
+            process_allocator.terminate()
+            try:
+                five_seconds = 5
+                process_allocator.wait(timeout=five_seconds)
+            except subprocess.TimeoutExpired:
+                process_allocator.kill()
+
+
+class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        cloudpickle.register_pickle_by_value(sys.modules[TestActor.__module__])
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cloudpickle.unregister_pickle_by_value(sys.modules[TestActor.__module__])
+
+    def assert_computed_world_size(
+        self, computed: ValueMesh[int], expected_world_size: int
+    ) -> None:
+        expected_world_sizes = {
+            rank: expected_world_size for rank in range(0, expected_world_size)
+        }
+        computed_world_sizes = {p.rank: v for p, v in list(computed.flatten("rank"))}
+        self.assertDictEqual(expected_world_sizes, computed_world_sizes)
+
+    async def test_call_allocate_twice(self) -> None:
+        class DeletingAllocInitializer(StaticRemoteAllocInitializer):
+            """test initializer that removes the last address from the list each time initialize_alloc() is called
+            used to test that the state of the initializer is preserved across calls to allocate()
+            """
+
+            async def initialize_alloc(self) -> list[str]:
+                alloc = await super().initialize_alloc()
+                self.addrs.pop(-1)
+                return alloc
+
+        with remote_process_allocator() as host1, remote_process_allocator() as host2:
+            initializer = DeletingAllocInitializer(host1, host2)
+
+            allocator = RemoteAllocator(
+                world_id="test_remote_allocator",
+                initializer=initializer,
+                heartbeat_interval=_100_MILLISECONDS,
+            )
+
+            spec = AllocSpec(AllocConstraints(), host=1, gpu=1)
+
+            await allocator.allocate(spec)
+            self.assertEqual([host1], initializer.addrs)
+
+            await allocator.allocate(spec)
+            self.assertEqual([], initializer.addrs)
+
+    async def test_throws_when_initializer_returns_empty_addrs(self) -> None:
+        class EmptyAllocInitializer(StaticRemoteAllocInitializer):
+            """test initializer that returns an empty list of addresses"""
+
+            async def initialize_alloc(self) -> list[str]:
+                return []
+
+        empty_initializer = EmptyAllocInitializer()
+        with self.assertRaisesRegex(
+            RuntimeError, r"initializer must return non-empty list of addresses"
+        ):
+            allocator = RemoteAllocator(
+                world_id="test_remote_allocator",
+                initializer=empty_initializer,
+                heartbeat_interval=_100_MILLISECONDS,
+            )
+            await allocator.allocate(AllocSpec(AllocConstraints(), host=1, gpu=1))
+
+    async def test_allocate_2d_mesh(self) -> None:
+        hosts = 2
+        gpus = 4
+        world_size = hosts * gpus
+        spec = AllocSpec(AllocConstraints(), host=hosts, gpu=gpus)
+
+        # create 2x process-allocators (on their own bind addresses) to simulate 2 hosts
+        with remote_process_allocator() as host1, remote_process_allocator() as host2:
+            allocator = RemoteAllocator(
+                world_id="test_remote_allocator",
+                initializer=StaticRemoteAllocInitializer(host1, host2),
+                heartbeat_interval=_100_MILLISECONDS,
+            )
+            alloc = await allocator.allocate(spec)
+            proc_mesh = await ProcMesh.from_alloc(alloc)
+            actor = await proc_mesh.spawn("test_actor", TestActor)
+
+            values = await actor.compute_world_size.call(
+                master_addr="::",
+                master_port=get_free_port(),
+            )
+
+            self.assert_computed_world_size(values, world_size)
+
+    async def test_stacked_1d_meshes(self) -> None:
+        # create two stacked actor meshes on the same host
+        # each actor mesh running on separate process-allocators
+
+        with remote_process_allocator() as host1_a, remote_process_allocator() as host1_b:
+            allocator_a = RemoteAllocator(
+                world_id="a",
+                initializer=StaticRemoteAllocInitializer(host1_a),
+                heartbeat_interval=_100_MILLISECONDS,
+            )
+            allocator_b = RemoteAllocator(
+                world_id="b",
+                initializer=StaticRemoteAllocInitializer(host1_b),
+                heartbeat_interval=_100_MILLISECONDS,
+            )
+
+            spec_a = AllocSpec(AllocConstraints(), host=1, gpu=2)
+            spec_b = AllocSpec(AllocConstraints(), host=1, gpu=6)
+
+            proc_mesh_a = await ProcMesh.from_alloc(await allocator_a.allocate(spec_a))
+            proc_mesh_b = await ProcMesh.from_alloc(await allocator_b.allocate(spec_b))
+
+            actor_a = await proc_mesh_a.spawn("actor_a", TestActor)
+            actor_b = await proc_mesh_b.spawn("actor_b", TestActor)
+
+            results_a = await actor_a.compute_world_size.call(
+                master_addr="::", master_port=get_free_port()
+            )
+            results_b = await actor_b.compute_world_size.call(
+                master_addr="::", master_port=get_free_port()
+            )
+
+            self.assert_computed_world_size(results_a, 2)  # a is a 1x2 mesh
+            self.assert_computed_world_size(results_b, 6)  # b is a 1x6 mesh

{torchmonarch_nightly-2025.6.10.dist-info → torchmonarch_nightly-2025.6.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchmonarch-nightly
-Version: 2025.6.10
+Version: 2025.6.11
 Summary: Monarch: Single controller library
 Author: Meta
 Author-email: oncall+monarch@xmail.facebook.com

{torchmonarch_nightly-2025.6.10.dist-info → torchmonarch_nightly-2025.6.11.dist-info}/RECORD

@@ -1,16 +1,16 @@
 monarch/__init__.py,sha256=iUvWHc0-7Q2tovRoRxOIiA3TsefMXCbWl-jEfQ2djew,6897
-monarch/_rust_bindings.so,sha256=0-svsKnUJboaOBd5i-LOfpHiRRAgVLX_1Hq_YYREQi8,39756680
+monarch/_rust_bindings.so,sha256=g2tlum6iqfdR4KRkVhp_BwUmlz0tYUSITNVaJjSNitE,40645720
 monarch/_testing.py,sha256=MN8DK1e-wzV0-R_nFW1b_7-O5oKfWvZ12BMGD4Z7PQk,6755
-monarch/actor_mesh.py,sha256=AKdjPg3FM6Yt35uFPBnP7fNVEu6busu5BXVWLwjU2A4,23000
-monarch/allocator.py,sha256=_2DKFP9pSD33zDgH7xZJC8Tq7BQrCeQEUmMB7_xCT0Y,1784
-monarch/bootstrap_main.py,sha256=SYTOz-pTXiJNk78PPD5HAOJDSb8t2JfitRWdmWB3ogo,2559
+monarch/actor_mesh.py,sha256=4I8xp_XIM6KZJY_jXVjJ8tPW2l1J4a6ZhrknU7zKbAk,23947
+monarch/allocator.py,sha256=ylvYTf31o-PT385cYJPhi17uNbC4yl_RAraqD0fVe4g,4112
+monarch/bootstrap_main.py,sha256=EYaTMA1lxy2213L_04drTKlJvZQjzNdD3jeUHiqSBJc,2578
 monarch/cached_remote_function.py,sha256=kYdB6r4OHx_T_uX4q3tCNcp1t2DJwF8tPTIahUiT2pU,8785
 monarch/fetch.py,sha256=61jxo7sx4QNUTkc0_rF5NaJROen4tKbAaiIjrXWLOvg,1705
 monarch/future.py,sha256=lcdFEe7m1shYPPuvZ1RkS6JUIChEKGBWe3v7x_nu4Hg,731
 monarch/gradient_generator.py,sha256=Rl3dmXGceTdCc1mYBg2JciR88ywGPnW7TVkL86KwqEA,6366
 monarch/memory.py,sha256=ol86dBhFAJqg78iF25-BuK0wuwj1onR8FIioZ_B0gjw,1377
 monarch/mesh_controller.py,sha256=A3G8Z5S0w3mjCVI2r6YGM6K3BUs3ZHU8PFo6kCaYTU4,8615
-monarch/monarch_controller,sha256=Q1eR_EVJqDQLrJZ_6p1ldxVDAU1OmN5lSSuctDcaAFY,20396832
+monarch/monarch_controller,sha256=41B7zLv7M7_CSmChN5bfvVrygi2VeBhMDcNQXlnbVZU,20394376
 monarch/notebook.py,sha256=zu9MKDFKf1-rCM2TqFSRJjMBeiWuKcJSyUFLvoZRQzs,25949
 monarch/opaque_module.py,sha256=oajOu_WD1hD4hxE8HDdO-tvWY7KDHWd7VaAhJEa5L2I,10446
 monarch/opaque_object.py,sha256=IVpll4pyuKZMo_EnPh4s0qnx8RlAcJrJ1yoLX6E75wQ,2782
@@ -132,6 +132,7 @@ tests/error_test_binary.py,sha256=64H-ucdkQ2i7GD8sidStl227cOy7gyeqvO4kTm1y7Ic,48
 tests/sleep_binary.py,sha256=XfLYaAfwm9xgzM-svs8fhAeFhwYIg6SyVEnx4e6wbUw,1009
 tests/test_actor_error.py,sha256=z3Sf4lteUggTryPLOhRKJ55v0MwVK3a7QN7-U2U9iJg,7484
 tests/test_alloc.py,sha256=D6DdQbtOZEvvnnc7LV-WyWFMk0Xb77eblH6Oz90zJTA,745
+tests/test_allocator.py,sha256=dqQbQyOjOX3JgnHIPT0iawT0wMeFztbLCYjK2tl8GcI,8149
 tests/test_coalescing.py,sha256=-KtAWzTaeXbyzltplfojavx0iFeeZnvej-tFTlu2p5k,15616
 tests/test_controller.py,sha256=yxuVp2DG3TDKJlwuE3cFm9dbWMlbrYtG1uHfvVWRYbw,30935
 tests/test_device_mesh.py,sha256=DrbezYOM0thfP9MgLXb5-F0VoLOmSz5GR0GwjR_3bE4,5290
@@ -150,9 +151,9 @@ tests/simulator/test_profiling.py,sha256=TGYCfzTLdkpIwnOuO6KApprmrgPIRQe60KRX3wk
 tests/simulator/test_simulator.py,sha256=LO8lA0ssY-OGEBL5ipEu74f97Y765TEwfUOv-DtIptM,14568
 tests/simulator/test_task.py,sha256=ipqBDuDAysuo1xOB9S5psaFvwe6VATD43IovCTSs0t4,2327
 tests/simulator/test_worker.py,sha256=QrWWIJ3HDgDLkBPRc2mwYPlOQoXQcj1qRfc0WUfKkFY,3507
-torchmonarch_nightly-2025.6.10.dist-info/licenses/LICENSE,sha256=e0Eotbf_rHOYPuEUlppIbvwy4SN98CZnl_hqwvbDA4Q,1530
-torchmonarch_nightly-2025.6.10.dist-info/METADATA,sha256=DR1GtSFqtqsjhKWi38uGcvhw2p3ycHYSOwDmsErwLj0,2772
-torchmonarch_nightly-2025.6.10.dist-info/WHEEL,sha256=_wZSFk0d90K9wOBp8Q-UGxshyiJ987JoPiyUBNC6VLk,104
-torchmonarch_nightly-2025.6.10.dist-info/entry_points.txt,sha256=sqfQ16oZqjEvttUI-uj9BBXIIE6jt05bYFSmy-2hyXI,106
-torchmonarch_nightly-2025.6.10.dist-info/top_level.txt,sha256=E-ZssZzyM17glpVrh-S9--qJ-w9p2EjuYOuNw9tQ4Eg,33
-torchmonarch_nightly-2025.6.10.dist-info/RECORD,,
+torchmonarch_nightly-2025.6.11.dist-info/licenses/LICENSE,sha256=e0Eotbf_rHOYPuEUlppIbvwy4SN98CZnl_hqwvbDA4Q,1530
+torchmonarch_nightly-2025.6.11.dist-info/METADATA,sha256=SCdAxETtVZ5ESzbLepOp6mf1L4G-HSYVkjdRFT7D0kg,2772
+torchmonarch_nightly-2025.6.11.dist-info/WHEEL,sha256=_wZSFk0d90K9wOBp8Q-UGxshyiJ987JoPiyUBNC6VLk,104
+torchmonarch_nightly-2025.6.11.dist-info/entry_points.txt,sha256=sqfQ16oZqjEvttUI-uj9BBXIIE6jt05bYFSmy-2hyXI,106
+torchmonarch_nightly-2025.6.11.dist-info/top_level.txt,sha256=E-ZssZzyM17glpVrh-S9--qJ-w9p2EjuYOuNw9tQ4Eg,33
+torchmonarch_nightly-2025.6.11.dist-info/RECORD,,