torchmonarch-nightly 2025.6.27__cp311-cp311-manylinux2014_x86_64.whl

monarch/common/function.py

@@ -0,0 +1,160 @@
+# 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
+import sys
+import warnings
+from logging import Logger
+
+# pyre-ignore
+from pickle import _getattribute, PickleError, whichmodule
+from types import BuiltinFunctionType, FunctionType
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    NamedTuple,
+    Optional,
+    Protocol,
+    runtime_checkable,
+)
+
+import cloudpickle
+
+logger: Logger = logging.getLogger(__name__)
+
+
+@runtime_checkable
+class ResolvableFunction(Protocol):
+    def resolve(self) -> Callable: ...
+
+
+ConvertsToResolvable = Any
+
+
+def _string_resolver(arg: Any) -> Optional[ResolvableFunction]:
+    if isinstance(arg, str) and "." in arg:
+        return ResolvableFunctionFromPath(arg)
+
+
+def _torch_resolver(arg: Any) -> Optional[ResolvableFunction]:
+    import torch
+
+    if isinstance(arg, torch._ops.OpOverload):
+        return ResolvableFunctionFromPath("torch.ops." + str(arg))
+
+
+def function_to_import_path(arg: BuiltinFunctionType | FunctionType) -> Optional[str]:
+    # code replicated from pickler to check if we
+    # would successfully be able to pickle this function.
+    name = getattr(arg, "__qualname__", None)
+    if name is None:
+        name = arg.__name__
+    try:
+        # pyre-ignore
+        module_name = whichmodule(arg, name)
+        __import__(module_name, level=0)
+        module = sys.modules[module_name]
+        if module_name == "__main__":
+            return None  # the workers will not have the same main
+
+        # pytest installs its own custom loaders that do not
+        # survive process creation
+        try:
+            if "pytest" in module.__loader__.__class__.__module__:
+                return None
+        except AttributeError:
+            pass
+
+        # pyre-ignore
+        obj2, parent = _getattribute(module, name)
+        # support annotations that cover up the global impl
+        if obj2 is arg or getattr(obj2, "_remote_impl", None) is arg:
+            return f"{module_name}.{name}"
+    except (PickleError, ImportError, KeyError, AttributeError):
+        pass
+    return None
+
+
+def _function_resolver(arg: Any):
+    if isinstance(arg, (FunctionType, BuiltinFunctionType)):
+        if path := function_to_import_path(arg):
+            return ResolvableFunctionFromPath(path)
+
+
+def _cloudpickle_resolver(arg: Any):
+    # @lint-ignore PYTHONPICKLEISBAD
+    return ResolvableFromCloudpickle(cloudpickle.dumps(arg))
+
+
+resolvers = [
+    _torch_resolver,
+    _string_resolver,
+    _function_resolver,
+    _cloudpickle_resolver,
+]
+
+
+_cached_resolvers = {}
+
+
+def maybe_resolvable_function(arg: Any) -> Optional[ResolvableFunction]:
+    if arg == "__test_panic":
+        return ResolvableFunctionFromPath("__test_panic")
+    r = _cached_resolvers.get(arg)
+    if r is not None:
+        return r
+    for resolver in resolvers:
+        r = resolver(arg)
+        if r is not None:
+            _cached_resolvers[arg] = r
+            return r
+    return None
+
+
+def resolvable_function(arg: ConvertsToResolvable) -> ResolvableFunction:
+    if isinstance(arg, ResolvableFunction):
+        return arg
+    r = maybe_resolvable_function(arg)
+    if r is None:
+        raise ValueError(f"Unsupported target for a remote call: {arg!r}")
+    return r
+
+
+class ResolvableFunctionFromPath(NamedTuple):
+    path: str
+
+    def resolve(self):
+        first, *parts = self.path.split(".")
+        if first == "torch":
+            function = importlib.import_module("torch")
+            for p in parts:
+                function = getattr(function, p)
+            assert isinstance(function, Callable)
+        else:
+            modulename, funcname = self.path.rsplit(".", 1)
+            module = importlib.import_module(modulename)
+            function = getattr(module, funcname)
+            # support annotations that cover up the global impl
+            actual = getattr(function, "_remote_impl", None)
+            return function if actual is None else actual
+        return function
+
+    def __str__(self):
+        return self.path
+
+
+class ResolvableFromCloudpickle(NamedTuple):
+    data: bytes
+
+    def resolve(self):
+        # @lint-ignore PYTHONPICKLEISBAD
+        return cloudpickle.loads(self.data)
+
+
+Propagator = Any

