feat: extend discoverability through CLI

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/<env>/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/<name>/config.yaml`` or ``resources_servers/<env>/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 ?

nemo_gym/agent_registry.py

@@ -44,9 +44,11 @@
 from omegaconf import OmegaConf
 
 from nemo_gym import PARENT_DIR
+from nemo_gym.discovery import discover_components
 
 
-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,12 @@ def discover_agents(agents_dir: Path = AGENTS_DIR) -> Dict[str, AgentEntry]:
         )
 
     return agents
+
+
+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 (``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)

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, discover_components
 from nemo_gym.global_config import (
     POLICY_MODEL_KEY_NAME,
     GlobalConfigDictParser,
@@ -34,47 +33,12 @@
 )
 
 
-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)
+BENCHMARKS_SUBDIR = "benchmarks"
+BENCHMARKS_DIR = PARENT_DIR / BENCHMARKS_SUBDIR
 
 
 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
@@ -136,11 +100,35 @@ 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``.
+
+    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]:
+    """Sorted config paths under one dir that declare a benchmark, discovered by content.
+
+    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 []
+    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_in_dir(benchmarks_dir: Path) -> Dict[str, BenchmarkConfig]:
+    """Map benchmark name -> :class:`BenchmarkConfig` for every benchmark config under one 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)
@@ -156,11 +144,20 @@ def _load_benchmarks_from_config_paths(config_paths: List[Path]) -> Dict[str, Be
         if not maybe_bc:
             continue
 
-        benchmarks_dict[maybe_bc.name] = maybe_bc
+        benchmarks_dict[_benchmark_config_name(config_path.relative_to(benchmarks_dir))] = maybe_bc
 
     return benchmarks_dict
 
 
+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
+    (``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)
+
+
 # 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

nemo_gym/cli/agents.py

@@ -14,26 +14,53 @@
 # 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 (
+    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,
     GlobalConfigDictParserConfig,
     get_global_config_dict,
 )
 
 
-def list_agents() -> None:
-    """CLI command: list discovered agent harnesses and how each composes (Pattern A vs B).
+def _inspect_agent(name: str, agents: dict, global_config_dict) -> None:
+    """Render the ``gym list agents <name>`` 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,
+    )
 
-    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.
+
+def list_agents() -> None:
+    """List discovered agent harnesses and how each composes (Pattern A vs. self-contained B), or inspect one
+    by name (``gym list agents <name>``). 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(
@@ -44,6 +71,16 @@ def list_agents() -> None:
 
     agents = discover_agents()
 
+    name = global_config_dict.get(COMPONENT_NAME_KEY_NAME)
+    if name:
+        _inspect_agent(name, agents, global_config_dict)
+        return
+
+    # `gym search agents <query>` 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 = [
             {
@@ -59,10 +96,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")