torchmonarch-nightly 2025.6.27__cp313-cp313-manylinux2014_x86_64.whl

monarch/worker/compiled_block.py

@@ -0,0 +1,270 @@
+# 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
+from collections import defaultdict
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, TYPE_CHECKING
+
+import torch.fx
+from monarch.common.messages import DependentOnError
+from monarch.common.tree import tree_map
+from torch.fx.proxy import GraphAppendingTracer
+
+from .lines import Lines
+
+if TYPE_CHECKING:
+    from .worker import Cell, Stream
+
+logger = logging.getLogger(__name__)
+
+
+class Symbol:
+    def __init__(self, name: str):
+        self.name = name
+
+    def __repr__(self):
+        return self.name
+
+
+@dataclass
+class ErrorContext:
+    ident: Optional[int]
+    fallback_resume_offset: int
+
+
+class _BlockTracer(GraphAppendingTracer):
+    def __init__(self, ctx: "ErrorContext", graph: torch.fx.Graph):
+        super().__init__(graph)
+        self._error_context = ctx
+
+    def create_node(self, *args, **kwargs):
+        n = super().create_node(*args, **kwargs)
+        n.context = self._error_context
+        return n
+
+
+class CompiledBlock:
+    def __init__(self):
+        self.graphs: Dict["Stream", torch.fx.Graph] = defaultdict(
+            lambda: torch.fx.Graph()
+        )
+        self.used_formals: Set[int] = set()
+        self.used_results: Set[int] = set()
+        self.results: Dict[torch.fx.Node, int] = {}
+        self.fallback: Dict["Stream", List[Callable]] = defaultdict(list)
+        self.recording_stream: Optional["Stream"] = None
+        self.defined_borrows = {}  # dict not set to preserve order
+        self.defined_cells: Dict["Cell", int] = {}  # dict not set to preserve order
+        self.mutated_cells: Dict["Cell", "Stream"] = {}
+        self.current_context = ErrorContext(None, 0)
+        self.impls: Dict["Stream", Callable] = {}
+
+    def call_function(self, *args, **kwargs):
+        n = self.recording_graph.call_function(*args, **kwargs)
+        n.context = self.current_context
+        return n
+
+    def define_formal(self, stream: "Stream", argument_index: int):
+        self.used_formals.add(argument_index)
+        n = self.graphs[stream].call_module("formal", (argument_index,))
+        # pyre-ignore
+        n.context = self.current_context
+        return n
+
+    def define_result(self, node: torch.fx.Node, output_index: int):
+        self.used_results.add(output_index)
+        self.results[node] = output_index
+
+    def input_cell(self, cell: "Cell"):
+        n = self.recording_graph.call_module("input_cell", (cell,))
+        n.context = self.current_context
+        return n
+
+    def proxy(self, n: torch.fx.Node):
+        return torch.fx.Proxy(n, _BlockTracer(self.current_context, n.graph))
+
+    def mutates(self, results: Sequence["Cell"]):
+        for r in results:
+            if r not in self.defined_cells:
+                assert self.recording_stream is not None
+                self.mutated_cells[r] = self.recording_stream
+
+    @property
+    def recording_graph(self):
+        return self.graphs[self.recording_stream]
+
+    @contextmanager
+    def record_to(self, stream: "Stream"):
+        orig, self.recording_stream = self.recording_stream, stream
+        ctx = ErrorContext(None, len(self.fallback[stream]))
+        orig_context, self.current_context = self.current_context, ctx
+        try:
+            yield
+        finally:
+            self.recording_stream = orig
+            self.current_context = orig_context
+
+    def emit_stream(self, stream: "Stream"):
+        # Generated function looks like this:
+
+        # def fn(actuals: List["Cell"], outputs: List["Cell"]):
+        #  a, b, c, d, e, f, g, e = EXTERNAL # global variable bound to all the values we just want to bind to this code
+        #  try:
+        #     a = cell0.get()
+        #     b = actuals[0].get()
+        #     r = a + b
+        #     outputs[2].set(r)
+        #     t = r + r
+        #     r2 = r + t
+        #     outputs[4].set(r2)
+        #
+        #  except Exception as e:
+        #    # error recovery, fallback to
+        #    # interpreter code that can handle some values failing
+        #    return fallback(locals())
+
+        graph: torch.fx.Graph = self.graphs[stream]
+        fallback_functions = self.fallback[stream]
+
+        external: List[Any] = []
+        external_names: List[str] = []
+        external_id_to_name: Dict[int, str] = {}
+
+        def arg_map(x):
+            if isinstance(x, torch.fx.Node):
+                return x
+            elif id(x) in external_id_to_name:
+                return external_id_to_name[id(x)]
+            else:
+                candidate = getattr(x, "__name__", "external")
+                sym = Symbol(graph._graph_namespace.create_name(candidate, None))
+                external_names.append(sym.name)
+                external.append(x)
+                external_id_to_name[id(x)] = sym
+                return sym
+
+        lines = Lines()
+        body = Lines()
+
+        def fallback(results, exc):
+            lineno = exc.__traceback__.tb_lineno
+            error_context: ErrorContext = lines.get_context(lineno)
+            # report new errors and set
+            # defined identifiers for currently failing
+            # op to DependentOnError
+            if not isinstance(exc, DependentOnError):
+                if error_context is None or error_context.ident is None:
+                    raise exc
+                exc = stream.report_error(
+                    stream.current_recording,
+                    error_context.ident,
+                    exc,
+                )
+
+            # set exceptionson all the values this stream was responsible  for.
+            # this is the explicitly passed cell outputs, and all
+            # the cells we mutated.
+            for c, s in self.mutated_cells.items():
+                if s is stream:
+                    c.set(exc)
+            for r, i in self.results.items():
+                if r.graph is graph:
+                    results[i].set(exc)
+
+            # some ops we have to run despite errors such as
+            # borrows, collectives, send_tensor
+            # we run these universally here.
+            # Note that all of these are ok loading from cells with dependent
+            # on error status.
+            inst_range = range(
+                error_context.fallback_resume_offset, len(fallback_functions)
+            )
+            for inst in inst_range:
+                fallback_functions[inst]()
+
+        fallback_sym = arg_map(fallback)
+
+        # figure out the last use of each node that isn't
+        # live out, so that we appropriatelly `del` the variable.
+        seen = {r for r in self.results.keys() if r.graph is graph}
+        last_uses = defaultdict(list)
+        for node in reversed(graph.nodes):
+            for n in node.all_input_nodes:
+                if n not in seen:
+                    last_uses[node].append(n)
+                    seen.add(n)
+
+        # generate the repeat body
+        for node in graph.nodes:
+            if node.op == "call_module":
+                if node.target == "input_cell":
+                    # each input goes into the prologue where we issue a load from the
+                    # cell it came from.
+                    (cell_obj,) = node.args
+                    cell = arg_map(cell_obj)
+                    with body.context(node.context):
+                        body.emit(f"    {node.name} = {cell}.get()")
+                elif node.target == "formal":
+                    with body.context(node.context):
+                        (i,) = node.args
+                        body.emit(f"    {node.name} = actuals[{i}].get()")
+            else:
+                assert node.op == "call_function"
+                fn = arg_map(node.target)
+                args, kwargs = tree_map(arg_map, (node.args, node.kwargs))
+                all = [
+                    *(repr(a) for a in args),
+                    *(f"{k}={repr(v)}" for k, v in kwargs.items()),
+                ]
+                assign = ""
+                if node in seen:
+                    assign = f"{node.name} = "
+                with body.context(node.context):
+                    body.emit(f"    {assign}{fn}({', '.join(all)})")
+                    # some inputs to this node may no longer be used in the body
+                    # of the loop. We explicitly del them so their lifetime
+                    # is not longer than it was originally without compilation.
+                    to_delete = [repr(d) for d in last_uses[node]]
+                    if to_delete:
+                        body.emit(f"    del {', '.join(to_delete)}")
+        for r, i in self.results.items():
+            if r.graph is not graph:
+                continue
+            body.emit(f"    results[{i}].set({r})")
+
+        lines.emit("def impl(results, actuals):")
+        lines.emit(f"  {', '.join(external_names)} = EXTERNAL")
+        lines.emit("  _exception = None")
+        lines.emit("  try:")
+        lines.emit_lines(body)
+        lines.emit("  except Exception as e:")
+        lines.emit("    _exception = e")
+        # we do not call `fallback` inside of the exception block because we
+        # do not want future exceptions to have the stack trace of e attached.
+        lines.emit(
+            f"  if _exception is not None: return {fallback_sym}(results, _exception)"
+        )
+
+        gbls = {"EXTERNAL": external}
+        text = lines.text()
+        logger.debug(f"Compiled\n{text}")
+        exec(lines.text(), gbls)
+        return gbls["impl"]
+
+    def emit(self):
+        self.impls = {stream: self.emit_stream(stream) for stream in self.graphs.keys()}
+
+        # fallback functions for borrows/reduce/send read directly from these cells
+        # we need to make sure they are set to errors so that they work correctly.
+        # it always gets an error value
+        err = DependentOnError(-1)
+        for cell in self.defined_cells:
+            cell.set(err)
+
+        self.defined_cells.clear()

