[monarch] Enable pickling HostMesh and ProcMesh from inside tokio thread in certain cases

Previously, both `HostMesh` and `ProcMesh` stored their underlying rust mesh in an asynchronous `Shared[...]` pytokio object, so any attempts to access the rust meshes from a tokio thread would need to be called from a coroutine. This is a problem for pickling, which needs to be synchronous, and would therefore have to call `Shared[...].block_on()` to get the underlying rust mesh.

This diff makes it so that when the internal `Shared[...]` task for `HostMesh` and `ProcMesh` completes, the result is stored so that it can be accessed without the use of `block_on()`. This enables pickling `HostMesh` and `ProcMesh` inside tokio threads, as long as their backing rust meshes have finished initializing -- this will always be the case inside actor endpoints, since the `HostMesh` and `ProcMesh` both need to be done initializing by the time the actor endpoint runs.

Differential Revision: [D84195494](https://our.internmc.facebook.com/intern/diff/D84195494/)

ghstack-source-id: 315055002
Pull Request resolved: #1473

Author: samlurye

python/monarch/_src/actor/bootstrap.py

@@ -114,5 +114,9 @@ def attach_to_workers(
     host_mesh: PythonTask[HyHostMesh] = _attach_to_workers(workers_tasks, name=name)
     extent = Extent(["hosts"], [len(workers)])
     return HostMesh(
-        host_mesh.spawn(), extent.region, stream_logs=True, is_fake_in_process=False
+        host_mesh.spawn(),
+        extent.region,
+        stream_logs=True,
+        is_fake_in_process=False,
+        _initialized_hy_host_mesh=None,
     )

python/monarch/_src/actor/v1/host_mesh.py

@@ -96,11 +96,13 @@ def __init__(
         region: Region,
         stream_logs: bool,
         is_fake_in_process: bool,
+        _initialized_hy_host_mesh: Optional[HyHostMesh],
     ) -> None:
         self._hy_host_mesh = hy_host_mesh
         self._region = region
         self._stream_logs = stream_logs
         self._is_fake_in_process = is_fake_in_process
+        self._initialized_host_mesh = _initialized_hy_host_mesh
 
     @classmethod
     def allocate_nonblocking(
@@ -122,13 +124,24 @@ async def task() -> HyHostMesh:
                 bootstrap_cmd,
             )
 
-        return cls(
-            PythonTask.from_coroutine(task()).spawn(),
+        hy_host_mesh_task = PythonTask.from_coroutine(task()).spawn()
+
+        hm: HostMesh = cls(
+            hy_host_mesh_task,
             extent.region,
             alloc.stream_logs,
             isinstance(allocator, LocalAllocator),
+            None,
         )
 
+        async def task(hy_host_mesh_task: Shared[HyHostMesh]) -> HyHostMesh:
+            hm._initialized_host_mesh = await hy_host_mesh_task
+            return hm._initialized_host_mesh
+
+        hm._hy_host_mesh = PythonTask.from_coroutine(task(hy_host_mesh_task)).spawn()
+
+        return hm
+
     def spawn_procs(
         self,
         per_host: Dict[str, int] | None = None,
@@ -187,15 +200,25 @@ def _new_with_shape(self, shape: Shape) -> "HostMesh":
         if shape.region == self._region:
             return self
 
+        initialized_hm: Optional[HyHostMesh] = (
+            None
+            if self._initialized_host_mesh is None
+            else self._initialized_host_mesh.sliced(shape.region)
+        )
+
         async def task() -> HyHostMesh:
-            hy_host_mesh = await self._hy_host_mesh
-            return hy_host_mesh.sliced(shape.region)
+            return (
+                initialized_hm
+                if initialized_hm
+                else (await self._hy_host_mesh).sliced(shape.region)
+            )
 
         return HostMesh(
             PythonTask.from_coroutine(task()).spawn(),
             shape.region,
             self.stream_logs,
             self.is_fake_in_process,
+            initialized_hm,
         )
 
     @property
@@ -222,11 +245,12 @@ async def task() -> HyHostMesh:
             region,
             stream_logs,
             is_fake_in_process,
+            hy_host_mesh,
         )
 
     def __reduce_ex__(self, protocol: ...) -> Tuple[Any, Tuple[Any, ...]]:
         return HostMesh._from_initialized_hy_host_mesh, (
-            self._hy_host_mesh.block_on(),
+            self._initialized_mesh(),
             self._region,
             self.stream_logs,
             self.is_fake_in_process,
@@ -238,12 +262,18 @@ def is_fake_in_process(self) -> bool:
 
     def __eq__(self, other: "HostMesh") -> bool:
         return (
-            self._hy_host_mesh.block_on() == other._hy_host_mesh.block_on()
+            self._initialized_mesh() == other._initialized_mesh()
             and self._region == other._region
             and self.stream_logs == other.stream_logs
             and self.is_fake_in_process == other.is_fake_in_process
         )
 
+    def _initialized_mesh(self) -> HyHostMesh:
+        if self._initialized_host_mesh is None:
+            self._hy_host_mesh.block_on()
+            assert self._initialized_host_mesh is not None
+        return self._initialized_host_mesh
+
 
 def fake_in_process_host(name: str) -> "HostMesh":
     """

python/monarch/_src/actor/v1/proc_mesh.py

@@ -81,6 +81,7 @@ def __init__(
         host_mesh: "HostMesh",
         region: Region,
         root_region: Region,
+        _initialized_hy_proc_mesh: Optional[HyProcMesh],
         _device_mesh: Optional["DeviceMesh"] = None,
     ) -> None:
         _proc_mesh_registry.add(self)
@@ -91,6 +92,7 @@ def __init__(
         self._maybe_device_mesh = _device_mesh
         self._logging_manager = LoggingManager()
         self._controller_controller: Optional["_ControllerController"] = None
+        self._initialized_proc_mesh = _initialized_hy_proc_mesh
 
     @property
     def initialized(self) -> Future[Literal[True]]:
@@ -138,14 +140,25 @@ def _new_with_shape(self, shape: Shape) -> "ProcMesh":
             else self._maybe_device_mesh._new_with_shape(shape)
         )
 
+        initialized_pm: Optional[HyProcMesh] = (
+            None
+            if self._initialized_proc_mesh is None
+            else self._initialized_proc_mesh.sliced(shape.region)
+        )
+
         async def task() -> HyProcMesh:
-            return (await self._proc_mesh).sliced(shape.region)
+            return (
+                initialized_pm
+                if initialized_pm
+                else (await self._proc_mesh).sliced(shape.region)
+            )
 
         return ProcMesh(
             PythonTask.from_coroutine(task()).spawn(),
             self._host_mesh,
             shape.region,
             self._root_region,
+            initialized_pm,
             _device_mesh=device_mesh,
         )
 
@@ -188,7 +201,7 @@ def from_host_mesh(
         setup: Callable[[], None] | None = None,
         _attach_controller_controller: bool = True,
     ) -> "ProcMesh":
-        pm = ProcMesh(hy_proc_mesh, host_mesh, region, region)
+        pm = ProcMesh(hy_proc_mesh, host_mesh, region, region, None)
 
         if _attach_controller_controller:
             instance = context().actor_instance
@@ -211,6 +224,7 @@ async def task(
             stream_log_to_client: bool,
         ) -> HyProcMesh:
             hy_proc_mesh = await hy_proc_mesh_task
+            pm._initialized_proc_mesh = hy_proc_mesh
 
             await pm._logging_manager.init(hy_proc_mesh, stream_log_to_client)
 
@@ -367,11 +381,14 @@ async def task() -> HyProcMesh:
             host_mesh,
             region,
             root_region,
+            _initialized_hy_proc_mesh=hy_proc_mesh,
         )
 
     def __reduce_ex__(self, protocol: ...) -> Tuple[Any, Tuple[Any, ...]]:
         return ProcMesh._from_initialized_hy_proc_mesh, (
-            self._proc_mesh.block_on(),
+            self._initialized_proc_mesh
+            if self._initialized_proc_mesh
+            else self._proc_mesh.block_on(),
             self._host_mesh,
             self._region,
             self._root_region,

python/tests/test_proc_mesh.py

@@ -8,7 +8,10 @@
 
 from typing import cast
 
+import cloudpickle
+
 import pytest
+from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
 from monarch._rust_bindings.monarch_hyperactor.shape import Extent, Shape, Slice
 from monarch._src.actor.actor_mesh import Actor, ActorMesh, context, ValueMesh
 from monarch._src.actor.endpoint import endpoint
@@ -137,3 +140,20 @@ def test_nested_meshes() -> None:
         assert value == point.rank + 1
     for point, value in res_1:
         assert value == point.rank
+
+
+@pytest.mark.timeout(60)
+async def test_pickle_initialized_proc_mesh_in_tokio_thread() -> None:
+    host = create_local_host_mesh("host", Extent(["hosts"], [2]))
+    proc = host.spawn_procs(per_host={"gpus": 2})
+
+    async def task():
+        cloudpickle.dumps(proc)
+
+    await proc.initialized
+    PythonTask.from_coroutine(task()).block_on()
+
+    async def task():
+        cloudpickle.dumps(proc.slice(gpus=0, hosts=0))
+
+    PythonTask.from_coroutine(task()).block_on()