ci: parallelize the server test suite (in-process concurrency, ~17min → faster locally + CI)

.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,33 +147,31 @@ 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
           # ca-certificates is there to support curl and mitigate `curl: (77) error setting certificate file: /etc/ssl/certs/ca-certificates.crt`
           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

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)}"""
 

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: