feat(feature-store): Add lineage registration to DatasetBuilder

sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py

@@ -3,8 +3,9 @@
 """Dataset Builder for FeatureStore."""
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 import datetime
+import logging
 
 import pandas as pd
 
@@ -18,6 +19,8 @@
     run_athena_query,
 )
 
+logger = logging.getLogger(__name__)
+
 _DEFAULT_CATALOG = "AwsDataCatalog"
 _DEFAULT_DATABASE = "sagemaker_featurestore"
 
@@ -258,6 +261,8 @@ class DatasetBuilder:
     _event_time_starting_timestamp: datetime.datetime = field(default=None, init=False)
     _event_time_ending_timestamp: datetime.datetime = field(default=None, init=False)
     _feature_groups_to_be_merged: List[FeatureGroupToBeMerged] = field(default_factory=list, init=False)
+    _register_as_dataset: bool = False
+    _source_feature_groups: List = field(default_factory=list)
 
     @classmethod
     def create(
@@ -269,6 +274,7 @@ def create(
         event_time_identifier_feature_name: str = None,
         included_feature_names: List[str] = None,
         kms_key_id: str = None,
+        register_as_dataset: bool = False,
     ) -> "DatasetBuilder":
         """Create a DatasetBuilder for generating a Dataset.
 
@@ -298,6 +304,7 @@ def create(
             _event_time_identifier_feature_name=event_time_identifier_feature_name,
             _included_feature_names=included_feature_names,
             _kms_key_id=kms_key_id,
+            _register_as_dataset=register_as_dataset,
         )
 
     def with_feature_group(
@@ -336,6 +343,7 @@ def with_feature_group(
                 feature_name_in_target, join_comparator, join_type,
             )
         )
+        self._source_feature_groups.append(feature_group)
         return self
 
     def point_in_time_accurate_join(self) -> "DatasetBuilder":
@@ -779,3 +787,86 @@ def _construct_join_condition(self, fg: FeatureGroupToBeMerged, suffix: str) ->
             )
 
         return join
+
+    def _collect_source_feature_group_arns(self) -> List[str]:
+        """Collect and deduplicate Feature Group ARNs from base and merged FGs."""
+        arns = []
+        # Base FG
+        if isinstance(self._base, FeatureGroup):
+            base_arn = getattr(self._base, "feature_group_arn", None)
+            if base_arn:
+                arns.append(base_arn)
+        # Merged FGs
+        for fg in self._source_feature_groups:
+            fg_arn = getattr(fg, "feature_group_arn", None)
+            if fg_arn and fg_arn not in arns:
+                arns.append(fg_arn)
+        return arns
+
+    def _register_as_hub_content_dataset(
+        self, csv_path: str, query_execution_id: Optional[str] = None
+    ) -> None:
+        """Register the output CSV as a SM Dataset (HubContent) for lineage tracking.
+
+        This is a best-effort operation: if it fails due to missing permissions
+        (AccessDeniedException), a warning is logged and the method returns without
+        raising — the primary workflow (returning the CSV) is not affected.
+
+        Args:
+            csv_path: S3 path of the generated CSV file.
+            query_execution_id: Athena query execution ID (for provenance tracking).
+        """
+        source_fg_arns = self._collect_source_feature_group_arns()
+        if not source_fg_arns:
+            logger.warning(
+                "register_as_dataset=True but no Feature Group ARNs found. "
+                "Skipping dataset registration."
+            )
+            return
+
+        # Generate a dataset name from the base FG name + timestamp
+        base_name = ""
+        if isinstance(self._base, FeatureGroup):
+            base_name = getattr(self._base, "feature_group_name", "dataset")
+        else:
+            base_name = "dataframe-dataset"
+        timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+        dataset_name = f"fs-{base_name}-{timestamp}"
+
+        try:
+            from sagemaker.ai_registry.dataset import DataSet
+
+            DataSet.create(
+                name=dataset_name,
+                source=csv_path,
+                content_metadata={
+                    "SourceFeatureGroups": source_fg_arns,
+                    "ExtractionMethod": "FeatureStoreDatasetBuilder",
+                    "AthenaQueryExecutionId": query_execution_id or "",
+                },
+                description=f"Dataset extracted from Feature Groups: {', '.join(source_fg_arns)}",
+                sagemaker_session=self._sagemaker_session,
+                wait=False,
+            )
+            logger.info(
+                "Registered dataset '%s' as SM Dataset (HubContent) with source FGs: %s",
+                dataset_name,
+                source_fg_arns,
+            )
+        except Exception as e:
+            # Graceful fallback: log warning, don't block the primary workflow
+            error_msg = str(e)
+            if "AccessDenied" in error_msg or "not authorized" in error_msg.lower():
+                logger.warning(
+                    "Unable to register dataset as HubContent due to missing permissions "
+                    "(sagemaker:ImportHubContent). Lineage will not be created for this "
+                    "dataset extraction. To enable lineage, add sagemaker:ImportHubContent "
+                    "permission to your execution role. Error: %s",
+                    error_msg,
+                )
+            else:
+                logger.warning(
+                    "Failed to register dataset as HubContent. Lineage will not be created. "
+                    "Error: %s",
+                    error_msg,
+                )

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py

@@ -345,3 +345,178 @@ def test_to_csv_raises_for_invalid_base(self, mock_session):
 
         with pytest.raises(ValueError, match="must be either"):
             builder.to_csv_file()
+
+
+class TestDatasetBuilderRegisterAsDataset:
+    """Tests for the register_as_dataset lineage feature."""
+
+    @pytest.fixture
+    def mock_session(self):
+        return Mock()
+
+    @pytest.fixture
+    def mock_feature_group(self):
+        fg = MagicMock(spec=FeatureGroup)
+        fg.feature_group_name = "customers-fg"
+        fg.feature_group_arn = "arn:aws:sagemaker:us-west-2:123456789012:feature-group/customers-fg"
+        return fg
+
+    def test_register_as_dataset_default_false(self, mock_session):
+        """Default register_as_dataset is False."""
+        builder = DatasetBuilder.create(
+            base=MagicMock(spec=FeatureGroup),
+            output_path="s3://bucket/output",
+            session=mock_session,
+        )
+        assert builder._register_as_dataset is False
+
+    def test_register_as_dataset_true_sets_flag(self, mock_session):
+        """register_as_dataset=True is stored correctly."""
+        builder = DatasetBuilder.create(
+            base=MagicMock(spec=FeatureGroup),
+            output_path="s3://bucket/output",
+            session=mock_session,
+            register_as_dataset=True,
+        )
+        assert builder._register_as_dataset is True
+
+    def test_collect_source_fg_arns_from_base(self, mock_session, mock_feature_group):
+        """Collects base FG ARN."""
+        builder = DatasetBuilder(
+            _sagemaker_session=mock_session,
+            _base=mock_feature_group,
+            _output_path="s3://bucket/output",
+            _register_as_dataset=True,
+        )
+        arns = builder._collect_source_feature_group_arns()
+        assert arns == ["arn:aws:sagemaker:us-west-2:123456789012:feature-group/customers-fg"]
+
+    def test_collect_source_fg_arns_with_merged_fg(self, mock_session, mock_feature_group):
+        """Collects base + merged FG ARNs."""
+        merged_fg = MagicMock(spec=FeatureGroup)
+        merged_fg.feature_group_arn = "arn:aws:sagemaker:us-west-2:123456789012:feature-group/orders-fg"
+
+        builder = DatasetBuilder(
+            _sagemaker_session=mock_session,
+            _base=mock_feature_group,
+            _output_path="s3://bucket/output",
+            _register_as_dataset=True,
+        )
+        builder._source_feature_groups.append(merged_fg)
+
+        arns = builder._collect_source_feature_group_arns()
+        assert len(arns) == 2
+        assert "arn:aws:sagemaker:us-west-2:123456789012:feature-group/customers-fg" in arns
+        assert "arn:aws:sagemaker:us-west-2:123456789012:feature-group/orders-fg" in arns
+
+    def test_collect_source_fg_arns_deduplicates(self, mock_session, mock_feature_group):
+        """Doesn't duplicate ARNs if same FG used twice."""
+        builder = DatasetBuilder(
+            _sagemaker_session=mock_session,
+            _base=mock_feature_group,
+            _output_path="s3://bucket/output",
+            _register_as_dataset=True,
+        )
+        builder._source_feature_groups.append(mock_feature_group)
+
+        arns = builder._collect_source_feature_group_arns()
+        assert len(arns) == 1
+
+    def test_collect_source_fg_arns_dataframe_base_empty(self, mock_session):
+        """DataFrame base has no FG ARN."""
+        df = pd.DataFrame({"id": [1], "event_time": ["2024-01-01"]})
+        builder = DatasetBuilder(
+            _sagemaker_session=mock_session,
+            _base=df,
+            _output_path="s3://bucket/output",
+            _record_identifier_feature_name="id",
+            _event_time_identifier_feature_name="event_time",
+            _register_as_dataset=True,
+        )
+        arns = builder._collect_source_feature_group_arns()
+        assert arns == []
+
+    def test_register_hub_content_called_on_success(self, mock_session, mock_feature_group):
+        """DataSet.create is called with correct params."""
+        builder = DatasetBuilder(
+            _sagemaker_session=mock_session,
+            _base=mock_feature_group,
+            _output_path="s3://bucket/output",
+            _register_as_dataset=True,
+        )
+
+        with patch(
+            "sagemaker.ai_registry.dataset.DataSet.create",
+        ) as mock_create:
+            builder._register_as_hub_content_dataset(
+                csv_path="s3://bucket/output/result.csv",
+                query_execution_id="abc-123",
+            )
+            mock_create.assert_called_once()
+            call_kwargs = mock_create.call_args[1]
+            assert "customers-fg" in call_kwargs["name"]
+            assert call_kwargs["source"] == "s3://bucket/output/result.csv"
+            assert call_kwargs["content_metadata"]["SourceFeatureGroups"] == [
+                "arn:aws:sagemaker:us-west-2:123456789012:feature-group/customers-fg"
+            ]
+            assert call_kwargs["content_metadata"]["ExtractionMethod"] == "FeatureStoreDatasetBuilder"
+            assert call_kwargs["content_metadata"]["AthenaQueryExecutionId"] == "abc-123"
+            assert call_kwargs["sagemaker_session"] == mock_session
+            assert call_kwargs["wait"] is False
+
+    def test_register_graceful_on_access_denied(self, mock_session, mock_feature_group, caplog):
+        """AccessDeniedException logs warning, doesn't raise."""
+        builder = DatasetBuilder(
+            _sagemaker_session=mock_session,
+            _base=mock_feature_group,
+            _output_path="s3://bucket/output",
+            _register_as_dataset=True,
+        )
+
+        with patch(
+            "sagemaker.ai_registry.dataset.DataSet.create",
+        ) as mock_create:
+            mock_create.side_effect = Exception("AccessDenied: not authorized")
+            # Should NOT raise
+            builder._register_as_hub_content_dataset(
+                csv_path="s3://bucket/output/result.csv",
+            )
+
+    def test_register_graceful_on_generic_error(self, mock_session, mock_feature_group):
+        """Generic errors log warning, don't raise."""
+        builder = DatasetBuilder(
+            _sagemaker_session=mock_session,
+            _base=mock_feature_group,
+            _output_path="s3://bucket/output",
+            _register_as_dataset=True,
+        )
+
+        with patch(
+            "sagemaker.ai_registry.dataset.DataSet.create",
+        ) as mock_create:
+            mock_create.side_effect = Exception("Some service error")
+            # Should NOT raise
+            builder._register_as_hub_content_dataset(
+                csv_path="s3://bucket/output/result.csv",
+            )
+
+    def test_register_skipped_when_no_fg_arns(self, mock_session):
+        """Skips registration when no FG ARNs available (DataFrame base, no merged FGs)."""
+        df = pd.DataFrame({"id": [1], "event_time": ["2024-01-01"]})
+        builder = DatasetBuilder(
+            _sagemaker_session=mock_session,
+            _base=df,
+            _output_path="s3://bucket/output",
+            _record_identifier_feature_name="id",
+            _event_time_identifier_feature_name="event_time",
+            _register_as_dataset=True,
+        )
+
+        with patch(
+            "sagemaker.ai_registry.dataset.DataSet.create",
+        ) as mock_create:
+            builder._register_as_hub_content_dataset(
+                csv_path="s3://bucket/output/result.csv",
+            )
+            # Should NOT call DataSet.create since no FG ARNs
+            mock_create.assert_not_called()

sagemaker-train/src/sagemaker/ai_registry/dataset.py

@@ -232,6 +232,7 @@ def create(
         role: Optional[str] = None,
         domain_id: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
+        content_metadata: Optional[dict] = None,
     ) -> "DataSet":
         """Create a new DataSet Hub AIR entity.
 
@@ -251,6 +252,9 @@ def create(
                 environment; supply it explicitly when creating datasets outside Studio
                 (e.g. from a laptop or CI) so they still appear in the target domain.
             sagemaker_session: Optional SageMaker session. If not provided, uses default session.
+            content_metadata: Optional metadata dict (PascalCase keys) to include in the
+                HubContent document. Used by Feature Store lineage to pass source FG ARNs.
+                When provided, file format validation is skipped.
 
         Returns:
             DataSet: The created dataset instance
@@ -268,8 +272,9 @@ def create(
         if domain_id is None:
             domain_id = _get_current_domain_id(sagemaker_session)
 
-        # Validate dataset file
-        cls._validate_dataset_file(source)
+        # Validate dataset file (skip for Feature Store metadata-only datasets)
+        if content_metadata is None:
+            cls._validate_dataset_file(source)
         sagemaker_session = TrainDefaults.get_sagemaker_session(sagemaker_session=sagemaker_session)
         role = TrainDefaults.get_role(role=role, sagemaker_session=sagemaker_session)
 
@@ -281,18 +286,19 @@ def create(
             s3_prefix = s3_key  # Use full path including filename
             method = DataSetMethod.GENERATED
 
-            # Download and validate format
-            with tempfile.NamedTemporaryFile(
-                delete=False, suffix=os.path.splitext(s3_key)[1]
-            ) as tmp_file:
-                local_path = tmp_file.name
+            # Download and validate format (skip for Feature Store datasets)
+            if content_metadata is None:
+                with tempfile.NamedTemporaryFile(
+                    delete=False, suffix=os.path.splitext(s3_key)[1]
+                ) as tmp_file:
+                    local_path = tmp_file.name
 
-            try:
-                AIRHub.download_from_s3(source, local_path)
-                cls._validate_dataset_format(local_path)
-            finally:
-                if os.path.exists(local_path):
-                    os.remove(local_path)
+                try:
+                    AIRHub.download_from_s3(source, local_path)
+                    cls._validate_dataset_format(local_path)
+                finally:
+                    if os.path.exists(local_path):
+                        os.remove(local_path)
         else:
             # Local file - upload to S3
             bucket_name = _get_default_bucket()
@@ -313,6 +319,7 @@ def create(
             conversation_id=DATASET_DEFAULT_CONVERSATION_ID,  # Required for now, needs cleanup
             conversation_checkpoint_id=DATASET_DEFAULT_CHECKPOINT_ID,
             dependencies=[],
+            content_metadata=content_metadata,
         )
 
         document_str = hub_content_document.to_json()