Allow stop to be used as a blocking or non-blocking method (#546)

Summary:
Pull Request resolved: #546

Allow the stop method to also be blocking.

Reviewed By: highker

Differential Revision: D78362041

fbshipit-source-id: 096d069e7fee7cf704255296dacecdfe91b2a36e

Author: jayasi

monarch_hyperactor/src/proc_mesh.rs

@@ -227,6 +227,32 @@ impl PyProcMesh {
             .borrow()
             .map_err(|_| PyRuntimeError::new_err("`ProcMesh` has already been stopped"))
     }
+
+    async fn stop_mesh(
+        inner: SharedCell<TrackedProcMesh>,
+        proc_events: SharedCell<Mutex<ProcEvents>>,
+    ) -> Result<(), anyhow::Error> {
+        // "Take" the proc mesh wrapper.  Once we do, it should be impossible for new
+        // actor meshes to be spawned.
+        let tracked_proc_mesh = inner.take().await.map_err(|e| {
+            PyRuntimeError::new_err(format!("`ProcMesh` has already been stopped: {}", e))
+        })?;
+        let (proc_mesh, children) = tracked_proc_mesh.into_inner();
+
+        // Now we discard all in-flight actor meshes.  After this, the `ProcMesh` should be "unused".
+        children.discard_all().await?;
+
+        // Finally, take ownership of the inner proc mesh, which will allowing dropping it.
+        let _proc_mesh = proc_mesh.take().await?;
+
+        // Grab the alloc back from `ProcEvents` and use that to stop the mesh.
+        let proc_events_taken = proc_events.take().await?;
+        let mut alloc = proc_events_taken.into_inner().into_alloc();
+
+        alloc.stop_and_wait().await?;
+
+        anyhow::Ok(())
+    }
 }
 
 #[pymethods]
@@ -350,32 +376,28 @@ impl PyProcMesh {
         Ok(self.try_inner()?.shape().clone().into())
     }
 
-    fn stop<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        let tracked_proc_mesh = self.inner.clone();
+    fn stop_nonblocking<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        // Clone the necessary fields from self to avoid capturing self in the async block
+        let inner = self.inner.clone();
         let proc_events = self.proc_events.clone();
+
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            async {
-                // "Take" the proc mesh wrapper.  Once we do, it should be impossible for new
-                // actor meshes to be spawned.
-                let (proc_mesh, children) = tracked_proc_mesh
-                    .take()
-                    .await
-                    .map_err(|_| PyRuntimeError::new_err("`ProcMesh` has already been stopped"))?
-                    .into_inner();
-                // Now we discard all in-flight actor meshes.  After this, the `ProcMesh` should be "unused".
-                children.discard_all().await?;
-                // Finally, take ownership of the inner proc mesh, which will allowing dropping it.
-                let _proc_mesh = proc_mesh.take().await?;
-                // Grab the alloc back from `ProcEvents` and use that to stop the mesh.
-                let mut alloc = proc_events.take().await?.into_inner().into_alloc();
-                alloc.stop_and_wait().await?;
-
-                anyhow::Ok(())
-            }
-            .await?;
+            Self::stop_mesh(inner, proc_events).await?;
             PyResult::Ok(())
         })
     }
+
+    fn stop_blocking<'py>(&self, py: Python<'py>) -> PyResult<()> {
+        // Clone the necessary fields from self to avoid capturing self in the async block
+        let inner = self.inner.clone();
+        let proc_events = self.proc_events.clone();
+
+        signal_safe_block_on(py, async move {
+            Self::stop_mesh(inner, proc_events)
+                .await
+                .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))
+        })?
+    }
 }
 
 /// A keepalive token that aborts a task only after the last clone

python/monarch/_rust_bindings/monarch_hyperactor/proc_mesh.pyi

@@ -81,12 +81,18 @@ class ProcMesh:
         """
         ...
 
-    async def stop(self) -> None:
+    async def stop_nonblocking(self) -> None:
         """
         Stop the proc mesh.
         """
         ...
 
+    def stop_blocking(self) -> None:
+        """
+        Stop the proc mesh. Blocks until the mesh is fully stopped.
+        """
+        ...
+
     def __repr__(self) -> str: ...
 
 @final

python/monarch/_src/actor/proc_mesh.py

@@ -285,15 +285,25 @@ def logging_option(self, stream_to_client: bool = False) -> None:
             )
         self._logging_mesh_client.set_mode(stream_to_client)
 
-    async def stop(self) -> None:
-        await self._proc_mesh.stop()
-        self._stopped = True
-
     async def __aenter__(self) -> "ProcMesh":
         if self._stopped:
             raise RuntimeError("`ProcMesh` has already been stopped")
         return self
 
+    def stop(self) -> Future[None]:
+        async def _stop_nonblocking() -> None:
+            await self._proc_mesh.stop_nonblocking()
+            self._stopped = True
+
+        def _stop_blocking() -> None:
+            self._proc_mesh.stop_blocking()
+            self._stopped = True
+
+        return Future(
+            lambda: _stop_nonblocking(),
+            lambda: _stop_blocking(),
+        )
+
     async def __aexit__(
         self, exc_type: object, exc_val: object, exc_tb: object
     ) -> None:

python/tests/test_allocator.py

@@ -204,6 +204,24 @@ async def test_allocate_2d_mesh(self) -> None:
 
             self.assert_computed_world_size(values, world_size)
 
+    async def test_stop_proc_mesh_blocking(self) -> None:
+        spec = AllocSpec(AllocConstraints(), host=2, gpu=4)
+        with remote_process_allocator() as host1, remote_process_allocator() as host2:
+            allocator = RemoteAllocator(
+                world_id="test_remote_allocator",
+                initializer=StaticRemoteAllocInitializer(host1, host2),
+                heartbeat_interval=_100_MILLISECONDS,
+            )
+            alloc = await allocator.allocate(spec)
+            proc_mesh = await ProcMesh.from_alloc(alloc)
+            actor = proc_mesh.spawn("test_actor", TestActor).get()
+            proc_mesh.stop().get()
+            with self.assertRaises(
+                RuntimeError, msg="`ProcMesh` has already been stopped"
+            ):
+                proc_mesh.spawn("test_actor", TestActor).get()
+            del actor
+
     async def test_stop_proc_mesh(self) -> None:
         spec = AllocSpec(AllocConstraints(), host=2, gpu=4)