diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py index 39a8bc9f5c..12c79b380c 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py +++ b/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, + ) diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py index 4cd2e47e08..039251546e 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py +++ b/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() diff --git a/sagemaker-train/src/sagemaker/ai_registry/dataset.py b/sagemaker-train/src/sagemaker/ai_registry/dataset.py index 86d8f907e6..df79a73265 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/dataset.py +++ b/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,13 +272,27 @@ 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) # Parse S3 URL to extract bucket and prefix - if source.startswith("s3://"): + if content_metadata is not None: + # Feature Store datasets are CSVs, not LLM training formats — skip format validation + if source.startswith("s3://"): + parsed = urlparse(source) + bucket_name = parsed.netloc + s3_key = parsed.path.lstrip("/") + s3_prefix = s3_key + method = DataSetMethod.GENERATED + else: + bucket_name = _get_default_bucket() + s3_prefix = _get_default_s3_prefix(name) + method = DataSetMethod.UPLOADED + AIRHub.upload_to_s3(bucket_name, s3_prefix, source) + elif source.startswith("s3://"): parsed = urlparse(source) bucket_name = parsed.netloc s3_key = parsed.path.lstrip("/") @@ -313,6 +331,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() diff --git a/sagemaker-train/src/sagemaker/ai_registry/dataset_utils.py b/sagemaker-train/src/sagemaker/ai_registry/dataset_utils.py index 6ec192881a..3d617a1ba3 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/dataset_utils.py +++ b/sagemaker-train/src/sagemaker/ai_registry/dataset_utils.py @@ -64,6 +64,7 @@ def __init__( conversation_id: Optional[str] = None, conversation_checkpoint_id: Optional[str] = None, dependencies: Optional[List[str]] = None, + content_metadata: Optional[dict] = None, ): self.dataset_type = dataset_type self.dataset_role_arn = dataset_role_arn @@ -74,6 +75,7 @@ def __init__( self.conversation_id = conversation_id self.conversation_checkpoint_id = conversation_checkpoint_id self.dependencies = dependencies or [] + self.content_metadata = content_metadata def to_json(self) -> str: """Convert to JSON string.""" @@ -93,6 +95,8 @@ def to_json(self) -> str: if self.conversation_checkpoint_id: content["ConversationCheckpointId"] = self.conversation_checkpoint_id content["Dependencies"] = self.dependencies + if self.content_metadata: + content["ContentMetadata"] = self.content_metadata return json.dumps(content) diff --git a/sagemaker-train/tests/unit/ai_registry/test_dataset.py b/sagemaker-train/tests/unit/ai_registry/test_dataset.py index f92f131aec..28641567a6 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_dataset.py +++ b/sagemaker-train/tests/unit/ai_registry/test_dataset.py @@ -390,3 +390,107 @@ def test_create_version_failure(self, mock_air_hub): result = dataset.create_version("s3://bucket/new-data") assert result is False + + +class TestDataSetCreateWithContentMetadata: + """Tests for DataSet.create() with content_metadata (Feature Store lineage).""" + + @patch('sagemaker.train.defaults.resolve_and_validate_role', return_value="arn:aws:iam::123456789012:role/SageMakerRole") + @patch('sagemaker.core.helper.session_helper.Session') + @patch('sagemaker.train.common_utils.finetune_utils._get_current_domain_id') + @patch('sagemaker.ai_registry.dataset.AIRHub') + def test_create_skips_validation_when_content_metadata_provided( + self, mock_air_hub, mock_get_domain_id, mock_session, mock_resolve_role + ): + """DataSet.create() skips file/format validation when content_metadata is set.""" + mock_get_domain_id.return_value = None + mock_session_instance = Mock() + mock_session_instance.get_caller_identity_arn.return_value = ( + "arn:aws:iam::123456789012:role/SageMakerRole" + ) + mock_session.return_value = mock_session_instance + + mock_air_hub.get_hub_name.return_value = "test-hub" + mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} + mock_air_hub.describe_hub_content.return_value = { + RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn", + RESPONSE_KEY_HUB_CONTENT_VERSION: "1", + "CreationTime": "2024-01-01", + "LastModifiedTime": "2024-01-01", + } + + with patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_file') as mock_validate_file, \ + patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_format') as mock_validate_format, \ + patch('sagemaker.ai_registry.dataset.DataSet.wait'): + + DataSet.create( + name="fs-test-dataset", + source="s3://bucket/output/result.csv", + content_metadata={ + "SourceFeatureGroups": ["arn:aws:sagemaker:us-west-2:123:feature-group/fg1"], + "ExtractionMethod": "FeatureStoreDatasetBuilder", + "AthenaQueryExecutionId": "abc-123", + }, + sagemaker_session=mock_session_instance, + wait=False, + ) + + # Validation should be skipped + mock_validate_file.assert_not_called() + mock_validate_format.assert_not_called() + + @patch('sagemaker.train.defaults.resolve_and_validate_role', return_value="arn:aws:iam::123456789012:role/SageMakerRole") + @patch('sagemaker.core.helper.session_helper.Session') + @patch('sagemaker.train.common_utils.finetune_utils._get_current_domain_id') + @patch('sagemaker.ai_registry.dataset.AIRHub') + def test_create_passes_content_metadata_to_document( + self, mock_air_hub, mock_get_domain_id, mock_session, mock_resolve_role + ): + """DataSet.create() includes ContentMetadata in the hub content document.""" + mock_get_domain_id.return_value = None + mock_session_instance = Mock() + mock_session_instance.get_caller_identity_arn.return_value = ( + "arn:aws:iam::123456789012:role/SageMakerRole" + ) + mock_session.return_value = mock_session_instance + + mock_air_hub.get_hub_name.return_value = "test-hub" + mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} + mock_air_hub.describe_hub_content.return_value = { + RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn", + RESPONSE_KEY_HUB_CONTENT_VERSION: "1", + "CreationTime": "2024-01-01", + "LastModifiedTime": "2024-01-01", + } + + with patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_file'), \ + patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_format'), \ + patch('sagemaker.ai_registry.dataset.DataSet.wait'): + + metadata = { + "SourceFeatureGroups": [ + "arn:aws:sagemaker:us-west-2:123:feature-group/fg1", + "arn:aws:sagemaker:us-west-2:123:feature-group/fg2", + ], + "ExtractionMethod": "FeatureStoreDatasetBuilder", + "AthenaQueryExecutionId": "query-xyz", + } + + DataSet.create( + name="fs-test-dataset", + source="s3://bucket/output/result.csv", + content_metadata=metadata, + sagemaker_session=mock_session_instance, + wait=False, + ) + + # Verify import_hub_content was called with document containing ContentMetadata + mock_air_hub.import_hub_content.assert_called_once() + call_kwargs = mock_air_hub.import_hub_content.call_args[1] + document = json.loads(call_kwargs["hub_content_document"]) + assert "ContentMetadata" in document + assert document["ContentMetadata"]["SourceFeatureGroups"] == [ + "arn:aws:sagemaker:us-west-2:123:feature-group/fg1", + "arn:aws:sagemaker:us-west-2:123:feature-group/fg2", + ] + assert document["ContentMetadata"]["ExtractionMethod"] == "FeatureStoreDatasetBuilder" diff --git a/sagemaker-train/tests/unit/ai_registry/test_dataset_utils.py b/sagemaker-train/tests/unit/ai_registry/test_dataset_utils.py index ecd1728220..a0aad09d65 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_dataset_utils.py +++ b/sagemaker-train/tests/unit/ai_registry/test_dataset_utils.py @@ -160,3 +160,38 @@ def test_to_json_none_dependencies(self): parsed = json.loads(json_str) assert parsed["Dependencies"] == [] + + def test_to_json_with_content_metadata(self): + """Test JSON serialization includes ContentMetadata when provided.""" + metadata = { + "SourceFeatureGroups": [ + "arn:aws:sagemaker:us-west-2:123456789012:feature-group/my-fg" + ], + "ExtractionMethod": "FeatureStoreDatasetBuilder", + "AthenaQueryExecutionId": "abc-123", + } + doc = DataSetHubContentDocument( + dataset_s3_bucket="test-bucket", + dataset_s3_prefix="datasets/test.csv", + content_metadata=metadata, + ) + json_str = doc.to_json() + parsed = json.loads(json_str) + + assert "ContentMetadata" in parsed + assert parsed["ContentMetadata"]["SourceFeatureGroups"] == [ + "arn:aws:sagemaker:us-west-2:123456789012:feature-group/my-fg" + ] + assert parsed["ContentMetadata"]["ExtractionMethod"] == "FeatureStoreDatasetBuilder" + assert parsed["ContentMetadata"]["AthenaQueryExecutionId"] == "abc-123" + + def test_to_json_without_content_metadata(self): + """Test JSON serialization omits ContentMetadata when None.""" + doc = DataSetHubContentDocument( + dataset_s3_bucket="test-bucket", + dataset_s3_prefix="datasets/test.csv", + ) + json_str = doc.to_json() + parsed = json.loads(json_str) + + assert "ContentMetadata" not in parsed