Pdhote/fix model download (Azure#38977)

* fix resolving path_prefix from asset url * add unit test * add info changelog * path prefix from asset-url
kingernupur · Dec 24, 2024 · 693084d · 693084d
1 parent 8bff427
commit 693084d
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md
@@ -5,6 +5,7 @@
 
 ### Bugs Fixed
   - Fixed disableLocalAuthentication handling while creating amlCompute
+  - Fixed model download issue
 
 ## 1.23.0 (2024-12-05)
 

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_storage_utils.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_storage_utils.py
@@ -191,10 +191,7 @@ def get_ds_name_and_path_prefix(asset_uri: str, registry_name: Optional[str] = N
     if registry_name:
         try:
             split_paths = re.findall(STORAGE_URI_REGEX, asset_uri)
-            if split_paths[0][3] == "":
-                path_prefix = split_paths[0][2]
-            else:
-                path_prefix = split_paths[0][3]
+            path_prefix = split_paths[0][-1]
         except Exception as e:
             msg = "Registry asset URI could not be parsed."
             raise MlException(message=msg, no_personal_data_message=msg) from e

diff --git a/sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_storage_utils.py b/sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_storage_utils.py
@@ -101,6 +101,7 @@ def test_storage_uri_to_prefix(
         reg_uri_1 = "https://ccccccccddddd345.blob.core.windows.net/demoregist-16d33653-20bf-549b-a3c1-17d975359581/ExperimentRun/dcid.5823bbb4-bb28-497c-b9f2-1ff3a0778b10/model"
         reg_uri_2 = "https://ccccccccccc1978ccc.blob.core.windows.net/demoregist-b46fb119-d3f8-5994-a971-a9c730227846/LocalUpload/0c225a0230907e61c00ea33eac35a54d/model.pkl"
         reg_uri_3 = "https://ccccccccddr546ddd.blob.core.windows.net/some-reg-9717e928-33c2-50c2-90f5-f410b12b8727/sklearn_regression_model.pkl"
+        reg_uri_4 = "https://ccccccccddr546ddd.blob.core.windows.net/some-reg-c42e6002-ea61-510a-9414-bc4504b74f9f"
         workspace_uri_1 = "azureml://subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/000000000000000/workspaces/some_test_3/datastores/workspaceblobstore/paths/LocalUpload/26960525964086056a7301dd061fb9be/lightgbm_mlflow_model"
 
         assert get_ds_name_and_path_prefix(reg_uri_1, "registry_name") == (
@@ -112,6 +113,7 @@ def test_storage_uri_to_prefix(
             "LocalUpload/0c225a0230907e61c00ea33eac35a54d/model.pkl",
         )
         assert get_ds_name_and_path_prefix(reg_uri_3, "registry_name") == (None, "sklearn_regression_model.pkl")
+        assert get_ds_name_and_path_prefix(reg_uri_4, "registry_name") == (None, "")
         assert get_ds_name_and_path_prefix(workspace_uri_1) == (
             "workspaceblobstore",
             "LocalUpload/26960525964086056a7301dd061fb9be/lightgbm_mlflow_model",