Merge branch 'main' into config-composer

Resolve PR #1673 conflicts with main (now includes #1804/#1805/#1806). The
gym env compose work coexists with #1805's install-root _asset_config_path and the
agent asset selector in nemo_gym/cli/main.py; both `env compose` and `env validate`
remain registered. Auto-merge verified: no conflict markers, 167 unit tests pass
(only the pre-existing TestDidYouMean argparse failures), and `gym env compose
--benchmark gsm8k --agent simple_agent` + `gym env validate` both work.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

Author: wprazuch

nemo_gym/__init__.py

@@ -37,6 +37,31 @@
 
 sys.path.append(str(PARENT_DIR))
 
+
+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.
+
+    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)
+    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
+
+
 # TODO: Maybe eventually we want an override for OMP_NUM_THREADS ?
 
 # Turn off HF tokenizers paralellism

nemo_gym/cli/main.py

@@ -21,7 +21,7 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 
-from nemo_gym import WORKING_DIR
+from nemo_gym import PARENT_DIR, WORKING_DIR
 
 
 VERSION_TARGET = "nemo_gym.cli.general:version"
@@ -173,22 +173,34 @@ 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 WORKING_DIR (built-ins) first, then any user-registered --search-dir roots.
+    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.
     """
     parent, subdir, default_flavor = _ASSETS[flag]
     server_name, _, config_flavor = value.partition("/")
     config_flavor = config_flavor or default_flavor or server_name
     config_dir = f"{parent}/{server_name}/{subdir}".rstrip("/")
     path = f"{config_dir}/{config_flavor}.yaml"
 
-    # Match in WORKING_DIR (built-ins) and every --search-dir root; dedupe roots that resolve to the same file.
-    roots = [WORKING_DIR, *(Path(d) for d in search_dirs)]
+    # 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)
     matches: list[Path] = []
 
     for root in roots:
         candidate = root / path
         if candidate.exists():
-            matches.append(candidate.resolve())
+            resolved = candidate.resolve()
+            if resolved not in matches:
+                matches.append(resolved)
 
     if len(matches) > 1:
         matches_str = ", ".join(f"`{m}`" for m in matches)

nemo_gym/prompt.py

@@ -27,7 +27,7 @@
 import yaml
 from pydantic import BaseModel, Field
 
-from nemo_gym import PARENT_DIR
+from nemo_gym import _resolve_under_cwd_or_install
 from nemo_gym.config_types import BaseNeMoGymCLIConfig
 
 
@@ -38,26 +38,18 @@ class PromptConfig(BaseModel):
     system: Optional[str] = None
 
 
-def _resolve_path(path: str) -> Path:
-    """Resolve a path relative to the Gym root (PARENT_DIR), consistent with config_paths resolution."""
-    p = Path(path)
-    if not p.is_absolute():
-        p = PARENT_DIR / p
-    return p
-
-
 @lru_cache(maxsize=64)
 def load_prompt_config(path: str) -> PromptConfig:
     """Load and validate a YAML prompt config file.
 
-    Relative paths are resolved against the Gym root directory (``PARENT_DIR``),
-    consistent with how ``config_paths`` and other Gym paths are resolved.
+    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.
 
     Returns a ``PromptConfig`` with required ``user`` and optional ``system`` fields.
     Each value is a string template with ``{placeholder}`` syntax.
     Results are cached so the same file is only parsed once.
     """
-    resolved = _resolve_path(path)
+    resolved = _resolve_under_cwd_or_install(path)
     with open(resolved) as f:
         data = yaml.safe_load(f)
     return PromptConfig.model_validate(data)
@@ -125,7 +117,7 @@ def materialize_prompts(input_jsonl: str, prompt_config: str, output_jsonl: str)
         output_jsonl: Path to write materialized JSONL (with responses_create_params.input).
     """
     prompt_cfg = load_prompt_config(prompt_config)
-    resolved_prompt_path = str(_resolve_path(prompt_config))
+    resolved_prompt_path = str(_resolve_under_cwd_or_install(prompt_config))
     output_path = Path(output_jsonl)
     output_path.parent.mkdir(parents=True, exist_ok=True)
 

nemo_gym/train_data_utils.py

@@ -25,6 +25,7 @@
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 from tqdm.auto import tqdm
 
