[Feature] Add Azure Blob Storage support for RunAI Model Streamer (#34614)

Signed-off-by: hasethuraman <hsethuraman@microsoft.com>

vllm/config/vllm.py

@@ -1574,8 +1574,9 @@ class VllmConfig:
                 "runai_streamer_sharded",
             ):
                 raise ValueError(
-                    f"To load a model from S3, 'load_format' "
-                    f"must be 'runai_streamer' or 'runai_streamer_sharded', "
+                    f"To load a model from object storage (S3/GCS/Azure), "
+                    f"'load_format' must be 'runai_streamer' or "
+                    f"'runai_streamer_sharded', "
                     f"but got '{self.load_config.load_format}'. "
                     f"Model: {self.model_config.model}"
                 )

vllm/model_executor/model_loader/runai_streamer_loader.py

@@ -21,7 +21,7 @@ from vllm.transformers_utils.runai_utils import is_runai_obj_uri, list_safetenso
 class RunaiModelStreamerLoader(BaseModelLoader):
     """
     Model loader that can load safetensors
-    files from local FS or S3 bucket.
+    files from local FS, S3, GCS, or Azure Blob Storage.
     """
 
     def __init__(self, load_config: LoadConfig):

vllm/transformers_utils/runai_utils.py

@@ -13,7 +13,7 @@ from vllm.utils.import_utils import PlaceholderModule
 
 logger = init_logger(__name__)
 
-SUPPORTED_SCHEMES = ["s3://", "gs://"]
+SUPPORTED_SCHEMES = ["s3://", "gs://", "az://"]
 
 try:
     from runai_model_streamer import list_safetensors as runai_list_safetensors

vllm/transformers_utils/utils.py

@@ -23,8 +23,12 @@ def is_gcs(model_or_path: str) -> bool:
     return model_or_path.lower().startswith("gs://")
 
 
+def is_azure(model_or_path: str) -> bool:
+    return model_or_path.lower().startswith("az://")
+
+
 def is_cloud_storage(model_or_path: str) -> bool:
-    return is_s3(model_or_path) or is_gcs(model_or_path)
+    return is_s3(model_or_path) or is_gcs(model_or_path) or is_azure(model_or_path)
 
 
 def without_trust_remote_code(kwargs: dict[str, Any]) -> dict[str, Any]: