feat: extend gym list and gym search commands

nemo_gym/benchmarks.py

@@ -39,7 +39,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
@@ -101,29 +101,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 _is_benchmark_config(config_path: Path) -> bool:
@@ -161,7 +148,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() -> Dict[str, BenchmarkConfig]:

nemo_gym/cli/agents.py

@@ -14,23 +14,24 @@
 # 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,
     GlobalConfigDictParserConfig,
     get_global_config_dict,
 )
 
 
 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(
@@ -41,6 +42,16 @@ def list_agents() -> None:
 
     agents = discover_agents()
 
+    # `gym search agents <query>` reuses this command, narrowing to fuzzy matches on
+    # name + description + 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.description or "", *entry.variants)
+        }
+
     if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False):
         payload = [
             {
@@ -56,10 +67,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")

nemo_gym/cli/env.py

@@ -38,14 +38,15 @@
 
 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, 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,
     JSON_OUTPUT_KEY_NAME,
     NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME,
     NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME,
     NEMO_GYM_RESERVED_TOP_LEVEL_KEYS,
+    QUERY_KEY_NAME,
     GlobalConfigDictParser,
     GlobalConfigDictParserConfig,
     get_global_config_dict,
@@ -919,9 +920,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:
 
@@ -940,6 +940,16 @@ def list_environments() -> None:
 
     environments = discover_environments()
 
+    # `gym search environments <query>` reuses this command, narrowing to fuzzy matches on
+    # name + domain + description.
+    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 "", env.description or "")
+        }
+
     if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False):
         print(
             json.dumps(
@@ -952,10 +962,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")

nemo_gym/cli/eval.py

@@ -14,27 +14,24 @@
 # limitations under the License.
 
 import asyncio
-import difflib
 import importlib
 import json
 from copy import deepcopy
 from multiprocessing import Pool
 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 (
@@ -46,31 +43,11 @@
     get_first_server_config_dict,
     get_global_config_dict,
 )
-from nemo_gym.reward_profile import RewardProfileConfig, RewardProfiler
-from nemo_gym.rollout_collection import (
-    E2ERolloutCollectionConfig,
-    RolloutAggregationConfig,
-    RolloutAggregationHelper,
-    RolloutCollectionConfig,
-    RolloutCollectionHelper,
-    loads_jsonl_line,
-)
-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
+# NOTE: `reward_profile`, `rollout_collection`, and `train_data_utils` are imported lazily inside the run/aggregate/
+# profile commands below: they pull in heavy deps (wandb, mlflow, anthropic) that the fast `list`/`search`
+# commands in this module must not pay for on every invocation.
 
 
 def list_benchmarks() -> None:
@@ -94,11 +71,13 @@ def list_benchmarks() -> None:
     metadata = {name: read_config_metadata(bench.path) for name, bench in benchmarks.items()}
 
     # `gym search <query>` reuses this command, narrowing the listing to fuzzy matches
-    # across the benchmark name and domain.
+    # across the benchmark config name, its dataset name, domain, and description.
     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 "", metadata[name][1] or "")
         }
 
     if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False):
@@ -116,11 +95,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 = (
@@ -306,6 +281,13 @@ def prepare_benchmark() -> None:
 
 @exit_cleanly_on_config_error
 def e2e_rollout_collection():  # pragma: no cover
+    from nemo_gym.rollout_collection import (
+        E2ERolloutCollectionConfig,
+        RolloutCollectionConfig,
+        RolloutCollectionHelper,
+    )
+    from nemo_gym.train_data_utils import TrainDataProcessor
+
     global_config_dict = get_global_config_dict()
 
     # Ensure we have the right config first thing
@@ -384,6 +366,8 @@ def e2e_rollout_collection():  # pragma: no cover
 
 @exit_cleanly_on_config_error
 def collect_rollouts():  # pragma: no cover
+    from nemo_gym.rollout_collection import RolloutCollectionConfig, RolloutCollectionHelper
+
     config = RolloutCollectionConfig.model_validate(get_global_config_dict())
     rch = RolloutCollectionHelper()
 
@@ -392,6 +376,8 @@ def collect_rollouts():  # pragma: no cover
 
 @exit_cleanly_on_config_error
 def aggregate_rollouts():  # pragma: no cover
+    from nemo_gym.rollout_collection import RolloutAggregationConfig, RolloutAggregationHelper
+
     config = RolloutAggregationConfig.model_validate(get_global_config_dict())
     rah = RolloutAggregationHelper()
 
@@ -400,6 +386,9 @@ def aggregate_rollouts():  # pragma: no cover
 
 @exit_cleanly_on_config_error
 def reward_profile():  # pragma: no cover
+    from nemo_gym.reward_profile import RewardProfileConfig, RewardProfiler
+    from nemo_gym.rollout_collection import loads_jsonl_line
+
     config = RewardProfileConfig.model_validate(get_global_config_dict())
 
     if not Path(config.materialized_inputs_jsonl_fpath).exists():