remove asyncio.lock (#1391)

Summary:
Pull Request resolved: #1391

so things to know:
- Actors are run in asyncio loop, so ok to assume loop is running
- but any lock is blocking, and we can use the monarch Future to call await multiple times on common future, results will be:
(1) enqueue Future  to event loop
(2) return reference to enqueued Future
(3) return Future result (immediate) if already completed

Reviewed By: allenwang28

Differential Revision: D83666787

fbshipit-source-id: 2d6fd4d2873e31b05235cce9c3bb736848240f54

Author: dstaay-fb

python/monarch/_src/rdma/rdma.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-unsafe
-import asyncio
 import ctypes
 import functools
 import logging
@@ -118,27 +117,25 @@ def _get_addr_and_size(buf: torch.Tensor | memoryview) -> tuple[int, int]:
 
 class RdmaController(Actor):
     def __init__(self) -> None:
-        self._managers: Dict[ProcMesh, _RdmaManager] = {}
-        self._lock = asyncio.Lock()
+        self._manager_futures: Dict[ProcMesh, Future[_RdmaManager]] = {}
 
     @endpoint
     async def init_rdma_on_mesh(self, proc_mesh: ProcMesh) -> None:
         # Note: RdmaController acts as coordinator and can run on any node
         # The RDMA support check should happen on the target proc_mesh nodes, not on RdmaController's node
 
-        if proc_mesh in self._managers:
-            return
-
-        async with self._lock:
-            if proc_mesh not in self._managers:
-                self._managers[proc_mesh] = none_throws(
-                    await Future(
-                        coro=_RdmaManager.create_rdma_manager_nonblocking(
-                            await Future(coro=proc_mesh._proc_mesh.task())
-                        )
-                    )
+        if proc_mesh not in self._manager_futures:
+
+            async def create_manager() -> _RdmaManager:
+                proc_mesh_result = await Future(coro=proc_mesh._proc_mesh.task())
+                return none_throws(
+                    await _RdmaManager.create_rdma_manager_nonblocking(proc_mesh_result)
                 )
 
+            self._manager_futures[proc_mesh] = Future(coro=create_manager())
+
+        await self._manager_futures[proc_mesh]
+
 
 @functools.cache
 def _check_cuda_expandable_segments_enabled() -> bool: