Skip to content

Commit 05aa8a8

Browse files
vmoensvmoens
committed
[CI] Separate GPU and CPU tests with pytest markers (#3404)
Add pytest.mark.gpu to tests that require CUDA, and update run_all.sh to filter tests based on whether running on GPU or CPU machines. Changes: - Register 'gpu' marker in pytest.ini and conftest.py - Add pytest.mark.gpu to ~30 tests that explicitly require CUDA - Update run_all.sh to use GPU_MARKER_FILTER: - GPU jobs (CU_VERSION != cpu): run only pytest.mark.gpu tests - CPU jobs (CU_VERSION = cpu): run all tests except pytest.mark.gpu This significantly reduces GPU machine usage by running only GPU-requiring tests on expensive GPU runners (~30 tests instead of ~2000+). Tests that can run on either device will run on CPU machines only. The optimization can be disabled by setting TORCHRL_GPU_FILTER=0. ghstack-source-id: cf14d1e Pull-Request: #3404
1 parent cfeb61a commit 05aa8a8

File tree

14 files changed

+133
-0
lines changed

14 files changed

+133
-0
lines changed

‎.github/unittest/linux/scripts/run_all.sh‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,28 @@ fi
270270

271271
TORCHRL_TEST_SUITE="${TORCHRL_TEST_SUITE:-all}" # all|distributed|nondistributed
272272

273+
# GPU test filtering: Run GPU-only tests on GPU machines, CPU-only tests on CPU machines.
274+
# This avoids running ~2000+ tests on expensive GPU machines when only ~30 require GPU.
275+
# Tests are marked with @pytest.mark.gpu if they require CUDA.
276+
#
277+
# Set TORCHRL_GPU_FILTER=0 to disable this optimization and run all tests.
278+
#
279+
# We use an array to handle the marker expression properly (avoids quoting issues).
280+
GPU_MARKER_FILTER=()
281+
if [ "${TORCHRL_GPU_FILTER:-1}" = "1" ]; then
282+
if [ "${CU_VERSION:-}" == cpu ]; then
283+
# CPU job: run only tests that do NOT require GPU
284+
GPU_MARKER_FILTER=(-m 'not gpu')
285+
echo "GPU filtering enabled: Running CPU-only tests (excluding @pytest.mark.gpu)"
286+
else
287+
# GPU job: run only tests that require GPU
288+
GPU_MARKER_FILTER=(-m gpu)
289+
echo "GPU filtering enabled: Running GPU-only tests (@pytest.mark.gpu)"
290+
fi
291+
else
292+
echo "GPU filtering disabled: Running all tests"
293+
fi
294+
273295
export PYTORCH_TEST_WITH_SLOW='1'
274296
python -m torch.utils.collect_env
275297

@@ -292,6 +314,7 @@ run_distributed_tests() {
292314
local json_report_args="--json-report --json-report-file=${json_report_dir}/test-results-distributed.json --json-report-indent=2"
293315

294316
# Run both test_distributed.py and test_rb_distributed.py (both use torch.distributed)
317+
# Note: distributed tests always run on GPU, no need for GPU_MARKER_FILTER here
295318
python .github/unittest/helpers/coverage_run_parallel.py -m pytest test/test_distributed.py test/test_rb_distributed.py \
296319
${json_report_args} \
297320
--instafail --durations 200 -vv --capture no \
@@ -327,12 +350,14 @@ run_non_distributed_tests() {
327350
1)
328351
echo "Running shard 1: test_transforms.py only"
329352
python .github/unittest/helpers/coverage_run_parallel.py -m pytest test/test_transforms.py \
353+
"${GPU_MARKER_FILTER[@]}" \
330354
${json_report_args} \
331355
${common_args}
332356
;;
333357
2)
334358
echo "Running shard 2: test_envs.py and test_collectors.py"
335359
python .github/unittest/helpers/coverage_run_parallel.py -m pytest test/test_envs.py test/test_collectors.py \
360+
"${GPU_MARKER_FILTER[@]}" \
336361
${json_report_args} \
337362
${common_args}
338363
;;
@@ -344,13 +369,15 @@ run_non_distributed_tests() {
344369
--ignore test/test_envs.py \
345370
--ignore test/test_collectors.py \
346371
${xdist_args} \
372+
"${GPU_MARKER_FILTER[@]}" \
347373
${json_report_args} \
348374
${common_args}
349375
;;
350376
all|"")
351377
echo "Running all tests (no sharding)"
352378
python .github/unittest/helpers/coverage_run_parallel.py -m pytest test \
353379
${common_ignores} \
380+
"${GPU_MARKER_FILTER[@]}" \
354381
${json_report_args} \
355382
${common_args}
356383
;;

‎.github/workflows/validate-test-partitioning.yml‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
# Validates that GPU/CPU test partitioning covers all tests
2+
#
3+
# This workflow ensures that:
4+
# 1. Tests marked with @pytest.mark.gpu + tests not marked = all tests
5+
# 2. No tests are accidentally excluded from CI
6+
#
7+
# Runs on PRs to catch partitioning issues before merge.
8+
name: Validate Test Partitioning
9+
10+
on:
11+
pull_request:
12+
push:
13+
branches: [main, nightly]
14+
workflow_dispatch:
15+
16+
concurrency:
17+
group: ${{ github.workflow }}-${{ github.ref }}
18+
cancel-in-progress: true
19+
20+
jobs:
21+
validate:
22+
runs-on: ubuntu-latest
23+
timeout-minutes: 10
24+
steps:
25+
- uses: actions/checkout@v4
26+
27+
- name: Set up Python
28+
uses: actions/setup-python@v5
29+
with:
30+
python-version: "3.12"
31+
32+
- name: Install minimal dependencies
33+
run: |
34+
pip install pytest tensordict
35+
pip install torch --index-url https://download.pytorch.org/whl/cpu
36+
# Skip editable install - just add to PYTHONPATH for test collection
37+
38+
- name: Validate test partitioning
39+
env:
40+
PYTHONPATH: ${{ github.workspace }}
41+
run: |
42+
set -e
43+
44+
echo "=================================================="
45+
echo " TEST PARTITIONING VALIDATION"
46+
echo "=================================================="
47+
48+
# Collect test counts
49+
# Note: We ignore test_loggers.py due to torchvision operator issues on CPU-only
50+
ALL=$(pytest --collect-only -q test/ --ignore test/test_loggers.py 2>/dev/null | tail -1 | grep -oE "^[0-9]+")
51+
GPU=$(pytest --collect-only -q -m gpu test/ --ignore test/test_loggers.py 2>/dev/null | tail -1 | grep -oE "^[0-9]+")
52+
CPU=$(pytest --collect-only -q -m "not gpu" test/ --ignore test/test_loggers.py 2>/dev/null | tail -1 | grep -oE "^[0-9]+")
53+
54+
echo ""
55+
echo "Total tests: $ALL"
56+
echo "GPU tests (@pytest.mark.gpu): $GPU"
57+
echo "CPU tests (not gpu): $CPU"
58+
echo "GPU + CPU: $((GPU + CPU))"
59+
echo ""
60+
61+
# Validate: GPU + CPU should equal ALL
62+
if [ "$((GPU + CPU))" -eq "$ALL" ]; then
63+
echo "✅ PASS: Test partitioning is valid!"
64+
echo " All tests are accounted for."
65+
else
66+
echo "❌ FAIL: Test partitioning mismatch!"
67+
echo " GPU ($GPU) + CPU ($CPU) = $((GPU + CPU)), but total is $ALL"
68+
echo ""
69+
echo " This means some tests are either:"
70+
echo " - Missing the @pytest.mark.gpu marker (if they require CUDA)"
71+
echo " - Being excluded unintentionally"
72+
exit 1
73+
fi

‎pytest.ini‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ addopts =
66
--tb=native
77
markers =
88
unity_editor
9+
slow: mark test as slow to run
10+
gpu: mark test as requiring a GPU (CUDA device)
911
testpaths =
1012
test
1113
xfail_strict = True

‎test/compile/test_compile_collectors.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ def test_compiled_policy(self, collector_cls, compile_policy, device):
7777
collector.shutdown()
7878
del collector
7979

80+
@pytest.mark.gpu
8081
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
8182
@pytest.mark.parametrize(
8283
"collector_cls",

‎test/conftest.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,9 @@ def pytest_runtest_setup(item):
148148

149149
def pytest_configure(config):
150150
config.addinivalue_line("markers", "slow: mark test as slow to run")
151+
config.addinivalue_line(
152+
"markers", "gpu: mark test as requiring a GPU (CUDA device)"
153+
)
151154

152155

153156
def pytest_collection_modifyitems(config, items):

‎test/llm/test_llm_updaters.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ def get_open_port():
7272
)
7373

7474

75+
@pytest.mark.gpu
7576
@pytest.mark.skipif(not _has_transformers, reason="missing transformers dependencies")
7677
@pytest.mark.skipif(not _has_vllm, reason="missing vllm dependencies")
7778
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@@ -415,6 +416,7 @@ def test_local_llm_specific_features(self, target_vllm_engine):
415416
"See LLM_TEST_ISSUES.md for details.",
416417
strict=False,
417418
)
419+
@pytest.mark.gpu
418420
@pytest.mark.skipif(not _has_ray, reason="missing ray dependencies")
419421
@pytest.mark.skipif(not _has_vllm, reason="missing vllm dependencies")
420422
@pytest.mark.skipif(not _has_transformers, reason="missing transformers dependencies")
@@ -611,6 +613,7 @@ def test_weight_sync_vllm_collective_ray(self, request):
611613
ray.shutdown()
612614

613615

616+
@pytest.mark.gpu
614617
@pytest.mark.xfail(
615618
reason="AsyncVLLM tests fail due to Ray placement group timeout. "
616619
"See LLM_TEST_ISSUES.md for details.",

‎test/llm/test_vllm.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ def sampling_params():
3939
class TestAsyncVLLMIntegration:
4040
"""Integration tests for AsyncVLLM with real models."""
4141

42+
@pytest.mark.gpu
4243
@pytest.mark.skipif(not _has_vllm, reason="vllm not available")
4344
@pytest.mark.skipif(not _has_ray, reason="ray not available")
4445
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@@ -111,6 +112,7 @@ def test_vllm_api_compatibility(self, sampling_params):
111112
finally:
112113
service.shutdown()
113114

115+
@pytest.mark.gpu
114116
@pytest.mark.skipif(not _has_vllm, reason="vllm not available")
115117
@pytest.mark.skipif(not _has_ray, reason="ray not available")
116118
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

‎test/llm/test_wrapper.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2104,6 +2104,7 @@ def test_log_probs_consistency(
21042104
"See LLM_TEST_ISSUES.md for details.",
21052105
strict=False,
21062106
)
2107+
@pytest.mark.gpu
21072108
@pytest.mark.skipif(not _has_vllm, reason="vllm not available")
21082109
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
21092110
def test_sync_async_vllm_strict_equivalence(

‎test/test_collectors.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2333,6 +2333,7 @@ def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
23332333
def _set_seed(self, seed: int | None) -> None:
23342334
...
23352335

2336+
@pytest.mark.gpu
23362337
@pytest.mark.skipif(not torch.cuda.is_available(), reason="no cuda device")
23372338
@pytest.mark.parametrize("env_device", ["cuda:0", "cpu"])
23382339
@pytest.mark.parametrize("storing_device", [None, "cuda:0", "cpu"])
@@ -2371,6 +2372,7 @@ def test_no_synchronize(self, env_device, storing_device, no_cuda_sync):
23712372
assert u == i, i
23722373
mock_synchronize.assert_not_called()
23732374

2375+
@pytest.mark.gpu
23742376
@pytest.mark.parametrize("device", ["cuda", "cpu"])
23752377
@pytest.mark.parametrize("storing_device", ["cuda", "cpu"])
23762378
@pytest.mark.skipif(not torch.cuda.is_available(), reason="no cuda device found")
@@ -3162,6 +3164,7 @@ def test_multi_collector_consistency(
31623164
assert_allclose_td(c2.unsqueeze(0), d2)
31633165

31643166

3167+
@pytest.mark.gpu
31653168
@pytest.mark.skipif(
31663169
not torch.cuda.is_available() and (not has_mps()),
31673170
reason="No casting if no cuda",
@@ -3363,6 +3366,7 @@ def test_param_sync_mixed_device(
33633366
col.shutdown()
33643367
del col
33653368

3369+
@pytest.mark.gpu
33663370
@pytest.mark.skipif(
33673371
not torch.cuda.is_available() or torch.cuda.device_count() < 3,
33683372
reason="requires at least 3 CUDA devices",

‎test/test_envs.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -597,6 +597,7 @@ def test_auto_spec(self, env_type):
597597
env.auto_specs_(policy, tensordict=td.copy(), observation_key=obs_vals)
598598
env.check_env_specs(tensordict=td.copy())
599599

600+
@pytest.mark.gpu
600601
@pytest.mark.skipif(not torch.cuda.device_count(), reason="No cuda device found.")
601602
@pytest.mark.parametrize("break_when_any_done", [True, False])
602603
def test_auto_cast_to_device(self, break_when_any_done):
@@ -1526,6 +1527,7 @@ def test_parallel_env_with_policy(
15261527
# env_serial.close()
15271528
env0.close(raise_if_closed=False)
15281529

1530+
@pytest.mark.gpu
15291531
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
15301532
@pytest.mark.parametrize("heterogeneous", [False, True])
15311533
def test_transform_env_transform_no_device(
@@ -1638,6 +1640,7 @@ def test_parallel_env_custom_method(self, parallel, maybe_fork_ParallelEnv):
16381640
finally:
16391641
env.close(raise_if_closed=False)
16401642

1643+
@pytest.mark.gpu
16411644
@pytest.mark.skipif(not torch.cuda.device_count(), reason="no cuda to test on")
16421645
@pytest.mark.skipif(not _has_gym, reason="no gym")
16431646
@pytest.mark.parametrize("frame_skip", [4])
@@ -1742,6 +1745,7 @@ def test_parallel_env_cast(
17421745
env_serial.close(raise_if_closed=False)
17431746
env0.close(raise_if_closed=False)
17441747

1748+
@pytest.mark.gpu
17451749
@pytest.mark.skipif(not _has_gym, reason="no gym")
17461750
@pytest.mark.skipif(not torch.cuda.device_count(), reason="no cuda device detected")
17471751
@pytest.mark.parametrize("frame_skip", [4])
@@ -2726,6 +2730,7 @@ def test_marl_group_type(group_type):
27262730
check_marl_grouping(group_type.get_group_map(agent_names), agent_names)
27272731

27282732

2733+
@pytest.mark.gpu
27292734
@pytest.mark.skipif(not torch.cuda.device_count(), reason="No cuda device")
27302735
class TestConcurrentEnvs:
27312736
"""Concurrent parallel envs on multiple procs can interfere."""

0 commit comments

Comments
 (0)