Add support for Azure DNS zone endpoints (kserve#3819)

* Add support for Azure DNS zone endpoints Signed-off-by: tjandy98 <[email protected]> * Add test cases for Azure Blob and File Share URI pattern matching Signed-off-by: tjandy98 <[email protected]> * flake8 Signed-off-by: tjandy98 <[email protected]> * black Signed-off-by: tjandy98 <[email protected]> --------- Signed-off-by: tjandy98 <[email protected]>
yuzisun · Aug 17, 2024 · 61b5e7e · 61b5e7e
1 parent 2fae6e2
commit 61b5e7e
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 4 deletions.
diff --git a/python/kserve/kserve/storage/storage.py b/python/kserve/kserve/storage/storage.py
@@ -46,8 +46,14 @@
 _S3_PREFIX = "s3://"
 _HDFS_PREFIX = "hdfs://"
 _WEBHDFS_PREFIX = "webhdfs://"
-_AZURE_BLOB_RE = "https://(.+?).blob.core.windows.net/(.+)"
-_AZURE_FILE_RE = "https://(.+?).file.core.windows.net/(.+)"
+_AZURE_BLOB_RE = [
+    "https://(.+?).blob.core.windows.net/(.+)",
+    "https://(.+?).z[0-9]{2}.blob.storage.azure.net/(.+)",
+]
+_AZURE_FILE_RE = [
+    "https://(.+?).file.core.windows.net/(.+)",
+    "https://(.+?).z[0-9]{2}.file.storage.azure.net/(.+)",
+]
 _LOCAL_PREFIX = "file://"
 _URI_RE = "https?://(.+)/(.+)"
 _HTTP_PREFIX = "http(s)://"
@@ -93,9 +99,9 @@ def download(uri: str, out_dir: str = None) -> str:
                 model_dir = Storage._download_s3(uri, out_dir)
             elif uri.startswith(_HDFS_PREFIX) or uri.startswith(_WEBHDFS_PREFIX):
                 model_dir = Storage._download_hdfs(uri, out_dir)
-            elif re.search(_AZURE_BLOB_RE, uri):
+            elif any(re.search(pattern, uri) for pattern in _AZURE_BLOB_RE):
                 model_dir = Storage._download_azure_blob(uri, out_dir)
-            elif re.search(_AZURE_FILE_RE, uri):
+            elif any(re.search(pattern, uri) for pattern in _AZURE_FILE_RE):
                 model_dir = Storage._download_azure_file_share(uri, out_dir)
             elif re.search(_URI_RE, uri):
                 model_dir = Storage._download_from_uri(uri, out_dir)

diff --git a/python/kserve/kserve/storage/test/test_storage.py b/python/kserve/kserve/storage/test/test_storage.py
@@ -303,3 +303,33 @@ def test_unpack_zip_file():
     Storage._unpack_archive_file(tar_file, mimetype, out_dir)
     assert os.path.exists(os.path.join(out_dir, "model.pth"))
     os.remove(os.path.join(out_dir, "model.pth"))
+
+
+@mock.patch(STORAGE_MODULE + ".Storage._download_azure_blob")
+def test_download_azure_blob_called_with_matching_uri(mock_download_azure_blob):
+    azure_blob_uris = [
+        "https://accountname.blob.core.windows.net/container/some/blob/",
+        "https://accountname.z20.blob.storage.azure.net/container/some/blob/",
+    ]
+
+    for uri in azure_blob_uris:
+        Storage.download(uri, out_dir="dest_path")
+
+    expected_calls = [mock.call(uri, "dest_path") for uri in azure_blob_uris]
+    mock_download_azure_blob.assert_has_calls(expected_calls)
+
+
+@mock.patch(STORAGE_MODULE + ".Storage._download_azure_file_share")
+def test_download_azure_file_share_called_with_matching_uri(
+    mock_download_azure_file_share,
+):
+    azure_file_uris = [
+        "https://accountname.file.core.windows.net/container/some/blob",
+        "https://accountname.z20.file.storage.azure.net/container/some/blob",
+    ]
+
+    for uri in azure_file_uris:
+        Storage.download(uri, out_dir="dest_path")
+
+    expected_calls = [mock.call(uri, "dest_path") for uri in azure_file_uris]
+    mock_download_azure_file_share.assert_has_calls(expected_calls)