torchmonarch-nightly 2025.6.27__cp311-cp311-manylinux2014_x86_64.whl

monarch/__init__.py

@@ -0,0 +1,189 @@
+# 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
+
+from importlib import import_module as _import_module
+from typing import TYPE_CHECKING
+
+# Import before monarch to pre-load torch DSOs as, in exploded wheel flows,
+# our RPATHs won't correctly find them.
+import torch  # noqa: F401
+
+# submodules of monarch should not be imported in this
+# top-level file because it will cause them to get
+# loaded even if they are not actually being used.
+# for instance if we import monarch.common.functions,
+# we might not want to also import monarch.common.tensor,
+# which recursively imports torch.
+
+# Instead to expose functionality as part of the
+# monarch.* API, import it inside the TYPE_CHECKING
+# guard (so typechecker works), and then add it
+# to the _public_api dict and __all__ list. These
+# entries will get loaded on demand.
+
+
+if TYPE_CHECKING:
+    from monarch import timer
+    from monarch.allocator import LocalAllocator, ProcessAllocator
+    from monarch.common._coalescing import coalescing
+
+    from monarch.common.device_mesh import (
+        DeviceMesh,
+        get_active_mesh,
+        no_mesh,
+        RemoteProcessGroup,
+        slice_mesh,
+        to_mesh,
+    )
+
+    from monarch.common.function import resolvers as function_resolvers
+
+    from monarch.common.future import Future
+
+    from monarch.common.invocation import RemoteException
+    from monarch.common.opaque_ref import OpaqueRef
+    from monarch.common.pipe import create_pipe, Pipe, remote_generator
+    from monarch.common.remote import remote
+    from monarch.common.selection import Selection
+    from monarch.common.shape import NDSlice, Shape
+    from monarch.common.stream import get_active_stream, Stream
+    from monarch.common.tensor import reduce, reduce_, Tensor
+    from monarch.fetch import fetch_shard, inspect, show
+    from monarch.future import ActorFuture
+    from monarch.gradient_generator import grad_function, grad_generator
+    from monarch.notebook import mast_mesh, reserve_torchx as mast_reserve
+    from monarch.python_local_mesh import python_local_mesh
+    from monarch.rust_backend_mesh import (
+        rust_backend_mesh,
+        rust_backend_meshes,
+        rust_mast_mesh,
+    )
+    from monarch.rust_local_mesh import local_mesh, local_meshes, SocketType
+    from monarch.simulator.config import set_meta  # noqa
+    from monarch.simulator.interface import Simulator
+    from monarch.world_mesh import world_mesh
+
+
+_public_api = {
+    "coalescing": ("monarch.common._coalescing", "coalescing"),
+    "remote": ("monarch.common.remote", "remote"),
+    "DeviceMesh": ("monarch.common.device_mesh", "DeviceMesh"),
+    "get_active_mesh": ("monarch.common.device_mesh", "get_active_mesh"),
+    "no_mesh": ("monarch.common.device_mesh", "no_mesh"),
+    "RemoteProcessGroup": ("monarch.common.device_mesh", "RemoteProcessGroup"),
+    "function_resolvers": ("monarch.common.function", "resolvers"),
+    "Future": ("monarch.common.future", "Future"),
+    "RemoteException": ("monarch.common.invocation", "RemoteException"),
+    "Shape": ("monarch.common.shape", "Shape"),
+    "NDSlice": ("monarch.common.shape", "NDSlice"),
+    "Selection": ("monarch.common.selection", "Selection"),
+    "OpaqueRef": ("monarch.common.opaque_ref", "OpaqueRef"),
+    "create_pipe": ("monarch.common.pipe", "create_pipe"),
+    "Pipe": ("monarch.common.pipe", "Pipe"),
+    "remote_generator": ("monarch.common.pipe", "remote_generator"),
+    "get_active_stream": ("monarch.common.stream", "get_active_stream"),
+    "Stream": ("monarch.common.stream", "Stream"),
+    "Tensor": ("monarch.common.tensor", "Tensor"),
+    "reduce": ("monarch.common.tensor", "reduce"),
+    "reduce_": ("monarch.common.tensor", "reduce_"),
+    "to_mesh": ("monarch.common.device_mesh", "to_mesh"),
+    "slice_mesh": ("monarch.common.device_mesh", "slice_mesh"),
+    "call_on_shard_and_fetch": ("monarch.fetch", "call_on_shard_and_fetch"),
+    "fetch_shard": ("monarch.fetch", "fetch_shard"),
+    "inspect": ("monarch.fetch", "inspect"),
+    "show": ("monarch.fetch", "show"),
+    "grad_function": ("monarch.gradient_generator", "grad_function"),
+    "grad_generator": ("monarch.gradient_generator", "grad_generator"),
+    "python_local_mesh": ("monarch.python_local_mesh", "python_local_mesh"),
+    "mast_mesh": ("monarch.notebook", "mast_mesh"),
+    "mast_reserve": ("monarch.notebook", "reserve_torchx"),
+    "rust_backend_mesh": ("monarch.rust_backend_mesh", "rust_backend_mesh"),
+    "rust_backend_meshes": ("monarch.rust_backend_mesh", "rust_backend_meshes"),
+    "local_mesh": ("monarch.rust_local_mesh", "local_mesh"),
+    "local_meshes": ("monarch.rust_local_mesh", "local_meshes"),
+    "SocketType": ("monarch.rust_local_mesh", "SocketType"),
+    "rust_mast_mesh": ("monarch.rust_backend_mesh", "rust_mast_mesh"),
+    "set_meta": ("monarch.simulator.config", "set_meta"),
+    "Simulator": ("monarch.simulator.interface", "Simulator"),
+    "world_mesh": ("monarch.world_mesh", "world_mesh"),
+    "timer": ("monarch.timer", "timer"),
+    "ProcessAllocator": ("monarch.allocator", "ProcessAllocator"),
+    "LocalAllocator": ("monarch.allocator", "LocalAllocator"),
+    "ActorFuture": ("monarch.future", "ActorFuture"),
+    "builtins": ("monarch.builtins", "builtins"),
+}
+
+
+def __getattr__(name):
+    if name in _public_api:
+        module_path, attr_name = _public_api[name]
+        module = _import_module(module_path)
+        result = getattr(module, attr_name)
+        globals()[name] = result
+        return result
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+try:
+    from __manifest__ import fbmake  # noqa
+
+    IN_PAR = True
+except ImportError:
+    IN_PAR = False
+
+# we have to explicitly list this rather than just take the keys of the _public_api
+# otherwise tools think the imports are unused
+__all__ = [
+    "coalescing",
+    "DeviceMesh",
+    "get_active_mesh",
+    "no_mesh",
+    "remote",
+    "RemoteProcessGroup",
+    "function_resolvers",
+    "Future",
+    "RemoteException",
+    "Shape",
+    "Selection",
+    "NDSlice",
+    "OpaqueRef",
+    "create_pipe",
+    "Pipe",
+    "remote_generator",
+    "get_active_stream",
+    "Stream",
+    "Tensor",
+    "reduce",
+    "reduce_",
+    "to_mesh",
+    "slice_mesh",
+    "call_on_shard_and_fetch",
+    "fetch_shard",
+    "inspect",
+    "show",
+    "grad_function",
+    "grad_generator",
+    "python_local_mesh",
+    "mast_mesh",
+    "mast_reserve",
+    "rust_backend_mesh",
+    "rust_backend_meshes",
+    "local_mesh",
+    "local_meshes",
+    "SocketType",
+    "rust_mast_mesh",
+    "set_meta",
+    "Simulator",
+    "world_mesh",
+    "timer",
+    "ProcessAllocator",
+    "LocalAllocator",
+    "ActorFuture",
+    "builtins",
+]
+assert sorted(__all__) == sorted(_public_api)

monarch/_monarch/__init__.py

@@ -0,0 +1,5 @@
+# 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.

monarch/_monarch/hyperactor/__init__.py

@@ -0,0 +1,58 @@
+# 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 monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage
+
+from monarch._rust_bindings.monarch_hyperactor.alloc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    LocalAllocatorBase,
+)
+
+from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox, PortId
+
+from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ActorId,
+    Alloc,
+    AllocConstraints,
+    AllocSpec,
+    init_proc,
+    Proc,
+    Serialized,
+)
+
+from monarch._rust_bindings.monarch_hyperactor.shape import (  # @manual=//monarch/monarch_extension:monarch_extension
+    Shape,
+)
+
+__all__ = [
+    "init_proc",
+    "Actor",
+    "ActorId",
+    "ActorHandle",
+    "Alloc",
+    "AllocSpec",
+    "PortId",
+    "Proc",
+    "Serialized",
+    "PickledMessage",
+    "PickledMessageClientActor",
+    "PythonMessage",
+    "Mailbox",
+    "PortHandle",
+    "PortReceiver",
+    "OncePortHandle",
+    "OncePortReceiver",
+    "Alloc",
+    "AllocSpec",
+    "AllocConstraints",
+    "ProcMesh",
+    "PythonActorMesh",
+    "ProcessAllocatorBase",
+    "Shape",
+    "Selection",
+    "LocalAllocatorBase",
+]

monarch/_monarch/selection/__init__.py

@@ -0,0 +1,13 @@
+# 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_hyperactor.selection import (  # @manual=//monarch/monarch_extension:monarch_extension
+    Selection,
+)
+
+__all__ = [
+    "Selection",
+]

monarch/_monarch/worker/debugger.py

@@ -0,0 +1,117 @@
+# 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 bdb
+import io
+import logging
+import pdb  # noqa
+import sys
+from typing import cast, Optional
+
+from monarch._rust_bindings.monarch_extension import debugger
+from monarch._rust_bindings.monarch_messages.debugger import DebuggerAction
+
+logger = logging.getLogger(__name__)
+
+
+def _set_trace(*, header=None):
+    ds = PdbWrapper(header)
+    ds.set_trace()
+
+
+class PdbWrapper(pdb.Pdb):
+    def __init__(self, header: Optional[str]):
+        self._actor = debugger.PdbActor()
+        self.header = header
+        super().__init__(
+            # pyre-ignore
+            stdout=WriteWrapper(self._actor),
+            stdin=ReadWrapper.create(self._actor),
+        )
+        self._first = True
+
+    def setup(self, *args, **kwargs):
+        r = super().setup(*args, **kwargs)
+        if self._first:
+            self._first = False
+            # when we enter the debugger, we want to present the user's stack frame
+            # not the nested one inside session.run. This means that the local
+            # variables are what gets printed, etc. To do this
+            # we first execute up 2 to get to that frame.
+            self.do_up(2)
+        return r
+
+    def set_continue(self) -> None:
+        r = super().set_continue()
+        if not self.breaks:
+            # no more breakpoints so this debugger will not
+            # be used again, and we detach from the controller io.
+            self._actor.send(DebuggerAction.Detach())
+            self._actor.drain_and_stop()
+            # break cycle with itself before we exit
+            self.stdin = sys.stdin
+            self.stdout = sys.stdout
+        return r
+
+    def set_trace(self):
+        self._actor.send(DebuggerAction.Paused())
+        message = self._actor.receive()
+        # we give the controller the option to ignore this request to debug
+        # by issuing a "detach" message immediately.
+        if isinstance(message, DebuggerAction.Detach):
+            return
+        elif isinstance(message, DebuggerAction.Attach):
+            pass
+        else:
+            raise RuntimeError(f"unexpected debugger message {message}")
+        if self.header:
+            self.message(self.header)
+        super().set_trace()
+
+    def set_quit(self):
+        self._actor.send(DebuggerAction.Detach())
+        self._actor.drain_and_stop()
+        super().set_quit()
+
+
+class ReadWrapper(io.RawIOBase):
+    def __init__(self, actor: debugger.PdbActor):
+        self._actor = actor
+
+    def readinto(self, b):
+        self._actor.send(DebuggerAction.Read(len(b)))
+        response = self._actor.receive()
+        if isinstance(response, DebuggerAction.Detach):
+            raise bdb.BdbQuit
+        assert isinstance(response, DebuggerAction.Write)
+        response = cast(DebuggerAction.Write, response)
+        payload = debugger.get_bytes_from_write_action(response)
+        assert len(payload) <= len(b)
+        b[: len(payload)] = payload
+        return len(payload)
+
+    def readable(self) -> bool:
+        return True
+
+    @classmethod
+    def create(cls, actor: debugger.PdbActor):
+        return io.TextIOWrapper(io.BufferedReader(cls(actor)))
+
+
+class WriteWrapper:
+    def __init__(self, actor: debugger.PdbActor):
+        self._actor = actor
+
+    def writable(self) -> bool:
+        return True
+
+    def write(self, s: str):
+        self._actor.send(DebuggerAction.Write(s.encode()))
+
+    def flush(self):
+        pass

