torchmonarch-nightly 2025.6.27__cp311-cp311-manylinux2014_x86_64.whl

monarch/allocator.py

@@ -0,0 +1,220 @@
+# 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 abc
+import logging
+from typing import final, Optional
+
+from monarch import ActorFuture as Future
+from monarch._rust_bindings.hyperactor_extension.alloc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    Alloc,
+    AllocSpec,
+)
+
+from monarch._rust_bindings.monarch_hyperactor.alloc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    LocalAllocatorBase,
+    ProcessAllocatorBase,
+    RemoteAllocatorBase,
+)
+
+ALLOC_LABEL_PROC_MESH_NAME = "procmesh.monarch.meta.com/name"
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@final
+class ProcessAllocator(ProcessAllocatorBase):
+    """
+    An allocator that allocates by spawning local processes.
+    """
+
+    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),
+        )
+
+
+@final
+class LocalAllocator(LocalAllocatorBase):
+    """
+    An allocator that allocates by spawning actors into the current process.
+    """
+
+    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),
+        )
+
+
+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, match_labels: dict[str, str]) -> 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.
+
+        Arguments:
+        - `match_labels`: The match labels specified in `AllocSpec.AllocConstraints`. Initializer implementations
+            can read specific labels for matching a set of hosts that will service `allocate()` 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, match_labels: dict[str, str]) -> list[str]:
+        _ = match_labels  # Suppress unused variable warning
+        return list(self.addrs)
+
+
+class TorchXRemoteAllocInitializer(RemoteAllocInitializer):
+    """
+    For monarch runtimes running as a job on a supported scheduler.
+    Such runtimes are typically launched using the monarch CLI (e.g `monarch create --scheduler slurm ...`).
+
+    Returns the server addresses of a specific monarch runtime by using TorchX's status API
+    to get the hostnames of the nodes.
+    """
+
+    def __init__(
+        self,
+        server_handle: str,
+        /,
+        transport: Optional[str] = None,
+        port: Optional[int] = None,
+    ) -> None:
+        """
+        NOTE: If `transport` and `port` specified, they are used over the `transport` and `port`
+          information that is tagged as metadata on the server's job. This is useful in two specific
+          situations:
+            1) The job was NOT created wit monarch CLI (hence no metadata tags exist)
+            2) The scheduler does not support job metadata tagging
+
+        Arguments:
+        - `server_handle`: points to a monarch runtime. Of the form `{scheduler}://{namespace}/{job_id}`.
+             the `{namespace}` can be empty if not configured (e.g. `slurm:///1234` - notice the triple slashes).
+        - `transport`: the channel transport that should be used to connect to the remote process allocator address
+        - `port`: the port that the remote process allocator is running on
+
+        """
+        self.server_handle = server_handle
+        self.transport = transport
+        self.port = port
+
+    async def initialize_alloc(self, match_labels: dict[str, str]) -> list[str]:
+        # lazy import since torchx-fb is not included in `fbcode//monarch/python/monarch:monarch.whl`
+        # nor any of the base conda environments
+        from monarch.tools.commands import server_ready
+
+        mesh_name = match_labels.get(ALLOC_LABEL_PROC_MESH_NAME)
+
+        server = await server_ready(self.server_handle)
+
+        # job does not exist or it is in a terminal state (SUCCEEDED, FAILED, CANCELLED)
+        if not (server and server.is_running):
+            raise ValueError(
+                f"{self.server_handle} does not exist or is in a terminal state"
+            )
+
+        if not mesh_name:
+            logger.info(
+                "no match label `%s` specified in alloc constraints",
+                ALLOC_LABEL_PROC_MESH_NAME,
+            )
+
+            num_meshes = len(server.meshes)
+
+            if num_meshes == 1:
+                logger.info(
+                    "found a single proc mesh `%s` in %s, will allocate on it",
+                    server.meshes[0].name,
+                    self.server_handle,
+                )
+            else:
+                raise RuntimeError(
+                    f"{num_meshes} proc meshes in {self.server_handle},"
+                    f" please specify the mesh name as a match label `{ALLOC_LABEL_PROC_MESH_NAME}`"
+                    f" in allocation constraints of the alloc spec"
+                )
+            mesh = server.meshes[0]
+        else:
+            mesh = server.get_mesh_spec(mesh_name)
+
+        server_addrs = mesh.server_addrs(self.transport, self.port)
+
+        logger.info(
+            "initializing alloc on remote allocator addresses: %s", server_addrs
+        )
+        return server_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

@@ -0,0 +1,59 @@
+# 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.
+
+"""
+This is the main function for the boostrapping a new process using a ProcessAllocator.
+"""
+
+import asyncio
+import importlib.resources
+import logging
+import os
+import sys
+
+# 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())
+
+
+if __name__ == "__main__":
+    invoke_main()  # pragma: no cover

monarch/builtins/__init__.py

@@ -0,0 +1,14 @@
+# 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
+"""
+Builtins for Monarch is a set of remote function defintions for PyTorch functions and other utilities.
+"""
+
+from .log import log_remote, set_logging_level_remote
+
+__all__ = ["log_remote", "set_logging_level_remote"]

monarch/builtins/log.py

@@ -0,0 +1,22 @@
+# 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 logging
+
+from monarch.common.remote import remote
+
+
+logger = logging.getLogger(__name__)
+
+
+@remote(propagate="inspect")
+def log_remote(*args, level: int = logging.WARNING, **kwargs) -> None:
+    logger.log(level, *args, **kwargs)
+
+
+@remote(propagate="inspect")
+def set_logging_level_remote(level: int) -> None:
+    logger.setLevel(level)

monarch/builtins/random.py

@@ -0,0 +1,68 @@
+# 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
+from typing import Callable
+
+import torch
+from monarch.common.remote import remote
+
+
+@remote(propagate="inspect")
+def set_manual_seed_remote(seed: int, process_idx: int = 0) -> None:
+    torch.manual_seed(seed ^ process_idx)
+
+
+@remote(propagate=lambda: torch.zeros(1))
+def get_rng_state_remote() -> torch.Tensor:
+    return torch.get_rng_state()
+
+
+@remote(propagate="inspect")
+def set_rng_state_remote(new_state: torch.Tensor) -> None:
+    torch.set_rng_state(new_state)
+
+
+def _run_no_return(f: Callable) -> None:
+    f()
+    return None
+
+
+# TODO: return result when uint64 is supported from remote function
+@remote(propagate=lambda: _run_no_return(torch.seed))
+def seed_remote() -> None:
+    torch.seed()
+
+
+# same underlying implementation as seed_remote (torch.seed)
+# TODO: return result when uint64 is supported from remote function
+@remote(propagate=lambda: _run_no_return(torch.random.seed))
+def random_seed_remote() -> None:
+    torch.random.seed()
+
+
+@remote(propagate="inspect")
+def manual_seed_cuda_remote(seed: int) -> None:
+    torch.cuda.manual_seed(seed)
+
+
+@remote(propagate="inspect")
+def manual_seed_all_cuda_remote(seed: int) -> None:
+    torch.cuda.manual_seed_all(seed)
+
+
+@remote(propagate=lambda: [torch.zeros(1)])
+def get_rng_state_all_cuda_remote() -> list[torch.Tensor]:
+    return torch.cuda.get_rng_state_all()
+
+
+@remote(propagate="inspect")
+def set_rng_state_all_cuda_remote(states: list[torch.Tensor]) -> None:
+    torch.cuda.set_rng_state_all(states)
+
+
+# initial_seed may sometimes return a uint64 which currenly can't be unwrapped by the framework
+# def initial_seed_remote() -> int: ...

monarch/cached_remote_function.py

@@ -0,0 +1,257 @@
+# 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 logging
+
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Type, Union
+
+import torch
+from monarch.common.process_group import SingleControllerProcessGroupWrapper
+
+from monarch.common.remote import DummyProcessGroup, remote, RemoteProcessGroup
+
+from torch import autograd
+from torch.utils._pytree import tree_flatten, tree_unflatten
+
+logger = logging.getLogger(__name__)
+
+
+def _controller_autograd_function_forward(
+    autograd_function_class: Type[autograd.Function],
+):
+    """
+    Decorator for authoring a controller remote function wrapper that wraps an autograd.Function forward.
+    Sets up the autograd.function.FunctionCtx() to send over the wire and sets up the original ctx
+    with the ctx_tensors and ctx attributes.
+    """
+
+    def decorator(func):
+        def wrapper(ctx, *args):
+            # Need dummy context because cannot pickle autograd.FunctionBackward
+            wire_ctx = autograd.function.FunctionCtx()
+            # Track arg tensors that have requires_grad
+            arg_tensors, _ = tree_flatten(args)
+            wire_ctx.args_requires_grads = []
+            for i, arg in enumerate(arg_tensors):
+                if isinstance(arg, torch.Tensor) and arg.requires_grad:
+                    wire_ctx.args_requires_grads.append(i)
+            out, ctx_attrs, ctx_tensors = func(
+                autograd_function_class.__module__,
+                autograd_function_class.__name__,
+                wire_ctx,
+                *args,
+            )
+            if ctx is None:
+                return out
+            ctx.save_for_backward(*ctx_tensors)
+            ctx.attr_names = ctx_attrs.keys()
+            ctx.pg_names = []
+            dim_to_remote_group = {}
+            for arg in args:
+                if isinstance(arg, RemoteProcessGroup):
+                    dim_to_remote_group[arg.dims] = arg
+            for name, v in ctx_attrs.items():
+                if isinstance(v, DummyProcessGroup):
+                    setattr(ctx, name, dim_to_remote_group[v.dims])
+                    ctx.pg_names.append(name)
+                else:
+                    setattr(ctx, name, v)
+
+            return out
+
+        return wrapper
+
+    return decorator
+
+
+def _controller_autograd_function_backward(
+    autograd_function_class: Type[autograd.Function],
+):
+    """
+    Decorator for authoring a controller remote function wrapper that wraps an autograd.Function backward.
+    Manually sets up wire_ctx with ctx tensors and attributes.
+    """
+
+    def decorator(func):
+        def wrapper(ctx, *grad_outputs):
+            # Manually set up wire_ctx with ctx tensors and attributes
+            wire_ctx = autograd.function.FunctionCtx()
+            # send over tensor references with ctx_tensors
+            ctx_tensors = ctx.saved_tensors
+            wire_ctx.save_for_backward(ctx_tensors)
+            for name in ctx.attr_names:
+                setattr(wire_ctx, name, getattr(ctx, name))
+            process_groups = {name: getattr(ctx, name) for name in ctx.pg_names}
+
+            return func(
+                autograd_function_class.__module__,
+                autograd_function_class.__name__,
+                wire_ctx,
+                ctx_tensors,
+                # explicitly pass pg to worker
+                process_groups,
+                *grad_outputs,
+            )
+
+        return wrapper
+
+    return decorator
+
+
+@contextmanager
+def manage_grads(list_of_tensors, indices):
+    try:
+        for i in indices:
+            assert list_of_tensors[i].is_leaf, "can't have non-leaf tensors on worker"
+            list_of_tensors[i].requires_grad = True
+        yield list_of_tensors
+    finally:
+        for i in indices:
+            list_of_tensors[i].requires_grad = False
+
+
+def worker_autograd_function_forward(
+    module_name: str,
+    class_name: str,
+    ctx: autograd.function.FunctionCtx,
+    *args,
+    **kwargs,
+):
+    # Capture initial state of ctx attributes
+    before = set()
+    before.add("to_save")
+    for attr in dir(ctx):
+        if not attr.startswith("_"):
+            before.add(attr)
+
+    # Set tensors that require grad from additional arg
+    flatten_args, spec = tree_flatten(args)
+    # pyre-ignore
+    with manage_grads(flatten_args, ctx.args_requires_grads) as args_with_grad:
+        args = tree_unflatten(args_with_grad, spec)
+
+        # Call the original forward function
+        module = importlib.import_module(module_name)
+        class_ = getattr(module, class_name)
+        with torch.no_grad():
+            out = class_.forward(ctx, *args, **kwargs)
+
+        # Capture state of ctx attributes after the function call
+        after = set()
+        for attr in dir(ctx):
+            if not attr.startswith("_"):
+                after.add(attr)
+        ctx_attrs = {attr: getattr(ctx, attr) for attr in after - before}
+        ctx_attrs["ctx_requires_grads"] = []
+
+        if not hasattr(ctx, "to_save"):
+            to_save = []
+        else:
+            # pyre-ignore
+            for idx, t in enumerate(ctx.to_save):
+                # generally, workers should not have requires_grad set. Set to correct state after
+                # but record requires_grad for next forward
+                if isinstance(t, torch.Tensor) and t.requires_grad and t.is_leaf:
+                    t.requires_grad = False
+                    ctx_attrs["ctx_requires_grads"].append(idx)
+            to_save = ctx.to_save
+    return out, ctx_attrs, to_save
+
+
+def worker_autograd_function_backward(
+    module_name: str,
+    class_name: str,
+    ctx: autograd.function.FunctionCtx,
+    ctx_tensors: List[torch.Tensor],
+    process_groups: Dict[
+        str, Union[SingleControllerProcessGroupWrapper, DummyProcessGroup]
+    ],
+    *grad_outputs: torch.Tensor,
+):
+    # set correct requires_grad state pre backward
+    # pyre-ignore
+    with manage_grads(ctx_tensors, ctx.ctx_requires_grads) as ctx_grad_tensors:
+        # for i in ctx.ctx_requires_grads:
+        #     ctx_tensors[i].requires_grad = True
+        if ctx_grad_tensors:
+            # pyre-ignore
+            ctx.saved_tensors = ctx_grad_tensors
+        for name, v in process_groups.items():
+            setattr(ctx, name, v)
+        # Call the original backward function
+        module = importlib.import_module(module_name)
+        class_ = getattr(module, class_name)
+        with torch.no_grad():
+            out = class_.backward(ctx, *grad_outputs)
+    return out
+
+
+forward_remote_fn = remote(
+    "monarch.cached_remote_function.worker_autograd_function_forward"
+)
+
+backward_remote_fn = remote(
+    "monarch.cached_remote_function.worker_autograd_function_backward"
+)
+
+
+class RemoteAutogradFunction(autograd.Function):
+    """
+    New autograd.Function (custom forward/backward) that will run on the worker as a UDF RemoteFunction
+
+
+    Example::
+        my_remote_autograd_function = remote_autograd_function(my_custom_autograd_function)
+    """
+
+    @staticmethod
+    def forward(ctx, *args):
+        raise NotImplementedError()
+
+    @staticmethod
+    def backward(ctx, *grads):
+        raise NotImplementedError()
+
+
+def remote_autograd_function(
+    target_class: Type[autograd.Function], name: Optional[str] = None
+) -> Type[RemoteAutogradFunction]:
+    """
+    Returns a new autograd.Function (custom forward/backward) that will run on the worker as a UDF RemoteFunction
+    Logic is done on the controller (e.g., Dtensors set up and saved for backward).
+    The autograd.function.FunctionCtx() is sent over the wire to the worker.
+    Special handling is done for ctx_tensors, requires_grad fo tensors and process groups.
+
+    Args:
+        target_class: autograd.Function class to be run remotely
+        name: name of the new autograd.Function to be called on the worker
+    """
+    if issubclass(target_class, RemoteAutogradFunction):
+        logging.warning(
+            f"{target_class} is already a autograd.Function UDF! You are likely monkey-patching too many times"
+        )
+        return target_class
+    assert issubclass(
+        target_class, autograd.Function
+    ), f"{target_class} is not a torch.autograd.Function!"
+    if name is None:
+        name = f"Remote_{target_class.__name__}"
+
+    return type(
+        name,
+        (RemoteAutogradFunction,),
+        {
+            "forward": staticmethod(
+                _controller_autograd_function_forward(target_class)(forward_remote_fn)
+            ),
+            "backward": staticmethod(
+                _controller_autograd_function_backward(target_class)(backward_remote_fn)
+            ),
+        },
+    )

monarch/code_sync.py

@@ -0,0 +1,10 @@
+# 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.
+
+from monarch._rust_bindings.monarch_extension.code_sync import (  # noqa: F401
+    RemoteWorkspace,
+    RsyncMeshClient,
+)

monarch/common/_C.pyi

@@ -0,0 +1,11 @@
+# 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
+
+def patch_cuda() -> None: ...
+def mock_cuda() -> None: ...
+def unmock_cuda() -> None: ...