torchmonarch-nightly 2025.6.4__cp310-cp310-manylinux2014_x86_64.whl

monarch/random.py

@@ -0,0 +1,61 @@
+# 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 os
+from typing import NamedTuple, Tuple
+
+import torch
+from monarch.common.remote import remote
+from monarch.common.tensor import Tensor
+
+
+class State(NamedTuple):
+    cpu: Tensor
+    cuda: Tensor
+
+
+@remote(
+    propagate=lambda: (
+        torch.empty(5056, dtype=torch.uint8),
+        torch.empty(16, dtype=torch.uint8),
+    )
+)
+def _get_state() -> Tuple[torch.Tensor, torch.Tensor]:
+    return (torch.get_rng_state(), torch.cuda.get_rng_state())
+
+
+@remote(propagate=lambda state: None)
+def set_state(state: Tuple[Tensor, Tensor]):
+    cpu, device = state
+    torch.set_rng_state(cpu)
+    torch.cuda.set_rng_state(device)
+
+
+@remote(propagate=lambda _: None)
+def _manual_seed(seed: torch.Tensor):
+    torch.manual_seed(seed.item())
+
+
+@remote(propagate=lambda: None)
+def make_deterministic():
+    torch.use_deterministic_algorithms(True)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    # env var for deterministic CuBLAS
+    # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def get_state() -> State:
+    return State(*_get_state())
+
+
+def new_state(seed: Tensor) -> State:
+    orig = get_state()
+    _manual_seed(seed)
+    mine = get_state()
+    set_state(orig)
+    return mine

monarch/rdma.py

@@ -0,0 +1,190 @@
+# 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 ctypes
+
+import traceback
+
+from dataclasses import dataclass
+from traceback import extract_tb, StackSummary
+from typing import cast, Dict, Optional, Tuple
+
+import torch
+
+from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
+
+from monarch.actor_mesh import (
+    _ActorMeshRefImpl,
+    Actor,
+    ActorMeshRef,
+    endpoint,
+    MonarchContext,
+)
+
+
+@dataclass
+class LocalRDMARecord:
+    data: torch.Tensor
+
+
+_local_buffers: Dict[int, "LocalRDMARecord"] = {}
+
+
+def _get_bytes(storage: torch.Tensor, offset: int, size: int) -> bytearray:
+    """Extracts a bytearray from a 1D, 1byte per item tensor."""
+    if offset + size > storage.numel():
+        raise ValueError(f"Read out of range: {offset + size} > {storage.size()}")
+    addr = storage.data_ptr()
+    if storage.device.type != "cpu":
+        result = bytearray(size)
+        result_tensor = torch.frombuffer(
+            result,
+            dtype=torch.uint8,
+        )
+        source_tensor = storage[offset:]
+        result_tensor.copy_(source_tensor)
+    else:
+        ctypes_array = (ctypes.c_byte * size).from_address(addr)
+        result = bytearray(ctypes_array)
+    return result
+
+
+class RDMAManager(Actor):
+    @staticmethod
+    def on_proc(proc_id: str) -> "RDMAManager":
+        ctx = MonarchContext.get()
+        return cast(
+            RDMAManager,
+            ActorMeshRef(
+                RDMAManager,
+                _ActorMeshRefImpl.from_actor_id(
+                    ctx.mailbox,
+                    ActorId.from_string(f"{proc_id}.rdma_manager[0]"),
+                ),
+                ctx.mailbox,
+            ),
+        )
+
+    @endpoint
+    async def drop(self, addr: int) -> None:
+        if addr in _local_buffers:
+            del _local_buffers[addr]
+
+    @endpoint
+    async def fetch(self, addr: int, offset: int, nbytes: int) -> bytearray:
+        if addr not in _local_buffers:
+            raise ValueError(f"Unknown buffer {addr}")
+        storage = _local_buffers[addr].data
+        return _get_bytes(storage, offset, nbytes)
+
+    @endpoint
+    async def put(self, addr: int, offset: int, bytes: bytearray) -> None:
+        if addr not in _local_buffers:
+            raise ValueError(f"Unknown buffer {addr}")
+        storage = _local_buffers[addr].data
+        storage[offset : offset + len(bytes)] = torch.frombuffer(
+            bytes, dtype=storage.dtype
+        )
+
+
+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")
+    if t.dtype != torch.uint8:
+        raise ValueError(f"Tensor must be uint8, got {t.dtype}")
+    if not t.is_contiguous():
+        raise ValueError("Tensor must be contiguous")
+
+
+class RDMABuffer:
+    def __init__(self, data: torch.Tensor) -> None:
+        """
+        RDMABuffer only supports 1D contiguous tensors that are 1 byte per item.
+
+        To create a 1 byte, 1D view, use t.view(torch.uint8).flatten()
+
+        TODO: Create TensorBuffer, which will be main user API supporting non-contiguous , multi-byte-per-elment tensors
+        """
+        _assert_tensor_is_1d_contiguous_uint8(data)
+        assert data.storage_offset() == 0
+        storage = data.untyped_storage()
+        self.addr: int = storage.data_ptr()
+        self.begin = 0
+        self.end: int = storage.size()
+        self.proc_id: str = MonarchContext.get().proc_id
+        self.local_data: object = None
+        _local_buffers[self.addr] = LocalRDMARecord(data)
+
+    def drop(self) -> None:
+        if self.proc_id is None:
+            del _local_buffers[self.addr]
+            return
+        rmda_actor = RDMAManager.on_proc(self.proc_id)
+        # pyre-ignore[16]: Undefined attribute [16]: `Endpoint` has no attribute `cast`.
+        rmda_actor.drop.cast(self.addr)
+
+    def __getstate__(self) -> Tuple[int, int, int, Optional[str]]:
+        proc_id = self.proc_id
+        # locally created RDMABuffer being set remotely,
+        # record its proc_id so we know how to establish connections to it
+        if proc_id is None:
+            proc_id = MonarchContext.get().proc_id
+        return (self.addr, self.begin, self.end, proc_id)
+
+    def __setstate__(self, state: Tuple[int, int, int, str]) -> None:
+        self.local_data = None
+        self.addr, self.begin, self.end, self.proc_id = state
+
+    async def read_into(self, dst: torch.Tensor, offset: int = 0) -> None:
+        """
+        Read data from the RDMABuffer into a destination tensor.
+
+        The destination tensor must be contiguous and 1 byte per item.
+        """
+        _assert_tensor_is_1d_contiguous_uint8(dst)
+        bytes = await RDMAManager.on_proc(self.proc_id).fetch.call_one(
+            self.addr, offset, dst.numel()
+        )
+        dst.copy_(torch.frombuffer(bytes, dtype=torch.uint8))
+
+    async def write(self, src: torch.Tensor, offset: int = 0) -> None:
+        """
+        Write data from a source tensor into the RDMABuffer.
+
+        The source tensor must be contiguous and 1 byte per item.
+        """
+        _assert_tensor_is_1d_contiguous_uint8(src)
+        bytes = _get_bytes(
+            src,
+            cast(int, src.storage_offset()),
+            src.numel(),
+        )
+        await RDMAManager.on_proc(self.proc_id).put.call_one(self.addr, offset, bytes)
+
+
+class ActorMeshRefCallFailedException(Exception):
+    """
+    Deterministic problem with the user's code.
+    For example, an OOM resulting in trying to allocate too much GPU memory, or violating
+    some invariant enforced by the various APIs.
+    """
+
+    def __init__(
+        self,
+        exception: Exception,
+        message: str = "A remote service call has failed asynchronously.",
+    ) -> None:
+        self.exception = exception
+        self.actor_mesh_ref_frames: StackSummary = extract_tb(exception.__traceback__)
+        self.message = message
+
+    def __str__(self) -> str:
+        exe = str(self.exception)
+        actor_mesh_ref_tb = "".join(traceback.format_list(self.actor_mesh_ref_frames))
+        return (
+            f"{self.message}\n"
+            f"Traceback of where the service call failed (most recent call last):\n{actor_mesh_ref_tb}{type(self.exception).__name__}: {exe}"
+        )

monarch/remote_class.py

@@ -0,0 +1,114 @@
+# 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 importlib
+import itertools
+from typing import Any, Dict
+
+from monarch.common import device_mesh
+from monarch.common.remote import remote
+
+
+class ControllerRemoteClass:
+    """
+    This class simplifies the creation and management of remote classes.  It serves as
+    the controller side of a remote class architecture.  Classes that are intended to be
+    controlled remotely should inherit from this class. The constructor of the inheriting
+    class must invoke `super().__init__()` with the path to the remote class that will be
+    used on the worker nodes.  Methods that are intended for remote execution must be
+    decorated with `ControllerRemoteClass.remote_method`.
+
+    Note: This class is designed for use by the controller developer only and should
+    not be directly used in model code.
+
+    Example usage:
+
+    class ControllerMyClass(ControllerRemoteClass):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__("my_package.my_class", *args, **kwargs)
+
+        @ControllerRemoteClass.remote_method
+        def some_method(self, *args, **kwargs) -> None:
+            # This method is intended for remote execution and does nothing locally.
+            pass
+    """
+
+    _counter = itertools.count()
+
+    def __init__(self, cls_path: str, *args, **kwargs) -> None:
+        self.ident = next(ControllerRemoteClass._counter)
+        self.cls_path = cls_path
+        self.mesh = device_mesh._active
+        _controller_remote_class_method(
+            cls_path, "remote_init", self.ident, *args, **kwargs
+        )
+
+    def __del__(self) -> None:
+        mesh = getattr(self, "mesh", None)
+        if mesh is not None and not mesh.client._shutdown:
+            with self.mesh.activate():
+                _controller_remote_class_method(
+                    self.cls_path,
+                    "remote_del",
+                    self.ident,
+                )
+
+    @staticmethod
+    def remote_method(fn):
+        def wrapper(self, *args, **kwargs) -> None:
+            _controller_remote_class_method(
+                self.cls_path, "remote_method", self.ident, fn.__name__, *args, **kwargs
+            )
+
+        return wrapper
+
+
+# Add the logic as a separate private function instead of adding ita to
+# ResolvableFunctionFromPath. This avoids users to using this directly.
+_controller_remote_class_method = remote(
+    "monarch.remote_class._remote_class_method", propagate="inspect"
+)
+
+
+def _remote_class_method(cls_path: str, method_name: str, *args, **kwargs) -> None:
+    modulename, classname = cls_path.rsplit(".", 1)
+    module = importlib.import_module(modulename)
+    cls = getattr(module, classname)
+    method = getattr(cls, method_name)
+    method(*args, **kwargs)
+
+
+class WorkerRemoteClass:
+    """
+    This class is designed to be used alongside ``ControllerRemoteClass`` and represents
+    the worker-side of a remote class architecture. Instances of this class should just
+    mimic standard Python classes, with the notable exception that all methods must
+    return None -- the current RemoteClass architecture does not support methods that
+    return values.
+
+    The `ident` attribute is used for tracking object instances created via `remote_init`.
+    This tracking is necessary because the remote function would otherwise lose the
+    reference to the object.
+    """
+
+    _objects: Dict[int, Any] = {}
+
+    @classmethod
+    def remote_init(cls, ident: int, *args, **kwargs) -> None:
+        WorkerRemoteClass._objects[ident] = cls(*args, **kwargs)
+
+    @classmethod
+    def remote_del(cls, ident) -> None:
+        WorkerRemoteClass._objects.pop(ident)
+
+    @classmethod
+    def remote_method(cls, ident: int, method_name, *args, **kwargs) -> None:
+        instance = WorkerRemoteClass._objects[ident]
+        assert (
+            cls == instance.__class__
+        ), "Mismatched class type {cls} {instance.__class__}"
+        getattr(instance, method_name)(*args, **kwargs)

monarch/rust_backend_mesh.py