monarch/common/function_caching.py

@@ -0,0 +1,164 @@
+# 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 itertools
+from typing import Any, Dict, Hashable, List, NamedTuple, Optional, Sequence, Tuple
+
+import torch
+from torch import autograd
+from torch.utils._pytree import tree_flatten, TreeSpec
+
+
+class AliasOf(NamedTuple):
+    group: int  # 0 -this group, -1 - the parent, -2 - parent's parent, etc.
+    offset: int
+
+
+class Storage(NamedTuple):
+    numel: int
+
+
+# Hashable pattern for recreating tensors
+# Each tensor either creates its own Storage
+# or is an AliasOf another tensor either earlier in this list,
+# or in one of the parent lists.
+# parent lists are used to represent other collections of tensors
+# for instance if this pattern is for outputs of a function
+# parents might contains lists of inputs to the function and captured
+# globals as two separate lists.
+class TensorGroupPattern(NamedTuple):
+    entries: Tuple["PatternEntry", ...]
+
+    def empty(self, parents: List[List[torch.Tensor]]) -> List[torch.Tensor]:
+        tensors = []
+        for entry in self.entries:
+            match entry.storage:
+                case AliasOf(group=group, offset=offset):
+                    base = tensors[offset] if group == 0 else parents[group][offset]
+                case Storage(numel=numel):
+                    base = torch.empty(
+                        (numel,),
+                        dtype=entry.dtype,
+                        layout=entry.layout,
+                        device=entry.device,
+                    )
+                case _:
+                    raise ValueError("unexpected storage")
+            t = torch.as_strided(base, entry.size, entry.stride, entry.storage_offset)
+            tensors.append(t)
+        return tensors
+
+
+class PatternEntry(NamedTuple):
+    size: Tuple[int, ...]
+    stride: Tuple[int, ...]
+    storage_offset: int
+    dtype: torch.dtype
+    layout: torch.layout
+    device: torch.device
+    storage: AliasOf | Storage
+
+
+# Takes a list of tensors and computes the pattern of aliasing that
+# would reconstruct the group. If `parent` is specified aliases
+# are also computed with respect to that group and its parents.
+# new storage is only specified is a tensor's storage was not
+# seen in any parent or previously in a group.
+class TensorGroup:
+    def __init__(
+        self,
+        tensors: Sequence[torch.Tensor],
+        parent: Optional["TensorGroup"] = None,
+    ):
+        self.parent = parent
+        self.tensors = tensors
+        self.storage_dict: Dict[torch.UntypedStorage, int] = {}
+
+        def create_entry(i: int, t: torch.Tensor):
+            storage = t.untyped_storage()
+            numel = t.untyped_storage().size() // t.element_size()
+            alias = self._find_alias(storage)
+            if alias is None:
+                self.storage_dict[storage] = i
+                alias = Storage(numel)
+
+            return PatternEntry(
+                tuple(t.size()),
+                tuple(t.stride()),
+                int(t.storage_offset()),
+                t.dtype,
+                t.layout,
+                t.device,
+                alias,
+            )
+
+        self.pattern = TensorGroupPattern(
+            tuple(create_entry(i, t) for i, t in enumerate(tensors))
+        )
+
+    def _find_alias(self, storage: torch.UntypedStorage) -> Optional[AliasOf]:
+        grp = self
+        for i in itertools.count():
+            if storage in grp.storage_dict:
+                return AliasOf(-i, grp.storage_dict[storage])
+            if grp.parent is None:
+                return None
+            grp = grp.parent
+
+
+class TensorPlaceholder:
+    pass
+
+
+# singleton to represent where tensors go in a pytree
+tensor_placeholder = TensorPlaceholder()
+
+
+def _to_placeholder(x):
+    if isinstance(x, torch.Tensor):
+        return tensor_placeholder
+    return x
+
+
+def _remove_ctx(x):
+    if isinstance(x, autograd.function.FunctionCtx):
+        return None
+    return x
+
+
+# customizable set of filters to handle data types that appear
+# in functions that one wants to support in cached functions
+key_filters = [_to_placeholder, _remove_ctx]
+
+
+def _filter_key(v: Any):
+    for filter in key_filters:
+        v = filter(v)
+    return v
+
+
+class HashableTreeSpec(NamedTuple):
+    type: Any
+    context: Any
+    children_specs: Tuple["HashableTreeSpec", ...]
+
+    @staticmethod
+    def from_treespec(t: "TreeSpec"):
+        return HashableTreeSpec(
+            t.type,
+            tuple(t.context) if isinstance(t.context, list) else t.context,
+            tuple(HashableTreeSpec.from_treespec(child) for child in t.children_specs),
+        )
+
+
+def hashable_tensor_flatten(args, kwargs) -> Tuple[List[torch.Tensor], Hashable]:
+    values, spec = tree_flatten((args, kwargs))
+    tensors = [t for t in values if isinstance(t, torch.Tensor)]
+    key: Hashable = (
+        tuple(_filter_key(v) for v in values),
+        HashableTreeSpec.from_treespec(spec),
+    )
+    return tensors, key

