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 12, 01:00 UTC
Validating Runai Model Streamer Integration with S3 Object Storage (#29320)
Signed-off-by: Noa Neria <noa@run.ai>
Noa Neria committed
Dec 4, 2025 at 12:04 UTC
6366c098d7c76120b6a55a6829a2649c727a2862
13 files changed
+100
-11
docker/Dockerfile
+1
-1
@@ -580,7 +580,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
580
else \
581
BITSANDBYTES_VERSION="0.46.1"; \
582
fi; \
583
- uv pip install --system accelerate hf_transfer modelscope "bitsandbytes>=${BITSANDBYTES_VERSION}" 'timm>=1.0.17' 'runai-model-streamer[s3,gcs]>=0.15.0'
583
+ uv pip install --system accelerate hf_transfer modelscope "bitsandbytes>=${BITSANDBYTES_VERSION}" 'timm>=1.0.17' 'runai-model-streamer[s3,gcs]>=0.15.3'
584
585
ENV VLLM_USAGE_SOURCE production-docker-image
586
requirements/nightly_torch_test.txt
+1
-1
@@ -42,6 +42,6 @@ tritonclient==2.51.0
42
43
numba == 0.61.2 # Required for N-gram speculative decoding
44
numpy
45
-runai-model-streamer[s3,gcs]==0.15.0
45
+runai-model-streamer[s3,gcs]==0.15.3
46
fastsafetensors>=0.1.10
47
pydantic>=2.12 # 2.11 leads to error on python 3.13
requirements/rocm.txt
+1
-1
@@ -12,7 +12,7 @@ tensorizer==2.10.1
12
packaging>=24.2
13
setuptools>=77.0.3,<80.0.0
14
setuptools-scm>=8
15
-runai-model-streamer[s3,gcs]==0.15.0
15
+runai-model-streamer[s3,gcs]==0.15.3
16
conch-triton-kernels==1.2.1
17
timm>=1.0.17
18
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@d6f998a03432b2452f8de2bb5cefb5af9795d459
requirements/test.in
+1
-1
@@ -51,7 +51,7 @@ tritonclient==2.51.0
51
arctic-inference == 0.1.1 # Required for suffix decoding test
52
numba == 0.61.2 # Required for N-gram speculative decoding
53
numpy
54
-runai-model-streamer[s3,gcs]==0.15.0
54
+runai-model-streamer[s3,gcs]==0.15.3
55
fastsafetensors>=0.1.10
56
pydantic>=2.12 # 2.11 leads to error on python 3.13
57
decord==0.6.0
requirements/test.txt
+3
-3
@@ -965,11 +965,11 @@ rsa==4.9.1
965
# via google-auth
966
rtree==1.4.0
967
# via torchgeo
968
-runai-model-streamer==0.15.0
968
+runai-model-streamer==0.15.3
969
# via -r requirements/test.in
970
-runai-model-streamer-gcs==0.15.0
970
+runai-model-streamer-gcs==0.15.3
971
# via runai-model-streamer
972
-runai-model-streamer-s3==0.15.0
972
+runai-model-streamer-s3==0.15.3
973
# via runai-model-streamer
974
s3transfer==0.10.3
975
# via boto3
setup.py
+1
-1
@@ -797,7 +797,7 @@ setup(
797
"bench": ["pandas", "matplotlib", "seaborn", "datasets"],
798
"tensorizer": ["tensorizer==2.10.1"],
799
"fastsafetensors": ["fastsafetensors >= 0.1.10"],
800
- "runai": ["runai-model-streamer[s3,gcs] >= 0.15.0"],
800
+ "runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
801
"audio": [
802
"librosa",
803
"soundfile",
tests/model_executor/model_loader/runai_streamer_loader/__init__.py
renamed
tests/model_executor/model_loader/runai_streamer_loader/conftest.py
new
+39
@@ -0,0 +1,39 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
5
+from vllm.v1.executor import UniProcExecutor
6
+from vllm.v1.worker.worker_base import WorkerWrapperBase
7
+
8
+
9
+# This is a dummy executor for patching in test_runai_model_streamer_s3.py.
10
+# We cannot use vllm_runner fixture here, because it spawns worker process.
11
+# The worker process reimports the patched entities, and the patch is not applied.
12
+class RunaiDummyExecutor(UniProcExecutor):
13
+ def _init_executor(self) -> None:
14
+ distributed_init_method = get_distributed_init_method(get_ip(), get_open_port())
15
+
16
+ local_rank = 0
17
+ rank = 0
18
+ is_driver_worker = True
19
+
20
+ device_info = self.vllm_config.device_config.device.__str__().split(":")
21
+ if len(device_info) > 1:
22
+ local_rank = int(device_info[1])
23
+
24
+ worker_rpc_kwargs = dict(
25
+ vllm_config=self.vllm_config,
26
+ local_rank=local_rank,
27
+ rank=rank,
28
+ distributed_init_method=distributed_init_method,
29
+ is_driver_worker=is_driver_worker,
30
+ )
31
+
32
+ wrapper_kwargs = {
33
+ "vllm_config": self.vllm_config,
34
+ }
35
+
36
+ self.driver_worker = WorkerWrapperBase(**wrapper_kwargs)
37
+
38
+ self.collective_rpc("init_worker", args=([worker_rpc_kwargs],))
39
+ self.collective_rpc("init_device")
tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py
renamed
tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py
new
+52
@@ -0,0 +1,52 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from pathlib import Path
5
+
6
+from huggingface_hub import snapshot_download
7
+from runai_model_streamer.safetensors_streamer.streamer_mock import StreamerPatcher
8
+
9
+from vllm.engine.arg_utils import EngineArgs
10
+
11
+from .conftest import RunaiDummyExecutor
12
+
13
+load_format = "runai_streamer"
14
+test_model = "openai-community/gpt2"
15
+
16
+
17
+def test_runai_model_loader_download_files_s3_mocked_with_patch(
18
+ vllm_runner,
19
+ tmp_path: Path,
20
+ monkeypatch,
21
+):
22
+ patcher = StreamerPatcher(str(tmp_path))
23
+
24
+ test_mock_s3_model = "s3://my-mock-bucket/gpt2/"
25
+
26
+ # Download model from HF
27
+ mock_model_dir = f"{tmp_path}/gpt2"
28
+ snapshot_download(repo_id=test_model, local_dir=mock_model_dir)
29
+
30
+ monkeypatch.setattr(
31
+ "vllm.transformers_utils.runai_utils.runai_list_safetensors",
32
+ patcher.shim_list_safetensors,
33
+ )
34
+ monkeypatch.setattr(
35
+ "vllm.transformers_utils.runai_utils.runai_pull_files",
36
+ patcher.shim_pull_files,
37
+ )
38
+ monkeypatch.setattr(
39
+ "vllm.model_executor.model_loader.weight_utils.SafetensorsStreamer",
40
+ patcher.create_mock_streamer,
41
+ )
42
+
43
+ engine_args = EngineArgs(
44
+ model=test_mock_s3_model,
45
+ load_format=load_format,
46
+ tensor_parallel_size=1,
47
+ )
48
+
49
+ vllm_config = engine_args.create_engine_config()
50
+
51
+ executor = RunaiDummyExecutor(vllm_config)
52
+ executor.driver_worker.load_model()
tests/model_executor/model_loader/runai_streamer_loader/test_runai_utils.py
renamed
tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py
renamed
vllm/transformers_utils/runai_utils.py
+1
-3
@@ -18,9 +18,7 @@ SUPPORTED_SCHEMES = ["s3://", "gs://"]
18
try:
19
from runai_model_streamer import list_safetensors as runai_list_safetensors
20
from runai_model_streamer import pull_files as runai_pull_files
21
-except (ImportError, OSError):
22
- # see https://github.com/run-ai/runai-model-streamer/issues/26
23
- # OSError will be raised on arm64 platform
21
+except ImportError:
22
runai_model_streamer = PlaceholderModule("runai_model_streamer") # type: ignore[assignment]
23
runai_pull_files = runai_model_streamer.placeholder_attr("pull_files")
24
runai_list_safetensors = runai_model_streamer.placeholder_attr("list_safetensors")