ci: parallelize the server test suite (in-process concurrency, ~17min → faster locally + CI) - #1577
Merged
Merged
Conversation
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 <wprazuch@nvidia.com>
wprazuch
changed the title
ko3n1g
approved these changes
Jun 12, 2026
Contributor
|
/ok to test fcbad0e |
ko3n1g
enabled auto-merge (squash)
June 12, 2026 11:20
Contributor
|
/ok to test 9c6db4c |
wprazuch
changed the title
wprazuch
added a commit
that referenced
this pull request
Jun 22, 2026
…n → ~2-3min) (#1577) The **Unit tests → `Test`** check runs `ng_test_all`, which iterates all ~114 server modules **strictly sequentially** in a single runner (build venv → pytest → delete venv per module), taking **~22 min**. This is the long pole on every PR that touches core files (which trigger the full suite via the existing change-detection). **`ng_test_all` sharding** (`nemo_gym/cli.py`): new `+num_shards` / `+shard_index` options. Modules are partitioned **round-robin over a sorted list** so heavy modules (longmt_eval, the scientific servers) spread evenly across shards. The `fail_on_total_and_test_mismatch` check still runs against the full (unsharded) module list. Extracted `_select_shard()` with unit tests. **Workflow restructure** (`.github/workflows/unit-tests.yml`): - `detect` — change classification (unchanged logic), exposed as job outputs. - `Test` — core library unit tests (`ng_dev_test`) + the changed-servers path. (Fast now — the full server loop moved out.) - `Server suite (shard N)` — `ng_test_all +num_shards=8 +shard_index=N` across an **8-way matrix** (`fail-fast: false`), only on a full run. - `Server suite` — aggregation job that gates on all shards (single required check). **Result: full-suite wall-time ~22 min → ~2-3 min.** Same total compute, parallelized. Also pins uv to **0.11.19** (0.11.20 has a resolver regression that drops pinned deps — see #1576; this overlaps #1576 on `unit-tests.yml`, so merge #1576 first or fold them). Server tests moved out of the `Test` job into the sharded jobs. To preserve gating, add **`Server suite`** as a required status check alongside **`Test`**. (`Server suite` uses `if: always()` so it reports green when the full suite isn't needed for a change, and red if any shard fails — so it's safe to require unconditionally.) - `_select_shard` unit tests: complete + disjoint partition, balanced sizes, round-robin stride, out-of-range guard. `pytest tests/unit_tests/test_cli.py` 14/14, ruff clean, workflow YAML parses. - The workflow itself can only be fully validated once it runs on CI. Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
ng_test_allruns all ~114 server modules strictly sequentially in one process (build venv →pytest → delete venv per module), ~17 min. This is the long pole on every full-suite PR.
Approach (per review)
Run the suite concurrently within a single
ng_test_allinvocation rather than shardingacross CI matrix jobs — so local
ng_test_allgets the same speedup, not just CI.ng_test_all: new+max_concurrency(default1= sequential, fully backwards-compatible).Modules run in a
ThreadPoolExecutor; each still runs in its own isolated subprocess + venv, soisolation is unchanged. Per-module output is captured and printed atomically so concurrent
logs don't interleave. Extracted
_run_module_tests(per-module worker) and_run_module_tests_all(orchestration) — both unit-tested.run_command: optionalcapture=to pipe combined stdout/stderr in text mode. Strictlyadditive — the non-capture path is byte-identical (existing tests unchanged).
unit-tests.yml: single job (no matrix). Runsng_test_all +max_concurrency=${TEST_CONCURRENCY|4}onruns-on: ${vars.TEST_RUNNER||ubuntu-latest}.The team can point
TEST_RUNNERat a larger runner and bumpTEST_CONCURRENCYvia repovariables — no code change needed. Keeps the uv 0.11.19 pin.
Why concurrency helps even on a 2-core runner
Per-module time is dominated by venv build (network/disk I/O) + pytest startup, so running several
modules at once overlaps the waiting. A larger
TEST_RUNNERamplifies it further.Testing
_run_module_tests_all: sequential order preserved; concurrent runs every module exactly once;concurrency genuinely overlaps (in-flight > 1).
run_commandcapture branch: pipes stdout+STDOUT, text mode (asserted); non-capture pathunchanged.
pytest tests/unit_tests/test_cli.py tests/unit_tests/test_cli_setup_command.py54/54, ruff clean,workflow YAML parses. The full concurrent
ng_test_allis validated on CI.Notes
+num_shards).TEST_RUNNER(larger runner label) andTEST_CONCURRENCYto dial in the speedup. Defaults are safe (ubuntu-latest, concurrency 4).unit-tests.yml; merge ci: pin uv to 0.11.19 (0.11.20 resolver regression breaks the test suite) #1576 first or fold.