Enable rdma tests on v1 (meta-pytorch#1507)

Summary:
Pull Request resolved: meta-pytorch#1507

Enable rdma tests on v1. Also ensure proc mesh is initialized in `_ensure_init_rdma_manager` to prevent tokio runtime blocking issue. The tests in `test_rdma_unit.py` actually weren't running *at all*, even for v0, because whoever wrote them used a synchronous function decorator to wrap their async function tests.
ghstack-source-id: 315627106

Reviewed By: zdevito, mariusae

Differential Revision: D84395240

fbshipit-source-id: 53782e655f691587b4587694f507d13d6356d4a5

Author: samlurye

hyperactor_mesh/src/resource.rs

@@ -72,8 +72,10 @@ pub enum Status {
     /// The resource is stopped.
     Stopped,
     /// The resource has failed, with an error message.
+    #[strum(to_string = "Failed({0})")]
     Failed(String),
     /// The resource has been declared failed after a timeout.
+    #[strum(to_string = "Timeout({0:?})")]
     Timeout(Duration),
 }
 

python/monarch/_src/actor/proc_mesh.py

@@ -161,6 +161,13 @@ def __eq__(self, other: object) -> bool:
     def _proc_mesh(self) -> Shared["HyProcMeshV0"]:
         return _deref_proc_mesh(self)._proc_mesh
 
+    @property
+    def initialized(self) -> Future[Literal[True]]:
+        async def task() -> Literal[True]:
+            return True
+
+        return Future(coro=task())
+
 
 _proc_mesh_lock: threading.Lock = threading.Lock()
 _proc_mesh_key: int = 0

python/monarch/_src/rdma/rdma.py

@@ -55,6 +55,10 @@ def is_rdma_available():
 @functools.cache
 def _ensure_init_rdma_manager() -> Shared[None]:
     async def task() -> None:
+        # Ensure the proc mesh is initialized before we can send it over the wire,
+        # since pickling the proc mesh before it is initiliazed would block the
+        # tokio runtime and cause a panic.
+        await context().actor_instance.proc_mesh.initialized
         await (
             await get_or_spawn_controller("rdma_controller", RdmaController)
         ).init_rdma_on_mesh.call_one(none_throws(context().actor_instance.proc_mesh))

python/tests/rdma_load_test.py

@@ -11,7 +11,6 @@
 import statistics
 import time
 
-import pytest
 
 # parse up front to extract env variables.
 args = None
@@ -63,16 +62,10 @@
 
 # pyre-ignore
 import torch
-from monarch._src.actor.v1 import enabled as v1_enabled
 from monarch.actor import Actor, endpoint, this_host
 from monarch.rdma import RDMABuffer
 
 
-pytestmark: pytest.MarkDecorator = pytest.mark.skipif(
-    v1_enabled, reason="ENABLE ME ASAP ONCE V1 RDMA LANDS"
-)
-
-
 class RDMATest(Actor):
     def __init__(
         self, device: str = "cpu", operation: str = "write", size_mb: int = 64

python/tests/test_rdma.py

@@ -12,14 +12,10 @@
 
 import pytest
 import torch
-from monarch._src.actor.v1 import enabled as v1_enabled
 from monarch.actor import Actor, current_rank, endpoint, this_host
 from monarch.rdma import is_rdma_available, RDMAAction, RDMABuffer
 
 
-pytestmark = pytest.mark.skipif(v1_enabled, reason="ENABLE ASAP ONCE V1 RDMA LANDS")
-
-
 needs_cuda = pytest.mark.skipif(
     not torch.cuda.is_available(),
     reason="CUDA not available",

python/tests/test_rdma_unit.py

@@ -86,14 +86,10 @@ async def test_new_operation():
 import pytest
 
 import torch
-from monarch._src.actor.v1 import enabled as v1_enabled
 from monarch.actor import Actor, endpoint, this_host
 from monarch.rdma import is_rdma_available, RDMABuffer
 
 
-pytestmark = pytest.mark.skipif(v1_enabled, reason="ENABLE ASAP ONCE V1 RDMA LANDS")
-
-
 TIMEOUT = 60  # 60 seconds
 
 
@@ -546,8 +542,8 @@ def _test_with_all_data(func):
     @pytest.mark.parametrize("controller_device", CONTROLLER_DEVICES)
     @pytest.mark.parametrize("receiver_device", RECEIVER_DEVICES)
     @pytest.mark.asyncio
-    def marked(dtype, data_getter, controller_device, receiver_device):
-        return func(dtype, data_getter, controller_device, receiver_device)
+    async def marked(dtype, data_getter, controller_device, receiver_device):
+        return await func(dtype, data_getter, controller_device, receiver_device)
 
     return marked
 
@@ -560,8 +556,8 @@ def _test_with_no_data(func):
     @pytest.mark.parametrize("controller_device", CONTROLLER_DEVICES)
     @pytest.mark.parametrize("receiver_device", RECEIVER_DEVICES)
     @pytest.mark.asyncio
-    def marked(dtype, data_getter, controller_device, receiver_device):
-        return func(dtype, data_getter, controller_device, receiver_device)
+    async def marked(dtype, data_getter, controller_device, receiver_device):
+        return await func(dtype, data_getter, controller_device, receiver_device)
 
     return marked
 

python/tests/test_rdma_unsupported.py

@@ -14,13 +14,9 @@
 """
 
 import pytest
-from monarch._src.actor.v1 import enabled as v1_enabled
 from monarch.rdma import is_rdma_available
 
 
-pytestmark = pytest.mark.skipif(v1_enabled, reason="ENABLE ASAP ONCE V1 RDMA LANDS")
-
-
 needs_no_rdma = pytest.mark.skipif(
     is_rdma_available(),
     reason="RDMA is available, test only runs on systems without RDMA support",