monarch/worker/debugger.py

@@ -0,0 +1,125 @@
+# 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 Optional, TYPE_CHECKING
+
+from monarch.common import messages
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from .worker import Stream, Worker
+
+_orig_set_trace = pdb.set_trace
+
+
+def _set_trace(*, header=None):
+    from .worker import _tls
+
+    stream = _tls.stream
+    if stream is None:
+        _orig_set_trace(header=header)
+    ds = PdbWrapper(stream, header)
+    ds.set_trace()
+
+
+class PdbWrapper(pdb.Pdb):
+    def __init__(self, stream: "Stream", header: Optional[str]):
+        self.stream = stream
+        self.worker: "Worker" = self.stream.worker
+        self.header = header
+        # pyre-ignore
+        super().__init__(stdout=WriteWrapper(self), stdin=ReadWrapper.create(self))
+        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._send("detach")
+            # break cycle with itself before we exit
+            self.stdin = sys.stdin
+            self.stdout = sys.stdout
+        return r
+
+    def _send(self, action):
+        self.worker.schedule(
+            lambda: self.worker.q.send(messages.DebuggerMessage(self.stream.id, action))
+        )
+
+    def set_trace(self):
+        self._send("paused")
+        message = self.stream.debugger_queue.get()
+        # we give the controller the option to ignore this request to debug
+        # by issuing a "detach" message immediately.
+        match message:
+            case "attach":
+                pass
+            case "detach":
+                return
+            case other:
+                raise RuntimeError(f"unexpected debugger message {other}")
+        if self.header:
+            self.message(self.header)
+        super().set_trace()
+
+
+class ReadWrapper(io.RawIOBase):
+    def __init__(self, session: "PdbWrapper"):
+        self.session = session
+
+    def readinto(self, b):
+        self.session._send(messages.DebuggerRead(len(b)))
+        response = self.session.stream.debugger_queue.get()
+        if response == "detach":
+            # this gets injected by the worker event loop to
+            # get the worker thread to exit on an Exit command.
+            raise bdb.BdbQuit
+        assert isinstance(response, messages.DebuggerWrite) and len(
+            response.payload
+        ) <= len(b)
+        b[: len(response.payload)] = response.payload
+        return len(response.payload)
+
+    def readable(self) -> bool:
+        return True
+
+    @classmethod
+    def create(cls, session: "PdbWrapper"):
+        return io.TextIOWrapper(io.BufferedReader(cls(session)))
+
+
+class WriteWrapper:
+    def __init__(self, session: "PdbWrapper"):
+        self.session = session
+
+    def writable(self) -> bool:
+        return True
+
+    def write(self, s: str):
+        self.session._send(messages.DebuggerWrite(s.encode()))
+
+    def flush(self):
+        pass