monarch/common/future.py

@@ -0,0 +1,168 @@
+# 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 math
+import os
+import subprocess
+from typing import (
+    Any,
+    Callable,
+    cast,
+    Generic,
+    Optional,
+    Sequence,
+    TYPE_CHECKING,
+    TypeVar,
+)
+
+from monarch_supervisor import TTL
+
+if TYPE_CHECKING:
+    from monarch.common.client import Client
+
+from .invocation import RemoteException
+
+logger = logging.getLogger(__name__)
+
+try:
+    PYSPY_REPORT_INTERVAL: Optional[float] = float(
+        os.environ["CONTROLLER_PYSPY_REPORT_INTERVAL"]
+    )
+except KeyError:
+    PYSPY_REPORT_INTERVAL = None
+
+
+def _split(elems, cond):
+    trues = []
+    falses = []
+    for elem in elems:
+        if cond(elem):
+            trues.append(elem)
+        else:
+            falses.append(elem)
+    return trues, falses
+
+
+def _periodic_TTL(interval: Optional[float]) -> Callable[[], float]:
+    if interval is None:
+        return lambda: math.inf
+
+    ttl = TTL(interval)
+
+    def _remaining():
+        nonlocal ttl
+        rem = ttl()
+        if rem == 0:
+            ttl = TTL(interval)
+        return rem
+
+    return _remaining
+
+
+T = TypeVar("T")
+
+
+class Future(Generic[T]):
+    def __init__(self, client: "Client"):
+        self._client = client
+        self._status = "incomplete"
+        self._callbacks = None
+        self._result: T | Exception | None = None
+
+    def _set_result(self, r):
+        assert self._status == "incomplete"
+        self._result = r
+        self._status = "exception" if isinstance(r, RemoteException) else "complete"
+        if self._callbacks:
+            for cb in self._callbacks:
+                try:
+                    cb(self)
+                except Exception:
+                    logger.exception("exception in controller's Future callback")
+        self._callbacks = None
+        self._client = None
+
+    def _wait(self, timeout: Optional[float]):
+        if self._status != "incomplete":
+            return True
+
+        assert self._client is not None
+
+        # see if the future is done already
+        # and we just haven't processed the messages
+        while self._client.handle_next_message(0):
+            if self._status != "incomplete":
+                return True
+
+        ttl = TTL(timeout)
+        ttl_pyspy = _periodic_TTL(PYSPY_REPORT_INTERVAL)
+        while self._status == "incomplete" and _wait(self._client, ttl, ttl_pyspy):
+            ...
+
+        return self._status != "incomplete"
+
+    def result(self, timeout: Optional[float] = None) -> T:
+        if not self._wait(timeout):
+            raise TimeoutError()
+        if self._status == "exception":
+            raise cast(Exception, self._result)
+        return cast(T, self._result)
+
+    def done(self) -> bool:
+        return self._wait(0)
+
+    def exception(self, timeout: Optional[float] = None):
+        if not self._wait(timeout):
+            raise TimeoutError()
+        return self._result if self._status == "exception" else None
+
+    def add_callback(self, callback):
+        if not self._callbacks:
+            self._callbacks = [callback]
+        else:
+            self._callbacks.append(callback)
+
+
+def _wait(client: "Client", ttl: Callable[[], float], ttl_pyspy: Callable[[], float]):
+    remaining = ttl()
+    pyspy_remaining = ttl_pyspy()
+    if pyspy_remaining == 0:
+        try:
+            logging.warning(
+                f"future has not finished in {PYSPY_REPORT_INTERVAL} seconds (remaining time to live is {remaining}), py-spying process to debug."
+            )
+            subprocess.run(["py-spy", "dump", "-s", "-p", str(os.getpid())])
+        except FileNotFoundError:
+            logging.warning("py-spy is not installed.")
+    timeout = min(remaining, pyspy_remaining)
+    client.handle_next_message(timeout=None if timeout == math.inf else timeout)
+    return remaining > 0
+
+
+def stream(futures: Sequence[Future], timeout: Optional[float] = None):
+    """Stream the provided futures as they complete.
+
+    If a timeout is provided, it applies to the completion of the entire set of futures.
+    """
+    assert len(futures) > 0
+
+    ttl = TTL(timeout)
+    pyspy_ttl = _periodic_TTL(PYSPY_REPORT_INTERVAL)
+
+    assert (
+        len({f._client for f in futures if f._client is not None}) <= 1
+    ), "all futures must be from the same controller"
+
+    todo = futures
+    while True:
+        done, todo = _split(todo, lambda f: f._status != "incomplete")
+        for f in done:
+            yield f
+
+        if len(todo) == 0 or not _wait(todo[0]._client, ttl, pyspy_ttl):
+            break

