Allow proc_mesh.stop() after actor_mesh.stop() (#877)

Summary:
Pull Request resolved: #877

Currently calling proc_mesh.stop() will fail if an actor_mesh spawned on this proc_mesh has been stopped.

This is because the python-binding wrapper for ProcMesh (TrackedProcMesh) holds onto a map of SharedCells containing each RootActorMesh spawned on it. When an PythonActorMesh is stopped, the RootActorMesh is dropped, but does not update the state of the TrackedProcMesh to remove this. When a ProcMesh is stopped, it will attempt to discard every SharedCell containing it's RootActorMeshes, but will return an error if any RootActorMesh has been dropped (The error we are seeing)

We can make it such that stopping the PythonActorMesh will update the state of the TrackedProcMesh, but this would involve plumbing through references of TrackedProcMesh and related state to PythonActorMesh

A simpler solution, it to simply allow the TrackedProcMesh to discard the errors related to attempting to discard a SharedCell where the RootActorMesh has already been dropped

Reviewed By: ahmadsharif1

Differential Revision: D80269385

Author: thomasywang

hyperactor_mesh/src/shared_cell.rs

@@ -14,6 +14,7 @@ use std::sync::atomic::Ordering;
 
 use async_trait::async_trait;
 use dashmap::DashMap;
+use futures::future::join_all;
 use futures::future::try_join_all;
 use preempt_rwlock::OwnedPreemptibleRwLockReadGuard;
 use preempt_rwlock::PreemptibleRwLock;
@@ -219,6 +220,16 @@ impl SharedCellPool {
         .await?;
         Ok(())
     }
+
+    /// Run `take` on all cells in the pool and immediately drop them or produce an error if the cell has already been taken
+    pub async fn discard_or_error_all(self) -> Vec<Result<(), EmptyCellError>> {
+        join_all(
+            self.map
+                .iter()
+                .map(|r| async move { r.value().discard().await }),
+        )
+        .await
+    }
 }
 
 /// Trait to facilitate storing `SharedCell`s of different types in a single pool.

monarch_hyperactor/src/proc_mesh.rs

@@ -233,7 +233,9 @@ impl PyProcMesh {
         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?;
+        // Discarding actor meshes that have been individually stopped will result in an expected error
+        // which we can safely ignore
+        children.discard_or_error_all().await;
 
         // Finally, take ownership of the inner proc mesh, which will allowing dropping it.
         let _proc_mesh = proc_mesh.take().await?;

python/tests/test_python_actors.py

@@ -1012,6 +1012,15 @@ async def test_actor_mesh_stop(self) -> None:
         await am_2.print.call("hello 3")
         await am_2.log.call("hello 4")
 
+        await pm.stop()
+
+    async def test_proc_mesh_stop_after_actor_mesh_stop(self) -> None:
+        pm = proc_mesh(gpus=2)
+        am = await pm.spawn("printer", Printer)
+
+        await cast(ActorMesh, am).stop()
+        await pm.stop()
+
 
 class PortedActor(Actor):
     @endpoint(explicit_response_port=True)