diff --git a/src/sagemaker/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py b/src/sagemaker/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py index 2f87a761fe..fe6bfae478 100644 --- a/src/sagemaker/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py +++ b/src/sagemaker/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py @@ -59,7 +59,9 @@ def create_pipeline_version_context( pipeline_name, last_update_time ), context_type=f"{PIPELINE_VERSION_CONTEXT_TYPE}-{pipeline_name}", - source_uri=pipeline_arn, + # Suffix the version key so this context's source URI is distinct from the + # pipeline context's, which is also the pipeline ARN. + source_uri=f"{pipeline_arn}/{last_update_time}", source_type=last_update_time, properties={ PIPELINE_NAME_KEY: pipeline_name, diff --git a/tests/unit/sagemaker/feature_store/feature_processor/lineage/test_constants.py b/tests/unit/sagemaker/feature_store/feature_processor/lineage/test_constants.py index c3da980257..bd2208014b 100644 --- a/tests/unit/sagemaker/feature_store/feature_processor/lineage/test_constants.py +++ b/tests/unit/sagemaker/feature_store/feature_processor/lineage/test_constants.py @@ -382,7 +382,9 @@ def read_data(self, spark, params) -> DataFrame: context_arn=f"{PIPELINE_NAME}-version-context-arn", context_name=f"sm-fs-fe-{PIPELINE_NAME}-{LAST_UPDATE_TIME}-fep-ver", context_type=f"FeatureEngineeringPipelineVersion-{PIPELINE_NAME}", - source=ContextSource(source_uri=PIPELINE_ARN, source_types=LAST_UPDATE_TIME), + source=ContextSource( + source_uri=f"{PIPELINE_ARN}/{LAST_UPDATE_TIME}", source_types=LAST_UPDATE_TIME + ), properties={"PipelineName": PIPELINE_NAME, "LastUpdateTime": LAST_UPDATE_TIME}, ) diff --git a/tests/unit/sagemaker/feature_store/feature_processor/lineage/test_pipeline_version_lineage_entity_handler.py b/tests/unit/sagemaker/feature_store/feature_processor/lineage/test_pipeline_version_lineage_entity_handler.py index d6c376f50b..f20cd05cbc 100644 --- a/tests/unit/sagemaker/feature_store/feature_processor/lineage/test_pipeline_version_lineage_entity_handler.py +++ b/tests/unit/sagemaker/feature_store/feature_processor/lineage/test_pipeline_version_lineage_entity_handler.py @@ -42,7 +42,7 @@ def test_create_pipeline_version_context(): create_method.assert_called_with( context_name=f"sm-fs-fe-{PIPELINE_NAME}-{LAST_UPDATE_TIME}-fep-ver", context_type=f"FeatureEngineeringPipelineVersion-{PIPELINE_NAME}", - source_uri=PIPELINE_ARN, + source_uri=f"{PIPELINE_ARN}/{LAST_UPDATE_TIME}", source_type=LAST_UPDATE_TIME, properties={ "PipelineName": PIPELINE_NAME,