monarch/common/invocation.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 traceback
+from typing import Any, List, Optional, Tuple
+
+from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ActorId,
+)
+
+
+Seq = int
+
+
+class DeviceException(Exception):
+    """
+    Non-deterministic failure in the underlying worker, controller or its infrastructure.
+    For example, a worker may enter a crash loop, or its GPU may be lost
+    """
+
+    def __init__(
+        self,
+        exception: Exception,
+        frames: List[traceback.FrameSummary],
+        source_actor_id: ActorId,
+        message: str,
+    ):
+        self.exception = exception
+        self.frames = frames
+        self.source_actor_id = source_actor_id
+        self.message = message
+
+    def __str__(self):
+        try:
+            exe = str(self.exception)
+            worker_tb = "".join(traceback.format_list(self.frames))
+            return (
+                f"{self.message}\n"
+                f"Traceback of the failure on worker (most recent call last):\n{worker_tb}{type(self.exception).__name__}: {exe}"
+            )
+        except Exception as e:
+            print(e)
+            return "oops"
+
+
+class RemoteException(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,
+        seq: Seq,
+        exception: Exception,
+        controller_frame_index: Optional[int],
+        controller_frames: Optional[List[traceback.FrameSummary]],
+        worker_frames: List[traceback.FrameSummary],
+        source_actor_id: ActorId,
+        message="A remote function has failed asynchronously.",
+    ):
+        self.exception = exception
+        self.worker_frames = worker_frames
+        self.message = message
+        self.seq = seq
+        self.controller_frame_index = controller_frame_index
+        self.source_actor_id = source_actor_id
+        self.controller_frames = controller_frames
+
+    def __str__(self):
+        try:
+            exe = str(self.exception)
+            worker_tb = "".join(traceback.format_list(self.worker_frames))
+            controller_tb = (
+                "".join(traceback.format_list(self.controller_frames))
+                if self.controller_frames is not None
+                else "  <not related to a specific invocation>\n"
+            )
+            return (
+                f"{self.message}\n"
+                f"Traceback of where the remote function was issued on controller (most recent call last):\n{controller_tb}"
+                f"Traceback of where the remote function failed on worker (most recent call last):\n{worker_tb}{type(self.exception).__name__}: {exe}"
+            )
+        except Exception as e:
+            print(e)
+            return "oops"
+
+
+class Invocation:
+    def __init__(self, seq: Seq):
+        self.seq = seq
+        self.users: Optional[set["Invocation"]] = set()
+        self.failure: Optional[RemoteException] = None
+        self.fut_value: Any = None
+
+    def __repr__(self):
+        return f"<Invocation {self.seq}>"
+
+    def fail(self, remote_exception: RemoteException):
+        if self.failure is None or self.failure.seq > remote_exception.seq:
+            self.failure = remote_exception
+            return True
+        return False
+
+    def add_user(self, r: "Invocation"):
+        if self.users is not None:
+            self.users.add(r)
+        if self.failure is not None:
+            r.fail(self.failure)
+
+    def complete(self) -> Tuple[Any, Optional[RemoteException]]:
+        """
+        Complete the current invocation.
+        Return the result and exception tuple.
+        """
+        # after completion we no longer need to inform users of failures
+        # since they will just immediately get the value during add_user
+        self.users = None
+
+        return (self.fut_value if self.failure is None else None, self.failure)