feat(cli): judge / auxiliary-model resources-server scaffold template (M6g)

.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
@@ -132,21 +134,28 @@ 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
+          # 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.
+          curl -LsSf https://astral.sh/uv/0.11.19/install.sh | 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'
+        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/benchmarks.py

@@ -14,26 +14,19 @@
 # limitations under the License.
 """Benchmark discovery and preparation utilities."""
 
-import importlib
-from glob import glob
-from multiprocessing import Pool
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional
 
-import rich
 from omegaconf import DictConfig, OmegaConf
-from pydantic import BaseModel, Field
-from rich.table import Table
-from tqdm.auto import tqdm
+from pydantic import BaseModel
 
 from nemo_gym import PARENT_DIR
-from nemo_gym.config_types import BaseNeMoGymCLIConfig, BenchmarkDatasetConfig
+from nemo_gym.config_types import BenchmarkDatasetConfig
 from nemo_gym.global_config import (
     POLICY_MODEL_KEY_NAME,
     GlobalConfigDictParser,
     GlobalConfigDictParserConfig,
     get_first_server_config_dict,
-    get_global_config_dict,
 )
 
 
@@ -105,179 +98,3 @@ def _load_benchmarks_from_config_paths(config_paths: List[Path]) -> Dict[str, Be
         benchmarks_dict[maybe_bc.name] = maybe_bc
 
     return benchmarks_dict
-
-
-def list_benchmarks() -> None:
-    """CLI command: list available benchmarks."""
-    global_config_dict = get_global_config_dict(
-        global_config_dict_parser_config=GlobalConfigDictParserConfig(
-            initial_global_config_dict=GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT,
-        )
-    )
-    BaseNeMoGymCLIConfig.model_validate(global_config_dict)
-
-    assert BENCHMARKS_DIR.exists(), "Missing benchmarks directory"
-
-    config_paths = glob("**/config.yaml", root_dir=BENCHMARKS_DIR, recursive=True)
-    config_paths = [BENCHMARKS_DIR / p for p in config_paths]
-    config_paths = sorted(config_paths)
-
-    benchmarks = _load_benchmarks_from_config_paths(config_paths)
-
-    if not benchmarks:
-        rich.print("[yellow]No benchmarks found.[/yellow]")
-        rich.print(f"Expected benchmarks directory: {BENCHMARKS_DIR}")
-        return
-
-    table = Table(title=f"Available benchmarks in NeMo Gym ({len(benchmarks)})")
-    table.add_column("Benchmark name")
-    table.add_column("Agent name")
-    table.add_column("Num repeats")
-
-    for name, bench in benchmarks.items():
-        table.add_row(name, bench.agent_name, str(bench.num_repeats))
-
-    rich.print(table)
-
-
-class PrepareBenchmarkConfig(BaseNeMoGymCLIConfig):
-    """
-    Prepare benchmark data by running the benchmark's prepare.py script.
-
-    The benchmark is identified from a config_paths entry pointing to a
-    benchmarks/*/config.yaml file.
-
-    Examples:
-
-    ```bash
-    ng_prepare_benchmark "+config_paths=[benchmarks/aime24/config.yaml]"
-    ```
-    """
-
-    use_cached_prepared_benchmarks: bool = Field(
-        default=False, description="Skip benchmark preparation if the prepared file is already present"
-    )
-    num_prepare_benchmark_processes: int = Field(
-        default=1, description="Number of processes to parallelize benchmark preparation"
-    )
-
-
-def _multiprocess_benchmark_prepare_fn(args):
-    benchmark_config: BenchmarkConfig
-    prepare_module_path: str
-    (benchmark_config, prepare_module_path) = args
-
-    print(f"Preparing benchmark: {benchmark_config.name}")
-
-    module = importlib.import_module(prepare_module_path)
-    output_fpath = module.prepare()
-    assert output_fpath.absolute() == benchmark_config.dataset.jsonl_fpath.absolute(), (
-        f"Expected the actual prepared dataset output fpath to match the jsonl_fpath set in the config. Instead got {output_fpath=} jsonl_fpath={benchmark_config.dataset.jsonl_fpath}"
-    )
-    print(f"Benchmark data prepared at: {output_fpath}")
-
-
-def prepare_benchmark() -> None:
-    """CLI command: prepare benchmark data."""
-    global_config_dict = get_global_config_dict(
-        global_config_dict_parser_config=GlobalConfigDictParserConfig(
-            initial_global_config_dict=GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT,
-        )
-    )
-    prepare_benchmark_config = PrepareBenchmarkConfig.model_validate(global_config_dict)
-
-    benchmarks_dict: Dict[str, BenchmarkConfig] = dict()
-    for server_instance_name in global_config_dict:
-        server_config = global_config_dict[server_instance_name]
-        if not isinstance(server_config, (dict, DictConfig)) or "responses_api_agents" not in server_config:
-            continue
-
-        inner_server_config = get_first_server_config_dict(global_config_dict, server_instance_name)
-
-        datasets: List[BenchmarkDatasetConfig] = []
-        for dataset in inner_server_config.get("datasets") or []:
-            if dataset["type"] != "benchmark":
-                continue
-
-            datasets.append(BenchmarkDatasetConfig.model_validate(dataset))
-
-        if len(datasets) < 1:
-            continue
-
-        assert len(datasets) == 1, (
-            f"Expected 1 benchmark dataset for `{server_instance_name}`, but found {len(datasets)}!"
-        )
-
-        dataset = datasets[0]
-
-        benchmarks_dict[server_instance_name] = BenchmarkConfig(
-            name=dataset.name,
-            path=Path(""),
-            agent_name=server_instance_name,
-            num_repeats=dataset.num_repeats,
-            dataset=dataset,
-        )
-
-    assert benchmarks_dict, (
-        'No benchmark config found in config_paths. Pass a benchmark config, e.g.: "+config_paths=[benchmarks/aime24/config.yaml]"'
-    )
-
-    # Validate all benchmarks before preparing any
-    prepare_script_missing: List[BenchmarkConfig] = []
-    prepare_function_missing: List[BenchmarkConfig] = []
-
-    validated: List[Tuple[BenchmarkConfig, str]] = []
-    already_prepared: List[BenchmarkConfig] = []
-    for benchmark_config in benchmarks_dict.values():
-        prepare_script_path = benchmark_config.dataset.prepare_script
-        if not prepare_script_path.exists():
-            prepare_script_missing.append(benchmark_config)
-            continue
-
-        prepare_module_path = ".".join(prepare_script_path.with_suffix("").parts)
-        module = importlib.import_module(prepare_module_path)
-        if not hasattr(module, "prepare"):
-            prepare_function_missing.append(benchmark_config)
-            continue
-
-        is_already_prepared = benchmark_config.dataset.jsonl_fpath.exists()
-        if prepare_benchmark_config.use_cached_prepared_benchmarks and is_already_prepared:
-            already_prepared.append(benchmark_config)
-            continue
-
-        validated.append((benchmark_config, prepare_module_path))
-
-    if already_prepared:
-        already_prepared_str = "".join(f"- {bc.name}: {bc.dataset.jsonl_fpath}\n" for bc in already_prepared)
-        already_prepared_str = f"""The following benchmarks have already been prepared. Since `use_cached_prepared_benchmarks=true`, we will skip re-preparation of those benchmarks.
-        {already_prepared_str}"""
-        print(already_prepared_str)
-
-    errors_to_print = ""
-    if prepare_script_missing:
-        prepare_script_missing_str = "".join(
-            f"- {bc.name}: {bc.dataset.prepare_script}\n" for bc in prepare_script_missing
-        )
-        errors_to_print += f"""The following benchmarks are missing a valid prepare script:
-{prepare_script_missing_str}
-"""
-    if prepare_function_missing:  # pragma: no cover
-        prepare_function_missing_str = "".join(
-            f"- {bc.name}: {bc.dataset.prepare_script}\n" for bc in prepare_function_missing
-        )
-        errors_to_print += f"""The following benchmarks have a prepare script, but are missing the prepare function:
-{prepare_function_missing_str}
-"""
-    if errors_to_print:
-        errors_to_print = f"""Did not prepare any benchmarks due to benchmark config errors.
-{errors_to_print}"""
-        raise RuntimeError(errors_to_print)
-
-    # Prepare after all validations pass
-    if prepare_benchmark_config.num_prepare_benchmark_processes > 1:  # pragma: no cover
-        with Pool(processes=prepare_benchmark_config.num_prepare_benchmark_processes) as pool:
-            results = pool.imap_unordered(_multiprocess_benchmark_prepare_fn, validated)
-            list(tqdm(results, total=len(validated)))
-    else:
-        results = map(_multiprocess_benchmark_prepare_fn, validated)
-        list(tqdm(results, total=len(validated)))

nemo_gym/cli/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo_gym/cli/dataset.py

@@ -0,0 +1,96 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo_gym.config_types import (
+    DeleteJsonlDatasetGitlabConfig,
+    DownloadJsonlDatasetGitlabConfig,
+    DownloadJsonlDatasetHuggingFaceConfig,
+    UploadJsonlDatasetGitlabConfig,
+    UploadJsonlDatasetHuggingFaceConfig,
+    UploadJsonlDatasetHuggingFaceMaybeDeleteConfig,
+)
+from nemo_gym.dataset_orchestrator import (
+    delete_jsonl_dataset_from_gitlab,
+    upload_jsonl_dataset_to_hf_maybe_delete,
+)
+from nemo_gym.gitlab_utils import download_jsonl_dataset, upload_jsonl_dataset
+from nemo_gym.global_config import GlobalConfigDictParserConfig, get_global_config_dict
+from nemo_gym.hf_utils import download_hf_dataset_as_jsonl
+from nemo_gym.prompt import MaterializePromptsConfig, materialize_prompts
+from nemo_gym.train_data_utils import TrainDataProcessor
+
+
+def upload_jsonl_dataset_cli() -> None:  # pragma: no cover
+    global_config = get_global_config_dict()
+    config = UploadJsonlDatasetGitlabConfig.model_validate(global_config)
+    upload_jsonl_dataset(config)
+
+
+def download_jsonl_dataset_cli() -> None:  # pragma: no cover
+    global_config = get_global_config_dict()
+    config = DownloadJsonlDatasetGitlabConfig.model_validate(global_config)
+    download_jsonl_dataset(config)
+
+
+def upload_jsonl_dataset_to_hf_cli() -> None:  # pragma: no cover
+    global_config = get_global_config_dict()
+    config = UploadJsonlDatasetHuggingFaceMaybeDeleteConfig.model_validate(global_config)
+    upload_jsonl_dataset_to_hf_maybe_delete(config, delete_from_gitlab=config.delete_from_gitlab)
+
+
+def download_jsonl_dataset_from_hf_cli() -> None:  # pragma: no cover
+    global_config = get_global_config_dict()
+    config = DownloadJsonlDatasetHuggingFaceConfig.model_validate(global_config)
+
+    if config.artifact_fpath:
+        print(f"Downloading file '{config.artifact_fpath}' from '{config.repo_id}'...")
+    else:
+        print(f"Downloading '{config.split or 'all'}' split(s) from '{config.repo_id}'...")
+
+    download_hf_dataset_as_jsonl(config)
+
+
+def delete_jsonl_dataset_from_gitlab_cli() -> None:  # pragma: no cover
+    global_config = get_global_config_dict()
+    config = DeleteJsonlDatasetGitlabConfig.model_validate(global_config)
+    delete_jsonl_dataset_from_gitlab(config.dataset_name)
+
+
+def upload_jsonl_dataset_to_hf_and_delete_gitlab_cli() -> None:  # pragma: no cover
+    global_config = get_global_config_dict()
+    config = UploadJsonlDatasetHuggingFaceConfig.model_validate(global_config)
+    upload_jsonl_dataset_to_hf_maybe_delete(config, delete_from_gitlab=True)
+
+
+def materialize_prompts_cli() -> None:  # pragma: no cover
+    """CLI entry point for ng_materialize_prompts."""
+    global_config_dict = get_global_config_dict(
+        global_config_dict_parser_config=GlobalConfigDictParserConfig(
+            initial_global_config_dict=GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT,
+        )
+    )
+    config = MaterializePromptsConfig.model_validate(global_config_dict)
+    materialize_prompts(config.input_jsonl_fpath, config.prompt_config, config.output_jsonl_fpath)
+
+
+def prepare_data():  # pragma: no cover
+    global_config_dict = get_global_config_dict(
+        global_config_dict_parser_config=GlobalConfigDictParserConfig(
+            initial_global_config_dict=GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT,
+        )
+    )
+
+    data_processor = TrainDataProcessor()
+    data_processor.run(global_config_dict)