torchmonarch-nightly 2025.8.2__cp312-cp312-manylinux2014_x86_64.whl → 2025.9.3__cp312-cp312-manylinux2014_x86_64.whl

monarch/_src/actor/python_extension_methods.py

@@ -0,0 +1,90 @@
+# 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 importlib
+
+from typing import cast, Type, TypeVar
+
+
+T = TypeVar("T")
+
+
+class PatchRustClass:
+    def __init__(self, rust_class: Type):
+        self.rust_class = rust_class
+
+    def __call__(self, python_class: Type[T]) -> Type[T]:
+        rust_name = f"{self.rust_class.__module__}.{self.rust_class.__name__}"
+        python_name = f"{python_class.__module__}.{python_class.__name__}"
+        if rust_name != python_name:
+            raise ValueError(f"mismatched type names {rust_name} != {python_name}")
+        for name, implementation in python_class.__dict__.items():
+            if hasattr(self.rust_class, name):
+                # do not patch in the stub methods that
+                # are already defined by the rust implementation
+                continue
+            if not callable(implementation) and not isinstance(
+                implementation, property
+            ):
+                continue
+            setattr(self.rust_class, name, implementation)
+        return cast(Type[T], self.rust_class)
+
+
+def rust_struct(name: str) -> PatchRustClass:
+    """
+    When we bind a rust struct into Python, it is sometimes faster to implement
+    parts of the desired Python API in Python. It is also easier to understand
+    what the class does in terms of these methods.
+
+    We also want to avoid having to wrap rust objects in another layer of python objects
+    because:
+    * wrappers double the python overhead
+    * it is easy to confuse which level of wrappers and API takes, especially
+      along the python<->rust boundary.
+
+
+    To avoid wrappers we first define the class in pyo3. Lets say we add a class
+    monarch_hyperactor::actor_mesh::TestClass which we will want to extend with python methods in
+    the monarch/actor/_src/actor_mesh.py. In rust we will define the class as
+
+         #[pyclass(name = "TestClass", module = "monarch._src.actor_mesh")]
+         struct TestClass {}
+         #[pymethods]
+         impl TestClass {
+            fn hello(&self) {
+                println!("hello");
+            }
+         }
+
+    Then rather than writing typing stubs in a pyi file we write the stub code directly in
+    monarch/actor/_src/actor_mesh.py along with any helper methods:
+
+        @rust_struct("monarch_hyperactor::actor_mesh::TestClass")
+        class TestClass:
+            def hello(self) -> None:
+                ...
+            def hello_world(self) -> None:
+                self.hello()
+                print("world")
+
+    This class annotation then merges the python extension methods with the rust
+    class implementation. Any rust code that returns the TestClass will have the `hello_world`
+    extension method attached. Python typechecking always things TestClass is the python code,
+    so typing works.
+
+    It is ok to have the pyclass module not match where it is defined because (1) we patch it into the right place
+    to make sure pickling works, and (2) the rust_struct annotation points directly to where to find the rust code,
+    and will be discovered by goto line in the IDE.
+    """
+
+    *modules, name = name.split("::")
+    module_name = ".".join(modules)
+    module = importlib.import_module(f"monarch._rust_bindings.{module_name}")
+
+    rust_class = getattr(module, name)
+
+    return PatchRustClass(rust_class)

monarch/_src/actor/shape.py

@@ -10,7 +10,7 @@ from abc import ABC, abstractmethod
 
 from typing import Dict, Generator, Sequence, Tuple, Union
 
-from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice
+from monarch._rust_bindings.monarch_hyperactor.shape import Extent, Shape, Slice
 
 from typing_extensions import Self
 
@@ -224,5 +224,12 @@ class MeshTrait(ABC):
     def sizes(self) -> dict[str, int]:
         return dict(zip(self._labels, self._ndslice.sizes))
 
+    @property
+    def extent(self) -> "Extent":
+        return Extent(self._labels, self._ndslice.sizes)
+
+    def __len__(self) -> int:
+        return len(self._ndslice)
+
 
 __all__ = ["NDSlice", "Shape", "MeshTrait"]

monarch/_src/actor/source_loader.py