+from nemo_gym import _resolve_under_cwd_or_install
 from nemo_gym.base_resources_server import BaseRunRequest
 from nemo_gym.config_types import (
     AGENT_REF_KEY,
@@ -424,7 +425,9 @@ def load_datasets(
         local_datasets_not_found: Dict[str, List[DatasetConfig]] = defaultdict(list)
         for c in server_instance_configs:
             for d in c.datasets:
-                jsonl_fpath = Path(d.jsonl_fpath)
+                # Read check: a built-in dataset path resolves under the install root too, so a
+                # bundled example dataset counts as "found" (no download) from an external cwd.
+                jsonl_fpath = _resolve_under_cwd_or_install(d.jsonl_fpath)
                 if jsonl_fpath.exists():
                     local_datasets_found[c.name].append(d)
                 else:
@@ -523,7 +526,7 @@ def _iter_dataset_lines(self, dataset_config: DatasetConfig):
             )
 
         # Don't load everything into memory at once. Throw things away immediately.
-        with open(dataset_config.jsonl_fpath) as f:
+        with open(_resolve_under_cwd_or_install(dataset_config.jsonl_fpath)) as f:
             for line in tqdm(f, desc=f"{dataset_config.jsonl_fpath}"):
                 for _ in range(repeats):
                     yield line
@@ -661,6 +664,9 @@ def validate_samples_and_aggregate_metrics(
                     else:
                         continue
 
+                # Ensure the artifact dir exists: the metrics file is written next to the dataset's
+                # (cwd-relative) jsonl_fpath, which may not exist yet when collating from a fresh cwd.
+                metrics_fpath.parent.mkdir(parents=True, exist_ok=True)
                 with open(metrics_fpath, "w") as f:
                     json.dump(aggregate_metrics_dict, f, indent=4)
 
@@ -706,6 +712,9 @@ def _collate_samples_single_type(
 
                 data_path = Path(d.jsonl_fpath)
                 prepare_path = data_path.with_name(f"{data_path.stem}_prepare.jsonl")
+                # Create the artifact dir if needed (the prepared file is written next to the
+                # cwd-relative jsonl_fpath, which may not exist when collating from a fresh cwd).
+                prepare_path.parent.mkdir(parents=True, exist_ok=True)
                 with open(prepare_path, "w") as target:
                     for line in self._iter_dataset_lines(d):
                         d = json.loads(line)

pyproject.toml

@@ -329,6 +329,12 @@ exclude-dependencies = [
 ]
 
 [tool.setuptools.package-data]
+# The built-in asset trees (resources_servers, responses_api_agents, responses_api_models,
+# benchmarks, environments — all under [tool.setuptools.packages.find]) ship their configs, READMEs,
+# and example data into the wheel via the setuptools-scm VCS file-finder, which includes git-tracked
+# files. (Adding `environments` to packages.find above is what makes that tree ship at all.) Bulk
+# train/validation JSONL is intentionally not committed to git, so it is excluded automatically.
+# Only nemo_gym's own non-package resource scripts need an explicit entry here.
 nemo_gym = ["resources/*.py"]
 
 [tool.setuptools.dynamic]
@@ -445,6 +451,8 @@ include = [
     "responses_api_agents.*",
     "responses_api_models",
     "responses_api_models.*",
+    "environments",
+    "environments.*",
     "nemo_gym",
     "nemo_gym.*",
 ]

tests/unit_tests/test_cli_main.py

@@ -958,6 +958,75 @@ def test_did_you_mean_spans_user_dir(self, monkeypatch: MonkeyPatch, tmp_path, c
         assert "Did you mean `mybench`?" in capsys.readouterr().err
 
 
+class TestInstallRootResolution:
+    """Built-in assets resolve from the Gym install root regardless of cwd (REQ C2/C5).
+
+    In a wheel install PARENT_DIR is site-packages (where the asset trees are installed) while the
+    user runs from an unrelated project dir, so WORKING_DIR/cwd is not the install root. The selector
+    must still find built-ins under PARENT_DIR. Editable-from-repo keeps working because PARENT_DIR,
+    WORKING_DIR, and cwd all coincide and dedupe to a single root.
+    """
+
+    def _make_resources_server(self, root, name: str = "foo") -> None:
+        config_dir = root / "resources_servers" / name / "configs"
+        config_dir.mkdir(parents=True)
+        (config_dir / f"{name}.yaml").write_text("{}\n")
+
+    def test_builtin_resolves_from_install_root_when_cwd_differs(self, monkeypatch: MonkeyPatch, tmp_path) -> None:
+        install_root = tmp_path / "site-packages"
+        user_cwd = tmp_path / "my-project"
+        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.chdir(user_cwd)
+
+        resolved = cli_main._asset_config_path("resources-server", "foo")
+        assert resolved == str(install_root / "resources_servers" / "foo" / "configs" / "foo.yaml")
+
+    def test_user_cwd_asset_resolves_when_not_builtin(self, monkeypatch: MonkeyPatch, tmp_path) -> None:
+        install_root = tmp_path / "site-packages"
+        user_cwd = tmp_path / "my-project"
+        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.chdir(user_cwd)
+
+        resolved = cli_main._asset_config_path("resources-server", "myenv")
+        assert resolved == str(user_cwd / "resources_servers" / "myenv" / "configs" / "myenv.yaml")
+
+    def test_same_name_in_install_root_and_cwd_is_ambiguous(self, monkeypatch: MonkeyPatch, tmp_path) -> None:
+        # A user asset shadowing a built-in of the same name is ambiguous; they must disambiguate with --config.
+        install_root = tmp_path / "site-packages"
+        user_cwd = tmp_path / "my-project"
+        install_root.mkdir()
+        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.chdir(user_cwd)
+
+        with pytest.raises(ValueError, match="ambiguous"):
+            cli_main._asset_config_path("resources-server", "foo")
+
+    def test_editable_layout_single_root_not_self_ambiguous(self, monkeypatch: MonkeyPatch, tmp_path) -> None:
+        # Editable install: PARENT_DIR == WORKING_DIR == cwd. The same file found via all three roots
+        # must dedupe to one match, not raise a spurious ambiguity error.
+        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.chdir(repo_root)
+
+        resolved = cli_main._asset_config_path("resources-server", "foo")
+        assert resolved == str(repo_root / "resources_servers" / "foo" / "configs" / "foo.yaml")
+
+
 class TestListEnvironmentsRouting:
     def test_list_environments_dispatches(self, monkeypatch: MonkeyPatch) -> None:
         target, overrides = _dispatch_for(monkeypatch, ["list", "environments"])

tests/unit_tests/test_prompt.py

@@ -18,10 +18,9 @@
 import pytest
 import yaml
 
-from nemo_gym import PARENT_DIR
+from nemo_gym import PARENT_DIR, _resolve_under_cwd_or_install
 from nemo_gym.prompt import (
     PromptConfig,
-    _resolve_path,
     apply_prompt_to_row,
     fill_prompt,
     load_prompt_config,
@@ -73,21 +72,29 @@ def test_caching(self, tmp_path):
         assert result1 is result2
 
 
-class TestResolvePath:
-    def test_absolute_path(self, tmp_path):
-        p = tmp_path / "prompt.yaml"
-        p.write_text("")
-        assert _resolve_path(str(p)) == p
+class TestResolveUnderCwdOrInstall:
+    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_relative_path_resolves_to_parent_dir(self, tmp_path):
-        # Create a file relative to PARENT_DIR to test resolution
-        test_file = PARENT_DIR / "test_resolve_path_temp.yaml"
-        test_file.write_text("user: '{q}'")
+    def test_cwd_preferred_over_install_root(self, tmp_path, monkeypatch):
+        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.chdir(tmp_path)
+        rel = "test_resolve_install_fallback.yaml"
+        install_file = PARENT_DIR / rel
+        install_file.write_text("{}")
         try:
-            resolved = _resolve_path("test_resolve_path_temp.yaml")
-            assert resolved.exists()
+            assert _resolve_under_cwd_or_install(rel) == install_file
         finally:
-            test_file.unlink()
+            install_file.unlink()
+
+    def test_missing_returns_cwd_candidate(self, tmp_path, monkeypatch):
+        monkeypatch.chdir(tmp_path)
+        assert _resolve_under_cwd_or_install("nope.yaml") == tmp_path / "nope.yaml"
 
 
 class TestFillPrompt:

tests/unit_tests/test_train_data_utils.py

@@ -21,6 +21,7 @@
 
 import nemo_gym.global_config
 import nemo_gym.train_data_utils
+from nemo_gym import _resolve_under_cwd_or_install
 from nemo_gym.config_types import DatasetConfig, ResponsesAPIAgentServerInstanceConfig
 from nemo_gym.global_config import DictConfig, GlobalConfigDictParser
 from nemo_gym.train_data_utils import (
@@ -576,7 +577,9 @@ def custom_open(filename, mode="r"):
                 write_filenames.append(filename)
                 return mock_write_file()
 
-            if filename == "resources_servers/example_multi_step/data/example.jsonl":
+            if Path(filename) == _resolve_under_cwd_or_install(
+                "resources_servers/example_multi_step/data/example.jsonl"
+            ):
                 return original_open(filename, mode)
             elif filename == Path("resources_servers/example_multi_step/data/example_metrics.json"):
                 with original_open(filename, mode) as f:
@@ -1117,6 +1120,26 @@ def custom_open(filename, mode="r"):
             Path("example.jsonl"),
         ]
 
+    def test_collate_creates_missing_parent_dir(self, tmp_path: Path, monkeypatch: MonkeyPatch) -> None:
+        # The prepared-output file is written next to the dataset's jsonl_fpath. When that dir does not
+        # exist in the cwd (e.g. a built-in dataset resolved from the install root while collating from
+        # a fresh cwd), the write must create the parent instead of crashing with FileNotFoundError.
+        missing_dir = tmp_path / "does" / "not" / "exist"
+        assert not missing_dir.exists()
+        cfg = _make_agent_instance_config(
+            "ex", [{"name": "example", "type": "example", "jsonl_fpath": str(missing_dir / "data.jsonl")}]
+        )
+        processor = TrainDataProcessor()
+        # Bypass the dataset read so the source file isn't needed; we're exercising the write path.
+        monkeypatch.setattr(processor, "_iter_dataset_lines", lambda d: iter(['{"foo": "bar"}']))
+
+        paths = processor._collate_samples_single_type("example", [cfg])
+
+        prepare_path = missing_dir / "data_prepare.jsonl"
+        assert paths == [prepare_path]
+        assert prepare_path.exists()  # parent dir auto-created; write did not crash
+        assert json.loads(prepare_path.read_text().strip())["foo"] == "bar"
+
     def test_collate_samples_metrics_conflict_raises_ValueError(self, monkeypatch: MonkeyPatch) -> None:
         write_filenames_to_mock = dict()