diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py index bbeb2c1b2d..52179da4bf 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py @@ -60,7 +60,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/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_constants.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_constants.py index 7e5d33a054..0271c55ed5 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_constants.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_constants.py @@ -386,7 +386,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/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_pipeline_version_lineage_entity_handler.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_pipeline_version_lineage_entity_handler.py index 52e65749be..a0a777d810 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_pipeline_version_lineage_entity_handler.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/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,