Add async context manager to ProcMesh for stop (#447)

Summary:
Pull Request resolved: #447

Now that a ProcMesh has a `stop` function which closes the external processes,
add `__aenter__` and `__aexit__` functions to be used as an async context manager.

Reviewed By: pablorfb-meta

Differential Revision: D77877888

fbshipit-source-id: ab2260725b4800f265160707a0a5885d4d2385f6

Author: dulinriley

python/monarch/proc_mesh.py

@@ -88,6 +88,7 @@ def __init__(
         self._rsync_mesh_client: Optional[RsyncMeshClient] = None
         self._auto_reload_actor: Optional[AutoReloadActor] = None
         self._maybe_device_mesh: Optional[DeviceMesh] = _device_mesh
+        self._stopped = False
         if _mock_shape is None:
             self._rdma_manager = self._spawn_blocking("rdma_manager", RDMAManager)
 
@@ -247,6 +248,33 @@ async def sync_workspace(self, auto_reload: bool = False) -> None:
 
     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
+
+    async def __aexit__(
+        self, exc_type: object, exc_val: object, exc_tb: object
+    ) -> None:
+        # In case there are multiple nested "async with" statements, we only
+        # want it to close once.
+        if not self._stopped:
+            await self.stop()
+
+    # Finalizer to check if the proc mesh was closed properly.
+    def __del__(self) -> None:
+        if not self._stopped:
+            import warnings
+
+            warnings.warn(
+                f"unstopped ProcMesh {self!r}",
+                ResourceWarning,
+                stacklevel=2,
+                source=self,
+            )
+            # Cannot call stop here because it is async.
 
 
 async def local_proc_mesh_nonblocking(

python/tests/test_allocator.py

@@ -218,6 +218,63 @@ async def test_stop_proc_mesh(self) -> None:
             # now we doing casting without accessing the wrapped type.
             del actor
 
+    async def test_stop_proc_mesh_context_manager(self) -> None:
+        spec = AllocSpec(AllocConstraints(), host=2, gpu=4)
+
+        # create 2x process-allocators (on their own bind addresses) to simulate 2 hosts
+        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)
+            with self.assertRaises(ValueError, msg="foo"):
+                async with proc_mesh:
+                    actor = await proc_mesh.spawn("test_actor", TestActor)
+                    # Ensure that proc mesh is stopped when context manager exits.
+                    raise ValueError("foo")
+
+            with self.assertRaises(
+                RuntimeError, msg="`ProcMesh` has already been stopped"
+            ):
+                await proc_mesh.spawn("test_actor", TestActor)
+
+            # TODO(agallagher): It'd be nice to test that this just fails
+            # immediately, trying to access the wrapped actor mesh, but right
+            # now we doing casting without accessing the wrapped type.
+            del actor
+
+    async def test_stop_proc_mesh_context_manager_multiple_times(self) -> None:
+        spec = AllocSpec(AllocConstraints(), host=2, gpu=4)
+
+        # create 2x process-allocators (on their own bind addresses) to simulate 2 hosts
+        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)
+            # We can nest multiple context managers on the same mesh, the innermost
+            # one closes the mesh and it cannot be used after that.
+            async with proc_mesh:
+                async with proc_mesh:
+                    actor = await proc_mesh.spawn("test_actor", TestActor)
+
+                with self.assertRaises(
+                    RuntimeError, msg="`ProcMesh` has already been stopped"
+                ):
+                    await proc_mesh.spawn("test_actor", TestActor)
+                # Exiting a second time should not raise an error.
+
+            # TODO(agallagher): It'd be nice to test that this just fails
+            # immediately, trying to access the wrapped actor mesh, but right
+            # now we doing casting without accessing the wrapped type.
+            del actor
+
     async def test_stacked_1d_meshes(self) -> None:
         # create two stacked actor meshes on the same host
         # each actor mesh running on separate process-allocators