Skip to content

Commit 01ce63d

Browse files
committed
[monarch] make test_python_actor.py work on mac
Properly skip things that require tensor engine and/or cuda ``` USE_TENSOR_ENGINE=0 pip install --no-build-isolation -e . pytest python/tests/test_python_actors.py ``` ghstack-source-id: 0d553e6 Pull-Request: #269
1 parent a2d4277 commit 01ce63d

File tree

1 file changed

+22
-1
lines changed

1 file changed

+22
-1
lines changed

python/tests/test_python_actors.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,12 @@
2929
)
3030
from monarch.debugger import init_debugging
3131

32-
from monarch.mesh_controller import spawn_tensor_engine
32+
from monarch._rust_bindings import has_tensor_engine
33+
34+
if has_tensor_engine():
35+
from monarch.mesh_controller import spawn_tensor_engine
36+
else:
37+
spawn_tensor_engine = None
3338

3439
from monarch.proc_mesh import local_proc_mesh, proc_mesh
3540
from monarch.rdma import RDMABuffer
@@ -114,6 +119,10 @@ async def get_buffer(self):
114119
return self.buffer
115120

116121

122+
@pytest.mark.skipif(
123+
not torch.cuda.is_available(),
124+
reason="CUDA not available",
125+
)
117126
async def test_proc_mesh_rdma():
118127
proc = await proc_mesh(gpus=1)
119128
server = await proc.spawn("server", ParameterServer)
@@ -282,6 +291,10 @@ async def update_weights(self):
282291
), f"{torch.sum(self.generator.weight.data)=}, {self.step=}"
283292

284293

294+
@pytest.mark.skipif(
295+
not torch.cuda.is_available(),
296+
reason="CUDA not available",
297+
)
285298
async def test_gpu_trainer_generator():
286299
trainer_proc = await proc_mesh(gpus=1)
287300
gen_proc = await proc_mesh(gpus=1)
@@ -311,6 +324,10 @@ async def test_sync_actor():
311324
assert r == 5
312325

313326

327+
@pytest.mark.skipif(
328+
not torch.cuda.is_available(),
329+
reason="CUDA not available",
330+
)
314331
def test_gpu_trainer_generator_sync() -> None:
315332
trainer_proc = proc_mesh(gpus=1).get()
316333
gen_proc = proc_mesh(gpus=1).get()
@@ -391,6 +408,10 @@ def check(module, path):
391408
check(bindings, "monarch._rust_bindings")
392409

393410

411+
@pytest.mark.skipif(
412+
not has_tensor_engine(),
413+
reason="Tensor engine not available",
414+
)
394415
@pytest.mark.skipif(
395416
torch.cuda.device_count() < 2,
396417
reason="Not enough GPUs, this test requires at least 2 GPUs",

0 commit comments

Comments
 (0)