feat: enable external plugins for all CLI commands

nemo_gym/__init__.py

@@ -12,9 +12,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
+import 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 +38,90 @@
 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(*, sys_path: List[Path] | None = None) -> List[Path]:
+    """Ordered, de-duplicated roots to look for a Gym component/artifact under: the ``NEMO_GYM_EXTRA_ROOTS``
+    roots first, then cwd and the install root (``PARENT_DIR``, the built-ins) last.
+
+    Earlier roots win on a name collision, so user components shadow built-ins. De-duplicated by resolved
+    path, since cwd/install root coincide in an editable checkout. The single source of truth for where Gym
+    looks for components — used by path resolution (:func:`_resolve_under_cwd_or_install`), module imports
+    (:func:`_augment_sys_path`) and the ``gym list``/``gym search`` discovery functions.
+
+    ``sys_path`` weaves existing ``sys.path`` entries in for import precedence: they slot in after the extra
+    roots but before cwd/built-ins. Any entry that is itself cwd or the install root is dropped so those two
+    always stay last — matching the file-lookup order even in a wheel install, where the install root is
+    already on ``sys.path`` as ``site-packages``.
+    """
+    # cwd and the built-ins are always searched last; drop them from the woven sys.path entries so they
+    # can't get pinned mid-path (e.g. site-packages == PARENT_DIR in a wheel install).
+    trailing = [Path.cwd(), PARENT_DIR]
+    trailing_resolved = {root.resolve() for root in trailing}
+    middle = [entry for entry in (sys_path or []) if entry.resolve() not in trailing_resolved]
+    candidates = [*_extra_roots(), *middle, *trailing]
+    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.
+
+    Extra roots go to the front and cwd/``PARENT_DIR`` (the built-ins) to the end, so import precedence
+    matches file resolution under :func:`component_search_roots` (a plugin shadows a same-named Gym module).
+    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).
+    """
+    new_path = [str(root) for root in component_search_roots(sys_path=(Path(path) for path in sys.path))]
+    logging.debug(f"Replacing sys.path {sys.path} with {new_path}")
+    sys.path[:] = new_path
+
+
+_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,18 @@
 # 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
 
+import yaml
 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,43 +34,8 @@
 )
 
 
-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):
@@ -161,6 +126,53 @@ def _load_benchmarks_from_config_paths(config_paths: List[Path]) -> Dict[str, Be
     return benchmarks_dict
 
 
+def _is_benchmark_config(config_path: Path) -> bool:
+    """True if the config declares a `type: benchmark` dataset anywhere in its structure.
+
+    A raw single-file parse (no `config_paths`/interpolation resolution), so it's format-agnostic and can't
+    fail on includes. An unparseable file is kept (returns True) so the resolve step surfaces a diagnostic.
+    """
+
+    def declares(node: object) -> bool:
+        if isinstance(node, dict):
+            return node.get("type") == "benchmark" or any(declares(v) for v in node.values())
+        if isinstance(node, list):
+            return any(declares(v) for v in node)
+        return False
+
+    try:
+        return declares(yaml.safe_load(config_path.read_text(errors="ignore")))
+    except yaml.YAMLError:
+        return True
+
+
+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. :func:`_is_benchmark_config` 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 _is_benchmark_config(p))
+
+
+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() -> 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

@@ -28,12 +28,9 @@
 
 
 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(

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, print_rich_table
 from nemo_gym.config_types import BaseNeMoGymCLIConfig
@@ -71,14 +71,14 @@
 def _resolve_server_dir(rel_path: Path) -> Path:
     """Resolve a relative server dir (e.g. ``resources_servers/<name>``) 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):
@@ -618,14 +618,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():
@@ -921,11 +921,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(
@@ -953,7 +956,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():