From fcbad0eab5195f95846e93e13bd5a1c55aa2a66e Mon Sep 17 00:00:00 2001 From: Wojciech Prazuch Date: Thu, 11 Jun 2026 12:24:22 +0200 Subject: [PATCH] ci: shard the full server test suite across 8 parallel runners ng_test_all ran all ~114 server modules sequentially in one job (~17 min). Add +num_shards/+shard_index to ng_test_all (round-robin over sorted module paths so heavy modules spread evenly) and split the Unit tests workflow: - detect: change classification (outputs run_full/run_servers) - Test: core library unit tests (ng_dev_test) + the changed-servers path - Server suite (shard N): full ng_test_all sharded across an 8-way matrix - Server suite: aggregation gate over the shards (single required check) Cuts full-suite wall-time from ~17 min to ~2-3 min. Also pins uv to 0.11.19 (0.11.20 resolver regression drops pinned deps). Branch protection note: add 'Server suite' as a required check alongside 'Test' to preserve server-suite gating (server tests moved out of the Test job). Signed-off-by: Wojciech Prazuch --- .github/workflows/unit-tests.yml | 113 ++++++++++++++++++++++++++----- nemo_gym/cli.py | 42 +++++++++++- tests/unit_tests/test_cli.py | 33 +++++++++ 3 files changed, 169 insertions(+), 19 deletions(-) diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index 85124757a4..0056f3ed27 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -23,10 +23,21 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true +env: + # Number of parallel shards for the full server suite. Keep in sync with the matrix below. + NUM_SHARDS: "8" + # Pin uv: 0.11.20 has a resolver regression that silently drops pinned deps from + # `uv pip install -r requirements.txt`. 0.11.19 is the latest known-good version. + UV_INSTALL_URL: "https://astral.sh/uv/0.11.19/install.sh" + jobs: - test: - name: Test + detect: + name: Detect changes runs-on: ubuntu-latest + outputs: + run_full: ${{ steps.changes.outputs.run_full }} + run_servers: ${{ steps.changes.outputs.run_servers }} + server_all_changed_files: ${{ steps.changed-files.outputs.server_all_changed_files }} steps: - name: Checkout repository uses: actions/checkout@v6 @@ -114,8 +125,20 @@ jobs: fi echo "============================================" + # Core library unit tests (run once on a full run) + the changed-servers path. The full + # server suite is sharded into the parallel `server-suite` matrix below. + test: + name: Test + needs: detect + if: needs.detect.outputs.run_full == 'true' || needs.detect.outputs.run_servers == 'true' + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v6 + with: + submodules: 'recursive' + - name: Cache uv dependencies - if: steps.changes.outputs.run_full == 'true' || steps.changes.outputs.run_servers == 'true' uses: actions/cache@v4 with: path: ~/.cache/uv @@ -124,7 +147,6 @@ jobs: uv-${{ runner.os }}- - name: Setup for test - if: steps.changes.outputs.run_full == 'true' || steps.changes.outputs.run_servers == 'true' run: | sudo apt-get update # Curl is required for setup_nvidia.sh to download uv @@ -132,25 +154,24 @@ jobs: sudo apt-get install -y --no-install-recommends git curl ca-certificates # The flow below should be used and synced with any Docker or container related flows. There is no script here to keep it 100% explicit. # This is how we test and this is how you should use/consume. - curl -LsSf https://astral.sh/uv/install.sh | sh + curl -LsSf "$UV_INSTALL_URL" | sh uv venv --python 3.12 source .venv/bin/activate uv sync --extra dev - name: Test - if: steps.changes.outputs.run_full == 'true' || steps.changes.outputs.run_servers == 'true' run: | source .venv/bin/activate - # Full suite: core library tests + all server tests - if [[ "${{ steps.changes.outputs.run_full }}" == "true" ]]; then - echo "Running full test suite" + # Full suite: core library unit tests run here once; all server tests run in the + # parallel `server-suite` matrix. + if [[ "${{ needs.detect.outputs.run_full }}" == "true" ]]; then + echo "Running core library unit tests (server suite runs in the sharded matrix)" ng_dev_test - ng_test_all +fail_on_total_and_test_mismatch=true +delete_venvs_after_each_test=true - # Server-only: test only the changed servers - elif [[ "${{ steps.changes.outputs.run_servers }}" == "true" ]]; then - CHANGED_SERVERS=$(echo "${{ steps.changed-files.outputs.server_all_changed_files }}" | \ + # Server-only: test only the changed servers (typically a small set, no sharding needed) + elif [[ "${{ needs.detect.outputs.run_servers }}" == "true" ]]; then + CHANGED_SERVERS=$(echo "${{ needs.detect.outputs.server_all_changed_files }}" | \ tr ' ' '\n' | cut -d'/' -f1-2 | sort -u) echo "Testing changed servers:" @@ -174,8 +195,68 @@ jobs: echo "Some server tests failed" exit 1 fi + fi - # Docs-only: nothing to test - else - echo "No tests to run" + # Full server suite, split across NUM_SHARDS parallel runners. Only runs on a full run. + server-suite: + name: Server suite (shard ${{ matrix.shard }}) + needs: detect + if: needs.detect.outputs.run_full == 'true' + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + shard: [0, 1, 2, 3, 4, 5, 6, 7] + steps: + - name: Checkout repository + uses: actions/checkout@v6 + with: + submodules: 'recursive' + + - name: Cache uv dependencies + uses: actions/cache@v4 + with: + path: ~/.cache/uv + key: uv-${{ runner.os }}-${{ hashFiles('uv.lock') }} + restore-keys: | + uv-${{ runner.os }}- + + - name: Setup for test + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends git curl ca-certificates + curl -LsSf "$UV_INSTALL_URL" | sh + uv venv --python 3.12 + source .venv/bin/activate + uv sync --extra dev + + - name: Server tests (shard ${{ matrix.shard }}/${{ env.NUM_SHARDS }}) + run: | + source .venv/bin/activate + ng_test_all \ + +fail_on_total_and_test_mismatch=true \ + +delete_venvs_after_each_test=true \ + +num_shards=${{ env.NUM_SHARDS }} \ + +shard_index=${{ matrix.shard }} + + # Single aggregated gate for the sharded server suite — make this a required check + # (in addition to "Test") instead of the individual shard jobs. `always()` so it also + # produces a green result when the full suite isn't required for this change. + server-suite-result: + name: Server suite + needs: [detect, server-suite] + if: always() + runs-on: ubuntu-latest + steps: + - name: Aggregate shard results + run: | + if [[ "${{ needs.detect.outputs.run_full }}" != "true" ]]; then + echo "Full server suite not required for this change; nothing to aggregate." + exit 0 + fi + if [[ "${{ needs.server-suite.result }}" == "success" ]]; then + echo "All server-suite shards passed." + exit 0 fi + echo "One or more server-suite shards failed (result: ${{ needs.server-suite.result }})." + exit 1 diff --git a/nemo_gym/cli.py b/nemo_gym/cli.py index e3df40025a..b38c029f36 100644 --- a/nemo_gym/cli.py +++ b/nemo_gym/cli.py @@ -618,6 +618,31 @@ class TestAllConfig(BaseNeMoGymCLIConfig): default=False, description="Delete each server venv after its tests have been run (default: False).", ) + num_shards: int = Field( + default=1, + ge=1, + description="Total number of shards to split the server suite across (default: 1 = no sharding). " + "Used to parallelize the suite across CI runners.", + ) + shard_index: int = Field( + default=0, + ge=0, + description="Which shard (0-based) this invocation runs; must be < num_shards (default: 0).", + ) + + +def _select_shard(dir_paths: List[Path], shard_index: int, num_shards: int) -> List[Path]: + """Deterministically select this shard's subset of modules. + + Round-robin (stride) over a sorted list spreads heavy modules across shards more evenly than + contiguous chunks, which balances wall-time when the suite is parallelized across CI runners. + """ + if num_shards <= 1: + return dir_paths + assert 0 <= shard_index < num_shards, ( + f"shard_index ({shard_index}) must be in [0, num_shards) for num_shards={num_shards}" + ) + return sorted(dir_paths, key=str)[shard_index::num_shards] def test_all(): # pragma: no cover @@ -635,6 +660,15 @@ def test_all(): # pragma: no cover dir_paths = [p for p in dir_paths if (p / "README.md").exists()] print(f"Found {len(dir_paths)} modules to test:{_display_list_of_paths(dir_paths)}\n") + # Keep the full list for the total-vs-tested mismatch check below, then narrow to this shard. + full_dir_paths = dir_paths + dir_paths = _select_shard(dir_paths, test_all_config.shard_index, test_all_config.num_shards) + if test_all_config.num_shards > 1: + print( + f"Shard {test_all_config.shard_index + 1}/{test_all_config.num_shards}: " + f"testing {len(dir_paths)} of {len(full_dir_paths)} modules:{_display_list_of_paths(dir_paths)}\n" + ) + tests_passed: List[Path] = [] tests_failed: List[Path] = [] tests_missing: List[Path] = [] @@ -710,10 +744,12 @@ def test_all(): # pragma: no cover """) if test_all_config.fail_on_total_and_test_mismatch: - extra_candidates = [p for p in candidate_dir_paths if Path(p) not in dir_paths] + # Compare against the full (unsharded) module list — every module must be testable + # regardless of how many shards we split the run into. + extra_candidates = [p for p in candidate_dir_paths if Path(p) not in full_dir_paths] assert ( - len(candidate_dir_paths) == len(dir_paths) - ), f"""Mismatch on the number of total modules found ({len(candidate_dir_paths)}) and the number of actual modules tested ({len(dir_paths)})! + len(candidate_dir_paths) == len(full_dir_paths) + ), f"""Mismatch on the number of total modules found ({len(candidate_dir_paths)}) and the number of actual modules tested ({len(full_dir_paths)})! Extra candidate paths:{_display_list_of_paths(extra_candidates)}""" diff --git a/tests/unit_tests/test_cli.py b/tests/unit_tests/test_cli.py index 74f7fce578..d6208b5de4 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -31,12 +31,45 @@ _GRACEFUL_SHUTDOWN_TIMEOUT_SEC, RunConfig, RunHelper, + _select_shard, display_help, init_resources_server, ) from nemo_gym.config_types import ResourcesServerInstanceConfig +class TestSelectShard: + def test_no_sharding_returns_all(self) -> None: + paths = [Path(f"resources_servers/s{i}") for i in range(5)] + assert _select_shard(paths, shard_index=0, num_shards=1) == paths + + def test_round_robin_partition_is_complete_and_disjoint(self) -> None: + paths = [Path(f"resources_servers/s{i:02d}") for i in range(10)] + num_shards = 4 + shards = [_select_shard(paths, i, num_shards) for i in range(num_shards)] + # Every module appears in exactly one shard, and the union is the full sorted set. + flattened = [p for shard in shards for p in shard] + assert sorted(flattened, key=str) == sorted(paths, key=str) + assert len(flattened) == len(set(flattened)) == len(paths) + # Round-robin stride: shard 0 gets indices 0,4,8 of the sorted list. + assert shards[0] == [ + Path("resources_servers/s00"), + Path("resources_servers/s04"), + Path("resources_servers/s08"), + ] + + def test_balanced_sizes(self) -> None: + paths = [Path(f"resources_servers/s{i:02d}") for i in range(10)] + sizes = sorted(len(_select_shard(paths, i, 4)) for i in range(4)) + # 10 across 4 shards -> sizes differ by at most 1. + assert sizes[-1] - sizes[0] <= 1 + + def test_shard_index_out_of_range_raises(self) -> None: + paths = [Path("resources_servers/s0")] + with raises(AssertionError): + _select_shard(paths, shard_index=4, num_shards=4) + + # TODO: Eventually we want to add more tests to ensure that the CLI flows do not break class TestCLI: def test_sanity(self) -> None: