fix: gym dataset collate creates the artifact dir when missing (run-from-wheel)

nemo_gym/train_data_utils.py

@@ -664,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)
 
@@ -709,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)

tests/unit_tests/test_train_data_utils.py

@@ -1120,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()