torchmonarch-nightly 2025.6.13__cp310-cp310-manylinux2014_x86_64.whl → 2025.6.14__cp310-cp310-manylinux2014_x86_64.whl

monarch/actor_mesh.py

@@ -276,11 +276,11 @@ class Endpoint(Generic[P, R]):
         return self.choose(*args, **kwargs)
 
     def call(self, *args: P.args, **kwargs: P.kwargs) -> "Future[ValueMesh[R]]":
-        p: PortId
-        r: PortReceiver[R]
-        p, r = port(self)
+        p: Port[R]
+        r: RankedPortReceiver[R]
+        p, r = ranked_port(self)
         # pyre-ignore
-        send(self, args, kwargs, port=p, rank_in_response=True)
+        send(self, args, kwargs, port=p)
 
         async def process() -> ValueMesh[R]:
             results: List[R] = [None] * len(self._actor_mesh)  # pyre-fixme[9]
@@ -375,9 +375,8 @@ def send(
     endpoint: Endpoint[P, R],
     args: Tuple[Any, ...],
     kwargs: Dict[str, Any],
-    port: "Optional[PortId]" = None,
+    port: "Optional[Port]" = None,
     selection: Selection = "all",
-    rank_in_response: bool = False,
 ) -> None:
     """
     Fire-and-forget broadcast invocation of the endpoint across all actors in the mesh.
@@ -386,7 +385,10 @@ def send(
     """
     endpoint._signature.bind(None, *args, **kwargs)
     message = PythonMessage(
-        endpoint._name, _pickle((args, kwargs)), port, rank_in_response
+        endpoint._name,
+        _pickle((args, kwargs)),
+        None if port is None else port._port,
+        None,
     )
     endpoint._actor_mesh.cast(message, selection)
 
@@ -408,18 +410,16 @@ def endpoint(
     return EndpointProperty(method)
 
 
-class Port:
-    def __init__(self, port: PortId, mailbox: Mailbox, rank_in_response: bool) -> None:
+class Port(Generic[R]):
+    def __init__(self, port: PortId, mailbox: Mailbox, rank: Optional[int]) -> None:
         self._port = port
         self._mailbox = mailbox
-        self._rank_in_response = rank_in_response
+        self._rank = rank
 
-    def send(self, method: str, obj: object) -> None:
-        if self._rank_in_response:
-            obj = (MonarchContext.get().point.rank, obj)
+    def send(self, method: str, obj: R) -> None:
         self._mailbox.post(
             self._port,
-            PythonMessage(method, _pickle(obj), None),
+            PythonMessage(method, _pickle(obj), None, self._rank),
         )
 
 
@@ -428,12 +428,21 @@ class Port:
 # and handles concerns is different.
 def port(
     endpoint: Endpoint[P, R], once: bool = False
-) -> Tuple["PortId", "PortReceiver[R]"]:
+) -> Tuple["Port[R]", "PortReceiver[R]"]:
     handle, receiver = (
         endpoint._mailbox.open_once_port() if once else endpoint._mailbox.open_port()
     )
     port_id: PortId = handle.bind()
-    return port_id, PortReceiver(endpoint._mailbox, receiver)
+    return Port(port_id, endpoint._mailbox, rank=None), PortReceiver(
+        endpoint._mailbox, receiver
+    )
+
+
+def ranked_port(
+    endpoint: Endpoint[P, R], once: bool = False
+) -> Tuple["Port[R]", "RankedPortReceiver[R]"]:
+    p, receiver = port(endpoint, once)
+    return p, RankedPortReceiver[R](receiver._mailbox, receiver._receiver)
 
 
 class PortReceiver(Generic[R]):
@@ -458,18 +467,20 @@ class PortReceiver(Generic[R]):
             return payload
         else:
             assert msg.method == "exception"
-            if isinstance(payload, tuple):
-                # If the payload is a tuple, it's because we requested the rank
-                # to be included in the response; just ignore it.
-                raise payload[1]
-            else:
-                # pyre-ignore
-                raise payload
+            # pyre-ignore
+            raise payload
 
     def recv(self) -> "Future[R]":
         return Future(lambda: self._recv(), self._blocking_recv)
 
 
+class RankedPortReceiver(PortReceiver[Tuple[int, R]]):
+    def _process(self, msg: PythonMessage) -> Tuple[int, R]:
+        if msg.rank is None:
+            raise ValueError("RankedPort receiver got a message without a rank")
+        return msg.rank, super()._process(msg)
+
+
 singleton_shape = Shape([], NDSlice(offset=0, sizes=[], strides=[]))
 
 
@@ -493,7 +504,7 @@ class _Actor:
         panic_flag: PanicFlag,
     ) -> Optional[Coroutine[Any, Any, Any]]:
         port = (
-            Port(message.response_port, mailbox, message.rank_in_response)
+            Port(message.response_port, mailbox, rank)
             if message.response_port
             else None
         )

tests/test_python_actors.py

@@ -8,6 +8,7 @@ import asyncio
 import operator
 import os
 import re
+import threading
 from types import ModuleType
 from unittest.mock import AsyncMock, patch
 
@@ -549,3 +550,44 @@ async def test_debug() -> None:
 
         with pytest.raises(monarch.actor_mesh.ActorError, match="ValueError: bad rank"):
             await fut
+
+
+class TLSActor(Actor):
+    """An actor that manages thread-local state."""
+
+    def __init__(self):
+        self.local = threading.local()
+        self.local.value = 0
+
+    @endpoint
+    def increment(self):
+        self.local.value += 1
+
+    @endpoint
+    async def increment_async(self):
+        self.local.value += 1
+
+    @endpoint
+    def get(self):
+        return self.local.value
+
+    @endpoint
+    async def get_async(self):
+        return self.local.value
+
+
+async def test_actor_tls() -> None:
+    """Test that thread-local state is respected."""
+    pm = await proc_mesh(gpus=1)
+    am = await pm.spawn("tls", TLSActor)
+    await am.increment.call_one()
+    # TODO(suo): TLS is NOT preserved across async/sync endpoints, because currently
+    # we run async endpoints on a different thread than sync ones.
+    # Will fix this in a followup diff.
+
+    # await am.increment_async.call_one()
+    await am.increment.call_one()
+    # await am.increment_async.call_one()
+
+    assert 2 == await am.get.call_one()
+    # assert 4 == await am.get_async.call_one()

{torchmonarch_nightly-2025.6.13.dist-info → torchmonarch_nightly-2025.6.14.dist-info}/METADATA

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

{torchmonarch_nightly-2025.6.13.dist-info → torchmonarch_nightly-2025.6.14.dist-info}/RECORD

@@ -1,7 +1,7 @@
 monarch/__init__.py,sha256=iUvWHc0-7Q2tovRoRxOIiA3TsefMXCbWl-jEfQ2djew,6897
-monarch/_rust_bindings.so,sha256=FJb4gGPNDWqT1nPkxEYSX4hEsIbjb_v8Oa0RDwMcH5A,40302936
+monarch/_rust_bindings.so,sha256=iHKENTGdgZf4uRXWLL61j9Z7z-MOAIKabPo5JD0b6dw,40611888
 monarch/_testing.py,sha256=jOIOG6jcZBzvEvG_DwSnwCkaMVXvSun6sJAG6nXemww,7859
-monarch/actor_mesh.py,sha256=8Ih3CIArLTyZmWSHppXm5N2WlAjmGXpaQhkkFtjJFxc,25351
+monarch/actor_mesh.py,sha256=nAW65WFEWMJWCv8zuH9GSOyTNXwFN8QNqZxMZTuSYxw,25537
 monarch/allocator.py,sha256=ylvYTf31o-PT385cYJPhi17uNbC4yl_RAraqD0fVe4g,4112
 monarch/bootstrap_main.py,sha256=RCUQhJk07yMFiKp6HzQuqZFUpkgsT9kVEyimiwjn6_E,1827
 monarch/cached_remote_function.py,sha256=kYdB6r4OHx_T_uX4q3tCNcp1t2DJwF8tPTIahUiT2pU,8785
