include pytest-asyncio and timeout

Differential Revision: D75028810

Author: amirafzali

.github/workflows/test.yml

@@ -55,7 +55,7 @@ jobs:
         pip install pyzmq requests numpy pyre-extensions
 
         # Test dependencies
-        pip install pytest cloudpickle
+        pip install pytest cloudpickle pytest-timeout pytest-asyncio
 
         # Build and install monarch
         python setup.py install

README.md

@@ -28,7 +28,7 @@ pip install setuptools-rust
 # install torch, can use conda or build it yourself or whatever
 pip install torch
 # install other deps, see pyproject.toml for latest
-pip install pyzmq requests numpy pyre-extensions pytest-timeout cloudpickle
+pip install pyzmq requests numpy pyre-extensions pytest-timeout cloudpickle pytest-asyncio
 
 # install the package
 python setup.py install

python/tests/_monarch/test_hyperactor.py

@@ -1,4 +1,5 @@
 # pyre-strict
+# pyre-ignore-all-errors[56]
 
 import multiprocessing
 import os
@@ -73,12 +74,14 @@ def test_fn() -> None:
     assert code == signal.SIGTERM, code
 
 
+@pytest.mark.asyncio
 async def test_allocator() -> None:
     spec = AllocSpec(AllocConstraints(), replica=2)
     allocator = monarch.LocalAllocator()
     _ = await allocator.allocate(spec)
 
 
+@pytest.mark.asyncio
 async def test_proc_mesh() -> None:
     spec = AllocSpec(AllocConstraints(), replica=2)
     allocator = monarch.LocalAllocator()
@@ -87,6 +90,7 @@ async def test_proc_mesh() -> None:
     assert str(proc_mesh) == "<ProcMesh { shape: {replica=2} }>"
 
 
+@pytest.mark.asyncio
 async def test_actor_mesh() -> None:
     spec = AllocSpec(AllocConstraints(), replica=2)
     allocator = monarch.LocalAllocator()

python/tests/test_python_actors.py

@@ -1,5 +1,7 @@
 import operator
 
+import pytest
+
 import torch
 
 from monarch.proc_mesh import local_proc_mesh, proc_mesh
@@ -53,6 +55,7 @@ async def get_grad_buffer(self) -> torch.Tensor:
         return self.grad_buffer
 
 
+@pytest.mark.asyncio
 async def test_choose():
     proc = await local_proc_mesh(gpus=2)
     v = await proc.spawn("counter", Counter, 3)
@@ -64,6 +67,7 @@ async def test_choose():
     assert result == result2
 
 
+@pytest.mark.asyncio
 async def test_stream():
     proc = await local_proc_mesh(gpus=2)
     v = await proc.spawn("counter2", Counter, 3)
@@ -93,6 +97,7 @@ async def get_buffer(self):
         return self.buffer
 
 
+@pytest.mark.asyncio
 async def test_proc_mesh_rdma():
     proc = await proc_mesh(gpus=1)
     server = await proc.spawn("server", ParameterServer)
@@ -155,6 +160,7 @@ async def get(self, to: To):
         return [x async for x in to.whoami.stream()]
 
 
+@pytest.mark.asyncio
 async def test_mesh_passed_to_mesh():
     proc = await local_proc_mesh(gpus=2)
     f = await proc.spawn("from", From)
@@ -164,6 +170,7 @@ async def test_mesh_passed_to_mesh():
     assert all[0] != all[1]
 
 
+@pytest.mark.asyncio
 async def test_mesh_passed_to_mesh_on_different_proc_mesh():
     proc = await local_proc_mesh(gpus=2)
     proc2 = await local_proc_mesh(gpus=2)
@@ -174,6 +181,7 @@ async def test_mesh_passed_to_mesh_on_different_proc_mesh():
     assert all[0] != all[1]
 
 
+@pytest.mark.asyncio
 async def test_actor_slicing():
     proc = await local_proc_mesh(gpus=2)
     proc2 = await local_proc_mesh(gpus=2)
@@ -189,6 +197,7 @@ async def test_actor_slicing():
     assert result[0] == result[1]
 
 
+@pytest.mark.asyncio
 async def test_aggregate():
     proc = await local_proc_mesh(gpus=2)
     counter = await proc.spawn("counter", Counter, 1)
@@ -204,6 +213,7 @@ async def run(self, fn):
         return fn()
 
 
+@pytest.mark.asyncio
 async def test_rank_size():
     proc = await local_proc_mesh(gpus=2)
     r = await proc.spawn("runit", RunIt)
@@ -261,6 +271,7 @@ async def update_weights(self):
         ), f"{torch.sum(self.generator.weight.data)=}, {self.step=}"
 
 
+@pytest.mark.asyncio
 async def test_gpu_trainer_generator():
     trainer_proc = await proc_mesh(gpus=1)
     gen_proc = await proc_mesh(gpus=1)
@@ -282,6 +293,7 @@ def sync_endpoint(self, a_counter: Counter):
         return a_counter.value.choose().get()
 
 
+@pytest.mark.asyncio
 async def test_sync_actor():
     proc = await local_proc_mesh(gpus=2)
     a = await proc.spawn("actor", SyncActor)