Propagate DataBlend revisions through pretrain SFT planning

Author: intern_nem_dev_2

src/nemotron/data_prep/core/work_items.py

@@ -50,6 +50,9 @@ class DatasetWorkItem:
     sample: str | int | None
     sample_seed: int
 
+    # Source revision propagated from DataBlend.Dataset for deterministic HF discovery.
+    revision: str | None = None
+
     # Resolved tokenizer config (for plan creation)
     tokenizer_config: dict = field(default_factory=dict)
 
@@ -95,6 +98,9 @@ class SftDatasetWorkItem:
     sample: str | int | None
     sample_seed: int
 
+    # Source revision propagated from DataBlend.Dataset for deterministic HF discovery.
+    revision: str | None = None
+
     # Resolved tokenizer config (for plan creation)
     tokenizer_config: dict = field(default_factory=dict)
 

src/nemotron/data_prep/recipes/pretrain.py

@@ -47,13 +47,12 @@
 import hashlib
 import json
 import time
+from collections.abc import Callable, Mapping
 from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Mapping
+from typing import TYPE_CHECKING, Any
 
 import cosmos_xenna.pipelines.v1 as pipelines_v1
-
-from collections.abc import Callable
 from fsspec import AbstractFileSystem
 
 from nemotron.data_prep.config import (
@@ -64,10 +63,11 @@
     ObservabilityConfig,
     TokenizerConfig,
 )
-from nemotron.data_prep.observability import pipeline_wandb_hook
-from nemotron.data_prep.utils.filesystem import ensure_dir, get_filesystem, write_json
 from nemotron.data_prep.core.finalize import scan_dataset_receipts
 from nemotron.data_prep.core.planning import PlanRequest, resolve_tokenizer, verify_binidx_output
+from nemotron.data_prep.core.work_items import DatasetWorkItem, ShardWorkItem
+from nemotron.data_prep.observability import pipeline_wandb_hook
+from nemotron.data_prep.recipes.execution_mode import ExecutionModeRequest, resolve_execution_mode
 from nemotron.data_prep.stages import (
     BinIdxTokenizationStage,
     BinIdxTokenizationStageConfig,
@@ -77,9 +77,8 @@
     PlanStage,
     PlanStageConfig,
 )
+from nemotron.data_prep.utils.filesystem import ensure_dir, get_filesystem, write_json
 from nemotron.data_prep.utils.hf_env import detect_hf_env_vars
-from nemotron.data_prep.core.work_items import DatasetWorkItem, ShardWorkItem
-from nemotron.data_prep.recipes.execution_mode import ExecutionModeRequest, resolve_execution_mode
 
 if TYPE_CHECKING:
     from nemotron.data_prep.blend import DataBlend
@@ -122,6 +121,7 @@ def to_plan_request(self, item: DatasetWorkItem) -> PlanRequest:
                 weight=item.weight,
                 split=item.split,
                 subset=item.subset,
+                revision=item.revision,
                 text_field=item.text_field,
             ),
             num_shards=item.num_shards,
