machichima commented on code in PR #745:
URL: https://github.com/apache/mahout/pull/745#discussion_r2636967236
##########
qdp/qdp-python/tests/test_bindings.py:
##########
@@ -50,6 +50,36 @@ def test_dlpack_device():
assert device_info == (2, 0), "Expected (2, 0) for CUDA device 0"
[email protected]
+def test_dlpack_device_id_non_zero():
+ """Test device_id propagation for non-zero devices (requires multi-GPU)."""
+ pytest.importorskip("torch")
+ import torch
+ from mahout_qdp import QdpEngine
+
+ # Check if multiple GPUs are available
+ if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+ pytest.skip("Multi-GPU setup required for this test")
Review Comment:
nit: It's better to use `pytest.mark.skipif` to skip before any setup runs
and make it clearer that the test is skipped in the first glance
```py
@pytest.mark.skipif(
not __import__("torch").cuda.is_available() or
__import__("torch").cuda.device_count() < 2,
reason="Multi-GPU setup required for this test"
)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
