feat(config): unify dataset source via discriminated source: block (FEP-1025)

nemo_gym/cli/env.py

@@ -745,7 +745,7 @@ def init_resources_server():  # pragma: no cover
         jsonl_fpath: resources_servers/{server_type_name}/data/train.jsonl   # local data file for this split
         num_repeats: 1                          # times to repeat each example (e.g. for pass@k / mean@k)
         license: Apache 2.0                     # required for train/validation; must be an allowed license string
-        # to fetch this split from a registry instead, add gitlab_identifier: or huggingface_identifier:
+        # to fetch this split from a registry instead, add a source: block (type: gitlab | huggingface)
       - name: validation
         type: validation
         jsonl_fpath: resources_servers/{server_type_name}/data/validation.jsonl

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
@@ -366,12 +367,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[
@@ -394,6 +420,62 @@ 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.
+
+        The unified `source:` block is mutually exclusive with the legacy identifiers. The two
+        legacy identifiers may still be set together (a gitlab-primary / huggingface-fallback pair
+        selected at download time by `config.data_source`) for backward compatibility. A legacy
+        identifier emits a deprecation warning and, when a single backend is given, is 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.
+        """
+        legacy_specified = [
+            name
+            for name, value in (
+                ("gitlab_identifier", self.gitlab_identifier),
+                ("huggingface_identifier", self.huggingface_identifier),
+            )
+            if value is not None
+        ]
+        if self.source is not None and legacy_specified:
+            raise ValueError(
+                f"Specify a dataset source once for '{self.name}': set only one of "
+                f"['source', {', '.join(repr(name) for name in legacy_specified)}]. "
+                "Prefer the unified `source:` block."
+            )
+
+        if self.source is not None:
+            # `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
+
+        if not legacy_specified:
+            return self
+
+        warnings.warn(
+            f"{' and '.join(f'`{name}`' for name in legacy_specified)} "
+            f"{'is' if len(legacy_specified) == 1 else 'are'} deprecated for dataset "
+            f"'{self.name}'; prefer the unified `source:` block.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        # Mirror a single legacy identifier into `source`. When both are set (primary + fallback),
+        # the single discriminated `source:` can't represent both, so leave it unset and keep the
+        # legacy fields as the source of truth.
+        if len(legacy_specified) == 1:
+            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())
+
+        return self
+
 
 class BenchmarkDatasetConfig(BaseModel):
     name: str

resources_servers/example_multi_step/configs/example_multi_step.yaml

@@ -19,18 +19,28 @@ example_multi_step_simple_agent:
       - name: train
         type: train
         jsonl_fpath: resources_servers/example_multi_step/data/train.jsonl
-        gitlab_identifier:
+        # Unified dataset source. `type` selects the backend; the remaining
+        # fields are backend-specific. This replaces the deprecated gitlab_identifier:/
+        # huggingface_identifier: blocks (which still work but are no longer recommended).
+        source:
+          type: gitlab
           dataset_name: example_multi_step
           version: 0.0.1
           artifact_fpath: train.jsonl
         license: Apache 2.0
       - name: validation
         type: validation
         jsonl_fpath: resources_servers/example_multi_step/data/validation.jsonl
-        gitlab_identifier:
+        source:
+          type: gitlab
           dataset_name: example_multi_step
           version: 0.0.1
           artifact_fpath: validation.jsonl
+        # A HuggingFace-hosted split would instead look like:
+        #   source:
+        #     type: huggingface
+        #     repo_id: nvidia/example_multi_step
+        #     artifact_fpath: validation.jsonl   # optional
         license: Apache 2.0
       - name: example
         type: example

resources_servers/mcqa/configs/mcqa.yaml

@@ -31,7 +31,10 @@ mcqa_simple_agent:
       - name: validation
         type: validation
         jsonl_fpath: resources_servers/mcqa/data/validation.jsonl
-        huggingface_identifier:
+        # Unified dataset source. `type` selects the backend; here the split
+        # is pulled from the HuggingFace Hub. `artifact_fpath` is optional for huggingface sources.
+        source:
+          type: huggingface
           repo_id: nvidia/Nemotron-RL-knowledge-mcqa
         license: Apache 2.0
       - name: example

tests/unit_tests/test_cli.py

@@ -152,6 +152,12 @@ def test_init_resources_server_includes_domain(self) -> None:
                 # This should not raise an assertion error about missing domain
                 instance_config = ResourcesServerInstanceConfig.model_validate(full_config_dict)
                 assert instance_config is not None
+
+                # The generated config points users at the unified `source:` identifier, not the
+                # deprecated gitlab_identifier/huggingface_identifier.
+                assert "source:" in config_text
+                assert "gitlab_identifier" not in config_text
+                assert "huggingface_identifier" not in config_text
         finally:
             # Clean up the test server directory
             if server_path.exists():

tests/unit_tests/test_dataset_source.py

@@ -0,0 +1,136 @@
+# 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_both_legacy_identifiers_together_is_allowed(self) -> None:
+        # A gitlab-primary / huggingface-fallback pair (backend chosen at download time) must stay
+        # valid; the single discriminated `source:` can't represent both, so it is left unset.
+        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",
+                    },
+                    huggingface_identifier={"repo_id": "org/dataset"},
+                )
+            )
+
+        assert cfg.source is None
+        assert cfg.gitlab_identifier is not None
+        assert cfg.huggingface_identifier is not None
+
+    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,