@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a3e2e250f0

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 17:00 UTC

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

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

Hari committed Mar 15, 2026 at 17:08 UTC a3e2e250f09d7a347cfdccfe2f7b593edd1b7bce

14 files changed +75 -17

docker/Dockerfile

+2 -2

@@ -620,7 +620,7 @@ RUN set -eux; \

620 ARG BITSANDBYTES_VERSION_X86=0.46.1

621 ARG BITSANDBYTES_VERSION_ARM64=0.42.0

622 ARG TIMM_VERSION=">=1.0.17"

623 -ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.3"

623 +ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"

624 RUN --mount=type=cache,target=/root/.cache/uv \

625 if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \

626 BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_ARM64}"; \

@@ -628,7 +628,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \

628 BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \

629 fi; \

630 uv pip install --system accelerate hf_transfer modelscope \

631 - "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs]${RUNAI_MODEL_STREAMER_VERSION}"

631 + "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"

632

633 # ============================================================

634 # VLLM INSTALLATION (depends on build stage)

docker/versions.json

+1 -1

@@ -83,7 +83,7 @@

83 "default": ">=1.0.17"

84 },

85 "RUNAI_MODEL_STREAMER_VERSION": {

86 - "default": ">=0.15.3"

86 + "default": ">=0.15.7"

87 }

88 }

89 }

docs/models/extensions/runai_model_streamer.md

+10

@@ -31,6 +31,16 @@ vllm serve gs://core-llm/Llama-3-8b \

31 --load-format runai_streamer

32 ```

34 +To run model from Azure Blob Storage run:

35 +

36 +```bash

37 +AZURE_STORAGE_ACCOUNT_NAME=<account> \

38 +vllm serve az://<container>/<model-path> \

39 + --load-format runai_streamer

40 +```

41 +

42 +Authentication uses `DefaultAzureCredential`, which supports `az login`, managed identity, environment variables (`AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_CLIENT_SECRET`), and other methods.

43 +

44 To run model from a S3 compatible object store run:

46 ```bash

requirements/nightly_torch_test.txt

+1 -1

@@ -42,7 +42,7 @@ tritonclient>=2.51.0

43 numba == 0.61.2 # Required for N-gram speculative decoding

44 numpy

45 -runai-model-streamer[s3,gcs]==0.15.3

45 +runai-model-streamer[s3,gcs,azure]==0.15.7

46 fastsafetensors>=0.2.2

47 instanttensor>=0.1.5

48 pydantic>=2.12 # 2.11 leads to error on python 3.13

requirements/rocm.txt

+1 -1

@@ -15,7 +15,7 @@ tensorizer==2.10.1

15 packaging>=24.2

16 setuptools>=77.0.3,<80.0.0

17 setuptools-scm>=8

18 -runai-model-streamer[s3,gcs]==0.15.3

18 +runai-model-streamer[s3,gcs,azure]==0.15.7

19 conch-triton-kernels==1.2.1

20 timm>=1.0.17

21 # amd-quark: required for Quark quantization on ROCm

requirements/test.in

+1 -1

@@ -56,7 +56,7 @@ grpcio-reflection==1.78.0

56 arctic-inference == 0.1.1 # Required for suffix decoding test

57 numba == 0.61.2 # Required for N-gram speculative decoding

58 numpy

59 -runai-model-streamer[s3,gcs]==0.15.3

59 +runai-model-streamer[s3,gcs,azure]==0.15.7

60 fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage

61 instanttensor>=0.1.5

62 pydantic>=2.12 # 2.11 leads to error on python 3.13

requirements/test.txt

+38 -5

@@ -64,6 +64,14 @@ audioread==3.0.1

64 # via librosa

65 av==16.1.0

66 # via -r requirements/test.in

67 +azure-core==1.38.2

68 + # via

69 + # azure-identity

70 + # azure-storage-blob

71 +azure-identity==1.25.2

72 + # via runai-model-streamer-azure

73 +azure-storage-blob==12.28.0

74 + # via runai-model-streamer-azure

75 backoff==2.2.1

76 # via

77 # -r requirements/test.in

@@ -103,8 +111,10 @@ certifi==2024.8.30

111 # rasterio

112 # requests

113 # sentry-sdk

106 -cffi==1.17.1

107 - # via soundfile

114 +cffi==2.0.0

115 + # via

116 + # cryptography

117 + # soundfile

118 chardet==5.2.0

119 # via mbstrdecoder

120 charset-normalizer==3.4.0

@@ -148,6 +158,12 @@ coverage==7.10.6

158 # via pytest-cov

159 cramjam==2.9.0

160 # via fastparquet

161 +cryptography==46.0.5

162 + # via

163 + # azure-identity

164 + # azure-storage-blob

165 + # msal

166 + # pyjwt

167 cuda-bindings==12.9.4

168 # via torch

169 cuda-pathfinder==1.3.3

@@ -379,6 +395,8 @@ iniconfig==2.0.0

395 # via pytest

396 instanttensor==0.1.5

397 # via -r requirements/test.in

398 +isodate==0.7.2

399 + # via azure-storage-blob

400 isoduration==20.11.0

401 # via jsonschema

402 isort==5.13.2

@@ -492,6 +510,12 @@ more-itertools==10.5.0

510 # via lm-eval

511 mpmath==1.3.0

512 # via sympy

513 +msal==1.34.0

514 + # via

515 + # azure-identity

516 + # msal-extensions

517 +msal-extensions==1.3.1

518 + # via azure-identity

519 msgpack==1.1.0

520 # via

521 # librosa

@@ -828,6 +852,8 @@ pydantic-extra-types==2.10.5

852 # via mistral-common

853 pygments==2.18.0

854 # via rich

855 +pyjwt==2.11.0

856 + # via msal