@@ -0,0 +1,280 @@
+# 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 logging
+import time
+from logging import Logger
+from typing import Any, Callable, Optional, Protocol
+
+from monarch._rust_bindings.monarch_extension.client import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ClientActor,
+    SystemSnapshotFilter,
+)
+
+from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ActorId,
+    init_proc,
+    Proc,
+)
+from monarch.common.client import Client
+from monarch.common.device_mesh import DeviceMesh, DeviceMeshStatus
+from monarch.common.invocation import DeviceException, RemoteException
+from monarch.common.mast import MastJob
+from monarch.common.shape import NDSlice
+from monarch.controller.rust_backend.controller import RustController
+
+TORCHX_MAST_TASK_GROUP_NAME = "script"
+
+logger: Logger = logging.getLogger(__name__)
+
+# A world tuple contains a worker world name and a controller actor id
+# The pair forms a functional world that can be used to create a device mesh
+MeshWorld = tuple[str, ActorId]
+
+# Taken from //monarch/controller/src/bootstrap.rs
+WORLD_WORKER_LABEL = "world.monarch.meta.com/worker"
+WORLD_CONTROLLER_LABEL = "world.monarch.meta.com/controllerActorId"
+WORLD_CONTROLLER_IP = "world.monarch.meta.com/ip_addr"
+
+
+class IBootstrap(Protocol):
+    def get_mesh_worlds(self) -> list[MeshWorld]:
+        """Returns the list of mesh worlds."""
+        ...
+
+    def kill_mesh(self, mesh_world: MeshWorld) -> None:
+        """Kills a mesh in a bootstrap instance."""
+        ...
+
+    def spawn_mesh(self, mesh_world: MeshWorld) -> None:
+        """Spawns a mesh in a bootstrap instance."""
+        ...
+
+
+class IPoolDeviceMeshProvider(Protocol):
+    def new_mesh(self, timeout_in_sec: Optional[int] = None) -> DeviceMesh:
+        raise NotImplementedError()
+
+
+class PoolDeviceMeshProvider:
+    """
+    Given a client actor, the device mesh provider discovers and keeps track of
+    the world status and provides a device mesh given a healthy world.
+    """
+
+    def __init__(
+        self,
+        hosts: int,
+        gpus: int,
+        proc: Proc,
+    ) -> None:
+        self._hosts = hosts
+        self._gpus = gpus
+        self._mesh_map: dict[MeshWorld, DeviceMesh | None] = {}
+        self._proc = proc
+        # Root client is not used to create device meshes.
+        # It is only used to pull the world status.
+        self._root_client: ClientActor = ClientActor(
+            proc=self._proc,
+            actor_name="root_client",  # The client name really doesn't matter
+        )
+
+    def new_mesh(self, timeout_in_sec: Optional[int] = None) -> DeviceMesh:
+        """
+        Creates a new device mesh based on the current world status.
+        If no healthy world is found, the call will block until a healthy world is found
+        or timeout_in_sec is reached.xtimeout_in_sec being None indicates no timeout.
+        """
+
+        logger.info("Trying to allocate a new mesh in its desired world...")
+
+        def _create_exit(
+            client: Client,
+        ) -> Callable[[Optional[RemoteException | DeviceException | Exception]], None]:
+            def _exit(
+                error: Optional[RemoteException | DeviceException | Exception] = None,
+            ) -> None:
+                client.shutdown(True, error)
+
+            return _exit
+
+        def _is_world_healthy(world_status: dict[str, str], target_world: str) -> bool:
+            return (
+                target_world in world_status
+                and DeviceMeshStatus(world_status[target_world])
+                == DeviceMeshStatus.LIVE
+            )
+
+        now = time.time()
+        while timeout_in_sec is None or time.time() - now < timeout_in_sec:
+            # Pull the fresh world status
+            self._refresh_worlds()
+            world_status = self._root_client.world_status()
+            self._remove_evicted_worlds(world_status)
+
+            # Find the next available world
+            for mesh_world, mesh in self._mesh_map.items():
+                if mesh is not None:
+                    # Mesh has been allocated to this world, skip
+                    continue
+
+                worker_world, controller_id = mesh_world
+                controller_world = controller_id.world_name
+
+                if (not _is_world_healthy(world_status, worker_world)) or (
+                    not _is_world_healthy(world_status, controller_world)
+                ):
+                    # Either controller world is not ready or worker world is not ready
+                    continue
+
+                # Create a new device mesh
+                backend_ctrl = RustController(
+                    proc=self._proc,
+                    client_actor=ClientActor.new_with_parent(
+                        self._proc, self._root_client.actor_id
+                    ),
+                    controller_id=controller_id,
+                    worker_world_name=worker_world,
+                )
+                client = Client(backend_ctrl, self._hosts * self._gpus, self._gpus)
+
+                # TODO: we need to consider hosts and gpus constraints as well
+                dm = DeviceMesh(
+                    client,
+                    NDSlice(
+                        offset=0,
+                        sizes=[self._hosts, self._gpus],
+                        strides=[self._gpus, 1],
+                    ),
+                    ("host", "gpu"),
+                    worker_world,
+                )
+                dm.exit = _create_exit(client)
+                self._mesh_map[mesh_world] = dm
+
+                logger.info("Mesh successfully allocated in world: %s", worker_world)
+
+                return dm
+
+            # TODO(T216841374): Change to healthy world push based checks
+            sleep_sec = 0.05
+            logger.debug(f"No healthy world found, sleeping for {sleep_sec}s...")
+            time.sleep(sleep_sec)
+
+        raise TimeoutError(f"Could not find a healthy world in {timeout_in_sec}s!")
+
+    def _refresh_worlds(self) -> None:
+        system_snapshot = self._root_client.world_state(
+            filter=SystemSnapshotFilter(world_labels={WORLD_WORKER_LABEL: "1"})
+        )
+        for world_id, world_snapshot in system_snapshot.items():
+            if WORLD_CONTROLLER_LABEL not in world_snapshot.labels:
+                continue
+            controller_actor_id = ActorId.from_string(
+                world_snapshot.labels[WORLD_CONTROLLER_LABEL]
+            )
+            world_tuple = (world_id, controller_actor_id)
+            if world_tuple not in self._mesh_map:
+                logger.debug(f"Discovered new worker world {world_id}")
+                self._mesh_map[world_tuple] = None
+
+    def _remove_evicted_worlds(self, world_status: dict[str, str]) -> None:
+        """
+        Go through the mesh map and remove the world that has already been evicted by the system.
+        """
+        mesh_worlds_to_remove = []
+        for mesh_world, _ in self._mesh_map.items():
+            worker_world, controller_id = mesh_world
+            controller_world = controller_id.world_name
+
+            if (
+                world_status.get(worker_world) is None
+                or world_status.get(controller_world) is None
+            ):
+                logger.debug(f"Removing Evicted world {mesh_world}")
+                mesh_worlds_to_remove.append(mesh_world)
+
+        for mesh_world in mesh_worlds_to_remove:
+            self._mesh_map.pop(mesh_world)
+
+
+def rust_mast_mesh(
+    job_name: str, system_port: int = 29500, **kwargs: Any
+) -> DeviceMesh:
+    job = MastJob(job_name, TORCHX_MAST_TASK_GROUP_NAME)
+    if not job.is_running():
+        job.wait_for_running(10 * 60)
+    hostnames = job.get_hostnames()
+    system_addr = f"metatls!{hostnames[0]}.facebook.com:{system_port}"
+    return rust_backend_mesh(
+        system_addr,
+        **kwargs,
+    )
+
+
+def rust_backend_mesh(
+    system_addr: str,
+    hosts: int,
+    gpus: int,
+) -> DeviceMesh:
+    dms = rust_backend_meshes(
+        system_addr,
+        hosts,
+        gpus,
+        requested_meshes=1,
+    )
+    assert len(dms) == 1
+    return dms[0]
+
+
+def rust_backend_meshes(
+    system_addr: str,
+    hosts: int,
+    gpus: int,
+    requested_meshes: int = 1,
+) -> list[DeviceMesh]:
+    """
+    Given system system_addr, discover worlds registered and create a device mesh per
+    world with hosts and gpus. The call will block until requested_meshes
+    are discovered and created, or 1200s timeout is reached.
+    Args:
+        system_addr: the system address to connect to.
+        hosts: number of hosts to create the device mesh with.
+        gpus: number of gpus to create the device mesh with.
+        requested_meshes: the minimum number of meshes to create.
+    """
+    mesh_provider = rust_backend_mesh_provider(system_addr, hosts, gpus)
+    dms: list[DeviceMesh] = []
+
+    # Given a client actor and a list of world names, wait for all the worlds to be ready.
+    max_timeout_in_sec = 1200
+    start_time = time.time()
+    while True:
+        if time.time() - start_time > max_timeout_in_sec:
+            raise TimeoutError(
+                f"Timeout ({max_timeout_in_sec} sec) waiting for all worlds to be ready."
+            )
+        mesh = mesh_provider.new_mesh()
+        dms.append(mesh)
+        if len(dms) == requested_meshes:
+            return dms
+
+
+def rust_backend_mesh_provider(
+    system_addr: str,
+    hosts: int,
+    gpus: int,
+    client_proc_id: str = "client[0]",
+    # pyre-fixme[11]: Annotation `DeviceMeshProvider` is not defined as a type.
+) -> PoolDeviceMeshProvider:
+    proc: Proc = init_proc(
+        proc_id=client_proc_id,
+        bootstrap_addr=system_addr,
+        timeout=5,
+        supervision_update_interval=5,
+    )
+    return PoolDeviceMeshProvider(hosts, gpus, proc)