torchmonarch-nightly 2025.6.19__cp310-cp310-manylinux2014_x86_64.whl → 2025.6.20__cp310-cp310-manylinux2014_x86_64.whl

monarch/common/client.py

@@ -41,6 +41,8 @@ from monarch.common import messages
 from monarch.common.borrows import Borrow, StorageAliases
 from monarch.common.controller_api import LogMessage, MessageResult, TController
 from monarch.common.device_mesh import DeviceMesh
+
+from monarch.common.future import Future
 from monarch.common.invocation import DeviceException, RemoteException, Seq
 from monarch.common.recording import flatten_messages, Recording
 
@@ -52,9 +54,6 @@ from monarch.common.tree import tree_map
 
 from . import _coalescing
 
-if TYPE_CHECKING:
-    from monarch.common.future import Future
-
 
 logger = logging.getLogger(__name__)
 
@@ -447,6 +446,39 @@ class Client:
     def mesh_state(self) -> WorldState:
         return self.inner.worker_world_state()
 
+    def fetch(
+        self,
+        mesh: "DeviceMesh",
+        stream: "StreamRef",
+        shard,
+        preprocess_message,
+        args,
+        kwargs,
+        defs: Tuple["Tensor", ...],
+        uses: Tuple["Tensor", ...],
+    ) -> "Future":
+        fut = Future(self)
+        ident = self.new_node(defs, uses, fut)
+        process = mesh._process(shard)
+        self.send(
+            process,
+            messages.SendValue(
+                ident,
+                None,
+                defs,
+                preprocess_message,
+                args,
+                kwargs,
+                stream,
+            ),
+        )
+        # we have to ask for status updates
+        # from workers to be sure they have finished
+        # enough work to count this future as finished,
+        # and all potential errors have been reported
+        self._request_status()
+        return fut
+
 
 def tree_map_refs(first_ref: int, tree):
     def translate_id(ref: int) -> int:

monarch/common/remote.py

@@ -21,6 +21,7 @@ from typing import (
     overload,
     Protocol,
     Tuple,
+    TYPE_CHECKING,
     TypeVar,
 )
 
@@ -30,6 +31,9 @@ import torch
 
 from monarch.common import _coalescing, device_mesh, messages, stream
 
+if TYPE_CHECKING:
+    from monarch.common.client import Client
+
 from monarch.common.device_mesh import RemoteProcessGroup
 from monarch.common.fake import fake_call
 
@@ -173,30 +177,19 @@ def _call_on_shard_and_fetch(
         propagator, rfunction, args, kwargs, ambient_mesh, stream._active
     )
 
-    client = mesh.client
+    client: "Client" = mesh.client
     if _coalescing.is_active(client):
         raise NotImplementedError("NYI: fetching results during a coalescing block")
