Skip to content

Commit 6abd5b6

Browse files
authored
Fix: MLFlow path resolution (#5503)
1 parent 10b683a commit 6abd5b6

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

sagemaker-serve/src/sagemaker/serve/model_builder_utils.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1391,12 +1391,14 @@ def _get_artifact_path(self, mlflow_model_path: str) -> str:
13911391
if "@" in mlflow_model_path:
13921392
_, model_name_and_alias, artifact_uri = mlflow_model_path.split("/", 2)
13931393
model_name, model_alias = model_name_and_alias.split("@")
1394-
model_metadata = mlflow_client.get_model_version_by_alias(model_name, model_alias)
1394+
model_version_info = mlflow_client.get_model_version_by_alias(model_name, model_alias)
1395+
source = mlflow_client.get_model_version_download_uri(
1396+
model_name, model_version_info.version
1397+
)
13951398
else:
13961399
_, model_name, model_version, artifact_uri = mlflow_model_path.split("/", 3)
1397-
model_metadata = mlflow_client.get_model_version(model_name, model_version)
1400+
source = mlflow_client.get_model_version_download_uri(model_name, model_version)
13981401

1399-
source = model_metadata.source
14001402
if not source.endswith("/"):
14011403
source += "/"
14021404
return source + artifact_uri

0 commit comments

Comments
 (0)