diff --git a/nemo_gym/train_data_utils.py b/nemo_gym/train_data_utils.py index 054654ecc5..ae48c45d50 100644 --- a/nemo_gym/train_data_utils.py +++ b/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) diff --git a/tests/unit_tests/test_train_data_utils.py b/tests/unit_tests/test_train_data_utils.py index 3e3bce0a89..cad0a0e5f0 100644 --- a/tests/unit_tests/test_train_data_utils.py +++ b/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()