ci: parallelize ng_test_all in-process (gym env test / ng_test_all)

Ports #1596 onto the unified gym CLI (#1434). Adds +max_concurrency to
ng_test_all: each module still runs in its own isolated subprocess/venv, but
up to N run concurrently via a ThreadPoolExecutor, so local ng_test_all and CI
both speed up (no matrix sharding). run_command gains an additive capture=
param so concurrent module output is collected and printed atomically.

Workflow runs on ${{ vars.TEST_RUNNER || 'ubuntu-latest' }} with
TEST_CONCURRENCY (default 8). Targets martas/1434.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

Author: wprazuch

.github/workflows/unit-tests.yml

@@ -26,7 +26,9 @@ concurrency:
 jobs:
   test:
     name: Test
-    runs-on: ubuntu-latest
+    # Runs on a standard runner by default. Set the `TEST_RUNNER` repo/org variable to a larger
+    # runner label (e.g. a multi-core runner) to give the concurrent server suite more cores.
+    runs-on: ${{ vars.TEST_RUNNER || 'ubuntu-latest' }}
     steps:
       - name: Checkout repository
         uses: actions/checkout@v6
@@ -139,14 +141,19 @@ jobs:
 
       - name: Test
         if: steps.changes.outputs.run_full == 'true' || steps.changes.outputs.run_servers == 'true'
+        env:
+          # How many module test suites to run concurrently. Each module still runs in its own
+          # isolated subprocess/venv. Defaults to 8; override via the `TEST_CONCURRENCY` repo/org
+          # variable (e.g. lower it if a runner hits memory pressure, raise it on a larger runner).
+          TEST_CONCURRENCY: ${{ vars.TEST_CONCURRENCY || '8' }}
         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"
             ng_dev_test
-            ng_test_all +fail_on_total_and_test_mismatch=true +delete_venvs_after_each_test=true
+            ng_test_all +fail_on_total_and_test_mismatch=true +delete_venvs_after_each_test=true +max_concurrency=${TEST_CONCURRENCY}
 
           # Server-only: test only the changed servers
           elif [[ "${{ steps.changes.outputs.run_servers }}" == "true" ]]; then

nemo_gym/cli/env.py

@@ -17,6 +17,8 @@
 import json
 import os
 import shlex
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
 from glob import glob
 from os import makedirs
 from os.path import exists
@@ -501,11 +503,13 @@ def _validate_data_single(test_config: TestConfig) -> None:  # pragma: no cover
     print(f"The data for {test_config.dir_path} has been successfully validated!")
 
 
-def _test_single(test_config: TestConfig, global_config_dict: DictConfig) -> Popen:  # pragma: no cover
+def _test_single(
+    test_config: TestConfig, global_config_dict: DictConfig, capture: bool = False
+) -> Popen:  # pragma: no cover
     # Eventually we may want more sophisticated testing here, but this is sufficient for now.
     prefix = test_config.entrypoint.replace("/", "\\/")
     command = f"""{setup_env_command(test_config.dir_path, global_config_dict, prefix)} && pytest"""
-    return run_command(command, test_config.dir_path)
+    return run_command(command, test_config.dir_path, capture=capture)
 
 
 def test():  # pragma: no cover
@@ -549,6 +553,71 @@ class TestAllConfig(BaseNeMoGymCLIConfig):
         default=False,
         description="Delete each server venv after its tests have been run (default: False).",
     )
+    max_concurrency: int = Field(
+        default=1,
+        ge=1,
+        description="How many module test suites to run concurrently (default: 1 = sequential). "
+        "Each module still runs in its own isolated subprocess/venv; raising this parallelizes the "
+        "suite on one machine (CI runner or local), bounded by available cores/IO.",
+    )
+
+
+@dataclass
+class _ModuleTestResult:
+    dir_path: Path
+    return_code: int
+    data_validation_failed: bool
+    elapsed_s: float
+    output: Optional[str]  # captured combined stdout/stderr when run concurrently, else None
+
+
+def _run_module_tests(
+    dir_path: Path, global_config_dict: DictConfig, delete_venv: bool, capture: bool
+) -> _ModuleTestResult:  # pragma: no cover
+    """Run one module's test suite (build venv -> pytest -> data validation -> optional cleanup).
+
+    Safe to call from multiple threads concurrently: each module builds its own venv and runs in
+    its own subprocess, and `_validate_data_single` only reads files. When `capture` is set, the
+    subprocess output is collected so the caller can print it without interleaving.
+    """
+    start_time = time()
+    test_config = TestConfig(entrypoint=str(dir_path), should_validate_data=True)
+    proc = _test_single(test_config, global_config_dict, capture=capture)
+
+    if capture:
+        output, _ = proc.communicate()
+        return_code = proc.returncode
+    else:
+        return_code = proc.wait()
+        output = None
+
+    data_validation_failed = False
+    try:
+        _validate_data_single(test_config)
+    except AssertionError:
+        data_validation_failed = True
+
+    if delete_venv:
+        rmtree(dir_path / ".venv", ignore_errors=True)
+
+    return _ModuleTestResult(dir_path, return_code, data_validation_failed, time() - start_time, output)
+
+
+def _run_module_tests_all(run_one, dir_paths: List[Path], max_concurrency: int) -> List[_ModuleTestResult]:
+    """Apply `run_one(dir_path)` over every module, up to `max_concurrency` at a time.
+
+    Returns results in completion order. With max_concurrency == 1 this is a plain sequential loop;
+    otherwise modules run concurrently in a thread pool (the heavy work happens in subprocesses).
+    """
+    if max_concurrency <= 1:
+        return [run_one(dir_path) for dir_path in tqdm(dir_paths, desc="Running tests")]
+
+    results: List[_ModuleTestResult] = []
+    with ThreadPoolExecutor(max_workers=max_concurrency) as executor:
+        futures = [executor.submit(run_one, dir_path) for dir_path in dir_paths]
+        for future in tqdm(as_completed(futures), total=len(dir_paths), desc="Running tests"):
+            results.append(future.result())
+    return results
 
 
 def test_all():  # pragma: no cover
@@ -566,22 +635,29 @@ 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")
 
+    max_concurrency = test_all_config.max_concurrency
+    capture = max_concurrency > 1
+    if capture:
+        print(f"Running up to {max_concurrency} module test suites concurrently.\n")
+
+    def run_one(dir_path: Path) -> _ModuleTestResult:
+        return _run_module_tests(dir_path, global_config_dict, test_all_config.delete_venvs_after_each_test, capture)
+
+    results = _run_module_tests_all(run_one, dir_paths, max_concurrency)
+
     tests_passed: List[Path] = []
     tests_failed: List[Path] = []
     tests_missing: List[Path] = []
     data_validation_failed: List[Path] = []
     times_taken: List[Tuple[float, Path]] = []
-    for dir_path in tqdm(dir_paths, desc="Running tests"):
-        start_time = time()
-
-        test_config = TestConfig(
-            entrypoint=str(dir_path),
-            should_validate_data=True,  # Test all always validates data.
-        )
-        proc = _test_single(test_config, global_config_dict)
-        return_code = proc.wait()
-
-        match return_code:
+    for result in results:
+        dir_path = result.dir_path
+        # When run concurrently each module's output was captured; print it atomically here so
+        # logs stay readable instead of interleaving across modules.
+        if result.output is not None:
+            print(f"\n===== {dir_path} =====\n{result.output}")
+
+        match result.return_code:
             case 0:
                 tests_passed.append(dir_path)
             case 1 | 2:
