From f74fd75c62c9b9063cbf5f4879b5899630fd8c60 Mon Sep 17 00:00:00 2001 From: ruiliann666 <141953824+ruiliann666@users.noreply.github.com> Date: Mon, 13 Jan 2025 16:32:04 -0800 Subject: [PATCH 1/2] Raise error early for MLFlow in local container mode --- src/sagemaker/serve/builder/model_builder.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py index a7a518105c..5584f43bc6 100644 --- a/src/sagemaker/serve/builder/model_builder.py +++ b/src/sagemaker/serve/builder/model_builder.py @@ -883,6 +883,16 @@ def build( # pylint: disable=R0911 if mode: self.mode = mode + + if self.mode == Mode.LOCAL_CONTAINER and ( + self.model_metadata and ( + "MLFLOW_MODEL_PATH" in self.model_metadata or + "MLFLOW_TRACKING_ARN" in self.model_metadata + ) + ): + raise ValueError("MLflow model support is not available in Local Container mode. " + "Please use SageMaker Endpoint mode for MLflow models.") + if role_arn: self.role_arn = role_arn From 804123c2655c45f1bc0c1532e5ed272a1b94dc1f Mon Sep 17 00:00:00 2001 From: Ruilian Gao Date: Tue, 14 Jan 2025 19:31:35 +0000 Subject: [PATCH 2/2] Improve error message --- src/sagemaker/serve/builder/model_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py index 5584f43bc6..f94136b403 100644 --- a/src/sagemaker/serve/builder/model_builder.py +++ b/src/sagemaker/serve/builder/model_builder.py @@ -890,8 +890,8 @@ def build( # pylint: disable=R0911 "MLFLOW_TRACKING_ARN" in self.model_metadata ) ): - raise ValueError("MLflow model support is not available in Local Container mode. " - "Please use SageMaker Endpoint mode for MLflow models.") + raise ValueError("MLflow model support is not available for Mode.LOCAL_CONTAINER. " + "Please use Mode.SAGEMAKER_ENDPOINT for MLflow models.") if role_arn: self.role_arn = role_arn