diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index 85124757a4..fc3c5a1aa5 100644 --- a/.github/workflows/unit-tests.yml +++ b/.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,13 +134,20 @@ 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 @@ -146,7 +155,7 @@ jobs: 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 diff --git a/nemo_gym/benchmarks.py b/nemo_gym/benchmarks.py index 9ba1e9fafd..7d57c6a2f6 100644 --- a/nemo_gym/benchmarks.py +++ b/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))) diff --git a/nemo_gym/cli/__init__.py b/nemo_gym/cli/__init__.py new file mode 100644 index 0000000000..3159bfe656 --- /dev/null +++ b/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. diff --git a/nemo_gym/cli/dataset.py b/nemo_gym/cli/dataset.py new file mode 100644 index 0000000000..eb3081ed4c --- /dev/null +++ b/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) diff --git a/nemo_gym/cli/dev.py b/nemo_gym/cli/dev.py new file mode 100644 index 0000000000..2134d45d7a --- /dev/null +++ b/nemo_gym/cli/dev.py @@ -0,0 +1,37 @@ +# 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 subprocess import Popen + +from nemo_gym.config_types import BaseNeMoGymCLIConfig +from nemo_gym.global_config import get_global_config_dict + + +def dev_test(): # pragma: no cover + """ + Run core NeMo Gym tests with coverage reporting (runs pytest with --cov flag). + + Examples: + + ```bash + ng_dev_test + ``` + """ + global_config_dict = get_global_config_dict() + # Just here for help + BaseNeMoGymCLIConfig.model_validate(global_config_dict) + + proc = Popen("pytest --cov=. --durations=10", shell=True) + exit(proc.wait()) diff --git a/nemo_gym/cli.py b/nemo_gym/cli/env.py similarity index 79% rename from nemo_gym/cli.py rename to nemo_gym/cli/env.py index e3df40025a..7953a2afdc 100644 --- a/nemo_gym/cli.py +++ b/nemo_gym/cli/env.py @@ -12,16 +12,14 @@ # 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. + import asyncio import json import os -import platform import shlex -import sys -from copy import deepcopy +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass from glob import glob -from importlib.metadata import entry_points -from importlib.metadata import version as md_version from os import makedirs from os.path import exists from pathlib import Path @@ -32,27 +30,26 @@ from time import sleep, time from typing import Dict, List, Optional, Tuple -import psutil import rich import uvicorn from devtools import pprint -from omegaconf import DictConfig, OmegaConf, open_dict +from omegaconf import DictConfig, OmegaConf from pydantic import Field from rich.table import Table from tqdm.auto import tqdm -from nemo_gym import PARENT_DIR, ROOT_DIR, __version__ +from nemo_gym import PARENT_DIR, ROOT_DIR from nemo_gym.cli_setup_command import run_command, setup_env_command from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( DRY_RUN_KEY_NAME, + JSON_OUTPUT_KEY_NAME, NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME, NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME, NEMO_GYM_RESERVED_TOP_LEVEL_KEYS, GlobalConfigDictParserConfig, get_global_config_dict, ) -from nemo_gym.rollout_collection import E2ERolloutCollectionConfig, RolloutCollectionConfig, RolloutCollectionHelper from nemo_gym.server_status import StatusCommand from nemo_gym.server_utils import ( HEAD_SERVER_KEY_NAME, @@ -62,7 +59,6 @@ ServerStatus, initialize_ray, ) -from nemo_gym.train_data_utils import TrainDataProcessor # Grace period after SIGINT before escalating to SIGKILL. Kept short so Ctrl-C is responsive. @@ -428,69 +424,6 @@ def run( rh.run_forever() -def e2e_rollout_collection(): # pragma: no cover - global_config_dict = get_global_config_dict() - - # Ensure we have the right config first thing - e2e_rollout_collection_config = E2ERolloutCollectionConfig.model_validate(global_config_dict) - - # Prepare data - data_processor_config_dict = deepcopy(global_config_dict) - with open_dict(data_processor_config_dict): - data_processor_config_dict["should_download"] = True - data_processor_config_dict["mode"] = "train_preparation" - - output_fpath = Path(e2e_rollout_collection_config.output_jsonl_fpath) - data_process_output_dir = output_fpath.parent / "preprocessed_datasets" - data_processor_config_dict["output_dirpath"] = str(data_process_output_dir) - - input_jsonl_fpath = data_process_output_dir / f"{e2e_rollout_collection_config.split}.jsonl" - should_skip_data_processing = ( - e2e_rollout_collection_config.reuse_existing_data_preparation and input_jsonl_fpath.exists() - ) - if not should_skip_data_processing: - if e2e_rollout_collection_config.reuse_existing_data_preparation: - print( - f"Even though the `reuse_existing_data_preparation=true` flag was set, we will still do data preparation since the final input jsonl fpath `{input_jsonl_fpath}` does not exist yet" - ) - - data_processor = TrainDataProcessor() - data_processor.run(data_processor_config_dict) - else: - print( - f"Skipping data preparation since `reuse_existing_data_preparation=true` and the final input jsonl fpath `{input_jsonl_fpath}` already exists" - ) - - # Convert to RolloutCollectionConfig - rollout_collection_config_dict = deepcopy(global_config_dict) - with open_dict(rollout_collection_config_dict): - assert input_jsonl_fpath.exists(), input_jsonl_fpath - rollout_collection_config_dict["input_jsonl_fpath"] = str(input_jsonl_fpath) - - rollout_collection_config = RolloutCollectionConfig.model_validate( - OmegaConf.to_container(rollout_collection_config_dict) - ) - - rh = RunHelper() - rh.start(None) - - rch = RolloutCollectionHelper() - - print( - f"""Output artifacts: -1. Preprocessed datasets: {data_processor_config_dict["output_dirpath"]} -2. Dataset file used for rollout collection: {rollout_collection_config_dict["input_jsonl_fpath"]} -3. Rollout collection results file: {output_fpath} -""" - ) - try: - asyncio.run(rch.run_from_config(rollout_collection_config)) - except KeyboardInterrupt: - pass - finally: - rh.shutdown() - - def _validate_data_single(test_config: TestConfig) -> None: # pragma: no cover if not test_config.should_validate_data: return @@ -570,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 @@ -618,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 @@ -635,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: @@ -659,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)") @@ -721,32 +721,19 @@ def test_all(): # pragma: no cover exit(1) -def dev_test(): # pragma: no cover - """ - Run core NeMo Gym tests with coverage reporting (runs pytest with --cov flag). - - Examples: - - ```bash - ng_dev_test - ``` - """ - global_config_dict = get_global_config_dict() - # Just here for help - BaseNeMoGymCLIConfig.model_validate(global_config_dict) - - proc = Popen("pytest --cov=. --durations=10", shell=True) - exit(proc.wait()) - - def init_resources_server(): # pragma: no cover """ Initialize a new resources server with template files and directory structure. + Pass `+template=judge` to scaffold an LLM-as-judge / auxiliary-model verifier (a server that + calls another model — judge, reward model, or subagent — from `verify()`) instead of the + default basic verifier. + Examples: ```bash ng_init_resources_server +entrypoint=resources_servers/my_server + ng_init_resources_server +entrypoint=resources_servers/my_judge +template=judge ``` """ config_dict = get_global_config_dict() @@ -767,38 +754,76 @@ def init_resources_server(): # pragma: no cover configs_dirpath = dirpath / "configs" makedirs(configs_dirpath) + template = str(config_dict.get("template", "basic")) + if template not in ("basic", "judge"): + print(f"Unknown template '{template}'. Choose one of: basic, judge.") + exit() + + # The resources-server config block, and which app/test templates to scaffold, vary by template. + resources_server_body = """ # Module (relative to this server dir) that defines the FastAPI app. + entrypoint: app.py + # Task category, used by `gym list environments`. See the Domain enum for valid values. + domain: other""" + app_template_fname = "resources_server_template.py" + test_template_fname = "resources_server_test_template.py" + if template == "judge": + # LLM-as-judge / auxiliary-model verifier: the resources server calls another model + # (judge, reward model, or subagent) from verify(). See the judge app template. + resources_server_body += """ + # The auxiliary model (judge / reward model / subagent) this verifier calls. Wire `name` + # to a model server you pass at run time, like `policy_model` is for the agent. + judge_model_server: + type: responses_api_models + name: judge_model + # Base Responses API params for the judge; `input` is filled in per task by the server. + judge_responses_create_params: + model: judge_model + input: [] + max_output_tokens: 1024""" + app_template_fname = "judge_resources_server_template.py" + test_template_fname = "judge_resources_server_test_template.py" + config_fpath = configs_dirpath / f"{server_type_name}.yaml" with open(config_fpath, "w") as f: - f.write(f"""{server_type_name}_resources_server: + f.write(f"""# Resources server: implements verification and any task-specific tools/state. +{server_type_name}_resources_server: {server_type}: {server_type_name}: - entrypoint: app.py - domain: other +{resources_server_body} +# Agent server: drives the model and talks to the resources server above. {server_type_name}_simple_agent: responses_api_agents: simple_agent: entrypoint: app.py + # The resources server this agent uses for tools + verification. resources_server: type: resources_servers name: {server_type_name}_resources_server + # The model the agent drives. `policy_model` is a magic name resolved at run time + # from your --model-name/--model-url flags (or a model-server config). model_server: type: responses_api_models name: policy_model + # One entry per dataset split. `source:` declares where the JSONL is fetched from; + # `type` selects the backend (gitlab | huggingface). Omit `source` for a purely local file. datasets: - name: train type: train jsonl_fpath: resources_servers/{server_type_name}/data/train.jsonl num_repeats: 1 - gitlab_identifier: + source: + type: gitlab dataset_name: {server_type_name} version: 0.0.1 artifact_fpath: train.jsonl + # A license is required for train/validation splits. license: Apache 2.0 - name: validation type: validation jsonl_fpath: resources_servers/{server_type_name}/data/validation.jsonl num_repeats: 1 - gitlab_identifier: + source: + type: gitlab dataset_name: {server_type_name} version: 0.0.1 artifact_fpath: validation.jsonl @@ -810,7 +835,7 @@ def init_resources_server(): # pragma: no cover """) app_fpath = dirpath / "app.py" - with open(ROOT_DIR / "resources/resources_server_template.py") as f: + with open(ROOT_DIR / f"resources/{app_template_fname}") as f: app_template = f.read() app_content = app_template.replace("ExampleMultiStep", server_type_title) with open(app_fpath, "w") as f: @@ -820,7 +845,7 @@ def init_resources_server(): # pragma: no cover makedirs(tests_dirpath) tests_fpath = tests_dirpath / "test_app.py" - with open(ROOT_DIR / "resources/resources_server_test_template.py") as f: + with open(ROOT_DIR / f"resources/{test_template_fname}") as f: tests_template = f.read() tests_content = tests_template.replace("ExampleMultiStep", server_type_title) tests_content = tests_content.replace("from app", f"from resources_servers.{server_type_name}.app") @@ -887,39 +912,17 @@ def dump_config(): # pragma: no cover print(OmegaConf.to_yaml(global_config_dict, resolve=True)) -def display_help(): - """ - Display a list of available NeMo Gym CLI commands. - - Examples: - - ```bash - ng_help - ``` - """ - global_config_dict = get_global_config_dict() - # Just here for help - BaseNeMoGymCLIConfig.model_validate(global_config_dict) - - eps = entry_points().select(group="console_scripts") - project_scripts = {ep.name: ep.value for ep in eps if ep.name.startswith(("nemo_gym_", "ng_"))} - rich.print("""Run a command with `+h=true` or `+help=true` to see more detailed information! - -[bold]Available CLI scripts[/bold] ------------------""") - for script in project_scripts: - if not script.startswith("ng_"): - continue - - print(script) - - def status(): # pragma: no cover global_config_dict = get_global_config_dict() BaseNeMoGymCLIConfig.model_validate(global_config_dict) status_cmd = StatusCommand() servers = status_cmd.discover_servers() + + if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): + print(json.dumps([server.model_dump(mode="json") for server in servers])) + return + status_cmd.display_status(servers) @@ -965,93 +968,3 @@ def pip_list(): # pragma: no cover proc = run_command(command, dir_path) return_code = proc.wait() exit(return_code) - - -class VersionConfig(BaseNeMoGymCLIConfig): - """ - Display gym version and system information. - - Examples: - - ```bash - # Display version information - ng_version - - # Output as JSON - ng_version +json=true - ``` - """ - - json_format: bool = Field(default=False, alias="json", description="Output in JSON format for programmatic use.") - - -def version(): # pragma: no cover - """Display gym version and system information.""" - global_config_dict = get_global_config_dict() - config = VersionConfig.model_validate(global_config_dict) - - json_output = config.json_format - - version_info = { - "nemo_gym": __version__, - "python": platform.python_version(), - "python_path": sys.executable, - "installation_path": str(PARENT_DIR), - } - - key_deps = [ - "openai", - "ray", - ] - - dependencies = {dep: md_version(dep) for dep in key_deps} - - version_info["dependencies"] = dependencies - - # System info - version_info["system"] = { - "os": f"{platform.system()} {platform.release()}", - "platform": platform.platform(), - "architecture": platform.machine(), - "processor": platform.processor() or "unknown", - "cpus": os.cpu_count(), - } - - # Memory info - mem = psutil.virtual_memory() - version_info["system"]["memory_gb"] = round(mem.total / (1024**3), 2) - - # Output - if json_output: - print(json.dumps(version_info)) - else: - output = f"""\ -NeMo Gym v{version_info["nemo_gym"]} -Python {version_info["python"]} ({version_info["python_path"]}) -Installation: {version_info["installation_path"]}""" - - if "dependencies" in version_info: - deps_lines = "\n".join(f" {dep}: {ver}" for dep, ver in version_info["dependencies"].items()) - sys_info = version_info["system"] - output += f""" - -Key Dependencies: -{deps_lines} - -System: - OS: {sys_info["os"]} - Platform: {sys_info["platform"]} - Architecture: {sys_info["architecture"]} - Processor: {sys_info["processor"]} - CPUs: {sys_info["cpus"]} - Memory: {sys_info["memory_gb"]} GB""" - - print(output) - - -def reinstall(): # pragma: no cover - global_config_dict = get_global_config_dict() - # Just here for help - BaseNeMoGymCLIConfig.model_validate(global_config_dict) - - Popen("uv sync --extra dev --group docs", shell=True).communicate() diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py new file mode 100644 index 0000000000..9a6ce87e09 --- /dev/null +++ b/nemo_gym/cli/eval.py @@ -0,0 +1,428 @@ +# 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. + +import asyncio +import difflib +import importlib +import json +from copy import deepcopy +from glob import glob +from multiprocessing import Pool +from pathlib import Path +from typing import Dict, List, Tuple + +import orjson +import rich +from omegaconf import DictConfig, OmegaConf, open_dict +from pydantic import Field +from rich.table import Table +from tqdm.auto import tqdm + +from nemo_gym.benchmarks import BENCHMARKS_DIR, BenchmarkConfig, _load_benchmarks_from_config_paths +from nemo_gym.cli.env import RunHelper +from nemo_gym.config_types import BaseNeMoGymCLIConfig, BenchmarkDatasetConfig +from nemo_gym.global_config import ( + JSON_OUTPUT_KEY_NAME, + POLICY_MODEL_KEY_NAME, + ROLLOUT_INDEX_KEY_NAME, + TASK_INDEX_KEY_NAME, + GlobalConfigDictParser, + GlobalConfigDictParserConfig, + get_first_server_config_dict, + get_global_config_dict, +) +from nemo_gym.reward_profile import RewardProfileConfig, RewardProfiler +from nemo_gym.rollout_collection import ( + E2ERolloutCollectionConfig, + RolloutAggregationConfig, + RolloutAggregationHelper, + RolloutCollectionConfig, + RolloutCollectionHelper, +) +from nemo_gym.train_data_utils import TrainDataProcessor + + +def _fuzzy_matches(query: str, *fields: str) -> bool: + """Whether `query` fuzzily matches any of `fields`: a substring or a close difflib match (token-aware).""" + needle = query.lower() + for field in fields: + if not field: + continue + haystack = field.lower() + if needle in haystack: + return True + tokens = haystack.replace("_", " ").replace("-", " ").split() + if difflib.get_close_matches(needle, [haystack, *tokens], n=1, cutoff=0.6): + return True + return False + + +def _benchmark_extras(bench: BenchmarkConfig) -> tuple[str, list[str]]: + """Resolve a benchmark's config to its `(domain, extra search terms)`. + + `BenchmarkConfig` flattens away the resource server name, the resource server `domain`, and the + dataset names. We re-resolve the config with the same parser `BenchmarkConfig` uses (so chained + `config_paths` / `_inherit_from` are applied) and read those fields back out for the domain column + and richer `gym search` matching. + """ + initial_config_dict = OmegaConf.load(bench.path) + if POLICY_MODEL_KEY_NAME not in initial_config_dict: + initial_config_dict = OmegaConf.merge( + initial_config_dict, GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT + ) + resolved = GlobalConfigDictParser().parse_no_environment(initial_global_config_dict=initial_config_dict) + + domain = "" + terms: list[str] = [] + for instance_name in resolved: + instance = resolved[instance_name] + if not isinstance(instance, (dict, DictConfig)): + continue + + resource_servers = instance.get("resources_servers") + if resource_servers: + terms.append(instance_name) # e.g. aime24_math_with_judge_resources_server + for rs_name, rs_config in resource_servers.items(): + terms.append(rs_name) # e.g. math_with_judge + found_domain = (rs_config or {}).get("domain") + if found_domain: + domain = str(found_domain) + + agents = instance.get("responses_api_agents") + if agents: + for agent_config in agents.values(): + for dataset in (agent_config or {}).get("datasets") or []: + if (dataset or {}).get("name"): + terms.append(dataset["name"]) + + if domain: + terms.append(domain) + return domain, terms + + +def list_benchmarks() -> None: + """CLI command: list available benchmarks, optionally filtered by a `query` (the `gym search` entry point).""" + 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) + + # Resolve the domain + richer search terms once per benchmark, for the domain column and `gym search`. + extras = {name: _benchmark_extras(bench) for name, bench in benchmarks.items()} + + # `gym search ` reuses this command, narrowing the listing to fuzzy matches across the + # benchmark name, agent name, resource server name, dataset names, and domain. + query = global_config_dict.get("query") + if query: + benchmarks = { + name: bench + for name, bench in benchmarks.items() + if _fuzzy_matches(query, name, bench.agent_name, *extras[name][1]) + } + + if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): + payload = [ + { + "name": name, + "agent_name": bench.agent_name, + "domain": extras[name][0], + "num_repeats": bench.num_repeats, + } + for name, bench in benchmarks.items() + ] + print(json.dumps(payload)) + return + + if not benchmarks: + if query: + rich.print(f"[yellow]No benchmarks match '{query}'.[/yellow]") + return + rich.print("[yellow]No benchmarks found.[/yellow]") + rich.print(f"Expected benchmarks directory: {BENCHMARKS_DIR}") + return + + title = ( + f"Benchmarks matching '{query}' ({len(benchmarks)})" + if query + else f"Available benchmarks in NeMo Gym ({len(benchmarks)})" + ) + table = Table(title=title) + table.add_column("Benchmark name") + table.add_column("Domain") + table.add_column("Agent name") + table.add_column("Num repeats") + + for name, bench in benchmarks.items(): + table.add_row(name, extras[name][0], 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))) + + +def e2e_rollout_collection(): # pragma: no cover + global_config_dict = get_global_config_dict() + + # Ensure we have the right config first thing + e2e_rollout_collection_config = E2ERolloutCollectionConfig.model_validate(global_config_dict) + + # Prepare data + data_processor_config_dict = deepcopy(global_config_dict) + with open_dict(data_processor_config_dict): + data_processor_config_dict["should_download"] = True + data_processor_config_dict["mode"] = "train_preparation" + + output_fpath = Path(e2e_rollout_collection_config.output_jsonl_fpath) + data_process_output_dir = output_fpath.parent / "preprocessed_datasets" + data_processor_config_dict["output_dirpath"] = str(data_process_output_dir) + + input_jsonl_fpath = data_process_output_dir / f"{e2e_rollout_collection_config.split}.jsonl" + should_skip_data_processing = ( + e2e_rollout_collection_config.reuse_existing_data_preparation and input_jsonl_fpath.exists() + ) + if not should_skip_data_processing: + if e2e_rollout_collection_config.reuse_existing_data_preparation: + print( + f"Even though the `reuse_existing_data_preparation=true` flag was set, we will still do data preparation since the final input jsonl fpath `{input_jsonl_fpath}` does not exist yet" + ) + + data_processor = TrainDataProcessor() + data_processor.run(data_processor_config_dict) + else: + print( + f"Skipping data preparation since `reuse_existing_data_preparation=true` and the final input jsonl fpath `{input_jsonl_fpath}` already exists" + ) + + # Convert to RolloutCollectionConfig + rollout_collection_config_dict = deepcopy(global_config_dict) + with open_dict(rollout_collection_config_dict): + assert input_jsonl_fpath.exists(), input_jsonl_fpath + rollout_collection_config_dict["input_jsonl_fpath"] = str(input_jsonl_fpath) + + rollout_collection_config = RolloutCollectionConfig.model_validate( + OmegaConf.to_container(rollout_collection_config_dict) + ) + + rh = RunHelper() + rh.start(None) + + rch = RolloutCollectionHelper() + + print( + f"""Output artifacts: +1. Preprocessed datasets: {data_processor_config_dict["output_dirpath"]} +2. Dataset file used for rollout collection: {rollout_collection_config_dict["input_jsonl_fpath"]} +3. Rollout collection results file: {output_fpath} +""" + ) + try: + asyncio.run(rch.run_from_config(rollout_collection_config)) + except KeyboardInterrupt: + pass + finally: + rh.shutdown() + + +def collect_rollouts(): # pragma: no cover + config = RolloutCollectionConfig.model_validate(get_global_config_dict()) + rch = RolloutCollectionHelper() + + asyncio.run(rch.run_from_config(config)) + + +def aggregate_rollouts(): # pragma: no cover + config = RolloutAggregationConfig.model_validate(get_global_config_dict()) + rah = RolloutAggregationHelper() + + asyncio.run(rah.run_from_config(config)) + + +def reward_profile(): # pragma: no cover + config = RewardProfileConfig.model_validate(get_global_config_dict()) + + with open(config.materialized_inputs_jsonl_fpath) as f: + rows = list(map(orjson.loads, f)) + + with open(config.rollouts_jsonl_fpath) as f: + results = list(map(orjson.loads, f)) + + # Results may be out of order. + results.sort(key=lambda r: (r[TASK_INDEX_KEY_NAME], r[ROLLOUT_INDEX_KEY_NAME])) + + rp = RewardProfiler() + group_level_metrics, agent_level_metrics = rp.profile_from_data( + rows, results, allow_partial_rollouts=config.allow_partial_rollouts + ) + completion_summary = rp.profile_completion_summary(rows, results) + reward_profiling_fpath, agent_level_metrics_fpath = rp.write_to_disk( + group_level_metrics, agent_level_metrics, Path(config.rollouts_jsonl_fpath) + ) + + print(f"""Profiling outputs: +Reward profile completion: {completion_summary["completed_rollout_rows"]}/{completion_summary["expected_rollout_rows"]} rollout rows ({completion_summary["reward_profile_completion_pct"]:.2f}%) +Input rows: {completion_summary["total_input_rows"]} total; {completion_summary["complete_input_rows"]} complete; {completion_summary["partial_input_rows"]} partial; {completion_summary["missing_input_rows"]} without rollouts dropped from output. +Reward profiling outputs: {reward_profiling_fpath} +Agent-level metrics: {agent_level_metrics_fpath}""") diff --git a/nemo_gym/cli/general.py b/nemo_gym/cli/general.py new file mode 100644 index 0000000000..3c577e8e88 --- /dev/null +++ b/nemo_gym/cli/general.py @@ -0,0 +1,145 @@ +# 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. + +import json +import os +import platform +import sys +from importlib.metadata import entry_points +from importlib.metadata import version as md_version +from subprocess import Popen + +import psutil +import rich +from pydantic import Field + +from nemo_gym import PARENT_DIR, __version__ +from nemo_gym.config_types import BaseNeMoGymCLIConfig +from nemo_gym.global_config import JSON_OUTPUT_KEY_NAME, get_global_config_dict + + +def display_help_legacy(): + """ + Display a list of available NeMo Gym CLI commands. + + Examples: + + ```bash + ng_help + ``` + """ + global_config_dict = get_global_config_dict() + # Just here for help + BaseNeMoGymCLIConfig.model_validate(global_config_dict) + + eps = entry_points().select(group="console_scripts") + project_scripts = {ep.name: ep.value for ep in eps if ep.name.startswith(("nemo_gym_", "ng_"))} + rich.print("""Run a command with `+h=true` or `+help=true` to see more detailed information! + +[bold]Available CLI scripts[/bold] +-----------------""") + for script in project_scripts: + if not script.startswith("ng_"): + continue + + print(script) + + +class VersionConfig(BaseNeMoGymCLIConfig): + """ + Display gym version and system information. + + Examples: + + ```bash + # Display version information + ng_version + + # Output as JSON + ng_version +json=true + ``` + """ + + json_format: bool = Field(default=False, alias="json", description="Output in JSON format for programmatic use.") + + +def version(): # pragma: no cover + """Display gym version and system information.""" + global_config_dict = get_global_config_dict() + # Just here for help. + VersionConfig.model_validate(global_config_dict) + + version_info = { + "nemo_gym": __version__, + "python": platform.python_version(), + "python_path": sys.executable, + "installation_path": str(PARENT_DIR), + } + + key_deps = [ + "openai", + "ray", + ] + + dependencies = {dep: md_version(dep) for dep in key_deps} + + version_info["dependencies"] = dependencies + + # System info + version_info["system"] = { + "os": f"{platform.system()} {platform.release()}", + "platform": platform.platform(), + "architecture": platform.machine(), + "processor": platform.processor() or "unknown", + "cpus": os.cpu_count(), + } + + # Memory info + mem = psutil.virtual_memory() + version_info["system"]["memory_gb"] = round(mem.total / (1024**3), 2) + + if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): + print(json.dumps(version_info)) + else: + output = f"""\ +NeMo Gym v{version_info["nemo_gym"]} +Python {version_info["python"]} ({version_info["python_path"]}) +Installation: {version_info["installation_path"]}""" + + if "dependencies" in version_info: + deps_lines = "\n".join(f" {dep}: {ver}" for dep, ver in version_info["dependencies"].items()) + sys_info = version_info["system"] + output += f""" + +Key Dependencies: +{deps_lines} + +System: + OS: {sys_info["os"]} + Platform: {sys_info["platform"]} + Architecture: {sys_info["architecture"]} + Processor: {sys_info["processor"]} + CPUs: {sys_info["cpus"]} + Memory: {sys_info["memory_gb"]} GB""" + + print(output) + + +def reinstall(): # pragma: no cover + global_config_dict = get_global_config_dict() + # Just here for help + BaseNeMoGymCLIConfig.model_validate(global_config_dict) + + Popen("uv sync --extra dev --group docs", shell=True).communicate() diff --git a/nemo_gym/cli/legacy.py b/nemo_gym/cli/legacy.py new file mode 100644 index 0000000000..837f4b52ba --- /dev/null +++ b/nemo_gym/cli/legacy.py @@ -0,0 +1,79 @@ +# 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. +"""Backward-compatibility shim for the legacy ``ng_*`` / ``nemo_gym_*`` commands. + +Every legacy console script points here; the script name (``sys.argv[0]``) identifies which +command was invoked. We print a one-time deprecation notice mapping it to the new ``gym`` +command, then re-enter the ``gym`` router so the command keeps working (REQ 7). +""" + +import sys +from pathlib import Path + +from nemo_gym.cli.main import dispatch +from nemo_gym.cli.main import main as gym_main + + +# Legacy command (ng_/nemo_gym_ prefix stripped) -> equivalent `gym` subcommand tokens. +LEGACY = { + "run": ["env", "run"], + "test": ["env", "test"], + "test_all": ["env", "test"], + "dev_test": ["dev", "test"], + "init_resources_server": ["env", "init"], + "list_benchmarks": ["list", "benchmarks"], + "prepare_benchmark": ["eval", "prepare"], + "collect_rollouts": ["eval", "run", "--no-serve"], + "e2e_collect_rollouts": ["eval", "run"], + "aggregate_rollouts": ["eval", "aggregate"], + "materialize_prompts": ["dataset", "render"], + "reward_profile": ["eval", "profile"], + "upload_dataset_to_gitlab": ["dataset", "upload", "--storage", "gitlab"], + "download_dataset_from_gitlab": ["dataset", "download", "--storage", "gitlab"], + "prepare_data": ["dataset", "collate"], + "upload_dataset_to_hf": ["dataset", "upload"], + "download_dataset_from_hf": ["dataset", "download"], + "gitlab_to_hf_dataset": ["dataset", "migrate"], + "delete_dataset_from_gitlab": ["dataset", "rm"], + "dump_config": ["env", "resolve"], + "help": ["--help"], + "status": ["env", "status"], + "pip_list": ["env", "packages"], + "version": ["--version"], +} + + +def main() -> None: + alias = Path(sys.argv[0]).name + key = alias.removeprefix("nemo_gym_").removeprefix("ng_") + + # `reinstall` has no `gym` equivalent (`gym install` was dropped); point users at the uv command it runs. + if key == "reinstall": + print( + f"⚠ `{alias}` is deprecated and will be removed in a future release; " + f"run `uv sync --extra dev --group docs` instead.", + file=sys.stderr, + ) + dispatch("nemo_gym.cli.general:reinstall", sys.argv[1:]) + return + + tokens = LEGACY[key] + print( + f"⚠ `{alias}` is deprecated and will be removed in a future release; use `gym {' '.join(tokens)}` instead.", + file=sys.stderr, + ) + # Re-enter the gym router with the equivalent subcommand, preserving the user's Hydra overrides. + sys.argv = [sys.argv[0], *tokens, *sys.argv[1:]] + gym_main() diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py new file mode 100644 index 0000000000..0eace3b678 --- /dev/null +++ b/nemo_gym/cli/main.py @@ -0,0 +1,551 @@ +# 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. +import argparse +import difflib +import importlib +import re +import sys +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field +from pathlib import Path + +from nemo_gym import WORKING_DIR + + +VERSION_TARGET = "nemo_gym.cli.general:version" + + +def _did_you_mean(value: str, candidates: Iterable[str]) -> str: + """A ` Did you mean \\`X\\`?` fragment for the closest candidate to `value`, or `""` if none is close enough.""" + matches = difflib.get_close_matches(value, list(candidates), n=1) + return f" Did you mean `{matches[0]}`?" if matches else "" + + +class _GymArgumentParser(argparse.ArgumentParser): + """ArgumentParser that appends a difflib "did you mean?" hint to invalid-choice errors. + + Covers mistyped commands/groups and bad --flag choices (e.g. --storage), since argparse validates all of them + as choices against the registry baked into the parser. + """ + + def error(self, message: str) -> None: + match = re.search(r"invalid choice: '([^']+)' \(choose from (.+)\)", message) + if match: + typo, choices = match.group(1), re.findall(r"'([^']+)'", match.group(2)) + message += _did_you_mean(typo, choices) + super().error(message) + + +@dataclass(frozen=True) +class Flag: + # Register this flag's argument(s) on a command's subparser. + register: Callable[[argparse.ArgumentParser], None] + # Turn the parsed value into leading Hydra override tokens (default: contributes nothing). + translate_to_hydra: Callable[[argparse.Namespace], list[str]] = lambda args: [] + + +@dataclass(frozen=True) +class Command: + # What to run: either a "module:function" string (lazily imported and called with no args), + # or a callable(args, overrides) that owns dispatch (e.g. picks the target from parsed flags). + target: str | Callable[[argparse.Namespace, list[str]], None] + # One-line help shown in the parent listing and atop this command's own --help. + summary: str | None = None + # Flags this command accepts; reusable ones (e.g. CONFIG) are shared across commands. + flags: tuple[Flag, ...] = field(default_factory=tuple) + + +def dispatch(target: str, overrides: list[str]) -> None: + module_path, func_name = target.split(":") + # Drop the parsed command tokens so the downstream Hydra parsing only sees overrides. + sys.argv = [sys.argv[0], *overrides] + func = getattr(importlib.import_module(module_path), func_name) + func() + + +def _value_flag( + name: str, hydra_key: str, flag_help: str, *, aliases: tuple[str, ...] = (), choices: tuple[str, ...] | None = None +) -> Flag: + """A `--name VALUE` flag that maps to the Hydra override `+=VALUE` (omitted when unset).""" + dest = name.replace("-", "_") + return Flag( + register=lambda p: p.add_argument(f"--{name}", *aliases, dest=dest, choices=choices, help=flag_help), + translate_to_hydra=lambda args: ( + [f"+{hydra_key}={getattr(args, dest)}"] if getattr(args, dest) is not None else [] + ), + ) + + +def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag: + """A `--name` store_true flag that maps to the Hydra override `+=true` when set.""" + dest = name.replace("-", "_") + return Flag( + register=lambda p: p.add_argument(f"--{name}", action="store_true", help=flag_help), + translate_to_hydra=lambda args: [f"+{hydra_key}=true"] if getattr(args, dest) else [], + ) + + +# Shared flag: load Gym config files. Reused by every command that reads server/benchmark configs. +CONFIG = Flag( + register=lambda p: p.add_argument( + "--config", + action="append", + metavar="PATH", + help="Config file to load; repeatable. Maps to +config_paths=[...].", + ), + translate_to_hydra=lambda args: [f"+config_paths=[{','.join(args.config)}]"] if args.config else [], +) + +# Shared flag: select the storage backend. Reused by `dataset upload` and `dataset download`. +STORAGE = Flag( + register=lambda p: p.add_argument( + "--storage", choices=("hf", "gitlab"), default="hf", help="Storage backend (default: hf)." + ) +) + +# Shared model-server flags. Reused by commands that spin up / target a model server (`eval run`, `env run`). +# --model is the served model identifier across all backends: an API model name, an HF id, or a local checkpoint +# path, interpreted per --model-type (e.g. a path/HF id to serve with local_vllm_model). +MODEL = _value_flag( + "model", + "policy_model_name", + "Model name, HF id, or local checkpoint path (interpreted per --model-type).", + aliases=("-m",), +) +MODEL_URL = _value_flag("model-url", "policy_base_url", "Model server base URL.") +MODEL_API_KEY = _value_flag("model-api-key", "policy_api_key", "Model server API key.") + +# Shared flag: select a single resource server by name. Reused by `env test`, `env init`, and `env packages`. +RESOURCE_SERVER = Flag( + register=lambda p: p.add_argument("--resource-server", metavar="NAME", help="Name of the resource server."), + translate_to_hydra=lambda args: ( + [f"+entrypoint=resources_servers/{args.resource_server}"] if args.resource_server else [] + ), +) + +# Shared flag: emit machine-readable JSON instead of human output. Reused by reporting commands (version, list, +# env status). The reserved `json` config key is read centrally via nemo_gym.cli.output.emit. +JSON = _bool_flag("json", "json", "Output as JSON for programmatic use.") + +# Positional search query for `gym search`; surfaced to the listing command as the `query` config key. +QUERY = Flag( + register=lambda p: p.add_argument("query", metavar="QUERY", help="Substring to match against component names."), + translate_to_hydra=lambda args: [f"+query={args.query}"] if getattr(args, "query", None) else [], +) + + +# Asset selector flag -> (parent dir, configs subdir, default config flavor). All accept `name` or `name/flavor`, +# resolving to `//[/].yaml`. A None default flavor falls back to the server name. +_ASSETS = { + "benchmark": ("benchmarks", "", "config"), + "resource-server": ("resources_servers", "configs", None), + "model-type": ("responses_api_models", "configs", None), +} + + +def _asset_config_path(flag: str, value: str, search_dirs: tuple[str, ...] = ()) -> str: + """Map a named asset (`name` or `name/flavor`) to its config path. + + Searches WORKING_DIR (built-ins) first, then any user-registered --search-dir roots. + """ + parent, subdir, default_flavor = _ASSETS[flag] + server_name, _, config_flavor = value.partition("/") + config_flavor = config_flavor or default_flavor or server_name + config_dir = f"{parent}/{server_name}/{subdir}".rstrip("/") + path = f"{config_dir}/{config_flavor}.yaml" + + # Match in WORKING_DIR (built-ins) and every --search-dir root; dedupe roots that resolve to the same file. + roots = [WORKING_DIR, *(Path(d) for d in search_dirs)] + matches: list[Path] = [] + + for root in roots: + candidate = root / path + if candidate.exists(): + matches.append(candidate.resolve()) + + if len(matches) > 1: + matches_str = ", ".join(f"`{m}`" for m in matches) + raise ValueError( + f"`--{flag} {value}` is ambiguous: it matches multiple configs ({matches_str}). " + f"Pass the intended config directly with `--config ` instead." + ) + if matches: + return str(matches[0]) + + # No match: suggest the closest real name across all roots (a config flavor when the server exists, else a + # server name) and report the full paths that were searched. + available = ", ".join(set(f"`{(root / config_dir).resolve()}`" for root in roots if (root / config_dir).is_dir())) + typo = config_flavor + candidates = [p.stem for root in roots for p in (root / config_dir).glob("*.yaml")] + + if len(candidates) == 0: + available = ", ".join(set(f"`{(root / parent).resolve()}`" for root in roots if (root / parent).is_dir())) + typo = server_name + candidates = [ + child.name + for root in roots + if (root / parent).is_dir() + for child in (root / parent).iterdir() + if child.is_dir() + ] + + raise ValueError( + f"`--{flag} {value}` was specified which implies config `{path}`, which does not exist.{_did_you_mean(typo, candidates)} " + f"See available {flag} configs in {available}." + ) + + +def _asset_selector(flag: str) -> Flag: + """A `-- NAME` selector that resolves the named asset to a config and adds it to +config_paths.""" + dest = flag.replace("-", "_") + return Flag( + register=lambda p: p.add_argument(f"--{flag}", metavar="NAME", help=f"Load the named {flag} config."), + translate_to_hydra=lambda args: ( + [ + f"+config_paths=[{_asset_config_path(flag, getattr(args, dest), tuple(getattr(args, 'search_dir', None) or ()))}]" + ] + if getattr(args, dest) + else [] + ), + ) + + +BENCHMARK = _asset_selector("benchmark") +RESOURCE_SERVER_CONFIG = _asset_selector("resource-server") +MODEL_TYPE = _asset_selector("model-type") + +# Shared flag: register extra root dirs to search for named components. Consumed by the asset selectors above +# (not emitted as a Hydra override). Reused by every command that accepts a -- NAME selector. +SEARCH_DIR = Flag( + register=lambda p: p.add_argument( + "--search-dir", + action="append", + metavar="DIR", + help="Extra root directory to search for named components; repeatable.", + ), +) + + +def _merge_config_paths(overrides: list[str]) -> list[str]: + """Coalesce all `+config_paths=[...]` tokens (from --config and asset selectors) into one (Hydra rejects dupes).""" + prefix = "+config_paths=[" + paths: list[str] = [] + rest: list[str] = [] + for token in overrides: + if token.startswith(prefix) and token.endswith("]"): + paths.extend(p for p in token[len(prefix) : -1].split(",") if p) + else: + rest.append(token) + return ([f"+config_paths=[{','.join(paths)}]"] if paths else []) + rest + + +def _eval_run(args: argparse.Namespace, overrides: list[str]) -> None: + target = "nemo_gym.cli.eval:collect_rollouts" if args.no_serve else "nemo_gym.cli.eval:e2e_rollout_collection" + dispatch(target, overrides) + + +def _env_test(args: argparse.Namespace, overrides: list[str]) -> None: + # Run a single server's tests if +entrypoint was passed. No need to check for + # --resource-server because it is translated to +entrypoint in the flag definition. + + has_entrypoint = any(override.lstrip("+").split("=", 1)[0] == "entrypoint" for override in overrides) + dispatch("nemo_gym.cli.env:test" if has_entrypoint else "nemo_gym.cli.env:test_all", overrides) + + +def _dataset_upload(args: argparse.Namespace, overrides: list[str]) -> None: + targets = { + "hf": "nemo_gym.cli.dataset:upload_jsonl_dataset_to_hf_cli", + "gitlab": "nemo_gym.cli.dataset:upload_jsonl_dataset_cli", + } + dispatch(targets[args.storage], overrides) + + +def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: + targets = { + "hf": "nemo_gym.cli.dataset:download_jsonl_dataset_from_hf_cli", + "gitlab": "nemo_gym.cli.dataset:download_jsonl_dataset_cli", + } + dispatch(targets[args.storage], overrides) + + +# One-line help for each command group, shown in `gym --help`. +GROUPS = { + "list": "List available components. As of now, only benchmarks are available.", + "dataset": "Manage datasets.", + "env": "Develop and run environments.", + "eval": "Run evaluations.", + "dev": "Contributor helpers.", +} + +COMMANDS = { + "list benchmarks": Command( + target="nemo_gym.cli.eval:list_benchmarks", summary="List available benchmarks.", flags=(JSON,) + ), + "search": Command( + target="nemo_gym.cli.eval:list_benchmarks", + summary="Search available components (currently benchmarks) by name; like `list` filtered to a query.", + flags=(QUERY, JSON), + ), + "dataset upload": Command( + target=_dataset_upload, + summary="Upload a prepared dataset to HF (default) or GitLab.", + flags=( + STORAGE, + _value_flag("input", "input_jsonl_fpath", "Local JSONL file to upload.", aliases=("-i",)), + _value_flag("name", "dataset_name", "Dataset name."), + # GitLab stores it as `version`, HF as `revision`; emit both and let each backend keep its own. + Flag( + register=lambda p: p.add_argument( + "--revision", dest="revision", help="Dataset revision (version) to download." + ), + translate_to_hydra=lambda args: ( + # we set both version and revision because GitLab and HF use different keys + # and we extra="ignore" so it's safe to set both + [f"+version={args.revision}", f"+revision={args.revision}"] if args.revision is not None else [] + ), + ), + _value_flag("split", "split", "Dataset split (HF only)."), + _bool_flag("create-pr", "create_pr", "Open a pull request instead of committing directly (HF only)."), + ), + ), + "dataset download": Command( + target=_dataset_download, + summary="Download a dataset from HF (default) or GitLab.", + flags=( + STORAGE, + _value_flag("repo-id", "repo_id", "HF repo id, e.g. org/dataset (HF only)."), + _value_flag("name", "dataset_name", "Dataset name (GitLab only)."), + # NOTE(martas): HF download does not allow to specify revision + _value_flag("revision", "version", "Dataset version (GitLab only)."), + _value_flag( + "artifact", "artifact_fpath", "Remote file to fetch (GitLab: required; HF: optional raw file)." + ), + _value_flag("output", "output_fpath", "Local destination file.", aliases=("-o",)), + _value_flag( + "output-dir", "output_dirpath", "Local destination directory; needed for all splits (HF only)." + ), + _value_flag("split", "split", "Dataset split (HF only)."), + ), + ), + "dataset rm": Command( + target="nemo_gym.cli.dataset:delete_jsonl_dataset_from_gitlab_cli", + summary="Delete a dataset from GitLab.", + flags=(_value_flag("name", "dataset_name", "Name of the dataset to delete."),), + ), + "dataset migrate": Command( + target="nemo_gym.cli.dataset:upload_jsonl_dataset_to_hf_and_delete_gitlab_cli", + summary="Transfer a dataset from GitLab to HF.", + flags=( + _value_flag("input", "input_jsonl_fpath", "Local JSONL file to upload to HF.", aliases=("-i",)), + _value_flag("name", "dataset_name", "Dataset name."), + _value_flag("revision", "revision", "Dataset revision (HF)."), + _value_flag("split", "split", "Dataset split."), + _bool_flag("create-pr", "create_pr", "Open a pull request instead of committing directly."), + ), + ), + "dataset render": Command( + target="nemo_gym.cli.dataset:materialize_prompts_cli", + summary="Generate a dataset preview.", + flags=( + _value_flag("input", "input_jsonl_fpath", "Raw input JSONL file.", aliases=("-i",)), + _value_flag("prompt-config", "prompt_config", "Prompt template YAML to apply."), + _value_flag("output", "output_jsonl_fpath", "Output JSONL file.", aliases=("-o",)), + ), + ), + "dataset collate": Command( + target="nemo_gym.cli.dataset:prepare_data", + summary="Validate and collate the dataset.", + flags=( + CONFIG, + RESOURCE_SERVER_CONFIG, + SEARCH_DIR, + _value_flag("mode", "mode", "Data preparation mode.", choices=("train_preparation", "example_validation")), + _value_flag("output-dir", "output_dirpath", "Output directory for the prepared data."), + _bool_flag("download", "should_download", "Download source datasets before collating."), + ), + ), + "env init": Command( + target="nemo_gym.cli.env:init_resources_server", + summary="Scaffold config for a new server, benchmark, or agent.", + flags=(RESOURCE_SERVER,), + ), + "env resolve": Command( + target="nemo_gym.cli.env:dump_config", + summary="Resolve the final config from configs, flags, and overrides.", + flags=(CONFIG,), + ), + "env packages": Command( + target="nemo_gym.cli.env:pip_list", + summary="Print pip packages for the selected resource server.", + flags=( + RESOURCE_SERVER, + _bool_flag("outdated", "outdated", "List only outdated packages."), + Flag( + register=lambda p: p.add_argument( + "--json", action="store_true", help="Output the package list as JSON." + ), + translate_to_hydra=lambda args: ["+format=json"] if args.json else [], + ), + ), + ), + "env test": Command( + target=_env_test, + summary="Test the resource server(s); runs all if no resource server is given.", + flags=(RESOURCE_SERVER,), + ), + "env run": Command( + target="nemo_gym.cli.env:run", + summary="Start the servers.", + flags=(CONFIG, RESOURCE_SERVER_CONFIG, MODEL_TYPE, SEARCH_DIR, MODEL, MODEL_URL, MODEL_API_KEY), + ), + "env status": Command(target="nemo_gym.cli.env:status", summary="Print the server status.", flags=(JSON,)), + "eval prepare": Command( + target="nemo_gym.cli.eval:prepare_benchmark", + summary="Prepare benchmark data and dump it to disk.", + flags=(CONFIG, BENCHMARK, SEARCH_DIR), + ), + "eval run": Command( + target=_eval_run, + summary="Collate data, start servers, and collect rollouts.", + flags=( + CONFIG, + BENCHMARK, + RESOURCE_SERVER_CONFIG, + MODEL_TYPE, + SEARCH_DIR, + Flag( + register=lambda p: p.add_argument( + "--no-serve", + action="store_true", + help="Collect against already-running servers instead of starting them.", + ) + ), + _bool_flag("resume", "resume_from_cache", "Resume from cached rollouts instead of recollecting."), + _value_flag("agent", "agent_name", "Agent to collect rollouts with.", aliases=("-a",)), + _value_flag("input", "input_jsonl_fpath", "Input tasks JSONL file.", aliases=("-i",)), + _value_flag("output", "output_jsonl_fpath", "Output rollouts JSONL file.", aliases=("-o",)), + _value_flag("limit", "limit", "Maximum number of tasks to run."), + _value_flag("num-repeats", "num_repeats", "Number of rollouts per task."), + _value_flag("prompt-config", "prompt_config", "Prompt template YAML to apply."), + _value_flag("concurrency", "num_samples_in_parallel", "Maximum number of concurrent samples."), + _value_flag("split", "split", "Dataset split to use (train, validation, or benchmark)."), + MODEL, + MODEL_URL, + MODEL_API_KEY, + _value_flag("temperature", "responses_create_params.temperature", "Sampling temperature."), + _value_flag("top-p", "responses_create_params.top_p", "Nucleus sampling top-p."), + _value_flag("max-output-tokens", "responses_create_params.max_output_tokens", "Maximum output tokens."), + ), + ), + "eval aggregate": Command( + target="nemo_gym.cli.eval:aggregate_rollouts", + summary="Aggregate sharded rollout results.", + flags=( + CONFIG, + _value_flag( + "output", + "output_jsonl_fpath", + "Path for the merged rollouts and aggregate-metrics file.", + aliases=("-o",), + ), + ), + ), + "eval profile": Command( + target="nemo_gym.cli.eval:reward_profile", + summary="Compute a reward profile from rollouts.", + flags=( + _value_flag( + "inputs", "materialized_inputs_jsonl_fpath", "Materialized inputs JSONL fed to rollout collection." + ), + _value_flag("rollouts", "rollouts_jsonl_fpath", "Rollouts JSONL produced by collection."), + ), + ), + "dev test": Command(target="nemo_gym.cli.dev:dev_test", summary="Run NeMo Gym's unit tests."), +} + + +def _add_leaf(subparsers: argparse._SubParsersAction, name: str, command: Command) -> None: + leaf = subparsers.add_parser(name, help=command.summary, description=command.summary) + # `_parser=leaf` so error reporting (and flag "did you mean?" hints) uses this command's own options/prog. + leaf.set_defaults(_command=command, _parser=leaf) + leaf.add_argument("-v", "--verbose", action="store_true", help="Set logging level to DEBUG.") + for flag in command.flags: + flag.register(leaf) + + +def build_parser() -> argparse.ArgumentParser: + # _GymArgumentParser propagates to every subparser (argparse defaults parser_class to type(self)). + parser = _GymArgumentParser(prog="gym", add_help=True) + parser.add_argument("--version", action="store_true", help="Show the NeMo Gym version and exit.") + parser.add_argument("--json", action="store_true", help="With --version, output as JSON.") + parser.set_defaults(_parser=parser) + + subparsers = parser.add_subparsers() + groups: dict[str, argparse._SubParsersAction] = {} + + for command_name, command in COMMANDS.items(): + parts = command_name.split() + if len(parts) == 1: + _add_leaf(subparsers, parts[0], command) + continue + + group_name, action_name = parts + if group_name not in groups: + group_parser = subparsers.add_parser( + group_name, help=GROUPS.get(group_name), description=GROUPS.get(group_name) + ) + group_parser.set_defaults(_parser=group_parser) + groups[group_name] = group_parser.add_subparsers() + _add_leaf(groups[group_name], action_name, command) + + return parser + + +def main() -> None: + parser = build_parser() + args, overrides = parser.parse_known_args() + + # Hydra overrides never start with "-" so we treat them as unknown flags. + unknown_flags = [token for token in overrides if token.startswith("-")] + if unknown_flags: + error_parser = getattr(args, "_parser", parser) + known_options = [opt for action in error_parser._actions for opt in action.option_strings] + hints = "".join(_did_you_mean(flag.split("=", 1)[0], known_options) for flag in unknown_flags) + error_parser.error(f"unrecognized arguments: {' '.join(unknown_flags)}{hints}") + + if args.version: + dispatch(VERSION_TARGET, ["+json=true", *overrides] if args.json else overrides) + return + + command = getattr(args, "_command", None) + if command is None: + args._parser.print_help() + sys.exit(1) + + try: + translated = [token for flag in command.flags for token in flag.translate_to_hydra(args)] + except ValueError as exc: + sys.stderr.write(f"{getattr(args, '_parser', parser).prog}: error: {exc}\n") + sys.exit(2) + + # --config and the asset selectors all emit +config_paths; coalesce them into one token. + overrides = _merge_config_paths(translated + overrides) + # --verbose flows through the config (as +verbose=true) so it reaches spun-up servers, not just this process. + if getattr(args, "verbose", False): + overrides = ["+verbose=true", *overrides] + if callable(command.target): + command.target(args, overrides) + else: + dispatch(command.target, overrides) diff --git a/nemo_gym/cli_setup_command.py b/nemo_gym/cli_setup_command.py index 9e5a72181d..4632dd3cc0 100644 --- a/nemo_gym/cli_setup_command.py +++ b/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, ) diff --git a/nemo_gym/config_types.py b/nemo_gym/config_types.py index 22c1949884..d3646c810e 100644 --- a/nemo_gym/config_types.py +++ b/nemo_gym/config_types.py @@ -12,10 +12,11 @@ # 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. +import warnings from argparse import ArgumentParser from enum import Enum from pathlib import Path -from typing import Any, ClassVar, Dict, List, Literal, Optional, Set, Tuple, Union +from typing import Annotated, Any, ClassVar, Dict, List, Literal, Optional, Set, Tuple, Union import rich from omegaconf import DictConfig, OmegaConf @@ -358,12 +359,37 @@ def check_output_path(self) -> "DownloadJsonlDatasetHuggingFaceConfig": DatasetType = Union[Literal["train"], Literal["validation"], Literal["example"]] +class GitlabDatasetSource(BaseModel): + """Unified ``source:`` for a dataset fetched from the GitLab model registry.""" + + type: Literal["gitlab"] + dataset_name: str + version: str + artifact_fpath: str + + +class HuggingFaceDatasetSource(BaseModel): + """Unified ``source:`` for a dataset fetched from the HuggingFace Hub.""" + + type: Literal["huggingface"] + repo_id: str + artifact_fpath: Optional[str] = None + + +# One discriminated `source:` block replaces the parallel gitlab_identifier / huggingface_identifier +# fields; `type` selects the backend so it's unambiguous which fields apply. +DatasetSource = Annotated[Union[GitlabDatasetSource, HuggingFaceDatasetSource], Field(discriminator="type")] + + class DatasetConfig(BaseModel): name: str type: DatasetType jsonl_fpath: str num_repeats: int = Field(default=1, ge=1) + # Unified, self-describing dataset source. Prefer this over the legacy *_identifier fields below. + source: Optional[DatasetSource] = None + # Deprecated: kept working (and back-filled from/into `source`) for backward compatibility. gitlab_identifier: Optional[JsonlDatasetGitlabIdentifer] = None huggingface_identifier: Optional[JsonlDatasetHuggingFaceIdentifer] = None license: Optional[ @@ -386,6 +412,54 @@ def check_train_validation_sets(self) -> "DatasetConfig": return self + @model_validator(mode="after") + def normalize_dataset_source(self) -> "DatasetConfig": + """Reconcile the unified `source:` with the legacy `*_identifier` fields. + + Exactly one source may be specified. A legacy identifier is accepted (with a deprecation + warning) and mirrored into `source`; conversely a `source:` is mirrored back into the + matching legacy field so existing consumers that read `gitlab_identifier`/ + `huggingface_identifier` keep working. + """ + specified = [ + name + for name, value in ( + ("source", self.source), + ("gitlab_identifier", self.gitlab_identifier), + ("huggingface_identifier", self.huggingface_identifier), + ) + if value is not None + ] + if len(specified) > 1: + raise ValueError( + f"Specify a dataset source once for '{self.name}': set only one of {specified}. " + "Prefer the unified `source:` block." + ) + if not specified: + return self + + if self.source is None: + # Legacy identifier was used: mirror it into `source` and nudge toward the new field. + if self.gitlab_identifier is not None: + self.source = GitlabDatasetSource(type="gitlab", **self.gitlab_identifier.model_dump()) + else: + self.source = HuggingFaceDatasetSource(type="huggingface", **self.huggingface_identifier.model_dump()) + warnings.warn( + f"`{specified[0]}` is deprecated for dataset '{self.name}'; use the unified " + f"`source:` block (type: {self.source.type}).", + DeprecationWarning, + stacklevel=2, + ) + else: + # `source:` was used: back-fill the matching legacy field for existing consumers. + fields = self.source.model_dump(exclude={"type"}) + if isinstance(self.source, GitlabDatasetSource): + self.gitlab_identifier = JsonlDatasetGitlabIdentifer(**fields) + else: + self.huggingface_identifier = JsonlDatasetHuggingFaceIdentifer(**fields) + + return self + class BenchmarkDatasetConfig(BaseModel): name: str diff --git a/nemo_gym/dataset_orchestrator.py b/nemo_gym/dataset_orchestrator.py index 80401564d7..a2cf9bdb14 100644 --- a/nemo_gym/dataset_orchestrator.py +++ b/nemo_gym/dataset_orchestrator.py @@ -15,15 +15,11 @@ from typing import Union from nemo_gym.config_types import ( - DeleteJsonlDatasetGitlabConfig, - DownloadJsonlDatasetHuggingFaceConfig, UploadJsonlDatasetHuggingFaceConfig, UploadJsonlDatasetHuggingFaceMaybeDeleteConfig, ) from nemo_gym.gitlab_utils import delete_model_from_gitlab, is_model_in_gitlab -from nemo_gym.hf_utils import download_hf_dataset_as_jsonl from nemo_gym.hf_utils import upload_jsonl_dataset as upload_jsonl_dataset_to_hf -from nemo_gym.server_utils import get_global_config_dict def delete_jsonl_dataset_from_gitlab(gitlab_model_name: str) -> None: # pragma: no cover @@ -57,33 +53,3 @@ def upload_jsonl_dataset_to_hf_maybe_delete( if delete_from_gitlab: delete_jsonl_dataset_from_gitlab(gitlab_model_name) - - -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 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 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) diff --git a/nemo_gym/gitlab_utils.py b/nemo_gym/gitlab_utils.py index 23c7a2898f..e83b85787d 100644 --- a/nemo_gym/gitlab_utils.py +++ b/nemo_gym/gitlab_utils.py @@ -74,12 +74,6 @@ def upload_jsonl_dataset( """) -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( config: DownloadJsonlDatasetGitlabConfig, ) -> None: # pragma: no cover @@ -100,12 +94,6 @@ def download_jsonl_dataset( f.write(response.content.decode()) -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 is_model_in_gitlab(model_name: str) -> bool: # pragma: no cover client = create_mlflow_client() diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index a614cecace..006808a48f 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -12,6 +12,8 @@ # 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. +import logging +import sys from argparse import ArgumentParser from collections import defaultdict from copy import deepcopy @@ -68,6 +70,8 @@ COPY_KEY_NAME = "_copy" DELETE_KEY_KEY_NAME = "_delete_key" NEMO_GYM_LOG_DIR_KEY_NAME = "nemo_gym_log_dir" +VERBOSE_KEY_NAME = "verbose" +JSON_OUTPUT_KEY_NAME = "json" NEMO_GYM_RESERVED_TOP_LEVEL_KEYS = [ CONFIG_PATHS_KEY_NAME, ENTRYPOINT_KEY_NAME, @@ -90,6 +94,8 @@ INHERIT_FROM_KEY_NAME, COPY_KEY_NAME, NEMO_GYM_LOG_DIR_KEY_NAME, + VERBOSE_KEY_NAME, + JSON_OUTPUT_KEY_NAME, ] # Data keys @@ -183,6 +189,12 @@ def inner_hydra_wrapper(cfg: DictConfig) -> DictConfig: inner_hydra_wrapper() + # Hydra installs a console log handler on stdout; move it to stderr so command stdout stays machine-readable + # (e.g. `gym ... --json`). Diagnostics belong on stderr; only the requested data goes to stdout. + for handler in logging.getLogger().handlers: + if isinstance(handler, logging.StreamHandler) and getattr(handler, "stream", None) is sys.stdout: + handler.setStream(sys.stderr) + global_config_dict: DictConfig = config_list[0] return global_config_dict @@ -624,6 +636,7 @@ def get_global_config_dict( _GLOBAL_CONFIG_DICT = global_config_dict + _apply_verbosity(global_config_dict) return global_config_dict set_global_config_dict( @@ -631,9 +644,18 @@ def get_global_config_dict( global_config_dict_parser_cls=global_config_dict_parser_cls, ) + _apply_verbosity(_GLOBAL_CONFIG_DICT) return _GLOBAL_CONFIG_DICT +def _apply_verbosity(global_config_dict: DictConfig) -> None: + """Set logging to DEBUG when `verbose` is in the config. Runs in the CLI process and, because the + config dict is forwarded to every spun-up server, in each server process too.""" + if global_config_dict.get(VERBOSE_KEY_NAME): + logging.basicConfig(level=logging.DEBUG) + logging.getLogger().setLevel(logging.DEBUG) + + def set_global_config_dict( global_config_dict_parser_config: Optional[GlobalConfigDictParserConfig] = None, global_config_dict_parser_cls: Type[GlobalConfigDictParser] = GlobalConfigDictParser, diff --git a/nemo_gym/prompt.py b/nemo_gym/prompt.py index 8a4d10f5bc..fae834d7ca 100644 --- a/nemo_gym/prompt.py +++ b/nemo_gym/prompt.py @@ -29,7 +29,6 @@ from nemo_gym import PARENT_DIR from nemo_gym.config_types import BaseNeMoGymCLIConfig -from nemo_gym.global_config import GlobalConfigDictParserConfig, get_global_config_dict class PromptConfig(BaseModel): @@ -162,14 +161,3 @@ class MaterializePromptsConfig(BaseNeMoGymCLIConfig): input_jsonl_fpath: str = Field(description="Raw JSONL data (no responses_create_params.input).") prompt_config: str = Field(description="Path to prompt YAML file to apply.") output_jsonl_fpath: str = Field(description="Output path for materialized JSONL with populated prompts.") - - -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) diff --git a/nemo_gym/resources/judge_resources_server_template.py b/nemo_gym/resources/judge_resources_server_template.py new file mode 100644 index 0000000000..482df69adc --- /dev/null +++ b/nemo_gym/resources/judge_resources_server_template.py @@ -0,0 +1,117 @@ +# 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. +import asyncio +from contextlib import nullcontext +from typing import List, Optional + +from fastapi import FastAPI +from pydantic import Field + +from nemo_gym.base_resources_server import ( + BaseResourcesServerConfig, + BaseVerifyRequest, + BaseVerifyResponse, + SimpleResourcesServer, +) +from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import ( + NeMoGymEasyInputMessage, + NeMoGymResponse, + NeMoGymResponseCreateParamsNonStreaming, +) +from nemo_gym.server_utils import get_response_json + + +class ExampleMultiStepResourcesServerConfig(BaseResourcesServerConfig): + # The auxiliary model in the verification loop — a judge, reward model, or subagent. It is + # wired by name to a model server you pass at run time (see configs/.yaml), exactly like + # `policy_model` is for the agent. + judge_model_server: ModelServerRef + # Base Responses API params for the auxiliary model; `input` is filled in per verify() call. + judge_responses_create_params: NeMoGymResponseCreateParamsNonStreaming + # Optional system prompt and the per-task user prompt. Placeholders: {question}, {answer}. + judge_system_message: Optional[str] = None + judge_prompt_template: str = ( + "Question:\n{question}\n\nAnswer:\n{answer}\n\n" + "Reply with a single number from 0 to 1 scoring how correct the answer is." + ) + # Bound concurrent calls to the auxiliary model so a large rollout batch can't overwhelm it. + # Set to None to disable limiting. + judge_endpoint_max_concurrency: Optional[int] = Field(default=64) + + +class ExampleMultiStepResourcesServer(SimpleResourcesServer): + config: ExampleMultiStepResourcesServerConfig + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + if self.config.judge_endpoint_max_concurrency is not None: + self._judge_concurrency = asyncio.Semaphore(self.config.judge_endpoint_max_concurrency) + else: + self._judge_concurrency = nullcontext() + + def setup_webserver(self) -> FastAPI: + app = super().setup_webserver() + + # Additional server routes go here! e.g.: + # app.post("/get_weather")(self.get_weather) + + return app + + async def _score_with_judge(self, question: str, answer: str) -> float: + """Call the auxiliary model with the per-task prompt and read a reward from its reply.""" + params = self.config.judge_responses_create_params.model_copy(deep=True) + + messages: List[NeMoGymEasyInputMessage] = [] + if self.config.judge_system_message: + messages.append(NeMoGymEasyInputMessage(role="system", content=self.config.judge_system_message)) + prompt = self.config.judge_prompt_template.format(question=question, answer=answer) + messages.append(NeMoGymEasyInputMessage(role="user", content=prompt)) + params.input = messages + + async with self._judge_concurrency: + response = await self.server_client.post( + server_name=self.config.judge_model_server.name, + url_path="/v1/responses", + json=params, + ) + judge_response = NeMoGymResponse.model_validate(await get_response_json(response)) + return self._parse_reward(judge_response) + + @staticmethod + def _parse_reward(judge_response: NeMoGymResponse) -> float: + """Read a clamped 0..1 score from the judge's reply. + + This default expects the judge to reply with a leading number; replace it with your own + parsing (a verdict label, a JSON field, etc.). + """ + try: + text = judge_response.output[-1].content[-1].text + return max(0.0, min(1.0, float(text.strip().split()[0]))) + except (AttributeError, IndexError, ValueError): + return 0.0 + + async def verify(self, body: BaseVerifyRequest) -> BaseVerifyResponse: + # TODO: pull the question and the model's answer for this task out of `body` + # (e.g. the last user message in body.responses_create_params.input and the model's + # final output), then score them with the auxiliary model. + question = "" + answer = "" + reward = await self._score_with_judge(question, answer) + return BaseVerifyResponse(**body.model_dump(), reward=reward) + + +if __name__ == "__main__": + ExampleMultiStepResourcesServer.run_webserver() diff --git a/nemo_gym/resources/judge_resources_server_test_template.py b/nemo_gym/resources/judge_resources_server_test_template.py new file mode 100644 index 0000000000..eaa2326fb1 --- /dev/null +++ b/nemo_gym/resources/judge_resources_server_test_template.py @@ -0,0 +1,49 @@ +# 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 unittest.mock import MagicMock + +from app import ExampleMultiStepResourcesServer, ExampleMultiStepResourcesServerConfig + +from nemo_gym.server_utils import ServerClient + + +def _make_config() -> ExampleMultiStepResourcesServerConfig: + return ExampleMultiStepResourcesServerConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="", + judge_model_server={"type": "responses_api_models", "name": "judge_model"}, + judge_responses_create_params={"model": "judge_model", "input": []}, + ) + + +class TestApp: + def test_sanity(self) -> None: + ExampleMultiStepResourcesServer(config=_make_config(), server_client=MagicMock(spec=ServerClient)) + + def test_parse_reward_reads_leading_number(self) -> None: + server = ExampleMultiStepResourcesServer(config=_make_config(), server_client=MagicMock(spec=ServerClient)) + + judge_response = MagicMock() + judge_response.output[-1].content[-1].text = "0.75 — mostly correct" + assert server._parse_reward(judge_response) == 0.75 + + def test_parse_reward_falls_back_to_zero_on_unparseable_reply(self) -> None: + server = ExampleMultiStepResourcesServer(config=_make_config(), server_client=MagicMock(spec=ServerClient)) + + judge_response = MagicMock() + judge_response.output[-1].content[-1].text = "no score here" + assert server._parse_reward(judge_response) == 0.0 diff --git a/nemo_gym/reward_profile.py b/nemo_gym/reward_profile.py index e7bc8eb6d6..de395f4798 100644 --- a/nemo_gym/reward_profile.py +++ b/nemo_gym/reward_profile.py @@ -31,7 +31,6 @@ AGENT_REF_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, TASK_INDEX_KEY_NAME, - get_global_config_dict, ) @@ -701,31 +700,3 @@ def compute_aggregate_metrics( agent_metrics=serialized_agent, key_metrics=key_metrics, ) - - -def reward_profile(): # pragma: no cover - config = RewardProfileConfig.model_validate(get_global_config_dict()) - - with open(config.materialized_inputs_jsonl_fpath) as f: - rows = list(map(orjson.loads, f)) - - with open(config.rollouts_jsonl_fpath) as f: - results = list(map(orjson.loads, f)) - - # Results may be out of order. - results.sort(key=lambda r: (r[TASK_INDEX_KEY_NAME], r[ROLLOUT_INDEX_KEY_NAME])) - - rp = RewardProfiler() - group_level_metrics, agent_level_metrics = rp.profile_from_data( - rows, results, allow_partial_rollouts=config.allow_partial_rollouts - ) - completion_summary = rp.profile_completion_summary(rows, results) - reward_profiling_fpath, agent_level_metrics_fpath = rp.write_to_disk( - group_level_metrics, agent_level_metrics, Path(config.rollouts_jsonl_fpath) - ) - - print(f"""Profiling outputs: -Reward profile completion: {completion_summary["completed_rollout_rows"]}/{completion_summary["expected_rollout_rows"]} rollout rows ({completion_summary["reward_profile_completion_pct"]:.2f}%) -Input rows: {completion_summary["total_input_rows"]} total; {completion_summary["complete_input_rows"]} complete; {completion_summary["partial_input_rows"]} partial; {completion_summary["missing_input_rows"]} without rollouts dropped from output. -Reward profiling outputs: {reward_profiling_fpath} -Agent-level metrics: {agent_level_metrics_fpath}""") diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index 86ae07d1ff..ed53bf4904 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -44,7 +44,6 @@ from nemo_gym.server_utils import ( GlobalAIOHTTPAsyncClientConfig, ServerClient, - get_global_config_dict, get_response_json, is_global_aiohttp_client_request_debug_enabled, is_global_aiohttp_client_setup, @@ -617,13 +616,6 @@ def setup_server_client( return server_client -def collect_rollouts(): # pragma: no cover - config = RolloutCollectionConfig.model_validate(get_global_config_dict()) - rch = RolloutCollectionHelper() - - asyncio.run(rch.run_from_config(config)) - - class RolloutAggregationConfig(BaseNeMoGymCLIConfig): """ Aggregate metrics across rollout shards produced by `ng_collect_rollouts +disable_aggregation=true`. @@ -719,10 +711,3 @@ async def run_from_config(self, config: RolloutAggregationConfig) -> Optional[Pa Aggregate metrics: {aggregate_metrics_fpath}""") return aggregate_metrics_fpath - - -def aggregate_rollouts(): # pragma: no cover - config = RolloutAggregationConfig.model_validate(get_global_config_dict()) - rah = RolloutAggregationHelper() - - asyncio.run(rah.run_from_config(config)) diff --git a/nemo_gym/server_status.py b/nemo_gym/server_status.py index cfa65eb067..b4166cb94d 100644 --- a/nemo_gym/server_status.py +++ b/nemo_gym/server_status.py @@ -12,6 +12,7 @@ # 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. +import logging from time import time from typing import List @@ -21,6 +22,9 @@ from nemo_gym.server_utils import ServerClient, ServerInstanceDisplayConfig, ServerStatus +logger = logging.getLogger(__name__) + + class StatusCommand: """Main class to check server status""" @@ -73,10 +77,10 @@ def discover_servers(self) -> List[ServerInstanceDisplayConfig]: return servers except (requests.RequestException, ConnectionError) as e: - print(f""" -Could not connect to head server: {e} -Is the head server running? Start it with: `ng_run` - """) + logger.warning( + "Could not connect to head server: %s. Is the head server running? Start it with: `gym env run`", + e, + ) return [] def display_status(self, servers: List[ServerInstanceDisplayConfig]) -> None: diff --git a/nemo_gym/train_data_utils.py b/nemo_gym/train_data_utils.py index 5191741b37..2ea6e4842d 100644 --- a/nemo_gym/train_data_utils.py +++ b/nemo_gym/train_data_utils.py @@ -40,7 +40,6 @@ from nemo_gym.global_config import ( HF_TOKEN_KEY_NAME, GlobalConfigDictParser, - GlobalConfigDictParserConfig, get_global_config_dict, ) from nemo_gym.hf_utils import ( @@ -823,14 +822,3 @@ def validate_backend_credentials(backend: str) -> tuple[bool, str]: ) return True, "" - - -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) diff --git a/pyproject.toml b/pyproject.toml index ac4f4ffacc..260609237b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -311,87 +311,93 @@ Download = "https://github.com/NVIDIA-NeMo/Gym/releases" Homepage = "https://github.com/NVIDIA-NeMo/Gym" [project.scripts] + +gym = "nemo_gym.cli.main:main" +ng = "nemo_gym.cli.main:main" + +########################## DEPRECATED COMMANDS ########################## + # Run/test scripts for servers. -nemo_gym_run = "nemo_gym.cli:run" -ng_run = "nemo_gym.cli:run" -nemo_gym_test = "nemo_gym.cli:test" -ng_test = "nemo_gym.cli:test" -nemo_gym_test_all = "nemo_gym.cli:test_all" -ng_test_all = "nemo_gym.cli:test_all" +nemo_gym_run = "nemo_gym.cli.legacy:main" +ng_run = "nemo_gym.cli.legacy:main" +nemo_gym_test = "nemo_gym.cli.legacy:main" +ng_test = "nemo_gym.cli.legacy:main" +nemo_gym_test_all = "nemo_gym.cli.legacy:main" +ng_test_all = "nemo_gym.cli.legacy:main" # Development convenience aliases. -nemo_gym_dev_test = "nemo_gym.cli:dev_test" -ng_dev_test = "nemo_gym.cli:dev_test" -nemo_gym_init_resources_server = "nemo_gym.cli:init_resources_server" -ng_init_resources_server = "nemo_gym.cli:init_resources_server" +nemo_gym_dev_test = "nemo_gym.cli.legacy:main" +ng_dev_test = "nemo_gym.cli.legacy:main" +nemo_gym_init_resources_server = "nemo_gym.cli.legacy:main" +ng_init_resources_server = "nemo_gym.cli.legacy:main" # Benchmarks -nemo_gym_list_benchmarks = "nemo_gym.benchmarks:list_benchmarks" -ng_list_benchmarks = "nemo_gym.benchmarks:list_benchmarks" -nemo_gym_prepare_benchmark = "nemo_gym.benchmarks:prepare_benchmark" -ng_prepare_benchmark = "nemo_gym.benchmarks:prepare_benchmark" +nemo_gym_list_benchmarks = "nemo_gym.cli.legacy:main" +ng_list_benchmarks = "nemo_gym.cli.legacy:main" +nemo_gym_prepare_benchmark = "nemo_gym.cli.legacy:main" +ng_prepare_benchmark = "nemo_gym.cli.legacy:main" # Rollout collection -nemo_gym_collect_rollouts = "nemo_gym.rollout_collection:collect_rollouts" -ng_collect_rollouts = "nemo_gym.rollout_collection:collect_rollouts" -nemo_gym_e2e_collect_rollouts = "nemo_gym.cli:e2e_rollout_collection" -ng_e2e_collect_rollouts = "nemo_gym.cli:e2e_rollout_collection" -nemo_gym_aggregate_rollouts = "nemo_gym.rollout_collection:aggregate_rollouts" -ng_aggregate_rollouts = "nemo_gym.rollout_collection:aggregate_rollouts" +nemo_gym_collect_rollouts = "nemo_gym.cli.legacy:main" +ng_collect_rollouts = "nemo_gym.cli.legacy:main" +nemo_gym_e2e_collect_rollouts = "nemo_gym.cli.legacy:main" +ng_e2e_collect_rollouts = "nemo_gym.cli.legacy:main" +nemo_gym_aggregate_rollouts = "nemo_gym.cli.legacy:main" +ng_aggregate_rollouts = "nemo_gym.cli.legacy:main" # Prompt materialization -nemo_gym_materialize_prompts = "nemo_gym.prompt:materialize_prompts_cli" -ng_materialize_prompts = "nemo_gym.prompt:materialize_prompts_cli" +nemo_gym_materialize_prompts = "nemo_gym.cli.legacy:main" +ng_materialize_prompts = "nemo_gym.cli.legacy:main" # Reward profiling -nemo_gym_reward_profile = "nemo_gym.reward_profile:reward_profile" -ng_reward_profile = "nemo_gym.reward_profile:reward_profile" +nemo_gym_reward_profile = "nemo_gym.cli.legacy:main" +ng_reward_profile = "nemo_gym.cli.legacy:main" # Dataset management -nemo_gym_upload_dataset_to_gitlab = "nemo_gym.gitlab_utils:upload_jsonl_dataset_cli" -ng_upload_dataset_to_gitlab = "nemo_gym.gitlab_utils:upload_jsonl_dataset_cli" -nemo_gym_download_dataset_from_gitlab = "nemo_gym.gitlab_utils:download_jsonl_dataset_cli" -ng_download_dataset_from_gitlab = "nemo_gym.gitlab_utils:download_jsonl_dataset_cli" -nemo_gym_prepare_data = "nemo_gym.train_data_utils:prepare_data" -ng_prepare_data = "nemo_gym.train_data_utils:prepare_data" +nemo_gym_upload_dataset_to_gitlab = "nemo_gym.cli.legacy:main" +ng_upload_dataset_to_gitlab = "nemo_gym.cli.legacy:main" +nemo_gym_download_dataset_from_gitlab = "nemo_gym.cli.legacy:main" +ng_download_dataset_from_gitlab = "nemo_gym.cli.legacy:main" +nemo_gym_prepare_data = "nemo_gym.cli.legacy:main" +ng_prepare_data = "nemo_gym.cli.legacy:main" # HF dataset upload and download -nemo_gym_upload_dataset_to_hf = "nemo_gym.dataset_orchestrator:upload_jsonl_dataset_to_hf_cli" -ng_upload_dataset_to_hf = "nemo_gym.dataset_orchestrator:upload_jsonl_dataset_to_hf_cli" -nemo_gym_download_dataset_from_hf = "nemo_gym.dataset_orchestrator:download_jsonl_dataset_from_hf_cli" -ng_download_dataset_from_hf = "nemo_gym.dataset_orchestrator:download_jsonl_dataset_from_hf_cli" +nemo_gym_upload_dataset_to_hf = "nemo_gym.cli.legacy:main" +ng_upload_dataset_to_hf = "nemo_gym.cli.legacy:main" +nemo_gym_download_dataset_from_hf = "nemo_gym.cli.legacy:main" +ng_download_dataset_from_hf = "nemo_gym.cli.legacy:main" # Gitlab -> HF (upload then delete) -nemo_gym_gitlab_to_hf_dataset = "nemo_gym.dataset_orchestrator:upload_jsonl_dataset_to_hf_and_delete_gitlab_cli" -ng_gitlab_to_hf_dataset = "nemo_gym.dataset_orchestrator:upload_jsonl_dataset_to_hf_and_delete_gitlab_cli" +nemo_gym_gitlab_to_hf_dataset = "nemo_gym.cli.legacy:main" +ng_gitlab_to_hf_dataset = "nemo_gym.cli.legacy:main" # Manually delete from Gitlab -nemo_gym_delete_dataset_from_gitlab = "nemo_gym.dataset_orchestrator:delete_jsonl_dataset_from_gitlab_cli" -ng_delete_dataset_from_gitlab = "nemo_gym.dataset_orchestrator:delete_jsonl_dataset_from_gitlab_cli" +nemo_gym_delete_dataset_from_gitlab = "nemo_gym.cli.legacy:main" +ng_delete_dataset_from_gitlab = "nemo_gym.cli.legacy:main" # Configuration utils -nemo_gym_dump_config = "nemo_gym.cli:dump_config" -ng_dump_config = "nemo_gym.cli:dump_config" +nemo_gym_dump_config = "nemo_gym.cli.legacy:main" +ng_dump_config = "nemo_gym.cli.legacy:main" # Display help -nemo_gym_help = "nemo_gym.cli:display_help" -ng_help = "nemo_gym.cli:display_help" +nemo_gym_help = "nemo_gym.cli.legacy:main" +ng_help = "nemo_gym.cli.legacy:main" # Server status -nemo_gym_status = "nemo_gym.cli:status" -ng_status = "nemo_gym.cli:status" +nemo_gym_status = "nemo_gym.cli.legacy:main" +ng_status = "nemo_gym.cli.legacy:main" # Environment-specific uv pip list -nemo_gym_pip_list = "nemo_gym.cli:pip_list" -ng_pip_list = "nemo_gym.cli:pip_list" +nemo_gym_pip_list = "nemo_gym.cli.legacy:main" +ng_pip_list = "nemo_gym.cli.legacy:main" # Display version -nemo_gym_version = "nemo_gym.cli:version" -ng_version = "nemo_gym.cli:version" +nemo_gym_version = "nemo_gym.cli.legacy:main" +ng_version = "nemo_gym.cli.legacy:main" # Re-install Gym and dependencies -nemo_gym_reinstall = "nemo_gym.cli:reinstall" -ng_reinstall = "nemo_gym.cli:reinstall" +nemo_gym_reinstall = "nemo_gym.cli.legacy:main" +ng_reinstall = "nemo_gym.cli.legacy:main" [tool.setuptools.packages.find] where = ["."] @@ -449,6 +455,9 @@ omit = [ "/tmp/*", "benchmarks/*", "environments/*/prepare.py", + # Scaffolding templates copied verbatim into generated servers; exercised by the bundled + # test template inside a scaffolded server, not imported in-process. + "nemo_gym/resources/*_template.py", ] data_file = "results/.coverage" concurrency = ["thread", "multiprocessing"] diff --git a/resources_servers/graphwalks/data/example_rollouts.jsonl b/resources_servers/graphwalks/data/example_rollouts.jsonl index 347e9707f0..5b4cf654b6 100644 --- a/resources_servers/graphwalks/data/example_rollouts.jsonl +++ b/resources_servers/graphwalks/data/example_rollouts.jsonl @@ -2,3 +2,4 @@ {"responses_create_params":{"background":null,"include":null,"input":[{"content":"\nYou will be given a graph as a list of directed edges. All nodes are at least degree 1. You will also get a description of an operation to perform on the graph.\nYour job is to execute the operation on the graph and return the set of nodes that the operation results in. If asked for a breadth-first search (BFS), only return the nodes that are both reachable and exactly at that depth (not nodes at intermediate depths), and do not return the starting node. If asked for the parents of a node, only return the nodes that have an edge leading to the given node, do not return the given node itself.\n\nHere is an example:\n\n\nThe graph has the following edges:\nuvwx -> alke\nabcd -> uvwx\nabcd -> efgh\nefgh -> uvwx\n\nExample 1:\nOperation:\nPerform a BFS from node abcd and return only the nodes at exactly depth 1 (not nodes at intermediate depths).\nFinal Answer: [uvwx, efgh]\n\nExample 2:\nOperation:\nPerform a BFS from node alke and return only the nodes at exactly depth 1 (not nodes at intermediate depths).\nFinal Answer: []\n\nExample 3:\nOperation:\nFind the parents of node uvwx.\nFinal Answer: [abcd, efgh]\n\nExample 4:\nOperation:\nFind the parents of node abcd.\nFinal Answer: []\n\n\nHere is the graph to operate on:\nThe graph has the following edges:\ncfcd208495 -> 45c48cce2e\nc4ca4238a0 -> c74d97b01e\nc81e728d9d -> 1c383cd30b\neccbc87e4b -> d645920e39\na87ff679a2 -> e369853df7\ne4da3b7fbb -> c74d97b01e\n1679091c5a -> 70efdf2ec9\n8f14e45fce -> d645920e39\nc9f0f895fb -> aab3238922\n45c48cce2e -> 1f0e3dad99\nd3d9446802 -> d645920e39\n6512bd43d9 -> 34173cb38f\nc20ad4d76f -> eccbc87e4b\nc51ce410c1 -> 182be0c5cd\naab3238922 -> a1d0c6e83f\n9bf31c7ff0 -> a5771bce93\nc74d97b01e -> 70efdf2ec9\n70efdf2ec9 -> aab3238922\n6f4922f455 -> c4ca4238a0\n1f0e3dad99 -> f7177163c8\n98f1370821 -> c20ad4d76f\n3c59dc048e -> aab3238922\nb6d767d2f8 -> 4e732ced34\n37693cfc74 -> e369853df7\n1ff1de7740 -> 33e75ff09d\n8e296a067a -> 3c59dc048e\n4e732ced34 -> c81e728d9d\n02e74f10e0 -> 6ea9ab1baa\n33e75ff09d -> c74d97b01e\n6ea9ab1baa -> 98f1370821\n34173cb38f -> 37693cfc74\nc16a5320fa -> 33e75ff09d\n6364d3f0f4 -> 3c59dc048e\n182be0c5cd -> a87ff679a2\ne369853df7 -> 8e296a067a\n1c383cd30b -> a5771bce93\n19ca14e7ea -> d3d9446802\na5bfc9e079 -> 1c383cd30b\na5771bce93 -> 98f1370821\nd67d8ab4f4 -> d3d9446802\nd645920e39 -> 9bf31c7ff0\n3416a75f4c -> c9f0f895fb\na1d0c6e83f -> 182be0c5cd\n17e62166fc -> 182be0c5cd\nf7177163c8 -> c16a5320fa\n6c8349cc72 -> 17e62166fc\n\n\nOperation:\nFind the parents of node 3416a75f4c.\n\nYou should reason through the operation step by step, and then return the set of nodes that the operation results in. Return your final answer as a list of nodes in the very last line of your response. For example, if the operation returns the set of nodes [node1, node2, node3], your response should be:\nFinal Answer: [node1, node2, node3]\nIf the operation returns the empty set, your response should be:\nFinal Answer: []\n","role":"user","type":"message"}],"instructions":null,"max_output_tokens":null,"max_tool_calls":null,"metadata":null,"model":null,"parallel_tool_calls":true,"previous_response_id":null,"prompt":null,"reasoning":null,"service_tier":null,"store":null,"temperature":null,"text":null,"tool_choice":"auto","tools":[],"top_logprobs":null,"top_p":null,"truncation":null,"user":null,"stream":null},"response":{"id":"resp_22bb0cb613854aa9bc4084321c6119c1","created_at":1780041657.0,"error":null,"incomplete_details":null,"instructions":null,"metadata":null,"model":"/lustre/fsw/portfolios/llmservice/users/jiaqiz/results/ultra-v3-posttraining-mopd/production-ultra-mopd-studentpurplestep66_blend20260524_tp8_cp32_ep64_pp1_gpp1_pps1024_gbs1024-20260524-jiaqiz/eval/step_42/hf","object":"response","output":[{"id":"rs_839f758f5089444cb677cd5fb2ccc120","summary":[{"text":"We are given a graph as a list of directed edges. We need to find the parents of node \"3416a75f4c\". Parents are nodes that have an edge leading to the given node. So we need to find all edges where the target is \"3416a75f4c\". Let's scan the edges.\n\nEdges list:\ncfcd208495 -> 45c48cce2e\nc4ca4238a0 -> c74d97b01e\nc81e728d9d -> 1c383cd30b\neccbc87e4b -> d645920e39\na87ff679a2 -> e369853df7\ne4da3b7fbb -> c74d97b01e\n1679091c5a -> 70efdf2ec9\n8f14e45fce -> d645920e39\nc9f0f895fb -> aab3238922\n45c48cce2e -> 1f0e3dad99\nd3d9446802 -> d645920e39\n6512bd43d9 -> 34173cb38f\nc20ad4d76f -> eccbc87e4b\nc51ce410c1 -> 182be0c5cd\naab3238922 -> a1d0c6e83f\n9bf31c7ff0 -> a5771bce93\nc74d97b01e -> 70efdf2ec9\n70efdf2ec9 -> aab3238922\n6f4922f455 -> c4ca4238a0\n1f0e3dad99 -> f7177163c8\n98f1370821 -> c20ad4d76f\n3c59dc048e -> aab3238922\nb6d767d2f8 -> 4e732ced34\n37693cfc74 -> e369853df7\n1ff1de7740 -> 33e75ff09d\n8e296a067a -> 3c59dc048e\n4e732ced34 -> c81e728d9d\n02e74f10e0 -> 6ea9ab1baa\n33e75ff09d -> c74d97b01e\n6ea9ab1baa -> 98f1370821\n34173cb38f -> 37693cfc74\nc16a5320fa -> 33e75ff09d\n6364d3f0f4 -> 3c59dc048e\n182be0c5cd -> a87ff679a2\ne369853df7 -> 8e296a067a\n1c383cd30b -> a5771bce93\n19ca14e7ea -> d3d9446802\na5bfc9e079 -> 1c383cd30b\na5771bce93 -> 98f1370821\nd67d8ab4f4 -> d3d9446802\nd645920e39 -> 9bf31c7ff0\n3416a75f4c -> c9f0f895fb\na1d0c6e83f -> 182be0c5cd\n17e62166fc -> 182be0c5cd\nf7177163c8 -> c16a5320fa\n6c8349cc72 -> 17e62166fc\n\nWe are looking for edges where the target is \"3416a75f4c\". Scanning the list, I see only one edge that has \"3416a75f4c\" as source: \"3416a75f4c -> c9f0f895fb\". But we need target. There is no edge with target \"3416a75f4c\". So the set of parents is empty.\n\nThus final answer: [].\n\nBut let's double-check: The operation says \"Find the parents of node 3416a75f4c.\" Parents are nodes that have an edge leading to the given node. So we need edges where the second node is 3416a75f4c. None exist. So empty set.\n\nReturn as list: [].\n\nFinal Answer: []","type":"summary_text"}],"type":"reasoning","encrypted_content":null},{"id":"msg_0f9f7a0dfaf6404cbef0474e2706a04b","content":[{"annotations":[],"text":"Final Answer: []","type":"output_text","logprobs":null}],"role":"assistant","status":"completed","type":"message"}],"parallel_tool_calls":true,"temperature":null,"tool_choice":"auto","tools":[],"top_p":null,"background":null,"conversation":null,"max_output_tokens":null,"max_tool_calls":null,"previous_response_id":null,"prompt":null,"prompt_cache_key":null,"reasoning":null,"safety_identifier":null,"service_tier":null,"status":null,"text":null,"top_logprobs":null,"truncation":null,"usage":{"input_tokens":1384,"input_tokens_details":{"cached_tokens":0},"output_tokens":1211,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2595},"user":null},"reward":1.0,"expected_answer":"[]","problem_type":"parents","n_tokens":1369,"prompt_chars":2688,"f1":1.0,"parse_failed":false,"predicted_nodes":[],"_ng_task_index":5,"_ng_rollout_index":0,"agent_ref":{"type":"responses_api_agents","name":"graphwalks_n3_1m_benchmark_simple_agent"}} {"responses_create_params":{"background":null,"include":null,"input":[{"content":"\nYou will be given a graph as a list of directed edges. All nodes are at least degree 1. You will also get a description of an operation to perform on the graph.\nYour job is to execute the operation on the graph and return the set of nodes that the operation results in. If asked for a breadth-first search (BFS), only return the nodes that are both reachable and exactly at that depth (not nodes at intermediate depths), and do not return the starting node. If asked for the parents of a node, only return the nodes that have an edge leading to the given node, do not return the given node itself.\n\nHere is an example:\n\n\nThe graph has the following edges:\nuvwx -> alke\nabcd -> uvwx\nabcd -> efgh\nefgh -> uvwx\n\nExample 1:\nOperation:\nPerform a BFS from node abcd and return only the nodes at exactly depth 1 (not nodes at intermediate depths).\nFinal Answer: [uvwx, efgh]\n\nExample 2:\nOperation:\nPerform a BFS from node alke and return only the nodes at exactly depth 1 (not nodes at intermediate depths).\nFinal Answer: []\n\nExample 3:\nOperation:\nFind the parents of node uvwx.\nFinal Answer: [abcd, efgh]\n\nExample 4:\nOperation:\nFind the parents of node abcd.\nFinal Answer: []\n\n\nHere is the graph to operate on:\nThe graph has the following edges:\ncfcd208495 -> 1679091c5a\ncfcd208495 -> c4ca4238a0\ncfcd208495 -> c4ca4238a0\ncfcd208495 -> e4da3b7fbb\ncfcd208495 -> c4ca4238a0\nc4ca4238a0 -> 1679091c5a\nc4ca4238a0 -> 1679091c5a\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> c81e728d9d\nc4ca4238a0 -> c81e728d9d\nc81e728d9d -> 8f14e45fce\nc81e728d9d -> 8f14e45fce\nc81e728d9d -> c9f0f895fb\nc81e728d9d -> c4ca4238a0\nc81e728d9d -> eccbc87e4b\neccbc87e4b -> cfcd208495\neccbc87e4b -> c81e728d9d\neccbc87e4b -> c4ca4238a0\neccbc87e4b -> c81e728d9d\neccbc87e4b -> c81e728d9d\na87ff679a2 -> e4da3b7fbb\na87ff679a2 -> c4ca4238a0\na87ff679a2 -> cfcd208495\na87ff679a2 -> a87ff679a2\na87ff679a2 -> e4da3b7fbb\ne4da3b7fbb -> c81e728d9d\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> c81e728d9d\ne4da3b7fbb -> c9f0f895fb\ne4da3b7fbb -> c81e728d9d\n1679091c5a -> 1679091c5a\n1679091c5a -> c81e728d9d\n1679091c5a -> a87ff679a2\n1679091c5a -> c4ca4238a0\n1679091c5a -> 8f14e45fce\n8f14e45fce -> 8f14e45fce\n8f14e45fce -> a87ff679a2\n8f14e45fce -> c81e728d9d\n8f14e45fce -> c81e728d9d\n8f14e45fce -> c4ca4238a0\nc9f0f895fb -> a87ff679a2\nc9f0f895fb -> c9f0f895fb\nc9f0f895fb -> c9f0f895fb\nc9f0f895fb -> cfcd208495\nc9f0f895fb -> eccbc87e4b\n\n\nOperation:\nFind the parents of node cfcd208495.\n\nYou should reason through the operation step by step, and then return the set of nodes that the operation results in. Return your final answer as a list of nodes in the very last line of your response. For example, if the operation returns the set of nodes [node1, node2, node3], your response should be:\nFinal Answer: [node1, node2, node3]\nIf the operation returns the empty set, your response should be:\nFinal Answer: []\n","role":"user","type":"message"}],"instructions":null,"max_output_tokens":null,"max_tool_calls":null,"metadata":null,"model":null,"parallel_tool_calls":true,"previous_response_id":null,"prompt":null,"reasoning":null,"service_tier":null,"store":null,"temperature":null,"text":null,"tool_choice":"auto","tools":[],"top_logprobs":null,"top_p":null,"truncation":null,"user":null,"stream":null},"response":{"id":"resp_6b3568a7e3574ec686a30a88119ea2fa","created_at":1780041661.0,"error":null,"incomplete_details":null,"instructions":null,"metadata":null,"model":"/lustre/fsw/portfolios/llmservice/users/jiaqiz/results/ultra-v3-posttraining-mopd/production-ultra-mopd-studentpurplestep66_blend20260524_tp8_cp32_ep64_pp1_gpp1_pps1024_gbs1024-20260524-jiaqiz/eval/step_42/hf","object":"response","output":[{"id":"rs_cbad1103e8bf46a0b994d1cd18ef48b7","summary":[{"text":"We are given a graph as a list of directed edges. We need to find the parents of node cfcd208495. Parents are nodes that have an edge leading to the given node. So we need to find all nodes X such that there is an edge X -> cfcd208495.\n\nWe have the list of edges. Let's parse them and find all edges where the target is cfcd208495.\n\nEdges list:\n\ncfcd208495 -> 1679091c5a\ncfcd208495 -> c4ca4238a0\ncfcd208495 -> c4ca4238a0\ncfcd208495 -> e4da3b7fbb\ncfcd208495 -> c4ca4238a0\nc4ca4238a0 -> 1679091c5a\nc4ca4238a0 -> 1679091c5a\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> c81e728d9d\nc4ca4238a0 -> c81e728d9d\nc81e728d9d -> 8f14e45fce\nc81e728d9d -> 8f14e45fce\nc81e728d9d -> c9f0f895fb\nc81e728d9d -> c4ca4238a0\nc81e728d9d -> eccbc87e4b\neccbc87e4b -> cfcd208495\neccbc87e4b -> c81e728d9d\neccbc87e4b -> c4ca4238a0\neccbc87e4b -> c81e728d9d\neccbc87e4b -> c81e728d9d\na87ff679a2 -> e4da3b7fbb\na87ff679a2 -> c4ca4238a0\na87ff679a2 -> cfcd208495\na87ff679a2 -> a87ff679a2\na87ff679a2 -> e4da3b7fbb\ne4da3b7fbb -> c81e728d9d\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> c81e728d9d\ne4da3b7fbb -> c9f0f895fb\ne4da3b7fbb -> c81e728d9d\n1679091c5a -> 1679091c5a\n1679091c5a -> c81e728d9d\n1679091c5a -> a87ff679a2\n1679091c5a -> c4ca4238a0\n1679091c5a -> 8f14e45fce\n8f14e45fce -> 8f14e45fce\n8f14e45fce -> a87ff679a2\n8f14e45fce -> c81e728d9d\n8f14e45fce -> c81e728d9d\n8f14e45fce -> c4ca4238a0\nc9f0f895fb -> a87ff679a2\nc9f0f895fb -> c9f0f895fb\nc9f0f895fb -> c9f0f895fb\nc9f0f895fb -> cfcd208495\nc9f0f895fb -> eccbc87e4b\n\nNow find edges where target is cfcd208495:\n\n- eccbc87e4b -> cfcd208495\n- a87ff679a2 -> cfcd208495\n- e4da3b7fbb -> cfcd208495\n- c9f0f895fb -> cfcd208495\n\nAlso check if any other edges: maybe cfcd208495 -> cfcd208495? Not present. So parents are: eccbc87e4b, a87ff679a2, e4da3b7fbb, c9f0f895fb.\n\nWe should return them as a list. Order? The example returns in the order they appear? In example 3, parents of uvwx: [abcd, efgh] which matches the order of edges: abcd -> uvwx appears first, then efgh -> uvwx appears later. In our list, edges appear in order: eccbc87e4b -> cfcd208495 (first occurrence), then a87ff679a2 -> cfcd208495, then e4da3b7fbb -> cfcd208495, then c9f0f895fb -> cfcd208495. So order: eccbc87e4b, a87ff679a2, e4da3b7fbb, c9f0f895fb.\n\nWe should output as a list. The problem says \"return the set of nodes\" but example shows list. We'll output as list.\n\nFinal Answer: [eccbc87e4b, a87ff679a2, e4da3b7fbb, c9f0f895fb]","type":"summary_text"}],"type":"reasoning","encrypted_content":null},{"id":"msg_93248f6f633741d7b3278cae2ba26f4d","content":[{"annotations":[],"text":"The parents of a node are the nodes that have a directed edge leading to that node. Examining the given edges, the edges that point to `cfcd208495` are:\n\n- `eccbc87e4b -> cfcd208495`\n- `a87ff679a2 -> cfcd208495`\n- `e4da3b7fbb -> cfcd208495`\n- `c9f0f895fb -> cfcd208495`\n\nThus, the parents of `cfcd208495` are these four nodes.\n\nFinal Answer: [eccbc87e4b, a87ff679a2, e4da3b7fbb, c9f0f895fb]","type":"output_text","logprobs":null}],"role":"assistant","status":"completed","type":"message"}],"parallel_tool_calls":true,"temperature":null,"tool_choice":"auto","tools":[],"top_p":null,"background":null,"conversation":null,"max_output_tokens":null,"max_tool_calls":null,"previous_response_id":null,"prompt":null,"prompt_cache_key":null,"reasoning":null,"safety_identifier":null,"service_tier":null,"status":null,"text":null,"top_logprobs":null,"truncation":null,"usage":{"input_tokens":1344,"input_tokens_details":{"cached_tokens":0},"output_tokens":1611,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2955},"user":null},"reward":1.0,"expected_answer":"[\"a87ff679a2\", \"c9f0f895fb\", \"e4da3b7fbb\", \"eccbc87e4b\"]","problem_type":"parents","n_tokens":1329,"prompt_chars":2663,"f1":1.0,"parse_failed":false,"predicted_nodes":["eccbc87e4b","a87ff679a2","e4da3b7fbb","c9f0f895fb"],"_ng_task_index":2,"_ng_rollout_index":0,"agent_ref":{"type":"responses_api_agents","name":"graphwalks_n3_1m_benchmark_simple_agent"}} {"responses_create_params":{"background":null,"include":null,"input":[{"content":"\nYou will be given a graph as a list of directed edges. All nodes are at least degree 1. You will also get a description of an operation to perform on the graph.\nYour job is to execute the operation on the graph and return the set of nodes that the operation results in. If asked for a breadth-first search (BFS), only return the nodes that are both reachable and exactly at that depth (not nodes at intermediate depths), and do not return the starting node. If asked for the parents of a node, only return the nodes that have an edge leading to the given node, do not return the given node itself.\n\nHere is an example:\n\n\nThe graph has the following edges:\nuvwx -> alke\nabcd -> uvwx\nabcd -> efgh\nefgh -> uvwx\n\nExample 1:\nOperation:\nPerform a BFS from node abcd and return only the nodes at exactly depth 1 (not nodes at intermediate depths).\nFinal Answer: [uvwx, efgh]\n\nExample 2:\nOperation:\nPerform a BFS from node alke and return only the nodes at exactly depth 1 (not nodes at intermediate depths).\nFinal Answer: []\n\nExample 3:\nOperation:\nFind the parents of node uvwx.\nFinal Answer: [abcd, efgh]\n\nExample 4:\nOperation:\nFind the parents of node abcd.\nFinal Answer: []\n\n\nHere is the graph to operate on:\nThe graph has the following edges:\ncfcd208495 -> 45c48cce2e\ncfcd208495 -> c81e728d9d\ncfcd208495 -> cfcd208495\ncfcd208495 -> d3d9446802\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> cfcd208495\nc4ca4238a0 -> 8f14e45fce\nc81e728d9d -> e4da3b7fbb\nc81e728d9d -> d3d9446802\nc81e728d9d -> 1679091c5a\nc81e728d9d -> c4ca4238a0\neccbc87e4b -> eccbc87e4b\neccbc87e4b -> 8f14e45fce\neccbc87e4b -> 1679091c5a\neccbc87e4b -> e4da3b7fbb\na87ff679a2 -> c81e728d9d\na87ff679a2 -> 45c48cce2e\na87ff679a2 -> c4ca4238a0\na87ff679a2 -> e4da3b7fbb\ne4da3b7fbb -> c4ca4238a0\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> 8f14e45fce\n1679091c5a -> cfcd208495\n1679091c5a -> eccbc87e4b\n1679091c5a -> cfcd208495\n1679091c5a -> d3d9446802\n8f14e45fce -> c81e728d9d\n8f14e45fce -> 45c48cce2e\n8f14e45fce -> c81e728d9d\n8f14e45fce -> 1679091c5a\nc9f0f895fb -> eccbc87e4b\nc9f0f895fb -> 45c48cce2e\nc9f0f895fb -> cfcd208495\nc9f0f895fb -> c4ca4238a0\n45c48cce2e -> d3d9446802\n45c48cce2e -> 8f14e45fce\n45c48cce2e -> 8f14e45fce\n45c48cce2e -> d3d9446802\nd3d9446802 -> eccbc87e4b\nd3d9446802 -> e4da3b7fbb\nd3d9446802 -> a87ff679a2\nd3d9446802 -> 45c48cce2e\n\n\nOperation:\nFind the parents of node a87ff679a2.\n\nYou should reason through the operation step by step, and then return the set of nodes that the operation results in. Return your final answer as a list of nodes in the very last line of your response. For example, if the operation returns the set of nodes [node1, node2, node3], your response should be:\nFinal Answer: [node1, node2, node3]\nIf the operation returns the empty set, your response should be:\nFinal Answer: []\n","role":"user","type":"message"}],"instructions":null,"max_output_tokens":null,"max_tool_calls":null,"metadata":null,"model":null,"parallel_tool_calls":true,"previous_response_id":null,"prompt":null,"reasoning":null,"service_tier":null,"store":null,"temperature":null,"text":null,"tool_choice":"auto","tools":[],"top_logprobs":null,"top_p":null,"truncation":null,"user":null,"stream":null},"response":{"id":"resp_49c351616734407a8520348626394e6b","created_at":1780041667.0,"error":null,"incomplete_details":null,"instructions":null,"metadata":null,"model":"/lustre/fsw/portfolios/llmservice/users/jiaqiz/results/ultra-v3-posttraining-mopd/production-ultra-mopd-studentpurplestep66_blend20260524_tp8_cp32_ep64_pp1_gpp1_pps1024_gbs1024-20260524-jiaqiz/eval/step_42/hf","object":"response","output":[{"id":"rs_dd9106bbf0424062ad9f781cd58ac20b","summary":[{"text":"We are given a graph as a list of directed edges. We need to find the parents of node a87ff679a2. Parents are nodes that have an edge leading to the given node. So we need to look through all edges and find those where the target is a87ff679a2. The source nodes of those edges are the parents.\n\nLet's list edges and check for target = a87ff679a2.\n\nEdges:\ncfcd208495 -> 45c48cce2e\ncfcd208495 -> c81e728d9d\ncfcd208495 -> cfcd208495\ncfcd208495 -> d3d9446802\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> cfcd208495\nc4ca4238a0 -> 8f14e45fce\nc81e728d9d -> e4da3b7fbb\nc81e728d9d -> d3d9446802\nc81e728d9d -> 1679091c5a\nc81e728d9d -> c4ca4238a0\neccbc87e4b -> eccbc87e4b\neccbc87e4b -> 8f14e45fce\neccbc87e4b -> 1679091c5a\neccbc87e4b -> e4da3b7fbb\na87ff679a2 -> c81e728d9d\na87ff679a2 -> 45c48cce2e\na87ff679a2 -> c4ca4238a0\na87ff679a2 -> e4da3b7fbb\ne4da3b7fbb -> c4ca4238a0\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> 8f14e45fce\n1679091c5a -> cfcd208495\n1679091c5a -> eccbc87e4b\n1679091c5a -> cfcd208495\n1679091c5a -> d3d9446802\n8f14e45fce -> c81e728d9d\n8f14e45fce -> 45c48cce2e\n8f14e45fce -> c81e728d9d\n8f14e45fce -> 1679091c5a\nc9f0f895fb -> eccbc87e4b\nc9f0f895fb -> 45c48cce2e\nc9f0f895fb -> cfcd208495\nc9f0f895fb -> c4ca4238a0\n45c48cce2e -> d3d9446802\n45c48cce2e -> 8f14e45fce\n45c48cce2e -> 8f14e45fce\n45c48cce2e -> d3d9446802\nd3d9446802 -> eccbc87e4b\nd3d9446802 -> e4da3b7fbb\nd3d9446802 -> a87ff679a2\nd3d9446802 -> 45c48cce2e\n\nNow find edges where target is a87ff679a2. Scanning:\n\n- d3d9446802 -> a87ff679a2 (yes)\nAny others? Let's check each line:\ncfcd208495 -> ... no\nc4ca4238a0 -> ... no\nc81e728d9d -> ... no\neccbc87e4b -> ... no\na87ff679a2 -> ... (source, not target)\ne4da3b7fbb -> ... no\n1679091c5a -> ... no\n8f14e45fce -> ... no\nc9f0f895fb -> ... no\n45c48cce2e -> ... no\nd3d9446802 -> a87ff679a2 (found)\n\nOnly one parent: d3d9446802.\n\nThus answer: [d3d9446802]\n\nWe should return as a list. The instruction: \"Return your final answer as a list of nodes in the very last line of your response. For example, if the operation returns the set of nodes [node1, node2, node3], your response should be: Final Answer: [node1, node2, node3]\"\n\nSo we output: Final Answer: [d3d9446802]\n\nBut careful: The example output uses brackets and commas. We'll follow that.\n\nWe'll also include reasoning, but final answer must be on last line.","type":"summary_text"}],"type":"reasoning","encrypted_content":null},{"id":"msg_03feedc42c474ab1ab5663cf13384892","content":[{"annotations":[],"text":"The parents of a node are the nodes that have a directed edge pointing to it. Scanning the given edges for those where the target is `a87ff679a2`, we find only one such edge:\n\n- `d3d9446802 -> a87ff679a2`\n\nNo other edges have `a87ff679a2` as the destination. Therefore, the only parent is `d3d9446802`.\n\nFinal Answer: [d3d9446802]","type":"output_text","logprobs":null}],"role":"assistant","status":"completed","type":"message"}],"parallel_tool_calls":true,"temperature":null,"tool_choice":"auto","tools":[],"top_p":null,"background":null,"conversation":null,"max_output_tokens":null,"max_tool_calls":null,"previous_response_id":null,"prompt":null,"prompt_cache_key":null,"reasoning":null,"safety_identifier":null,"service_tier":null,"status":null,"text":null,"top_logprobs":null,"truncation":null,"usage":{"input_tokens":1323,"input_tokens_details":{"cached_tokens":0},"output_tokens":1464,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2787},"user":null},"reward":1.0,"expected_answer":"[\"d3d9446802\"]","problem_type":"parents","n_tokens":1308,"prompt_chars":2638,"f1":1.0,"parse_failed":false,"predicted_nodes":["d3d9446802"],"_ng_task_index":11,"_ng_rollout_index":0,"agent_ref":{"type":"responses_api_agents","name":"graphwalks_n3_1m_benchmark_simple_agent"}} +{"responses_create_params":{"background":null,"include":null,"input":[{"content":"\nYou will be given a graph as a list of directed edges. All nodes are at least degree 1. You will also get a description of an operation to perform on the graph.\nYour job is to execute the operation on the graph and return the set of nodes that the operation results in. If asked for a breadth-first search (BFS), only return the nodes that are both reachable and exactly at that depth (not nodes at intermediate depths), and do not return the starting node. If asked for the parents of a node, only return the nodes that have an edge leading to the given node, do not return the given node itself.\n\nHere is an example:\n\n\nThe graph has the following edges:\nuvwx -> alke\nabcd -> uvwx\nabcd -> efgh\nefgh -> uvwx\n\nExample 1:\nOperation:\nPerform a BFS from node abcd and return only the nodes at exactly depth 1 (not nodes at intermediate depths).\nFinal Answer: [uvwx, efgh]\n\nExample 2:\nOperation:\nPerform a BFS from node alke and return only the nodes at exactly depth 1 (not nodes at intermediate depths).\nFinal Answer: []\n\nExample 3:\nOperation:\nFind the parents of node uvwx.\nFinal Answer: [abcd, efgh]\n\nExample 4:\nOperation:\nFind the parents of node abcd.\nFinal Answer: []\n\n\nHere is the graph to operate on:\nThe graph has the following edges:\ncfcd208495 -> 45c48cce2e\ncfcd208495 -> c81e728d9d\ncfcd208495 -> cfcd208495\ncfcd208495 -> d3d9446802\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> cfcd208495\nc4ca4238a0 -> 8f14e45fce\nc81e728d9d -> e4da3b7fbb\nc81e728d9d -> d3d9446802\nc81e728d9d -> 1679091c5a\nc81e728d9d -> c4ca4238a0\neccbc87e4b -> eccbc87e4b\neccbc87e4b -> 8f14e45fce\neccbc87e4b -> 1679091c5a\neccbc87e4b -> e4da3b7fbb\na87ff679a2 -> c81e728d9d\na87ff679a2 -> 45c48cce2e\na87ff679a2 -> c4ca4238a0\na87ff679a2 -> e4da3b7fbb\ne4da3b7fbb -> c4ca4238a0\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> 8f14e45fce\n1679091c5a -> cfcd208495\n1679091c5a -> eccbc87e4b\n1679091c5a -> cfcd208495\n1679091c5a -> d3d9446802\n8f14e45fce -> c81e728d9d\n8f14e45fce -> 45c48cce2e\n8f14e45fce -> c81e728d9d\n8f14e45fce -> 1679091c5a\nc9f0f895fb -> eccbc87e4b\nc9f0f895fb -> 45c48cce2e\nc9f0f895fb -> cfcd208495\nc9f0f895fb -> c4ca4238a0\n45c48cce2e -> d3d9446802\n45c48cce2e -> 8f14e45fce\n45c48cce2e -> 8f14e45fce\n45c48cce2e -> d3d9446802\nd3d9446802 -> eccbc87e4b\nd3d9446802 -> e4da3b7fbb\nd3d9446802 -> a87ff679a2\nd3d9446802 -> 45c48cce2e\n\n\nOperation:\nFind the parents of node a87ff679a2.\n\nYou should reason through the operation step by step, and then return the set of nodes that the operation results in. Return your final answer as a list of nodes in the very last line of your response. For example, if the operation returns the set of nodes [node1, node2, node3], your response should be:\nFinal Answer: [node1, node2, node3]\nIf the operation returns the empty set, your response should be:\nFinal Answer: []\n","role":"user","type":"message"}],"instructions":null,"max_output_tokens":null,"max_tool_calls":null,"metadata":null,"model":null,"parallel_tool_calls":true,"previous_response_id":null,"prompt":null,"reasoning":null,"service_tier":null,"store":null,"temperature":null,"text":null,"tool_choice":"auto","tools":[],"top_logprobs":null,"top_p":null,"truncation":null,"user":null,"stream":null},"response":{"id":"resp_49c351616734407a8520348626394e6b","created_at":1780041667.0,"error":null,"incomplete_details":null,"instructions":null,"metadata":null,"model":"/lustre/fsw/portfolios/llmservice/users/jiaqiz/results/ultra-v3-posttraining-mopd/production-ultra-mopd-studentpurplestep66_blend20260524_tp8_cp32_ep64_pp1_gpp1_pps1024_gbs1024-20260524-jiaqiz/eval/step_42/hf","object":"response","output":[{"id":"rs_dd9106bbf0424062ad9f781cd58ac20b","summary":[{"text":"We are given a graph as a list of directed edges. We need to find the parents of node a87ff679a2. Parents are nodes that have an edge leading to the given node. So we need to look through all edges and find those where the target is a87ff679a2. The source nodes of those edges are the parents.\n\nLet's list edges and check for target = a87ff679a2.\n\nEdges:\ncfcd208495 -> 45c48cce2e\ncfcd208495 -> c81e728d9d\ncfcd208495 -> cfcd208495\ncfcd208495 -> d3d9446802\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> e4da3b7fbb\nc4ca4238a0 -> cfcd208495\nc4ca4238a0 -> 8f14e45fce\nc81e728d9d -> e4da3b7fbb\nc81e728d9d -> d3d9446802\nc81e728d9d -> 1679091c5a\nc81e728d9d -> c4ca4238a0\neccbc87e4b -> eccbc87e4b\neccbc87e4b -> 8f14e45fce\neccbc87e4b -> 1679091c5a\neccbc87e4b -> e4da3b7fbb\na87ff679a2 -> c81e728d9d\na87ff679a2 -> 45c48cce2e\na87ff679a2 -> c4ca4238a0\na87ff679a2 -> e4da3b7fbb\ne4da3b7fbb -> c4ca4238a0\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> cfcd208495\ne4da3b7fbb -> 8f14e45fce\n1679091c5a -> cfcd208495\n1679091c5a -> eccbc87e4b\n1679091c5a -> cfcd208495\n1679091c5a -> d3d9446802\n8f14e45fce -> c81e728d9d\n8f14e45fce -> 45c48cce2e\n8f14e45fce -> c81e728d9d\n8f14e45fce -> 1679091c5a\nc9f0f895fb -> eccbc87e4b\nc9f0f895fb -> 45c48cce2e\nc9f0f895fb -> cfcd208495\nc9f0f895fb -> c4ca4238a0\n45c48cce2e -> d3d9446802\n45c48cce2e -> 8f14e45fce\n45c48cce2e -> 8f14e45fce\n45c48cce2e -> d3d9446802\nd3d9446802 -> eccbc87e4b\nd3d9446802 -> e4da3b7fbb\nd3d9446802 -> a87ff679a2\nd3d9446802 -> 45c48cce2e\n\nNow find edges where target is a87ff679a2. Scanning:\n\n- d3d9446802 -> a87ff679a2 (yes)\nAny others? Let's check each line:\ncfcd208495 -> ... no\nc4ca4238a0 -> ... no\nc81e728d9d -> ... no\neccbc87e4b -> ... no\na87ff679a2 -> ... (source, not target)\ne4da3b7fbb -> ... no\n1679091c5a -> ... no\n8f14e45fce -> ... no\nc9f0f895fb -> ... no\n45c48cce2e -> ... no\nd3d9446802 -> a87ff679a2 (found)\n\nOnly one parent: d3d9446802.\n\nThus answer: [d3d9446802]\n\nWe should return as a list. The instruction: \"Return your final answer as a list of nodes in the very last line of your response. For example, if the operation returns the set of nodes [node1, node2, node3], your response should be: Final Answer: [node1, node2, node3]\"\n\nSo we output: Final Answer: [d3d9446802]\n\nBut careful: The example output uses brackets and commas. We'll follow that.\n\nWe'll also include reasoning, but final answer must be on last line.","type":"summary_text"}],"type":"reasoning","encrypted_content":null},{"id":"msg_03feedc42c474ab1ab5663cf13384892","content":[{"annotations":[],"text":"The parents of a node are the nodes that have a directed edge pointing to it. Scanning the given edges for those where the target is `a87ff679a2`, we find only one such edge:\n\n- `d3d9446802 -> a87ff679a2`\n\nNo other edges have `a87ff679a2` as the destination. Therefore, the only parent is `d3d9446802`.\n\nFinal Answer: [d3d9446802]","type":"output_text","logprobs":null}],"role":"assistant","status":"completed","type":"message"}],"parallel_tool_calls":true,"temperature":null,"tool_choice":"auto","tools":[],"top_p":null,"background":null,"conversation":null,"max_output_tokens":null,"max_tool_calls":null,"previous_response_id":null,"prompt":null,"prompt_cache_key":null,"reasoning":null,"safety_identifier":null,"service_tier":null,"status":null,"text":null,"top_logprobs":null,"truncation":null,"usage":{"input_tokens":1323,"input_tokens_details":{"cached_tokens":0},"output_tokens":1464,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2787},"user":null},"reward":1.0,"expected_answer":"[\"d3d9446802\"]","problem_type":"parents","n_tokens":1308,"prompt_chars":2638,"f1":1.0,"parse_failed":false,"predicted_nodes":["d3d9446802"],"_ng_task_index":11,"_ng_rollout_index":0,"agent_ref":{"type":"responses_api_agents","name":"graphwalks_n3_1m_benchmark_simple_agent"}} diff --git a/responses_api_models/local_vllm_model/configs/local_vllm_model.yaml b/responses_api_models/local_vllm_model/configs/local_vllm_model.yaml new file mode 100644 index 0000000000..09ff537d96 --- /dev/null +++ b/responses_api_models/local_vllm_model/configs/local_vllm_model.yaml @@ -0,0 +1,26 @@ +# Generic local vLLM deployment. Serves the checkpoint passed via `--model ` +# (i.e. `policy_model_name`). Set parallelism to match your node before running. +# +# Example: +# gym eval run --benchmark aime24 --model-type local_vllm_model --model Qwen/Qwen3-30B-A3B-Instruct-2507 +policy_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: ${policy_model_name} + return_token_id_information: false + uses_reasoning_parser: false + + # If your model is downloaded at ~/.cache/huggingface/hub/models----, set hf_home to ~/.cache/huggingface. + hf_home: null + + vllm_serve_env_vars: {} + + # vLLM parallelism is sensitive; tensor_parallel_size * pipeline_parallel_size GPUs are used per model instance. + # Override per node, e.g. `++policy_model.responses_api_models.local_vllm_model.vllm_serve_kwargs.tensor_parallel_size=8`. + vllm_serve_kwargs: + data_parallel_size: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + trust_remote_code: true + gpu_memory_utilization: 0.9 diff --git a/tests/unit_tests/test_benchmarks.py b/tests/unit_tests/test_benchmarks.py index 02e8ecb9cb..8ac1de1806 100644 --- a/tests/unit_tests/test_benchmarks.py +++ b/tests/unit_tests/test_benchmarks.py @@ -19,7 +19,7 @@ from omegaconf import OmegaConf from yaml import safe_load -from nemo_gym.benchmarks import list_benchmarks, prepare_benchmark +from nemo_gym.cli.eval import _benchmark_extras, _fuzzy_matches, list_benchmarks, prepare_benchmark def _mock_global_config(config: dict = None): @@ -29,18 +29,119 @@ def _mock_global_config(config: dict = None): class TestListBenchmarks: def test_lists_found_benchmarks(self, capsys) -> None: - with patch("nemo_gym.benchmarks.get_global_config_dict", return_value=_mock_global_config()): + with patch("nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config()): list_benchmarks() assert "aime24" in capsys.readouterr().out def test_no_benchmarks(self, capsys) -> None: with ( - patch("nemo_gym.benchmarks.get_global_config_dict", return_value=_mock_global_config()), - patch("nemo_gym.benchmarks._load_benchmarks_from_config_paths", return_value={}), + patch("nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config()), + patch("nemo_gym.cli.eval._load_benchmarks_from_config_paths", return_value={}), ): list_benchmarks() assert "No benchmarks found" in capsys.readouterr().out + def test_json_output(self, capsys) -> None: + import json + + bench = MagicMock(agent_name="my_agent", num_repeats=4) + with ( + patch("nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config({"json": True})), + patch("nemo_gym.cli.eval._load_benchmarks_from_config_paths", return_value={"my_bench": bench}), + patch("nemo_gym.cli.eval._benchmark_extras", return_value=("math", ["math_with_judge"])), + ): + list_benchmarks() + assert json.loads(capsys.readouterr().out) == [ + {"name": "my_bench", "agent_name": "my_agent", "domain": "math", "num_repeats": 4} + ] + + def test_json_output_empty(self, capsys) -> None: + import json + + with ( + patch("nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config({"json": True})), + patch("nemo_gym.cli.eval._load_benchmarks_from_config_paths", return_value={}), + ): + list_benchmarks() + assert json.loads(capsys.readouterr().out) == [] + + +class TestFuzzyMatches: + def test_substring_matches(self) -> None: + assert _fuzzy_matches("math", "math_with_judge") + + def test_token_typo_matches(self) -> None: + # `aimee` is a near-miss for the `aime` token in `aime24`. + assert _fuzzy_matches("aimee", "aime24") + + def test_matches_against_agent_field(self) -> None: + assert _fuzzy_matches("judge", "aime24", "math_with_judge_agent") + + def test_skips_empty_fields(self) -> None: + assert not _fuzzy_matches("math", "", None) + + def test_no_match(self) -> None: + assert not _fuzzy_matches("zzznomatch", "aime24", "math_with_judge") + + +class TestBenchmarkExtras: + def test_resolves_domain_and_terms_from_real_config(self) -> None: + from nemo_gym.benchmarks import BENCHMARKS_DIR, BenchmarkConfig + + bench = BenchmarkConfig.from_config_path(BENCHMARKS_DIR / "aime24" / "config.yaml") + domain, terms = _benchmark_extras(bench) + + assert domain == "math" + # The resource server's inner name and the dataset name are recovered for search. + assert "math_with_judge" in terms + assert "aime24" in terms + + +class TestSearchBenchmarks: + # Map each benchmark name to the (domain, extra terms) its config would resolve to. + # As in the real resolver, `domain` is also included among the search terms. + EXTRAS = { + "aime24": ("math", ["aime24_math_resources_server", "math_with_judge", "aime24", "math"]), + "gpqa_diamond": ("science", ["gpqa_resources_server", "gpqa", "gpqa_diamond", "science"]), + } + + def _bench(self, key: str): + bench = MagicMock(agent_name="my_agent", num_repeats=1) + bench.config_key = key # let the patched _benchmark_extras find the right entry + return bench + + def _benchmarks(self) -> dict: + return {name: self._bench(name) for name in self.EXTRAS} + + def _run(self, query: str, benchmarks: dict, capsys) -> str: + with ( + patch("nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config({"query": query})), + patch("nemo_gym.cli.eval._load_benchmarks_from_config_paths", return_value=benchmarks), + patch("nemo_gym.cli.eval._benchmark_extras", side_effect=lambda b: self.EXTRAS[b.config_key]), + ): + list_benchmarks() + return capsys.readouterr().out + + def test_query_filters_by_name(self, capsys) -> None: + out = self._run("aime", self._benchmarks(), capsys) + assert "aime24" in out + assert "gpqa" not in out + + def test_query_matches_domain(self, capsys) -> None: + # "science" only appears via gpqa's domain, not its name/agent. + out = self._run("science", self._benchmarks(), capsys) + assert "gpqa_diamond" in out + assert "aime24" not in out + + def test_query_matches_resource_server(self, capsys) -> None: + # "judge" only appears via aime24's resource server name. + out = self._run("judge", self._benchmarks(), capsys) + assert "aime24" in out + assert "gpqa_diamond" not in out + + def test_query_no_match_message(self, capsys) -> None: + assert "No benchmarks match 'zzz'" in self._run("zzz", self._benchmarks(), capsys) + class TestPrepareBenchmark: def _make_bench_dir(self, tmp_path: Path, name: str = "fake_bench") -> tuple[Path, Path]: @@ -73,13 +174,13 @@ def test_calls_prepare(self, tmp_path: Path) -> None: with ( patch( - "nemo_gym.benchmarks.get_global_config_dict", + "nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config( {"config_paths": [str(config_path)], **safe_load(config_path.read_text())} ), ), - patch("nemo_gym.benchmarks.BENCHMARKS_DIR", bench_dir.parent), - patch("nemo_gym.benchmarks.importlib.import_module", return_value=mock_module), + patch("nemo_gym.cli.eval.BENCHMARKS_DIR", bench_dir.parent), + patch("nemo_gym.cli.eval.importlib.import_module", return_value=mock_module), ): prepare_benchmark() mock_module.prepare.assert_called_once() @@ -90,12 +191,12 @@ def test_missing_prepare_py(self, tmp_path: Path) -> None: with ( patch( - "nemo_gym.benchmarks.get_global_config_dict", + "nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config( {"config_paths": [str(config_path)], **safe_load(config_path.read_text())} ), ), - patch("nemo_gym.benchmarks.BENCHMARKS_DIR", bench_dir.parent), + patch("nemo_gym.cli.eval.BENCHMARKS_DIR", bench_dir.parent), ): with pytest.raises(RuntimeError, match="The following benchmarks are missing a valid prepare script"): prepare_benchmark() @@ -107,13 +208,13 @@ def test_missing_prepare_function(self, tmp_path: Path) -> None: with ( patch( - "nemo_gym.benchmarks.get_global_config_dict", + "nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config( {"config_paths": [str(config_path)], **safe_load(config_path.read_text())} ), ), - patch("nemo_gym.benchmarks.BENCHMARKS_DIR", bench_dir.parent), - patch("nemo_gym.benchmarks.importlib.import_module", return_value=mock_module), + patch("nemo_gym.cli.eval.BENCHMARKS_DIR", bench_dir.parent), + patch("nemo_gym.cli.eval.importlib.import_module", return_value=mock_module), ): with pytest.raises( AssertionError, @@ -124,10 +225,10 @@ def test_missing_prepare_function(self, tmp_path: Path) -> None: def test_no_benchmark_in_config_paths(self) -> None: with ( patch( - "nemo_gym.benchmarks.get_global_config_dict", + "nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config({"config_paths": ["resources_servers/foo/configs/foo.yaml"]}), ), - patch("nemo_gym.benchmarks._load_benchmarks_from_config_paths", return_value={}), + patch("nemo_gym.cli.eval._load_benchmarks_from_config_paths", return_value={}), ): with pytest.raises(AssertionError, match="No benchmark config found in config_paths"): prepare_benchmark() @@ -141,7 +242,7 @@ def test_caching_sanity(self, tmp_path: Path) -> None: with ( patch( - "nemo_gym.benchmarks.get_global_config_dict", + "nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config( { "use_cached_prepared_benchmarks": True, @@ -150,8 +251,8 @@ def test_caching_sanity(self, tmp_path: Path) -> None: } ), ), - patch("nemo_gym.benchmarks.BENCHMARKS_DIR", bench_dir.parent), - patch("nemo_gym.benchmarks.importlib.import_module", return_value=mock_module), + patch("nemo_gym.cli.eval.BENCHMARKS_DIR", bench_dir.parent), + patch("nemo_gym.cli.eval.importlib.import_module", return_value=mock_module), ): prepare_benchmark() diff --git a/tests/unit_tests/test_cli.py b/tests/unit_tests/test_cli.py index 74f7fce578..7053b7c144 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -15,6 +15,7 @@ import shutil import sys import tomllib +import warnings from importlib import import_module from io import StringIO from pathlib import Path @@ -22,19 +23,20 @@ from unittest.mock import MagicMock, patch from omegaconf import OmegaConf -from pytest import MonkeyPatch, raises +from pytest import MonkeyPatch import nemo_gym.global_config from nemo_gym import PARENT_DIR -from nemo_gym.cli import ( +from nemo_gym.cli.env import ( _FORCE_KILL_REAP_TIMEOUT_SEC, _GRACEFUL_SHUTDOWN_TIMEOUT_SEC, RunConfig, RunHelper, - display_help, + _run_module_tests_all, init_resources_server, ) -from nemo_gym.config_types import ResourcesServerInstanceConfig +from nemo_gym.cli.general import display_help_legacy +from nemo_gym.config_types import DatasetConfig, ResourcesServerInstanceConfig # TODO: Eventually we want to add more tests to ensure that the CLI flows do not break @@ -42,32 +44,16 @@ class TestCLI: def test_sanity(self) -> None: RunConfig(entrypoint="", name="") - def test_pyproject_scripts(self) -> None: + def test_pyproject_scripts_are_importable(self) -> None: + """Every console-script entry point must resolve to an importable callable.""" pyproject_path = PARENT_DIR / "pyproject.toml" with pyproject_path.open("rb") as f: pyproject_data = tomllib.load(f) - project_scripts = pyproject_data["project"]["scripts"] - - for script_name, import_path in project_scripts.items(): - # Dedupe `nemo_gym_*` from `ng_*` commands - if not script_name.startswith("ng_"): - continue - - # We only test `+h=true` and not `+help=true` - print(f"Running `{script_name} +h=true`") - + for script_name, import_path in pyproject_data["project"]["scripts"].items(): module, fn = import_path.split(":") - fn = getattr(import_module(module), fn) - - with MonkeyPatch.context() as mp: - mp.setattr(nemo_gym.global_config, "_GLOBAL_CONFIG_DICT", OmegaConf.create({"h": True})) - - text_trap = StringIO() - mp.setattr(sys, "stdout", text_trap) - - with raises(SystemExit): - fn() + target = getattr(import_module(module), fn) + assert callable(target), f"{script_name} -> {import_path} is not callable" def test_display_help_discovers_scripts(self) -> None: with MonkeyPatch.context() as mp: @@ -76,7 +62,7 @@ def test_display_help_discovers_scripts(self) -> None: text_trap = StringIO() mp.setattr(sys, "stdout", text_trap) - display_help() + display_help_legacy() output = text_trap.getvalue() assert "ng_help" in output @@ -137,11 +123,63 @@ def test_init_resources_server_includes_domain(self) -> None: # This should not raise an assertion error about missing domain instance_config = ResourcesServerInstanceConfig.model_validate(full_config_dict) assert instance_config is not None + + # The generated config carries inline field documentation (friction #7). + config_text = config_file.read_text() + assert "# Task category" in config_text + assert "policy_model" in config_text and "magic name" in config_text + + # The scaffold emits the canonical `source:` block, not the deprecated + # `gitlab_identifier`, so a fresh server starts on the recommended schema. + datasets = config_dict[f"{server_name}_simple_agent"]["responses_api_agents"]["simple_agent"][ + "datasets" + ] + train = next(d for d in datasets if d["name"] == "train") + assert "gitlab_identifier" not in train + assert train["source"]["type"] == "gitlab" + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + DatasetConfig.model_validate(OmegaConf.to_container(train, resolve=True)) finally: # Clean up the test server directory if server_path.exists(): shutil.rmtree(server_path) + def test_init_resources_server_judge_template(self) -> None: + """`+template=judge` scaffolds an auxiliary-model verifier (config + judge app/tests).""" + + server_name = "test_cli_judge_server" + entrypoint = f"resources_servers/{server_name}" + server_path = Path(entrypoint).resolve() + + if server_path.exists(): + shutil.rmtree(server_path) + + try: + with MonkeyPatch.context() as mp: + mp.setattr( + nemo_gym.global_config, + "_GLOBAL_CONFIG_DICT", + OmegaConf.create({"entrypoint": entrypoint, "template": "judge"}), + ) + + init_resources_server() + + config_dict = OmegaConf.load(server_path / "configs" / f"{server_name}.yaml") + server_config = config_dict[f"{server_name}_resources_server"]["resources_servers"][server_name] + # The judge template wires the auxiliary model into the resources server block. + assert server_config["judge_model_server"]["name"] == "judge_model" + assert "judge_responses_create_params" in server_config + + # The judge app + its matching test template are scaffolded (not the basic ones). + app_text = (server_path / "app.py").read_text() + assert "judge_model_server" in app_text and "_score_with_judge" in app_text + test_text = (server_path / "tests" / "test_app.py").read_text() + assert "judge_responses_create_params" in test_text + finally: + if server_path.exists(): + shutil.rmtree(server_path) + def test_run_helper_prefers_cwd_server_over_install(self, tmp_path: Path) -> None: """ng_run should use a local CWD server dir instead of the installed one.""" # Create a fake local server dir in tmp_path (simulates user's own resources_servers/) @@ -227,3 +265,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 diff --git a/tests/unit_tests/test_cli_legacy.py b/tests/unit_tests/test_cli_legacy.py new file mode 100644 index 0000000000..cad72c0a61 --- /dev/null +++ b/tests/unit_tests/test_cli_legacy.py @@ -0,0 +1,55 @@ +# 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. + +import sys +import tomllib + +import pytest +from pytest import MonkeyPatch + +import nemo_gym.cli.legacy as legacy +from nemo_gym import PARENT_DIR + + +def _legacy_scripts() -> list[tuple[str, str]]: + """All console scripts whose name is a legacy ng_*/nemo_gym_* alias.""" + with (PARENT_DIR / "pyproject.toml").open("rb") as f: + scripts = tomllib.load(f)["project"]["scripts"] + return [(name, target) for name, target in scripts.items() if name.startswith(("ng_", "nemo_gym_"))] + + +LEGACY_SCRIPTS = _legacy_scripts() + + +class TestLegacyDeprecation: + """Remove these tests once the legacy commands are removed.""" + + def test_legacy_scripts_were_discovered(self) -> None: + # Guard so the parametrized test below can't pass vacuously if discovery breaks. + assert len(LEGACY_SCRIPTS) > 1 + + @pytest.mark.parametrize("name, target", LEGACY_SCRIPTS) + def test_legacy_command_shows_deprecation(self, monkeypatch: MonkeyPatch, capsys, name: str, target: str) -> None: + # Every legacy alias must route through the shim, which prints a deprecation notice and keeps working. + assert target == "nemo_gym.cli.legacy:main", f"{name} should route through the legacy shim" + + # Stub the actual execution paths so nothing real runs. + monkeypatch.setattr(legacy, "gym_main", lambda: None) + monkeypatch.setattr(legacy, "dispatch", lambda *a, **k: None) + monkeypatch.setattr(sys, "argv", [name]) + + legacy.main() + + assert "deprecated" in capsys.readouterr().err diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py new file mode 100644 index 0000000000..9f1f5cb7e4 --- /dev/null +++ b/tests/unit_tests/test_cli_main.py @@ -0,0 +1,859 @@ +# 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. +import logging +import sys + +import pytest +from pytest import MonkeyPatch + +import nemo_gym.cli.main as cli_main +import nemo_gym.global_config as gc +from nemo_gym import WORKING_DIR +from nemo_gym.cli.main import main +from nemo_gym.global_config import NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME + + +def _dispatch_for(monkeypatch: MonkeyPatch, argv: list[str]) -> tuple[str, list[str]]: + """Run the gym router for `argv` and return the (target, overrides) handed to dispatch.""" + captured: dict = {} + + def fake_dispatch(target: str, overrides: list[str]) -> None: + captured["target"] = target + captured["overrides"] = overrides + + monkeypatch.setattr(cli_main, "dispatch", fake_dispatch) + monkeypatch.setattr(sys, "argv", ["gym", *argv]) + main() + return captured["target"], captured["overrides"] + + +def _split_overrides(overrides: list[str]) -> tuple[set[str], set[str]]: + """Split overrides into (config paths, other overrides) as sets, so tests never assert ordering.""" + prefix = "+config_paths=[" + config_tokens = [o for o in overrides if o.startswith(prefix) and o.endswith("]")] + assert len(config_tokens) <= 1 # --config and the asset selectors coalesce into a single token + paths = set(config_tokens[0][len(prefix) : -1].split(",")) if config_tokens else set() + others = {o for o in overrides if o not in config_tokens} + return paths, others + + +# `gym ` -> the legacy ng_ function it dispatches to, for the config-accepting commands. +CONFIG_COMMANDS = [ + (["env", "run"], "nemo_gym.cli.env:run"), + (["env", "resolve"], "nemo_gym.cli.env:dump_config"), + (["eval", "prepare"], "nemo_gym.cli.eval:prepare_benchmark"), + (["eval", "aggregate"], "nemo_gym.cli.eval:aggregate_rollouts"), + (["eval", "run"], "nemo_gym.cli.eval:e2e_rollout_collection"), + (["dataset", "collate"], "nemo_gym.cli.dataset:prepare_data"), +] + + +class TestConfigFlag: + @pytest.mark.parametrize("command, expected_target", CONFIG_COMMANDS) + def test_config_becomes_config_paths(self, monkeypatch: MonkeyPatch, command, expected_target) -> None: + """`gym --config X` dispatches to ng_ with +config_paths=[X].""" + target, overrides = _dispatch_for(monkeypatch, [*command, "--config", "my.yaml"]) + assert target == expected_target + assert overrides == ["+config_paths=[my.yaml]"] + + def test_repeated_config_joined_into_one_list(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["env", "run", "--config", "a.yaml", "--config", "b.yaml"]) + + # We have this set of asserts to avoid asserting configs order in the string + assert len(overrides) == 1 + override = overrides[0] + assert override.startswith("+config_paths=[") + assert override.endswith("]") + assert "a.yaml" in override + assert "b.yaml" in override + + def test_config_is_prepended_before_passthrough_overrides(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["env", "run", "--config", "a.yaml", "+foo=bar"]) + assert len(overrides) == 2 + assert "+config_paths=[a.yaml]" in overrides + assert "+foo=bar" in overrides + + def test_without_config_no_config_paths_added(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["env", "run", "+foo=bar"]) + assert overrides == ["+foo=bar"] + + def test_config_rejected_on_non_config_command(self, monkeypatch: MonkeyPatch) -> None: + # `dataset rm` does not declare --config, so the router must reject it rather than leak it downstream. + monkeypatch.setattr(cli_main, "dispatch", lambda target, overrides: None) + monkeypatch.setattr(sys, "argv", ["gym", "dataset", "rm", "--config", "x.yaml"]) + with pytest.raises(SystemExit): + main() + + +class TestStorageFlag: + @pytest.mark.parametrize( + "argv, expected_target", + [ + (["dataset", "upload"], "nemo_gym.cli.dataset:upload_jsonl_dataset_to_hf_cli"), + (["dataset", "upload", "--storage", "hf"], "nemo_gym.cli.dataset:upload_jsonl_dataset_to_hf_cli"), + (["dataset", "upload", "--storage", "gitlab"], "nemo_gym.cli.dataset:upload_jsonl_dataset_cli"), + (["dataset", "download"], "nemo_gym.cli.dataset:download_jsonl_dataset_from_hf_cli"), + (["dataset", "download", "--storage", "hf"], "nemo_gym.cli.dataset:download_jsonl_dataset_from_hf_cli"), + (["dataset", "download", "--storage", "gitlab"], "nemo_gym.cli.dataset:download_jsonl_dataset_cli"), + ], + ) + def test_storage_selects_backend(self, monkeypatch: MonkeyPatch, argv, expected_target) -> None: + target, _ = _dispatch_for(monkeypatch, argv) + assert target == expected_target + + def test_storage_does_not_leak_into_overrides(self, monkeypatch: MonkeyPatch) -> None: + # --storage only selects the target; it must not appear in the Hydra overrides. + _, overrides = _dispatch_for(monkeypatch, ["dataset", "upload", "--storage", "gitlab", "+foo=bar"]) + assert overrides == ["+foo=bar"] + + def test_invalid_storage_value_is_rejected(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr(sys, "argv", ["gym", "dataset", "upload", "--storage", "s3"]) + with pytest.raises(SystemExit): + main() + + +class TestEvalRunFlags: + @pytest.mark.parametrize( + "flag_argv, expected_override", + [ + (["--agent", "my_agent"], "+agent_name=my_agent"), + (["-a", "my_agent"], "+agent_name=my_agent"), + (["--input", "in.jsonl"], "+input_jsonl_fpath=in.jsonl"), + (["-i", "in.jsonl"], "+input_jsonl_fpath=in.jsonl"), + (["--output", "out.jsonl"], "+output_jsonl_fpath=out.jsonl"), + (["-o", "out.jsonl"], "+output_jsonl_fpath=out.jsonl"), + (["--limit", "1024"], "+limit=1024"), + (["--num-repeats", "4"], "+num_repeats=4"), + (["--concurrency", "10"], "+num_samples_in_parallel=10"), + (["--prompt-config", "p.yaml"], "+prompt_config=p.yaml"), + (["--split", "benchmark"], "+split=benchmark"), + (["--model", "openai/gpt-oss-120b"], "+policy_model_name=openai/gpt-oss-120b"), + (["-m", "openai/gpt-oss-120b"], "+policy_model_name=openai/gpt-oss-120b"), + (["--model-url", "http://0.0.0.0:10240/v1"], "+policy_base_url=http://0.0.0.0:10240/v1"), + (["--model-api-key", "sk-your-api-key"], "+policy_api_key=sk-your-api-key"), + (["--temperature", "1.0"], "+responses_create_params.temperature=1.0"), + (["--top-p", "1.0"], "+responses_create_params.top_p=1.0"), + (["--max-output-tokens", "4096"], "+responses_create_params.max_output_tokens=4096"), + (["--resume"], "+resume_from_cache=true"), + ], + ) + def test_flag_maps_to_single_override(self, monkeypatch: MonkeyPatch, flag_argv, expected_override) -> None: + _, overrides = _dispatch_for(monkeypatch, ["eval", "run", *flag_argv]) + assert overrides == [expected_override] + + def test_unset_flags_contribute_nothing(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["eval", "run", "--agent", "x"]) + assert overrides == ["+agent_name=x"] + + def test_default_dispatches_e2e(self, monkeypatch: MonkeyPatch) -> None: + target, _ = _dispatch_for(monkeypatch, ["eval", "run"]) + assert target == "nemo_gym.cli.eval:e2e_rollout_collection" + + def test_no_serve_dispatches_collect_without_override(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for(monkeypatch, ["eval", "run", "--no-serve"]) + assert target == "nemo_gym.cli.eval:collect_rollouts" + assert overrides == [] + + def test_readme_collect_rollouts_example(self, monkeypatch: MonkeyPatch) -> None: + # From resources_servers/my_weather_tool README: + # ng_collect_rollouts +agent_name=... +input_jsonl_fpath=... +output_jsonl_fpath=... +limit=1024 +num_repeats=1 + target, overrides = _dispatch_for( + monkeypatch, + [ + "eval", + "run", + "--no-serve", + "--agent", + "my_weather_tool_simple_agent", + "--input", + "resources_servers/my_weather_tool/data/example.jsonl", + "--output", + "resources_servers/my_weather_tool/data/example_rollouts.jsonl", + "--limit", + "1024", + "--num-repeats", + "1", + ], + ) + assert target == "nemo_gym.cli.eval:collect_rollouts" + assert set(overrides) == { + "+agent_name=my_weather_tool_simple_agent", + "+input_jsonl_fpath=resources_servers/my_weather_tool/data/example.jsonl", + "+output_jsonl_fpath=resources_servers/my_weather_tool/data/example_rollouts.jsonl", + "+limit=1024", + "+num_repeats=1", + } + + def test_readme_model_and_sampling_example(self, monkeypatch: MonkeyPatch) -> None: + # From the gpt-oss eval example: ++policy_* and ++responses_create_params.* overrides. + _, overrides = _dispatch_for( + monkeypatch, + [ + "eval", + "run", + "--model", + "openai/gpt-oss-120b", + "--model-url", + "http://0.0.0.0:10240/v1", + "--model-api-key", + "dummy_key", + "--temperature", + "1.0", + "--top-p", + "1.0", + ], + ) + assert set(overrides) == { + "+policy_model_name=openai/gpt-oss-120b", + "+policy_base_url=http://0.0.0.0:10240/v1", + "+policy_api_key=dummy_key", + "+responses_create_params.temperature=1.0", + "+responses_create_params.top_p=1.0", + } + + def test_flags_compose_with_config_and_passthrough(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, + [ + "eval", + "run", + "--no-serve", + "--config", + "b.yaml", + "--agent", + "a", + "+responses_create_params.tool_choice=auto", + ], + ) + assert target == "nemo_gym.cli.eval:collect_rollouts" + assert "+config_paths=[b.yaml]" in overrides + assert "+agent_name=a" in overrides + assert "+responses_create_params.tool_choice=auto" in overrides # unknown +override passes through + + +class TestEnvTestResourceServerFlag: + def test_no_resource_server_runs_all(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for(monkeypatch, ["env", "test"]) + assert target == "nemo_gym.cli.env:test_all" + assert overrides == [] + + def test_resource_server_name_translates_to_entrypoint(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for(monkeypatch, ["env", "test", "--resource-server", "gpqa"]) + assert target == "nemo_gym.cli.env:test" + assert overrides == ["+entrypoint=resources_servers/gpqa"] + + def test_direct_entrypoint_override_also_runs_single(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for(monkeypatch, ["env", "test", "+entrypoint=resources_servers/gpqa"]) + assert target == "nemo_gym.cli.env:test" + assert overrides == ["+entrypoint=resources_servers/gpqa"] + + +class TestDatasetFlags: + def test_upload_hf_default(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, + ["dataset", "upload", "-i", "data/train.jsonl", "--name", "my_ds", "--split", "train", "--create-pr"], + ) + assert target == "nemo_gym.cli.dataset:upload_jsonl_dataset_to_hf_cli" + assert set(overrides) == { + "+input_jsonl_fpath=data/train.jsonl", + "+dataset_name=my_ds", + "+split=train", + "+create_pr=true", + } + + def test_upload_gitlab(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, + [ + "dataset", + "upload", + "--storage", + "gitlab", + "-i", + "data/train.jsonl", + "--name", + "my_ds", + "--revision", + "0.0.1", + ], + ) + assert target == "nemo_gym.cli.dataset:upload_jsonl_dataset_cli" + overrides.remove( + "+revision=0.0.1" + ) # we set both version and revision because GitLab and HF use different keys + assert set(overrides) == { + "+input_jsonl_fpath=data/train.jsonl", + "+dataset_name=my_ds", + "+version=0.0.1", + } + + def test_download_hf_default(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, + [ + "dataset", + "download", + "--repo-id", + "org/my_ds", + "--artifact", + "train.jsonl", + "--output-dir", + "./data", + "--split", + "train", + ], + ) + assert target == "nemo_gym.cli.dataset:download_jsonl_dataset_from_hf_cli" + assert set(overrides) == { + "+repo_id=org/my_ds", + "+artifact_fpath=train.jsonl", + "+output_dirpath=./data", + "+split=train", + } + + def test_download_gitlab(self, monkeypatch: MonkeyPatch) -> None: + # On download, --revision is GitLab-only and maps to +version (HF download has no revision field). + target, overrides = _dispatch_for( + monkeypatch, + [ + "dataset", + "download", + "--storage", + "gitlab", + "--name", + "my_ds", + "--revision", + "0.0.1", + "--artifact", + "train.jsonl", + "-o", + "./train.jsonl", + ], + ) + assert target == "nemo_gym.cli.dataset:download_jsonl_dataset_cli" + assert set(overrides) == { + "+dataset_name=my_ds", + "+version=0.0.1", + "+artifact_fpath=train.jsonl", + "+output_fpath=./train.jsonl", + } + + def test_rm(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for(monkeypatch, ["dataset", "rm", "--name", "my_ds"]) + assert target == "nemo_gym.cli.dataset:delete_jsonl_dataset_from_gitlab_cli" + assert overrides == ["+dataset_name=my_ds"] + + def test_migrate_revision_maps_to_hf_revision(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, + ["dataset", "migrate", "-i", "data/train.jsonl", "--name", "my_ds", "--revision", "r1", "--create-pr"], + ) + assert target == "nemo_gym.cli.dataset:upload_jsonl_dataset_to_hf_and_delete_gitlab_cli" + assert set(overrides) == { + "+input_jsonl_fpath=data/train.jsonl", + "+dataset_name=my_ds", + "+revision=r1", + "+create_pr=true", + } + + def test_render(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, ["dataset", "render", "-i", "raw.jsonl", "--prompt-config", "p.yaml", "-o", "out.jsonl"] + ) + assert target == "nemo_gym.cli.dataset:materialize_prompts_cli" + assert set(overrides) == { + "+input_jsonl_fpath=raw.jsonl", + "+prompt_config=p.yaml", + "+output_jsonl_fpath=out.jsonl", + } + + def test_collate(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, + [ + "dataset", + "collate", + "--config", + "c.yaml", + "--mode", + "train_preparation", + "--output-dir", + "./prep", + "--download", + ], + ) + assert target == "nemo_gym.cli.dataset:prepare_data" + assert set(overrides) == { + "+config_paths=[c.yaml]", + "+mode=train_preparation", + "+output_dirpath=./prep", + "+should_download=true", + } + + def test_bool_flags_omitted_when_unset(self, monkeypatch: MonkeyPatch) -> None: + # --create-pr not passed -> no +create_pr override leaks in. + _, overrides = _dispatch_for(monkeypatch, ["dataset", "upload", "--name", "my_ds"]) + assert overrides == ["+dataset_name=my_ds"] + + def test_collate_mode_rejects_invalid_choice(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr(sys, "argv", ["gym", "dataset", "collate", "--mode", "bogus"]) + with pytest.raises(SystemExit): + main() + + +class TestEvalAggregateFlags: + def test_output_flag(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for(monkeypatch, ["eval", "aggregate", "-o", "out.jsonl"]) + assert target == "nemo_gym.cli.eval:aggregate_rollouts" + assert overrides == ["+output_jsonl_fpath=out.jsonl"] + + +class TestEvalProfileFlags: + def test_profile_flags(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, ["eval", "profile", "--inputs", "in.jsonl", "--rollouts", "r.jsonl"] + ) + assert target == "nemo_gym.cli.eval:reward_profile" + assert set(overrides) == { + "+materialized_inputs_jsonl_fpath=in.jsonl", + "+rollouts_jsonl_fpath=r.jsonl", + } + + def test_profile_does_not_accept_config(self, monkeypatch: MonkeyPatch) -> None: + # reward_profile reads file paths, not config_paths, so --config is not offered and is rejected. + monkeypatch.setattr(cli_main, "dispatch", lambda target, overrides: None) + monkeypatch.setattr(sys, "argv", ["gym", "eval", "profile", "--config", "x.yaml"]) + with pytest.raises(SystemExit): + main() + + +class TestEnvRunFlags: + def test_model_flags(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for( + monkeypatch, + [ + "env", + "run", + "--config", + "c.yaml", + "--model", + "gpt", + "--model-url", + "http://x", + "--model-api-key", + "k", + ], + ) + assert target == "nemo_gym.cli.env:run" + assert set(overrides) == { + "+config_paths=[c.yaml]", + "+policy_model_name=gpt", + "+policy_base_url=http://x", + "+policy_api_key=k", + } + + +class TestModelFlag: + """`--model` is the single served-model identifier (name, HF id, or local checkpoint path) for any backend.""" + + def test_local_vllm_deployment_flow(self, monkeypatch: MonkeyPatch) -> None: + # The deployment invocation: select the local vLLM server type and pass the checkpoint to serve via --model. + _, overrides = _dispatch_for( + monkeypatch, + ["eval", "run", "--model-type", "local_vllm_model", "--model", "Qwen/Qwen3-8B"], + ) + paths, others = _split_overrides(overrides) + assert paths == {str(WORKING_DIR / "responses_api_models/local_vllm_model/configs/local_vllm_model.yaml")} + assert others == {"+policy_model_name=Qwen/Qwen3-8B"} + + def test_short_alias_on_env_run(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["env", "run", "-m", "/ckpt/path"]) + assert overrides == ["+policy_model_name=/ckpt/path"] + + +class TestEnvInitFlags: + def test_resource_server_translates_to_entrypoint(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for(monkeypatch, ["env", "init", "--resource-server", "my_server"]) + assert target == "nemo_gym.cli.env:init_resources_server" + assert overrides == ["+entrypoint=resources_servers/my_server"] + + +class TestEnvPackagesFlags: + def test_flags(self, monkeypatch: MonkeyPatch) -> None: + target, overrides = _dispatch_for(monkeypatch, ["env", "packages", "--resource-server", "gpqa", "--outdated"]) + assert target == "nemo_gym.cli.env:pip_list" + assert set(overrides) == { + "+entrypoint=resources_servers/gpqa", + "+outdated=true", + } + + +class TestJsonFlag: + @pytest.mark.parametrize( + "argv, expected_target", + [ + (["list", "benchmarks", "--json"], "nemo_gym.cli.eval:list_benchmarks"), + (["env", "status", "--json"], "nemo_gym.cli.env:status"), + ], + ) + def test_json_becomes_config_override(self, monkeypatch: MonkeyPatch, argv, expected_target) -> None: + # Reporting commands surface --json as the reserved `json` config key, read centrally by cli.output.emit. + target, overrides = _dispatch_for(monkeypatch, argv) + assert target == expected_target + assert overrides == ["+json=true"] + + def test_no_json_no_override(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["list", "benchmarks"]) + assert overrides == [] + + +class TestSearch: + def test_search_routes_to_list_with_query(self, monkeypatch: MonkeyPatch) -> None: + # `gym search ` reuses the benchmarks listing, passing the query as the `query` config key. + target, overrides = _dispatch_for(monkeypatch, ["search", "math"]) + assert target == "nemo_gym.cli.eval:list_benchmarks" + assert overrides == ["+query=math"] + + def test_search_json(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["search", "math", "--json"]) + assert set(overrides) == {"+query=math", "+json=true"} + + def test_version_json_dispatches_with_override(self, monkeypatch: MonkeyPatch) -> None: + # `gym --version --json` is the top-level path; it still forwards +json=true to the version command. + target, overrides = _dispatch_for(monkeypatch, ["--version", "--json"]) + assert target == "nemo_gym.cli.general:version" + assert overrides == ["+json=true"] + + def test_env_packages_json_maps_to_uv_format(self, monkeypatch: MonkeyPatch) -> None: + # env packages delegates to `uv pip list`, so --json maps onto its own --format=json rather than +json=true. + target, overrides = _dispatch_for(monkeypatch, ["env", "packages", "--resource-server", "mcqa", "--json"]) + assert target == "nemo_gym.cli.env:pip_list" + assert set(overrides) == {"+entrypoint=resources_servers/mcqa", "+format=json"} + + def test_env_packages_without_json(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["env", "packages", "--resource-server", "mcqa"]) + assert overrides == ["+entrypoint=resources_servers/mcqa"] + + +class TestVerboseFlag: + @pytest.mark.parametrize("flag", ["-v", "--verbose"]) + def test_verbose_injects_config_override(self, monkeypatch: MonkeyPatch, flag: str) -> None: + # --verbose flows through the config (so it reaches servers), not just the local logger. + _, overrides = _dispatch_for(monkeypatch, ["env", "status", flag]) + assert overrides == ["+verbose=true"] + + def test_no_verbose_no_override(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["env", "status"]) + assert overrides == [] + + def test_verbose_prepended_before_other_overrides(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["eval", "run", "--verbose", "--agent", "a", "+x=1"]) + assert "+verbose=true" in overrides + assert "+agent_name=a" in overrides + assert "+x=1" in overrides + + def test_config_verbose_sets_debug_on_load(self, monkeypatch: MonkeyPatch) -> None: + # The server-side path: a config carrying `verbose` (forwarded via env var) raises the log level. + monkeypatch.setattr(gc, "_GLOBAL_CONFIG_DICT", None) + monkeypatch.setenv(NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME, "verbose: true\nsome_server: {}\n") + root = logging.getLogger() + original = root.level + try: + root.setLevel(logging.WARNING) + gc.get_global_config_dict() + assert root.level == logging.DEBUG + finally: + root.setLevel(original) + + def test_config_without_verbose_keeps_level(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr(gc, "_GLOBAL_CONFIG_DICT", None) + monkeypatch.setenv(NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME, "some_server: {}\n") + root = logging.getLogger() + original = root.level + try: + root.setLevel(logging.WARNING) + gc.get_global_config_dict() + assert root.level == logging.WARNING + finally: + root.setLevel(original) + + +class TestAssetSelectors: + """Named selectors (--benchmark, --resource-server, --model-type) that resolve a name to a default config path. + + Each example mirrors a real invocation from the docs/READMEs, so the sugar stays faithful to the documented + config paths it replaces. The legacy `+config_paths=[...]` form each one is derived from is cited inline. + """ + + @pytest.mark.parametrize( + "argv, expected_config", + [ + # benchmarks/gsm8k/README.md: ng_prepare_benchmark "+config_paths=[benchmarks/gsm8k/config.yaml]" + (["eval", "prepare", "--benchmark", "gsm8k"], "benchmarks/gsm8k/config.yaml"), + # benchmarks/aime25-x/README.md: ng_prepare_benchmark "+config_paths=[benchmarks/aime25-x/config.yaml]" + (["eval", "prepare", "--benchmark", "aime25-x"], "benchmarks/aime25-x/config.yaml"), + # README.md / quickstart.mdx: resources_servers/mcqa/configs/mcqa.yaml + (["env", "run", "--resource-server", "mcqa"], "resources_servers/mcqa/configs/mcqa.yaml"), + # model-server/vllm.mdx: resources_servers/example_multi_step/configs/example_multi_step.yaml + ( + ["env", "run", "--resource-server", "example_multi_step"], + "resources_servers/example_multi_step/configs/example_multi_step.yaml", + ), + # README.md / quickstart.mdx: responses_api_models/openai_model/configs/openai_model.yaml + ( + ["env", "run", "--model-type", "openai_model"], + "responses_api_models/openai_model/configs/openai_model.yaml", + ), + # model-server/vllm.mdx: responses_api_models/vllm_model/configs/vllm_model.yaml + (["env", "run", "--model-type", "vllm_model"], "responses_api_models/vllm_model/configs/vllm_model.yaml"), + ], + ) + def test_name_resolves_to_config_path(self, monkeypatch: MonkeyPatch, argv, expected_config) -> None: + _, overrides = _dispatch_for(monkeypatch, argv) + assert overrides == [f"+config_paths=[{WORKING_DIR / (expected_config)}]"] + + def test_quickstart_resource_server_plus_model(self, monkeypatch: MonkeyPatch) -> None: + # README.md / quickstart.mdx: + # ng_run "+config_paths=[resources_servers/mcqa/configs/mcqa.yaml, + # responses_api_models/openai_model/configs/openai_model.yaml]" + target, overrides = _dispatch_for( + monkeypatch, ["env", "run", "--resource-server", "mcqa", "--model-type", "openai_model"] + ) + assert target == "nemo_gym.cli.env:run" + paths, others = _split_overrides(overrides) + assert paths == { + str(WORKING_DIR / "resources_servers/mcqa/configs/mcqa.yaml"), + str(WORKING_DIR / "responses_api_models/openai_model/configs/openai_model.yaml"), + } + assert others == set() + + def test_gpqa_benchmark_plus_model(self, monkeypatch: MonkeyPatch) -> None: + # benchmarks/gpqa/README.md: + # ng_run "+config_paths=[benchmarks/gpqa/config.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml]" + _, overrides = _dispatch_for(monkeypatch, ["eval", "run", "--benchmark", "gpqa", "--model-type", "vllm_model"]) + paths, others = _split_overrides(overrides) + assert paths == { + str(WORKING_DIR / "benchmarks/gpqa/config.yaml"), + str(WORKING_DIR / "responses_api_models/vllm_model/configs/vllm_model.yaml"), + } + assert others == set() + + def test_cli_reference_e2e_rollout_example(self, monkeypatch: MonkeyPatch) -> None: + # fern .../reference/cli-commands.mdx ng_e2e_collect_rollouts example: + # config_paths="responses_api_models/openai_model/configs/openai_model.yaml, + # resources_servers/math_with_judge/configs/math_with_judge.yaml" + # ng_e2e_collect_rollouts "+config_paths=[$config_paths]" + # ++output_jsonl_fpath=results/test_e2e_rollout_collection/aime24.jsonl ++split=validation + target, overrides = _dispatch_for( + monkeypatch, + [ + "eval", + "run", + "--resource-server", + "math_with_judge", + "--model-type", + "openai_model", + "--output", + "results/test_e2e_rollout_collection/aime24.jsonl", + "--split", + "validation", + ], + ) + assert target == "nemo_gym.cli.eval:e2e_rollout_collection" + paths, others = _split_overrides(overrides) + assert paths == { + str(WORKING_DIR / "resources_servers/math_with_judge/configs/math_with_judge.yaml"), + str(WORKING_DIR / "responses_api_models/openai_model/configs/openai_model.yaml"), + } + assert others == { + "+output_jsonl_fpath=results/test_e2e_rollout_collection/aime24.jsonl", + "+split=validation", + } + + def test_cli_reference_prepare_data_example(self, monkeypatch: MonkeyPatch) -> None: + # fern .../reference/cli-commands.mdx ng_prepare_data example: + # config_paths includes resources_servers/example_multi_step/configs/example_multi_step.yaml + # ng_prepare_data "+config_paths=[...]" +output_dirpath=data/example_multi_step +mode=example_validation + target, overrides = _dispatch_for( + monkeypatch, + [ + "dataset", + "collate", + "--resource-server", + "example_multi_step", + "--mode", + "example_validation", + "--output-dir", + "data/example_multi_step", + ], + ) + assert target == "nemo_gym.cli.dataset:prepare_data" + paths, others = _split_overrides(overrides) + assert paths == {str(WORKING_DIR / "resources_servers/example_multi_step/configs/example_multi_step.yaml")} + assert others == { + "+mode=example_validation", + "+output_dirpath=data/example_multi_step", + } + + def test_resource_server_flavor_syntax(self, monkeypatch: MonkeyPatch) -> None: + # `/` picks a named config inside the server's configs/ dir; math_with_judge ships several + # flavoured configs (see reference/faq.mdx, which pairs a math_with_judge dataset flavour for profiling). + _, overrides = _dispatch_for(monkeypatch, ["eval", "run", "--resource-server", "math_with_judge/dapo17k"]) + assert overrides == [ + f"+config_paths=[{WORKING_DIR / 'resources_servers/math_with_judge/configs/dapo17k.yaml'}]" + ] + + def test_benchmark_flavor_syntax(self, monkeypatch: MonkeyPatch) -> None: + # Benchmarks are flavoured too: flavor is a sibling `.yaml` (no configs/ dir), default `config`. + # e.g. benchmarks/finance_sec_search ships config_web_search.yaml alongside the default config.yaml. + _, overrides = _dispatch_for( + monkeypatch, ["eval", "prepare", "--benchmark", "finance_sec_search/config_web_search"] + ) + assert overrides == [f"+config_paths=[{WORKING_DIR / 'benchmarks/finance_sec_search/config_web_search.yaml'}]"] + + def test_selectors_merge_into_single_config_paths(self, monkeypatch: MonkeyPatch) -> None: + # --config and multiple asset selectors all feed one +config_paths list (Hydra rejects duplicates). + # _split_overrides asserts they coalesce into a single token. + _, overrides = _dispatch_for( + monkeypatch, + ["eval", "run", "--config", "extra.yaml", "--resource-server", "mcqa", "--model-type", "openai_model"], + ) + paths, others = _split_overrides(overrides) + assert paths == { + "extra.yaml", # raw --config value passes through verbatim; only name selectors get rooted + str(WORKING_DIR / "resources_servers/mcqa/configs/mcqa.yaml"), + str(WORKING_DIR / "responses_api_models/openai_model/configs/openai_model.yaml"), + } + assert others == set() + + def test_unknown_benchmark_errors_with_available_hint(self, monkeypatch: MonkeyPatch, capsys) -> None: + monkeypatch.setattr(cli_main, "dispatch", lambda target, overrides: None) + monkeypatch.setattr(sys, "argv", ["gym", "eval", "prepare", "--benchmark", "does_not_exist"]) + with pytest.raises(SystemExit): + main() + err = capsys.readouterr().err + assert "benchmarks/does_not_exist/config.yaml" in err + assert "does not exist" in err + assert "benchmarks/" in err + + def test_unknown_flavor_error_points_at_configs_dir(self, monkeypatch: MonkeyPatch, capsys) -> None: + # For a known server with an unknown flavor, the hint should point at that server's configs/ dir. + monkeypatch.setattr(cli_main, "dispatch", lambda target, overrides: None) + monkeypatch.setattr(sys, "argv", ["gym", "env", "run", "--resource-server", "mcqa/nope"]) + with pytest.raises(SystemExit): + main() + err = capsys.readouterr().err + assert "resources_servers/mcqa/configs/nope.yaml" in err + assert "resources_servers/mcqa/configs/" in err + + +class TestDidYouMean: + """difflib-backed "did you mean?" hints for mistyped commands, flags, and component names (proposal UX 4).""" + + def test_helper_suggests_close_match(self) -> None: + assert cli_main._did_you_mean("evl", ["list", "eval", "env"]) == " Did you mean `eval`?" + + def test_helper_silent_when_nothing_close(self) -> None: + assert cli_main._did_you_mean("zzzzzz", ["list", "eval", "env"]) == "" + + def _run_expecting_exit(self, monkeypatch: MonkeyPatch, capsys, argv: list[str]) -> str: + monkeypatch.setattr(cli_main, "dispatch", lambda target, overrides: None) + monkeypatch.setattr(sys, "argv", ["gym", *argv]) + with pytest.raises(SystemExit): + main() + return capsys.readouterr().err + + def test_mistyped_group(self, monkeypatch: MonkeyPatch, capsys) -> None: + err = self._run_expecting_exit(monkeypatch, capsys, ["evl"]) + assert "invalid choice: 'evl'" in err + assert "Did you mean `eval`?" in err + + def test_mistyped_action(self, monkeypatch: MonkeyPatch, capsys) -> None: + err = self._run_expecting_exit(monkeypatch, capsys, ["eval", "rnu"]) + assert "Did you mean `run`?" in err + + def test_mistyped_flag_choice(self, monkeypatch: MonkeyPatch, capsys) -> None: + # --storage validates choices, so the parser-level hint kicks in. + err = self._run_expecting_exit(monkeypatch, capsys, ["dataset", "upload", "--storage", "gitlb"]) + assert "Did you mean `gitlab`?" in err + + def test_misspelled_flag(self, monkeypatch: MonkeyPatch, capsys) -> None: + err = self._run_expecting_exit(monkeypatch, capsys, ["eval", "run", "--benchmrk", "aalcr"]) + assert "unrecognized arguments: --benchmrk" in err + assert "Did you mean `--benchmark`?" in err + + def test_misspelled_component_name(self, monkeypatch: MonkeyPatch, capsys) -> None: + err = self._run_expecting_exit(monkeypatch, capsys, ["eval", "prepare", "--benchmark", "aalcrr"]) + assert "Did you mean `aalcr`?" in err + + def test_misspelled_component_flavor(self, monkeypatch: MonkeyPatch, capsys) -> None: + err = self._run_expecting_exit( + monkeypatch, capsys, ["eval", "run", "--resource-server", "math_with_judge/dapo17"] + ) + assert "Did you mean `dapo17k`?" in err + + +class TestSearchDir: + """--search-dir registers extra roots that the name->config selectors also search (REQ 5).""" + + def _make_user_benchmark(self, tmp_path, name: str = "mybench") -> None: + bench_dir = tmp_path / "benchmarks" / name + bench_dir.mkdir(parents=True) + (bench_dir / "config.yaml").write_text("{}\n") + + def test_resolves_component_from_user_dir(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + self._make_user_benchmark(tmp_path) + _, overrides = _dispatch_for( + monkeypatch, ["eval", "prepare", "--benchmark", "mybench", "--search-dir", str(tmp_path)] + ) + # User-dir matches are returned rooted so Hydra can resolve them. + assert overrides == [f"+config_paths=[{tmp_path / 'benchmarks' / 'mybench' / 'config.yaml'}]"] + + def test_builtin_resolves_when_search_dir_lacks_it(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + # A built-in still resolves under WORKING_DIR when a --search-dir is provided that does not shadow it. + _, overrides = _dispatch_for( + monkeypatch, ["eval", "prepare", "--benchmark", "gsm8k", "--search-dir", str(tmp_path)] + ) + assert overrides == [f"+config_paths=[{WORKING_DIR / 'benchmarks/gsm8k/config.yaml'}]"] + + def test_ambiguous_match_errors(self, monkeypatch: MonkeyPatch, tmp_path, capsys) -> None: + # A built-in name also present in a --search-dir is ambiguous; the user must disambiguate with --config. + self._make_user_benchmark(tmp_path, name="gsm8k") # gsm8k also exists under WORKING_DIR + monkeypatch.setattr(cli_main, "dispatch", lambda target, overrides: None) + monkeypatch.setattr( + sys, "argv", ["gym", "eval", "prepare", "--benchmark", "gsm8k", "--search-dir", str(tmp_path)] + ) + with pytest.raises(SystemExit): + main() + err = capsys.readouterr().err + assert "ambiguous" in err + assert str(WORKING_DIR / "benchmarks" / "gsm8k" / "config.yaml") in err + assert str(tmp_path / "benchmarks" / "gsm8k" / "config.yaml") in err + + def test_search_dir_alone_emits_nothing(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + # --search-dir is consumed by the selectors; on its own it is not a Hydra override. + _, overrides = _dispatch_for(monkeypatch, ["eval", "prepare", "--search-dir", str(tmp_path)]) + assert overrides == [] + + def test_did_you_mean_spans_user_dir(self, monkeypatch: MonkeyPatch, tmp_path, capsys) -> None: + self._make_user_benchmark(tmp_path) + monkeypatch.setattr(cli_main, "dispatch", lambda target, overrides: None) + monkeypatch.setattr( + sys, "argv", ["gym", "eval", "prepare", "--benchmark", "mybnch", "--search-dir", str(tmp_path)] + ) + with pytest.raises(SystemExit): + main() + assert "Did you mean `mybench`?" in capsys.readouterr().err diff --git a/tests/unit_tests/test_cli_setup_command.py b/tests/unit_tests/test_cli_setup_command.py index b751c4472d..3effe3a625 100644 --- a/tests/unit_tests/test_cli_setup_command.py +++ b/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"}) diff --git a/tests/unit_tests/test_dataset_source.py b/tests/unit_tests/test_dataset_source.py new file mode 100644 index 0000000000..7f4a208b43 --- /dev/null +++ b/tests/unit_tests/test_dataset_source.py @@ -0,0 +1,117 @@ +# 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 pydantic import ValidationError +from pytest import raises, warns + +from nemo_gym.config_types import ( + DatasetConfig, + GitlabDatasetSource, + HuggingFaceDatasetSource, +) + + +def _dataset(**extra) -> dict: + return {"name": "ds", "type": "example", "jsonl_fpath": "data.jsonl", **extra} + + +class TestDatasetSource: + def test_source_gitlab_backfills_legacy_identifier(self) -> None: + cfg = DatasetConfig.model_validate( + _dataset( + source={ + "type": "gitlab", + "dataset_name": "my_dataset", + "version": "0.0.1", + "artifact_fpath": "train.jsonl", + } + ) + ) + + assert isinstance(cfg.source, GitlabDatasetSource) + # Existing consumers read the legacy field; it must be back-filled from `source`. + assert cfg.gitlab_identifier is not None + assert cfg.gitlab_identifier.dataset_name == "my_dataset" + assert cfg.gitlab_identifier.version == "0.0.1" + assert cfg.gitlab_identifier.artifact_fpath == "train.jsonl" + assert cfg.huggingface_identifier is None + + def test_source_huggingface_backfills_legacy_identifier(self) -> None: + cfg = DatasetConfig.model_validate( + _dataset(source={"type": "huggingface", "repo_id": "org/dataset", "artifact_fpath": "train.jsonl"}) + ) + + assert isinstance(cfg.source, HuggingFaceDatasetSource) + assert cfg.huggingface_identifier is not None + assert cfg.huggingface_identifier.repo_id == "org/dataset" + assert cfg.huggingface_identifier.artifact_fpath == "train.jsonl" + assert cfg.gitlab_identifier is None + + def test_legacy_gitlab_identifier_mirrors_into_source_with_warning(self) -> None: + with warns(DeprecationWarning, match="gitlab_identifier"): + cfg = DatasetConfig.model_validate( + _dataset( + gitlab_identifier={ + "dataset_name": "my_dataset", + "version": "0.0.1", + "artifact_fpath": "train.jsonl", + } + ) + ) + + assert isinstance(cfg.source, GitlabDatasetSource) + assert cfg.source.dataset_name == "my_dataset" + assert cfg.source.version == "0.0.1" + assert cfg.source.artifact_fpath == "train.jsonl" + # Legacy field stays populated so nothing that already reads it breaks. + assert cfg.gitlab_identifier is not None + + def test_legacy_huggingface_identifier_mirrors_into_source_with_warning(self) -> None: + with warns(DeprecationWarning, match="huggingface_identifier"): + cfg = DatasetConfig.model_validate(_dataset(huggingface_identifier={"repo_id": "org/dataset"})) + + assert isinstance(cfg.source, HuggingFaceDatasetSource) + assert cfg.source.repo_id == "org/dataset" + assert cfg.source.artifact_fpath is None + assert cfg.huggingface_identifier is not None + + def test_specifying_source_and_legacy_identifier_is_rejected(self) -> None: + with raises(ValidationError, match="set only one"): + DatasetConfig.model_validate( + _dataset( + source={ + "type": "gitlab", + "dataset_name": "my_dataset", + "version": "0.0.1", + "artifact_fpath": "train.jsonl", + }, + gitlab_identifier={ + "dataset_name": "my_dataset", + "version": "0.0.1", + "artifact_fpath": "train.jsonl", + }, + ) + ) + + def test_no_source_is_allowed(self) -> None: + cfg = DatasetConfig.model_validate(_dataset()) + + assert cfg.source is None + assert cfg.gitlab_identifier is None + assert cfg.huggingface_identifier is None + + def test_source_discriminator_selects_backend(self) -> None: + with raises(ValidationError): + # Missing repo_id for the huggingface branch. + DatasetConfig.model_validate(_dataset(source={"type": "huggingface", "artifact_fpath": "train.jsonl"})) diff --git a/tests/unit_tests/test_server_status.py b/tests/unit_tests/test_server_status.py index 609ddbeae9..3157afad8a 100644 --- a/tests/unit_tests/test_server_status.py +++ b/tests/unit_tests/test_server_status.py @@ -12,15 +12,15 @@ # 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. +import logging from io import StringIO from unittest.mock import MagicMock import requests from pytest import MonkeyPatch -from nemo_gym.cli import ServerInstanceDisplayConfig from nemo_gym.server_status import StatusCommand -from nemo_gym.server_utils import ServerClient +from nemo_gym.server_utils import ServerClient, ServerInstanceDisplayConfig class TestServerStatus: @@ -209,7 +209,7 @@ def test_discover_servers(self, monkeypatch: MonkeyPatch) -> None: first_call = mock_get.call_args_list[0] assert first_call[0][0] == "http://127.0.0.1:11000/server_instances" - def test_discover_servers_head_server_down(self, monkeypatch: MonkeyPatch, capsys) -> None: + def test_discover_servers_head_server_down(self, monkeypatch: MonkeyPatch, capsys, caplog) -> None: mock_head_config = MagicMock() mock_head_config.host = "127.0.0.1" mock_head_config.port = 11000 @@ -220,9 +220,11 @@ def test_discover_servers_head_server_down(self, monkeypatch: MonkeyPatch, capsy monkeypatch.setattr(requests, "get", mock_get) cmd = StatusCommand() - servers = cmd.discover_servers() + with caplog.at_level(logging.WARNING): + servers = cmd.discover_servers() assert len(servers) == 0 - captured = capsys.readouterr() - assert "Could not connect to head server" in captured.out - assert "ng_run" in captured.out + # The warning goes through logging (stderr), keeping stdout machine-readable for `--json`. + assert capsys.readouterr().out == "" + assert "Could not connect to head server" in caplog.text + assert "gym env run" in caplog.text diff --git a/tests/unit_tests/test_train_data_utils.py b/tests/unit_tests/test_train_data_utils.py index 0e57463c23..31ada07687 100644 --- a/tests/unit_tests/test_train_data_utils.py +++ b/tests/unit_tests/test_train_data_utils.py @@ -125,6 +125,7 @@ def test_load_and_validate_server_instance_configs_sanity(self, monkeypatch: Mon "type": "example", "jsonl_fpath": "resources_servers/example_multi_step/data/example.jsonl", "num_repeats": 1, + "source": None, "gitlab_identifier": None, "huggingface_identifier": None, "license": None,