torchmonarch-nightly 2025.6.4__cp310-cp310-manylinux2014_x86_64.whl

monarch/common/tensor_factory.py

@@ -0,0 +1,31 @@
+# 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 typing import NamedTuple, Tuple
+
+import torch
+
+
+class TensorFactory(NamedTuple):
+    size: Tuple[int, ...]
+    dtype: torch.dtype
+    layout: torch.layout
+    device: torch.device
+
+    @staticmethod
+    def from_tensor(t):
+        return TensorFactory(t.size(), t.dtype, t.layout, t.device)
+
+    def empty(self):
+        return torch.empty(
+            self.size, dtype=self.dtype, layout=self.layout, device=self.device
+        )
+
+    def zeros(self):
+        return torch.full(
+            self.size, 0, dtype=self.dtype, layout=self.layout, device=self.device
+        )

monarch/common/tree.py

@@ -0,0 +1,73 @@
+# 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 typing import Any, Callable, Protocol, Sequence, Tuple
+
+import torch.utils._pytree as _pytree
+from torch.utils._pytree import (
+    _get_node_type,
+    register_pytree_node,
+    SUPPORTED_NODES,
+    tree_flatten,
+    tree_map,
+    tree_unflatten,
+)
+
+
+def flatten(tree, cond):
+    r, spec = tree_flatten(tree)
+
+    # be careful to not capture values we return in
+    # 'trues'. We do not need them to reconstruct and do not want to
+    # extend their lifetime.
+    trues = []
+    falses = []
+    conds = []
+    for e in r:
+        c = cond(e)
+        (trues if c else falses).append(e)
+        conds.append(c)
+
+    def unflatten(n):
+        n_it = iter(n)
+        falses_it = iter(falses)
+        return tree_unflatten([next(n_it if c else falses_it) for c in conds], spec)
+
+    return trues, unflatten
+
+
+def flattener(tree, cond=None):
+    """
+    Produce a _traceable_ flattener routine from tree. That is, it produces code that can
+    flatten another object shaped the same as tree, but whose structure cannot
+    be introspected because it might be (e.g.) an fx proxy value.
+    """
+    if isinstance(tree, (tuple, list)):
+        flattens = [flattener(t, cond) for t in tree]
+        return lambda obj: [
+            f for i, flatten in enumerate(flattens) for f in flatten(obj[i])
+        ]
+    elif isinstance(tree, dict):
+        keys = tuple(tree.keys())
+        flattens = [flattener(t, cond) for t in tree.values()]
+        return lambda obj: [
+            f for k, flatten in zip(keys, flattens) for f in flatten(obj[k])
+        ]
+    elif _get_node_type(tree) in SUPPORTED_NODES:
+        flatten_fn = SUPPORTED_NODES[_get_node_type(tree)].flatten_fn
+        trees, _ = flatten_fn(tree)
+        flattens = [flattener(t, cond) for t in trees]
+
+        def the_flattener(obj):
+            trees, _ = flatten_fn(obj)
+            return [f for i, flatten in enumerate(flattens) for f in flatten(trees[i])]
+
+        return the_flattener
+    elif cond is None or cond(tree):
+        return lambda obj: [obj]
+    else:
+        return lambda obj: []

monarch/controller/__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-strict

monarch/controller/backend.py

@@ -0,0 +1,223 @@
+# 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
+
+from abc import ABC, abstractmethod
+from typing import List, NamedTuple, Optional, Sequence, Tuple
+
+from monarch.common import messages
+
+from monarch.common.shape import iter_ranks, Slices as Ranks
+from monarch_supervisor import (
+    Context,
+    FunctionCall,
+    Host,
+    Process,
+    ProcessExited as ProcessExitedMsg,
+)
+from torch.distributed import TCPStore
+
+
+logger = logging.getLogger(__name__)
+
+
+class Backend(ABC):
+    @abstractmethod
+    def send(self, ranks: Ranks, msg) -> None:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def recvready(self, timeout: Optional[float]) -> Sequence[Tuple[int, NamedTuple]]:
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def world_size(self):
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def gpu_per_host(self):
+        raise NotImplementedError()
+
+
+class ProcessBackend(Backend):
+    def __init__(
+        self,
+        ctx: Context,
+        hosts: List[Host],
+        gpu_per_host: int,
+        _processes=None,
+        _store=None,
+    ):
+        self.ctx = ctx
+        self.hosts = hosts
+        self.store = self._create_store() if _store is None else _store
+        self._gpu_per_host = gpu_per_host
+        self.worker_processes = (
+            self._create_pg(ctx, hosts, gpu_per_host, self.store)
+            if _processes is None
+            else _processes
+        )
+        self.exiting = False
+        self.process_to_rank = {p: p.rank for p in self.worker_processes}
+        self.live_processes_per_rank: List[List[Process]] = [
+            [p] for p in self.worker_processes
+        ]
+
+    @property
+    def world_size(self):
+        return len(self.worker_processes)
+
+    @property
+    def gpu_per_host(self) -> int:
+        return self._gpu_per_host
+
+    def send(self, ranks: Ranks, msg) -> None:
+        handler = getattr(self, msg.__class__.__name__, None)
+        if handler is not None:
+            handler(ranks, msg)
+        self._send(ranks, msg)
+
+    def _send(self, ranks: Ranks, msg):
+        # the intent is for this to be optimized as tree broadcast
+        # base on if members of tree nodes overlap with a slice.
+        for rank in iter_ranks(ranks):
+            self.worker_processes[rank].send(msg)
+
+    def CommandGroup(self, ranks: Ranks, msg: messages.CommandGroup):
+        for command in msg.commands:
+            handler = getattr(self, command.__class__.__name__, None)
+            if handler is not None:
+                handler(ranks, command)
+
+    def CreatePipe(self, ranks: Ranks, msg: messages.CreatePipe):
+        pipe_ranks = list(enumerate(iter_ranks(ranks)))
+        for i, rank in pipe_ranks:
+            # In general, pipes on different workers may need to have different behavior.
+            # For example, two data loader pipes operating on the same dataset should
+            # load different shards of the dataset. In order to do this, each pipe process
+            # on the worker needs to know the number of instances of the pipe (e.g. len(pipe_ranks))
+            # and its unique rank among all instances of the pipe (e.g., i).
+            proc = self.worker_processes[rank].host.create_process(
+                FunctionCall(
+                    "monarch.worker.worker.pipe_main",
+                    f"{msg.key}-{rank}",
+                    msg.max_messages,
+                ),
+                env={"CUDA_VISIBLE_DEVICES": ""},
+                name=f"pipe-{rank}",
+            )
+            self.live_processes_per_rank[rank].append(proc)
+            self.process_to_rank[proc] = rank
+
+    def ProcessExited(
+        self, sender: Process, msg: ProcessExitedMsg
+    ) -> List[Tuple[int, NamedTuple]]:
+        return self._process_exited(sender, msg.result)
+
+    def Restarted(
+        self, sender: Process, restarted: messages.Restarted
+    ) -> List[Tuple[int, NamedTuple]]:
+        return self._process_exited(sender, restarted.result)
+
+    def _process_exited(
+        self, sender: Process, result: int | Exception
+    ) -> List[Tuple[int, NamedTuple]]:
+        rank = self.process_to_rank[sender]
+        if result != 0:
+            if not self.exiting or self.worker_processes[rank] is sender:
+                kind = (
+                    "worker"
+                    if self.worker_processes[rank] is sender
+                    else "pipe_process"
+                )
+                raise RuntimeError(f"Unexpected {kind} exit on rank {rank}")
+
+        live_procs = self.live_processes_per_rank[rank]
+        live_procs.remove(sender)
+        if len(live_procs) == 0:
+            return [(rank, ProcessExitedMsg(0))]
+        return []
+
+    def Exit(self, ranks: Ranks, msg: messages.Exit):
+        self.exiting = True
+        for rank in iter_ranks(ranks):
+            # ideally we are more kind to these processes.
+            # but first we need to develop the API for asking them
+            # to suspend, restore, fast forward, rewind, etc.
+            worker = self.worker_processes[rank]
+            for proc in self.live_processes_per_rank[rank]:
+                if worker is not proc:
+                    proc.signal()
+            self.worker_processes[rank].send(msg)
+
+    def recvready(self, timeout: Optional[float]) -> Sequence[Tuple[int, NamedTuple]]:
+        result = []
+        for sender, msg in self.ctx.recvready(timeout):
+            handler = getattr(self, msg.__class__.__name__, None)
+            if handler is not None:
+                result.extend(handler(sender, msg))
+                continue
+            elif isinstance(sender, Process):
+                result.append((sender.rank, msg))
+            else:
+                logger.warning("TODO: ignoring non-worker message: %s %s", sender, msg)
+        return result
+
+    @staticmethod
+    def _create_store():
+        if os.environ.get("INSIDE_RE_WORKER"):
+            hostname = "localhost"
+        else:
+            hostname = socket.gethostname()
+
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as sock:
+            sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
+            sock.bind(("::", 0))
+            port = sock.getsockname()[1]
+            store = TCPStore(
+                hostname,
+                port,
+                is_master=True,
+                use_libuv=False,
+                master_listen_fd=sock.detach(),
+            )
+        return store
+
+    @staticmethod
+    def _create_pg(
+        ctx: Context, hosts: List[Host], gpu_per_host: int, store, _restartable=False
+    ):
+        env = {
+            # cuda event cache disabled pending fix for:
+            # https://github.com/pytorch/pytorch/issues/143470
+            "TORCH_NCCL_CUDA_EVENT_CACHE": "0",
+            # disable nonblocking comm until D68727854 lands.
+            "TORCH_NCCL_USE_COMM_NONBLOCKING": "0",
+            # supervisor_pipe is a unique ID per Host object,
+            # so it lets us put multiple processes on the same GPU.
+            "NCCL_HOSTID": "$SUPERVISOR_PIPE",
+            "STORE_HOSTNAME": store.host,
+            "STORE_PORT": str(store.port),
+        }
+        for name, value in os.environ.items():
+            if name.startswith("NCCL_") and name not in env:
+                env[name] = value
+        return ctx.create_process_group(
+            hosts,
+            FunctionCall(
+                "monarch.worker.worker.worker_main", _restartable=_restartable
+            ),
+            processes_per_host=gpu_per_host,
+            env=env,
+            name="worker",
+        )

