feat(config): unify dataset source via discriminated source: block

Replace the parallel gitlab_identifier / huggingface_identifier fields on
DatasetConfig with a single self-describing source: block (type selects the
backend). Legacy fields keep working: a legacy identifier is mirrored into
source (with a DeprecationWarning), and a source: is back-filled into the
matching legacy field so existing consumers that read the *_identifier fields
are unaffected. Specifying both is rejected.

Addresses FEP-1025 (reduce configuration friction, epic #1205).

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

Author: wprazuch

nemo_gym/config_types.py

@@ -12,10 +12,11 @@
 # 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 warnings
 from argparse import ArgumentParser
 from enum import Enum
 from pathlib import Path
-from typing import Any, ClassVar, Dict, List, Literal, Optional, Set, Tuple, Union
+from typing import Annotated, Any, ClassVar, Dict, List, Literal, Optional, Set, Tuple, Union
 
 import rich
 from omegaconf import DictConfig, OmegaConf
@@ -358,12 +359,37 @@ def check_output_path(self) -> "DownloadJsonlDatasetHuggingFaceConfig":
 DatasetType = Union[Literal["train"], Literal["validation"], Literal["example"]]
 
 
+class GitlabDatasetSource(BaseModel):
+    """Unified ``source:`` for a dataset fetched from the GitLab model registry."""
+
+    type: Literal["gitlab"]
+    dataset_name: str
+    version: str
+    artifact_fpath: str
+
+
+class HuggingFaceDatasetSource(BaseModel):
+    """Unified ``source:`` for a dataset fetched from the HuggingFace Hub."""
+
+    type: Literal["huggingface"]
+    repo_id: str
+    artifact_fpath: Optional[str] = None
+
+
+# One discriminated `source:` block replaces the parallel gitlab_identifier / huggingface_identifier
+# fields; `type` selects the backend so it's unambiguous which fields apply.
+DatasetSource = Annotated[Union[GitlabDatasetSource, HuggingFaceDatasetSource], Field(discriminator="type")]
+
+
 class DatasetConfig(BaseModel):
     name: str
     type: DatasetType
     jsonl_fpath: str
 
     num_repeats: int = Field(default=1, ge=1)
+    # Unified, self-describing dataset source. Prefer this over the legacy *_identifier fields below.
+    source: Optional[DatasetSource] = None
+    # Deprecated: kept working (and back-filled from/into `source`) for backward compatibility.
     gitlab_identifier: Optional[JsonlDatasetGitlabIdentifer] = None
     huggingface_identifier: Optional[JsonlDatasetHuggingFaceIdentifer] = None
     license: Optional[
@@ -386,6 +412,54 @@ def check_train_validation_sets(self) -> "DatasetConfig":
 
         return self
 
+    @model_validator(mode="after")
+    def normalize_dataset_source(self) -> "DatasetConfig":
+        """Reconcile the unified `source:` with the legacy `*_identifier` fields.
+
+        Exactly one source may be specified. A legacy identifier is accepted (with a deprecation
+        warning) and mirrored into `source`; conversely a `source:` is mirrored back into the
+        matching legacy field so existing consumers that read `gitlab_identifier`/
+        `huggingface_identifier` keep working.
+        """
+        specified = [
+            name
+            for name, value in (
+                ("source", self.source),
+                ("gitlab_identifier", self.gitlab_identifier),
+                ("huggingface_identifier", self.huggingface_identifier),
+            )
+            if value is not None
+        ]
+        if len(specified) > 1:
+            raise ValueError(
+                f"Specify a dataset source once for '{self.name}': set only one of {specified}. "
+                "Prefer the unified `source:` block."
+            )
+        if not specified:
+            return self
+
+        if self.source is None:
+            # Legacy identifier was used: mirror it into `source` and nudge toward the new field.
+            if self.gitlab_identifier is not None:
+                self.source = GitlabDatasetSource(type="gitlab", **self.gitlab_identifier.model_dump())
+            else:
+                self.source = HuggingFaceDatasetSource(type="huggingface", **self.huggingface_identifier.model_dump())
+            warnings.warn(
+                f"`{specified[0]}` is deprecated for dataset '{self.name}'; use the unified "
+                f"`source:` block (type: {self.source.type}).",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        else:
+            # `source:` was used: back-fill the matching legacy field for existing consumers.
+            fields = self.source.model_dump(exclude={"type"})
+            if isinstance(self.source, GitlabDatasetSource):
+                self.gitlab_identifier = JsonlDatasetGitlabIdentifer(**fields)
+            else:
+                self.huggingface_identifier = JsonlDatasetHuggingFaceIdentifer(**fields)
+
+        return self
+
 
 class BenchmarkDatasetConfig(BaseModel):
     name: str

tests/unit_tests/test_dataset_source.py

@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+from pydantic import ValidationError
+from pytest import raises, warns
+
+from nemo_gym.config_types import (
+    DatasetConfig,
+    GitlabDatasetSource,
+    HuggingFaceDatasetSource,
+)
+
+
+def _dataset(**extra) -> dict:
+    return {"name": "ds", "type": "example", "jsonl_fpath": "data.jsonl", **extra}
+
+
+class TestDatasetSource:
+    def test_source_gitlab_backfills_legacy_identifier(self) -> None:
+        cfg = DatasetConfig.model_validate(
+            _dataset(
+                source={
+                    "type": "gitlab",
+                    "dataset_name": "my_dataset",
+                    "version": "0.0.1",
+                    "artifact_fpath": "train.jsonl",
+                }
+            )
+        )
+
+        assert isinstance(cfg.source, GitlabDatasetSource)
+        # Existing consumers read the legacy field; it must be back-filled from `source`.
+        assert cfg.gitlab_identifier is not None
+        assert cfg.gitlab_identifier.dataset_name == "my_dataset"
+        assert cfg.gitlab_identifier.version == "0.0.1"
+        assert cfg.gitlab_identifier.artifact_fpath == "train.jsonl"
+        assert cfg.huggingface_identifier is None
+
+    def test_source_huggingface_backfills_legacy_identifier(self) -> None:
+        cfg = DatasetConfig.model_validate(
+            _dataset(source={"type": "huggingface", "repo_id": "org/dataset", "artifact_fpath": "train.jsonl"})
+        )
+
+        assert isinstance(cfg.source, HuggingFaceDatasetSource)
+        assert cfg.huggingface_identifier is not None
+        assert cfg.huggingface_identifier.repo_id == "org/dataset"
+        assert cfg.huggingface_identifier.artifact_fpath == "train.jsonl"
+        assert cfg.gitlab_identifier is None
+
+    def test_legacy_gitlab_identifier_mirrors_into_source_with_warning(self) -> None:
+        with warns(DeprecationWarning, match="gitlab_identifier"):
+            cfg = DatasetConfig.model_validate(
+                _dataset(
+                    gitlab_identifier={
+                        "dataset_name": "my_dataset",
+                        "version": "0.0.1",
+                        "artifact_fpath": "train.jsonl",
+                    }
+                )
+            )
+
+        assert isinstance(cfg.source, GitlabDatasetSource)
+        assert cfg.source.dataset_name == "my_dataset"
+        assert cfg.source.version == "0.0.1"
+        assert cfg.source.artifact_fpath == "train.jsonl"
+        # Legacy field stays populated so nothing that already reads it breaks.
+        assert cfg.gitlab_identifier is not None
+
+    def test_legacy_huggingface_identifier_mirrors_into_source_with_warning(self) -> None:
+        with warns(DeprecationWarning, match="huggingface_identifier"):
+            cfg = DatasetConfig.model_validate(_dataset(huggingface_identifier={"repo_id": "org/dataset"}))
+
+        assert isinstance(cfg.source, HuggingFaceDatasetSource)
+        assert cfg.source.repo_id == "org/dataset"
+        assert cfg.source.artifact_fpath is None
+        assert cfg.huggingface_identifier is not None
+
+    def test_specifying_source_and_legacy_identifier_is_rejected(self) -> None:
+        with raises(ValidationError, match="set only one"):
+            DatasetConfig.model_validate(
+                _dataset(
+                    source={
+                        "type": "gitlab",
+                        "dataset_name": "my_dataset",
+                        "version": "0.0.1",
+                        "artifact_fpath": "train.jsonl",
+                    },
+                    gitlab_identifier={
+                        "dataset_name": "my_dataset",
+                        "version": "0.0.1",
+                        "artifact_fpath": "train.jsonl",
+                    },
+                )
+            )
+
+    def test_no_source_is_allowed(self) -> None:
+        cfg = DatasetConfig.model_validate(_dataset())
+
+        assert cfg.source is None
+        assert cfg.gitlab_identifier is None
+        assert cfg.huggingface_identifier is None
+
+    def test_source_discriminator_selects_backend(self) -> None:
+        with raises(ValidationError):
+            # Missing repo_id for the huggingface branch.
+            DatasetConfig.model_validate(_dataset(source={"type": "huggingface", "artifact_fpath": "train.jsonl"}))

tests/unit_tests/test_train_data_utils.py

@@ -125,6 +125,7 @@ def test_load_and_validate_server_instance_configs_sanity(self, monkeypatch: Mon
                                 "type": "example",
                                 "jsonl_fpath": "resources_servers/example_multi_step/data/example.jsonl",
                                 "num_repeats": 1,
+                                "source": None,
                                 "gitlab_identifier": None,
                                 "huggingface_identifier": None,
                                 "license": None,