monarch/worker/lines.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-unsafe
+from contextlib import contextmanager
+from typing import Any, List
+
+
+class Lines:
+    """
+    Simple way to emit code where we track a per-line context object.
+    """
+
+    def __init__(self, context=None):
+        self._lines: List[str] = []
+        self._context: List[Any] = []
+        self._current_context = context
+
+    def get_context(self, lineno) -> Any:
+        return self._context[lineno - 1]
+
+    @contextmanager
+    def context(self, obj: Any):
+        old, self._current_context = self._current_context, obj
+        try:
+            yield
+        finally:
+            self._current_context = old
+
+    def emit(self, lines: str) -> None:
+        self._lines.extend(lines.split("\n"))
+        while len(self._context) < len(self._lines):
+            self._context.append(self._current_context)
+
+    def emit_lines(self, lines: "Lines") -> None:
+        """
+        Append another lines object on this one,
+        preserving its per-line context.
+        """
+        self._lines.extend(lines._lines)
+        self._context.extend(lines._context)
+
+    def text(self) -> str:
+        return "\n".join(self._lines)

monarch/worker/monitor.py

@@ -0,0 +1,53 @@
+# 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 math
+import queue
+import threading
+from typing import Callable, Optional, Tuple
+
+from monarch_supervisor import TTL
+
+
+class Monitor:
+    """A monitor is a thread that watches for reported events to expire."""
+
+    def __init__(self) -> None:
+        self.thread = threading.Thread(target=self._main, daemon=True, name="monitor")
+        self.events: queue.Queue[Tuple[Callable[[], None], Callable[[], float]]] = (
+            queue.Queue()
+        )
+        self.events.put((lambda: None, TTL(None)))
+
+    def start(self) -> None:
+        """Start the monitor thread."""
+        self.thread.start()
+
+    def _main(self) -> None:
+        debug, ttl = self.events.get()
+        while True:
+            try:
+                timeout = ttl()
+                next_debug, next_ttl = self.events.get(
+                    timeout=None if timeout == math.inf else timeout
+                )
+            except queue.Empty:
+                debug()
+                next_debug, next_ttl = self.events.get(timeout=None)
+
+            debug, ttl = next_debug, next_ttl
+
+    def __call__(
+        self,
+        debug_fn: Callable[[], None] = lambda: None,
+        timeout: Optional[float] = None,
+    ) -> None:
+        """Start a new event with the provided timeout.
+        If a timeout is specified, and a new event is not reported by before it expires,
+        the provided debug_fn is called."""
+        self.events.put((debug_fn, TTL(timeout)))