fix: lineage tracking bug (aws#1447)

* fix: lineage bug

* fix: lineage

* fix: add validation for tracking ARN input with MLflow input type

* fix: bug

* fix: unit tests

* fix: mock

* fix: args

---------

Co-authored-by: Jacky Lee <[email protected]>

Author: 2 people

src/sagemaker/serve/builder/model_builder.py

@@ -508,8 +508,8 @@ def _model_builder_register_wrapper(self, *args, **kwargs):
             _maintain_lineage_tracking_for_mlflow_model(
                 mlflow_model_path=self.model_metadata[MLFLOW_MODEL_PATH],
                 s3_upload_path=self.s3_upload_path,
-                tracking_server_arn=self.model_metadata.get(MLFLOW_TRACKING_ARN),
                 sagemaker_session=self.sagemaker_session,
+                tracking_server_arn=self.model_metadata.get(MLFLOW_TRACKING_ARN),
             )
         return new_model_package
 
@@ -579,8 +579,8 @@ def _model_builder_deploy_wrapper(
             _maintain_lineage_tracking_for_mlflow_model(
                 mlflow_model_path=self.model_metadata[MLFLOW_MODEL_PATH],
                 s3_upload_path=self.s3_upload_path,
-                tracking_server_arn=self.model_metadata.get(MLFLOW_TRACKING_ARN),
                 sagemaker_session=self.sagemaker_session,
+                tracking_server_arn=self.model_metadata.get(MLFLOW_TRACKING_ARN),
             )
         return predictor
 

src/sagemaker/serve/utils/lineage_utils.py

@@ -54,28 +54,28 @@
 
 def _load_artifact_by_source_uri(
     source_uri: str,
-    artifact_type: str,
     sagemaker_session: Session,
     source_types_to_match: Optional[List[str]] = None,
+    artifact_type: Optional[str] = None,
 ) -> Optional[ArtifactSummary]:
     """Load lineage artifact by source uri
 
     Arguments:
         source_uri (str): The s3 uri used for uploading transfomred model artifacts.
-        artifact_type (str): The type of the lineage artifact.
-        source_types_to_match (Optional[List[str]]): A list of source type values to match against
-            the artifact's source types. If provided, the artifact's source types must match this
-            list.
         sagemaker_session (Session): Session object which manages interactions
             with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
             function creates one using the default AWS configuration chain.
+        source_types_to_match (Optional[List[str]]): A list of source type values to match against
+            the artifact's source types. If provided, the artifact's source types must match this
+            list.
+        artifact_type (Optional[str]): The type of the lineage artifact.
 
     Returns:
         ArtifactSummary: The Artifact Summary for the provided S3 URI.
     """
     artifacts = Artifact.list(source_uri=source_uri, sagemaker_session=sagemaker_session)
     for artifact_summary in artifacts:
-        if artifact_summary.artifact_type == artifact_type:
+        if artifact_type is None or artifact_summary.artifact_type == artifact_type:
             if source_types_to_match:
                 if artifact_summary.source.source_types is not None:
                     artifact_source_types = [
@@ -109,7 +109,9 @@ def _poll_lineage_artifact(
     logger.info("Polling lineage artifact for model data in %s", s3_uri)
     start_time = time.time()
     while time.time() - start_time < LINEAGE_POLLER_MAX_TIMEOUT_SECS:
-        result = _load_artifact_by_source_uri(s3_uri, artifact_type, sagemaker_session)
+        result = _load_artifact_by_source_uri(
+            s3_uri, sagemaker_session, artifact_type=artifact_type
+        )
         if result is not None:
             return result
         time.sleep(LINEAGE_POLLER_INTERVAL_SECS)
@@ -124,12 +126,12 @@ def _get_mlflow_model_path_type(mlflow_model_path: str) -> str:
     Returns:
         str: Description of what the input string is identified as.
     """
-    mlflow_rub_id_pattern = MLFLOW_RUN_ID_REGEX
+    mlflow_run_id_pattern = MLFLOW_RUN_ID_REGEX
     mlflow_registry_id_pattern = MLFLOW_REGISTRY_PATH_REGEX
     sagemaker_arn_pattern = MODEL_PACKAGE_ARN_REGEX
     s3_pattern = S3_PATH_REGEX
 
-    if re.match(mlflow_rub_id_pattern, mlflow_model_path):
+    if re.match(mlflow_run_id_pattern, mlflow_model_path):
         return MLFLOW_RUN_ID
     if re.match(mlflow_registry_id_pattern, mlflow_model_path):
         return MLFLOW_REGISTRY_PATH
@@ -146,12 +148,14 @@ def _get_mlflow_model_path_type(mlflow_model_path: str) -> str:
 def _create_mlflow_model_path_lineage_artifact(
     mlflow_model_path: str,
     sagemaker_session: Session,
+    source_types_to_match: Optional[List[str]] = None,
 ) -> Optional[Artifact]:
     """Creates a lineage artifact for the given MLflow model path.
 
     Args:
         mlflow_model_path (str): The path to the MLflow model.
         sagemaker_session (Session): The SageMaker session object.
+        source_types_to_match (Optional[List[str]]): Artifact source types.
 
     Returns:
         Optional[Artifact]: The created lineage artifact, or None if an error occurred.
@@ -161,8 +165,17 @@ def _create_mlflow_model_path_lineage_artifact(
         model_builder_input_model_data_type=_artifact_name,
     )
     try:
+        source_types = [dict(SourceIdType="Custom", Value="ModelBuilderInputModelData")]
+        if source_types_to_match:
+            source_types += [
+                dict(SourceIdType="Custom", Value=source_type)
+                for source_type in source_types_to_match
+                if source_type != "ModelBuilderInputModelData"
+            ]
+
         return Artifact.create(
             source_uri=mlflow_model_path,
+            source_types=source_types,
             artifact_type=MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
             artifact_name=_artifact_name,
             properties=properties,
@@ -178,37 +191,38 @@ def _create_mlflow_model_path_lineage_artifact(
 
 def _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
     mlflow_model_path: str,
-    tracking_server_arn: str,
     sagemaker_session: Session,
+    tracking_server_arn: Optional[str] = None,
 ) -> Optional[Union[Artifact, ArtifactSummary]]:
     """Retrieves an existing artifact for the given MLflow model path or
 
     creates a new one if it doesn't exist.
 
     Args:
         mlflow_model_path (str): The path to the MLflow model.
-        tracking_server_arn (Optional[str]): The MLflow tracking server ARN.
         sagemaker_session (Session): Session object which manages interactions
             with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
             function creates one using the default AWS configuration chain.
-
+        tracking_server_arn (Optional[str]): The MLflow tracking server ARN.
 
     Returns:
         Optional[Union[Artifact, ArtifactSummary]]: The existing or newly created artifact,
             or None if an error occurred.
     """
-    match = re.match(TRACKING_SERVER_ARN_REGEX, tracking_server_arn)
-    mlflow_tracking_server_name = match.group(4)
-    describe_result = sagemaker_session.sagemaker_client.describe_mlflow_tracking_server(
-        MlflowTrackingServerName=mlflow_tracking_server_name
-    )
-    tracking_server_creation_time = describe_result["CreationTime"].strftime(
-        TRACKING_SERVER_CREATION_TIME_FORMAT
-    )
-    source_types_to_match = [tracking_server_arn, tracking_server_creation_time]
+    source_types_to_match = ["ModelBuilderInputModelData"]
+    input_type = _get_mlflow_model_path_type(mlflow_model_path)
+    if tracking_server_arn and input_type in [MLFLOW_RUN_ID, MLFLOW_REGISTRY_PATH]:
+        match = re.match(TRACKING_SERVER_ARN_REGEX, tracking_server_arn)
+        mlflow_tracking_server_name = match.group(4)
+        describe_result = sagemaker_session.sagemaker_client.describe_mlflow_tracking_server(
+            MlflowTrackingServerName=mlflow_tracking_server_name
+        )
+        tracking_server_creation_time = describe_result["CreationTime"].strftime(
+            TRACKING_SERVER_CREATION_TIME_FORMAT
+        )
+        source_types_to_match += [tracking_server_arn, tracking_server_creation_time]
     _loaded_artifact = _load_artifact_by_source_uri(
         mlflow_model_path,
-        MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
         sagemaker_session,
         source_types_to_match,
     )
@@ -217,6 +231,7 @@ def _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
     return _create_mlflow_model_path_lineage_artifact(
         mlflow_model_path,
         sagemaker_session,
+        source_types_to_match,
     )
 
 
@@ -261,18 +276,18 @@ def _add_association_between_artifacts(
 def _maintain_lineage_tracking_for_mlflow_model(
     mlflow_model_path: str,
     s3_upload_path: str,
-    tracking_server_arn: Optional[str],
     sagemaker_session: Session,
+    tracking_server_arn: Optional[str] = None,
 ) -> None:
     """Maintains lineage tracking for an MLflow model by creating or retrieving artifacts.
 
     Args:
         mlflow_model_path (str): The path to the MLflow model.
         s3_upload_path (str): The S3 path where the transformed model data is uploaded.
-        tracking_server_arn (Optional[str]): The MLflow tracking server ARN.
         sagemaker_session (Session): Session object which manages interactions
             with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
             function creates one using the default AWS configuration chain.
+        tracking_server_arn (Optional[str]): The MLflow tracking server ARN.
     """
     artifact_for_transformed_model_data = _poll_lineage_artifact(
         s3_uri=s3_upload_path,
@@ -283,8 +298,8 @@ def _maintain_lineage_tracking_for_mlflow_model(
         mlflow_model_artifact = (
             _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
                 mlflow_model_path=mlflow_model_path,
-                tracking_server_arn=tracking_server_arn,
                 sagemaker_session=sagemaker_session,
+                tracking_server_arn=tracking_server_arn,
             )
         )
         if mlflow_model_artifact:

tests/unit/sagemaker/serve/utils/test_lineage_utils.py

@@ -57,7 +57,7 @@ def test_load_artifact_by_source_uri(mock_artifact_list):
     mock_artifact_list.return_value = mock_artifacts
 
     result = _load_artifact_by_source_uri(
-        source_uri, LineageSourceEnum.MODEL_DATA.value, sagemaker_session
+        source_uri, sagemaker_session, artifact_type=LineageSourceEnum.MODEL_DATA.value
     )
 
     mock_artifact_list.assert_called_once_with(
@@ -79,7 +79,7 @@ def test_load_artifact_by_source_uri_no_match(mock_artifact_list):
     mock_artifact_list.return_value = mock_artifacts
 
     result = _load_artifact_by_source_uri(
-        source_uri, LineageSourceEnum.MODEL_DATA.value, sagemaker_session
+        source_uri, sagemaker_session, artifact_type=LineageSourceEnum.MODEL_DATA.value
     )
 
     mock_artifact_list.assert_called_once_with(
@@ -106,7 +106,7 @@ def test_poll_lineage_artifact_found(mock_load_artifact):
     assert result == mock_artifact
     mock_load_artifact.assert_has_calls(
         [
-            call(s3_uri, LineageSourceEnum.MODEL_DATA.value, sagemaker_session),
+            call(s3_uri, sagemaker_session, artifact_type=LineageSourceEnum.MODEL_DATA.value),
         ]
     )
 
@@ -166,6 +166,7 @@ def test_create_mlflow_model_path_lineage_artifact_success(
     mock_artifact_create, mock_get_mlflow_path_type
 ):
     mlflow_model_path = "runs:/Ab12Cd34/my-model"
+    mock_source_types = [dict(SourceIdType="Custom", Value="ModelBuilderInputModelData")]
     sagemaker_session = Mock(spec=Session)
     mock_artifact = Mock(spec=Artifact)
     mock_get_mlflow_path_type.return_value = "mlflow_run_id"
@@ -177,6 +178,7 @@ def test_create_mlflow_model_path_lineage_artifact_success(
     mock_get_mlflow_path_type.assert_called_once_with(mlflow_model_path)
     mock_artifact_create.assert_called_once_with(
         source_uri=mlflow_model_path,
+        source_types=mock_source_types,
         artifact_type=MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
         artifact_name="mlflow_run_id",
         properties={"model_builder_input_model_data_type": "mlflow_run_id"},
@@ -233,20 +235,20 @@ def test_retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact_exi
     mock_sagemaker_client.describe_mlflow_tracking_server.return_value = mock_describe_response
     sagemaker_session.sagemaker_client = mock_sagemaker_client
     mock_source_types_to_match = [
+        "ModelBuilderInputModelData",
         mock_tracking_server_arn,
         mock_creation_time.strftime(TRACKING_SERVER_CREATION_TIME_FORMAT),
     ]
     mock_artifact_summary = Mock(spec=ArtifactSummary)
     mock_load_artifact.return_value = mock_artifact_summary
 
     result = _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
-        mlflow_model_path, mock_tracking_server_arn, sagemaker_session
+        mlflow_model_path, sagemaker_session, mock_tracking_server_arn
     )
 
     assert result == mock_artifact_summary
     mock_load_artifact.assert_called_once_with(
         mlflow_model_path,
-        MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
         sagemaker_session,
         mock_source_types_to_match,
     )
@@ -269,6 +271,7 @@ def test_retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact_cre
     mock_sagemaker_client.describe_mlflow_tracking_server.return_value = mock_describe_response
     sagemaker_session.sagemaker_client = mock_sagemaker_client
     mock_source_types_to_match = [
+        "ModelBuilderInputModelData",
         mock_tracking_server_arn,
         mock_creation_time.strftime(TRACKING_SERVER_CREATION_TIME_FORMAT),
     ]
@@ -277,17 +280,18 @@ def test_retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact_cre
     mock_create_artifact.return_value = mock_artifact
 
     result = _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
-        mlflow_model_path, mock_tracking_server_arn, sagemaker_session
+        mlflow_model_path, sagemaker_session, mock_tracking_server_arn
     )
 
     assert result == mock_artifact
     mock_load_artifact.assert_called_once_with(
         mlflow_model_path,
-        MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
         sagemaker_session,
         mock_source_types_to_match,
     )
-    mock_create_artifact.assert_called_once_with(mlflow_model_path, sagemaker_session)
+    mock_create_artifact.assert_called_once_with(
+        mlflow_model_path, sagemaker_session, mock_source_types_to_match
+    )
 
 
 @patch("sagemaker.lineage.association.Association.create")
@@ -364,7 +368,7 @@ def test_maintain_lineage_tracking_for_mlflow_model_success(
     mock_retrieve_create_artifact.return_value = mock_mlflow_model_artifact
 
     _maintain_lineage_tracking_for_mlflow_model(
-        mlflow_model_path, s3_upload_path, mock_tracking_server_arn, sagemaker_session
+        mlflow_model_path, s3_upload_path, sagemaker_session, mock_tracking_server_arn
     )
 
     mock_poll_artifact.assert_called_once_with(
@@ -402,7 +406,7 @@ def test_maintain_lineage_tracking_for_mlflow_model_no_model_data_artifact(
     mock_retrieve_create_artifact.return_value = None
 
     _maintain_lineage_tracking_for_mlflow_model(
-        mlflow_model_path, s3_upload_path, mock_tracking_server_arn, sagemaker_session
+        mlflow_model_path, s3_upload_path, sagemaker_session, mock_tracking_server_arn
     )
 
     mock_poll_artifact.assert_called_once_with(