@@ -0,0 +1,45 @@
+# 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
+import functools
+import importlib
+import importlib.abc
+import linecache
+
+from monarch._src.actor.actor_mesh import _context, Actor
+from monarch._src.actor.endpoint import endpoint
+from monarch._src.actor.proc_mesh import get_or_spawn_controller
+from monarch._src.actor.sync_state import fake_sync_state
+
+
+class SourceLoaderController(Actor):
+    @endpoint
+    def get_source(self, filename: str) -> str:
+        return "".join(linecache.getlines(filename))
+
+
+@functools.cache
+def source_loader_controller() -> SourceLoaderController:
+    with fake_sync_state():
+        return get_or_spawn_controller("source_loader", SourceLoaderController).get()
+
+
+@functools.cache
+def load_remote_source(filename: str) -> str:
+    with fake_sync_state():
+        return source_loader_controller().get_source.call_one(filename).get()
+
+
+class RemoteImportLoader(importlib.abc.Loader):
+    def __init__(self, filename: str):
+        self._filename = filename
+
+    def get_source(self, _module_name: str) -> str:
+        if _context.get(None) is not None:
+            return load_remote_source(self._filename)
+        else:
+            raise ImportError(f"could not get source for {self._filename}")

monarch/_src/actor/telemetry/__init__.py

@@ -8,12 +8,184 @@
 
 
 import logging
+import warnings
+from typing import Optional, Sequence
+
+import opentelemetry.metrics as metrics  # @manual=fbsource//third-party/pypi/opentelemetry-api:opentelemetry-api
+import opentelemetry.trace as trace  # @manual=fbsource//third-party/pypi/opentelemetry-api:opentelemetry-api
 
 from monarch._rust_bindings.monarch_hyperactor.telemetry import (  # @manual=//monarch/monarch_extension:monarch_extension
     forward_to_tracing,
+    PyCounter,
+    PyHistogram,
+    PyUpDownCounter,
 )
+from monarch._src.actor.telemetry.rust_span_tracing import RustTracerProvider
+from opentelemetry.context import Context
+from opentelemetry.metrics import CallbackT
+from opentelemetry.util.types import Attributes
 
 
 class TracingForwarder(logging.Handler):
     def emit(self, record: logging.LogRecord) -> None:
         forward_to_tracing(record)
+
+
+class Counter(metrics.Counter):
+    inner: PyCounter
+
+    def __init__(self, name: str) -> None:
+        super().__init__(name)
+        self.inner = PyCounter(name)
+
+    def add(
+        self,
+        amount: int | float,
+        attributes: Optional[Attributes] = None,
+        context: Optional[Context] = None,
+    ) -> None:
+        return self.inner.add(int(amount))
+
+
+class UpDownCounter(metrics.UpDownCounter):
+    inner: PyUpDownCounter
+
+    def __init__(self, name: str) -> None:
+        super().__init__(name)
+        self.inner = PyUpDownCounter(name)
+
+    def add(
+        self,
+        amount: int | float,
+        attributes: Optional[Attributes] = None,
+        context: Optional[Context] = None,
+    ) -> None:
+        self.inner.add(int(amount))
+
+
+class Histogram(metrics.Histogram):
+    inner: PyHistogram
+
+    def __init__(self, name: str) -> None:
+        super().__init__(name)
+        self.inner = PyHistogram(name)
+
+    def record(
+        self,
+        amount: int | float,
+        attributes: Optional[Attributes] = None,
+        context: Optional[Context] = None,
+    ) -> None:
+        self.inner.record(amount)
+
+
+class Meter(metrics.Meter):
+    def create_counter(
+        self,
+        name: str,
+        unit: str = "",
+        description: str = "",
+    ) -> metrics.Counter:
+        return Counter(name)
+
+    def create_up_down_counter(
+        self,
+        name: str,
+        unit: str = "",
+        description: str = "",
+    ) -> metrics.UpDownCounter:
+        return UpDownCounter(name)
+
+    def create_observable_counter(
+        self,
+        name: str,
+        callbacks: Optional[Sequence[CallbackT]] = None,
+        unit: str = "",
+        description: str = "",
+    ) -> metrics.ObservableCounter:
+        raise NotImplementedError()
+
+    def create_histogram(
+        self,
+        name: str,
+        unit: str = "",
+        description: str = "",
+        *,
+        explicit_bucket_boundaries_advisory: Optional[Sequence[float]] = None,
+    ) -> metrics.Histogram:
+        return Histogram(name)
+
+    def create_gauge(  # type: ignore # pylint: disable=no-self-use
+        self,
+        name: str,
+        unit: str = "",
+        description: str = "",
+    ) -> metrics._Gauge:  # pyright: ignore[reportReturnType]
+        warnings.warn(
+            "create_gauge() is not implemented and will be a no-op", stacklevel=2
+        )
+        raise NotImplementedError()
+
+    def create_observable_gauge(
+        self,
+        name: str,
+        callbacks: Optional[Sequence[CallbackT]] = None,
+        unit: str = "",
+        description: str = "",
+    ) -> metrics.ObservableGauge:
+        raise NotImplementedError()
+
+    def create_observable_up_down_counter(
+        self,
+        name: str,
+        callbacks: Optional[Sequence[CallbackT]] = None,
+        unit: str = "",
+        description: str = "",
+    ) -> metrics.ObservableUpDownCounter:
+        raise NotImplementedError()
+
+
+class MeterProvider(metrics.MeterProvider):
+    def get_meter(
+        self,
+        name: str,
+        version: Optional[str] = None,
+        schema_url: Optional[str] = None,
+        attributes: Optional[Attributes] = None,
+    ) -> metrics.Meter:
+        return Meter(name, version, schema_url)
+
+
+def get_monarch_tracer() -> trace.Tracer:
+    """
+    Creates and returns a Monarch python tracer that logs to the Rust telemetry system.
+
+    Returns:
+        Tracer: A configured OpenTelemetry tracer for Monarch.
+
+    Usage:
+        tracer = get_monarch_tracer()
+        with tracer.start_as_current_span("span_name") as span:
+            # code here
+    """
+    install()
+    return trace.get_tracer("monarch.python.tracer")
+
+
+_INSTALLED = False
+
+METER: metrics.Meter = metrics.get_meter("monarch")
+
+
+def install() -> None:
+    global _INSTALLED
+    if _INSTALLED:
+        return
+
+    provider = RustTracerProvider()
+    trace.set_tracer_provider(provider)
+    metrics.set_meter_provider(MeterProvider())
+
+    global METER
+    METER = metrics.get_meter("monarch")
+    _INSTALLED = True

monarch/_src/actor/telemetry/rust_span_tracing.py

@@ -11,28 +11,23 @@ from contextlib import contextmanager
 from typing import Iterator, Mapping, Optional, Union
 
 import opentelemetry.util.types as types  # @manual=fbsource//third-party/pypi/opentelemetry-api:opentelemetry-api
-
 from monarch._rust_bindings.monarch_hyperactor.telemetry import (
     get_current_span_id,
     PySpan,
 )
-
 from opentelemetry import (  # @manual=fbsource//third-party/pypi/opentelemetry-api:opentelemetry-api
     trace,
 )
-from opentelemetry.trace import Tracer
 from opentelemetry.trace.status import Status, StatusCode
-from pyre_extensions import override
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
 class SpanWrapper(trace.Span):
-    def __init__(self, name: str) -> None:
+    def __init__(self, name: str, actor_id: Optional[str]) -> None:
         super().__init__()
-        self._span: PySpan | None = PySpan(name)
+        self._span: PySpan | None = PySpan(name, actor_id)
 
-    @override
     def end(self, end_time: Optional[int] = None) -> None:
         # since PySpan is not sendable, we need to make sure it is deallocated on this thread so it doesn't log warnings.
         s = self._span
@@ -94,7 +89,8 @@ class RustTracer(trace.Tracer):
         record_exception: bool = True,
         set_status_on_exception: bool = True,
     ) -> trace.Span:
-        return SpanWrapper(name)
+        actor_id = str(attributes.get("actor_id")) if attributes else None
+        return SpanWrapper(name, actor_id)
 
     @contextmanager
     # pyre-fixme[15]: `start_as_current_span` overrides method defined in `Tracer`
@@ -111,14 +107,14 @@ class RustTracer(trace.Tracer):
         set_status_on_exception: bool = True,
         end_on_exit: bool = True,
     ) -> Iterator[trace.Span]:
-        with SpanWrapper(name) as s:
+        actor_id = str(attributes.get("actor_id")) if attributes else None
+        with SpanWrapper(name, actor_id) as s:
             with trace.use_span(s):
                 yield s
             del s
 
 
 class RustTracerProvider(trace.TracerProvider):
-    @override
     def get_tracer(
         self,
         instrumenting_module_name: str,
@@ -128,32 +124,3 @@ class RustTracerProvider(trace.TracerProvider):
         **kwargs: object,
     ) -> trace.Tracer:
         return RustTracer()
-
-
-def get_monarch_tracer() -> Tracer:
-    """
-    Creates and returns a Monarch python tracer that logs to the Rust telemetry system.
-
-    Returns:
-        Tracer: A configured OpenTelemetry tracer for Monarch.
-
-    Usage:
-        tracer = get_monarch_tracer()
-        with tracer.start_as_current_span("span_name") as span:
-            # code here
-    """
-    install()
-    return trace.get_tracer("monarch.python.tracer")
-
-
-_INSTALLED = False
-
-
-def install() -> None:
-    global _INSTALLED
-    if _INSTALLED:
-        return
-
-    provider = RustTracerProvider()
-    trace.set_tracer_provider(provider)
-    _INSTALLED = True

monarch/_src/debug_cli/__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

monarch/_src/debug_cli/debug_cli.py

@@ -0,0 +1,43 @@
+# 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
+import argparse
+import logging
+import subprocess
+
+from monarch._src.actor.debugger.debugger import (
+    _get_debug_server_host,
+    _get_debug_server_port,
+)
+
+
+def run():
+    parser = argparse.ArgumentParser(description="Monarch Debug CLI")
+    parser.add_argument(
+        "--host",
+        type=str,
+        default=_get_debug_server_host(),
+        help="Hostname where the debug server is running",
+    )
+    parser.add_argument(
+        "--port",
+        type=str,
+        default=_get_debug_server_port(),
+        help="Port that the debug server is listening on",
+    )
+    args = parser.parse_args()
+
+    for cmd in ["ncat", "nc", "netcat"]:
+        try:
+            subprocess.run([cmd, f"{args.host}", f"{args.port}"], check=True)
+            return
+        except FileNotFoundError:
+            pass
+
+    logging.error(
+        "Could not find a suitable netcat binary. Please install one and try again."
+    )

monarch/_src/tensor_engine/rdma.py

@@ -4,19 +4,28 @@
 # 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
+import asyncio
+import functools
 import logging
 import warnings
 from typing import Optional
 
 import torch
+from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask, Shared
 
 try:
-    from monarch._rust_bindings.rdma import _RdmaBuffer
+    from monarch._rust_bindings.rdma import _RdmaBuffer, _RdmaManager
 except ImportError as e:
     logging.error("RDMA is not available: {}".format(e))
     raise e
-from monarch._src.actor.actor_mesh import MonarchContext
+from typing import Dict
+
+from monarch._src.actor.actor_mesh import Actor, context
+from monarch._src.actor.endpoint import endpoint
 from monarch._src.actor.future import Future
+from monarch._src.actor.proc_mesh import get_or_spawn_controller, ProcMesh
+from pyre_extensions import none_throws
 
 
 # RDMARead/WriteTransferWarnings are warnings that are only printed once per process.
@@ -37,6 +46,44 @@ def is_available():
     return _RdmaBuffer.rdma_supported()
 
 
+class RdmaController(Actor):
+    def __init__(self) -> None:
+        self._managers: Dict[ProcMesh, _RdmaManager] = {}
+        self._lock = asyncio.Lock()
+
+    @endpoint
+    async def init_rdma_on_mesh(self, proc_mesh: ProcMesh) -> None:
+        if not _RdmaBuffer.rdma_supported():
+            raise RuntimeError(
+                "Cannot spawn _RdmaManager because RDMA is not supported on this machine"
+            )
+
+        if proc_mesh in self._managers:
+            return
+
+        async with self._lock:
+            if proc_mesh not in self._managers:
+                self._managers[proc_mesh] = none_throws(
+                    await Future(
+                        coro=_RdmaManager.create_rdma_manager_nonblocking(
+                            await Future(coro=proc_mesh._proc_mesh.task())
+                        )
+                    )
+                )
+
+
+# Cached so that we don't have to call out to the root client every time,
+# which may be on a different host.
+@functools.cache
+def _ensure_init_rdma_manager() -> Shared[None]:
+    async def task() -> None:
+        await (
+            await get_or_spawn_controller("rdma_controller", RdmaController)
+        ).init_rdma_on_mesh.call_one(none_throws(context().actor_instance.proc_mesh))
+
+    return PythonTask.from_coroutine(task()).spawn()
+
+
 def _assert_tensor_is_1d_contiguous_uint8(t: torch.Tensor) -> None:
     if t.ndim != 1:
         raise ValueError(f"Tensor must be 1D, got {t.ndim}D")