@@ -183,7 +183,7 @@ def _normalize_tokenizer(tokenizer: TokenizerConfig | Mapping[str, Any] | str) -
 
 
 def setup_pretrain_run(
-    blend: "DataBlend",
+    blend: DataBlend,
     output_dir: str | Path,
     tokenizer: TokenizerConfig | Mapping[str, Any] | str,
     *,
@@ -234,6 +234,7 @@ def setup_pretrain_run(
                 "weight": d.weight,
                 "split": d.split,
                 "subset": d.subset,
+                "revision": d.revision,
                 "text_field": getattr(d, "text_field", None) or text_field_default,
             }
             for d in blend.datasets
@@ -270,6 +271,7 @@ def setup_pretrain_run(
                 weight=d.weight,
                 split=d.split,
                 subset=d.subset,
+                revision=d.revision,
                 text_field=getattr(d, "text_field", None) or text_field_default,
                 run_hash=run_hash,
                 run_dir=run_dir,
@@ -297,7 +299,7 @@ def setup_pretrain_run(
 
 def finalize_pretrain_run(
     context: PretrainRunContext,
-    blend: "DataBlend",
+    blend: DataBlend,
     output_dir: str | Path,
 ) -> FormatResult:
     """
@@ -360,7 +362,7 @@ def finalize_pretrain_run(
 
 
 def run_pretrain_pipeline(
-    blend: "DataBlend",
+    blend: DataBlend,
     output_dir: str | Path,
     tokenizer: TokenizerConfig | Mapping[str, Any] | str,
     *,

src/nemotron/data_prep/recipes/sft.py

@@ -46,13 +46,12 @@
 import hashlib
 import json
 import time
+from collections.abc import Callable, Mapping
 from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Mapping
+from typing import TYPE_CHECKING, Any
 
 import cosmos_xenna.pipelines.v1 as pipelines_v1
-
-from collections.abc import Callable
 from fsspec import AbstractFileSystem
 
 from nemotron.data_prep.config import (
@@ -63,10 +62,11 @@
     ObservabilityConfig,
     TokenizerConfig,
 )
-from nemotron.data_prep.observability import pipeline_wandb_hook
-from nemotron.data_prep.utils.filesystem import ensure_dir, get_filesystem, write_json
 from nemotron.data_prep.core.finalize import scan_dataset_receipts
 from nemotron.data_prep.core.planning import PlanRequest, resolve_tokenizer, verify_parquet_output
+from nemotron.data_prep.core.work_items import SftDatasetWorkItem, SftShardWorkItem
+from nemotron.data_prep.observability import pipeline_wandb_hook
+from nemotron.data_prep.recipes.execution_mode import ExecutionModeRequest, resolve_execution_mode
 from nemotron.data_prep.stages import (
     DownloadStage,
     DownloadStageConfig,
@@ -76,9 +76,8 @@
     PlanStage,
     SftPlanStageConfig,
 )
+from nemotron.data_prep.utils.filesystem import ensure_dir, get_filesystem, write_json
 from nemotron.data_prep.utils.hf_env import detect_hf_env_vars
-from nemotron.data_prep.core.work_items import SftDatasetWorkItem, SftShardWorkItem
-from nemotron.data_prep.recipes.execution_mode import ExecutionModeRequest, resolve_execution_mode
 
 if TYPE_CHECKING:
     from nemotron.data_prep.blend import DataBlend
@@ -121,6 +120,7 @@ def to_plan_request(self, item: SftDatasetWorkItem) -> PlanRequest:
                 weight=item.weight,
                 split=item.split,
                 subset=item.subset,
+                revision=item.revision,
                 text_field=item.messages_field,
             ),
             num_shards=item.num_shards,
@@ -190,7 +190,7 @@ def _normalize_tokenizer(tokenizer: TokenizerConfig | Mapping[str, Any] | str) -
 
 
 def setup_sft_run(
-    blend: "DataBlend",
+    blend: DataBlend,
     output_dir: str | Path,
     tokenizer: TokenizerConfig | Mapping[str, Any] | str,
     *,
@@ -257,6 +257,7 @@ def setup_sft_run(
                 "weight": d.weight,
                 "split": d.split,
                 "subset": d.subset,
+                "revision": d.revision,
                 "messages_field": getattr(d, "messages_field", None) or messages_field_default,
                 "tools_field": getattr(d, "tools_field", None) or tools_field_default,
             }
@@ -304,6 +305,7 @@ def setup_sft_run(
                 weight=d.weight,
                 split=d.split,
                 subset=d.subset,
+                revision=d.revision,
                 run_hash=run_hash,
                 run_dir=run_dir,
                 config_hash=config_hash,
@@ -340,7 +342,7 @@ def setup_sft_run(
 
 def finalize_sft_run(
     context: SftRunContext,
-    blend: "DataBlend",
+    blend: DataBlend,
     output_dir: str | Path,
 ) -> FormatResult:
     """
@@ -403,7 +405,7 @@ def finalize_sft_run(
 
 
 def run_sft_pipeline(
-    blend: "DataBlend",
+    blend: DataBlend,
     output_dir: str | Path,
     tokenizer: TokenizerConfig | Mapping[str, Any] | str,
     *,

src/nemotron/kit/artifacts/pretrain_blends.py

@@ -104,15 +104,15 @@ def get_input_uris(self) -> list[str]:
     @classmethod
     def from_result(
         cls,
-        format_result: "FormatResult",
-        blend: "DataBlend",
+        format_result: FormatResult,
+        blend: DataBlend,
         tokenizer_model: str,
         blend_json_path: str | Path,
         *,
         text_field_default: str = "text",
         elapsed_sec: float = 0.0,
         name: str | None = None,
-    ) -> "PretrainBlendsArtifact":
+    ) -> PretrainBlendsArtifact:
         """Create artifact from pipeline format result.
 
         This is a convenience constructor that builds the source_datasets
@@ -137,6 +137,7 @@ def from_result(
                 weight=d.weight,
                 split=d.split,
                 subset=d.subset,
+                revision=d.revision,
                 text_field=d.text_field or text_field_default,
             )
             for d in blend.datasets

src/nemotron/kit/artifacts/sft_data.py

@@ -174,6 +174,7 @@ def from_result(
                 weight=d.weight,
                 split=d.split,
                 subset=d.subset,
+                revision=d.revision,
                 text_field=d.text_field or messages_field_default,
             )
             for d in blend.datasets