@@ -11,7 +11,7 @@ monarch/future.py,sha256=lcdFEe7m1shYPPuvZ1RkS6JUIChEKGBWe3v7x_nu4Hg,731
 monarch/gradient_generator.py,sha256=Rl3dmXGceTdCc1mYBg2JciR88ywGPnW7TVkL86KwqEA,6366
 monarch/memory.py,sha256=ol86dBhFAJqg78iF25-BuK0wuwj1onR8FIioZ_B0gjw,1377
 monarch/mesh_controller.py,sha256=Xft2edk7rz8_PPe-iIUZ09P-j4JDPGADBGHBiuiZ7YY,10363
-monarch/monarch_controller,sha256=mE9pvcBDKwW_4zOZlO17PJDk7W6z5skzIX5rxHQfKOs,20238936
+monarch/monarch_controller,sha256=P6ijudU0UZcrZW34d1lafCnJ1Hf90yo764rnkEg45dI,20399544
 monarch/notebook.py,sha256=zu9MKDFKf1-rCM2TqFSRJjMBeiWuKcJSyUFLvoZRQzs,25949
 monarch/opaque_module.py,sha256=oajOu_WD1hD4hxE8HDdO-tvWY7KDHWd7VaAhJEa5L2I,10446
 monarch/opaque_object.py,sha256=IVpll4pyuKZMo_EnPh4s0qnx8RlAcJrJ1yoLX6E75wQ,2782
@@ -144,7 +144,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=MPdXtnj4ZeyAaecDFJMXdz29KvimF9iB3bASgoo6-iM,16201
+tests/test_python_actors.py,sha256=MzGeuhGVICzwiNaQt8SFCKyfwhNzdRzZ4s2rJxYbeoo,17283
 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
@@ -154,9 +154,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.13.dist-info/licenses/LICENSE,sha256=e0Eotbf_rHOYPuEUlppIbvwy4SN98CZnl_hqwvbDA4Q,1530
-torchmonarch_nightly-2025.6.13.dist-info/METADATA,sha256=WhintlKk3a9WRrjo-QLNntfi87q98I4gcZW_0f42q48,2772
-torchmonarch_nightly-2025.6.13.dist-info/WHEEL,sha256=_wZSFk0d90K9wOBp8Q-UGxshyiJ987JoPiyUBNC6VLk,104
-torchmonarch_nightly-2025.6.13.dist-info/entry_points.txt,sha256=sqfQ16oZqjEvttUI-uj9BBXIIE6jt05bYFSmy-2hyXI,106
-torchmonarch_nightly-2025.6.13.dist-info/top_level.txt,sha256=E-ZssZzyM17glpVrh-S9--qJ-w9p2EjuYOuNw9tQ4Eg,33
-torchmonarch_nightly-2025.6.13.dist-info/RECORD,,
+torchmonarch_nightly-2025.6.14.dist-info/licenses/LICENSE,sha256=e0Eotbf_rHOYPuEUlppIbvwy4SN98CZnl_hqwvbDA4Q,1530
+torchmonarch_nightly-2025.6.14.dist-info/METADATA,sha256=NZRNanfPhkjOMapMX_WCs34wphY1w-zccgncjQWmylk,2772
+torchmonarch_nightly-2025.6.14.dist-info/WHEEL,sha256=_wZSFk0d90K9wOBp8Q-UGxshyiJ987JoPiyUBNC6VLk,104
+torchmonarch_nightly-2025.6.14.dist-info/entry_points.txt,sha256=sqfQ16oZqjEvttUI-uj9BBXIIE6jt05bYFSmy-2hyXI,106
+torchmonarch_nightly-2025.6.14.dist-info/top_level.txt,sha256=E-ZssZzyM17glpVrh-S9--qJ-w9p2EjuYOuNw9tQ4Eg,33
+torchmonarch_nightly-2025.6.14.dist-info/RECORD,,