monarch/_monarch/worker/logging.py

@@ -0,0 +1,107 @@
+# 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 logging
+import os
+import socket
+import sys
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+## NOTE THIS FILE IS A DIRECT COPY OF ~/fbsource/fbcode/monarch/python/monarch_supervisor/logging.py
+## It is copied here at this time to avoid pulling in the monarch python supervisor as dependency since there's
+## an expectation that the other one will be removed soon.
+
+
+def _handle_unhandled_exception(*args):
+    logger.error("Uncaught exception", exc_info=args)
+
+
+_glog_level_to_abbr = {
+    "DEBUG": "V",  # V is for VERBOSE in glog
+    "INFO": "I",
+    "WARNING": "W",
+    "ERROR": "E",
+    "CRITICAL": "C",
+}
+
+
+def fix_exception_lines(tb_lines):
+    formatted_lines = []
+    for line in tb_lines:
+        # Replace the standard file and line format with the custom format
+        if line.startswith("  File"):
+            # Extract the filename and line number
+            parts = line.split(",")
+            file_info = parts[0].strip()[6:-1]  # Remove '  File "' and '"'
+            line_info = parts[1].strip()[5:]  # Remove 'line '
+            new_line = f"  File {file_info}:{line_info}"
+            if len(parts) > 2:
+                new_line += ", " + ",".join(parts[2:]).strip()
+            formatted_lines.append(new_line)
+        else:
+            formatted_lines.append(line.strip())
+    return formatted_lines
+
+
+class _Formatter(logging.Formatter):
+    def __init__(self, suffix):
+        self.suffix = suffix
+
+    def format(self, record):
+        message = record.getMessage()
+        asctime = self.formatTime(record, "%m%d %H:%M:%S")
+
+        lines = message.strip().split("\n")
+        if record.exc_info:
+            exc_info = fix_exception_lines(
+                self.formatException(record.exc_info).split("\n")
+            )
+            lines.extend(exc_info)
+        if record.stack_info:
+            stack_info = self.formatStack(record.stack_info)
+            lines.extend(stack_info.strip().split("\n"))
+
+        shortlevel = _glog_level_to_abbr.get(record.levelname, record.levelname[0])
+
+        prefix = (
+            f"{shortlevel}{asctime}.{int(record.msecs*1000):06d} "
+            f"{record.filename}:"
+            f"{record.lineno}]{self.suffix}"
+        )
+        return "\n".join(f"{prefix} {line}" for line in lines)
+
+
+def initialize_logging(process_name=None):
+    log_folder = os.environ.get("TORCH_MONARCH_LOG_FOLDER")
+    log_level = os.environ.get("TORCH_MONARCH_LOG_LEVEL", "INFO")
+    suffix = "" if process_name is None else f" {process_name}:"
+    handler = None
+    if log_folder is not None:
+        log_folder_path = Path(log_folder)
+        log_folder_path.mkdir(parents=True, exist_ok=True)
+        safe_process_name = (
+            process_name.replace("/", "_") if process_name else "logfile.log"
+        )
+        log_file_name = f"{safe_process_name}.log"
+        log_file_path = log_folder_path / log_file_name
+        handler = logging.FileHandler(log_file_path)
+    else:
+        handler = logging.StreamHandler()
+    handler.setFormatter(_Formatter(suffix))
+    handler.setLevel(log_level)
+    logging.root.setLevel(log_level)
+    logging.root.addHandler(handler)
+    sys.excepthook = _handle_unhandled_exception
+
+
+def gethostname():
+    """Get the hostname of the machine."""
+    hostname = socket.gethostname()
+    hostname = hostname.replace(".facebook.com", "")
+    return hostname