@@ -59,6 +106,10 @@ class RDMABuffer:
             is_available()
         ), "Tried to create an RDMABuffer, but RDMA is not available on this platform."
 
+        # We need to ensure that _RdmaManager is initialized at this point, because under the hood
+        # _RdmaBuffer.create_rdma_buffer_blocking relies on this being the case.
+        _ensure_init_rdma_manager().block_on()
+
         if data.device.type != "cpu":
             # TODO - CUDA support for RDMABuffer exists at the Rust layer, but
             # runs into issues with MR creation. For now, only support CPU tensors.
@@ -76,12 +127,12 @@ class RDMABuffer:
             storage = data.untyped_storage()
             addr: int = storage.data_ptr()
             size = storage.element_size() * data.numel()
-            ctx = MonarchContext.get()
+            ctx = context()
             self._buffer: _RdmaBuffer = _RdmaBuffer.create_rdma_buffer_blocking(
                 addr=addr,
                 size=size,
-                proc_id=ctx.proc_id,
-                client=ctx.mailbox,
+                proc_id=ctx.actor_instance.proc_id,
+                client=ctx.actor_instance._mailbox,
             )
         # TODO - specific exception
         except Exception as e:
@@ -120,10 +171,12 @@ class RDMABuffer:
                 f"offset + size ({offset + size}) must be <= dst.numel() ({dst.numel()})"
             )
 
-        local_proc_id = MonarchContext.get().proc_id
-        client = MonarchContext.get().mailbox
+        local_proc_id = context().actor_instance.proc_id
+        client = context().actor_instance._mailbox
 
         async def read_into_nonblocking() -> Optional[int]:
+            await _ensure_init_rdma_manager()
+
             res = await self._buffer.read_into(
                 addr=addr,
                 size=size,
@@ -167,10 +220,12 @@ class RDMABuffer:
                 f"size + offset ({size + offset}) must be <= src.numel() ({src.numel()})"
             )
 
-        local_proc_id = MonarchContext.get().proc_id
-        client = MonarchContext.get().mailbox
+        local_proc_id = context().actor_instance.proc_id
+        client = context().actor_instance._mailbox
 
         async def write_from_nonblocking() -> None:
+            await _ensure_init_rdma_manager()
+
             res = await self._buffer.write_from(
                 addr=addr,
                 size=size,

monarch/_testing.py

@@ -133,9 +133,7 @@ class TestingContext:
     ) -> Generator[DeviceMesh, None, None]:
         key = (num_hosts, gpu_per_host)
         if key not in self._proc_mesh_cache:
-            self._proc_mesh_cache[key] = proc_mesh(
-                hosts=num_hosts, gpus=gpu_per_host
-            ).get()
+            self._proc_mesh_cache[key] = proc_mesh(hosts=num_hosts, gpus=gpu_per_host)
 
         dm = spawn_tensor_engine(self._proc_mesh_cache[key])
         dm = dm.rename(hosts="host", gpus="gpu")

monarch/actor/__init__.py

@@ -4,6 +4,7 @@
 # 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
 """
 Monarch Actor API - Public interface for actor functionality.
 """
@@ -13,18 +14,29 @@ from monarch._src.actor.actor_mesh import (
     Actor,
     ActorError,
     as_endpoint,
+    Channel,
+    context,
     current_actor_name,
     current_rank,
     current_size,
     Point,
-    port,
+    Port,
+    PortReceiver,
     send,
     ValueMesh,
 )
+from monarch._src.actor.debugger.debugger import debug_controller
 from monarch._src.actor.endpoint import endpoint
 from monarch._src.actor.future import Future
+
+from monarch._src.actor.host_mesh import (
+    HostMesh,
+    hosts_from_config,
+    this_host,
+    this_proc,
+)
 from monarch._src.actor.proc_mesh import (
-    debug_client,
+    get_or_spawn_controller,
     local_proc_mesh,
     proc_mesh,
     ProcMesh,
@@ -45,9 +57,17 @@ __all__ = [
     "Point",
     "proc_mesh",
     "ProcMesh",
-    "port",
+    "Channel",
     "send",
     "sim_proc_mesh",
     "ValueMesh",
-    "debug_client",
+    "debug_controller",
+    "get_or_spawn_controller",
+    "this_host",
+    "this_proc",
+    "HostMesh",
+    "context",
+    "hosts_from_config",
+    "Port",
+    "PortReceiver",
 ]