[Fix] regex for RunId to handle empty artifact path and change mlflow plugin name (aws#1455)

* [Fix] run id regex pattern such that empty artifact path is handled

* Change mlflow plugin name as per legal team requirement

Author: jiapinw

src/sagemaker/serve/builder/model_builder.py

@@ -702,8 +702,10 @@ def _get_artifact_path(self, mlflow_model_path: str) -> str:
                     "but MLflow model path was provided." % MLFLOW_TRACKING_ARN,
                 )
 
-            if not importlib.util.find_spec("awsmlflow"):
-                raise ImportError("Unable to import awsmlflow, check if awsmlflow is installed")
+            if not importlib.util.find_spec("mlflow_sagemaker"):
+                raise ImportError(
+                    "Unable to import mlflow_sagemaker, check if mlflow_sagemaker is installed"
+                )
 
             import mlflow
 
@@ -928,13 +930,15 @@ def validate(self, model_dir: str) -> Type[bool]:
     def set_tracking_arn(self, arn: str):
         """Set tracking server ARN"""
         # TODO: support native MLflow URIs
-        if importlib.util.find_spec("awsmlflow"):
+        if importlib.util.find_spec("mlflow_sagemaker"):
             import mlflow
 
             mlflow.set_tracking_uri(arn)
             self.model_metadata[MLFLOW_TRACKING_ARN] = arn
         else:
-            raise ImportError("Unable to import awsmlflow, check if awsmlflow is installed")
+            raise ImportError(
+                "Unable to import mlflow_sagemaker, check if mlflow_sagemaker is installed"
+            )
 
     def _hf_schema_builder_init(self, model_task: str):
         """Initialize the schema builder for the given HF_TASK

src/sagemaker/serve/model_format/mlflow/constants.py

@@ -22,7 +22,7 @@
 MODEL_PACKAGE_ARN_REGEX = (
     r"^arn:aws:sagemaker:[a-z0-9\-]+:[0-9]{12}:model-package\/(.*?)(?:/(\d+))?$"
 )
-MLFLOW_RUN_ID_REGEX = r"^runs:/[a-zA-Z0-9]+/[/a-zA-Z0-9\-_\.]+$"
+MLFLOW_RUN_ID_REGEX = r"^runs:/[a-zA-Z0-9]+(/[a-zA-Z0-9\-_\.]*)+$"
 MLFLOW_REGISTRY_PATH_REGEX = r"^models:/[a-zA-Z0-9\-_\.]+[@/]?[a-zA-Z0-9\-_\.][/a-zA-Z0-9\-_\.]*$"
 S3_PATH_REGEX = r"^s3:\/\/[a-zA-Z0-9\-_\.]+(?:\/[a-zA-Z0-9\-_\/\.]*)?$"
 MLFLOW_TRACKING_ARN = "MLFLOW_TRACKING_ARN"

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -2438,7 +2438,7 @@ def test_set_tracking_arn_mlflow_not_installed(self):
         tracking_arn = "arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/test"
         self.assertRaisesRegex(
             ImportError,
-            "Unable to import awsmlflow, check if awsmlflow is installed",
+            "Unable to import mlflow_sagemaker, check if mlflow_sagemaker is installed",
             builder.set_tracking_arn,
             tracking_arn,
         )