From cf3e16b579a416cfe44b86995bbbfd8875864e13 Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Mon, 13 Jul 2026 14:54:06 +0200 Subject: [PATCH 01/13] chore: unify metadata and parsing code when listing envs and benchmarks Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/benchmarks.py | 62 +++++------- nemo_gym/cli/env.py | 2 +- nemo_gym/cli/eval.py | 73 ++++---------- nemo_gym/discovery.py | 145 ++++++++++++++++++++++++++++ nemo_gym/registry.py | 44 ++------- tests/unit_tests/test_benchmarks.py | 138 +++++--------------------- tests/unit_tests/test_discovery.py | 117 ++++++++++++++++++++++ tests/unit_tests/test_registry.py | 5 +- 8 files changed, 340 insertions(+), 246 deletions(-) create mode 100644 nemo_gym/discovery.py create mode 100644 tests/unit_tests/test_discovery.py diff --git a/nemo_gym/benchmarks.py b/nemo_gym/benchmarks.py index 254962589d..48ce084a35 100644 --- a/nemo_gym/benchmarks.py +++ b/nemo_gym/benchmarks.py @@ -14,18 +14,17 @@ # limitations under the License. """Benchmark discovery and preparation utilities.""" -import re import sys -from copy import deepcopy +from glob import glob from pathlib import Path from typing import Dict, List, Optional from omegaconf import DictConfig, OmegaConf -from omegaconf.errors import InterpolationKeyError from pydantic import BaseModel from nemo_gym import PARENT_DIR from nemo_gym.config_types import BenchmarkDatasetConfig +from nemo_gym.discovery import _parse_no_environment_tolerating_unset_values from nemo_gym.global_config import ( POLICY_MODEL_KEY_NAME, GlobalConfigDictParser, @@ -36,42 +35,6 @@ BENCHMARKS_DIR = PARENT_DIR / "benchmarks" -# Fills unset `???`/`${...}` values during listing: they reference runtime-only values (API keys, -# endpoints) not needed to identify a benchmark, so a placeholder lets the config still resolve. -_UNSET_VALUE_PLACEHOLDER = "__unset_for_listing__" - - -def _parse_no_environment_tolerating_unset_values(initial_config_dict: DictConfig) -> DictConfig: - """`parse_no_environment` for *listing*: fill unset `???` values and undefined `${...}` interpolations - with a placeholder so a benchmark referencing runtime-only values can still be identified. Never mutates - the caller's config; errors other than those two propagate. - - `???` is filled anywhere; `${...}` only where parse forces resolution (top-level and server sections), - not inside arbitrary non-server nested dicts — fine for listing, whose interpolations live in servers. - """ - working = deepcopy(initial_config_dict) # never mutate the caller's config - parser = GlobalConfigDictParser() - - # Fill all `???` leaves in one pass. The loop below only adds placeholder keys, so no new `???` appear. - for path in parser.collect_missing_value_paths(working): - OmegaConf.update(working, path, _UNSET_VALUE_PLACEHOLDER) - - # OmegaConf reports undefined `${...}` keys only one at a time (as InterpolationKeyError), so loop: - # inject a placeholder for each reported key and retry until it resolves. - injected: set[str] = set() - while True: - try: - return parser.parse_no_environment(initial_global_config_dict=working) - except InterpolationKeyError as e: - # The missing key name is only in the message text — omegaconf never stores it on an attribute - # (`e.key`/`e.full_key` point at the containing node), so a regex is the only way to read it. - match = re.search(r"Interpolation key '([^']+)'", str(e)) - key = match.group(1) if match else None - if not key or key in injected: - raise # can't identify/clear the missing key; let the caller decide (warn + skip) - injected.add(key) - working = OmegaConf.merge(DictConfig({key: _UNSET_VALUE_PLACEHOLDER}), working) - class BenchmarkConfig(BaseModel): name: str @@ -161,6 +124,27 @@ def _load_benchmarks_from_config_paths(config_paths: List[Path]) -> Dict[str, Be return benchmarks_dict +def _benchmark_config_paths(benchmarks_dir: Path) -> List[Path]: + """Sorted config paths under one dir that declare a benchmark, discovered by content. + + A config defines a benchmark iff it declares a `type: benchmark` dataset (see `BenchmarkConfig`), + regardless of its filename. So discovery is content-based: scan every yaml and keep the ones that + literally declare such a dataset. That text check is a cheap prefilter so we only pay the resolve + cost on real candidates (not every prompt/endpoint yaml), and it finds benchmarks whose config + isn't named `config.yaml` — e.g. tau2's `configs/*.yaml` and livecodebench's `cascade.yaml`. + Returns an empty list if the directory is missing. + """ + if not benchmarks_dir.is_dir(): + return [] + config_paths = [benchmarks_dir / p for p in glob("**/*.yaml", root_dir=benchmarks_dir, recursive=True)] + return sorted(p for p in config_paths if "type: benchmark" in p.read_text(errors="ignore")) + + +def discover_benchmarks() -> Dict[str, BenchmarkConfig]: + """Map benchmark name -> :class:`BenchmarkConfig` for every benchmark config under ``benchmarks/``.""" + return _load_benchmarks_from_config_paths(_benchmark_config_paths(BENCHMARKS_DIR)) + + # Backward-compatibility shims (CLI refactor): these symbols moved to `nemo_gym.cli.eval`. # Re-exported lazily to avoid a circular import; accessing them emits a DeprecationWarning. from nemo_gym.cli._compat import moved_attr_getter # noqa: E402 diff --git a/nemo_gym/cli/env.py b/nemo_gym/cli/env.py index 8485cbd77d..9b11afce5f 100644 --- a/nemo_gym/cli/env.py +++ b/nemo_gym/cli/env.py @@ -953,7 +953,7 @@ def list_environments() -> None: return table = Table(title=f"Available environments in NeMo Gym ({len(environments)})") - table.add_column("Environment") + table.add_column("Name") table.add_column("Domain") table.add_column("Description") for name, environment in environments.items(): diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index 04ca1684f4..34dfb8287c 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -18,7 +18,6 @@ import importlib import json from copy import deepcopy -from glob import glob from multiprocessing import Pool from pathlib import Path from typing import Any, Dict, List, Tuple @@ -32,15 +31,14 @@ from nemo_gym.benchmarks import ( BENCHMARKS_DIR, BenchmarkConfig, - _load_benchmarks_from_config_paths, - _parse_no_environment_tolerating_unset_values, + discover_benchmarks, ) from nemo_gym.cli.env import RunHelper from nemo_gym.cli.utils import exit_cleanly_on_config_error, print_rich_table from nemo_gym.config_types import BaseNeMoGymCLIConfig, BenchmarkDatasetConfig, ConfigError, ConfigPathNotFoundError +from nemo_gym.discovery import read_config_metadata from nemo_gym.global_config import ( JSON_OUTPUT_KEY_NAME, - POLICY_MODEL_KEY_NAME, QUERY_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, TASK_INDEX_KEY_NAME, @@ -75,40 +73,12 @@ def _fuzzy_matches(query: str, *fields: str) -> bool: return False -def _benchmark_domain(bench: BenchmarkConfig) -> str: - """Resolve a benchmark's config to its `domain` (for the domain column and `gym search`). +def list_benchmarks() -> None: + """CLI command: list available benchmarks, optionally filtered by a `query` (the `gym search` entry point). - `BenchmarkConfig` flattens away the `domain`, so we re-resolve the config with the tolerant listing - parser (so chained `config_paths` / `_inherit_from` are applied) and read the field back out. `domain` - may be declared on any server config — a resources server (e.g. `aime24`) or an agent (e.g. `tau2`) — - so we scan every server group. + A benchmark is a specific kind of environment, so it shares `gym list environments`' columns (name, + domain, description) and reads them through the same `read_config_metadata` helper. """ - 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 = _parse_no_environment_tolerating_unset_values(initial_config_dict) - - for instance_name in resolved: - instance = resolved[instance_name] - if not isinstance(instance, (dict, DictConfig)): - continue - - for group_key in ("resources_servers", "responses_api_agents", "responses_api_models"): - servers = instance.get(group_key) - if not servers: - continue - for server_config in servers.values(): - found_domain = (server_config or {}).get("domain") - if found_domain: - return str(found_domain) - - return "" - - -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, @@ -116,34 +86,28 @@ def list_benchmarks() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - assert BENCHMARKS_DIR.exists(), "Missing benchmarks directory" - - # A config defines a benchmark iff it declares a `type: benchmark` dataset (see `BenchmarkConfig`), - # regardless of its filename. So discovery is content-based: scan every yaml and keep the ones that - # literally declare such a dataset. That text check is a cheap prefilter so we only pay the resolve - # cost on real candidates (not every prompt/endpoint yaml), and it finds benchmarks whose config - # isn't named `config.yaml` — e.g. tau2's `configs/*.yaml` and livecodebench's `cascade.yaml`. - config_paths = [BENCHMARKS_DIR / p for p in glob("**/*.yaml", root_dir=BENCHMARKS_DIR, recursive=True)] - config_paths = sorted(p for p in config_paths if "type: benchmark" in p.read_text(errors="ignore")) - - benchmarks = _load_benchmarks_from_config_paths(config_paths) + benchmarks = discover_benchmarks() - # Resolve the domain once per benchmark, for the domain column and `gym search`. - domains = {name: _benchmark_domain(bench) for name, bench in benchmarks.items()} + # Resolve domain + description once per benchmark, via the shared component-metadata reader — + # the same one `gym list environments` uses — for the columns and `gym search`. + metadata = {name: read_config_metadata(bench.path) for name, bench in benchmarks.items()} # `gym search ` reuses this command, narrowing the listing to fuzzy matches # across the benchmark name and domain. query = global_config_dict.get(QUERY_KEY_NAME) if query: - benchmarks = {name: bench for name, bench in benchmarks.items() if _fuzzy_matches(query, name, domains[name])} + benchmarks = { + name: bench for name, bench in benchmarks.items() if _fuzzy_matches(query, name, metadata[name][0] or "") + } if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): payload = [ { "name": name, "agent_name": bench.agent_name, - "domain": domains[name], + "domain": metadata[name][0] or "", "num_repeats": bench.num_repeats, + "description": metadata[name][1] or "", } for name, bench in benchmarks.items() ] @@ -164,13 +128,16 @@ def list_benchmarks() -> None: else f"Available benchmarks in NeMo Gym ({len(benchmarks)})" ) table = Table(title=title) - table.add_column("Benchmark name") + # Shared environment columns first (name, domain, description), then benchmark-specific ones. + table.add_column("Name") table.add_column("Domain") + table.add_column("Description") table.add_column("Agent name") table.add_column("Num repeats") for name, bench in benchmarks.items(): - table.add_row(name, domains[name], bench.agent_name, str(bench.num_repeats)) + domain, description = metadata[name] + table.add_row(name, domain or "", description or "", bench.agent_name, str(bench.num_repeats)) print_rich_table(table) diff --git a/nemo_gym/discovery.py b/nemo_gym/discovery.py new file mode 100644 index 0000000000..83766e0a90 --- /dev/null +++ b/nemo_gym/discovery.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. +"""Shared component-discovery utilities for the ``gym list``/``gym search`` commands. + +An environment and a benchmark (a benchmark is a specific kind of environment) are listed with the +same columns, read the same way. This module holds the config-reading code both listings share so the +per-component registries (``registry.py``, ``benchmarks.py``, ``agent_registry.py``) can depend on it +without depending on each other. It only reads config files — it never starts servers — and only +imports lower-level config utilities, so it sits below those registries in the import graph. +""" + +import re +from copy import deepcopy +from pathlib import Path +from typing import Optional, Tuple + +from omegaconf import DictConfig, OmegaConf +from omegaconf.errors import InterpolationKeyError + +from nemo_gym.global_config import ( + POLICY_MODEL_KEY_NAME, + GlobalConfigDictParser, + GlobalConfigDictParserConfig, +) + + +# Fills unset `???`/`${...}` values during listing: they reference runtime-only values (API keys, +# endpoints) not needed to identify a component, so a placeholder lets the config still resolve. +_UNSET_VALUE_PLACEHOLDER = "__unset_for_listing__" + +# Server groups a component's `domain`/`description` may be declared on. `domain` can sit on a +# resources server (e.g. `aime24`), an agent (e.g. `tau2`), or in principle a model server. +_SERVER_GROUP_KEYS = ("resources_servers", "responses_api_agents", "responses_api_models") + + +def _parse_no_environment_tolerating_unset_values(initial_config_dict: DictConfig) -> DictConfig: + """`parse_no_environment` for *listing*: fill unset `???` values and undefined `${...}` interpolations + with a placeholder so a component referencing runtime-only values can still be identified. Never mutates + the caller's config; errors other than those two propagate. + + `???` is filled anywhere; `${...}` only where parse forces resolution (top-level and server sections), + not inside arbitrary non-server nested dicts — fine for listing, whose interpolations live in servers. + """ + working = deepcopy(initial_config_dict) # never mutate the caller's config + parser = GlobalConfigDictParser() + + # Fill all `???` leaves in one pass. The loop below only adds placeholder keys, so no new `???` appear. + for path in parser.collect_missing_value_paths(working): + OmegaConf.update(working, path, _UNSET_VALUE_PLACEHOLDER) + + # OmegaConf reports undefined `${...}` keys only one at a time (as InterpolationKeyError), so loop: + # inject a placeholder for each reported key and retry until it resolves. + injected: set[str] = set() + while True: + try: + return parser.parse_no_environment(initial_global_config_dict=working) + except InterpolationKeyError as e: + # The missing key name is only in the message text — omegaconf never stores it on an attribute + # (`e.key`/`e.full_key` point at the containing node), so a regex is the only way to read it. + match = re.search(r"Interpolation key '([^']+)'", str(e)) + key = match.group(1) if match else None + if not key or key in injected: + raise # can't identify/clear the missing key; let the caller decide (warn + skip) + injected.add(key) + working = OmegaConf.merge(DictConfig({key: _UNSET_VALUE_PLACEHOLDER}), working) + + +def _scan_servers_for_metadata(container) -> Tuple[Optional[str], Optional[str]]: + """Best-effort ``(domain, description)`` from a config mapping, scanning every server group. + + Reads the first ``domain`` and the first ``description`` found across all server instances. Defensive + against malformed shapes (non-mapping top level, a server group that isn't a dict) so it never raises. + """ + domain: Optional[str] = None + description: Optional[str] = None + if not isinstance(container, (dict, DictConfig)): + return None, None + for instance in container.values(): + if not isinstance(instance, (dict, DictConfig)): + continue + for group_key in _SERVER_GROUP_KEYS: + servers = instance.get(group_key) + if not isinstance(servers, (dict, DictConfig)): + continue + for server_config in servers.values(): + if not isinstance(server_config, (dict, DictConfig)): + continue + if domain is None and server_config.get("domain"): + domain = str(server_config["domain"]) + if description is None and server_config.get("description"): + description = str(server_config["description"]) + return domain, description + + +def read_config_metadata(config_path: Path) -> Tuple[Optional[str], Optional[str]]: + """Shared listing metadata reader: ``(domain, description)`` for an environment *or* benchmark config. + + A benchmark is a specific kind of environment, so both are read the same way. Two-pass, because the two + kinds declare metadata differently: + + 1. **Raw (non-resolving) scan** of the config as written. Environment configs declare ``domain``/ + ``description`` inline, and reading them without resolution is safe even though an environment config + references model/agent servers defined elsewhere (resolving it in isolation would raise). + 2. **Resolving fallback**, only for whatever the raw scan left unset. Benchmark configs inherit their + metadata via ``config_paths``/``_inherit_from``, so it isn't present until the config is resolved. + Uses the listing-tolerant parser (unset runtime values are placeholdered); on any failure — e.g. an + environment config that references servers defined elsewhere — whatever the raw scan found is kept. + + Never raises: an unreadable or unresolvable config yields ``(None, None)``. + """ + try: + raw = OmegaConf.to_container(OmegaConf.load(config_path), resolve=False, throw_on_missing=False) + except Exception: + raw = None + domain, description = _scan_servers_for_metadata(raw) + if domain is not None and description is not None: + return domain, description + + try: + initial_config_dict = OmegaConf.load(config_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 = _parse_no_environment_tolerating_unset_values(initial_config_dict) + except Exception: + return domain, description + + resolved_domain, resolved_description = _scan_servers_for_metadata(resolved) + return ( + domain if domain is not None else resolved_domain, + description if description is not None else resolved_description, + ) diff --git a/nemo_gym/registry.py b/nemo_gym/registry.py index 81a9069c5e..53cb1da128 100644 --- a/nemo_gym/registry.py +++ b/nemo_gym/registry.py @@ -20,17 +20,18 @@ ``gym list environments``. Resolving a name to a config path for *running* is handled by the CLI's generic ``--environment`` asset selector, so this module is intentionally discovery-only. -Discovery only reads config files; it never resolves interpolations or starts servers, so it is -safe to call even when secrets/API keys referenced by a config are not set in the environment. +Discovery only reads config files and never starts servers. Its ``domain``/``description`` metadata is +read via the shared :func:`~nemo_gym.discovery.read_config_metadata` reader (the same one benchmark +listing uses), which tolerates unset secrets/API keys referenced by a config, so discovery is safe to +call even when those are not set in the environment. """ from dataclasses import dataclass from pathlib import Path -from typing import Dict, Optional, Tuple - -from omegaconf import OmegaConf +from typing import Dict, Optional from nemo_gym import PARENT_DIR +from nemo_gym.discovery import read_config_metadata ENVIRONMENTS_DIR = PARENT_DIR / "environments" @@ -48,37 +49,6 @@ class EnvironmentEntry: domain: Optional[str] = None -def _read_metadata(config_path: Path) -> Tuple[Optional[str], Optional[str]]: - """Best-effort ``(description, domain)`` from the config's resources_servers entry. - - Reads without resolving interpolations or missing values so a config that references an unset - key (e.g. an API key) still yields metadata instead of raising. - """ - try: - container = OmegaConf.to_container(OmegaConf.load(config_path), resolve=False, throw_on_missing=False) - except Exception: - return None, None - - if not isinstance(container, dict): - return None, None - - for top_level_value in container.values(): - if not isinstance(top_level_value, dict): - continue - resources_servers = top_level_value.get("resources_servers") - if not isinstance(resources_servers, dict): - continue - for server_config in resources_servers.values(): - if isinstance(server_config, dict): - description = server_config.get("description") - domain = server_config.get("domain") - return ( - description if isinstance(description, str) else None, - domain if isinstance(domain, str) else None, - ) - return None, None - - def discover_environments(environments_dir: Path = ENVIRONMENTS_DIR) -> Dict[str, EnvironmentEntry]: """Map environment name -> :class:`EnvironmentEntry` for every ``/config.yaml``. @@ -93,7 +63,7 @@ def discover_environments(environments_dir: Path = ENVIRONMENTS_DIR) -> Dict[str if not (child.is_dir() and config_path.is_file()): continue - description, domain = _read_metadata(config_path) + domain, description = read_config_metadata(config_path) environments[child.name] = EnvironmentEntry( name=child.name, config_path=config_path, diff --git a/tests/unit_tests/test_benchmarks.py b/tests/unit_tests/test_benchmarks.py index 4dc0926452..f535066beb 100644 --- a/tests/unit_tests/test_benchmarks.py +++ b/tests/unit_tests/test_benchmarks.py @@ -20,7 +20,7 @@ from omegaconf import OmegaConf from yaml import safe_load -from nemo_gym.cli.eval import _benchmark_domain, _fuzzy_matches, list_benchmarks, prepare_benchmark +from nemo_gym.cli.eval import _fuzzy_matches, list_benchmarks, prepare_benchmark def _mock_global_config(config: dict = None): @@ -34,10 +34,12 @@ def test_lists_found_benchmarks(self, capsys) -> None: list_benchmarks() assert "aime24" in capsys.readouterr().out - def test_discovers_by_type_benchmark_not_filename(self, tmp_path, capsys) -> None: + def test_discovers_by_type_benchmark_not_filename(self, tmp_path) -> None: # Discovery is content-based (a `type: benchmark` dataset), not filename-based: any yaml that # declares such a dataset is a candidate (e.g. tau2's `configs/tau2.yaml`), and yamls that don't # are skipped — regardless of filename. + from nemo_gym.benchmarks import _benchmark_config_paths + (tmp_path / "standard").mkdir() (tmp_path / "standard" / "config.yaml").write_text("x:\n datasets:\n - type: benchmark\n") (tmp_path / "flavored" / "configs").mkdir(parents=True) @@ -45,25 +47,13 @@ def test_discovers_by_type_benchmark_not_filename(self, tmp_path, capsys) -> Non (tmp_path / "notbench").mkdir() (tmp_path / "notbench" / "config.yaml").write_text("x:\n prompt_config: hi.yaml\n") # no benchmark dataset - captured = {} - - def fake_load(paths): - captured["paths"] = {str(p.relative_to(tmp_path)) for p in paths} - return {} - - with ( - patch("nemo_gym.cli.eval.get_global_config_dict", return_value=_mock_global_config()), - patch("nemo_gym.cli.eval.BENCHMARKS_DIR", tmp_path), - patch("nemo_gym.cli.eval._load_benchmarks_from_config_paths", side_effect=fake_load), - ): - list_benchmarks() - - assert captured["paths"] == {"standard/config.yaml", "flavored/configs/myflavor.yaml"} + found = {str(p.relative_to(tmp_path)) for p in _benchmark_config_paths(tmp_path)} + assert found == {"standard/config.yaml", "flavored/configs/myflavor.yaml"} def test_no_benchmarks(self, capsys) -> None: with ( 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={}), + patch("nemo_gym.cli.eval.discover_benchmarks", return_value={}), ): list_benchmarks() assert "No benchmarks found" in capsys.readouterr().out @@ -74,12 +64,18 @@ def test_json_output(self, capsys) -> None: 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_domain", return_value="math"), + patch("nemo_gym.cli.eval.discover_benchmarks", return_value={"my_bench": bench}), + patch("nemo_gym.cli.eval.read_config_metadata", return_value=("math", "a description")), ): list_benchmarks() assert json.loads(capsys.readouterr().out) == [ - {"name": "my_bench", "agent_name": "my_agent", "domain": "math", "num_repeats": 4} + { + "name": "my_bench", + "agent_name": "my_agent", + "domain": "math", + "num_repeats": 4, + "description": "a description", + } ] def test_json_output_empty(self, capsys) -> None: @@ -87,7 +83,7 @@ def test_json_output_empty(self, capsys) -> None: 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={}), + patch("nemo_gym.cli.eval.discover_benchmarks", return_value={}), ): list_benchmarks() assert json.loads(capsys.readouterr().out) == [] @@ -135,64 +131,7 @@ def test_every_repo_benchmark_appears_in_listing(self, capsys) -> None: ) -class TestTolerantInterpolationParse: - # Unset `???` values and unresolved `${...}` interpolations reference runtime-only values that aren't - # needed to identify a benchmark; listing fills them with a placeholder so the config still resolves. - def _resolve(self, d: dict): - from nemo_gym.benchmarks import _parse_no_environment_tolerating_unset_values - - return _parse_no_environment_tolerating_unset_values(OmegaConf.create(d)) - - @property - def _placeholder(self) -> str: - from nemo_gym.benchmarks import _UNSET_VALUE_PLACEHOLDER - - return _UNSET_VALUE_PLACEHOLDER - - def test_single_interpolation(self) -> None: - resolved = self._resolve({"foo": "${bar}"}) - assert resolved["foo"] == self._placeholder - - def test_single_missing_value(self) -> None: - resolved = self._resolve({"foo": "???"}) - assert resolved["foo"] == self._placeholder - - def test_mix(self) -> None: - # A mix across nested dicts: resolvable literals (incl. nested) pass through untouched, while an - # undefined `${...}` interpolation and unset `???` values (incl. nested) are filled with the - # placeholder. - resolved = self._resolve( - { - "name": "my_bench", - "num_repeats": 3, - "api_key": "${some_api_key}", - "server": { - "endpoint": "https://example.com", - "nested": { - "enabled": True, - "token": "???", - }, - }, - } - ) - # Correct key-value pairs are unmodified. - assert resolved["name"] == "my_bench" - assert resolved["num_repeats"] == 3 - assert resolved["server"]["endpoint"] == "https://example.com" - assert resolved["server"]["nested"]["enabled"] is True - # Undefined `${...}` and unset `???` values are filled. - assert resolved["api_key"] == self._placeholder - assert resolved["server"]["nested"]["token"] == self._placeholder - - def test_does_not_mutate_input(self) -> None: - from nemo_gym.benchmarks import _parse_no_environment_tolerating_unset_values - - cfg = OmegaConf.create({"foo": "???", "bar": "${baz}"}) - before = OmegaConf.to_container(cfg, resolve=False, throw_on_missing=False) - _parse_no_environment_tolerating_unset_values(cfg) - after = OmegaConf.to_container(cfg, resolve=False, throw_on_missing=False) - assert after == before == {"foo": "???", "bar": "${baz}"} - +class TestBenchmarkConfigStrictParsing: def test_strict_is_the_default_and_does_not_tolerate_unresolved_values(self) -> None: # The tolerance is listing-only: `from_initial_config_dict` defaults to strict, so other workflows # still get a hard error on an unresolved `${...}` rather than a silent placeholder. @@ -229,32 +168,6 @@ def test_no_match(self) -> None: assert not _fuzzy_matches("zzznomatch", "aime24", "math_with_judge") -class TestBenchmarkDomain: - def test_resolves_domain_from_real_config(self) -> None: - from nemo_gym.benchmarks import BENCHMARKS_DIR, BenchmarkConfig - - bench = BenchmarkConfig.from_config_path(BENCHMARKS_DIR / "aime24" / "config.yaml") - - assert _benchmark_domain(bench) == "math" - - def test_resolves_domain_defined_on_agent(self, tmp_path: Path) -> None: - # `domain` can be declared on the agent (responses_api_agents..domain) rather than on a - # resources server, as the tau2 config does. - config_path = tmp_path / "config.yaml" - config_path.write_text( - """tau2_agent: - responses_api_agents: - tau2: - entrypoint: app.py - domain: agent -""" - ) - bench = MagicMock() - bench.path = config_path - - assert _benchmark_domain(bench) == "agent" - - class TestSearchBenchmarks: # Map each benchmark name to the `domain` its config would resolve to. DOMAINS = { @@ -264,7 +177,7 @@ class TestSearchBenchmarks: def _bench(self, key: str): bench = MagicMock(agent_name="my_agent", num_repeats=1) - bench.config_key = key # let the patched _benchmark_domain find the right entry + bench.path = key # the patched read_config_metadata keys off the path to find the domain return bench def _benchmarks(self) -> dict: @@ -273,8 +186,8 @@ def _benchmarks(self) -> dict: 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_domain", side_effect=lambda b: self.DOMAINS[b.config_key]), + patch("nemo_gym.cli.eval.discover_benchmarks", return_value=benchmarks), + patch("nemo_gym.cli.eval.read_config_metadata", side_effect=lambda path: (self.DOMAINS[path], None)), ): list_benchmarks() return capsys.readouterr().out @@ -382,12 +295,9 @@ def test_missing_prepare_function(self, tmp_path: Path, capsys) -> None: assert "Expected the actual prepared dataset output fpath to match the jsonl_fpath set in the config" in out def test_no_benchmark_in_config_paths(self, capsys) -> None: - with ( - patch( - "nemo_gym.cli.eval.get_global_config_dict", - return_value=_mock_global_config({"config_paths": ["resources_servers/foo/configs/foo.yaml"]}), - ), - patch("nemo_gym.cli.eval._load_benchmarks_from_config_paths", return_value={}), + with patch( + "nemo_gym.cli.eval.get_global_config_dict", + return_value=_mock_global_config({"config_paths": ["resources_servers/foo/configs/foo.yaml"]}), ): with pytest.raises(SystemExit) as exc_info: prepare_benchmark() diff --git a/tests/unit_tests/test_discovery.py b/tests/unit_tests/test_discovery.py new file mode 100644 index 0000000000..64e2268b43 --- /dev/null +++ b/tests/unit_tests/test_discovery.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 pathlib import Path + +from omegaconf import OmegaConf + +from nemo_gym.discovery import ( + _UNSET_VALUE_PLACEHOLDER, + _parse_no_environment_tolerating_unset_values, + read_config_metadata, +) + + +class TestTolerantInterpolationParse: + # Unset `???` values and unresolved `${...}` interpolations reference runtime-only values that aren't + # needed to identify a component; listing fills them with a placeholder so the config still resolves. + def _resolve(self, d: dict): + return _parse_no_environment_tolerating_unset_values(OmegaConf.create(d)) + + def test_single_interpolation(self) -> None: + resolved = self._resolve({"foo": "${bar}"}) + assert resolved["foo"] == _UNSET_VALUE_PLACEHOLDER + + def test_single_missing_value(self) -> None: + resolved = self._resolve({"foo": "???"}) + assert resolved["foo"] == _UNSET_VALUE_PLACEHOLDER + + def test_mix(self) -> None: + # A mix across nested dicts: resolvable literals (incl. nested) pass through untouched, while an + # undefined `${...}` interpolation and unset `???` values (incl. nested) are filled with the + # placeholder. + resolved = self._resolve( + { + "name": "my_bench", + "num_repeats": 3, + "api_key": "${some_api_key}", + "server": { + "endpoint": "https://example.com", + "nested": { + "enabled": True, + "token": "???", + }, + }, + } + ) + # Correct key-value pairs are unmodified. + assert resolved["name"] == "my_bench" + assert resolved["num_repeats"] == 3 + assert resolved["server"]["endpoint"] == "https://example.com" + assert resolved["server"]["nested"]["enabled"] is True + # Undefined `${...}` and unset `???` values are filled. + assert resolved["api_key"] == _UNSET_VALUE_PLACEHOLDER + assert resolved["server"]["nested"]["token"] == _UNSET_VALUE_PLACEHOLDER + + def test_does_not_mutate_input(self) -> None: + cfg = OmegaConf.create({"foo": "???", "bar": "${baz}"}) + before = OmegaConf.to_container(cfg, resolve=False, throw_on_missing=False) + _parse_no_environment_tolerating_unset_values(cfg) + after = OmegaConf.to_container(cfg, resolve=False, throw_on_missing=False) + assert after == before == {"foo": "???", "bar": "${baz}"} + + +class TestReadConfigMetadata: + def test_reads_domain_and_description_from_real_benchmark_config(self) -> None: + # aime24 inherits both fields from its resources server via `config_paths`/`_inherit_from`, so + # this only resolves via the tolerant fallback parse. + from nemo_gym.benchmarks import BENCHMARKS_DIR + + domain, description = read_config_metadata(BENCHMARKS_DIR / "aime24" / "config.yaml") + + assert domain == "math" + assert description + + def test_reads_domain_defined_on_agent(self, tmp_path: Path) -> None: + # `domain` can be declared on the agent (responses_api_agents..domain) rather than on a + # resources server, as the tau2 config does. + config_path = tmp_path / "config.yaml" + config_path.write_text( + """tau2_agent: + responses_api_agents: + tau2: + entrypoint: app.py + domain: agent +""" + ) + + assert read_config_metadata(config_path)[0] == "agent" + + def test_reads_inline_metadata_without_resolving_external_server_refs(self, tmp_path: Path) -> None: + # Environment configs reference model/agent servers defined elsewhere; resolving one in isolation + # would raise. Inline domain/description must still be read from the raw config, no resolution. + config_path = tmp_path / "config.yaml" + config_path.write_text( + "env:\n" + " resources_servers:\n" + " env:\n" + " entrypoint: app.py\n" + " domain: rlhf\n" + " description: inline desc\n" + " judge_model_server:\n" + " type: responses_api_models\n" + " name: some_absent_model\n" + ) + + assert read_config_metadata(config_path) == ("rlhf", "inline desc") diff --git a/tests/unit_tests/test_registry.py b/tests/unit_tests/test_registry.py index be9ff1ee20..9358f3f251 100644 --- a/tests/unit_tests/test_registry.py +++ b/tests/unit_tests/test_registry.py @@ -88,8 +88,9 @@ def test_unparseable_or_metadataless_configs_still_discovered(self, tmp_path: Pa assert entry.description is None assert entry.domain is None - def test_metadata_does_not_resolve_interpolations(self, tmp_path: Path) -> None: - # A config referencing an unset interpolation must still be discoverable (no resolution). + def test_metadata_tolerates_unset_interpolations(self, tmp_path: Path) -> None: + # A config referencing an unset interpolation must still be discoverable: the shared metadata + # reader reads the inline `domain` from the raw config and tolerates the unresolved value. envs_dir = tmp_path / "environments" _make_env( envs_dir, From eb4f974c06345f140a4df2e6caac03626dbe4ade Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Tue, 14 Jul 2026 14:44:30 +0200 Subject: [PATCH 02/13] feat: add --search-dir to gym list and gym search Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/agent_registry.py | 22 +++++-- nemo_gym/benchmarks.py | 30 +++++---- nemo_gym/cli/agents.py | 12 ++-- nemo_gym/cli/env.py | 6 +- nemo_gym/cli/eval.py | 6 +- nemo_gym/cli/main.py | 52 +++++++++------ nemo_gym/discovery.py | 88 +++++++++++++++++-------- nemo_gym/global_config.py | 2 + nemo_gym/registry.py | 32 ++++++--- tests/unit_tests/test_agent_registry.py | 22 ++++--- tests/unit_tests/test_cli_main.py | 22 ++++--- tests/unit_tests/test_discovery.py | 51 ++++++++++++++ tests/unit_tests/test_registry.py | 28 ++++++-- 13 files changed, 266 insertions(+), 107 deletions(-) diff --git a/nemo_gym/agent_registry.py b/nemo_gym/agent_registry.py index 9b8958eb9b..4101fba347 100644 --- a/nemo_gym/agent_registry.py +++ b/nemo_gym/agent_registry.py @@ -39,14 +39,16 @@ from dataclasses import dataclass from pathlib import Path -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Sequence, Tuple, Union from omegaconf import OmegaConf from nemo_gym import PARENT_DIR +from nemo_gym.discovery import component_search_roots, merge_by_name -AGENTS_DIR = PARENT_DIR / "responses_api_agents" +AGENTS_SUBDIR = "responses_api_agents" +AGENTS_DIR = PARENT_DIR / AGENTS_SUBDIR AGENT_CONFIGS_SUBDIR = "configs" @@ -129,8 +131,8 @@ def _classify(config_paths: Tuple[Path, ...]) -> Tuple[bool, Optional[str]]: return not references_resources_server, description -def discover_agents(agents_dir: Path = AGENTS_DIR) -> Dict[str, AgentEntry]: - """Map agent name -> :class:`AgentEntry` for every agent dir under ``responses_api_agents/``. +def _discover_agents_in_dir(agents_dir: Path) -> Dict[str, AgentEntry]: + """Map agent name -> :class:`AgentEntry` for every agent dir under one ``responses_api_agents/`` dir. The name is the directory name. A directory is an agent if it has an ``app.py`` or at least one agent config. Returns an empty dict if the directory is missing. @@ -158,3 +160,15 @@ def discover_agents(agents_dir: Path = AGENTS_DIR) -> Dict[str, AgentEntry]: ) return agents + + +def discover_agents( + search_dirs: Optional[Union[Path, Sequence[Path]]] = None, +) -> Dict[str, AgentEntry]: + """Map agent name -> :class:`AgentEntry` for every discoverable agent dir. + + Scans the ``responses_api_agents/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` + root (``search_dirs`` + cwd + built-ins), merged so user agents shadow same-named built-ins. + ``search_dirs`` is one dir or a list. + """ + return merge_by_name(_discover_agents_in_dir(root / AGENTS_SUBDIR) for root in component_search_roots(search_dirs)) diff --git a/nemo_gym/benchmarks.py b/nemo_gym/benchmarks.py index 48ce084a35..45b701550f 100644 --- a/nemo_gym/benchmarks.py +++ b/nemo_gym/benchmarks.py @@ -17,14 +17,14 @@ import sys from glob import glob from pathlib import Path -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Sequence, Union from omegaconf import DictConfig, OmegaConf from pydantic import BaseModel from nemo_gym import PARENT_DIR from nemo_gym.config_types import BenchmarkDatasetConfig -from nemo_gym.discovery import _parse_no_environment_tolerating_unset_values +from nemo_gym.discovery import _parse_no_environment_tolerating_unset_values, component_search_roots, merge_by_name from nemo_gym.global_config import ( POLICY_MODEL_KEY_NAME, GlobalConfigDictParser, @@ -33,7 +33,8 @@ ) -BENCHMARKS_DIR = PARENT_DIR / "benchmarks" +BENCHMARKS_SUBDIR = "benchmarks" +BENCHMARKS_DIR = PARENT_DIR / BENCHMARKS_SUBDIR class BenchmarkConfig(BaseModel): @@ -127,12 +128,9 @@ def _load_benchmarks_from_config_paths(config_paths: List[Path]) -> Dict[str, Be def _benchmark_config_paths(benchmarks_dir: Path) -> List[Path]: """Sorted config paths under one dir that declare a benchmark, discovered by content. - A config defines a benchmark iff it declares a `type: benchmark` dataset (see `BenchmarkConfig`), - regardless of its filename. So discovery is content-based: scan every yaml and keep the ones that - literally declare such a dataset. That text check is a cheap prefilter so we only pay the resolve - cost on real candidates (not every prompt/endpoint yaml), and it finds benchmarks whose config - isn't named `config.yaml` — e.g. tau2's `configs/*.yaml` and livecodebench's `cascade.yaml`. - Returns an empty list if the directory is missing. + A config is a benchmark iff it declares a `type: benchmark` dataset, regardless of filename, so we scan + every yaml. The `type: benchmark` text check is a cheap prefilter (pay the resolve cost only on real + candidates) that also catches non-`config.yaml` names like tau2's `configs/*.yaml`. Empty if dir missing. """ if not benchmarks_dir.is_dir(): return [] @@ -140,9 +138,17 @@ def _benchmark_config_paths(benchmarks_dir: Path) -> List[Path]: return sorted(p for p in config_paths if "type: benchmark" in p.read_text(errors="ignore")) -def discover_benchmarks() -> Dict[str, BenchmarkConfig]: - """Map benchmark name -> :class:`BenchmarkConfig` for every benchmark config under ``benchmarks/``.""" - return _load_benchmarks_from_config_paths(_benchmark_config_paths(BENCHMARKS_DIR)) +def discover_benchmarks(search_dirs: Optional[Union[Path, Sequence[Path]]] = None) -> Dict[str, BenchmarkConfig]: + """Map benchmark name -> :class:`BenchmarkConfig` for every discoverable benchmark config. + + Scans the ``benchmarks/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` root + (``search_dirs`` + cwd + built-ins), merged so user benchmarks shadow same-named built-ins. + ``search_dirs`` is one dir or a list. + """ + return merge_by_name( + _load_benchmarks_from_config_paths(_benchmark_config_paths(root / BENCHMARKS_SUBDIR)) + for root in component_search_roots(search_dirs) + ) # Backward-compatibility shims (CLI refactor): these symbols moved to `nemo_gym.cli.eval`. diff --git a/nemo_gym/cli/agents.py b/nemo_gym/cli/agents.py index 07f99f0704..e4264e79f7 100644 --- a/nemo_gym/cli/agents.py +++ b/nemo_gym/cli/agents.py @@ -22,18 +22,16 @@ from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( JSON_OUTPUT_KEY_NAME, + SEARCH_DIR_KEY_NAME, GlobalConfigDictParserConfig, get_global_config_dict, ) def list_agents() -> None: - """CLI command: list discovered agent harnesses and how each composes (Pattern A vs B). - - Complements ``gym list benchmarks``: the asset selectors resolve a component *by name*, but only - this listing surfaces which agents are freely wireable into a separate environment (Pattern A) - versus self-contained harnesses that run with their own config (Pattern B) — the distinction the - config composer's compatibility guard relies on. + """List discovered agent harnesses and how each composes: freely wireable into a separate environment + (Pattern A) vs. self-contained harnesses that run with their own config (Pattern B). ``--search-dir`` + adds extra roots to scan on top of the cwd and built-ins. """ global_config_dict = get_global_config_dict( global_config_dict_parser_config=GlobalConfigDictParserConfig( @@ -42,7 +40,7 @@ def list_agents() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - agents = discover_agents() + agents = discover_agents(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): payload = [ diff --git a/nemo_gym/cli/env.py b/nemo_gym/cli/env.py index 9b11afce5f..5127c2ed8a 100644 --- a/nemo_gym/cli/env.py +++ b/nemo_gym/cli/env.py @@ -46,6 +46,7 @@ NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME, NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME, NEMO_GYM_RESERVED_TOP_LEVEL_KEYS, + SEARCH_DIR_KEY_NAME, GlobalConfigDictParser, GlobalConfigDictParserConfig, get_global_config_dict, @@ -921,11 +922,14 @@ def validate(): def list_environments() -> None: """List the environments available under environments/, by short name. + ``--search-dir`` adds extra roots to scan on top of the cwd and built-ins. + Examples: ```bash gym list environments gym list environments --json + gym list environments --search-dir /path/to/project ``` """ global_config_dict = get_global_config_dict( @@ -935,7 +939,7 @@ def list_environments() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - environments = discover_environments() + environments = discover_environments(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): print( diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index 34dfb8287c..8eaf996226 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -41,6 +41,7 @@ JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, + SEARCH_DIR_KEY_NAME, TASK_INDEX_KEY_NAME, GlobalConfigDictParserConfig, get_first_server_config_dict, @@ -77,7 +78,8 @@ def list_benchmarks() -> None: """CLI command: list available benchmarks, optionally filtered by a `query` (the `gym search` entry point). A benchmark is a specific kind of environment, so it shares `gym list environments`' columns (name, - domain, description) and reads them through the same `read_config_metadata` helper. + domain, description) and reads them through the same `read_config_metadata` helper. ``--search-dir`` + adds extra roots to scan on top of the cwd and built-ins. """ global_config_dict = get_global_config_dict( global_config_dict_parser_config=GlobalConfigDictParserConfig( @@ -86,7 +88,7 @@ def list_benchmarks() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - benchmarks = discover_benchmarks() + benchmarks = discover_benchmarks(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) # Resolve domain + description once per benchmark, via the shared component-metadata reader — # the same one `gym list environments` uses — for the columns and `gym search`. diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 06d67ceb27..018e859bd2 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -21,7 +21,7 @@ from dataclasses import dataclass, field from pathlib import Path -from nemo_gym import PARENT_DIR, WORKING_DIR +from nemo_gym.discovery import component_search_roots VERSION_TARGET = "nemo_gym.cli.general:version" @@ -164,10 +164,10 @@ def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag: 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 the Gym install root (`PARENT_DIR` — where the built-in asset trees live in both - editable and wheel installs), then the current working directory (the user's project), then - any user-registered --search-dir roots. Searching `PARENT_DIR` is what lets built-ins resolve - by name from an arbitrary cwd (e.g. a wheel install), not just from inside the repo checkout. + Searches the roots from :func:`~nemo_gym.discovery.component_search_roots` (``--search-dir`` + cwd + + install root), the same helper that backs `gym list`/`gym search`, so config resolution and discovery + agree on where components live. Searching the install root is what lets built-ins resolve by name from + an arbitrary cwd (e.g. a wheel install), not just inside the repo checkout. """ parent, subdir, default_flavor = _ASSETS[flag] server_name, _, config_flavor = value.partition("/") @@ -175,15 +175,7 @@ def _asset_config_path(flag: str, value: str, search_dirs: tuple[str, ...] = ()) config_dir = f"{parent}/{server_name}/{subdir}".rstrip("/") path = f"{config_dir}/{config_flavor}.yaml" - # Search the install root (built-ins) and the user's cwd / --search-dir roots; dedupe roots that - # resolve to the same directory so an editable install run from the repo root isn't searched twice. - seen_roots: set[Path] = set() - roots: list[Path] = [] - for root in (PARENT_DIR, WORKING_DIR, Path.cwd(), *(Path(d) for d in search_dirs)): - resolved_root = root.resolve() - if resolved_root not in seen_roots: - seen_roots.add(resolved_root) - roots.append(root) + roots = component_search_roots(search_dirs) matches: list[Path] = [] for root in roots: @@ -245,8 +237,8 @@ def _asset_selector(flag: str) -> Flag: RESOURCES_SERVER_CONFIG = _asset_selector("resources-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` for the asset selectors above: read straight from argv during config resolution, not +# emitted as a Hydra override. On every command that accepts a -- NAME selector. SEARCH_DIR = Flag( register=lambda p: p.add_argument( "--search-dir", @@ -256,6 +248,22 @@ def _asset_selector(flag: str) -> Flag: ), ) +# `--search-dir` for the `list`/`search` commands. Their targets are called with no args, so — like +# --json/--query — the roots reach them as the reserved `search_dir` config key, not read from argv. +DISCOVERY_SEARCH_DIR = Flag( + register=lambda p: p.add_argument( + "--search-dir", + action="append", + metavar="DIR", + help="Extra root directory to search for components; repeatable.", + ), + translate_to_hydra=lambda args: ( + [f"+search_dir=[{','.join(getattr(args, 'search_dir', None) or [])}]"] + if getattr(args, "search_dir", None) + else [] + ), +) + def _merge_config_paths(overrides: list[str]) -> list[str]: """Coalesce all `+config_paths=[...]` tokens (from --config and asset selectors) into one (Hydra rejects dupes).""" @@ -312,20 +320,24 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: # NOTE: none of the flags are argparse-required (every value can also be supplied as a Hydra `+key=value` override). COMMANDS = { "list benchmarks": Command( - target="nemo_gym.cli.eval:list_benchmarks", summary="List available benchmarks.", flags=(JSON,) + target="nemo_gym.cli.eval:list_benchmarks", + summary="List available benchmarks.", + flags=(JSON, DISCOVERY_SEARCH_DIR), ), "list environments": Command( - target="nemo_gym.cli.env:list_environments", summary="List available environments by name.", flags=(JSON,) + target="nemo_gym.cli.env:list_environments", + summary="List available environments by name.", + flags=(JSON, DISCOVERY_SEARCH_DIR), ), "list agents": Command( target="nemo_gym.cli.agents:list_agents", summary="List agent harnesses and how each composes (Pattern A vs self-contained B).", - flags=(JSON,), + flags=(JSON, DISCOVERY_SEARCH_DIR), ), "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), + flags=(QUERY, JSON, DISCOVERY_SEARCH_DIR), ), "dataset upload": Command( target=_dataset_upload, diff --git a/nemo_gym/discovery.py b/nemo_gym/discovery.py index 83766e0a90..2f643ee740 100644 --- a/nemo_gym/discovery.py +++ b/nemo_gym/discovery.py @@ -12,23 +12,22 @@ # 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. -"""Shared component-discovery utilities for the ``gym list``/``gym search`` commands. +"""Shared component-discovery helpers: which roots to scan for a component, how to resolve name +collisions across them, and how to read a component's ``(domain, description)``. -An environment and a benchmark (a benchmark is a specific kind of environment) are listed with the -same columns, read the same way. This module holds the config-reading code both listings share so the -per-component registries (``registry.py``, ``benchmarks.py``, ``agent_registry.py``) can depend on it -without depending on each other. It only reads config files — it never starts servers — and only -imports lower-level config utilities, so it sits below those registries in the import graph. +Lives below the per-component registries (``registry.py``, ``benchmarks.py``, ``agent_registry.py``) so +they can share it without depending on each other. Reads configs only; never starts servers. """ import re from copy import deepcopy from pathlib import Path -from typing import Optional, Tuple +from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TypeVar, Union from omegaconf import DictConfig, OmegaConf from omegaconf.errors import InterpolationKeyError +from nemo_gym import PARENT_DIR, WORKING_DIR from nemo_gym.global_config import ( POLICY_MODEL_KEY_NAME, GlobalConfigDictParser, @@ -36,6 +35,47 @@ ) +_T = TypeVar("_T") + + +def component_search_roots(search_dirs: Optional[Union[Path, Sequence[Path]]] = None) -> List[Path]: + """Ordered, de-duplicated roots to look for a Gym component under: any ``search_dirs`` (one dir or a + list, e.g. from ``--search-dir``), then cwd, then ``WORKING_DIR`` and the install root (``PARENT_DIR``, + the built-ins). + + Earlier roots win on a name collision (see :func:`merge_by_name`), so user components shadow built-ins. + De-duplicated by resolved path, since cwd/``WORKING_DIR``/install root coincide in an editable checkout. + The single source of truth for where Gym looks for components — used by both config resolution + (``_asset_config_path``) and the ``gym list``/``gym search`` discovery functions. + """ + if search_dirs is None: + extra: List[Path] = [] + elif isinstance(search_dirs, (str, Path)): + extra = [Path(search_dirs)] # a single dir + else: + extra = [Path(d) for d in search_dirs] # a list of dirs + candidates: List[Path] = [*extra, Path.cwd(), WORKING_DIR, PARENT_DIR] + roots: List[Path] = [] + seen: set[Path] = set() + for root in candidates: + resolved = root.resolve() + if resolved not in seen: + seen.add(resolved) + roots.append(root) + return roots + + +def merge_by_name(per_root: Iterable[Dict[str, _T]]) -> Dict[str, _T]: + """Merge per-root ``name -> entry`` mappings; earlier roots win on a collision (user shadows built-in), + matching :func:`component_search_roots` precedence. Insertion order preserved. + """ + merged: Dict[str, _T] = {} + for entries in per_root: + for name, entry in entries.items(): + merged.setdefault(name, entry) + return merged + + # Fills unset `???`/`${...}` values during listing: they reference runtime-only values (API keys, # endpoints) not needed to identify a component, so a placeholder lets the config still resolve. _UNSET_VALUE_PLACEHOLDER = "__unset_for_listing__" @@ -46,12 +86,9 @@ def _parse_no_environment_tolerating_unset_values(initial_config_dict: DictConfig) -> DictConfig: - """`parse_no_environment` for *listing*: fill unset `???` values and undefined `${...}` interpolations - with a placeholder so a component referencing runtime-only values can still be identified. Never mutates - the caller's config; errors other than those two propagate. - - `???` is filled anywhere; `${...}` only where parse forces resolution (top-level and server sections), - not inside arbitrary non-server nested dicts — fine for listing, whose interpolations live in servers. + """`parse_no_environment` for listing: fill unset `???` and undefined `${...}` values (runtime-only + things like API keys/endpoints) with a placeholder so the config still resolves enough to identify the + component. Never mutates the input; errors other than those two propagate. """ working = deepcopy(initial_config_dict) # never mutate the caller's config parser = GlobalConfigDictParser() @@ -78,10 +115,8 @@ def _parse_no_environment_tolerating_unset_values(initial_config_dict: DictConfi def _scan_servers_for_metadata(container) -> Tuple[Optional[str], Optional[str]]: - """Best-effort ``(domain, description)`` from a config mapping, scanning every server group. - - Reads the first ``domain`` and the first ``description`` found across all server instances. Defensive - against malformed shapes (non-mapping top level, a server group that isn't a dict) so it never raises. + """Best-effort ``(domain, description)`` from a config mapping: the first of each found across all + server groups. Defensive against malformed shapes, so it never raises. """ domain: Optional[str] = None description: Optional[str] = None @@ -105,20 +140,15 @@ def _scan_servers_for_metadata(container) -> Tuple[Optional[str], Optional[str]] def read_config_metadata(config_path: Path) -> Tuple[Optional[str], Optional[str]]: - """Shared listing metadata reader: ``(domain, description)`` for an environment *or* benchmark config. - - A benchmark is a specific kind of environment, so both are read the same way. Two-pass, because the two - kinds declare metadata differently: + """Shared ``(domain, description)`` reader for an environment *or* benchmark config. Two passes, because + the two declare metadata differently: - 1. **Raw (non-resolving) scan** of the config as written. Environment configs declare ``domain``/ - ``description`` inline, and reading them without resolution is safe even though an environment config - references model/agent servers defined elsewhere (resolving it in isolation would raise). - 2. **Resolving fallback**, only for whatever the raw scan left unset. Benchmark configs inherit their - metadata via ``config_paths``/``_inherit_from``, so it isn't present until the config is resolved. - Uses the listing-tolerant parser (unset runtime values are placeholdered); on any failure — e.g. an - environment config that references servers defined elsewhere — whatever the raw scan found is kept. + 1. Raw (non-resolving) scan — environment configs declare it inline, and this is safe even though they + reference servers defined elsewhere (resolving in isolation would raise). + 2. Resolving fallback for whatever's still unset — benchmark configs inherit it via + ``config_paths``/``_inherit_from``. Tolerates unset runtime values; on failure keeps the raw result. - Never raises: an unreadable or unresolvable config yields ``(None, None)``. + Never raises: an unreadable/unresolvable config yields ``(None, None)``. """ try: raw = OmegaConf.to_container(OmegaConf.load(config_path), resolve=False, throw_on_missing=False) diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 2eb1f67475..939dd16f44 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -89,6 +89,7 @@ QUERY_KEY_NAME = "query" OBSERVABILITY_ENABLED_KEY_NAME = "observability_enabled" MODEL_CALL_CAPTURE_DIR_KEY_NAME = "model_call_capture_dir" +SEARCH_DIR_KEY_NAME = "search_dir" NEMO_GYM_RESERVED_TOP_LEVEL_KEYS = [ CONFIG_PATHS_KEY_NAME, ENTRYPOINT_KEY_NAME, @@ -116,6 +117,7 @@ QUERY_KEY_NAME, OBSERVABILITY_ENABLED_KEY_NAME, MODEL_CALL_CAPTURE_DIR_KEY_NAME, + SEARCH_DIR_KEY_NAME, ] # Data keys diff --git a/nemo_gym/registry.py b/nemo_gym/registry.py index 53cb1da128..3e7906b1c3 100644 --- a/nemo_gym/registry.py +++ b/nemo_gym/registry.py @@ -20,21 +20,21 @@ ``gym list environments``. Resolving a name to a config path for *running* is handled by the CLI's generic ``--environment`` asset selector, so this module is intentionally discovery-only. -Discovery only reads config files and never starts servers. Its ``domain``/``description`` metadata is -read via the shared :func:`~nemo_gym.discovery.read_config_metadata` reader (the same one benchmark -listing uses), which tolerates unset secrets/API keys referenced by a config, so discovery is safe to -call even when those are not set in the environment. +Discovery only reads config files and never starts servers; ``domain``/``description`` come from the +shared :func:`~nemo_gym.discovery.read_config_metadata` reader, which tolerates unset secrets/API keys, +so it's safe to call even when those aren't set. """ from dataclasses import dataclass from pathlib import Path -from typing import Dict, Optional +from typing import Dict, Optional, Sequence, Union from nemo_gym import PARENT_DIR -from nemo_gym.discovery import read_config_metadata +from nemo_gym.discovery import component_search_roots, merge_by_name, read_config_metadata -ENVIRONMENTS_DIR = PARENT_DIR / "environments" +ENVIRONMENTS_SUBDIR = "environments" +ENVIRONMENTS_DIR = PARENT_DIR / ENVIRONMENTS_SUBDIR ENVIRONMENT_CONFIG_FILENAME = "config.yaml" @@ -49,8 +49,8 @@ class EnvironmentEntry: domain: Optional[str] = None -def discover_environments(environments_dir: Path = ENVIRONMENTS_DIR) -> Dict[str, EnvironmentEntry]: - """Map environment name -> :class:`EnvironmentEntry` for every ``/config.yaml``. +def _discover_environments_in_dir(environments_dir: Path) -> Dict[str, EnvironmentEntry]: + """Map environment name -> :class:`EnvironmentEntry` for every ``/config.yaml`` under one dir. The name is the directory name. Returns an empty dict if the directory is missing. """ @@ -73,3 +73,17 @@ def discover_environments(environments_dir: Path = ENVIRONMENTS_DIR) -> Dict[str ) return environments + + +def discover_environments( + search_dirs: Optional[Union[Path, Sequence[Path]]] = None, +) -> Dict[str, EnvironmentEntry]: + """Map environment name -> :class:`EnvironmentEntry` for every discoverable ``/config.yaml``. + + Scans the ``environments/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` root + (``search_dirs`` + cwd + built-ins), merged so user environments shadow same-named built-ins. + ``search_dirs`` is one dir or a list. + """ + return merge_by_name( + _discover_environments_in_dir(root / ENVIRONMENTS_SUBDIR) for root in component_search_roots(search_dirs) + ) diff --git a/tests/unit_tests/test_agent_registry.py b/tests/unit_tests/test_agent_registry.py index 7c2d8982d3..d2a0d27fdb 100644 --- a/tests/unit_tests/test_agent_registry.py +++ b/tests/unit_tests/test_agent_registry.py @@ -16,6 +16,7 @@ from nemo_gym.agent_registry import ( AgentEntry, + _discover_agents_in_dir, discover_agents, ) @@ -54,7 +55,7 @@ class TestDiscoverAgents: def test_discovers_and_classifies_pattern_a(self, tmp_path: Path) -> None: _make_agent(tmp_path, "simple_agent", configs={"simple_agent": _pattern_a()}) - agents = discover_agents(tmp_path) + agents = _discover_agents_in_dir(tmp_path) assert set(agents) == {"simple_agent"} entry = agents["simple_agent"] @@ -65,7 +66,7 @@ def test_discovers_and_classifies_pattern_a(self, tmp_path: Path) -> None: def test_classifies_pattern_b_as_not_composable(self, tmp_path: Path) -> None: _make_agent(tmp_path, "swe_agents", configs={"swebench": _pattern_b()}) - assert discover_agents(tmp_path)["swe_agents"].self_contained is True + assert _discover_agents_in_dir(tmp_path)["swe_agents"].self_contained is True def test_external_harness_agent_is_not_composable(self, tmp_path: Path) -> None: body = ( @@ -75,12 +76,12 @@ def test_external_harness_agent_is_not_composable(self, tmp_path: Path) -> None: _make_agent(tmp_path, "claude_code_agent", configs={"claude_code_agent": body}) # Has a resources_server but drives an external LLM harness -> not composable. - assert discover_agents(tmp_path)["claude_code_agent"].self_contained is True + assert _discover_agents_in_dir(tmp_path)["claude_code_agent"].self_contained is True def test_zero_config_agent_is_discovered_and_defaults_composable(self, tmp_path: Path) -> None: _make_agent(tmp_path, "aviary_agent", configs=None) # app.py only, no configs - entry = discover_agents(tmp_path)["aviary_agent"] + entry = _discover_agents_in_dir(tmp_path)["aviary_agent"] assert entry.config_paths == () assert entry.self_contained is False @@ -94,30 +95,33 @@ def test_multiple_variants_are_all_recorded(self, tmp_path: Path) -> None: }, ) - assert set(discover_agents(tmp_path)["langgraph_agent"].variants) == {"orchestrator_agent", "rewoo_agent"} + assert set(_discover_agents_in_dir(tmp_path)["langgraph_agent"].variants) == { + "orchestrator_agent", + "rewoo_agent", + } def test_non_agent_yaml_is_filtered_out(self, tmp_path: Path) -> None: # A configs/ file that is not a gym agent config (no responses_api_agents) is ignored; # the dir still counts as an agent because of app.py. _make_agent(tmp_path, "swe_agents", configs={"raw_harness": "agent:\n type: openhands\n"}) - entry = discover_agents(tmp_path)["swe_agents"] + entry = _discover_agents_in_dir(tmp_path)["swe_agents"] assert entry.config_paths == () def test_directory_without_app_or_configs_is_skipped(self, tmp_path: Path) -> None: (tmp_path / "not_an_agent").mkdir() (tmp_path / "loose_file.txt").write_text("x") - assert discover_agents(tmp_path) == {} + assert _discover_agents_in_dir(tmp_path) == {} def test_unparseable_config_does_not_crash_discovery(self, tmp_path: Path) -> None: _make_agent(tmp_path, "broken", configs={"broken": "responses_api_agents: [unclosed\n"}) # The bad file is skipped (not an agent config); the dir survives via app.py. - assert discover_agents(tmp_path)["broken"].config_paths == () + assert _discover_agents_in_dir(tmp_path)["broken"].config_paths == () def test_missing_directory_yields_no_agents(self, tmp_path: Path) -> None: - assert discover_agents(tmp_path / "nope") == {} + assert _discover_agents_in_dir(tmp_path / "nope") == {} class TestRealAgents: diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index 292edd6406..8c0904c891 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -1053,8 +1053,8 @@ def test_builtin_resolves_from_install_root_when_cwd_differs(self, monkeypatch: install_root.mkdir() user_cwd.mkdir() self._make_resources_server(install_root) # built-in only under the install root - monkeypatch.setattr(cli_main, "PARENT_DIR", install_root) - monkeypatch.setattr(cli_main, "WORKING_DIR", user_cwd) + monkeypatch.setattr("nemo_gym.discovery.PARENT_DIR", install_root) + monkeypatch.setattr("nemo_gym.discovery.WORKING_DIR", user_cwd) monkeypatch.chdir(user_cwd) resolved = cli_main._asset_config_path("resources-server", "foo") @@ -1066,8 +1066,8 @@ def test_user_cwd_asset_resolves_when_not_builtin(self, monkeypatch: MonkeyPatch install_root.mkdir() user_cwd.mkdir() self._make_resources_server(user_cwd, name="myenv") # exists only in the user's project - monkeypatch.setattr(cli_main, "PARENT_DIR", install_root) - monkeypatch.setattr(cli_main, "WORKING_DIR", user_cwd) + monkeypatch.setattr("nemo_gym.discovery.PARENT_DIR", install_root) + monkeypatch.setattr("nemo_gym.discovery.WORKING_DIR", user_cwd) monkeypatch.chdir(user_cwd) resolved = cli_main._asset_config_path("resources-server", "myenv") @@ -1081,8 +1081,8 @@ def test_same_name_in_install_root_and_cwd_is_ambiguous(self, monkeypatch: Monke user_cwd.mkdir() self._make_resources_server(install_root) self._make_resources_server(user_cwd) - monkeypatch.setattr(cli_main, "PARENT_DIR", install_root) - monkeypatch.setattr(cli_main, "WORKING_DIR", user_cwd) + monkeypatch.setattr("nemo_gym.discovery.PARENT_DIR", install_root) + monkeypatch.setattr("nemo_gym.discovery.WORKING_DIR", user_cwd) monkeypatch.chdir(user_cwd) with pytest.raises(ValueError, match="ambiguous"): @@ -1094,8 +1094,8 @@ def test_editable_layout_single_root_not_self_ambiguous(self, monkeypatch: Monke repo_root = tmp_path / "Gym" repo_root.mkdir() self._make_resources_server(repo_root) - monkeypatch.setattr(cli_main, "PARENT_DIR", repo_root) - monkeypatch.setattr(cli_main, "WORKING_DIR", repo_root) + monkeypatch.setattr("nemo_gym.discovery.PARENT_DIR", repo_root) + monkeypatch.setattr("nemo_gym.discovery.WORKING_DIR", repo_root) monkeypatch.chdir(repo_root) resolved = cli_main._asset_config_path("resources-server", "foo") @@ -1112,3 +1112,9 @@ def test_list_environments_json_dispatches(self, monkeypatch: MonkeyPatch) -> No target, overrides = _dispatch_for(monkeypatch, ["list", "environments", "--json"]) assert target == "nemo_gym.cli.env:list_environments" assert overrides == ["+json=true"] + + def test_search_dir_becomes_config_override(self, monkeypatch: MonkeyPatch) -> None: + # `--search-dir` (repeatable) reaches the no-arg list command as the reserved `search_dir` config + # key, read centrally from the resolved config — like --json/--query. + _, overrides = _dispatch_for(monkeypatch, ["list", "environments", "--search-dir", "/a", "--search-dir", "/b"]) + assert overrides == ["+search_dir=[/a,/b]"] diff --git a/tests/unit_tests/test_discovery.py b/tests/unit_tests/test_discovery.py index 64e2268b43..6b63a9bead 100644 --- a/tests/unit_tests/test_discovery.py +++ b/tests/unit_tests/test_discovery.py @@ -16,13 +16,64 @@ from omegaconf import OmegaConf +from nemo_gym import PARENT_DIR from nemo_gym.discovery import ( _UNSET_VALUE_PLACEHOLDER, _parse_no_environment_tolerating_unset_values, + component_search_roots, + merge_by_name, read_config_metadata, ) +class TestComponentSearchRoots: + def test_default_includes_cwd_and_install_root(self) -> None: + resolved = {root.resolve() for root in component_search_roots()} + assert Path.cwd().resolve() in resolved + assert PARENT_DIR.resolve() in resolved + + def test_search_dirs_take_precedence_and_keep_order(self, tmp_path: Path) -> None: + a = tmp_path / "a" + b = tmp_path / "b" + a.mkdir() + b.mkdir() + + roots = component_search_roots(search_dirs=[a, b]) + + assert roots[0] == a # explicit search dirs come first, in the given order + assert roots[1] == b + assert PARENT_DIR.resolve() in {root.resolve() for root in roots} # built-ins still scanned + + def test_accepts_a_single_dir(self, tmp_path: Path) -> None: + # `search_dirs` takes one dir or a list; a lone Path must be treated as that single root. + assert component_search_roots(tmp_path)[0] == tmp_path + + def test_dedupes_roots_by_resolved_path(self) -> None: + # Passing the install root as an explicit search dir must not scan it twice. + roots = component_search_roots(search_dirs=[PARENT_DIR]) + resolved = [root.resolve() for root in roots] + + assert resolved.count(PARENT_DIR.resolve()) == 1 + assert roots[0].resolve() == PARENT_DIR.resolve() # the explicit search dir still takes precedence + + +class TestMergeByName: + def test_merges_disjoint_roots(self) -> None: + assert merge_by_name([{"a": 1}, {"b": 2}]) == {"a": 1, "b": 2} + + def test_earlier_root_shadows_later_on_name_collision(self) -> None: + # A component found in an earlier root (e.g. the user's cwd) wins over a same-named one in a + # later root (e.g. a built-in) — the collision policy every `gym list` command relies on. + merged = merge_by_name([{"dup": "from_first"}, {"dup": "from_second", "other": "kept"}]) + assert merged == {"dup": "from_first", "other": "kept"} + + def test_preserves_order_within_and_across_roots(self) -> None: + assert list(merge_by_name([{"a": 1, "b": 2}, {"c": 3}])) == ["a", "b", "c"] + + def test_empty_input(self) -> None: + assert merge_by_name([]) == {} + + class TestTolerantInterpolationParse: # Unset `???` values and unresolved `${...}` interpolations reference runtime-only values that aren't # needed to identify a component; listing fills them with a placeholder so the config still resolves. diff --git a/tests/unit_tests/test_registry.py b/tests/unit_tests/test_registry.py index 9358f3f251..f9d8b837aa 100644 --- a/tests/unit_tests/test_registry.py +++ b/tests/unit_tests/test_registry.py @@ -14,7 +14,7 @@ # limitations under the License. from pathlib import Path -from nemo_gym.registry import discover_environments +from nemo_gym.registry import _discover_environments_in_dir, discover_environments def _make_env(environments_dir: Path, name: str, config_body: str) -> Path: @@ -50,7 +50,7 @@ def test_discovers_by_directory_name_with_metadata(self, tmp_path: Path) -> None _make_env(envs_dir, "alpha", _ENV_CONFIG.format(name="alpha")) _make_env(envs_dir, "beta", _ENV_CONFIG.format(name="beta")) - environments = discover_environments(envs_dir) + environments = _discover_environments_in_dir(envs_dir) assert set(environments) == {"alpha", "beta"} alpha = environments["alpha"] @@ -61,7 +61,7 @@ def test_discovers_by_directory_name_with_metadata(self, tmp_path: Path) -> None assert alpha.domain == "agent" def test_missing_directory_returns_empty(self, tmp_path: Path) -> None: - assert discover_environments(tmp_path / "does_not_exist") == {} + assert _discover_environments_in_dir(tmp_path / "does_not_exist") == {} def test_ignores_dirs_without_config_and_loose_files(self, tmp_path: Path) -> None: envs_dir = tmp_path / "environments" @@ -69,7 +69,7 @@ def test_ignores_dirs_without_config_and_loose_files(self, tmp_path: Path) -> No (envs_dir / "not_an_env").mkdir() # dir without a config.yaml (envs_dir / "__init__.py").write_text("") # loose file - assert set(discover_environments(envs_dir)) == {"real"} + assert set(_discover_environments_in_dir(envs_dir)) == {"real"} def test_unparseable_or_metadataless_configs_still_discovered(self, tmp_path: Path) -> None: # Configs without a parseable resources_servers block (or malformed YAML) must still be @@ -81,7 +81,7 @@ def test_unparseable_or_metadataless_configs_still_discovered(self, tmp_path: Pa _make_env(envs_dir, "rs_not_dict", "top:\n resources_servers: not_a_mapping\n") # rs not a dict _make_env(envs_dir, "broken", "key: [unclosed\n") # malformed YAML -> load raises - environments = discover_environments(envs_dir) + environments = _discover_environments_in_dir(envs_dir) assert set(environments) == {"no_rs", "top_list", "scalar_top", "rs_not_dict", "broken"} for entry in environments.values(): @@ -103,7 +103,7 @@ def test_metadata_tolerates_unset_interpolations(self, tmp_path: Path) -> None: " api_key: ${some_unset_key}\n", ) - environments = discover_environments(envs_dir) + environments = _discover_environments_in_dir(envs_dir) assert "needs_key" in environments assert environments["needs_key"].domain == "other" @@ -114,3 +114,19 @@ def test_workplace_assistant_is_discoverable(self) -> None: environments = discover_environments() assert "workplace_assistant" in environments assert environments["workplace_assistant"].config_path.name == "config.yaml" + + +class TestDiscoverEnvironmentsAcrossRoots: + def test_search_dirs_surface_user_environments_alongside_builtins(self, tmp_path: Path) -> None: + _make_env(tmp_path / "environments", "custom_env", _ENV_CONFIG.format(name="custom_env")) + + environments = discover_environments(search_dirs=[tmp_path]) + + assert "custom_env" in environments # a user-supplied environment is discovered + assert "workplace_assistant" in environments # ...alongside the built-ins + + def test_cwd_is_scanned_by_default(self, tmp_path: Path, monkeypatch) -> None: + _make_env(tmp_path / "environments", "cwd_env", _ENV_CONFIG.format(name="cwd_env")) + monkeypatch.chdir(tmp_path) + + assert "cwd_env" in discover_environments() From acb801a7589a8e8223a950eb2f471d5d1919cb42 Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Wed, 15 Jul 2026 09:31:31 +0200 Subject: [PATCH 03/13] chore: extract repeated code Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/agent_registry.py | 4 ++-- nemo_gym/benchmarks.py | 12 +++++++----- nemo_gym/discovery.py | 17 ++++++++++++++++- nemo_gym/registry.py | 6 ++---- 4 files changed, 27 insertions(+), 12 deletions(-) diff --git a/nemo_gym/agent_registry.py b/nemo_gym/agent_registry.py index 4101fba347..59294e35f1 100644 --- a/nemo_gym/agent_registry.py +++ b/nemo_gym/agent_registry.py @@ -44,7 +44,7 @@ from omegaconf import OmegaConf from nemo_gym import PARENT_DIR -from nemo_gym.discovery import component_search_roots, merge_by_name +from nemo_gym.discovery import discover_components AGENTS_SUBDIR = "responses_api_agents" @@ -171,4 +171,4 @@ def discover_agents( root (``search_dirs`` + cwd + built-ins), merged so user agents shadow same-named built-ins. ``search_dirs`` is one dir or a list. """ - return merge_by_name(_discover_agents_in_dir(root / AGENTS_SUBDIR) for root in component_search_roots(search_dirs)) + return discover_components(AGENTS_SUBDIR, _discover_agents_in_dir, search_dirs) diff --git a/nemo_gym/benchmarks.py b/nemo_gym/benchmarks.py index 45b701550f..3caf8881f4 100644 --- a/nemo_gym/benchmarks.py +++ b/nemo_gym/benchmarks.py @@ -24,7 +24,7 @@ from nemo_gym import PARENT_DIR from nemo_gym.config_types import BenchmarkDatasetConfig -from nemo_gym.discovery import _parse_no_environment_tolerating_unset_values, component_search_roots, merge_by_name +from nemo_gym.discovery import _parse_no_environment_tolerating_unset_values, discover_components from nemo_gym.global_config import ( POLICY_MODEL_KEY_NAME, GlobalConfigDictParser, @@ -138,6 +138,11 @@ def _benchmark_config_paths(benchmarks_dir: Path) -> List[Path]: return sorted(p for p in config_paths if "type: benchmark" in p.read_text(errors="ignore")) +def _discover_benchmarks_in_dir(benchmarks_dir: Path) -> Dict[str, BenchmarkConfig]: + """Map benchmark name -> :class:`BenchmarkConfig` for every benchmark config under one dir.""" + return _load_benchmarks_from_config_paths(_benchmark_config_paths(benchmarks_dir)) + + def discover_benchmarks(search_dirs: Optional[Union[Path, Sequence[Path]]] = None) -> Dict[str, BenchmarkConfig]: """Map benchmark name -> :class:`BenchmarkConfig` for every discoverable benchmark config. @@ -145,10 +150,7 @@ def discover_benchmarks(search_dirs: Optional[Union[Path, Sequence[Path]]] = Non (``search_dirs`` + cwd + built-ins), merged so user benchmarks shadow same-named built-ins. ``search_dirs`` is one dir or a list. """ - return merge_by_name( - _load_benchmarks_from_config_paths(_benchmark_config_paths(root / BENCHMARKS_SUBDIR)) - for root in component_search_roots(search_dirs) - ) + return discover_components(BENCHMARKS_SUBDIR, _discover_benchmarks_in_dir, search_dirs) # Backward-compatibility shims (CLI refactor): these symbols moved to `nemo_gym.cli.eval`. diff --git a/nemo_gym/discovery.py b/nemo_gym/discovery.py index 2f643ee740..ad5fbf38c0 100644 --- a/nemo_gym/discovery.py +++ b/nemo_gym/discovery.py @@ -22,7 +22,7 @@ import re from copy import deepcopy from pathlib import Path -from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TypeVar, Union +from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple, TypeVar, Union from omegaconf import DictConfig, OmegaConf from omegaconf.errors import InterpolationKeyError @@ -76,6 +76,21 @@ def merge_by_name(per_root: Iterable[Dict[str, _T]]) -> Dict[str, _T]: return merged +def discover_components( + subdir: str, + dir_scanning_fn: Callable[[Path], Dict[str, _T]], + search_dirs: Optional[Union[Path, Sequence[Path]]] = None, +) -> Dict[str, _T]: + """Run ``dir_scanning_fn`` on ``subdir`` of every :func:`component_search_roots` root and merge the results. + + The shared body of ``discover_environments``/``discover_agents``/``discover_models``/ + ``discover_benchmarks``: each passes its ``/`` subdir and a single-directory scan function, and + gets user-shadows-built-in merging (via :func:`merge_by_name`) for free. ``search_dirs`` is one dir or + a list. + """ + return merge_by_name(dir_scanning_fn(root / subdir) for root in component_search_roots(search_dirs)) + + # Fills unset `???`/`${...}` values during listing: they reference runtime-only values (API keys, # endpoints) not needed to identify a component, so a placeholder lets the config still resolve. _UNSET_VALUE_PLACEHOLDER = "__unset_for_listing__" diff --git a/nemo_gym/registry.py b/nemo_gym/registry.py index 3e7906b1c3..90123289db 100644 --- a/nemo_gym/registry.py +++ b/nemo_gym/registry.py @@ -30,7 +30,7 @@ from typing import Dict, Optional, Sequence, Union from nemo_gym import PARENT_DIR -from nemo_gym.discovery import component_search_roots, merge_by_name, read_config_metadata +from nemo_gym.discovery import discover_components, read_config_metadata ENVIRONMENTS_SUBDIR = "environments" @@ -84,6 +84,4 @@ def discover_environments( (``search_dirs`` + cwd + built-ins), merged so user environments shadow same-named built-ins. ``search_dirs`` is one dir or a list. """ - return merge_by_name( - _discover_environments_in_dir(root / ENVIRONMENTS_SUBDIR) for root in component_search_roots(search_dirs) - ) + return discover_components(ENVIRONMENTS_SUBDIR, _discover_environments_in_dir, search_dirs) From de0b55e7c27ff0f1a76db9194af58ccb086bfe3d Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Wed, 15 Jul 2026 10:01:17 +0200 Subject: [PATCH 04/13] feat: add gym list models Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/cli/main.py | 7 +- nemo_gym/cli/models.py | 65 +++++++++++++++++++ nemo_gym/model_registry.py | 85 +++++++++++++++++++++++++ tests/unit_tests/test_cli_main.py | 1 + tests/unit_tests/test_cli_models.py | 79 +++++++++++++++++++++++ tests/unit_tests/test_model_registry.py | 62 ++++++++++++++++++ 6 files changed, 298 insertions(+), 1 deletion(-) create mode 100644 nemo_gym/cli/models.py create mode 100644 nemo_gym/model_registry.py create mode 100644 tests/unit_tests/test_cli_models.py create mode 100644 tests/unit_tests/test_model_registry.py diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 018e859bd2..4e9a0a8d96 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -309,7 +309,7 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: # One-line help for each command group, shown in `gym --help`. GROUPS = { - "list": "List available components (benchmarks, agents, environments).", + "list": "List available components (benchmarks, environments, agents, models).", "dataset": "Manage datasets.", "env": "Develop and run environments.", "eval": "Run evaluations.", @@ -334,6 +334,11 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: summary="List agent harnesses and how each composes (Pattern A vs self-contained B).", flags=(JSON, DISCOVERY_SEARCH_DIR), ), + "list models": Command( + target="nemo_gym.cli.models:list_models", + summary="List model servers by the value to pass to --model-type.", + flags=(JSON, DISCOVERY_SEARCH_DIR), + ), "search": Command( target="nemo_gym.cli.eval:list_benchmarks", summary="Search available components (currently benchmarks) by name; like `list` filtered to a query.", diff --git a/nemo_gym/cli/models.py b/nemo_gym/cli/models.py new file mode 100644 index 0000000000..be0e060bd0 --- /dev/null +++ b/nemo_gym/cli/models.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 rich +from rich.table import Table + +from nemo_gym.cli.utils import print_rich_table +from nemo_gym.config_types import BaseNeMoGymCLIConfig +from nemo_gym.global_config import ( + JSON_OUTPUT_KEY_NAME, + SEARCH_DIR_KEY_NAME, + GlobalConfigDictParserConfig, + get_global_config_dict, +) +from nemo_gym.model_registry import discover_models + + +def list_models() -> None: + """List model servers, one row per ``--model-type`` value: ``Model`` is the token to pass (```` + for the default flavor, ``/`` for the rest); ``Model group`` is its model. ``--search-dir`` + adds extra roots on top of the cwd and built-ins. + """ + 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) + + models = discover_models(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + + # One row per passable `--model-type` value: `model` is the token, `model_group` its model. + rows = [ + {"model": model_type, "model_group": name} + for name, entry in models.items() + for model_type in entry.model_types + ] + + if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): + print(json.dumps(rows)) + return + + if not rows: + rich.print("No models found.") + return + + table = Table(title="NeMo Gym models") + table.add_column("Model", style="bold") + table.add_column("Model group") + for row in rows: + table.add_row(row["model"], row["model_group"]) + print_rich_table(table) diff --git a/nemo_gym/model_registry.py b/nemo_gym/model_registry.py new file mode 100644 index 0000000000..216033696d --- /dev/null +++ b/nemo_gym/model_registry.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. +"""Registry of model servers under ``responses_api_models//``. + +Maps each model dir to the config flavors it ships (``configs/.yaml``), so they can be +enumerated by the token passed to ``--model-type`` (see :attr:`ModelEntry.model_types`). Reads the +directory tree only; never loads a config. +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple, Union + +from nemo_gym import PARENT_DIR +from nemo_gym.discovery import discover_components + + +MODELS_SUBDIR = "responses_api_models" +MODELS_DIR = PARENT_DIR / MODELS_SUBDIR +MODEL_CONFIGS_SUBDIR = "configs" + + +@dataclass(frozen=True) +class ModelEntry: + """A discovered model server: its name, where it lives, and its config flavors.""" + + name: str + path: Path + config_paths: Tuple[Path, ...] # flavor config files, sorted (a model always ships at least one) + + @property + def variants(self) -> Dict[str, Path]: + """Map flavor name (config filename stem) -> config path.""" + return {path.stem: path for path in self.config_paths} + + @property + def model_types(self) -> List[str]: + """The tokens accepted by ``--model-type``: ```` for the flavor named after the model (the + default the selector resolves), ``/`` for the rest.""" + return [self.name if stem == self.name else f"{self.name}/{stem}" for stem in sorted(self.variants)] + + +def _discover_models_in_dir(models_dir: Path) -> Dict[str, ModelEntry]: + """Map model name -> :class:`ModelEntry` for every model dir under one ``responses_api_models/`` dir. + + The name is the directory name. A directory is a model iff it ships at least one ``configs/*.yaml`` — + the config a user passes to ``--model-type`` (a config-less dir has nothing to select, so it isn't + listed). Returns an empty dict if the directory is missing. + """ + models: Dict[str, ModelEntry] = {} + if not models_dir.is_dir(): + return models + + for child in sorted(models_dir.iterdir()): + if not child.is_dir(): + continue + configs_dir = child / MODEL_CONFIGS_SUBDIR + config_files = tuple(sorted(configs_dir.glob("*.yaml"))) if configs_dir.is_dir() else () + if not config_files: + continue + models[child.name] = ModelEntry(name=child.name, path=child, config_paths=config_files) + + return models + + +def discover_models(search_dirs: Optional[Union[Path, Sequence[Path]]] = None) -> Dict[str, ModelEntry]: + """Map model name -> :class:`ModelEntry` for every discoverable model server. + + Scans the ``responses_api_models/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` + root (``search_dirs`` + cwd + built-ins), merged so user models shadow same-named built-ins. + ``search_dirs`` is one dir or a list. + """ + return discover_components(MODELS_SUBDIR, _discover_models_in_dir, search_dirs) diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index 8c0904c891..399980cc42 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -645,6 +645,7 @@ class TestJsonFlag: [ (["list", "benchmarks", "--json"], "nemo_gym.cli.eval:list_benchmarks"), (["list", "agents", "--json"], "nemo_gym.cli.agents:list_agents"), + (["list", "models", "--json"], "nemo_gym.cli.models:list_models"), (["env", "status", "--json"], "nemo_gym.cli.env:status"), ], ) diff --git a/tests/unit_tests/test_cli_models.py b/tests/unit_tests/test_cli_models.py new file mode 100644 index 0000000000..d5b3436e59 --- /dev/null +++ b/tests/unit_tests/test_cli_models.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 +from pathlib import Path +from unittest.mock import patch + +from omegaconf import OmegaConf + +from nemo_gym.cli.models import list_models +from nemo_gym.model_registry import ModelEntry + + +def _mock_global_config(config: dict = None): + return OmegaConf.create(config or {}) + + +def _entry(name: str, flavors=()) -> ModelEntry: + path = Path("responses_api_models") / name + config_paths = tuple(path / "configs" / f"{f}.yaml" for f in flavors) + return ModelEntry(name=name, path=path, config_paths=config_paths) + + +_MODELS = { + "my_model": _entry("my_model", flavors=("my_model", "some_other_flavor")), + "another_model": _entry("another_model", flavors=("another_model",)), +} + + +class TestListModels: + def test_lists_per_variant_rows(self, capsys) -> None: + with ( + patch("nemo_gym.cli.models.get_global_config_dict", return_value=_mock_global_config()), + patch("nemo_gym.cli.models.discover_models", return_value=_MODELS), + ): + list_models() + out = capsys.readouterr().out + + variants = [token for entry in _MODELS.values() for token in entry.model_types] + assert len(variants) == 3 + # one data row per variant (data rows use the light "│"; the header uses the heavy "┃") + assert sum(1 for line in out.splitlines() if "│" in line) == 3 + for variant in variants: + assert variant in out + + def test_no_models(self, capsys) -> None: + with ( + patch("nemo_gym.cli.models.get_global_config_dict", return_value=_mock_global_config()), + patch("nemo_gym.cli.models.discover_models", return_value={}), + ): + list_models() + assert "No models found" in capsys.readouterr().out + + def test_json_output_is_per_variant_rows(self, capsys) -> None: + with ( + patch("nemo_gym.cli.models.get_global_config_dict", return_value=_mock_global_config({"json": True})), + patch("nemo_gym.cli.models.discover_models", return_value=_MODELS), + ): + list_models() + payload = json.loads(capsys.readouterr().out) + expected = [ + {"model": "my_model", "model_group": "my_model"}, + {"model": "my_model/some_other_flavor", "model_group": "my_model"}, + {"model": "another_model", "model_group": "another_model"}, + ] + assert len(payload) == len(expected) + for row in expected: + assert row in payload diff --git a/tests/unit_tests/test_model_registry.py b/tests/unit_tests/test_model_registry.py new file mode 100644 index 0000000000..305b6c6093 --- /dev/null +++ b/tests/unit_tests/test_model_registry.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 pathlib import Path + +from nemo_gym.model_registry import _discover_models_in_dir + + +def _make_model(models_dir: Path, name: str, *, app: bool = True, flavors=()) -> Path: + model_dir = models_dir / name + model_dir.mkdir(parents=True) + if app: + (model_dir / "app.py").write_text("# app\n") + if flavors: + configs_dir = model_dir / "configs" + configs_dir.mkdir() + for flavor in flavors: + (configs_dir / f"{flavor}.yaml").write_text("{}\n") + return model_dir + + +class TestDiscoverModels: + def test_discovers_by_directory_name(self, tmp_path: Path) -> None: + _make_model(tmp_path, "my_model", flavors=("my_model", "some_other_flavor")) + _make_model(tmp_path, "another_model", flavors=("another_model",)) + + models = _discover_models_in_dir(tmp_path) + + assert set(models) == {"my_model", "another_model"} + assert list(models["my_model"].variants) == ["my_model", "some_other_flavor"] + + def test_model_types_are_the_model_type_tokens(self, tmp_path: Path) -> None: + # The flavor named after the model is the default (bare ``); others are `/`. + _make_model(tmp_path, "my_model", flavors=("my_model", "some_other_flavor")) + + assert _discover_models_in_dir(tmp_path)["my_model"].model_types == [ + "my_model", + "my_model/some_other_flavor", + ] + + def test_dirs_without_a_config_are_skipped(self, tmp_path: Path) -> None: + # Only a dir that ships a config (something to pass to --model-type) is a model: a stray .egg-info, + # or a dir with just an app.py and no configs, has nothing selectable and is not listed. + (tmp_path / "my_model.egg-info").mkdir() + _make_model(tmp_path, "app_only_model", app=True, flavors=()) + _make_model(tmp_path, "another_model", flavors=("another_model",)) + + assert set(_discover_models_in_dir(tmp_path)) == {"another_model"} + + def test_missing_directory_yields_no_models(self, tmp_path: Path) -> None: + assert _discover_models_in_dir(tmp_path / "nope") == {} From b73fd16a7c228ffd7f845668ea91a86ae2e14f3e Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Wed, 15 Jul 2026 10:41:38 +0200 Subject: [PATCH 05/13] feat: add gym search Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/cli/agents.py | 18 ++++++++++----- nemo_gym/cli/env.py | 24 +++++++++++++++----- nemo_gym/cli/eval.py | 28 +++-------------------- nemo_gym/cli/main.py | 35 ++++++++++++++++++++++++----- nemo_gym/cli/models.py | 18 ++++++++++----- nemo_gym/cli/utils.py | 34 ++++++++++++++++++++++++++++ tests/unit_tests/test_benchmarks.py | 25 +-------------------- tests/unit_tests/test_cli.py | 23 +++++++++++++++++++ tests/unit_tests/test_cli_agents.py | 11 +++++++++ tests/unit_tests/test_cli_main.py | 19 ++++++++++++++-- tests/unit_tests/test_cli_models.py | 14 ++++++++++++ tests/unit_tests/test_cli_utils.py | 20 ++++++++++++++++- 12 files changed, 193 insertions(+), 76 deletions(-) diff --git a/nemo_gym/cli/agents.py b/nemo_gym/cli/agents.py index e4264e79f7..46d25cd845 100644 --- a/nemo_gym/cli/agents.py +++ b/nemo_gym/cli/agents.py @@ -14,14 +14,14 @@ # limitations under the License. import json -import rich from rich.table import Table from nemo_gym.agent_registry import discover_agents -from nemo_gym.cli.utils import print_rich_table +from nemo_gym.cli.utils import fuzzy_matches, print_no_matches, print_rich_table from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( JSON_OUTPUT_KEY_NAME, + QUERY_KEY_NAME, SEARCH_DIR_KEY_NAME, GlobalConfigDictParserConfig, get_global_config_dict, @@ -30,8 +30,9 @@ def list_agents() -> None: """List discovered agent harnesses and how each composes: freely wireable into a separate environment - (Pattern A) vs. self-contained harnesses that run with their own config (Pattern B). ``--search-dir`` - adds extra roots to scan on top of the cwd and built-ins. + (Pattern A) vs. self-contained harnesses that run with their own config (Pattern B). Optionally filtered + by a `query` (the `gym search agents` entry point). ``--search-dir`` adds extra roots on top of the cwd + and built-ins. """ global_config_dict = get_global_config_dict( global_config_dict_parser_config=GlobalConfigDictParserConfig( @@ -42,6 +43,11 @@ def list_agents() -> None: agents = discover_agents(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + # `gym search agents ` reuses this command, narrowing to fuzzy matches on name + variant names. + query = global_config_dict.get(QUERY_KEY_NAME) + if query: + agents = {name: entry for name, entry in agents.items() if fuzzy_matches(query, name, *entry.variants)} + if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): payload = [ { @@ -57,10 +63,10 @@ def list_agents() -> None: return if not agents: - rich.print("No agents found.") + print_no_matches("agents", query) return - table = Table(title="NeMo Gym agents") + table = Table(title=f"Agents matching '{query}'" if query else "NeMo Gym agents") table.add_column("agent", style="bold") table.add_column("composition") table.add_column("variants") diff --git a/nemo_gym/cli/env.py b/nemo_gym/cli/env.py index 5127c2ed8a..c29f575a04 100644 --- a/nemo_gym/cli/env.py +++ b/nemo_gym/cli/env.py @@ -38,7 +38,7 @@ from nemo_gym import PARENT_DIR, ROOT_DIR from nemo_gym.cli.setup_command import run_command, setup_env_command -from nemo_gym.cli.utils import exit_cleanly_on_config_error, print_rich_table +from nemo_gym.cli.utils import exit_cleanly_on_config_error, fuzzy_matches, print_no_matches, print_rich_table from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( DRY_RUN_KEY_NAME, @@ -46,6 +46,7 @@ NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME, NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME, NEMO_GYM_RESERVED_TOP_LEVEL_KEYS, + QUERY_KEY_NAME, SEARCH_DIR_KEY_NAME, GlobalConfigDictParser, GlobalConfigDictParserConfig, @@ -920,9 +921,8 @@ def validate(): def list_environments() -> None: - """List the environments available under environments/, by short name. - - ``--search-dir`` adds extra roots to scan on top of the cwd and built-ins. + """List the environments available under environments/, optionally filtered by a `query` (the + `gym search environments` entry point). ``--search-dir`` adds extra roots on top of the cwd and built-ins. Examples: @@ -941,6 +941,13 @@ def list_environments() -> None: environments = discover_environments(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + # `gym search environments ` reuses this command, narrowing to fuzzy matches on name + domain. + query = global_config_dict.get(QUERY_KEY_NAME) + if query: + environments = { + name: env for name, env in environments.items() if fuzzy_matches(query, name, env.domain or "") + } + if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): print( json.dumps( @@ -953,10 +960,15 @@ def list_environments() -> None: return if not environments: - rich.print("[yellow]No environments found.[/yellow]") + print_no_matches("environments", query) return - table = Table(title=f"Available environments in NeMo Gym ({len(environments)})") + title = ( + f"Environments matching '{query}' ({len(environments)})" + if query + else f"Available environments in NeMo Gym ({len(environments)})" + ) + table = Table(title=title) table.add_column("Name") table.add_column("Domain") table.add_column("Description") diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index 8eaf996226..370a428a6e 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -14,7 +14,6 @@ # limitations under the License. import asyncio -import difflib import importlib import json from copy import deepcopy @@ -22,19 +21,17 @@ from pathlib import Path from typing import Any, Dict, List, Tuple -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, discover_benchmarks, ) from nemo_gym.cli.env import RunHelper -from nemo_gym.cli.utils import exit_cleanly_on_config_error, print_rich_table +from nemo_gym.cli.utils import exit_cleanly_on_config_error, fuzzy_matches, print_no_matches, print_rich_table from nemo_gym.config_types import BaseNeMoGymCLIConfig, BenchmarkDatasetConfig, ConfigError, ConfigPathNotFoundError from nemo_gym.discovery import read_config_metadata from nemo_gym.global_config import ( @@ -59,21 +56,6 @@ 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.70): - return True - return False - - def list_benchmarks() -> None: """CLI command: list available benchmarks, optionally filtered by a `query` (the `gym search` entry point). @@ -99,7 +81,7 @@ def list_benchmarks() -> None: query = global_config_dict.get(QUERY_KEY_NAME) if query: benchmarks = { - name: bench for name, bench in benchmarks.items() if _fuzzy_matches(query, name, metadata[name][0] or "") + name: bench for name, bench in benchmarks.items() if fuzzy_matches(query, name, metadata[name][0] or "") } if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): @@ -117,11 +99,7 @@ def list_benchmarks() -> None: 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}") + print_no_matches("benchmarks", query) return title = ( diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 4e9a0a8d96..238db5da34 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -144,13 +144,36 @@ def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag: # global_config_dict.get(JSON_OUTPUT_KEY_NAME) (see general.py, eval.py, env.py). JSON = _bool_flag("json", "json", "Output as machine-readable JSON.") -# 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."), +# `gym search [] `: an optional component type plus the query. The query is surfaced to the +# chosen listing command as the reserved `query` config key; the type only picks which command to run +# (see `_search`). A lone positional is the query, defaulting to benchmarks — backward compatible. +_SEARCHABLE_TYPES = { + "benchmarks": "nemo_gym.cli.eval:list_benchmarks", + "environments": "nemo_gym.cli.env:list_environments", + "agents": "nemo_gym.cli.agents:list_agents", + "models": "nemo_gym.cli.models:list_models", +} + +SEARCH_TERMS = Flag( + register=lambda p: ( + p.add_argument( + "component_type", + nargs="?", + choices=list(_SEARCHABLE_TYPES), + help="Component type to search (default: benchmarks).", + ), + 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 [], ) +def _search(args: argparse.Namespace, overrides: list[str]) -> None: + """`gym search [] `: dispatch to the chosen type's listing command (default benchmarks), + which filters itself to the `query` config key already in `overrides`.""" + dispatch(_SEARCHABLE_TYPES[getattr(args, "component_type", None) or "benchmarks"], overrides) + + # 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 = { @@ -340,9 +363,9 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: flags=(JSON, DISCOVERY_SEARCH_DIR), ), "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, DISCOVERY_SEARCH_DIR), + target=_search, + summary="Search a component type (default benchmarks) by name; like `list` filtered to a query.", + flags=(SEARCH_TERMS, JSON, DISCOVERY_SEARCH_DIR), ), "dataset upload": Command( target=_dataset_upload, diff --git a/nemo_gym/cli/models.py b/nemo_gym/cli/models.py index be0e060bd0..270d6a66e4 100644 --- a/nemo_gym/cli/models.py +++ b/nemo_gym/cli/models.py @@ -14,13 +14,13 @@ # limitations under the License. import json -import rich from rich.table import Table -from nemo_gym.cli.utils import print_rich_table +from nemo_gym.cli.utils import fuzzy_matches, print_no_matches, print_rich_table from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( JSON_OUTPUT_KEY_NAME, + QUERY_KEY_NAME, SEARCH_DIR_KEY_NAME, GlobalConfigDictParserConfig, get_global_config_dict, @@ -30,8 +30,9 @@ def list_models() -> None: """List model servers, one row per ``--model-type`` value: ``Model`` is the token to pass (```` - for the default flavor, ``/`` for the rest); ``Model group`` is its model. ``--search-dir`` - adds extra roots on top of the cwd and built-ins. + for the default flavor, ``/`` for the rest); ``Model group`` is its model. Optionally + filtered by a `query` (the `gym search models` entry point). ``--search-dir`` adds extra roots on top of + the cwd and built-ins. """ global_config_dict = get_global_config_dict( global_config_dict_parser_config=GlobalConfigDictParserConfig( @@ -49,15 +50,20 @@ def list_models() -> None: for model_type in entry.model_types ] + # `gym search models ` reuses this command, narrowing to rows matching the token or its model. + query = global_config_dict.get(QUERY_KEY_NAME) + if query: + rows = [row for row in rows if fuzzy_matches(query, row["model"], row["model_group"])] + if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): print(json.dumps(rows)) return if not rows: - rich.print("No models found.") + print_no_matches("models", query) return - table = Table(title="NeMo Gym models") + table = Table(title=f"Models matching '{query}'" if query else "NeMo Gym models") table.add_column("Model", style="bold") table.add_column("Model group") for row in rows: diff --git a/nemo_gym/cli/utils.py b/nemo_gym/cli/utils.py index 28e943f229..8a6e767032 100644 --- a/nemo_gym/cli/utils.py +++ b/nemo_gym/cli/utils.py @@ -12,6 +12,40 @@ # 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 difflib +from typing import Optional + + +def print_no_matches(component_type: str, query: Optional[str]) -> None: + """Print the standard 'nothing to show' message for a `gym list`/`gym search` command. + + ``component_type`` is the plural noun (``benchmarks``, ``environments``, ...). Keeps the message + and styling identical across every listing. + """ + import rich + + if query: + rich.print(f"[yellow]No {component_type} match '{query}'.[/yellow]") + else: + rich.print(f"[yellow]No {component_type} found.[/yellow]") + + +def fuzzy_matches(query: str, *fields: str) -> bool: + """Whether `query` fuzzily matches any of `fields`: a substring or a close difflib match (token-aware). + + The shared matcher behind `gym search ` across every component type. + """ + 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.70): + return True + return False def print_rich_table(table) -> None: diff --git a/tests/unit_tests/test_benchmarks.py b/tests/unit_tests/test_benchmarks.py index f535066beb..067b8639b5 100644 --- a/tests/unit_tests/test_benchmarks.py +++ b/tests/unit_tests/test_benchmarks.py @@ -20,7 +20,7 @@ from omegaconf import OmegaConf from yaml import safe_load -from nemo_gym.cli.eval import _fuzzy_matches, list_benchmarks, prepare_benchmark +from nemo_gym.cli.eval import list_benchmarks, prepare_benchmark def _mock_global_config(config: dict = None): @@ -150,24 +150,6 @@ def test_strict_is_the_default_and_does_not_tolerate_unresolved_values(self) -> assert tolerated is None -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 TestSearchBenchmarks: # Map each benchmark name to the `domain` its config would resolve to. DOMAINS = { @@ -248,7 +230,6 @@ def test_calls_prepare(self, tmp_path: Path) -> None: {"config_paths": [str(config_path)], **safe_load(config_path.read_text())} ), ), - patch("nemo_gym.cli.eval.BENCHMARKS_DIR", bench_dir.parent), patch("nemo_gym.cli.eval.importlib.import_module", return_value=mock_module), ): prepare_benchmark() @@ -265,7 +246,6 @@ def test_missing_prepare_py(self, tmp_path: Path, capsys) -> None: {"config_paths": [str(config_path)], **safe_load(config_path.read_text())} ), ), - patch("nemo_gym.cli.eval.BENCHMARKS_DIR", bench_dir.parent), ): with pytest.raises(SystemExit) as exc_info: prepare_benchmark() @@ -285,7 +265,6 @@ def test_missing_prepare_function(self, tmp_path: Path, capsys) -> None: {"config_paths": [str(config_path)], **safe_load(config_path.read_text())} ), ), - 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(SystemExit) as exc_info: @@ -338,7 +317,6 @@ def test_no_prepare_script_args_does_not_error(self, tmp_path: Path) -> None: {"config_paths": [str(config_path)], **safe_load(config_path.read_text())} ), ), - patch("nemo_gym.cli.eval.BENCHMARKS_DIR", bench_dir.parent), patch("nemo_gym.cli.eval.importlib.import_module", return_value=mock_module), ): prepare_benchmark() @@ -363,7 +341,6 @@ def test_caching_sanity(self, tmp_path: Path) -> None: } ), ), - 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 31e1795aab..382f4946e4 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -423,3 +423,26 @@ def test_json_output(self, monkeypatch: MonkeyPatch, capsys) -> None: assert json.loads(capsys.readouterr().out) == [ {"name": "alpha", "domain": "agent", "description": "Alpha env"} ] + + def test_query_filters_environments(self, monkeypatch: MonkeyPatch, capsys) -> None: + # `gym search environments ` reuses this command via the `query` config key (matches name + domain). + beta = EnvironmentEntry( + name="beta", + config_path=Path("environments/beta/config.yaml"), + path=Path("environments/beta"), + description="Beta env", + domain="math", + ) + monkeypatch.setattr( + nemo_gym.cli.env, "get_global_config_dict", lambda **k: OmegaConf.create({"query": "alpha"}) + ) + monkeypatch.setattr( + nemo_gym.cli.env, "discover_environments", lambda *a, **k: {"alpha": self._ALPHA, "beta": beta} + ) + + list_environments() + + out = capsys.readouterr().out + assert "Environments matching 'alpha'" in out + assert "agent" in out # alpha's domain -> its row was rendered + assert "beta" not in out and "math" not in out # beta and its domain filtered out diff --git a/tests/unit_tests/test_cli_agents.py b/tests/unit_tests/test_cli_agents.py index 4a60669b2a..fdf0d382c2 100644 --- a/tests/unit_tests/test_cli_agents.py +++ b/tests/unit_tests/test_cli_agents.py @@ -75,3 +75,14 @@ def test_json_output(self, capsys) -> None: assert by_name["simple_agent"]["pattern"] == "A (composable)" assert by_name["swe_agents"]["self_contained"] is True assert by_name["swe_agents"]["variants"] == ["swebench_openhands"] + + def test_query_filters_agents(self, capsys) -> None: + # `gym search agents ` reuses this command via the `query` config key (name + variant names). + with ( + patch("nemo_gym.cli.agents.get_global_config_dict", return_value=_mock_global_config({"query": "swe"})), + patch("nemo_gym.cli.agents.discover_agents", return_value=_AGENTS), + ): + list_agents() + out = capsys.readouterr().out + assert "swe_agents" in out and "Agents matching" in out + assert "simple_agent" not in out diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index 399980cc42..7f3611fffc 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -661,12 +661,27 @@ def test_no_json_no_override(self, monkeypatch: MonkeyPatch) -> None: 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. + def test_search_query_only_defaults_to_benchmarks(self, monkeypatch: MonkeyPatch) -> None: + # `gym search ` (no type) reuses the benchmarks listing — backward compatible. target, overrides = _dispatch_for(monkeypatch, ["search", "math"]) assert target == "nemo_gym.cli.eval:list_benchmarks" assert overrides == ["+query=math"] + @pytest.mark.parametrize( + "component_type, expected_target", + [ + ("benchmarks", "nemo_gym.cli.eval:list_benchmarks"), + ("environments", "nemo_gym.cli.env:list_environments"), + ("agents", "nemo_gym.cli.agents:list_agents"), + ("models", "nemo_gym.cli.models:list_models"), + ], + ) + def test_search_type_routes_to_that_listing(self, monkeypatch, component_type, expected_target) -> None: + # `gym search ` runs that type's listing, filtered by the query. + target, overrides = _dispatch_for(monkeypatch, ["search", component_type, "swe"]) + assert target == expected_target + assert overrides == ["+query=swe"] + def test_search_json(self, monkeypatch: MonkeyPatch) -> None: _, overrides = _dispatch_for(monkeypatch, ["search", "math", "--json"]) assert set(overrides) == {"+query=math", "+json=true"} diff --git a/tests/unit_tests/test_cli_models.py b/tests/unit_tests/test_cli_models.py index d5b3436e59..cc86696711 100644 --- a/tests/unit_tests/test_cli_models.py +++ b/tests/unit_tests/test_cli_models.py @@ -62,6 +62,20 @@ def test_no_models(self, capsys) -> None: list_models() assert "No models found" in capsys.readouterr().out + def test_query_filters_rows(self, capsys) -> None: + # `gym search models ` reuses this command via the `query` config key (token + model group). + with ( + patch( + "nemo_gym.cli.models.get_global_config_dict", + return_value=_mock_global_config({"query": "some_other_flavor"}), + ), + patch("nemo_gym.cli.models.discover_models", return_value=_MODELS), + ): + list_models() + out = capsys.readouterr().out + assert "my_model/some_other_flavor" in out and "Models matching" in out + assert "another_model" not in out + def test_json_output_is_per_variant_rows(self, capsys) -> None: with ( patch("nemo_gym.cli.models.get_global_config_dict", return_value=_mock_global_config({"json": True})), diff --git a/tests/unit_tests/test_cli_utils.py b/tests/unit_tests/test_cli_utils.py index e0afcd7cf9..767e46cdc9 100644 --- a/tests/unit_tests/test_cli_utils.py +++ b/tests/unit_tests/test_cli_utils.py @@ -17,7 +17,7 @@ from rich.console import Console from rich.table import Table -from nemo_gym.cli.utils import print_rich_table +from nemo_gym.cli.utils import fuzzy_matches, print_rich_table # A cell value wider than Rich's 80-col non-TTY default, so a truncated render would ellipsize it. @@ -59,3 +59,21 @@ def test_not_truncated_regardless_of_ambient_width(self, capsys, monkeypatch) -> out = capsys.readouterr().out assert _LONG_NAME in out assert "…" not in 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_any_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") From b6a7149939874a2f4558502a6ab0ea81a47421bb Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Wed, 15 Jul 2026 11:07:11 +0200 Subject: [PATCH 06/13] feat: add gym list and gym search for resources servers Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/cli/main.py | 8 +- nemo_gym/cli/resources_servers.py | 73 +++++++++++++++ nemo_gym/resources_server_registry.py | 87 ++++++++++++++++++ tests/unit_tests/test_cli_main.py | 1 + .../unit_tests/test_cli_resources_servers.py | 90 +++++++++++++++++++ .../test_resources_server_registry.py | 67 ++++++++++++++ 6 files changed, 325 insertions(+), 1 deletion(-) create mode 100644 nemo_gym/cli/resources_servers.py create mode 100644 nemo_gym/resources_server_registry.py create mode 100644 tests/unit_tests/test_cli_resources_servers.py create mode 100644 tests/unit_tests/test_resources_server_registry.py diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 238db5da34..24877e6523 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -152,6 +152,7 @@ def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag: "environments": "nemo_gym.cli.env:list_environments", "agents": "nemo_gym.cli.agents:list_agents", "models": "nemo_gym.cli.models:list_models", + "resources-servers": "nemo_gym.cli.resources_servers:list_resources_servers", } SEARCH_TERMS = Flag( @@ -332,7 +333,7 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: # One-line help for each command group, shown in `gym --help`. GROUPS = { - "list": "List available components (benchmarks, environments, agents, models).", + "list": "List available components (benchmarks, environments, agents, models, resources-servers).", "dataset": "Manage datasets.", "env": "Develop and run environments.", "eval": "Run evaluations.", @@ -362,6 +363,11 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: summary="List model servers by the value to pass to --model-type.", flags=(JSON, DISCOVERY_SEARCH_DIR), ), + "list resources-servers": Command( + target="nemo_gym.cli.resources_servers:list_resources_servers", + summary="List resources servers (selectable with --resources-server) by name.", + flags=(JSON, DISCOVERY_SEARCH_DIR), + ), "search": Command( target=_search, summary="Search a component type (default benchmarks) by name; like `list` filtered to a query.", diff --git a/nemo_gym/cli/resources_servers.py b/nemo_gym/cli/resources_servers.py new file mode 100644 index 0000000000..b17ff5faa2 --- /dev/null +++ b/nemo_gym/cli/resources_servers.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 + +from rich.table import Table + +from nemo_gym.cli.utils import fuzzy_matches, print_no_matches, print_rich_table +from nemo_gym.config_types import BaseNeMoGymCLIConfig +from nemo_gym.global_config import ( + JSON_OUTPUT_KEY_NAME, + QUERY_KEY_NAME, + SEARCH_DIR_KEY_NAME, + GlobalConfigDictParserConfig, + get_global_config_dict, +) +from nemo_gym.resources_server_registry import discover_resources_servers + + +def list_resources_servers() -> None: + """List the resources servers selectable with ``--resources-server``, by short name. Optionally filtered + by a `query` (the `gym search resources-servers` entry point). ``--search-dir`` adds extra roots on top + of the cwd and built-ins. + """ + 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) + + servers = discover_resources_servers(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + + # `gym search resources-servers ` reuses this command, narrowing to fuzzy matches on name + domain. + query = global_config_dict.get(QUERY_KEY_NAME) + if query: + servers = {name: s for name, s in servers.items() if fuzzy_matches(query, name, s.domain or "")} + + if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): + print( + json.dumps( + [{"name": name, "domain": s.domain, "description": s.description} for name, s in servers.items()] + ) + ) + return + + if not servers: + print_no_matches("resources servers", query) + return + + title = ( + f"Resources servers matching '{query}' ({len(servers)})" + if query + else f"Available resources servers in NeMo Gym ({len(servers)})" + ) + table = Table(title=title) + table.add_column("Name") + table.add_column("Domain") + table.add_column("Description") + for name, server in servers.items(): + table.add_row(name, server.domain or "", server.description or "") + print_rich_table(table) diff --git a/nemo_gym/resources_server_registry.py b/nemo_gym/resources_server_registry.py new file mode 100644 index 0000000000..cf0d70f9da --- /dev/null +++ b/nemo_gym/resources_server_registry.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. +"""Registry of resources servers under ``resources_servers//``. + +A resources server (verifier + per-task state) is one *component* of an environment, selected by name +with ``--resources-server``. This module maps each server dir to its ``(domain, description)`` — read +the same way ``gym list environments``/``benchmarks`` read theirs (via +:func:`~nemo_gym.discovery.read_config_metadata`) — so they can be enumerated by name. +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Optional, Sequence, Union + +from nemo_gym import PARENT_DIR +from nemo_gym.discovery import discover_components, read_config_metadata + + +RESOURCES_SERVERS_SUBDIR = "resources_servers" +RESOURCES_SERVERS_DIR = PARENT_DIR / RESOURCES_SERVERS_SUBDIR +RESOURCES_SERVER_CONFIGS_SUBDIR = "configs" + + +@dataclass(frozen=True) +class ResourcesServerEntry: + """A discovered resources server: its name, where it lives, and lightweight metadata.""" + + name: str + config_path: Path # the config metadata was read from (the default `` flavor, else the first) + path: Path + description: Optional[str] = None + domain: Optional[str] = None + + +def _discover_resources_servers_in_dir(resources_servers_dir: Path) -> Dict[str, ResourcesServerEntry]: + """Map resources-server name -> :class:`ResourcesServerEntry` for every server dir under one dir. + + The name is the directory name. A directory is a resources server iff it ships at least one + ``configs/*.yaml`` (the config selected by ``--resources-server``). Metadata is read from the default + ``.yaml`` flavor when present, else the first config. Returns an empty dict if the dir is missing. + """ + servers: Dict[str, ResourcesServerEntry] = {} + if not resources_servers_dir.is_dir(): + return servers + + for child in sorted(resources_servers_dir.iterdir()): + if not child.is_dir(): + continue + configs_dir = child / RESOURCES_SERVER_CONFIGS_SUBDIR + config_files = sorted(configs_dir.glob("*.yaml")) if configs_dir.is_dir() else [] + if not config_files: + continue + metadata_config = next((c for c in config_files if c.stem == child.name), config_files[0]) + domain, description = read_config_metadata(metadata_config) + servers[child.name] = ResourcesServerEntry( + name=child.name, + config_path=metadata_config, + path=child, + description=description, + domain=domain, + ) + + return servers + + +def discover_resources_servers( + search_dirs: Optional[Union[Path, Sequence[Path]]] = None, +) -> Dict[str, ResourcesServerEntry]: + """Map resources-server name -> :class:`ResourcesServerEntry` for every discoverable server. + + Scans the ``resources_servers/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` + root (``search_dirs`` + cwd + built-ins), merged so user servers shadow same-named built-ins. + ``search_dirs`` is one dir or a list. + """ + return discover_components(RESOURCES_SERVERS_SUBDIR, _discover_resources_servers_in_dir, search_dirs) diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index 7f3611fffc..116b6ddbdd 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -674,6 +674,7 @@ def test_search_query_only_defaults_to_benchmarks(self, monkeypatch: MonkeyPatch ("environments", "nemo_gym.cli.env:list_environments"), ("agents", "nemo_gym.cli.agents:list_agents"), ("models", "nemo_gym.cli.models:list_models"), + ("resources-servers", "nemo_gym.cli.resources_servers:list_resources_servers"), ], ) def test_search_type_routes_to_that_listing(self, monkeypatch, component_type, expected_target) -> None: diff --git a/tests/unit_tests/test_cli_resources_servers.py b/tests/unit_tests/test_cli_resources_servers.py new file mode 100644 index 0000000000..063a505fc3 --- /dev/null +++ b/tests/unit_tests/test_cli_resources_servers.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 +from pathlib import Path +from unittest.mock import patch + +from omegaconf import OmegaConf + +from nemo_gym.cli.resources_servers import list_resources_servers +from nemo_gym.resources_server_registry import ResourcesServerEntry + + +def _mock_global_config(config: dict = None): + return OmegaConf.create(config or {}) + + +def _entry(name: str, domain: str, description: str) -> ResourcesServerEntry: + path = Path("resources_servers") / name + return ResourcesServerEntry( + name=name, config_path=path / "configs" / f"{name}.yaml", path=path, description=description, domain=domain + ) + + +_SERVERS = { + "mcqa": _entry("mcqa", "knowledge", "Multi-choice QA"), + "aviary": _entry("aviary", "math", "Math tasks"), +} + + +class TestListResourcesServers: + def test_lists_servers(self, capsys) -> None: + with ( + patch("nemo_gym.cli.resources_servers.get_global_config_dict", return_value=_mock_global_config()), + patch("nemo_gym.cli.resources_servers.discover_resources_servers", return_value=_SERVERS), + ): + list_resources_servers() + out = capsys.readouterr().out + assert "mcqa" in out and "knowledge" in out and "aviary" in out + + def test_no_servers(self, capsys) -> None: + with ( + patch("nemo_gym.cli.resources_servers.get_global_config_dict", return_value=_mock_global_config()), + patch("nemo_gym.cli.resources_servers.discover_resources_servers", return_value={}), + ): + list_resources_servers() + assert "No resources servers found" in capsys.readouterr().out + + def test_json_output(self, capsys) -> None: + with ( + patch( + "nemo_gym.cli.resources_servers.get_global_config_dict", + return_value=_mock_global_config({"json": True}), + ), + patch("nemo_gym.cli.resources_servers.discover_resources_servers", return_value=_SERVERS), + ): + list_resources_servers() + payload = json.loads(capsys.readouterr().out) + expected = [ + {"name": "mcqa", "domain": "knowledge", "description": "Multi-choice QA"}, + {"name": "aviary", "domain": "math", "description": "Math tasks"}, + ] + assert len(payload) == len(expected) + for row in expected: + assert row in payload + + def test_query_filters_servers(self, capsys) -> None: + # `gym search resources-servers ` reuses this command via the `query` config key (name + domain). + with ( + patch( + "nemo_gym.cli.resources_servers.get_global_config_dict", + return_value=_mock_global_config({"query": "math"}), + ), + patch("nemo_gym.cli.resources_servers.discover_resources_servers", return_value=_SERVERS), + ): + list_resources_servers() + out = capsys.readouterr().out + assert "aviary" in out and "Resources servers matching" in out + assert "mcqa" not in out and "knowledge" not in out diff --git a/tests/unit_tests/test_resources_server_registry.py b/tests/unit_tests/test_resources_server_registry.py new file mode 100644 index 0000000000..c430160c99 --- /dev/null +++ b/tests/unit_tests/test_resources_server_registry.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 pathlib import Path + +from nemo_gym.resources_server_registry import _discover_resources_servers_in_dir + + +def _make_resources_server(rs_dir: Path, name: str, *, flavors: dict) -> Path: + """`flavors` maps config stem -> (domain, description). At least one config makes a resources server.""" + server_dir = rs_dir / name + configs_dir = server_dir / "configs" + configs_dir.mkdir(parents=True) + for stem, (domain, description) in flavors.items(): + (configs_dir / f"{stem}.yaml").write_text( + f"{name}:\n resources_servers:\n {name}:\n" + f" entrypoint: app.py\n domain: {domain}\n description: {description}\n" + ) + return server_dir + + +class TestDiscoverResourcesServers: + def test_reads_metadata_from_default_flavor(self, tmp_path: Path) -> None: + # Metadata comes from the `.yaml` flavor (the one `--resources-server ` resolves to). + _make_resources_server( + tmp_path, + "my_server", + flavors={"my_server": ("knowledge", "The default"), "some_other_flavor": ("other", "A flavor")}, + ) + + entry = _discover_resources_servers_in_dir(tmp_path)["my_server"] + + assert entry.domain == "knowledge" + assert entry.description == "The default" + assert entry.config_path.name == "my_server.yaml" + + def test_reads_first_flavor_when_no_default(self, tmp_path: Path) -> None: + # No `.yaml`, so metadata is read from the first config alphabetically. + _make_resources_server( + tmp_path, "my_server", flavors={"beta": ("science", "Beta"), "alpha": ("math", "Alpha")} + ) + + entry = _discover_resources_servers_in_dir(tmp_path)["my_server"] + + assert entry.config_path.name == "alpha.yaml" + assert entry.domain == "math" + + def test_dirs_without_a_config_are_skipped(self, tmp_path: Path) -> None: + # A dir with no config (e.g. a stray .egg-info) is not a resources server. + (tmp_path / "my_server.egg-info").mkdir() + _make_resources_server(tmp_path, "real_server", flavors={"real_server": ("other", "Real")}) + + assert set(_discover_resources_servers_in_dir(tmp_path)) == {"real_server"} + + def test_missing_directory_yields_no_servers(self, tmp_path: Path) -> None: + assert _discover_resources_servers_in_dir(tmp_path / "nope") == {} From 73fb3621c6284f1f98273ece74a52fe2b815c75f Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Wed, 15 Jul 2026 13:33:47 +0200 Subject: [PATCH 07/13] chore: user per-flavor list of resources servers Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/resources_server_registry.py | 48 +++++++++++++------ .../test_resources_server_registry.py | 23 ++++++--- 2 files changed, 50 insertions(+), 21 deletions(-) diff --git a/nemo_gym/resources_server_registry.py b/nemo_gym/resources_server_registry.py index cf0d70f9da..7b0a2e5205 100644 --- a/nemo_gym/resources_server_registry.py +++ b/nemo_gym/resources_server_registry.py @@ -15,7 +15,7 @@ """Registry of resources servers under ``resources_servers//``. A resources server (verifier + per-task state) is one *component* of an environment, selected by name -with ``--resources-server``. This module maps each server dir to its ``(domain, description)`` — read +with ``--resources-server``. This module maps each config flavor to its ``(domain, description)`` — read the same way ``gym list environments``/``benchmarks`` read theirs (via :func:`~nemo_gym.discovery.read_config_metadata`) — so they can be enumerated by name. """ @@ -24,6 +24,8 @@ from pathlib import Path from typing import Dict, Optional, Sequence, Union +from omegaconf import OmegaConf + from nemo_gym import PARENT_DIR from nemo_gym.discovery import discover_components, read_config_metadata @@ -38,18 +40,31 @@ class ResourcesServerEntry: """A discovered resources server: its name, where it lives, and lightweight metadata.""" name: str - config_path: Path # the config metadata was read from (the default `` flavor, else the first) + config_path: Path path: Path description: Optional[str] = None domain: Optional[str] = None +def _config_defines_resources_server(config_path: Path) -> bool: + """True if a config declares a ``resources_servers`` block (vs a helper like a judge model). Never raises.""" + try: + raw = OmegaConf.to_container(OmegaConf.load(config_path), resolve=False, throw_on_missing=False) + except Exception: + return False + if not isinstance(raw, dict): + return False + return any( + isinstance(instance, dict) and isinstance(instance.get("resources_servers"), dict) for instance in raw.values() + ) + + def _discover_resources_servers_in_dir(resources_servers_dir: Path) -> Dict[str, ResourcesServerEntry]: - """Map resources-server name -> :class:`ResourcesServerEntry` for every server dir under one dir. + """Map resources-server name -> :class:`ResourcesServerEntry` for every flavor under one dir. - The name is the directory name. A directory is a resources server iff it ships at least one - ``configs/*.yaml`` (the config selected by ``--resources-server``). Metadata is read from the default - ``.yaml`` flavor when present, else the first config. Returns an empty dict if the dir is missing. + One entry per config flavor that declares a `resources_servers` block: `` for the default config + (`.yaml`) and `/` for the rest (`.yaml`). + Helper configs (no `resources_servers` block) are skipped. Empty dict if the dir is missing. """ servers: Dict[str, ResourcesServerEntry] = {} if not resources_servers_dir.is_dir(): @@ -62,15 +77,18 @@ def _discover_resources_servers_in_dir(resources_servers_dir: Path) -> Dict[str, config_files = sorted(configs_dir.glob("*.yaml")) if configs_dir.is_dir() else [] if not config_files: continue - metadata_config = next((c for c in config_files if c.stem == child.name), config_files[0]) - domain, description = read_config_metadata(metadata_config) - servers[child.name] = ResourcesServerEntry( - name=child.name, - config_path=metadata_config, - path=child, - description=description, - domain=domain, - ) + for config in config_files: + if not _config_defines_resources_server(config): + continue + name = child.name if config.stem == child.name else f"{child.name}/{config.stem}" + domain, description = read_config_metadata(config) + servers[name] = ResourcesServerEntry( + name=name, + config_path=config, + path=child, + description=description, + domain=domain, + ) return servers diff --git a/tests/unit_tests/test_resources_server_registry.py b/tests/unit_tests/test_resources_server_registry.py index c430160c99..fc94cf084a 100644 --- a/tests/unit_tests/test_resources_server_registry.py +++ b/tests/unit_tests/test_resources_server_registry.py @@ -39,22 +39,24 @@ def test_reads_metadata_from_default_flavor(self, tmp_path: Path) -> None: flavors={"my_server": ("knowledge", "The default"), "some_other_flavor": ("other", "A flavor")}, ) - entry = _discover_resources_servers_in_dir(tmp_path)["my_server"] + servers = _discover_resources_servers_in_dir(tmp_path) + assert set(servers) == {"my_server", "my_server/some_other_flavor"} # one token per flavor + entry = servers["my_server"] assert entry.domain == "knowledge" assert entry.description == "The default" assert entry.config_path.name == "my_server.yaml" - def test_reads_first_flavor_when_no_default(self, tmp_path: Path) -> None: - # No `.yaml`, so metadata is read from the first config alphabetically. + def test_flavor_tokens_when_no_default(self, tmp_path: Path) -> None: + # No `.yaml`, so each flavor is its own `/` token (no collapsing to one entry). _make_resources_server( tmp_path, "my_server", flavors={"beta": ("science", "Beta"), "alpha": ("math", "Alpha")} ) - entry = _discover_resources_servers_in_dir(tmp_path)["my_server"] + servers = _discover_resources_servers_in_dir(tmp_path) - assert entry.config_path.name == "alpha.yaml" - assert entry.domain == "math" + assert set(servers) == {"my_server/alpha", "my_server/beta"} + assert servers["my_server/alpha"].domain == "math" def test_dirs_without_a_config_are_skipped(self, tmp_path: Path) -> None: # A dir with no config (e.g. a stray .egg-info) is not a resources server. @@ -63,5 +65,14 @@ def test_dirs_without_a_config_are_skipped(self, tmp_path: Path) -> None: assert set(_discover_resources_servers_in_dir(tmp_path)) == {"real_server"} + def test_helper_configs_are_skipped(self, tmp_path: Path) -> None: + # A config with no `resources_servers` block (e.g. a judge model helper) is not a flavor. + _make_resources_server(tmp_path, "my_server", flavors={"my_server": ("knowledge", "Default")}) + (tmp_path / "my_server" / "configs" / "judge_model.yaml").write_text( + "judge_model:\n responses_api_models:\n m:\n x: 1\n" + ) + + assert set(_discover_resources_servers_in_dir(tmp_path)) == {"my_server"} + def test_missing_directory_yields_no_servers(self, tmp_path: Path) -> None: assert _discover_resources_servers_in_dir(tmp_path / "nope") == {} From afc910738580d9c1d83daa67b02db79009c0fa1a Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Wed, 15 Jul 2026 15:30:42 +0200 Subject: [PATCH 08/13] feat: add gym list for inspecting Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/cli/agents.py | 44 ++++++++- nemo_gym/cli/env.py | 54 +++++++++- nemo_gym/cli/eval.py | 44 ++++++++- nemo_gym/cli/main.py | 45 +++++---- nemo_gym/cli/models.py | 45 ++++++++- nemo_gym/cli/resources_servers.py | 45 ++++++++- nemo_gym/cli/utils.py | 69 ++++++++++++- nemo_gym/discovery.py | 36 ++++--- nemo_gym/global_config.py | 2 + nemo_gym/registry.py | 45 ++++++++- nemo_gym/resources_server_registry.py | 14 ++- tests/unit_tests/test_benchmarks.py | 45 +++++++++ tests/unit_tests/test_cli.py | 98 +++++++++++++++++++ tests/unit_tests/test_cli_agents.py | 42 ++++++++ tests/unit_tests/test_cli_main.py | 15 ++- tests/unit_tests/test_cli_models.py | 68 +++++++++++++ .../unit_tests/test_cli_resources_servers.py | 41 ++++++++ tests/unit_tests/test_cli_utils.py | 44 ++++++++- tests/unit_tests/test_registry.py | 29 ++++++ .../test_resources_server_registry.py | 18 ++++ 20 files changed, 781 insertions(+), 62 deletions(-) diff --git a/nemo_gym/cli/agents.py b/nemo_gym/cli/agents.py index 46d25cd845..c8a3b5428e 100644 --- a/nemo_gym/cli/agents.py +++ b/nemo_gym/cli/agents.py @@ -17,9 +17,16 @@ from rich.table import Table from nemo_gym.agent_registry import discover_agents -from nemo_gym.cli.utils import fuzzy_matches, print_no_matches, print_rich_table +from nemo_gym.cli.utils import ( + exit_unknown_component, + fuzzy_matches, + print_no_matches, + print_rich_table, + render_component_inspection, +) from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( + COMPONENT_NAME_KEY_NAME, JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, SEARCH_DIR_KEY_NAME, @@ -28,11 +35,33 @@ ) +def _inspect_agent(name: str, agents: dict, global_config_dict) -> None: + """Render the ``gym list agents `` inspect view for one agent (thin: no usage example).""" + entry = agents.get(name) + if entry is None: + exit_unknown_component(name, agents, "agent") + return + + details = { + "path": str(entry.path.resolve()), + "composition": "self-contained (B)" if entry.self_contained else "composable (A)", + } + if entry.variants: + details["variants"] = ", ".join(sorted(entry.variants)) + + render_component_inspection( + json_output=global_config_dict.get(JSON_OUTPUT_KEY_NAME, False), + name=name, + type_noun="agent", + description=entry.description, + details=details, + ) + + def list_agents() -> None: - """List discovered agent harnesses and how each composes: freely wireable into a separate environment - (Pattern A) vs. self-contained harnesses that run with their own config (Pattern B). Optionally filtered - by a `query` (the `gym search agents` entry point). ``--search-dir`` adds extra roots on top of the cwd - and built-ins. + """List discovered agent harnesses and how each composes (Pattern A vs. self-contained B), or inspect one + by name (``gym list agents ``). Optionally filtered by a `query` (the `gym search agents` entry + point). ``--search-dir`` adds extra roots on top of the cwd and built-ins. """ global_config_dict = get_global_config_dict( global_config_dict_parser_config=GlobalConfigDictParserConfig( @@ -43,6 +72,11 @@ def list_agents() -> None: agents = discover_agents(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) + if name: + _inspect_agent(name, agents, global_config_dict) + return + # `gym search agents ` reuses this command, narrowing to fuzzy matches on name + variant names. query = global_config_dict.get(QUERY_KEY_NAME) if query: diff --git a/nemo_gym/cli/env.py b/nemo_gym/cli/env.py index c29f575a04..32d8b71460 100644 --- a/nemo_gym/cli/env.py +++ b/nemo_gym/cli/env.py @@ -38,9 +38,17 @@ from nemo_gym import PARENT_DIR, ROOT_DIR from nemo_gym.cli.setup_command import run_command, setup_env_command -from nemo_gym.cli.utils import exit_cleanly_on_config_error, fuzzy_matches, print_no_matches, print_rich_table +from nemo_gym.cli.utils import ( + exit_cleanly_on_config_error, + exit_unknown_component, + fuzzy_matches, + print_no_matches, + print_rich_table, + render_component_inspection, +) from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( + COMPONENT_NAME_KEY_NAME, DRY_RUN_KEY_NAME, JSON_OUTPUT_KEY_NAME, NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME, @@ -52,7 +60,7 @@ GlobalConfigDictParserConfig, get_global_config_dict, ) -from nemo_gym.registry import discover_environments +from nemo_gym.registry import discover_environments, read_environment_details from nemo_gym.server_status import StatusCommand from nemo_gym.server_utils import ( HEAD_SERVER_KEY_NAME, @@ -920,14 +928,47 @@ def validate(): rich.print("[green]✓[/green] Config is valid.") +def _inspect_environment(name: str, environments: dict, global_config_dict) -> None: + """Render the ``gym list environments `` inspect view for one environment.""" + entry = environments.get(name) + if entry is None: + exit_unknown_component(name, environments, "environment") + return + + parsed = read_environment_details(entry.config_path) + details = {"config": str(entry.config_path.resolve())} + if parsed["resources_servers"]: + details["resources servers"] = ", ".join(parsed["resources_servers"]) + if parsed["agent"]: + details["agent"] = parsed["agent"] + if parsed["datasets"]: + details["datasets"] = ", ".join(parsed["datasets"]) + + description = parsed["description"] + if parsed["value"]: # surface `value` as a trailing line of the description + description = f"{description}\nValue: {parsed['value']}" if description else f"Value: {parsed['value']}" + + render_component_inspection( + json_output=global_config_dict.get(JSON_OUTPUT_KEY_NAME, False), + name=name, + type_noun="environment", + domain=parsed["domain"], + description=description, + details=details, + usage=f"gym env start --environment {name} --model-type vllm_model", + ) + + def list_environments() -> None: - """List the environments available under environments/, optionally filtered by a `query` (the - `gym search environments` entry point). ``--search-dir`` adds extra roots on top of the cwd and built-ins. + """List the environments under environments/, or inspect one by name (``gym list environments ``). + Optionally filtered by a `query` (the `gym search environments` entry point). ``--search-dir`` adds extra + roots on top of the cwd and built-ins. Examples: ```bash gym list environments + gym list environments calendar gym list environments --json gym list environments --search-dir /path/to/project ``` @@ -941,6 +982,11 @@ def list_environments() -> None: environments = discover_environments(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) + if name: + _inspect_environment(name, environments, global_config_dict) + return + # `gym search environments ` reuses this command, narrowing to fuzzy matches on name + domain. query = global_config_dict.get(QUERY_KEY_NAME) if query: diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index 370a428a6e..30308c9f2c 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -31,10 +31,18 @@ discover_benchmarks, ) from nemo_gym.cli.env import RunHelper -from nemo_gym.cli.utils import exit_cleanly_on_config_error, fuzzy_matches, print_no_matches, print_rich_table +from nemo_gym.cli.utils import ( + exit_cleanly_on_config_error, + exit_unknown_component, + fuzzy_matches, + print_no_matches, + print_rich_table, + render_component_inspection, +) from nemo_gym.config_types import BaseNeMoGymCLIConfig, BenchmarkDatasetConfig, ConfigError, ConfigPathNotFoundError from nemo_gym.discovery import read_config_metadata from nemo_gym.global_config import ( + COMPONENT_NAME_KEY_NAME, JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, @@ -56,8 +64,35 @@ from nemo_gym.train_data_utils import TrainDataProcessor +def _inspect_benchmark(name: str, benchmarks: dict, global_config_dict) -> None: + """Render the ``gym list benchmarks `` inspect view for one benchmark.""" + bench = benchmarks.get(name) + if bench is None: + exit_unknown_component(name, benchmarks, "benchmark") + return + + domain, description = read_config_metadata(bench.path) + details = { + "config": str(bench.path.resolve()), + "agent": bench.agent_name, + "num repeats": str(bench.num_repeats), + "dataset": str(bench.dataset.jsonl_fpath), + "prepare script": str(bench.dataset.prepare_script), + } + render_component_inspection( + json_output=global_config_dict.get(JSON_OUTPUT_KEY_NAME, False), + name=name, + type_noun="benchmark", + domain=domain, + description=description, + details=details, + usage=f"gym eval prepare --benchmark {name}\ngym eval run --benchmark {name} --model-type vllm_model", + ) + + def list_benchmarks() -> None: - """CLI command: list available benchmarks, optionally filtered by a `query` (the `gym search` entry point). + """List available benchmarks, or inspect one by name (``gym list benchmarks ``). Optionally filtered + by a `query` (the `gym search` entry point). A benchmark is a specific kind of environment, so it shares `gym list environments`' columns (name, domain, description) and reads them through the same `read_config_metadata` helper. ``--search-dir`` @@ -72,6 +107,11 @@ def list_benchmarks() -> None: benchmarks = discover_benchmarks(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) + if name: + _inspect_benchmark(name, benchmarks, global_config_dict) + return + # Resolve domain + description once per benchmark, via the shared component-metadata reader — # the same one `gym list environments` uses — for the columns and `gym search`. metadata = {name: read_config_metadata(bench.path) for name, bench in benchmarks.items()} diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 24877e6523..6eaeba7a3f 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -13,26 +13,20 @@ # 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 collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path +from nemo_gym.cli.utils import did_you_mean from nemo_gym.discovery import component_search_roots 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. @@ -47,7 +41,7 @@ def error(self, message: str) -> None: choices = re.findall(r"'([^']+)'", match.group(2)) if not choices: choices = [choice.strip() for choice in match.group(2).split(",")] - message += _did_you_mean(typo, choices) + message += did_you_mean(typo, choices) super().error(message) @@ -144,6 +138,15 @@ def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag: # global_config_dict.get(JSON_OUTPUT_KEY_NAME) (see general.py, eval.py, env.py). JSON = _bool_flag("json", "json", "Output as machine-readable JSON.") +# `gym list []`: an optional component name. When given, the listing command inspects that one +# component (surfaced as the reserved `component_name` config key) instead of listing all. +NAME = Flag( + register=lambda p: p.add_argument( + "name", nargs="?", metavar="NAME", help="Inspect a single component by name instead of listing all." + ), + translate_to_hydra=lambda args: [f"+component_name={args.name}"] if getattr(args, "name", None) else [], +) + # `gym search [] `: an optional component type plus the query. The query is surfaced to the # chosen listing command as the reserved `query` config key; the type only picks which command to run # (see `_search`). A lone positional is the query, defaulting to benchmarks — backward compatible. @@ -236,7 +239,7 @@ def _asset_config_path(flag: str, value: str, search_dirs: tuple[str, ...] = ()) ] raise ValueError( - f"`--{flag} {value}` was specified which implies config `{path}`, which does not exist.{_did_you_mean(typo, candidates)} " + 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}." ) @@ -345,28 +348,28 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: COMMANDS = { "list benchmarks": Command( target="nemo_gym.cli.eval:list_benchmarks", - summary="List available benchmarks.", - flags=(JSON, DISCOVERY_SEARCH_DIR), + summary="List or inspect available benchmarks.", + flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), ), "list environments": Command( target="nemo_gym.cli.env:list_environments", - summary="List available environments by name.", - flags=(JSON, DISCOVERY_SEARCH_DIR), + summary="List or inspect available environments.", + flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), ), "list agents": Command( target="nemo_gym.cli.agents:list_agents", - summary="List agent harnesses and how each composes (Pattern A vs self-contained B).", - flags=(JSON, DISCOVERY_SEARCH_DIR), + summary="List or inspect available agent harnesses.", + flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), ), "list models": Command( target="nemo_gym.cli.models:list_models", - summary="List model servers by the value to pass to --model-type.", - flags=(JSON, DISCOVERY_SEARCH_DIR), + summary="List or inspect available model servers.", + flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), ), "list resources-servers": Command( target="nemo_gym.cli.resources_servers:list_resources_servers", - summary="List resources servers (selectable with --resources-server) by name.", - flags=(JSON, DISCOVERY_SEARCH_DIR), + summary="List or inspect available resources servers.", + flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), ), "search": Command( target=_search, @@ -671,7 +674,7 @@ def main() -> None: 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) + 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: diff --git a/nemo_gym/cli/models.py b/nemo_gym/cli/models.py index 270d6a66e4..212128ba94 100644 --- a/nemo_gym/cli/models.py +++ b/nemo_gym/cli/models.py @@ -16,9 +16,16 @@ from rich.table import Table -from nemo_gym.cli.utils import fuzzy_matches, print_no_matches, print_rich_table +from nemo_gym.cli.utils import ( + exit_unknown_component, + fuzzy_matches, + print_no_matches, + print_rich_table, + render_component_inspection, +) from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( + COMPONENT_NAME_KEY_NAME, JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, SEARCH_DIR_KEY_NAME, @@ -28,11 +35,34 @@ from nemo_gym.model_registry import discover_models +def _inspect_model(name: str, models: dict, global_config_dict) -> None: + """Render the ``gym list models `` inspect view for one model (thin: no usage example). + + ``name`` may be a bare model or a ``/`` token; a valid flavor renders the model's + (main) inspection. + """ + model = name.split("/", 1)[0] + entry = models.get(model) + if entry is None or (name != model and name not in entry.model_types): + exit_unknown_component(name, [token for e in models.values() for token in e.model_types], "model") + return + + details = {"path": str(entry.path.resolve())} + if entry.model_types: + details["model-types"] = ", ".join(entry.model_types) + + render_component_inspection( + json_output=global_config_dict.get(JSON_OUTPUT_KEY_NAME, False), + name=model, + type_noun="model", + details=details, + ) + + def list_models() -> None: - """List model servers, one row per ``--model-type`` value: ``Model`` is the token to pass (```` - for the default flavor, ``/`` for the rest); ``Model group`` is its model. Optionally - filtered by a `query` (the `gym search models` entry point). ``--search-dir`` adds extra roots on top of - the cwd and built-ins. + """List model servers (one row per ``--model-type`` value: ``Model`` is the token to pass, ``Model group`` + its model), or inspect one by name (``gym list models ``). Optionally filtered by a `query` (the + `gym search models` entry point). ``--search-dir`` adds extra roots on top of the cwd and built-ins. """ global_config_dict = get_global_config_dict( global_config_dict_parser_config=GlobalConfigDictParserConfig( @@ -43,6 +73,11 @@ def list_models() -> None: models = discover_models(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) + if name: + _inspect_model(name, models, global_config_dict) + return + # One row per passable `--model-type` value: `model` is the token, `model_group` its model. rows = [ {"model": model_type, "model_group": name} diff --git a/nemo_gym/cli/resources_servers.py b/nemo_gym/cli/resources_servers.py index b17ff5faa2..dbc296e08a 100644 --- a/nemo_gym/cli/resources_servers.py +++ b/nemo_gym/cli/resources_servers.py @@ -16,22 +16,52 @@ from rich.table import Table -from nemo_gym.cli.utils import fuzzy_matches, print_no_matches, print_rich_table +from nemo_gym.cli.utils import ( + exit_unknown_component, + fuzzy_matches, + print_no_matches, + print_rich_table, + render_component_inspection, +) from nemo_gym.config_types import BaseNeMoGymCLIConfig from nemo_gym.global_config import ( + COMPONENT_NAME_KEY_NAME, JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, SEARCH_DIR_KEY_NAME, GlobalConfigDictParserConfig, get_global_config_dict, ) -from nemo_gym.resources_server_registry import discover_resources_servers +from nemo_gym.resources_server_registry import discover_resources_servers, read_resources_server_value + + +def _inspect_resources_server(name: str, servers: dict, global_config_dict) -> None: + """Render the ``gym list resources-servers `` inspect view for one server.""" + entry = servers.get(name) + if entry is None: + exit_unknown_component(name, servers, "resources server") + return + + value = read_resources_server_value(entry.config_path) + description = entry.description + if value: # surface `value` as a trailing line of the description + description = f"{description}\nValue: {value}" if description else f"Value: {value}" + + render_component_inspection( + json_output=global_config_dict.get(JSON_OUTPUT_KEY_NAME, False), + name=name, + type_noun="resources server", + domain=entry.domain, + description=description, + details={"config": str(entry.config_path.resolve())}, + usage=f"gym env start --resources-server {name} --model-type vllm_model", + ) def list_resources_servers() -> None: - """List the resources servers selectable with ``--resources-server``, by short name. Optionally filtered - by a `query` (the `gym search resources-servers` entry point). ``--search-dir`` adds extra roots on top - of the cwd and built-ins. + """List the resources servers selectable with ``--resources-server``, or inspect one by name + (``gym list resources-servers ``). Optionally filtered by a `query` (the + `gym search resources-servers` entry point). ``--search-dir`` adds extra roots on top of the cwd and built-ins. """ global_config_dict = get_global_config_dict( global_config_dict_parser_config=GlobalConfigDictParserConfig( @@ -42,6 +72,11 @@ def list_resources_servers() -> None: servers = discover_resources_servers(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) + if name: + _inspect_resources_server(name, servers, global_config_dict) + return + # `gym search resources-servers ` reuses this command, narrowing to fuzzy matches on name + domain. query = global_config_dict.get(QUERY_KEY_NAME) if query: diff --git a/nemo_gym/cli/utils.py b/nemo_gym/cli/utils.py index 8a6e767032..99d92a2471 100644 --- a/nemo_gym/cli/utils.py +++ b/nemo_gym/cli/utils.py @@ -13,7 +13,74 @@ # See the License for the specific language governing permissions and # limitations under the License. import difflib -from typing import Optional +import json +import sys +from typing import Dict, Iterable, Optional + + +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 "" + + +def exit_unknown_component(name: str, candidates: Iterable[str], type_label: str) -> None: + """Print an `unknown ''` error (with a did-you-mean hint) and exit nonzero.""" + import rich + + rich.print(f"[red]Unknown {type_label} '{name}'.[/red]" + did_you_mean(name, candidates)) + sys.exit(1) + + +def render_component_inspection( + *, + json_output: bool, + name: str, + type_noun: str, + domain: Optional[str] = None, + description: Optional[str] = None, + details: Dict[str, str], + usage: Optional[str] = None, +) -> None: + """Render the uniform ``gym list `` inspect view (or its ``--json`` payload). + + ``details`` is an ordered label -> value mapping (e.g. ``{"config": ..., "agent": ...}``). Text + sections (domain suffix, description, Details, Usage example) are omitted when empty. + """ + if json_output: + print( + json.dumps( + { + "name": name, + "type": type_noun, + "domain": domain, + "description": description, + "details": details, + "usage_example": usage, + } + ) + ) + return + + from rich.console import Console + from rich.markup import escape + + # The name and section titles are bold; all dynamic text is escaped so `[...]` in a description or + # value isn't parsed as Rich markup. + header = f"The [bold]{escape(name)}[/bold] {escape(type_noun)}" + if domain: + header += f" (domain: {escape(domain)})" + sections = [header] + if description: + sections.append(escape(description)) + if details: + body = "\n".join(f"{escape(label)}: {escape(str(val))}" for label, val in details.items()) + sections.append(f"[bold]Details:[/bold]\n{body}") + if usage: + sections.append(f"[bold]Usage example:[/bold]\n{escape(usage)}") + # `soft_wrap` so long descriptions/paths aren't reflowed to the console width (esp. when piped); + # `highlight=False` so only our explicit bold applies (no auto-styling of parens/numbers/paths). + Console().print("\n\n".join(sections), soft_wrap=True, highlight=False) def print_no_matches(component_type: str, query: Optional[str]) -> None: diff --git a/nemo_gym/discovery.py b/nemo_gym/discovery.py index ad5fbf38c0..418856047d 100644 --- a/nemo_gym/discovery.py +++ b/nemo_gym/discovery.py @@ -129,14 +129,15 @@ def _parse_no_environment_tolerating_unset_values(initial_config_dict: DictConfi working = OmegaConf.merge(DictConfig({key: _UNSET_VALUE_PLACEHOLDER}), working) -def _scan_servers_for_metadata(container) -> Tuple[Optional[str], Optional[str]]: - """Best-effort ``(domain, description)`` from a config mapping: the first of each found across all - server groups. Defensive against malformed shapes, so it never raises. +def iter_server_configs(container): + """Yield ``(group_key, server_name, server_config)`` for every server across all instances in a config. + + Walks a loaded config mapping (each top-level instance -> its ``resources_servers``/ + ``responses_api_agents``/``responses_api_models`` group -> each server). Defensive against malformed + shapes, so it never raises. The shared primitive behind metadata reads and the inspect deep-parse. """ - domain: Optional[str] = None - description: Optional[str] = None if not isinstance(container, (dict, DictConfig)): - return None, None + return for instance in container.values(): if not isinstance(instance, (dict, DictConfig)): continue @@ -144,13 +145,22 @@ def _scan_servers_for_metadata(container) -> Tuple[Optional[str], Optional[str]] servers = instance.get(group_key) if not isinstance(servers, (dict, DictConfig)): continue - for server_config in servers.values(): - if not isinstance(server_config, (dict, DictConfig)): - continue - if domain is None and server_config.get("domain"): - domain = str(server_config["domain"]) - if description is None and server_config.get("description"): - description = str(server_config["description"]) + for server_name, server_config in servers.items(): + if isinstance(server_config, (dict, DictConfig)): + yield group_key, server_name, server_config + + +def _scan_servers_for_metadata(container) -> Tuple[Optional[str], Optional[str]]: + """Best-effort ``(domain, description)`` from a config mapping: the first of each found across all + server groups. Never raises. + """ + domain: Optional[str] = None + description: Optional[str] = None + for _group_key, _server_name, server_config in iter_server_configs(container): + if domain is None and server_config.get("domain"): + domain = str(server_config["domain"]) + if description is None and server_config.get("description"): + description = str(server_config["description"]) return domain, description diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 939dd16f44..36e9fc9b34 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -90,6 +90,7 @@ OBSERVABILITY_ENABLED_KEY_NAME = "observability_enabled" MODEL_CALL_CAPTURE_DIR_KEY_NAME = "model_call_capture_dir" SEARCH_DIR_KEY_NAME = "search_dir" +COMPONENT_NAME_KEY_NAME = "component_name" NEMO_GYM_RESERVED_TOP_LEVEL_KEYS = [ CONFIG_PATHS_KEY_NAME, ENTRYPOINT_KEY_NAME, @@ -118,6 +119,7 @@ OBSERVABILITY_ENABLED_KEY_NAME, MODEL_CALL_CAPTURE_DIR_KEY_NAME, SEARCH_DIR_KEY_NAME, + COMPONENT_NAME_KEY_NAME, ] # Data keys diff --git a/nemo_gym/registry.py b/nemo_gym/registry.py index 90123289db..07805ffcef 100644 --- a/nemo_gym/registry.py +++ b/nemo_gym/registry.py @@ -27,10 +27,12 @@ from dataclasses import dataclass from pathlib import Path -from typing import Dict, Optional, Sequence, Union +from typing import Dict, List, Optional, Sequence, Union + +from omegaconf import DictConfig, OmegaConf from nemo_gym import PARENT_DIR -from nemo_gym.discovery import discover_components, read_config_metadata +from nemo_gym.discovery import discover_components, iter_server_configs, read_config_metadata ENVIRONMENTS_SUBDIR = "environments" @@ -85,3 +87,42 @@ def discover_environments( ``search_dirs`` is one dir or a list. """ return discover_components(ENVIRONMENTS_SUBDIR, _discover_environments_in_dir, search_dirs) + + +def read_environment_details(config_path: Path) -> Dict[str, object]: + """Deep-parse an environment config for the ``gym list environments `` inspect view. + + Returns ``domain``, ``description`` (via :func:`~nemo_gym.discovery.read_config_metadata`), plus + ``value``, ``resources_servers`` (names), ``agent`` (the agent type), and dataset ``names`` read from + the config's server blocks. Never raises: an unreadable config yields empty/None fields. + """ + domain, description = read_config_metadata(config_path) + try: + raw = OmegaConf.to_container(OmegaConf.load(config_path), resolve=False, throw_on_missing=False) + except Exception: + raw = None + + value: Optional[str] = None + resources_servers: List[str] = [] + agent: Optional[str] = None + datasets: List[str] = [] + for group_key, server_name, server_config in iter_server_configs(raw): + if group_key == "resources_servers": + resources_servers.append(server_name) + if value is None and server_config.get("value"): + value = str(server_config["value"]) + elif group_key == "responses_api_agents": + if agent is None: + agent = server_name + for dataset in server_config.get("datasets") or []: + if isinstance(dataset, (dict, DictConfig)) and dataset.get("name"): + datasets.append(str(dataset["name"])) + + return { + "domain": domain, + "description": description, + "value": value, + "resources_servers": resources_servers, + "agent": agent, + "datasets": datasets, + } diff --git a/nemo_gym/resources_server_registry.py b/nemo_gym/resources_server_registry.py index 7b0a2e5205..587d2fdf4c 100644 --- a/nemo_gym/resources_server_registry.py +++ b/nemo_gym/resources_server_registry.py @@ -27,7 +27,7 @@ from omegaconf import OmegaConf from nemo_gym import PARENT_DIR -from nemo_gym.discovery import discover_components, read_config_metadata +from nemo_gym.discovery import discover_components, iter_server_configs, read_config_metadata RESOURCES_SERVERS_SUBDIR = "resources_servers" @@ -103,3 +103,15 @@ def discover_resources_servers( ``search_dirs`` is one dir or a list. """ return discover_components(RESOURCES_SERVERS_SUBDIR, _discover_resources_servers_in_dir, search_dirs) + + +def read_resources_server_value(config_path: Path) -> Optional[str]: + """The ``value`` field declared on a resources server config (for the inspect view). Never raises.""" + try: + raw = OmegaConf.to_container(OmegaConf.load(config_path), resolve=False, throw_on_missing=False) + except Exception: + return None + for group_key, _server_name, server_config in iter_server_configs(raw): + if group_key == "resources_servers" and server_config.get("value"): + return str(server_config["value"]) + return None diff --git a/tests/unit_tests/test_benchmarks.py b/tests/unit_tests/test_benchmarks.py index 067b8639b5..c0e52db0ee 100644 --- a/tests/unit_tests/test_benchmarks.py +++ b/tests/unit_tests/test_benchmarks.py @@ -88,6 +88,51 @@ def test_json_output_empty(self, capsys) -> None: list_benchmarks() assert json.loads(capsys.readouterr().out) == [] + def test_inspect_benchmark_by_name(self, capsys) -> None: + bench = MagicMock(agent_name="my_agent", num_repeats=8) + bench.path = Path("benchmarks/aime24/config.yaml") + bench.dataset.jsonl_fpath = Path("benchmarks/aime24/data/aime24.jsonl") + bench.dataset.prepare_script = Path("benchmarks/aime24/prepare.py") + with ( + patch( + "nemo_gym.cli.eval.get_global_config_dict", + return_value=_mock_global_config({"component_name": "aime24"}), + ), + patch("nemo_gym.cli.eval.discover_benchmarks", return_value={"aime24": bench}), + patch("nemo_gym.cli.eval.read_config_metadata", return_value=("math", "AIME desc")), + ): + list_benchmarks() + out = capsys.readouterr().out + assert "The aime24 benchmark (domain: math)" in out and "AIME desc" in out + assert "agent: my_agent" in out and "num repeats: 8" in out + assert "gym eval prepare --benchmark aime24" in out + assert "gym eval run --benchmark aime24 --model-type vllm_model" in out + + def test_inspect_unknown_benchmark_exits(self, capsys) -> None: + with ( + patch( + "nemo_gym.cli.eval.get_global_config_dict", + return_value=_mock_global_config({"component_name": "aim24"}), + ), + patch("nemo_gym.cli.eval.discover_benchmarks", return_value={"aime24": MagicMock()}), + ): + with pytest.raises(SystemExit): + list_benchmarks() + out = capsys.readouterr().out + assert "Unknown benchmark 'aim24'" in out and "aime24" in out + + def test_inspect_shows_absolute_config_path(self, capsys) -> None: + # Real discovery: the config line must be the config's absolute path, not a cwd-relative one. + from nemo_gym.benchmarks import BENCHMARKS_DIR + + expected = (BENCHMARKS_DIR / "aime24" / "config.yaml").resolve() + with patch( + "nemo_gym.cli.eval.get_global_config_dict", + return_value=_mock_global_config({"component_name": "aime24"}), + ): + list_benchmarks() + assert f"config: {expected}" in capsys.readouterr().out + class TestLoadBenchmarksFromConfigPaths: def test_skips_configs_that_fail_to_resolve_with_warning(self, capsys) -> None: diff --git a/tests/unit_tests/test_cli.py b/tests/unit_tests/test_cli.py index 382f4946e4..8325c41733 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -446,3 +446,101 @@ def test_query_filters_environments(self, monkeypatch: MonkeyPatch, capsys) -> N assert "Environments matching 'alpha'" in out assert "agent" in out # alpha's domain -> its row was rendered assert "beta" not in out and "math" not in out # beta and its domain filtered out + + def test_inspect_environment_by_name(self, monkeypatch: MonkeyPatch, capsys) -> None: + monkeypatch.setattr( + nemo_gym.cli.env, "get_global_config_dict", lambda **k: OmegaConf.create({"component_name": "alpha"}) + ) + monkeypatch.setattr(nemo_gym.cli.env, "discover_environments", lambda *a, **k: {"alpha": self._ALPHA}) + monkeypatch.setattr( + nemo_gym.cli.env, + "read_environment_details", + lambda cfg: { + "domain": "agent", + "description": "Alpha env", + "value": "Some value", + "resources_servers": ["alpha_rs"], + "agent": "simple_agent", + "datasets": ["train", "example"], + }, + ) + + list_environments() + + out = capsys.readouterr().out + assert "The alpha environment (domain: agent)" in out + assert "Value: Some value" in out + assert "resources servers: alpha_rs" in out and "agent: simple_agent" in out + assert "datasets: train, example" in out + assert "gym env start --environment alpha --model-type vllm_model" in out + + def _mock_inspect_alpha(self, monkeypatch: MonkeyPatch, config: dict) -> None: + monkeypatch.setattr( + nemo_gym.cli.env, + "get_global_config_dict", + lambda **k: OmegaConf.create({"component_name": "alpha", **config}), + ) + monkeypatch.setattr(nemo_gym.cli.env, "discover_environments", lambda *a, **k: {"alpha": self._ALPHA}) + monkeypatch.setattr( + nemo_gym.cli.env, + "read_environment_details", + lambda cfg: { + "domain": "agent", + "description": "Alpha env", + "value": "Some value", + "resources_servers": [], + "agent": None, + "datasets": [], + }, + ) + + def test_inspect_folds_value_into_description(self, monkeypatch: MonkeyPatch, capsys) -> None: + # `value` is not a separate field: it is appended to the description, not surfaced on its own. + self._mock_inspect_alpha(monkeypatch, {"json": True}) + + list_environments() + + payload = json.loads(capsys.readouterr().out) + assert payload["description"] == "Alpha env\nValue: Some value" + assert "value" not in payload and "value" not in payload["details"] + + def test_inspect_json_output(self, monkeypatch: MonkeyPatch, capsys) -> None: + self._mock_inspect_alpha(monkeypatch, {"json": True}) + + list_environments() + + assert json.loads(capsys.readouterr().out) == { + "name": "alpha", + "type": "environment", + "domain": "agent", + "description": "Alpha env\nValue: Some value", + "details": {"config": str(self._ALPHA.config_path.resolve())}, + "usage_example": "gym env start --environment alpha --model-type vllm_model", + } + + def test_inspect_unknown_environment_exits(self, monkeypatch: MonkeyPatch, capsys) -> None: + monkeypatch.setattr( + nemo_gym.cli.env, "get_global_config_dict", lambda **k: OmegaConf.create({"component_name": "alfa"}) + ) + monkeypatch.setattr(nemo_gym.cli.env, "discover_environments", lambda *a, **k: {"alpha": self._ALPHA}) + + with raises(SystemExit): + list_environments() + + out = capsys.readouterr().out + assert "Unknown environment 'alfa'" in out and "alpha" in out + + def test_inspect_shows_absolute_config_path(self, monkeypatch: MonkeyPatch, capsys, tmp_path: Path) -> None: + # Real discovery (via --search-dir): the config line must be the config's absolute path. + cfg = tmp_path / "environments" / "my_env" / "config.yaml" + cfg.parent.mkdir(parents=True) + cfg.write_text("my_env:\n resources_servers:\n my_env:\n domain: agent\n description: D\n") + monkeypatch.setattr( + nemo_gym.cli.env, + "get_global_config_dict", + lambda **k: OmegaConf.create({"component_name": "my_env", "search_dir": [str(tmp_path)]}), + ) + + list_environments() + + assert f"config: {cfg.resolve()}" in capsys.readouterr().out diff --git a/tests/unit_tests/test_cli_agents.py b/tests/unit_tests/test_cli_agents.py index fdf0d382c2..4e0b6169e3 100644 --- a/tests/unit_tests/test_cli_agents.py +++ b/tests/unit_tests/test_cli_agents.py @@ -16,6 +16,7 @@ from pathlib import Path from unittest.mock import patch +import pytest from omegaconf import OmegaConf from nemo_gym.agent_registry import AgentEntry @@ -86,3 +87,44 @@ def test_query_filters_agents(self, capsys) -> None: out = capsys.readouterr().out assert "swe_agents" in out and "Agents matching" in out assert "simple_agent" not in out + + def test_inspect_agent_by_name(self, capsys) -> None: + with ( + patch( + "nemo_gym.cli.agents.get_global_config_dict", + return_value=_mock_global_config({"component_name": "swe_agents"}), + ), + patch("nemo_gym.cli.agents.discover_agents", return_value=_AGENTS), + ): + list_agents() + out = capsys.readouterr().out + assert "The swe_agents agent" in out + assert "composition: self-contained (B)" in out + assert "variants: swebench_openhands" in out + assert "SWE tasks" in out # description + assert "Usage example:" not in out # thin view + + def test_inspect_unknown_agent_exits(self, capsys) -> None: + with ( + patch( + "nemo_gym.cli.agents.get_global_config_dict", + return_value=_mock_global_config({"component_name": "swe_agent"}), + ), + patch("nemo_gym.cli.agents.discover_agents", return_value=_AGENTS), + ): + with pytest.raises(SystemExit): + list_agents() + out = capsys.readouterr().out + assert "Unknown agent 'swe_agent'" in out and "swe_agents" in out + + def test_inspect_shows_absolute_path(self, tmp_path: Path, capsys) -> None: + # Real discovery (via --search-dir): the path line must be the agent dir's absolute path. + agent_dir = tmp_path / "responses_api_agents" / "my_agent" + agent_dir.mkdir(parents=True) + (agent_dir / "app.py").write_text("") + with patch( + "nemo_gym.cli.agents.get_global_config_dict", + return_value=_mock_global_config({"component_name": "my_agent", "search_dir": [str(tmp_path)]}), + ): + list_agents() + assert f"path: {agent_dir.resolve()}" in capsys.readouterr().out diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index 116b6ddbdd..f64f5d8a02 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -955,10 +955,14 @@ 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`?" + from nemo_gym.cli.utils import did_you_mean + + assert 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"]) == "" + from nemo_gym.cli.utils import did_you_mean + + assert 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) @@ -1135,3 +1139,10 @@ def test_search_dir_becomes_config_override(self, monkeypatch: MonkeyPatch) -> N # key, read centrally from the resolved config — like --json/--query. _, overrides = _dispatch_for(monkeypatch, ["list", "environments", "--search-dir", "/a", "--search-dir", "/b"]) assert overrides == ["+search_dir=[/a,/b]"] + + def test_name_positional_becomes_component_name_override(self, monkeypatch: MonkeyPatch) -> None: + # `gym list ` reaches the listing command as the reserved `component_name` config key, + # switching it into inspect mode. + target, overrides = _dispatch_for(monkeypatch, ["list", "environments", "calendar"]) + assert target == "nemo_gym.cli.env:list_environments" + assert overrides == ["+component_name=calendar"] diff --git a/tests/unit_tests/test_cli_models.py b/tests/unit_tests/test_cli_models.py index cc86696711..ee21cef979 100644 --- a/tests/unit_tests/test_cli_models.py +++ b/tests/unit_tests/test_cli_models.py @@ -16,6 +16,7 @@ from pathlib import Path from unittest.mock import patch +import pytest from omegaconf import OmegaConf from nemo_gym.cli.models import list_models @@ -91,3 +92,70 @@ def test_json_output_is_per_variant_rows(self, capsys) -> None: assert len(payload) == len(expected) for row in expected: assert row in payload + + def test_inspect_model_by_name(self, capsys) -> None: + with ( + patch( + "nemo_gym.cli.models.get_global_config_dict", + return_value=_mock_global_config({"component_name": "my_model"}), + ), + patch("nemo_gym.cli.models.discover_models", return_value=_MODELS), + ): + list_models() + out = capsys.readouterr().out + assert "The my_model model" in out + assert "model-types: my_model, my_model/some_other_flavor" in out + assert "Usage example:" not in out # thin view + + def test_inspect_model_by_flavor_token(self, capsys) -> None: + # A valid `/` token renders the model's (main) inspection. + with ( + patch( + "nemo_gym.cli.models.get_global_config_dict", + return_value=_mock_global_config({"component_name": "my_model/some_other_flavor"}), + ), + patch("nemo_gym.cli.models.discover_models", return_value=_MODELS), + ): + list_models() + out = capsys.readouterr().out + assert "The my_model model" in out + assert "model-types: my_model, my_model/some_other_flavor" in out + + def test_inspect_unknown_model_exits(self, capsys) -> None: + with ( + patch( + "nemo_gym.cli.models.get_global_config_dict", + return_value=_mock_global_config({"component_name": "mymodel"}), + ), + patch("nemo_gym.cli.models.discover_models", return_value=_MODELS), + ): + with pytest.raises(SystemExit): + list_models() + out = capsys.readouterr().out + assert "Unknown model 'mymodel'" in out and "my_model" in out + + def test_inspect_unknown_flavor_exits(self, capsys) -> None: + # A `/` token with an invalid flavor is rejected. + with ( + patch( + "nemo_gym.cli.models.get_global_config_dict", + return_value=_mock_global_config({"component_name": "my_model/nope"}), + ), + patch("nemo_gym.cli.models.discover_models", return_value=_MODELS), + ): + with pytest.raises(SystemExit): + list_models() + out = capsys.readouterr().out + assert "Unknown model 'my_model/nope'" in out + + def test_inspect_shows_absolute_path(self, tmp_path: Path, capsys) -> None: + # Real discovery (via --search-dir): the path line must be the model dir's absolute path. + model_dir = tmp_path / "responses_api_models" / "my_model" + (model_dir / "configs").mkdir(parents=True) + (model_dir / "configs" / "my_model.yaml").write_text("my_model: {}\n") + with patch( + "nemo_gym.cli.models.get_global_config_dict", + return_value=_mock_global_config({"component_name": "my_model", "search_dir": [str(tmp_path)]}), + ): + list_models() + assert f"path: {model_dir.resolve()}" in capsys.readouterr().out diff --git a/tests/unit_tests/test_cli_resources_servers.py b/tests/unit_tests/test_cli_resources_servers.py index 063a505fc3..da32c2bb3f 100644 --- a/tests/unit_tests/test_cli_resources_servers.py +++ b/tests/unit_tests/test_cli_resources_servers.py @@ -16,6 +16,7 @@ from pathlib import Path from unittest.mock import patch +import pytest from omegaconf import OmegaConf from nemo_gym.cli.resources_servers import list_resources_servers @@ -88,3 +89,43 @@ def test_query_filters_servers(self, capsys) -> None: out = capsys.readouterr().out assert "aviary" in out and "Resources servers matching" in out assert "mcqa" not in out and "knowledge" not in out + + def test_inspect_resources_server_by_name(self, capsys) -> None: + with ( + patch( + "nemo_gym.cli.resources_servers.get_global_config_dict", + return_value=_mock_global_config({"component_name": "mcqa"}), + ), + patch("nemo_gym.cli.resources_servers.discover_resources_servers", return_value=_SERVERS), + patch("nemo_gym.cli.resources_servers.read_resources_server_value", return_value="Improve MMLU"), + ): + list_resources_servers() + out = capsys.readouterr().out + assert "The mcqa resources server (domain: knowledge)" in out + assert "Value: Improve MMLU" in out + assert "gym env start --resources-server mcqa --model-type vllm_model" in out + + def test_inspect_unknown_resources_server_exits(self, capsys) -> None: + with ( + patch( + "nemo_gym.cli.resources_servers.get_global_config_dict", + return_value=_mock_global_config({"component_name": "mcq"}), + ), + patch("nemo_gym.cli.resources_servers.discover_resources_servers", return_value=_SERVERS), + ): + with pytest.raises(SystemExit): + list_resources_servers() + out = capsys.readouterr().out + assert "Unknown resources server 'mcq'" in out and "mcqa" in out + + def test_inspect_shows_absolute_config_path(self, tmp_path: Path, capsys) -> None: + # Real discovery (via --search-dir): the config line must be the flavor config's absolute path. + cfg = tmp_path / "resources_servers" / "my_rs" / "configs" / "my_rs.yaml" + cfg.parent.mkdir(parents=True) + cfg.write_text("my_rs:\n resources_servers:\n my_rs:\n domain: knowledge\n description: D\n") + with patch( + "nemo_gym.cli.resources_servers.get_global_config_dict", + return_value=_mock_global_config({"component_name": "my_rs", "search_dir": [str(tmp_path)]}), + ): + list_resources_servers() + assert f"config: {cfg.resolve()}" in capsys.readouterr().out diff --git a/tests/unit_tests/test_cli_utils.py b/tests/unit_tests/test_cli_utils.py index 767e46cdc9..853f63a339 100644 --- a/tests/unit_tests/test_cli_utils.py +++ b/tests/unit_tests/test_cli_utils.py @@ -14,10 +14,16 @@ # limitations under the License. from unittest.mock import MagicMock, PropertyMock, patch +import pytest from rich.console import Console from rich.table import Table -from nemo_gym.cli.utils import fuzzy_matches, print_rich_table +from nemo_gym.cli.utils import ( + exit_unknown_component, + fuzzy_matches, + print_rich_table, + render_component_inspection, +) # A cell value wider than Rich's 80-col non-TTY default, so a truncated render would ellipsize it. @@ -77,3 +83,39 @@ def test_skips_empty_fields(self) -> None: def test_no_match(self) -> None: assert not fuzzy_matches("zzznomatch", "aime24", "math_with_judge") + + +class TestExitUnknownComponent: + def test_exits_nonzero_with_did_you_mean(self, capsys) -> None: + with pytest.raises(SystemExit) as exc: + exit_unknown_component("calndr", ["calendar", "arc_agi"], "environment") + assert exc.value.code == 1 + out = capsys.readouterr().out + assert "Unknown environment 'calndr'" in out and "calendar" in out + + +class TestRenderComponentInspection: + def test_full_text_view(self, capsys) -> None: + render_component_inspection( + json_output=False, + name="calendar", + type_noun="environment", + domain="agent", + description="A calendar env.\nValue: Improve scheduling", # value folded in by the caller + details={"config": "/abs/config.yaml", "agent": "simple_agent"}, + usage="gym env start --environment calendar --model-type vllm_model", + ) + out = capsys.readouterr().out + assert "The calendar environment (domain: agent)" in out + assert "A calendar env." in out and "Value: Improve scheduling" in out + assert "Details:\nconfig: /abs/config.yaml\nagent: simple_agent" in out + assert "Usage example:\ngym env start --environment calendar --model-type vllm_model" in out + + def test_omits_empty_sections(self, capsys) -> None: + # A thin view (model): no domain suffix, no description block, no usage. + render_component_inspection( + json_output=False, name="vllm_model", type_noun="model", details={"path": "/abs/vllm_model"} + ) + out = capsys.readouterr().out + assert "The vllm_model model" in out and "path: /abs/vllm_model" in out + assert "(domain" not in out and "Usage example:" not in out diff --git a/tests/unit_tests/test_registry.py b/tests/unit_tests/test_registry.py index f9d8b837aa..cba32e710f 100644 --- a/tests/unit_tests/test_registry.py +++ b/tests/unit_tests/test_registry.py @@ -130,3 +130,32 @@ def test_cwd_is_scanned_by_default(self, tmp_path: Path, monkeypatch) -> None: monkeypatch.chdir(tmp_path) assert "cwd_env" in discover_environments() + + +class TestReadEnvironmentDetails: + def test_extracts_resources_servers_agent_datasets_and_value(self, tmp_path: Path) -> None: + from nemo_gym.registry import read_environment_details + + cfg = tmp_path / "config.yaml" + cfg.write_text( + "env:\n" + " resources_servers:\n" + " my_rs:\n" + " domain: agent\n" + " description: Desc\n" + " value: The value\n" + "env_agent:\n" + " responses_api_agents:\n" + " simple_agent:\n" + " datasets:\n" + " - {name: train, type: train}\n" + " - {name: example, type: example}\n" + ) + + details = read_environment_details(cfg) + + assert details["domain"] == "agent" and details["description"] == "Desc" + assert details["value"] == "The value" + assert details["resources_servers"] == ["my_rs"] + assert details["agent"] == "simple_agent" + assert details["datasets"] == ["train", "example"] diff --git a/tests/unit_tests/test_resources_server_registry.py b/tests/unit_tests/test_resources_server_registry.py index fc94cf084a..e3065d440f 100644 --- a/tests/unit_tests/test_resources_server_registry.py +++ b/tests/unit_tests/test_resources_server_registry.py @@ -76,3 +76,21 @@ def test_helper_configs_are_skipped(self, tmp_path: Path) -> None: def test_missing_directory_yields_no_servers(self, tmp_path: Path) -> None: assert _discover_resources_servers_in_dir(tmp_path / "nope") == {} + + +class TestReadResourcesServerValue: + def test_reads_value_from_config(self, tmp_path: Path) -> None: + from nemo_gym.resources_server_registry import read_resources_server_value + + cfg = tmp_path / "my_server.yaml" + cfg.write_text("s:\n resources_servers:\n my_server:\n value: The value\n") + + assert read_resources_server_value(cfg) == "The value" + + def test_returns_none_when_no_value(self, tmp_path: Path) -> None: + from nemo_gym.resources_server_registry import read_resources_server_value + + cfg = tmp_path / "my_server.yaml" + cfg.write_text("s:\n resources_servers:\n my_server:\n domain: x\n") + + assert read_resources_server_value(cfg) is None From a30097151d6677b72f0185f07ea774601aa3d766 Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Thu, 16 Jul 2026 09:09:02 +0200 Subject: [PATCH 09/13] fix: list benchmarks by config name so they can be passed as --benchmark Originally we've used the dataset name. These two are usually identical, but for benchmarks with non-standard config location the output from `gym list command` was not usable as value for `--benchmark` flag in other commands. This commit fixes it. Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/benchmarks.py | 54 ++++++++++++++----------- nemo_gym/cli/eval.py | 6 ++- tests/unit_tests/test_benchmarks.py | 63 ++++++++++++++++++++++------- 3 files changed, 82 insertions(+), 41 deletions(-) diff --git a/nemo_gym/benchmarks.py b/nemo_gym/benchmarks.py index 3caf8881f4..28bb5308dc 100644 --- a/nemo_gym/benchmarks.py +++ b/nemo_gym/benchmarks.py @@ -38,7 +38,7 @@ class BenchmarkConfig(BaseModel): - name: str + name: str # this is a dataset name, not the config name (they are usually the same) path: Path agent_name: str num_repeats: int @@ -100,29 +100,16 @@ def from_initial_config_dict( ) -def _load_benchmarks_from_config_paths(config_paths: List[Path]) -> Dict[str, BenchmarkConfig]: - benchmarks_dict = dict() - for config_path in config_paths: - config_path = Path(config_path) +def _benchmark_config_name(rel_config_path: Path) -> str: + """The name of the benchmark config, given its path relative to ``benchmarks/``, sans ``.yaml``. - try: - # Listing has no runtime context, so tolerate unset runtime-only values. - maybe_bc = BenchmarkConfig.from_config_path(config_path, strict=False) - except Exception as e: - # Still unresolvable (e.g. a multi-benchmark suite) — skip with a warning rather than fail the - # whole listing, so it isn't silently invisible. - print( - f"Warning: skipping benchmark config '{config_path}': could not resolve it " - f"({type(e).__name__}: {str(e).splitlines()[0]}).", - file=sys.stderr, - ) - continue - if not maybe_bc: - continue - - benchmarks_dict[maybe_bc.name] = maybe_bc - - return benchmarks_dict + This is the identity we key benchmarks by, so a listed benchmark is always a valid ``--benchmark`` argument. + """ + rel = rel_config_path.with_suffix("") + parts = rel.parts + if len(parts) == 2 and parts[1] == "config": + return parts[0] + return rel.as_posix() def _benchmark_config_paths(benchmarks_dir: Path) -> List[Path]: @@ -140,7 +127,26 @@ def _benchmark_config_paths(benchmarks_dir: Path) -> List[Path]: def _discover_benchmarks_in_dir(benchmarks_dir: Path) -> Dict[str, BenchmarkConfig]: """Map benchmark name -> :class:`BenchmarkConfig` for every benchmark config under one dir.""" - return _load_benchmarks_from_config_paths(_benchmark_config_paths(benchmarks_dir)) + benchmarks_dict = dict() + for config_path in _benchmark_config_paths(benchmarks_dir): + try: + # Listing has no runtime context, so tolerate unset runtime-only values. + maybe_bc = BenchmarkConfig.from_config_path(config_path, strict=False) + except Exception as e: + # Still unresolvable (e.g. a multi-benchmark suite) — skip with a warning rather than fail the + # whole listing, so it isn't silently invisible. + print( + f"Warning: skipping benchmark config '{config_path}': could not resolve it " + f"({type(e).__name__}: {str(e).splitlines()[0]}).", + file=sys.stderr, + ) + continue + if not maybe_bc: + continue + + benchmarks_dict[_benchmark_config_name(config_path.relative_to(benchmarks_dir))] = maybe_bc + + return benchmarks_dict def discover_benchmarks(search_dirs: Optional[Union[Path, Sequence[Path]]] = None) -> Dict[str, BenchmarkConfig]: diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index 30308c9f2c..0e3895c5db 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -117,11 +117,13 @@ def list_benchmarks() -> None: metadata = {name: read_config_metadata(bench.path) for name, bench in benchmarks.items()} # `gym search ` reuses this command, narrowing the listing to fuzzy matches - # across the benchmark name and domain. + # across the benchmark cofnig name, its dataset name, and domain. query = global_config_dict.get(QUERY_KEY_NAME) if query: benchmarks = { - name: bench for name, bench in benchmarks.items() if fuzzy_matches(query, name, metadata[name][0] or "") + name: bench + for name, bench in benchmarks.items() + if fuzzy_matches(query, name, bench.name, metadata[name][0] or "") } if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False): diff --git a/tests/unit_tests/test_benchmarks.py b/tests/unit_tests/test_benchmarks.py index c0e52db0ee..bb505fb1b0 100644 --- a/tests/unit_tests/test_benchmarks.py +++ b/tests/unit_tests/test_benchmarks.py @@ -12,7 +12,6 @@ # 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 glob import glob from pathlib import Path from unittest.mock import MagicMock, patch @@ -134,11 +133,15 @@ def test_inspect_shows_absolute_config_path(self, capsys) -> None: assert f"config: {expected}" in capsys.readouterr().out -class TestLoadBenchmarksFromConfigPaths: - def test_skips_configs_that_fail_to_resolve_with_warning(self, capsys) -> None: +class TestDiscoverBenchmarksInDir: + def test_skips_configs_that_fail_to_resolve_with_warning(self, tmp_path: Path, capsys) -> None: # A candidate that still can't be resolved even with tolerance (e.g. a multi-benchmark suite) must # be skipped with a warning — not crash the whole listing, and not vanish silently. - from nemo_gym.benchmarks import BenchmarkConfig, _load_benchmarks_from_config_paths + from nemo_gym.benchmarks import BenchmarkConfig, _discover_benchmarks_in_dir + + for name in ("bad", "good"): + (tmp_path / name).mkdir() + (tmp_path / name / "config.yaml").write_text("x:\n datasets:\n - type: benchmark\n") good = MagicMock() good.name = "good_bench" @@ -146,36 +149,65 @@ def test_skips_configs_that_fail_to_resolve_with_warning(self, capsys) -> None: # Listing must resolve tolerantly (it scans files with no runtime context), so it opts out of strict. def fake_from_config_path(path, *, strict=True): assert strict is False - if Path(path).name == "bad.yaml": + if Path(path).parent.name == "bad": raise RuntimeError("cannot resolve without runtime values") return good with patch.object(BenchmarkConfig, "from_config_path", side_effect=fake_from_config_path): - result = _load_benchmarks_from_config_paths([Path("bad.yaml"), Path("good.yaml")]) + result = _discover_benchmarks_in_dir(tmp_path) - assert set(result) == {"good_bench"} + # The surviving benchmark is keyed by its config name (path under the dir, sans `.yaml`), not `dataset.name`. + assert set(result) == {"good"} err = capsys.readouterr().err - assert "Warning" in err and "bad.yaml" in err + assert "Warning" in err and "bad" in err def test_every_repo_benchmark_appears_in_listing(self, capsys) -> None: # Every config that declares a `type: benchmark` dataset must surface as its own listing entry — - # no silent drop from a name collision (the name-keyed dict is last-writer-wins) or a resolve - # failure. Mirrors the content-based discovery in `list_benchmarks`. - from nemo_gym.benchmarks import BENCHMARKS_DIR, _load_benchmarks_from_config_paths + # no silent drop from a name collision (the dict is last-writer-wins) or a resolve failure. Keying by + # the `--benchmark` name (path under `benchmarks/`) makes each config's key unique by construction, so + # none can shadow another. + from nemo_gym.benchmarks import BENCHMARKS_DIR, _benchmark_config_paths, _discover_benchmarks_in_dir - config_paths = [BENCHMARKS_DIR / p for p in glob("**/*.yaml", root_dir=BENCHMARKS_DIR, recursive=True)] - config_paths = sorted(p for p in config_paths if "type: benchmark" in p.read_text(errors="ignore")) + config_paths = _benchmark_config_paths(BENCHMARKS_DIR) assert config_paths, "no benchmark configs discovered under BENCHMARKS_DIR" - benchmarks = _load_benchmarks_from_config_paths(config_paths) + benchmarks = _discover_benchmarks_in_dir(BENCHMARKS_DIR) assert len(benchmarks) == len(config_paths), ( f"{len(config_paths)} benchmark config(s) discovered but only {len(benchmarks)} appear in the " - f"listing — a duplicate dataset name or resolve failure is hiding at least one.\n" + f"listing — a duplicate name or resolve failure is hiding at least one.\n" f"stderr:\n{capsys.readouterr().err}" ) +class TestBenchmarkConfigName: + @pytest.mark.parametrize( + "rel, expected", + [ + ("aime24/config.yaml", "aime24"), # `/config.yaml` shortens to `` + ("tau2/configs/tau2.yaml", "tau2/configs/tau2"), # a flavor keeps its full relative path + ("livecodebench/v5_2408_2502/config.yaml", "livecodebench/v5_2408_2502/config"), # nested: no shorten + ], + ) + def test_name_matches_the_benchmark_selector(self, rel: str, expected: str) -> None: + from nemo_gym.benchmarks import _benchmark_config_name + + assert _benchmark_config_name(Path(rel)) == expected + + def test_every_listed_token_round_trips_through_the_benchmark_selector(self) -> None: + # The point of keying by token: every value `gym list benchmarks` prints must resolve back to its own + # config via `--benchmark`, using the same `_asset_config_path` mapping the CLI uses. This covers the + # benchmarks whose `dataset.name` diverges from their on-disk path (e.g. tau2, livecodebench). + from nemo_gym.benchmarks import discover_benchmarks + from nemo_gym.cli.main import _asset_config_path + + benchmarks = discover_benchmarks() + assert benchmarks, "no benchmarks discovered" + for token, bench in benchmarks.items(): + resolved = Path(_asset_config_path("benchmark", token)) + assert resolved.resolve() == bench.path.resolve(), f"token {token!r} does not select its own config" + + class TestBenchmarkConfigStrictParsing: def test_strict_is_the_default_and_does_not_tolerate_unresolved_values(self) -> None: # The tolerance is listing-only: `from_initial_config_dict` defaults to strict, so other workflows @@ -204,6 +236,7 @@ class TestSearchBenchmarks: def _bench(self, key: str): bench = MagicMock(agent_name="my_agent", num_repeats=1) + bench.name = key # `dataset.name`; also fuzzy-matched by `gym search` bench.path = key # the patched read_config_metadata keys off the path to find the domain return bench From e59aaa6ba17d63a7730e1151589c2fb790b6aea0 Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Thu, 16 Jul 2026 09:53:30 +0200 Subject: [PATCH 10/13] fix: custom 'did you mean?' for benchmarks Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/cli/main.py | 54 ++++++++++++++++++++++--------- tests/unit_tests/test_cli_main.py | 17 +++++++++- 2 files changed, 54 insertions(+), 17 deletions(-) diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 6eaeba7a3f..746e6cf12f 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -221,25 +221,47 @@ def _asset_config_path(flag: str, value: str, search_dirs: tuple[str, ...] = ()) 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 + # No match: build a "did you mean?" hint and the roots searched + if flag == "benchmark": + # Benchmarks need special handling because some use non-standard config paths (arbitrary nesting), so + # the generic one-level flavor/sibling search below can't see them. + # Enumerate their real config names (the same values `gym list benchmarks` prints) instead. + from nemo_gym.benchmarks import _benchmark_config_name, _benchmark_config_paths + + config_names = { + _benchmark_config_name(p.relative_to(root / parent)) for root in roots - if (root / parent).is_dir() - for child in (root / parent).iterdir() - if child.is_dir() - ] + for p in _benchmark_config_paths(root / parent) + } + # A bare directory that only groups benchmarks (e.g. `livecodebench`) is not itself selectable, so point + # at the config names under it; otherwise fall back to a fuzzy match across every token. + under_dir = sorted(config_name for config_name in config_names if config_name.startswith(f"{value}/")) + hint = f" Did you mean `{min(under_dir, key=len)}`?" if under_dir else did_you_mean(value, config_names) + available = ", ".join(sorted(f"`{(root / parent).resolve()}`" for root in roots if (root / parent).is_dir())) + else: + # Suggest the closest real name across all roots: a config flavor when the server exists, else a server + # name, reporting the full paths that were searched in each case. + 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() + ] + + hint = did_you_mean(typo, candidates) raise ValueError( - f"`--{flag} {value}` was specified which implies config `{path}`, which does not exist.{did_you_mean(typo, candidates)} " + f"`--{flag} {value}` was specified which implies config `{path}`, which does not exist.{hint} " f"See available {flag} configs in {available}." ) diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index f64f5d8a02..61c3127c9e 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -1000,6 +1000,19 @@ def test_misspelled_component_flavor(self, monkeypatch: MonkeyPatch, capsys) -> ) assert "Did you mean `dapo17k`?" in err + def test_benchmark_group_dir_suggests_a_nested_token(self, monkeypatch: MonkeyPatch, capsys) -> None: + # `livecodebench` is a directory that only groups benchmarks (its configs are nested), so it is not + # itself a valid `--benchmark` value. The hint must point at a real nested token, not circle back to + # the bare directory name. + err = self._run_expecting_exit(monkeypatch, capsys, ["eval", "run", "--benchmark", "livecodebench"]) + assert "Did you mean `livecodebench/" in err + assert "Did you mean `livecodebench`?" not in err + + def test_benchmark_group_dir_with_flavor_configs_subdir(self, monkeypatch: MonkeyPatch, capsys) -> None: + # tau2 keeps its benchmarks under `configs/`; the hint should surface one of those tokens. + err = self._run_expecting_exit(monkeypatch, capsys, ["eval", "prepare", "--benchmark", "tau2"]) + assert "Did you mean `tau2/configs/" in err + class TestSearchDir: """--search-dir registers extra roots that the name->config selectors also search (REQ 5).""" @@ -1007,7 +1020,9 @@ class TestSearchDir: 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") + # A `type: benchmark` dataset so the config is discoverable as a real benchmark (the "did you mean" + # suggestions enumerate real benchmark tokens); resolution itself only needs the file to exist. + (bench_dir / "config.yaml").write_text("x:\n datasets:\n - type: benchmark\n") def test_resolves_component_from_user_dir(self, monkeypatch: MonkeyPatch, tmp_path) -> None: self._make_user_benchmark(tmp_path) From 68fe1086d45bbc3ad0f9c5dd52c654585d78716b Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Thu, 16 Jul 2026 13:01:07 +0200 Subject: [PATCH 11/13] fix: incorporate NEMO_GYM_EXTRA_ROOTS env var logic from PR #1264 by @gwarmstrong conversly to search_dir config field, env var allows to pass extra dirs deep into the code, including rollouts colletion, prompt templates, config paths and env.yaml loading. also it is propagated to any spawned process Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/agent_registry.py | 11 +- nemo_gym/benchmarks.py | 9 +- nemo_gym/cli/agents.py | 3 +- nemo_gym/cli/env.py | 3 +- nemo_gym/cli/eval.py | 3 +- nemo_gym/cli/main.py | 149 ++++++++++-------- nemo_gym/cli/models.py | 3 +- nemo_gym/cli/resources_servers.py | 3 +- nemo_gym/discovery.py | 32 ++-- nemo_gym/global_config.py | 2 - nemo_gym/model_registry.py | 9 +- nemo_gym/registry.py | 11 +- nemo_gym/resources_server_registry.py | 11 +- tests/unit_tests/test_cli.py | 6 +- tests/unit_tests/test_cli_agents.py | 8 +- tests/unit_tests/test_cli_main.py | 46 +++++- tests/unit_tests/test_cli_models.py | 8 +- .../unit_tests/test_cli_resources_servers.py | 8 +- tests/unit_tests/test_discovery.py | 29 ++-- tests/unit_tests/test_registry.py | 6 +- 20 files changed, 205 insertions(+), 155 deletions(-) diff --git a/nemo_gym/agent_registry.py b/nemo_gym/agent_registry.py index 59294e35f1..164ad3b884 100644 --- a/nemo_gym/agent_registry.py +++ b/nemo_gym/agent_registry.py @@ -39,7 +39,7 @@ from dataclasses import dataclass from pathlib import Path -from typing import Dict, Optional, Sequence, Tuple, Union +from typing import Dict, Optional, Tuple from omegaconf import OmegaConf @@ -162,13 +162,10 @@ def _discover_agents_in_dir(agents_dir: Path) -> Dict[str, AgentEntry]: return agents -def discover_agents( - search_dirs: Optional[Union[Path, Sequence[Path]]] = None, -) -> Dict[str, AgentEntry]: +def discover_agents() -> Dict[str, AgentEntry]: """Map agent name -> :class:`AgentEntry` for every discoverable agent dir. Scans the ``responses_api_agents/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` - root (``search_dirs`` + cwd + built-ins), merged so user agents shadow same-named built-ins. - ``search_dirs`` is one dir or a list. + root (``NEMO_GYM_EXTRA_ROOTS`` + cwd + built-ins), merged so user agents shadow same-named built-ins. """ - return discover_components(AGENTS_SUBDIR, _discover_agents_in_dir, search_dirs) + return discover_components(AGENTS_SUBDIR, _discover_agents_in_dir) diff --git a/nemo_gym/benchmarks.py b/nemo_gym/benchmarks.py index 28bb5308dc..e5a0f40371 100644 --- a/nemo_gym/benchmarks.py +++ b/nemo_gym/benchmarks.py @@ -17,7 +17,7 @@ import sys from glob import glob from pathlib import Path -from typing import Dict, List, Optional, Sequence, Union +from typing import Dict, List, Optional from omegaconf import DictConfig, OmegaConf from pydantic import BaseModel @@ -149,14 +149,13 @@ def _discover_benchmarks_in_dir(benchmarks_dir: Path) -> Dict[str, BenchmarkConf return benchmarks_dict -def discover_benchmarks(search_dirs: Optional[Union[Path, Sequence[Path]]] = None) -> Dict[str, BenchmarkConfig]: +def discover_benchmarks() -> Dict[str, BenchmarkConfig]: """Map benchmark name -> :class:`BenchmarkConfig` for every discoverable benchmark config. Scans the ``benchmarks/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` root - (``search_dirs`` + cwd + built-ins), merged so user benchmarks shadow same-named built-ins. - ``search_dirs`` is one dir or a list. + (``NEMO_GYM_EXTRA_ROOTS`` + cwd + built-ins), merged so user benchmarks shadow same-named built-ins. """ - return discover_components(BENCHMARKS_SUBDIR, _discover_benchmarks_in_dir, search_dirs) + return discover_components(BENCHMARKS_SUBDIR, _discover_benchmarks_in_dir) # Backward-compatibility shims (CLI refactor): these symbols moved to `nemo_gym.cli.eval`. diff --git a/nemo_gym/cli/agents.py b/nemo_gym/cli/agents.py index c8a3b5428e..a5800e9ba0 100644 --- a/nemo_gym/cli/agents.py +++ b/nemo_gym/cli/agents.py @@ -29,7 +29,6 @@ COMPONENT_NAME_KEY_NAME, JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, - SEARCH_DIR_KEY_NAME, GlobalConfigDictParserConfig, get_global_config_dict, ) @@ -70,7 +69,7 @@ def list_agents() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - agents = discover_agents(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + agents = discover_agents() name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) if name: diff --git a/nemo_gym/cli/env.py b/nemo_gym/cli/env.py index 32d8b71460..5e8ba51fa6 100644 --- a/nemo_gym/cli/env.py +++ b/nemo_gym/cli/env.py @@ -55,7 +55,6 @@ NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME, NEMO_GYM_RESERVED_TOP_LEVEL_KEYS, QUERY_KEY_NAME, - SEARCH_DIR_KEY_NAME, GlobalConfigDictParser, GlobalConfigDictParserConfig, get_global_config_dict, @@ -980,7 +979,7 @@ def list_environments() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - environments = discover_environments(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + environments = discover_environments() name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) if name: diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index 0e3895c5db..9105ad9209 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -46,7 +46,6 @@ JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, - SEARCH_DIR_KEY_NAME, TASK_INDEX_KEY_NAME, GlobalConfigDictParserConfig, get_first_server_config_dict, @@ -105,7 +104,7 @@ def list_benchmarks() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - benchmarks = discover_benchmarks(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + benchmarks = discover_benchmarks() name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) if name: diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 746e6cf12f..a37f39e8c5 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -14,16 +14,21 @@ # limitations under the License. import argparse import importlib +import logging +import os import re import sys from collections.abc import Callable +from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from nemo_gym.cli.utils import did_you_mean -from nemo_gym.discovery import component_search_roots +from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, component_search_roots +logger = logging.getLogger(__name__) + VERSION_TARGET = "nemo_gym.cli.general:version" @@ -188,13 +193,14 @@ def _search(args: argparse.Namespace, overrides: list[str]) -> None: } -def _asset_config_path(flag: str, value: str, search_dirs: tuple[str, ...] = ()) -> str: +def _asset_config_path(flag: str, value: str) -> str: """Map a named asset (`name` or `name/flavor`) to its config path. - Searches the roots from :func:`~nemo_gym.discovery.component_search_roots` (``--search-dir`` + cwd + - install root), the same helper that backs `gym list`/`gym search`, so config resolution and discovery - agree on where components live. Searching the install root is what lets built-ins resolve by name from - an arbitrary cwd (e.g. a wheel install), not just inside the repo checkout. + Searches the roots from :func:`~nemo_gym.discovery.component_search_roots` (``NEMO_GYM_EXTRA_ROOTS`` + + cwd + install root), the same helper that backs `gym list`/`gym search`, so config resolution and + discovery agree on where components live. ``--search-dir`` reaches here via ``NEMO_GYM_EXTRA_ROOTS`` (set + in ``main``). Searching the install root is what lets built-ins resolve by name from an arbitrary cwd + (e.g. a wheel install), not just inside the repo checkout. """ parent, subdir, default_flavor = _ASSETS[flag] server_name, _, config_flavor = value.partition("/") @@ -202,7 +208,7 @@ def _asset_config_path(flag: str, value: str, search_dirs: tuple[str, ...] = ()) config_dir = f"{parent}/{server_name}/{subdir}".rstrip("/") path = f"{config_dir}/{config_flavor}.yaml" - roots = component_search_roots(search_dirs) + roots = component_search_roots() matches: list[Path] = [] for root in roots: @@ -272,11 +278,7 @@ def _asset_selector(flag: str) -> Flag: 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 [] + [f"+config_paths=[{_asset_config_path(flag, getattr(args, dest))}]"] if getattr(args, dest) else [] ), ) @@ -286,31 +288,16 @@ def _asset_selector(flag: str) -> Flag: RESOURCES_SERVER_CONFIG = _asset_selector("resources-server") MODEL_TYPE = _asset_selector("model-type") -# `--search-dir` for the asset selectors above: read straight from argv during config resolution, not -# emitted as a Hydra override. On every command that accepts a -- NAME selector. +# `--search-dir`: extra component-search roots. `main()` folds these into the `NEMO_GYM_EXTRA_ROOTS` env +# var before dispatch (see there), so a single register-only flag suffices for every command — the roots +# reach discovery, the `-- NAME` selectors, deep path resolution, and spawned servers alike. 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.", - ), -) - -# `--search-dir` for the `list`/`search` commands. Their targets are called with no args, so — like -# --json/--query — the roots reach them as the reserved `search_dir` config key, not read from argv. -DISCOVERY_SEARCH_DIR = Flag( register=lambda p: p.add_argument( "--search-dir", action="append", metavar="DIR", help="Extra root directory to search for components; repeatable.", ), - translate_to_hydra=lambda args: ( - [f"+search_dir=[{','.join(getattr(args, 'search_dir', None) or [])}]"] - if getattr(args, "search_dir", None) - else [] - ), ) @@ -371,32 +358,32 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: "list benchmarks": Command( target="nemo_gym.cli.eval:list_benchmarks", summary="List or inspect available benchmarks.", - flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), + flags=(NAME, JSON, SEARCH_DIR), ), "list environments": Command( target="nemo_gym.cli.env:list_environments", summary="List or inspect available environments.", - flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), + flags=(NAME, JSON, SEARCH_DIR), ), "list agents": Command( target="nemo_gym.cli.agents:list_agents", summary="List or inspect available agent harnesses.", - flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), + flags=(NAME, JSON, SEARCH_DIR), ), "list models": Command( target="nemo_gym.cli.models:list_models", summary="List or inspect available model servers.", - flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), + flags=(NAME, JSON, SEARCH_DIR), ), "list resources-servers": Command( target="nemo_gym.cli.resources_servers:list_resources_servers", summary="List or inspect available resources servers.", - flags=(NAME, JSON, DISCOVERY_SEARCH_DIR), + flags=(NAME, JSON, SEARCH_DIR), ), "search": Command( target=_search, summary="Search a component type (default benchmarks) by name; like `list` filtered to a query.", - flags=(SEARCH_TERMS, JSON, DISCOVERY_SEARCH_DIR), + flags=(SEARCH_TERMS, JSON, SEARCH_DIR), ), "dataset upload": Command( target=_dataset_upload, @@ -687,6 +674,32 @@ def _handle_pydantic_validation_error(exc, parser: argparse.ArgumentParser) -> N parser.error(" ".join(parts) if parts else str(exc)) +@contextmanager +def _extra_roots_from_search_dir(search_dirs: list[str] | None): + """Set ``NEMO_GYM_EXTRA_ROOTS`` to ``--search-dir`` for the duration of the command, then restore. + + Setting the env var lets the roots reach every resolver (discovery, the -- selectors, + deep config/prompt/rollout resolution) and inherit into spawned server subprocesses. + The original value is restored (or the var unset) on exit so main() leaves no global side effect. + """ + if not search_dirs: + yield + return + original = os.environ.get(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME) + value = os.pathsep.join(search_dirs) + os.environ[NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME] = value + logger.debug(f"Set {NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME}={value} from --search-dir") + try: + yield + finally: + if original is None: + os.environ.pop(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, None) + logger.debug(f"Unset {NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME}") + else: + os.environ[NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME] = original + logger.debug(f"Restored {NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME}={original}") + + def main() -> None: parser = build_parser() args, overrides = parser.parse_known_args() @@ -699,34 +712,36 @@ def main() -> None: 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: - getattr(args, "_parser", parser).error(str(exc)) - - # --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] - - # Local import keeps `gym --help` (which returns before this point) free of pydantic's import cost; - # any real command loads pydantic anyway via its config's model_validate. - from pydantic import ValidationError - - try: - if callable(command.target): - command.target(args, overrides) - else: - dispatch(command.target, overrides) - except ValidationError as exc: - _handle_pydantic_validation_error(exc, getattr(args, "_parser", parser)) + # set NEMO_GYM_EXTRA_ROOTS from --search-dir for the duration of the command + with _extra_roots_from_search_dir(getattr(args, "search_dir", None)): + 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: + getattr(args, "_parser", parser).error(str(exc)) + + # --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] + + # Local import keeps `gym --help` (which returns before this point) free of pydantic's import cost; + # any real command loads pydantic anyway via its config's model_validate. + from pydantic import ValidationError + + try: + if callable(command.target): + command.target(args, overrides) + else: + dispatch(command.target, overrides) + except ValidationError as exc: + _handle_pydantic_validation_error(exc, getattr(args, "_parser", parser)) diff --git a/nemo_gym/cli/models.py b/nemo_gym/cli/models.py index 212128ba94..c5e8918705 100644 --- a/nemo_gym/cli/models.py +++ b/nemo_gym/cli/models.py @@ -28,7 +28,6 @@ COMPONENT_NAME_KEY_NAME, JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, - SEARCH_DIR_KEY_NAME, GlobalConfigDictParserConfig, get_global_config_dict, ) @@ -71,7 +70,7 @@ def list_models() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - models = discover_models(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + models = discover_models() name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) if name: diff --git a/nemo_gym/cli/resources_servers.py b/nemo_gym/cli/resources_servers.py index dbc296e08a..ea67b30c5c 100644 --- a/nemo_gym/cli/resources_servers.py +++ b/nemo_gym/cli/resources_servers.py @@ -28,7 +28,6 @@ COMPONENT_NAME_KEY_NAME, JSON_OUTPUT_KEY_NAME, QUERY_KEY_NAME, - SEARCH_DIR_KEY_NAME, GlobalConfigDictParserConfig, get_global_config_dict, ) @@ -70,7 +69,7 @@ def list_resources_servers() -> None: ) BaseNeMoGymCLIConfig.model_validate(global_config_dict) - servers = discover_resources_servers(search_dirs=global_config_dict.get(SEARCH_DIR_KEY_NAME)) + servers = discover_resources_servers() name = global_config_dict.get(COMPONENT_NAME_KEY_NAME) if name: diff --git a/nemo_gym/discovery.py b/nemo_gym/discovery.py index 418856047d..6adac0e2aa 100644 --- a/nemo_gym/discovery.py +++ b/nemo_gym/discovery.py @@ -19,10 +19,11 @@ they can share it without depending on each other. Reads configs only; never starts servers. """ +import os import re from copy import deepcopy from pathlib import Path -from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple, TypeVar, Union +from typing import Callable, Dict, Iterable, List, Optional, Tuple, TypeVar from omegaconf import DictConfig, OmegaConf from omegaconf.errors import InterpolationKeyError @@ -37,24 +38,27 @@ _T = TypeVar("_T") +# Extra component-search roots, `os.pathsep`-separated; same effect as repeating `--search-dir` +# (see :func:`component_search_roots`). +NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME = "NEMO_GYM_EXTRA_ROOTS" -def component_search_roots(search_dirs: Optional[Union[Path, Sequence[Path]]] = None) -> List[Path]: - """Ordered, de-duplicated roots to look for a Gym component under: any ``search_dirs`` (one dir or a - list, e.g. from ``--search-dir``), then cwd, then ``WORKING_DIR`` and the install root (``PARENT_DIR``, + +def component_search_roots() -> List[Path]: + """Ordered, de-duplicated roots to look for a Gym component under: the roots from the + ``NEMO_GYM_EXTRA_ROOTS`` env var, then cwd, then ``WORKING_DIR`` and the install root (``PARENT_DIR``, the built-ins). + ``NEMO_GYM_EXTRA_ROOTS`` is an ``os.pathsep``-separated list of extra roots — the single source of extra + roots. ``--search-dir`` is folded into it up front (see ``nemo_gym.cli.main.main``), so the flag and the + env var are one and the same channel here. + Earlier roots win on a name collision (see :func:`merge_by_name`), so user components shadow built-ins. De-duplicated by resolved path, since cwd/``WORKING_DIR``/install root coincide in an editable checkout. The single source of truth for where Gym looks for components — used by both config resolution (``_asset_config_path``) and the ``gym list``/``gym search`` discovery functions. """ - if search_dirs is None: - extra: List[Path] = [] - elif isinstance(search_dirs, (str, Path)): - extra = [Path(search_dirs)] # a single dir - else: - extra = [Path(d) for d in search_dirs] # a list of dirs - candidates: List[Path] = [*extra, Path.cwd(), WORKING_DIR, PARENT_DIR] + env_roots = [Path(d) for d in os.environ.get(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, "").split(os.pathsep) if d] + candidates: List[Path] = [*env_roots, Path.cwd(), WORKING_DIR, PARENT_DIR] roots: List[Path] = [] seen: set[Path] = set() for root in candidates: @@ -79,16 +83,14 @@ def merge_by_name(per_root: Iterable[Dict[str, _T]]) -> Dict[str, _T]: def discover_components( subdir: str, dir_scanning_fn: Callable[[Path], Dict[str, _T]], - search_dirs: Optional[Union[Path, Sequence[Path]]] = None, ) -> Dict[str, _T]: """Run ``dir_scanning_fn`` on ``subdir`` of every :func:`component_search_roots` root and merge the results. The shared body of ``discover_environments``/``discover_agents``/``discover_models``/ ``discover_benchmarks``: each passes its ``/`` subdir and a single-directory scan function, and - gets user-shadows-built-in merging (via :func:`merge_by_name`) for free. ``search_dirs`` is one dir or - a list. + gets user-shadows-built-in merging (via :func:`merge_by_name`) for free. """ - return merge_by_name(dir_scanning_fn(root / subdir) for root in component_search_roots(search_dirs)) + return merge_by_name(dir_scanning_fn(root / subdir) for root in component_search_roots()) # Fills unset `???`/`${...}` values during listing: they reference runtime-only values (API keys, diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 36e9fc9b34..d158c9ffe4 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -89,7 +89,6 @@ QUERY_KEY_NAME = "query" OBSERVABILITY_ENABLED_KEY_NAME = "observability_enabled" MODEL_CALL_CAPTURE_DIR_KEY_NAME = "model_call_capture_dir" -SEARCH_DIR_KEY_NAME = "search_dir" COMPONENT_NAME_KEY_NAME = "component_name" NEMO_GYM_RESERVED_TOP_LEVEL_KEYS = [ CONFIG_PATHS_KEY_NAME, @@ -118,7 +117,6 @@ QUERY_KEY_NAME, OBSERVABILITY_ENABLED_KEY_NAME, MODEL_CALL_CAPTURE_DIR_KEY_NAME, - SEARCH_DIR_KEY_NAME, COMPONENT_NAME_KEY_NAME, ] diff --git a/nemo_gym/model_registry.py b/nemo_gym/model_registry.py index 216033696d..30ff8c1941 100644 --- a/nemo_gym/model_registry.py +++ b/nemo_gym/model_registry.py @@ -21,7 +21,7 @@ from dataclasses import dataclass from pathlib import Path -from typing import Dict, List, Optional, Sequence, Tuple, Union +from typing import Dict, List, Tuple from nemo_gym import PARENT_DIR from nemo_gym.discovery import discover_components @@ -75,11 +75,10 @@ def _discover_models_in_dir(models_dir: Path) -> Dict[str, ModelEntry]: return models -def discover_models(search_dirs: Optional[Union[Path, Sequence[Path]]] = None) -> Dict[str, ModelEntry]: +def discover_models() -> Dict[str, ModelEntry]: """Map model name -> :class:`ModelEntry` for every discoverable model server. Scans the ``responses_api_models/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` - root (``search_dirs`` + cwd + built-ins), merged so user models shadow same-named built-ins. - ``search_dirs`` is one dir or a list. + root (``NEMO_GYM_EXTRA_ROOTS`` + cwd + built-ins), merged so user models shadow same-named built-ins. """ - return discover_components(MODELS_SUBDIR, _discover_models_in_dir, search_dirs) + return discover_components(MODELS_SUBDIR, _discover_models_in_dir) diff --git a/nemo_gym/registry.py b/nemo_gym/registry.py index 07805ffcef..8569bd8f50 100644 --- a/nemo_gym/registry.py +++ b/nemo_gym/registry.py @@ -27,7 +27,7 @@ from dataclasses import dataclass from pathlib import Path -from typing import Dict, List, Optional, Sequence, Union +from typing import Dict, List, Optional from omegaconf import DictConfig, OmegaConf @@ -77,16 +77,13 @@ def _discover_environments_in_dir(environments_dir: Path) -> Dict[str, Environme return environments -def discover_environments( - search_dirs: Optional[Union[Path, Sequence[Path]]] = None, -) -> Dict[str, EnvironmentEntry]: +def discover_environments() -> Dict[str, EnvironmentEntry]: """Map environment name -> :class:`EnvironmentEntry` for every discoverable ``/config.yaml``. Scans the ``environments/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` root - (``search_dirs`` + cwd + built-ins), merged so user environments shadow same-named built-ins. - ``search_dirs`` is one dir or a list. + (``NEMO_GYM_EXTRA_ROOTS`` + cwd + built-ins), merged so user environments shadow same-named built-ins. """ - return discover_components(ENVIRONMENTS_SUBDIR, _discover_environments_in_dir, search_dirs) + return discover_components(ENVIRONMENTS_SUBDIR, _discover_environments_in_dir) def read_environment_details(config_path: Path) -> Dict[str, object]: diff --git a/nemo_gym/resources_server_registry.py b/nemo_gym/resources_server_registry.py index 587d2fdf4c..467e859a18 100644 --- a/nemo_gym/resources_server_registry.py +++ b/nemo_gym/resources_server_registry.py @@ -22,7 +22,7 @@ from dataclasses import dataclass from pathlib import Path -from typing import Dict, Optional, Sequence, Union +from typing import Dict, Optional from omegaconf import OmegaConf @@ -93,16 +93,13 @@ def _discover_resources_servers_in_dir(resources_servers_dir: Path) -> Dict[str, return servers -def discover_resources_servers( - search_dirs: Optional[Union[Path, Sequence[Path]]] = None, -) -> Dict[str, ResourcesServerEntry]: +def discover_resources_servers() -> Dict[str, ResourcesServerEntry]: """Map resources-server name -> :class:`ResourcesServerEntry` for every discoverable server. Scans the ``resources_servers/`` subdir of every :func:`~nemo_gym.discovery.component_search_roots` - root (``search_dirs`` + cwd + built-ins), merged so user servers shadow same-named built-ins. - ``search_dirs`` is one dir or a list. + root (``NEMO_GYM_EXTRA_ROOTS`` + cwd + built-ins), merged so user servers shadow same-named built-ins. """ - return discover_components(RESOURCES_SERVERS_SUBDIR, _discover_resources_servers_in_dir, search_dirs) + return discover_components(RESOURCES_SERVERS_SUBDIR, _discover_resources_servers_in_dir) def read_resources_server_value(config_path: Path) -> Optional[str]: diff --git a/tests/unit_tests/test_cli.py b/tests/unit_tests/test_cli.py index 8325c41733..9685cc0230 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -46,6 +46,7 @@ ) from nemo_gym.cli.utils import exit_cleanly_on_config_error from nemo_gym.config_types import ConfigError, NoServerInstancesError, ResourcesServerInstanceConfig +from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.registry import EnvironmentEntry @@ -531,14 +532,15 @@ def test_inspect_unknown_environment_exits(self, monkeypatch: MonkeyPatch, capsy assert "Unknown environment 'alfa'" in out and "alpha" in out def test_inspect_shows_absolute_config_path(self, monkeypatch: MonkeyPatch, capsys, tmp_path: Path) -> None: - # Real discovery (via --search-dir): the config line must be the config's absolute path. + # Real discovery (via an extra root): the config line must be the config's absolute path. cfg = tmp_path / "environments" / "my_env" / "config.yaml" cfg.parent.mkdir(parents=True) cfg.write_text("my_env:\n resources_servers:\n my_env:\n domain: agent\n description: D\n") + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, str(tmp_path)) monkeypatch.setattr( nemo_gym.cli.env, "get_global_config_dict", - lambda **k: OmegaConf.create({"component_name": "my_env", "search_dir": [str(tmp_path)]}), + lambda **k: OmegaConf.create({"component_name": "my_env"}), ) list_environments() diff --git a/tests/unit_tests/test_cli_agents.py b/tests/unit_tests/test_cli_agents.py index 4e0b6169e3..3d881fb9ae 100644 --- a/tests/unit_tests/test_cli_agents.py +++ b/tests/unit_tests/test_cli_agents.py @@ -21,6 +21,7 @@ from nemo_gym.agent_registry import AgentEntry from nemo_gym.cli.agents import list_agents +from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME def _mock_global_config(config: dict = None): @@ -117,14 +118,15 @@ def test_inspect_unknown_agent_exits(self, capsys) -> None: out = capsys.readouterr().out assert "Unknown agent 'swe_agent'" in out and "swe_agents" in out - def test_inspect_shows_absolute_path(self, tmp_path: Path, capsys) -> None: - # Real discovery (via --search-dir): the path line must be the agent dir's absolute path. + def test_inspect_shows_absolute_path(self, tmp_path: Path, capsys, monkeypatch) -> None: + # Real discovery (via an extra root): the path line must be the agent dir's absolute path. agent_dir = tmp_path / "responses_api_agents" / "my_agent" agent_dir.mkdir(parents=True) (agent_dir / "app.py").write_text("") + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, str(tmp_path)) with patch( "nemo_gym.cli.agents.get_global_config_dict", - return_value=_mock_global_config({"component_name": "my_agent", "search_dir": [str(tmp_path)]}), + return_value=_mock_global_config({"component_name": "my_agent"}), ): list_agents() assert f"path: {agent_dir.resolve()}" in capsys.readouterr().out diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index 61c3127c9e..2f0de5e6f0 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import logging +import os import sys import types @@ -23,9 +24,18 @@ import nemo_gym.global_config as gc from nemo_gym import WORKING_DIR from nemo_gym.cli.main import main +from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.global_config import NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME +@pytest.fixture(autouse=True) +def _isolate_extra_roots_env(monkeypatch: MonkeyPatch): + # `main()` folds `--search-dir` into NEMO_GYM_EXTRA_ROOTS by mutating os.environ directly; delenv gives + # each test a clean baseline and restores the original on teardown (even after main() reassigns the key), + # so the roots never leak between tests. + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) + + 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 = {} @@ -1149,11 +1159,37 @@ def test_list_environments_json_dispatches(self, monkeypatch: MonkeyPatch) -> No assert target == "nemo_gym.cli.env:list_environments" assert overrides == ["+json=true"] - def test_search_dir_becomes_config_override(self, monkeypatch: MonkeyPatch) -> None: - # `--search-dir` (repeatable) reaches the no-arg list command as the reserved `search_dir` config - # key, read centrally from the resolved config — like --json/--query. - _, overrides = _dispatch_for(monkeypatch, ["list", "environments", "--search-dir", "/a", "--search-dir", "/b"]) - assert overrides == ["+search_dir=[/a,/b]"] + def test_search_dir_populates_extra_roots_env_during_command_then_restores(self, monkeypatch: MonkeyPatch) -> None: + # `--search-dir` (repeatable) is folded into NEMO_GYM_EXTRA_ROOTS for the duration of the command (no + # Hydra override) so the roots reach every resolver, then restored so main() leaves no global state. + seen = {} + + def fake_dispatch(target: str, overrides: list[str]) -> None: + seen["overrides"] = overrides + seen["env"] = os.environ.get(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME) + + monkeypatch.setattr(cli_main, "dispatch", fake_dispatch) + monkeypatch.setattr(sys, "argv", ["gym", "list", "environments", "--search-dir", "/a", "--search-dir", "/b"]) + main() + + assert seen["overrides"] == [] + assert seen["env"] == os.pathsep.join(["/a", "/b"]) # set while the command runs + assert NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME not in os.environ # restored (unset) after main() + + def test_search_dir_restores_pre_existing_extra_roots_env(self, monkeypatch: MonkeyPatch) -> None: + # A pre-existing NEMO_GYM_EXTRA_ROOTS is fully replaced by --search-dir for the command, then restored. + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, "/pre") + seen = {} + + def fake_dispatch(target: str, overrides: list[str]) -> None: + seen["env"] = os.environ.get(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME) + + monkeypatch.setattr(cli_main, "dispatch", fake_dispatch) + monkeypatch.setattr(sys, "argv", ["gym", "list", "environments", "--search-dir", "/a"]) + main() + + assert seen["env"] == "/a" # flag fully replaces the existing value for the command + assert os.environ[NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME] == "/pre" # original restored after main() def test_name_positional_becomes_component_name_override(self, monkeypatch: MonkeyPatch) -> None: # `gym list ` reaches the listing command as the reserved `component_name` config key, diff --git a/tests/unit_tests/test_cli_models.py b/tests/unit_tests/test_cli_models.py index ee21cef979..166b8d7844 100644 --- a/tests/unit_tests/test_cli_models.py +++ b/tests/unit_tests/test_cli_models.py @@ -20,6 +20,7 @@ from omegaconf import OmegaConf from nemo_gym.cli.models import list_models +from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.model_registry import ModelEntry @@ -148,14 +149,15 @@ def test_inspect_unknown_flavor_exits(self, capsys) -> None: out = capsys.readouterr().out assert "Unknown model 'my_model/nope'" in out - def test_inspect_shows_absolute_path(self, tmp_path: Path, capsys) -> None: - # Real discovery (via --search-dir): the path line must be the model dir's absolute path. + def test_inspect_shows_absolute_path(self, tmp_path: Path, capsys, monkeypatch) -> None: + # Real discovery (via an extra root): the path line must be the model dir's absolute path. model_dir = tmp_path / "responses_api_models" / "my_model" (model_dir / "configs").mkdir(parents=True) (model_dir / "configs" / "my_model.yaml").write_text("my_model: {}\n") + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, str(tmp_path)) with patch( "nemo_gym.cli.models.get_global_config_dict", - return_value=_mock_global_config({"component_name": "my_model", "search_dir": [str(tmp_path)]}), + return_value=_mock_global_config({"component_name": "my_model"}), ): list_models() assert f"path: {model_dir.resolve()}" in capsys.readouterr().out diff --git a/tests/unit_tests/test_cli_resources_servers.py b/tests/unit_tests/test_cli_resources_servers.py index da32c2bb3f..be47a4884f 100644 --- a/tests/unit_tests/test_cli_resources_servers.py +++ b/tests/unit_tests/test_cli_resources_servers.py @@ -20,6 +20,7 @@ from omegaconf import OmegaConf from nemo_gym.cli.resources_servers import list_resources_servers +from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.resources_server_registry import ResourcesServerEntry @@ -118,14 +119,15 @@ def test_inspect_unknown_resources_server_exits(self, capsys) -> None: out = capsys.readouterr().out assert "Unknown resources server 'mcq'" in out and "mcqa" in out - def test_inspect_shows_absolute_config_path(self, tmp_path: Path, capsys) -> None: - # Real discovery (via --search-dir): the config line must be the flavor config's absolute path. + def test_inspect_shows_absolute_config_path(self, tmp_path: Path, capsys, monkeypatch) -> None: + # Real discovery (via an extra root): the config line must be the flavor config's absolute path. cfg = tmp_path / "resources_servers" / "my_rs" / "configs" / "my_rs.yaml" cfg.parent.mkdir(parents=True) cfg.write_text("my_rs:\n resources_servers:\n my_rs:\n domain: knowledge\n description: D\n") + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, str(tmp_path)) with patch( "nemo_gym.cli.resources_servers.get_global_config_dict", - return_value=_mock_global_config({"component_name": "my_rs", "search_dir": [str(tmp_path)]}), + return_value=_mock_global_config({"component_name": "my_rs"}), ): list_resources_servers() assert f"config: {cfg.resolve()}" in capsys.readouterr().out diff --git a/tests/unit_tests/test_discovery.py b/tests/unit_tests/test_discovery.py index 6b63a9bead..4f29067b82 100644 --- a/tests/unit_tests/test_discovery.py +++ b/tests/unit_tests/test_discovery.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 os from pathlib import Path from omegaconf import OmegaConf @@ -19,6 +20,7 @@ from nemo_gym import PARENT_DIR from nemo_gym.discovery import ( _UNSET_VALUE_PLACEHOLDER, + NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, _parse_no_environment_tolerating_unset_values, component_search_roots, merge_by_name, @@ -32,29 +34,32 @@ def test_default_includes_cwd_and_install_root(self) -> None: assert Path.cwd().resolve() in resolved assert PARENT_DIR.resolve() in resolved - def test_search_dirs_take_precedence_and_keep_order(self, tmp_path: Path) -> None: + def test_env_var_roots_added_before_builtins(self, tmp_path: Path, monkeypatch) -> None: + # NEMO_GYM_EXTRA_ROOTS is the sole source of extra roots (--search-dir is folded into it). Its roots + # are searched ahead of cwd/built-ins, in listed order, so they can shadow them. a = tmp_path / "a" b = tmp_path / "b" a.mkdir() b.mkdir() + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, os.pathsep.join([str(a), str(b)])) - roots = component_search_roots(search_dirs=[a, b]) + roots = component_search_roots() - assert roots[0] == a # explicit search dirs come first, in the given order - assert roots[1] == b + assert roots[0] == a and roots[1] == b # os.pathsep-separated, order preserved assert PARENT_DIR.resolve() in {root.resolve() for root in roots} # built-ins still scanned - def test_accepts_a_single_dir(self, tmp_path: Path) -> None: - # `search_dirs` takes one dir or a list; a lone Path must be treated as that single root. - assert component_search_roots(tmp_path)[0] == tmp_path - - def test_dedupes_roots_by_resolved_path(self) -> None: - # Passing the install root as an explicit search dir must not scan it twice. - roots = component_search_roots(search_dirs=[PARENT_DIR]) + def test_dedupes_roots_by_resolved_path(self, monkeypatch) -> None: + # An extra root that resolves to the install root must not be scanned twice. + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, str(PARENT_DIR)) + roots = component_search_roots() resolved = [root.resolve() for root in roots] assert resolved.count(PARENT_DIR.resolve()) == 1 - assert roots[0].resolve() == PARENT_DIR.resolve() # the explicit search dir still takes precedence + assert roots[0].resolve() == PARENT_DIR.resolve() # the extra root still takes precedence + + def test_empty_or_unset_env_var_adds_nothing(self, tmp_path: Path, monkeypatch) -> None: + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, "") + assert component_search_roots()[0].resolve() == Path.cwd().resolve() class TestMergeByName: diff --git a/tests/unit_tests/test_registry.py b/tests/unit_tests/test_registry.py index cba32e710f..67836eeb63 100644 --- a/tests/unit_tests/test_registry.py +++ b/tests/unit_tests/test_registry.py @@ -14,6 +14,7 @@ # limitations under the License. from pathlib import Path +from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.registry import _discover_environments_in_dir, discover_environments @@ -117,10 +118,11 @@ def test_workplace_assistant_is_discoverable(self) -> None: class TestDiscoverEnvironmentsAcrossRoots: - def test_search_dirs_surface_user_environments_alongside_builtins(self, tmp_path: Path) -> None: + def test_extra_root_surfaces_user_environments_alongside_builtins(self, tmp_path: Path, monkeypatch) -> None: _make_env(tmp_path / "environments", "custom_env", _ENV_CONFIG.format(name="custom_env")) + monkeypatch.setenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, str(tmp_path)) - environments = discover_environments(search_dirs=[tmp_path]) + environments = discover_environments() assert "custom_env" in environments # a user-supplied environment is discovered assert "workplace_assistant" in environments # ...alongside the built-ins From fd2d92080f0fad0ed942b05594c4af87fdb08b28 Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Thu, 16 Jul 2026 13:10:02 +0200 Subject: [PATCH 12/13] fix: set logging to DEBUG on top of main so it applies to in-script logging too Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/cli/main.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index a37f39e8c5..18aff51c9a 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -704,6 +704,10 @@ def main() -> None: parser = build_parser() args, overrides = parser.parse_known_args() + if getattr(args, "verbose", False): + logging.basicConfig(level=logging.DEBUG) + logging.getLogger().setLevel(logging.DEBUG) + # 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: From 8c19560abedb8f8bab514bb1c1fdf53db434b7e6 Mon Sep 17 00:00:00 2001 From: Marta Stepniewska-Dziubinska Date: Thu, 16 Jul 2026 13:54:39 +0200 Subject: [PATCH 13/13] fix: use user-provided root dirs to resolve all paths (as did in PR #1264 from @gwarmstrong) Signed-off-by: Marta Stepniewska-Dziubinska --- nemo_gym/__init__.py | 87 +++++++++++++++---- nemo_gym/cli/env.py | 26 +++--- nemo_gym/cli/main.py | 4 +- nemo_gym/discovery.py | 35 +------- nemo_gym/global_config.py | 24 +++-- nemo_gym/prompt.py | 4 +- nemo_gym/rollout_collection.py | 8 +- tests/unit_tests/test_cli.py | 3 +- tests/unit_tests/test_cli_agents.py | 2 +- tests/unit_tests/test_cli_main.py | 19 ++-- tests/unit_tests/test_cli_models.py | 2 +- .../unit_tests/test_cli_resources_servers.py | 2 +- tests/unit_tests/test_discovery.py | 4 +- tests/unit_tests/test_global_config.py | 22 +++-- tests/unit_tests/test_prompt.py | 7 +- tests/unit_tests/test_registry.py | 2 +- 16 files changed, 139 insertions(+), 112 deletions(-) diff --git a/nemo_gym/__init__.py b/nemo_gym/__init__.py index 567e2faf76..50ccd931c1 100644 --- a/nemo_gym/__init__.py +++ b/nemo_gym/__init__.py @@ -12,9 +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 os import sys from os import environ from pathlib import Path +from typing import Callable, List, Optional, Union # /path/to/dir/Gym (PARENT_DIR) @@ -35,31 +37,80 @@ CACHE_DIR = WORKING_DIR / "cache" RESULTS_DIR = WORKING_DIR / "results" -sys.path.append(str(PARENT_DIR)) +# Extra component/artifact search roots, `os.pathsep`-separated. The single source of extra roots — the +# CLI's `--search-dir` is folded into this var (see nemo_gym.cli.main), so flag and env var are one channel. +# Read at call time so it reflects the current environment (incl. spawned server subprocesses). +NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME = "NEMO_GYM_EXTRA_ROOTS" -def _resolve_under_cwd_or_install(path) -> Path: - """Resolve a possibly-relative path for *reading* a built-in or user-supplied file. - Absolute paths are returned unchanged. A relative path is tried first under the current working - directory (the user's project), then under the Gym install root (``PARENT_DIR``) where built-in - assets live in both editable and wheel installs. This mirrors ``config_paths`` resolution, so a - repo-relative path like ``resources_servers//data/example.jsonl`` resolves by name from any - cwd. If neither exists the cwd candidate is returned so error messages point at the user's cwd. +def _extra_roots() -> List[Path]: + return [Path(d) for d in os.environ.get(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, "").split(os.pathsep) if d] - Use this for read paths only — never for write targets (e.g. metrics written next to a dataset), - which must stay relative to the user's writable cwd rather than the install root. + +def component_search_roots() -> List[Path]: + """Ordered, de-duplicated roots to look for a Gym component/artifact under: the ``NEMO_GYM_EXTRA_ROOTS`` + roots, then cwd, then ``WORKING_DIR`` and the install root (``PARENT_DIR``, the built-ins). + + Earlier roots win on a name collision, so user components shadow built-ins. De-duplicated by resolved + path, since cwd/``WORKING_DIR``/install root coincide in an editable checkout. The single source of truth + for where Gym looks for components — used by both path resolution (:func:`_resolve_under_cwd_or_install`) and the + ``gym list``/``gym search`` discovery functions. + """ + candidates = [*_extra_roots(), Path.cwd(), WORKING_DIR, PARENT_DIR] + roots: List[Path] = [] + seen: set[Path] = set() + for root in candidates: + resolved = root.resolve() + if resolved not in seen: + seen.add(resolved) + roots.append(root) + return roots + + +def _resolve_under_cwd_or_install( + rel_path: Union[str, Path], *, validator: Optional[Callable[[Path], bool]] = None +) -> Path: + """Resolve a possibly-relative path for *reading* a built-in or user-supplied file against the ordered + :func:`component_search_roots`. + + Absolute paths are returned unchanged. A relative path is returned rooted at the first root where it + exists (earliest-wins: extra roots > cwd > install root), so a repo-relative path like + ``benchmarks//config.yaml`` or ``resources_servers//data/example.jsonl`` resolves by name + from any cwd or plugin root. ``validator`` overrides the default ``Path.exists`` check — pass it when a + candidate is valid only if specific markers are present (e.g. a server dir must ship + ``requirements.txt`` or ``pyproject.toml``). If nothing matches, the path under the highest-priority + root is returned so error messages point at the user's own location. + + Use this for read paths only — never for write targets (e.g. metrics written next to a dataset), which + must stay relative to the user's writable cwd rather than the install root. """ - p = Path(path) + p = Path(rel_path) if p.is_absolute(): return p - cwd_candidate = Path.cwd() / p - if cwd_candidate.exists(): - return cwd_candidate - install_candidate = PARENT_DIR / p - if install_candidate.exists(): - return install_candidate - return cwd_candidate + is_valid = validator if validator is not None else Path.exists + roots = component_search_roots() + for root in roots: + candidate = root / p + if is_valid(candidate): + return candidate + return roots[0] / p + + +def _augment_sys_path() -> None: + """Put the artifact roots on ``sys.path`` so plugin modules (e.g. a benchmark's ``prepare.py``) import. + + Idempotent; reads ``NEMO_GYM_EXTRA_ROOTS`` at call time, so it can be re-run after ``--search-dir`` folds + roots into the env (see nemo_gym.cli.main). With no extra roots this is just ``append(PARENT_DIR)``, as + before. + """ + for root in [*_extra_roots(), PARENT_DIR]: + entry = str(root) + if entry not in sys.path: + sys.path.append(entry) + + +_augment_sys_path() # TODO: Maybe eventually we want an override for OMP_NUM_THREADS ? diff --git a/nemo_gym/cli/env.py b/nemo_gym/cli/env.py index 5e8ba51fa6..e7e4d7524c 100644 --- a/nemo_gym/cli/env.py +++ b/nemo_gym/cli/env.py @@ -36,7 +36,7 @@ from rich.table import Table from tqdm.auto import tqdm -from nemo_gym import PARENT_DIR, ROOT_DIR +from nemo_gym import PARENT_DIR, ROOT_DIR, _resolve_under_cwd_or_install, component_search_roots from nemo_gym.cli.setup_command import run_command, setup_env_command from nemo_gym.cli.utils import ( exit_cleanly_on_config_error, @@ -80,14 +80,14 @@ def _resolve_server_dir(rel_path: Path) -> Path: """Resolve a relative server dir (e.g. ``resources_servers/``) to an absolute path. - Checks the current working directory first (a user's local server), then falls back to the Gym - install root (``PARENT_DIR``) where built-in servers live in both editable and wheel installs. - This lets ``gym env test`` find and run built-in servers from any cwd, not just a repo checkout. + Searches NEMO_GYM_EXTRA_ROOTS, the current working directory (a user's local server), then the Gym + install root (``PARENT_DIR``) where built-in servers live in both editable and wheel installs. A + directory counts as a server only if it ships an install marker for one of our two venv setups. This + lets ``gym env test`` find and run built-in (and plugin) servers from any cwd, not just a repo checkout. """ - cwd_path = Path.cwd() / rel_path - if (cwd_path / "requirements.txt").exists() or (cwd_path / "pyproject.toml").exists(): - return cwd_path - return PARENT_DIR / rel_path + return _resolve_under_cwd_or_install( + rel_path, validator=lambda d: (d / "requirements.txt").exists() or (d / "pyproject.toml").exists() + ) class RunConfig(BaseNeMoGymCLIConfig): @@ -627,14 +627,14 @@ def test_all(): # pragma: no cover global_config_dict = get_global_config_dict() test_all_config = TestAllConfig.model_validate(global_config_dict) - # Discover server modules under both the cwd (a user's project) and the Gym install root - # (built-ins, which live under PARENT_DIR in editable and wheel installs). Entrypoints are kept - # relative; the cwd shadows the install root for same-named modules. This lets `gym env test` - # discover and run built-in servers from any cwd, not only a repo checkout. + # Discover server modules across every component-search root: NEMO_GYM_EXTRA_ROOTS (plugins), the cwd + # (a user's project), and the Gym install root (built-ins, under PARENT_DIR in editable and wheel + # installs). Entrypoints are kept relative; earlier roots shadow later ones for same-named modules. This + # lets `gym env test` discover and run built-in and plugin servers from any cwd, not only a repo checkout. server_type_dirs = ("resources_servers", "responses_api_agents", "responses_api_models") seen_rel_paths: set[str] = set() candidate_dir_paths: List[str] = [] - for root in (Path.cwd(), PARENT_DIR): + for root in component_search_roots(): for server_type_dir in server_type_dirs: for module_path in sorted((root / server_type_dir).glob("*")): if "pycache" in module_path.name or not module_path.is_dir(): diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 18aff51c9a..a205c6b15f 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -23,8 +23,8 @@ from dataclasses import dataclass, field from pathlib import Path +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, _augment_sys_path, component_search_roots from nemo_gym.cli.utils import did_you_mean -from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, component_search_roots logger = logging.getLogger(__name__) @@ -681,6 +681,7 @@ def _extra_roots_from_search_dir(search_dirs: list[str] | None): Setting the env var lets the roots reach every resolver (discovery, the -- selectors, deep config/prompt/rollout resolution) and inherit into spawned server subprocesses. The original value is restored (or the var unset) on exit so main() leaves no global side effect. + (sys.path is augmented with the new roots for plugin ``prepare.py`` imports and left as-is on exit.) """ if not search_dirs: yield @@ -688,6 +689,7 @@ def _extra_roots_from_search_dir(search_dirs: list[str] | None): original = os.environ.get(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME) value = os.pathsep.join(search_dirs) os.environ[NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME] = value + _augment_sys_path() # re-read env so --search-dir roots are importable (e.g. a benchmark prepare.py) logger.debug(f"Set {NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME}={value} from --search-dir") try: yield diff --git a/nemo_gym/discovery.py b/nemo_gym/discovery.py index 6adac0e2aa..e11fc80de7 100644 --- a/nemo_gym/discovery.py +++ b/nemo_gym/discovery.py @@ -19,16 +19,15 @@ they can share it without depending on each other. Reads configs only; never starts servers. """ -import os import re from copy import deepcopy from pathlib import Path -from typing import Callable, Dict, Iterable, List, Optional, Tuple, TypeVar +from typing import Callable, Dict, Iterable, Optional, Tuple, TypeVar from omegaconf import DictConfig, OmegaConf from omegaconf.errors import InterpolationKeyError -from nemo_gym import PARENT_DIR, WORKING_DIR +from nemo_gym import component_search_roots from nemo_gym.global_config import ( POLICY_MODEL_KEY_NAME, GlobalConfigDictParser, @@ -38,36 +37,6 @@ _T = TypeVar("_T") -# Extra component-search roots, `os.pathsep`-separated; same effect as repeating `--search-dir` -# (see :func:`component_search_roots`). -NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME = "NEMO_GYM_EXTRA_ROOTS" - - -def component_search_roots() -> List[Path]: - """Ordered, de-duplicated roots to look for a Gym component under: the roots from the - ``NEMO_GYM_EXTRA_ROOTS`` env var, then cwd, then ``WORKING_DIR`` and the install root (``PARENT_DIR``, - the built-ins). - - ``NEMO_GYM_EXTRA_ROOTS`` is an ``os.pathsep``-separated list of extra roots — the single source of extra - roots. ``--search-dir`` is folded into it up front (see ``nemo_gym.cli.main.main``), so the flag and the - env var are one and the same channel here. - - Earlier roots win on a name collision (see :func:`merge_by_name`), so user components shadow built-ins. - De-duplicated by resolved path, since cwd/``WORKING_DIR``/install root coincide in an editable checkout. - The single source of truth for where Gym looks for components — used by both config resolution - (``_asset_config_path``) and the ``gym list``/``gym search`` discovery functions. - """ - env_roots = [Path(d) for d in os.environ.get(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, "").split(os.pathsep) if d] - candidates: List[Path] = [*env_roots, Path.cwd(), WORKING_DIR, PARENT_DIR] - roots: List[Path] = [] - seen: set[Path] = set() - for root in candidates: - resolved = root.resolve() - if resolved not in seen: - seen.add(resolved) - roots.append(root) - return roots - def merge_by_name(per_root: Iterable[Dict[str, _T]]) -> Dict[str, _T]: """Merge per-root ``name -> entry`` mappings; earlier roots win on a collision (user shadows built-in), diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index d158c9ffe4..206984628d 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -36,7 +36,7 @@ from ray import __version__ as ray_version from wandb import Run -from nemo_gym import CACHE_DIR, PARENT_DIR, RESULTS_DIR, WORKING_DIR +from nemo_gym import CACHE_DIR, RESULTS_DIR, WORKING_DIR, _resolve_under_cwd_or_install, component_search_roots from nemo_gym.config_types import ( AlmostServerError, ConfigError, @@ -253,14 +253,12 @@ def load_extra_config_paths(self, config_paths: List[str]) -> Tuple[List[str], L for config_path in config_paths: original_entry = config_path config_path = Path(config_path) - # Check cwd first for user's local configs, then install location - searched_locations = [config_path] - if not config_path.is_absolute(): - cwd_path = Path.cwd() / config_path - install_path = PARENT_DIR / config_path - # cwd and the install root coincide when run from the repo; list each location once. - searched_locations = [cwd_path] if cwd_path == install_path else [cwd_path, install_path] - config_path = cwd_path if cwd_path.exists() else install_path + # Search NEMO_GYM_EXTRA_ROOTS, cwd, then the install root (see _resolve_under_cwd_or_install). + if config_path.is_absolute(): + searched_locations = [config_path] + else: + searched_locations = [root / config_path for root in component_search_roots()] + config_path = _resolve_under_cwd_or_install(original_entry) try: extra_config = _load_config_yaml(config_path) @@ -269,7 +267,8 @@ def load_extra_config_paths(self, config_paths: List[str]) -> Tuple[List[str], L raise ConfigPathNotFoundError( f"""config_paths entry '{original_entry}' was not found. Looked in: {searched} -Check the path is spelled correctly and is relative to your working directory or the Gym install root.""" +Check the path is spelled correctly and is relative to your working directory, an extra root +(NEMO_GYM_EXTRA_ROOTS / --search-dir), or the Gym install root.""" ) from e for new_config_path in extra_config.get(CONFIG_PATHS_KEY_NAME) or []: if new_config_path not in config_paths: @@ -563,12 +562,11 @@ def parse(self, parse_config: Optional[GlobalConfigDictParserConfig] = None) -> global_config_dict: DictConfig = OmegaConf.merge(initial_global_config_dict, global_config_dict) # Load the env.yaml config. We load it early so that people can use it to conveniently store config paths. - # Check cwd first for user's local env.yaml, then fall back to PARENT_DIR + # Search NEMO_GYM_EXTRA_ROOTS, cwd, then the install root. if parse_config.dotenv_path: dotenv_path = parse_config.dotenv_path else: - cwd_env_yaml = Path.cwd() / "env.yaml" - dotenv_path = cwd_env_yaml if cwd_env_yaml.exists() else PARENT_DIR / "env.yaml" + dotenv_path = _resolve_under_cwd_or_install("env.yaml") dotenv_extra_config = DictConfig({}) if dotenv_path.exists() and not parse_config.skip_load_from_dotenv: diff --git a/nemo_gym/prompt.py b/nemo_gym/prompt.py index 5512922e4c..dc5d48e304 100644 --- a/nemo_gym/prompt.py +++ b/nemo_gym/prompt.py @@ -42,8 +42,8 @@ class PromptConfig(BaseModel): def load_prompt_config(path: str) -> PromptConfig: """Load and validate a YAML prompt config file. - Relative paths are resolved against the current working directory first, then the Gym install - root, consistent with how ``config_paths`` and other Gym paths are resolved. + Relative paths are resolved against the component-search roots (extra roots, cwd, then the Gym install + root), consistent with how ``config_paths`` and other Gym paths are resolved. Returns a ``PromptConfig`` with required ``user`` and optional ``system`` fields. Each value is a string template with ``{placeholder}`` syntax. diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index 999f7ff2f0..956a15b0a0 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -31,7 +31,7 @@ from tqdm.asyncio import tqdm from wandb import Table -from nemo_gym import PARENT_DIR +from nemo_gym import _resolve_under_cwd_or_install from nemo_gym.base_resources_server import AggregateMetrics, AggregateMetricsRequest from nemo_gym.base_responses_api_model import ( clear_model_call_captures_for_rollouts, @@ -309,10 +309,8 @@ def _preprocess_rows_from_config(self, config: RolloutCollectionConfig) -> List[ f"{', '.join(s.name for s in skills_ref.skills)})" ) - _input_path = Path(config.input_jsonl_fpath) - if not _input_path.is_absolute(): - _cwd_path = Path.cwd() / _input_path - _input_path = _cwd_path if _cwd_path.exists() else PARENT_DIR / _input_path + # Search NEMO_GYM_EXTRA_ROOTS, cwd, then the install root. + _input_path = _resolve_under_cwd_or_install(config.input_jsonl_fpath) if not _input_path.exists(): raise ConfigPathNotFoundError( f"Input file not found: '{config.input_jsonl_fpath}' (--input). Check the path is spelled correctly." diff --git a/tests/unit_tests/test_cli.py b/tests/unit_tests/test_cli.py index 9685cc0230..99d417ee72 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -27,7 +27,7 @@ import nemo_gym.cli.env import nemo_gym.global_config -from nemo_gym import PARENT_DIR +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, PARENT_DIR from nemo_gym.cli.env import ( _FORCE_KILL_REAP_TIMEOUT_SEC, _GRACEFUL_SHUTDOWN_TIMEOUT_SEC, @@ -46,7 +46,6 @@ ) from nemo_gym.cli.utils import exit_cleanly_on_config_error from nemo_gym.config_types import ConfigError, NoServerInstancesError, ResourcesServerInstanceConfig -from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.registry import EnvironmentEntry diff --git a/tests/unit_tests/test_cli_agents.py b/tests/unit_tests/test_cli_agents.py index 3d881fb9ae..db4663c9f9 100644 --- a/tests/unit_tests/test_cli_agents.py +++ b/tests/unit_tests/test_cli_agents.py @@ -19,9 +19,9 @@ import pytest from omegaconf import OmegaConf +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.agent_registry import AgentEntry from nemo_gym.cli.agents import list_agents -from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME def _mock_global_config(config: dict = None): diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index 2f0de5e6f0..1d72eab39c 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -22,9 +22,8 @@ import nemo_gym.cli.main as cli_main import nemo_gym.global_config as gc -from nemo_gym import WORKING_DIR +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, WORKING_DIR from nemo_gym.cli.main import main -from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.global_config import NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME @@ -1099,8 +1098,8 @@ def test_builtin_resolves_from_install_root_when_cwd_differs(self, monkeypatch: install_root.mkdir() user_cwd.mkdir() self._make_resources_server(install_root) # built-in only under the install root - monkeypatch.setattr("nemo_gym.discovery.PARENT_DIR", install_root) - monkeypatch.setattr("nemo_gym.discovery.WORKING_DIR", user_cwd) + monkeypatch.setattr("nemo_gym.PARENT_DIR", install_root) + monkeypatch.setattr("nemo_gym.WORKING_DIR", user_cwd) monkeypatch.chdir(user_cwd) resolved = cli_main._asset_config_path("resources-server", "foo") @@ -1112,8 +1111,8 @@ def test_user_cwd_asset_resolves_when_not_builtin(self, monkeypatch: MonkeyPatch install_root.mkdir() user_cwd.mkdir() self._make_resources_server(user_cwd, name="myenv") # exists only in the user's project - monkeypatch.setattr("nemo_gym.discovery.PARENT_DIR", install_root) - monkeypatch.setattr("nemo_gym.discovery.WORKING_DIR", user_cwd) + monkeypatch.setattr("nemo_gym.PARENT_DIR", install_root) + monkeypatch.setattr("nemo_gym.WORKING_DIR", user_cwd) monkeypatch.chdir(user_cwd) resolved = cli_main._asset_config_path("resources-server", "myenv") @@ -1127,8 +1126,8 @@ def test_same_name_in_install_root_and_cwd_is_ambiguous(self, monkeypatch: Monke user_cwd.mkdir() self._make_resources_server(install_root) self._make_resources_server(user_cwd) - monkeypatch.setattr("nemo_gym.discovery.PARENT_DIR", install_root) - monkeypatch.setattr("nemo_gym.discovery.WORKING_DIR", user_cwd) + monkeypatch.setattr("nemo_gym.PARENT_DIR", install_root) + monkeypatch.setattr("nemo_gym.WORKING_DIR", user_cwd) monkeypatch.chdir(user_cwd) with pytest.raises(ValueError, match="ambiguous"): @@ -1140,8 +1139,8 @@ def test_editable_layout_single_root_not_self_ambiguous(self, monkeypatch: Monke repo_root = tmp_path / "Gym" repo_root.mkdir() self._make_resources_server(repo_root) - monkeypatch.setattr("nemo_gym.discovery.PARENT_DIR", repo_root) - monkeypatch.setattr("nemo_gym.discovery.WORKING_DIR", repo_root) + monkeypatch.setattr("nemo_gym.PARENT_DIR", repo_root) + monkeypatch.setattr("nemo_gym.WORKING_DIR", repo_root) monkeypatch.chdir(repo_root) resolved = cli_main._asset_config_path("resources-server", "foo") diff --git a/tests/unit_tests/test_cli_models.py b/tests/unit_tests/test_cli_models.py index 166b8d7844..473ea20797 100644 --- a/tests/unit_tests/test_cli_models.py +++ b/tests/unit_tests/test_cli_models.py @@ -19,8 +19,8 @@ import pytest from omegaconf import OmegaConf +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.cli.models import list_models -from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.model_registry import ModelEntry diff --git a/tests/unit_tests/test_cli_resources_servers.py b/tests/unit_tests/test_cli_resources_servers.py index be47a4884f..aad3f04c2d 100644 --- a/tests/unit_tests/test_cli_resources_servers.py +++ b/tests/unit_tests/test_cli_resources_servers.py @@ -19,8 +19,8 @@ import pytest from omegaconf import OmegaConf +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.cli.resources_servers import list_resources_servers -from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.resources_server_registry import ResourcesServerEntry diff --git a/tests/unit_tests/test_discovery.py b/tests/unit_tests/test_discovery.py index 4f29067b82..05437a9224 100644 --- a/tests/unit_tests/test_discovery.py +++ b/tests/unit_tests/test_discovery.py @@ -17,12 +17,10 @@ from omegaconf import OmegaConf -from nemo_gym import PARENT_DIR +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, PARENT_DIR, component_search_roots from nemo_gym.discovery import ( _UNSET_VALUE_PLACEHOLDER, - NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, _parse_no_environment_tolerating_unset_values, - component_search_roots, merge_by_name, read_config_metadata, ) diff --git a/tests/unit_tests/test_global_config.py b/tests/unit_tests/test_global_config.py index 2d487c57fd..51ecbcdaf0 100644 --- a/tests/unit_tests/test_global_config.py +++ b/tests/unit_tests/test_global_config.py @@ -22,7 +22,7 @@ import nemo_gym.global_config import nemo_gym.server_utils -from nemo_gym import CACHE_DIR, WORKING_DIR +from nemo_gym import CACHE_DIR, NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, WORKING_DIR from nemo_gym.config_types import ( AlmostServerError, ConfigError, @@ -1266,7 +1266,9 @@ def test_load_extra_config_paths_falls_back_to_parent_dir(self, monkeypatch: Mon cwd_dir = tmp_path / "cwd" cwd_dir.mkdir() monkeypatch.chdir(cwd_dir) - monkeypatch.setattr(nemo_gym.global_config, "PARENT_DIR", parent_dir) + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) + monkeypatch.setattr("nemo_gym.PARENT_DIR", parent_dir) + monkeypatch.setattr("nemo_gym.WORKING_DIR", parent_dir) config_paths, extra_configs = parser.load_extra_config_paths(["my_config.yaml"]) assert extra_configs[0]["my_key"] == "from_parent" @@ -1280,7 +1282,9 @@ def test_env_yaml_loaded_from_cwd(self, monkeypatch: MonkeyPatch, tmp_path: Path monkeypatch.chdir(tmp_path) empty_parent = tmp_path / "empty_parent" empty_parent.mkdir() - monkeypatch.setattr(nemo_gym.global_config, "PARENT_DIR", empty_parent) + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) + monkeypatch.setattr("nemo_gym.PARENT_DIR", empty_parent) + monkeypatch.setattr("nemo_gym.WORKING_DIR", empty_parent) parser = GlobalConfigDictParser() global_config_dict = parser.parse(GlobalConfigDictParserConfig(skip_load_from_cli=True)) @@ -1298,7 +1302,9 @@ def test_env_yaml_falls_back_to_parent_dir(self, monkeypatch: MonkeyPatch, tmp_p cwd_dir = tmp_path / "cwd" cwd_dir.mkdir() monkeypatch.chdir(cwd_dir) - monkeypatch.setattr(nemo_gym.global_config, "PARENT_DIR", parent_dir) + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) + monkeypatch.setattr("nemo_gym.PARENT_DIR", parent_dir) + monkeypatch.setattr("nemo_gym.WORKING_DIR", parent_dir) parser = GlobalConfigDictParser() global_config_dict = parser.parse(GlobalConfigDictParserConfig(skip_load_from_cli=True)) @@ -1321,7 +1327,9 @@ def test_load_extra_config_paths_missing_relative_lists_both_locations( cwd.mkdir() parent.mkdir() monkeypatch.chdir(cwd) - monkeypatch.setattr(nemo_gym.global_config, "PARENT_DIR", parent) + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) + monkeypatch.setattr("nemo_gym.PARENT_DIR", parent) + monkeypatch.setattr("nemo_gym.WORKING_DIR", parent) parser = GlobalConfigDictParser() with raises(ConfigPathNotFoundError) as exc_info: @@ -1337,7 +1345,9 @@ def test_load_extra_config_paths_missing_dedups_when_cwd_is_install_root( self, monkeypatch: MonkeyPatch, tmp_path: Path ) -> None: monkeypatch.chdir(tmp_path) - monkeypatch.setattr(nemo_gym.global_config, "PARENT_DIR", tmp_path) + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) + monkeypatch.setattr("nemo_gym.PARENT_DIR", tmp_path) + monkeypatch.setattr("nemo_gym.WORKING_DIR", tmp_path) parser = GlobalConfigDictParser() with raises(ConfigPathNotFoundError) as exc_info: diff --git a/tests/unit_tests/test_prompt.py b/tests/unit_tests/test_prompt.py index 1ebe3fb5a5..d2b72a6a65 100644 --- a/tests/unit_tests/test_prompt.py +++ b/tests/unit_tests/test_prompt.py @@ -18,7 +18,7 @@ import pytest import yaml -from nemo_gym import PARENT_DIR, _resolve_under_cwd_or_install +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, PARENT_DIR, _resolve_under_cwd_or_install from nemo_gym.prompt import ( PromptConfig, apply_prompt_to_row, @@ -77,12 +77,14 @@ def test_absolute_returned_unchanged(self, tmp_path): assert _resolve_under_cwd_or_install(str(tmp_path / "x.yaml")) == tmp_path / "x.yaml" def test_cwd_preferred_over_install_root(self, tmp_path, monkeypatch): + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) monkeypatch.chdir(tmp_path) (tmp_path / "rel.yaml").write_text("{}") assert _resolve_under_cwd_or_install("rel.yaml") == tmp_path / "rel.yaml" def test_falls_back_to_install_root(self, tmp_path, monkeypatch): # cwd lacks the file; a file present under the install root resolves there. + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) monkeypatch.chdir(tmp_path) rel = "test_resolve_install_fallback.yaml" install_file = PARENT_DIR / rel @@ -92,7 +94,8 @@ def test_falls_back_to_install_root(self, tmp_path, monkeypatch): finally: install_file.unlink() - def test_missing_returns_cwd_candidate(self, tmp_path, monkeypatch): + def test_missing_returns_highest_priority_root_candidate(self, tmp_path, monkeypatch): + monkeypatch.delenv(NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME, raising=False) monkeypatch.chdir(tmp_path) assert _resolve_under_cwd_or_install("nope.yaml") == tmp_path / "nope.yaml" diff --git a/tests/unit_tests/test_registry.py b/tests/unit_tests/test_registry.py index 67836eeb63..49f9d77344 100644 --- a/tests/unit_tests/test_registry.py +++ b/tests/unit_tests/test_registry.py @@ -14,7 +14,7 @@ # limitations under the License. from pathlib import Path -from nemo_gym.discovery import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME +from nemo_gym import NEMO_GYM_EXTRA_ROOTS_ENV_VAR_NAME from nemo_gym.registry import _discover_environments_in_dir, discover_environments