-    fut = Future(client)
-    ident = client.new_node(mutates, dtensors, fut)
-    process = mesh._process(shard)
-    client.send(
-        process,
-        messages.SendValue(
-            ident,
-            None,
-            mutates,
-            preprocess_message,
-            args,
-            kwargs,
-            stream._active._to_ref(client),
-        ),
+    return client.fetch(
+        mesh,
+        stream._active._to_ref(client),
+        shard,
+        preprocess_message,
+        args,
+        kwargs,
+        mutates,
+        dtensors,
     )
-    # we have to ask for status updates
-    # from workers to be sure they have finished
-    # enough work to count this future as finished,
-    # and all potential errors have been reported
-    client._request_status()
-    return fut
 
 
 @remote

monarch/future.py

@@ -5,21 +5,72 @@
 # LICENSE file in the root directory of this source tree.
 
 import asyncio
-from typing import Generator, Generic, TypeVar
+from functools import partial
+from typing import Generator, Generic, Optional, TypeVar
 
 R = TypeVar("R")
 
 
+def _incomplete(impl, self):
+    try:
+        return self._set_result(impl())
+    except Exception as e:
+        self._set_exception(e)
+        raise
+
+
+async def _aincomplete(impl, self):
+    try:
+        return self._set_result(await impl())
+    except Exception as e:
+        self._set_exception(e)
+        raise
+
+
 # TODO: consolidate with monarch.common.future
 class ActorFuture(Generic[R]):
     def __init__(self, impl, blocking_impl=None):
-        self._impl = impl
-        self._blocking_impl = blocking_impl
+        if blocking_impl is None:
+            blocking_impl = partial(asyncio.run, impl())
+        self._get = partial(_incomplete, blocking_impl)
+        self._aget = partial(_aincomplete, impl)
 
-    def get(self) -> R:
-        if self._blocking_impl is not None:
-            return self._blocking_impl()
-        return asyncio.run(self._impl())
+    def get(self, timeout: Optional[float] = None) -> R:
+        if timeout is not None:
+            return asyncio.run(asyncio.wait_for(self._aget(self), timeout))
+        return self._get(self)
 
     def __await__(self) -> Generator[R, None, R]:
-        return self._impl().__await__()
+        return self._aget(self).__await__()
+
+    def _set_result(self, result):
+        def f(self):
+            return result
+
+        async def af(self):
+            return result
+
+        self._get, self._aget = f, af
+        return result
+
+    def _set_exception(self, e):
+        def f(self):
+            raise e
+
+        async def af(self):
+            raise e
+
+        self._get, self._aget = f, af
+
+    # compatibility with old tensor engine Future objects
+    # hopefully we do not need done(), add_callback because
+    # they are harder to implement right.
+    def result(self, timeout: Optional[float] = None) -> R:
+        return self.get(timeout)
+
+    def exception(self, timeout: Optional[float] = None):
+        try:
+            self.get(timeout)
+            return None
+        except Exception as e:
+            return e

tests/test_python_actors.py

@@ -29,6 +29,7 @@ from monarch.actor_mesh import (
     MonarchContext,
 )
 from monarch.debugger import init_debugging
+from monarch.future import ActorFuture
 
 from monarch.mesh_controller import spawn_tensor_engine
 
@@ -672,3 +673,100 @@ async def test_async_concurrency():
     # actually concurrently processing messages.
     await am.no_more.call()
     await fut
+
+
+async def awaitit(f):
+    return await f
+
+
+def test_actor_future():
+    v = 0
+
+    async def incr():
+        nonlocal v
+        v += 1
+        return v
+
+    # can use async implementation from sync
+    # if no non-blocking is provided
+    f = ActorFuture(incr)
+    assert f.get() == 1
+    assert v == 1
+    assert f.get() == 1
+    assert asyncio.run(awaitit(f)) == 1
+
+    f = ActorFuture(incr)
+    assert asyncio.run(awaitit(f)) == 2
+    assert f.get() == 2
+
+    def incr2():
+        nonlocal v
+        v += 2
+        return v
+
+    # Use non-blocking optimization if provided
+    f = ActorFuture(incr, incr2)
+    assert f.get() == 4
+    assert asyncio.run(awaitit(f)) == 4
+
+    async def nope():
+        nonlocal v
+        v += 1
+        raise ValueError("nope")
+
+    f = ActorFuture(nope)
+
+    with pytest.raises(ValueError):
+        f.get()
+
+    assert v == 5
+
+    with pytest.raises(ValueError):
+        f.get()
+
+    assert v == 5
+
+    with pytest.raises(ValueError):
+        asyncio.run(awaitit(f))
+
+    assert v == 5
+
+    def nope():
+        nonlocal v
+        v += 1
+        raise ValueError("nope")
+
+    f = ActorFuture(incr, nope)
+
+    with pytest.raises(ValueError):
+        f.get()
+
+    assert v == 6
+
+    with pytest.raises(ValueError):
+        f.result()
+
+    assert f.exception() is not None
+
+    assert v == 6
+
+    with pytest.raises(ValueError):
+        asyncio.run(awaitit(f))
+
+    assert v == 6
+
+    async def seven():
+        return 7
+
+    f = ActorFuture(seven)
+
+    assert 7 == f.get(timeout=0.001)
+
+    async def neverfinish():
+        f = asyncio.Future()
+        await f
+
+    f = ActorFuture(neverfinish)
+
+    with pytest.raises(asyncio.exceptions.TimeoutError):
+        f.get(timeout=0.1)

{torchmonarch_nightly-2025.6.19.dist-info → torchmonarch_nightly-2025.6.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchmonarch-nightly
-Version: 2025.6.19
+Version: 2025.6.20
 Summary: Monarch: Single controller library
 Author: Meta
 Author-email: oncall+monarch@xmail.facebook.com

{torchmonarch_nightly-2025.6.19.dist-info → torchmonarch_nightly-2025.6.20.dist-info}/RECORD

@@ -1,5 +1,5 @@
 monarch/__init__.py,sha256=iUvWHc0-7Q2tovRoRxOIiA3TsefMXCbWl-jEfQ2djew,6897
-monarch/_rust_bindings.so,sha256=EUkkinIuX45ihfDu4ot656fOd0CxaepnmaZdUv0cOMY,41044112
+monarch/_rust_bindings.so,sha256=sDdg6RjptgNPmFnFiDAgv36k_Or_Kz47aYaZ2M5EAao,41088032
 monarch/_testing.py,sha256=jOIOG6jcZBzvEvG_DwSnwCkaMVXvSun6sJAG6nXemww,7859
 monarch/actor_mesh.py,sha256=m6QapbZHqYujXya28jW1II2wkBUV_nKGvxmWPSW9lsQ,24327
 monarch/allocator.py,sha256=UEaVLntH4xQ8Lr84TbgcXusvuK8FhSMJmav-omztUbw,4473
@@ -7,7 +7,7 @@ monarch/bootstrap_main.py,sha256=RCUQhJk07yMFiKp6HzQuqZFUpkgsT9kVEyimiwjn6_E,182
 monarch/cached_remote_function.py,sha256=kYdB6r4OHx_T_uX4q3tCNcp1t2DJwF8tPTIahUiT2pU,8785
 monarch/debugger.py,sha256=AdlvOG3X-9Pw9c1DLQYEy4vjEfh0ZtwtsNJEFLFzN8o,13312
 monarch/fetch.py,sha256=61jxo7sx4QNUTkc0_rF5NaJROen4tKbAaiIjrXWLOvg,1705
-monarch/future.py,sha256=lcdFEe7m1shYPPuvZ1RkS6JUIChEKGBWe3v7x_nu4Hg,731
+monarch/future.py,sha256=g1VYJl8ReBBS6VbikwWilnFqEr5qJDiSKid92AnWFV4,2058
 monarch/gradient_generator.py,sha256=Rl3dmXGceTdCc1mYBg2JciR88ywGPnW7TVkL86KwqEA,6366
 monarch/memory.py,sha256=ol86dBhFAJqg78iF25-BuK0wuwj1onR8FIioZ_B0gjw,1377
 monarch/mesh_controller.py,sha256=am1QP7dvn0OH1z9ADSKm41APs1HY_dHcBAhOVP-QDmE,10427
@@ -46,7 +46,7 @@ monarch/common/_device_utils.py,sha256=gBpl23wMjppVAEzzj8U9HyX-B7Bs2_3ftiMAkzUS4
 monarch/common/_tensor_to_table.py,sha256=yRjCNwvtl188Z1Dwkx3ZU-Bh2mwYnQ0Lnue2RAztwvc,5753
 monarch/common/base_tensor.py,sha256=ujRzR6lWaeCdPv2JX0vCR-VsCWn-3SHaJIkZH1Sw9FQ,1159
 monarch/common/borrows.py,sha256=7KR62xoUat1T6FyADsdHsxVAVIJDvfJWUnPO-xx277U,5307
-monarch/common/client.py,sha256=axo37s_z17nYQGOZG5fi_0zUEJ_8qw7INjs-Kw2vaVo,24937
+monarch/common/client.py,sha256=abYQqi-yFzG0ERvh3gMC5UgiWSezmM20kbxzalKpnf4,25806
 monarch/common/constants.py,sha256=ohvsVYMpfeWopv3KXDAeHWDFLukwc-OY37VRxpKNBE8,300
 monarch/common/context_manager.py,sha256=GOeyaFbyCqvQmkJ0oI7q6IxRd8_0mVyYKZRccI8iaug,1067
 monarch/common/controller_api.py,sha256=djGkK5aSd-V6pBkr3uBCXbfJv3OKf2o2VbBXJgFF2WI,3202
@@ -65,7 +65,7 @@ monarch/common/pipe.py,sha256=9pTf8--3yOv4HpnJEhgcmc_JM6Az4uL1y72TSQA55dw,5013
 monarch/common/process_group.py,sha256=FbJ_AJRZYFkvQ68L2naRq64J_aNuAKe5kO0MWdn_x74,1662
 monarch/common/recording.py,sha256=hoI9VY_FyW_xVx-jmfsKydqX5vW2GulwcDWsBdUVOm8,4637
 monarch/common/reference.py,sha256=O26lkzEeVwj0S1xEy-OLqdHVnACmmlbQCUmXRrW4n1Q,938
-monarch/common/remote.py,sha256=qZWXkShX20l07TseQSpVECh2yXZaVKYUvQXkeEM-zvY,9220
+monarch/common/remote.py,sha256=vklFYJvuaPpS8kAyFmRz-T-brfHvcZ1lPTC_-7DIwqM,8908
 monarch/common/selection.py,sha256=lpWFbZs3ArYy29e-53eoAVAjQFksf1RvZz9NvM0CUW4,308
 monarch/common/shape.py,sha256=B-7DI768ZhT8ECUNCJcI7DfCB7iDFGFH0r-HmXaAfcM,8296
 monarch/common/stream.py,sha256=_ejoxafHtdD10lLzznRCXKwrkZ_ZH9k_VTgiA5yfBrI,3583
@@ -145,7 +145,7 @@ tests/test_future.py,sha256=cXzaNi2YDwVyjR541ScXmgktX1YFsKzbl8wep0DMVbk,3032
 tests/test_grad_generator.py,sha256=p4Pm4kMEeGldt2jUVAkGKCB0mLccKI28pltH6OTGbQA,3412
 tests/test_mock_cuda.py,sha256=5hisElxeLJ5MHw3KM9gwxBiXiMaG-Rm382u3AsQcDOI,3068
 tests/test_pdb_actor.py,sha256=5KJhuhcZDPWMdjC6eAtDdwnz1W7jNFXvIrMSFaCWaPw,3858
-tests/test_python_actors.py,sha256=3ru2JsPQmaO7ppVX3-ls7JcvIeOgEmWWUsYKZCuBXPg,19256
+tests/test_python_actors.py,sha256=du0AiGiKtVHOLkDUKu6gV75eYf_NoHDKV6utKzrplz4,21010
 tests/test_remote_functions.py,sha256=5nxYB8dfA9NT9f9Od9O3htgQtPbiRNiXZ1Kgtn75sOQ,50056
 tests/test_rust_backend.py,sha256=94S3R995ZkyIhEiBsM5flcjf5X7bscEAHBtInbTRFe8,7776
 tests/test_signal_safe_block_on.py,sha256=bmal0XgzJowZXJV6T1Blow5a-vZluYWusCThLMGxyTE,3336
@@ -155,9 +155,9 @@ tests/simulator/test_profiling.py,sha256=TGYCfzTLdkpIwnOuO6KApprmrgPIRQe60KRX3wk
 tests/simulator/test_simulator.py,sha256=LO8lA0ssY-OGEBL5ipEu74f97Y765TEwfUOv-DtIptM,14568
 tests/simulator/test_task.py,sha256=ipqBDuDAysuo1xOB9S5psaFvwe6VATD43IovCTSs0t4,2327
 tests/simulator/test_worker.py,sha256=QrWWIJ3HDgDLkBPRc2mwYPlOQoXQcj1qRfc0WUfKkFY,3507
-torchmonarch_nightly-2025.6.19.dist-info/licenses/LICENSE,sha256=e0Eotbf_rHOYPuEUlppIbvwy4SN98CZnl_hqwvbDA4Q,1530
-torchmonarch_nightly-2025.6.19.dist-info/METADATA,sha256=2XYBEhTb9iSTFKhAGmq2Bg_AXwjQvcPj6CQmG4bBiLE,2772
-torchmonarch_nightly-2025.6.19.dist-info/WHEEL,sha256=_wZSFk0d90K9wOBp8Q-UGxshyiJ987JoPiyUBNC6VLk,104
-torchmonarch_nightly-2025.6.19.dist-info/entry_points.txt,sha256=sqfQ16oZqjEvttUI-uj9BBXIIE6jt05bYFSmy-2hyXI,106
-torchmonarch_nightly-2025.6.19.dist-info/top_level.txt,sha256=E-ZssZzyM17glpVrh-S9--qJ-w9p2EjuYOuNw9tQ4Eg,33
-torchmonarch_nightly-2025.6.19.dist-info/RECORD,,
+torchmonarch_nightly-2025.6.20.dist-info/licenses/LICENSE,sha256=e0Eotbf_rHOYPuEUlppIbvwy4SN98CZnl_hqwvbDA4Q,1530
+torchmonarch_nightly-2025.6.20.dist-info/METADATA,sha256=QKiDH01IYFpa492TDs5WzWeDRbjMKmpSAc3V9NpQ5YM,2772
+torchmonarch_nightly-2025.6.20.dist-info/WHEEL,sha256=_wZSFk0d90K9wOBp8Q-UGxshyiJ987JoPiyUBNC6VLk,104
+torchmonarch_nightly-2025.6.20.dist-info/entry_points.txt,sha256=sqfQ16oZqjEvttUI-uj9BBXIIE6jt05bYFSmy-2hyXI,106
+torchmonarch_nightly-2025.6.20.dist-info/top_level.txt,sha256=E-ZssZzyM17glpVrh-S9--qJ-w9p2EjuYOuNw9tQ4Eg,33
+torchmonarch_nightly-2025.6.20.dist-info/RECORD,,