@@ -590,21 +666,14 @@ def test_all():  # pragma: no cover
                 tests_missing.append(dir_path)
             case _:
                 raise ValueError(
-                    f"""Hit unrecognized exit code {return_code} while running tests for {dir_path}.
+                    f"""Hit unrecognized exit code {result.return_code} while running tests for {dir_path}.
 You can rerun just these tests using `ng_test +entrypoint={dir_path}` or run detailed tests via `cd {dir_path} && source .venv/bin/activate && pytest`."""
                 )
 
-        try:
-            _validate_data_single(test_config)
-        except AssertionError:
+        if result.data_validation_failed:
             data_validation_failed.append(dir_path)
 
-        if test_all_config.delete_venvs_after_each_test:
-            venv_path = dir_path / ".venv"
-            print(f"Deleting {venv_path} since `delete_venvs_after_each_test=true`")
-            rmtree(venv_path, ignore_errors=True)
-
-        times_taken.append((time() - start_time, dir_path))
+        times_taken.append((result.elapsed_s, dir_path))
 
     times_taken.sort(reverse=True)
     table = Table(title="Times taken per test (sorted from highest to lowest)")

nemo_gym/cli_setup_command.py

@@ -16,7 +16,7 @@
 import os
 from os import environ
 from pathlib import Path
-from subprocess import Popen
+from subprocess import PIPE, STDOUT, Popen
 from sys import stderr, stdout
 
 from omegaconf import DictConfig
@@ -172,7 +172,7 @@ def setup_env_command(dir_path: Path, global_config_dict: DictConfig, prefix: st
     return f"cd {dir_path} && {env_setup_cmd}"
 
 
-def run_command(command: str, working_dir_path: Path, server_name: str = "") -> Popen:
+def run_command(command: str, working_dir_path: Path, server_name: str = "", capture: bool = False) -> Popen:
     global_config_dict = get_global_config_dict()
 
     work_dir = f"{working_dir_path.absolute()}"
@@ -192,13 +192,24 @@ def run_command(command: str, working_dir_path: Path, server_name: str = "") ->
         log_path.parent.mkdir(parents=True, exist_ok=True)
         command = f"set -o pipefail; ({command}) 2>&1 | tee -a {log_path}"
 
-    redirect_stdout = stdout
-    redirect_stderr = stderr
+    # When capturing, pipe stdout+stderr into the process (text mode) so callers (e.g. concurrent
+    # test runs) can collect each command's output and print it atomically instead of interleaving
+    # streams. Otherwise inherit the parent's streams exactly as before.
+    if capture:
+        return Popen(
+            command,
+            executable="/bin/bash",
+            shell=True,
+            env=custom_env,
+            stdout=PIPE,
+            stderr=STDOUT,
+            text=True,
+        )
     return Popen(
         command,
         executable="/bin/bash",
         shell=True,
         env=custom_env,
-        stdout=redirect_stdout,
-        stderr=redirect_stderr,
+        stdout=stdout,
+        stderr=stderr,
     )

tests/unit_tests/test_cli.py

@@ -31,6 +31,7 @@
     _GRACEFUL_SHUTDOWN_TIMEOUT_SEC,
     RunConfig,
     RunHelper,
+    _run_module_tests_all,
     init_resources_server,
 )
 from nemo_gym.cli.general import display_help_legacy
@@ -211,3 +212,51 @@ def test_graceful_termination_does_not_kill(self) -> None:
         assert a.wait.call_count == 1
         assert b.wait.call_count == 1
         assert runner._processes == {}
+
+
+class TestRunModuleTestsAll:
+    def test_sequential_runs_all_in_order(self) -> None:
+        paths = [Path(f"resources_servers/s{i}") for i in range(5)]
+        seen: list[Path] = []
+
+        def run_one(p: Path) -> Path:
+            seen.append(p)
+            return p
+
+        results = _run_module_tests_all(run_one, paths, max_concurrency=1)
+        assert results == paths
+        assert seen == paths  # max_concurrency=1 preserves order
+
+    def test_concurrent_runs_every_module_exactly_once(self) -> None:
+        paths = [Path(f"resources_servers/s{i:02d}") for i in range(20)]
+
+        def run_one(p: Path) -> Path:
+            return p
+
+        results = _run_module_tests_all(run_one, paths, max_concurrency=8)
+        # Completion order is non-deterministic, but every module must run exactly once.
+        assert sorted(results, key=str) == sorted(paths, key=str)
+        assert len(results) == len(set(results)) == len(paths)
+
+    def test_concurrency_actually_overlaps(self) -> None:
+        import threading
+        from time import sleep
+
+        paths = [Path(f"resources_servers/s{i}") for i in range(8)]
+        lock = threading.Lock()
+        in_flight = 0
+        max_in_flight = 0
+
+        def run_one(p: Path) -> Path:
+            nonlocal in_flight, max_in_flight
+            with lock:
+                in_flight += 1
+                max_in_flight = max(max_in_flight, in_flight)
+            sleep(0.05)
+            with lock:
+                in_flight -= 1
+            return p
+
+        _run_module_tests_all(run_one, paths, max_concurrency=4)
+        # With a pool of 4, multiple modules must have been in flight simultaneously.
+        assert max_in_flight >= 2

tests/unit_tests/test_cli_setup_command.py

@@ -274,6 +274,29 @@ def test_sanity(self, monkeypatch: MonkeyPatch) -> None:
         actual_args = Popen_mock.call_args
         assert expected_args == actual_args
 
+    def test_capture_pipes_combined_output(self, monkeypatch: MonkeyPatch) -> None:
+        from subprocess import PIPE, STDOUT
+
+        Popen_mock, _ = self._setup(monkeypatch)
+
+        run_command(
+            command="my command",
+            working_dir_path=Path("/my path"),
+            capture=True,
+        )
+
+        # capture=True pipes stdout+stderr together in text mode so callers can collect output.
+        expected_args = call(
+            "my command",
+            executable="/bin/bash",
+            shell=True,
+            env={"PYTHONPATH": "/my path", "UV_CACHE_DIR": "default uv cache dir"},
+            stdout=PIPE,
+            stderr=STDOUT,
+            text=True,
+        )
+        assert Popen_mock.call_args == expected_args
+
     def test_custom_pythonpath(self, monkeypatch: MonkeyPatch) -> None:
         Popen_mock, get_global_config_dict_mock = self._setup(monkeypatch)
         monkeypatch.setattr(nemo_gym.cli_setup_command, "environ", {"PYTHONPATH": "existing pythonpath"})