Skip to content

Commit 08dd4d5

Browse files
pravali96Pravali Uppugunduri
andauthored
fix: use a distinct source URI for the FP pipeline version context (#6541)
Feature Processor created both the pipeline context and the pipeline version context with the pipeline ARN as their source URI. Consumers that key lineage graph nodes on source URI therefore collapsed the two contexts into a single node, and could not tell successive versions apart. Append the version key already passed to create_pipeline_version_context to the version context's source URI, so every version gets a URI distinct from the pipeline's and from other versions. The last update time is what already identifies a version throughout this module: it names the context and is stored as source_type and in properties["LastUpdateTime"]. SageMaker's own PipelineVersionId was not used because DescribePipeline, the only pipeline call Feature Processor makes, does not return it; reading it would add a ListPipelineVersions call to the lineage path and require that permission for all Feature Processor users. Reads are unaffected: load_pipeline_version_context resolves contexts by name, not by source URI. Contexts created before this change keep their old source URI. Co-authored-by: Pravali Uppugunduri <upravali@amazon.com>
1 parent b38c4d0 commit 08dd4d5

3 files changed

Lines changed: 7 additions & 3 deletions

File tree

‎src/sagemaker/feature_store/feature_processor/lineage/_pipeline_version_lineage_entity_handler.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,9 @@ def create_pipeline_version_context(
5959
pipeline_name, last_update_time
6060
),
6161
context_type=f"{PIPELINE_VERSION_CONTEXT_TYPE}-{pipeline_name}",
62-
source_uri=pipeline_arn,
62+
# Suffix the version key so this context's source URI is distinct from the
63+
# pipeline context's, which is also the pipeline ARN.
64+
source_uri=f"{pipeline_arn}/{last_update_time}",
6365
source_type=last_update_time,
6466
properties={
6567
PIPELINE_NAME_KEY: pipeline_name,

‎tests/unit/sagemaker/feature_store/feature_processor/lineage/test_constants.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -382,7 +382,9 @@ def read_data(self, spark, params) -> DataFrame:
382382
context_arn=f"{PIPELINE_NAME}-version-context-arn",
383383
context_name=f"sm-fs-fe-{PIPELINE_NAME}-{LAST_UPDATE_TIME}-fep-ver",
384384
context_type=f"FeatureEngineeringPipelineVersion-{PIPELINE_NAME}",
385-
source=ContextSource(source_uri=PIPELINE_ARN, source_types=LAST_UPDATE_TIME),
385+
source=ContextSource(
386+
source_uri=f"{PIPELINE_ARN}/{LAST_UPDATE_TIME}", source_types=LAST_UPDATE_TIME
387+
),
386388
properties={"PipelineName": PIPELINE_NAME, "LastUpdateTime": LAST_UPDATE_TIME},
387389
)
388390

‎tests/unit/sagemaker/feature_store/feature_processor/lineage/test_pipeline_version_lineage_entity_handler.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def test_create_pipeline_version_context():
4242
create_method.assert_called_with(
4343
context_name=f"sm-fs-fe-{PIPELINE_NAME}-{LAST_UPDATE_TIME}-fep-ver",
4444
context_type=f"FeatureEngineeringPipelineVersion-{PIPELINE_NAME}",
45-
source_uri=PIPELINE_ARN,
45+
source_uri=f"{PIPELINE_ARN}/{LAST_UPDATE_TIME}",
4646
source_type=LAST_UPDATE_TIME,
4747
properties={
4848
"PipelineName": PIPELINE_NAME,

0 commit comments

Comments
 (0)