monarch/controller/controller.py

@@ -0,0 +1,223 @@
+# 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 traceback
+from collections import deque
+from typing import Generator, List, NamedTuple, Optional, Sequence, Tuple, Union
+
+from monarch._rust_bindings.monarch_extension.client import (  # @manual=//monarch/monarch_extension:monarch_extension
+    DebuggerMessage,
+    WorldState,
+)
+
+from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ActorId,
+)
+
+from monarch.common import messages
+from monarch.common.controller_api import LogMessage, MessageResult
+from monarch.common.invocation import DeviceException, Seq
+from monarch.common.reference import Ref
+from monarch.common.shape import NDSlice
+from monarch.common.tensor import Tensor
+from monarch.controller import debugger
+
+from .backend import Backend
+from .history import History
+
+logger = logging.getLogger(__name__)
+
+
+class Controller:
+    def __init__(self, backend: Backend):
+        self._backend = backend
+        self._history = History(backend.world_size)
+        self._messages = deque()
+
+        self.exited = {}
+        self.active_debugger: Optional[Tuple[int, int]] = None
+        self.pending_debugger_sessions: deque[Tuple[int, int]] = deque()
+        # for current active session
+        self.pending_debugger_messages: deque[messages.DebuggerMessage] = deque()
+
+    def send(
+        self,
+        ranks: Union[NDSlice, List[NDSlice]],
+        msg: NamedTuple,
+    ) -> None:
+        self._backend.send(ranks, msg)
+
+    def next_message(
+        self, timeout: Optional[float]
+    ) -> Optional[MessageResult | LogMessage]:
+        if len(self._messages) == 0:
+            self._messages.extend(self._read_messages(timeout))
+        return self._messages.popleft() if len(self._messages) > 0 else None
+
+    def drop_refs(self, refs: Sequence[Ref]) -> None:
+        """
+        noop as this is used for the Rust controller to know when to gc invocations_for_ref for failed invocations
+        """
+        pass
+
+    def _read_messages(
+        self, timeout: Optional[float]
+    ) -> Generator[MessageResult, None, None]:
+        # XXX - how can we avoid always requesting status when waiting on futures?
+        #       we need to figure out what submesh we need to hear from before a future
+        #       is considered 'good'. This means not just waiting for the future value
+        #       but also for signal that any failures that could invalidate the future have
+        #       not happened. We could do better if tensors/collectives had an invalid bit
+        #       that we propagate. In real uses fetches might lag behind anyway so we would not
+        #       have to send out so many requests for current status.
+        for rank, value in self._backend.recvready(timeout):
+            yield from self._handle_message(rank, value)
+
+    def drain_and_stop(self) -> List[MessageResult | LogMessage | DebuggerMessage]:
+        messages = []
+        while self._messages:
+            messages.append(self._messages.popleft())
+        while len(self.exited) < self._backend.world_size:
+            messages.extend(self._read_messages(None))
+        return messages
+
+    def stop_mesh(self) -> None:
+        pass
+
+    def node(
+        self,
+        seq: Seq,
+        defs: Sequence["Tensor"],
+        uses: Sequence["Tensor"],
+    ) -> None:
+        self._history.ident(seq, defs, uses)
+
+    def _handle_message(self, sender, value) -> Generator[MessageResult, None, None]:
+        yield from getattr(self, value.__class__.__name__)(sender, *value)
+
+    def worker_world_state(self) -> WorldState:
+        # Eventhough not implemented, return needed so return value complies with type checking
+        assert 1 == 2, "not implemented"
+        return WorldState()
+
+    def ProcessExited(self, proc, result) -> Generator[MessageResult, None, None]:
+        if result != 0:
+            # XXX - this should start the failure recovery process
+            raise RuntimeError("Unexpected worker process exit")
+        self.exited[proc] = result
+        yield from []
+
+    def ProcessStarted(self, proc, pid) -> Generator[MessageResult, None, None]:
+        yield from []
+
+    def FetchResult(self, proc, ident, value) -> Generator[MessageResult, None, None]:
+        self._history.future_completed(ident, value)
+        yield from []
+
+    def RemoteFunctionFailed(
+        self,
+        proc,
+        failing_ident,
+        traceback_index,
+        exception: Exception,
+        worker_frames: List[traceback.FrameSummary],
+    ) -> Generator[MessageResult, None, None]:
+        self._history.propagate_failure(
+            failing_ident, traceback_index, exception, worker_frames
+        )
+        yield from self._history.rank_completed(proc, failing_ident)
+
+    def InternalException(
+        self,
+        proc,
+        exception: Exception,
+        worker_frames: List[traceback.FrameSummary],
+    ) -> Generator[MessageResult, None, None]:
+        yield MessageResult(
+            seq=0,  # will not be used
+            result=None,
+            error=DeviceException(
+                exception,
+                worker_frames,
+                ActorId.from_string("unknown[0].unknown[0]"),
+                message="A worker experienced an internal error.",
+            ),
+        )
+
+    def RemoteGeneratorFailed(
+        self,
+        proc,
+        exception: Exception,
+        frames: List[traceback.FrameSummary],
+    ) -> Generator[MessageResult, None, None]:
+        yield MessageResult(
+            seq=0,  # will not be used
+            result=None,
+            error=DeviceException(
+                exception=exception,
+                frames=frames,
+                source_actor_id=ActorId.from_string("unknown[0].unknown[0]"),
+                message="A remote generator failed.",
+            ),
+        )
+
+    def Status(
+        self, proc, first_uncompleted_ident
+    ) -> Generator[MessageResult, None, None]:
+        yield from self._history.rank_completed(proc, first_uncompleted_ident)
+
+    def DebuggerMessage(
+        self, proc, stream_id: int, action
+    ) -> Generator[MessageResult, None, None]:
+        if action == "paused":
+            self.pending_debugger_sessions.append((proc, stream_id))
+        else:
+            assert self.active_debugger == (proc, stream_id)
+            self.pending_debugger_messages.append(action)
+
+        if self.active_debugger is None:
+            yield from self._run_debugger_loop()
+
+    def _run_debugger_loop(self) -> Generator[MessageResult, None, None]:
+        # debug loop
+        while self.pending_debugger_sessions:
+            yield from self._run_debugger_session(
+                *self.pending_debugger_sessions.popleft()
+            )
+
+    def _run_debugger_session(
+        self, proc_id: int, stream_id: int
+    ) -> Generator[MessageResult, None, None]:
+        debugger.write(f"pdb attached to rank {proc_id}, stream {stream_id}\n")
+        self.active_debugger = (proc_id, stream_id)
+        try:
+            rank = NDSlice(offset=proc_id, sizes=[], strides=[])
+            self.send(rank, messages.DebuggerMessage(stream_id, "attach"))
+            while True:
+                while not self.pending_debugger_messages:
+                    # todo: eventually we should timeout
+                    yield from self._read_messages(None)
+                message = self.pending_debugger_messages.popleft()
+                match message:
+                    case "detach":
+                        break
+                    case messages.DebuggerRead(requested):
+                        self.send(
+                            rank,
+                            messages.DebuggerMessage(
+                                stream_id,
+                                messages.DebuggerWrite(debugger.read(requested)),
+                            ),
+                        )
+                    case messages.DebuggerWrite(payload):
+                        debugger.write(payload.decode())
+                    case other:
+                        raise RuntimeError(f"unexpected debugger message: {other}")
+        finally:
+            self.active_debugger = None
+            self.pending_debugger_messages.clear()

monarch/controller/debugger.py

@@ -0,0 +1,47 @@
+# 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 sys
+from typing import Optional
+
+_is_ipython: Optional[bool] = None
+
+
+def is_ipython() -> bool:
+    global _is_ipython
+    if _is_ipython is not None:
+        return _is_ipython
+    try:
+        from IPython import get_ipython
+
+        _is_ipython = get_ipython() is not None
+    except ImportError:
+        _is_ipython = False
+    return _is_ipython
+
+
+def write(msg: str) -> None:
+    sys.stdout.write(msg)
+    sys.stdout.flush()
+
+
+def read(requested_size: int) -> bytes:
+    if not is_ipython():
+        b = bytearray(requested_size)
+        bytes_read = sys.stdin.buffer.raw.readinto(b)
+        return bytes(b[:bytes_read])
+
+    # ipython doesn't have stdin directly connected
+    # so we need to use input() instead.
+    user_input = input() + "\n"
+    input_bytes = user_input.encode("utf-8")
+    num_bytes_to_write = len(input_bytes)
+    if requested_size < num_bytes_to_write:
+        raise RuntimeError(
+            f"Debugger input line too long, max length is {requested_size}"
+        )
+    return input_bytes[:num_bytes_to_write]

monarch/controller/history.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.
+
+# pyre-unsafe
+from collections import deque
+from typing import Generator, Sequence, TYPE_CHECKING
+
+from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ActorId,
+)
+
+from monarch.common.controller_api import MessageResult
+
+from monarch.common.invocation import Invocation, RemoteException, Seq
+
+if TYPE_CHECKING:
+    from monarch.common.tensor import Tensor
+
+
+class History:
+    def __init__(self, N):
+        self.first_uncompleted_ident = [0 for _ in range(N)]
+        self.min_first_uncompleted_ident = 0
+        self.invocations = deque[Invocation]()
+
+    def _invocation(
+        self,
+        seq: Seq,
+        defs: Sequence["Tensor"],
+        uses: Sequence["Tensor"],
+    ):
+        r = Invocation(seq)
+        for t in uses:
+            u = t._invocation
+            assert u is not None
+            u.add_user(r)
+        for t in defs:
+            t._invocation = r
+        return r
+
+    def ident(
+        self,
+        seq: Seq,
+        defs: Sequence["Tensor"],
+        uses: Sequence["Tensor"],
+    ):
+        invocation = self._invocation(seq, defs, uses)
+        self.invocations.append(invocation)
+
+    def propagate_failure(self, seq, traceback_index, exception, worker_frames):
+        invocation = self.invocations[seq - self.min_first_uncompleted_ident]
+        remote_exception = RemoteException(
+            seq,
+            exception,
+            traceback_index,
+            None,
+            worker_frames,
+            ActorId.from_string("unknown[0].unknown[0]"),
+        )
+        worklist = deque((invocation,))
+        while worklist:
+            invocation = worklist.popleft()
+            if invocation.fail(remote_exception):
+                worklist.extend(invocation.users)
+
+    def rank_completed(
+        self, rank, first_uncompleted_ident
+    ) -> Generator[MessageResult, None, None]:
+        # advance what our last completed action was, and
+        # trim the list of tracebacks if we no longer need them.
+        prev = self.first_uncompleted_ident[rank]
+        self.first_uncompleted_ident[rank] = first_uncompleted_ident
+        if prev == self.min_first_uncompleted_ident:
+            self.min_first_uncompleted_ident = min(self.first_uncompleted_ident)
+            for seq in range(prev, self.min_first_uncompleted_ident):
+                invocation = self.invocations.popleft()
+                assert seq == invocation.seq
+                result, error = invocation.complete()
+                yield MessageResult(
+                    seq=seq,
+                    result=result,
+                    error=error,
+                )
+
+    def future_completed(self, ident, value):
+        invocation = self.invocations[ident - self.min_first_uncompleted_ident]
+        invocation.fut_value = value