857 pyogrio==0.11.0

858 # via geopandas

859 pyparsing==3.2.0

@@ -945,6 +971,7 @@ regex==2024.9.11

971 # transformers

972 requests==2.32.3

973 # via

974 + # azure-core

975 # buildkite-test-collector

976 # datasets

977 # diffusers

@@ -957,6 +984,7 @@ requests==2.32.3

984 # lightly

985 # lm-eval

986 # mistral-common

987 + # msal

988 # mteb

989 # pooch

990 # ray

@@ -993,11 +1021,13 @@ rsa==4.9.1

1021 # via google-auth

1022 rtree==1.4.0

1023 # via torchgeo

996 -runai-model-streamer==0.15.3

1024 +runai-model-streamer==0.15.7

1025 # via -r requirements/test.in

998 -runai-model-streamer-gcs==0.15.3

1026 +runai-model-streamer-azure==0.15.7

1027 + # via runai-model-streamer

1028 +runai-model-streamer-gcs==0.15.7

1029 # via runai-model-streamer

1000 -runai-model-streamer-s3==0.15.3

1030 +runai-model-streamer-s3==0.15.7

1031 # via runai-model-streamer

1032 s3transfer==0.10.3

1033 # via boto3

@@ -1266,6 +1296,9 @@ typing-extensions==4.15.0

1296 # aiosignal

1297 # albumentations

1298 # alembic

1299 + # azure-core

1300 + # azure-identity

1301 + # azure-storage-blob

1302 # chz

1303 # fastapi

1304 # grpcio

setup.py

+1 -1

@@ -970,7 +970,7 @@ setup(

970 "tensorizer": ["tensorizer==2.10.1"],

971 "fastsafetensors": ["fastsafetensors >= 0.2.2"],

972 "instanttensor": ["instanttensor >= 0.1.5"],

973 - "runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],

973 + "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],

974 "audio": [

975 "librosa",

976 "scipy",

tests/model_executor/model_loader/runai_streamer_loader/test_runai_utils.py

@@ -19,6 +19,7 @@ from vllm.transformers_utils.runai_utils import (

19 def test_is_runai_obj_uri():

20 assert is_runai_obj_uri("gs://some-gcs-bucket/path")

21 assert is_runai_obj_uri("s3://some-s3-bucket/path")

22 + assert is_runai_obj_uri("az://some-azure-container/path")

23 assert not is_runai_obj_uri("nfs://some-nfs-path")

tests/transformers_utils/test_utils.py

@@ -11,6 +11,7 @@ from vllm.transformers_utils.gguf_utils import (

11 split_remote_gguf,

12 )

13 from vllm.transformers_utils.utils import (

14 + is_azure,

15 is_cloud_storage,

16 is_gcs,

17 is_s3,

@@ -31,9 +32,17 @@ def test_is_s3():

32 assert not is_s3("nfs://nfs-fqdn.local")

35 +def test_is_azure():

36 + assert is_azure("az://model-container/path")

37 + assert not is_azure("s3://model-path/path-to-model")

38 + assert not is_azure("/unix/local/path")

39 + assert not is_azure("nfs://nfs-fqdn.local")

40 +

41 +

42 def test_is_cloud_storage():

43 assert is_cloud_storage("gs://model-path")

44 assert is_cloud_storage("s3://model-path/path-to-model")

45 + assert is_cloud_storage("az://model-container/path")

46 assert not is_cloud_storage("/unix/local/path")

47 assert not is_cloud_storage("nfs://nfs-fqdn.local")

vllm/config/vllm.py

+3 -2

@@ -1574,8 +1574,9 @@ class VllmConfig:

1574 "runai_streamer_sharded",

1575 ):

1576 raise ValueError(

1577 - f"To load a model from S3, 'load_format' "

1578 - f"must be 'runai_streamer' or 'runai_streamer_sharded', "

1577 + f"To load a model from object storage (S3/GCS/Azure), "

1578 + f"'load_format' must be 'runai_streamer' or "

1579 + f"'runai_streamer_sharded', "

1580 f"but got '{self.load_config.load_format}'. "

1581 f"Model: {self.model_config.model}"

1582 )

vllm/model_executor/model_loader/runai_streamer_loader.py

+1 -1

@@ -21,7 +21,7 @@ from vllm.transformers_utils.runai_utils import is_runai_obj_uri, list_safetenso

21 class RunaiModelStreamerLoader(BaseModelLoader):

22 """

23 Model loader that can load safetensors

24 - files from local FS or S3 bucket.

24 + files from local FS, S3, GCS, or Azure Blob Storage.

25 """

27 def __init__(self, load_config: LoadConfig):

vllm/transformers_utils/runai_utils.py

+1 -1

@@ -13,7 +13,7 @@ from vllm.utils.import_utils import PlaceholderModule

14 logger = init_logger(__name__)

16 -SUPPORTED_SCHEMES = ["s3://", "gs://"]

16 +SUPPORTED_SCHEMES = ["s3://", "gs://", "az://"]

18 try:

19 from runai_model_streamer import list_safetensors as runai_list_safetensors

vllm/transformers_utils/utils.py

+5 -1

@@ -23,8 +23,12 @@ def is_gcs(model_or_path: str) -> bool:

23 return model_or_path.lower().startswith("gs://")

26 +def is_azure(model_or_path: str) -> bool:

27 + return model_or_path.lower().startswith("az://")

28 +

29 +

30 def is_cloud_storage(model_or_path: str) -> bool:

27 - return is_s3(model_or_path) or is_gcs(model_or_path)

31 + return is_s3(model_or_path) or is_gcs(model_or_path) or is_azure(model_or_path)

34 def without_trust_remote_code(kwargs: dict[str, Any]) -> dict[str, Any]: