Refactor model download to generate the path prefix for new registry model URI types (Azure#27838)

* fix and refactor model download getting prefix

Author: banibrata-de

sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_constants.py

@@ -40,3 +40,4 @@
 BLOB_STORAGE_CLIENT_NAME = "BlobStorageClient"
 GEN2_STORAGE_CLIENT_NAME = "Gen2StorageClient"
 DEFAULT_CONNECTION_TIMEOUT = 14400
+STORAGE_URI_REGEX = r"(https:\/\/([a-zA-Z0-9@:%_\\\-+~#?&=]+)[a-zA-Z0-9@:%._\\\-+~#?&=]+\.?)\/([a-zA-Z0-9@:%._\\\-+~#?&=]+)\/(.*)"

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_storage_utils.py

@@ -5,6 +5,7 @@
 import logging
 import re
 from typing import Tuple, Union
+from azure.ai.ml._artifacts._constants import STORAGE_URI_REGEX
 
 from azure.ai.ml._artifacts._blob_storage_helper import BlobStorageClient
 from azure.ai.ml._artifacts._fileshare_storage_helper import FileStorageClient
@@ -182,7 +183,19 @@ def get_artifact_path_from_storage_url(blob_url: str, container_name: dict) -> s
     return blob_url
 
 
-def get_ds_name_and_path_prefix(asset_uri: str) -> Tuple[str, str]:
-    ds_name = asset_uri.split("paths")[0].split("/")[-2]
-    path_prefix = asset_uri.split("paths")[1][1:]
+def get_ds_name_and_path_prefix(asset_uri: str, registry_name: str = None) -> Tuple[str, str]:
+    if registry_name:
+        try:
+            split_paths = re.findall(STORAGE_URI_REGEX, asset_uri)
+            path_prefix = split_paths[0][3]
+        except Exception:
+            raise Exception("Registry asset URI could not be parsed." )
+        ds_name = None
+    else:
+        try:
+            ds_name = asset_uri.split("paths")[0].split("/")[-2]
+            path_prefix = asset_uri.split("paths")[1][1:]
+        except Exception:
+            raise Exception("Workspace asset URI could not be parsed.")
+
     return ds_name, path_prefix

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py

@@ -305,6 +305,7 @@ def download(self, name: str, version: str, download_path: Union[PathLike, str]
         """
 
         model_uri = self.get(name=name, version=version).path
+        ds_name, path_prefix = get_ds_name_and_path_prefix(model_uri, self._registry_name)
         if self._registry_name:
             sas_uri = get_storage_details_for_registry_assets(
                 service_client=self._service_client,
@@ -316,10 +317,8 @@ def download(self, name: str, version: str, download_path: Union[PathLike, str]
                 uri=model_uri,
             )
             storage_client = get_storage_client(credential=None, storage_account=None, account_url=sas_uri)
-            path_prefix = model_uri.split("/")[-1]
 
         else:
-            ds_name, path_prefix = get_ds_name_and_path_prefix(model_uri)
             ds = self._datastore_operation.get(ds_name, include_secrets=True)
             acc_name = ds.account_name
 

sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_storage_utils.py

@@ -0,0 +1,45 @@
+
+import pytest
+
+from azure.ai.ml._utils._storage_utils import get_ds_name_and_path_prefix
+
+
+@pytest.mark.unittest
+class TestStorageUtils:
+    def test_storage_uri_to_prefix(
+        self,
+    ) -> None:
+        # These are the asset storage patterns supported for download
+        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'
+        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") == (None,'ExperimentRun/dcid.5823bbb4-bb28-497c-b9f2-1ff3a0778b10/model')
+        assert get_ds_name_and_path_prefix(reg_uri_2, "registry_name") == (None, '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(workspace_uri_1) == ('workspaceblobstore','LocalUpload/26960525964086056a7301dd061fb9be/lightgbm_mlflow_model')
+
+        
+    def test_storage_uri_to_prefix_malformed(
+        self,
+    ) -> None:
+        reg_uri_bad = 'https://ccccccccddd4512d.blob.core.windows.net/5823bbb4-bb28-497c-b9f2-1ff3a0778b10'
+        workspace_uri_bad = 'azureml://subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/000000000000000/workspaces/some_test_3/datastores/workspaceblobstore/path/LocalUpload/26960525964086056a7301dd061fb9be/lightgbm_mlflow_model'
+
+        with pytest.raises(Exception) as e:
+            get_ds_name_and_path_prefix(reg_uri_bad, "registry_name")
+        assert (
+            'Registry asset URI could not be parsed.'
+            in str(e.value)
+        )
+
+        with pytest.raises(Exception) as e:
+            get_ds_name_and_path_prefix(workspace_uri_bad)
+        assert (
+            'Workspace asset URI could not be parsed.'
+            in str(e.value)
+        )
+
+
+