A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][Model] Validate runai_streamer model_loader_extra_config (#45291)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Ting SUN committed
Jun 15, 2026 at 10:23 UTC
3d6ce816f02bfe7ac2d36ca4837e8e67179353d9
2 files changed
+75
-5
tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py
+44
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import os
5
import types
6
from unittest.mock import patch
7
@@ -78,3 +79,46 @@ def test_runai_passes_revision_by_name():
79
mock_idx.assert_called_once()
80
assert mock_idx.call_args.kwargs.get("revision") == "myrev"
81
assert "myrev" not in mock_idx.call_args.args
82
+
83
+
84
+def _runai_loader(extra):
85
+ return rsl.RunaiModelStreamerLoader(
86
+ LoadConfig(load_format="runai_streamer", model_loader_extra_config=extra)
87
+ )
88
+
89
+
90
+@pytest.mark.parametrize(
91
+ "extra, match",
92
+ [
93
+ ({"typo_key": 1}, "Unexpected extra config"),
94
+ ({"distributed": "yes"}, "distributed must be a bool"),
95
+ ({"concurrency": "16"}, "concurrency must be a positive integer"),
96
+ ({"concurrency": -1}, "concurrency must be a positive integer"),
97
+ ],
98
+)
99
+def test_runai_rejects_invalid_extra_config(extra, match):
100
+ # The loader used to silently drop unknown keys / wrong types / negatives.
101
+ with pytest.raises(ValueError, match=match):
102
+ _runai_loader(extra)
103
+
104
+
105
+def test_runai_accepts_valid_extra_config():
106
+ with patch.dict(os.environ, {}, clear=False):
107
+ os.environ.pop("RUNAI_STREAMER_CONCURRENCY", None)
108
+ os.environ.pop("RUNAI_STREAMER_MEMORY_LIMIT", None)
109
+ loader = _runai_loader(
110
+ {"distributed": True, "concurrency": 16, "memory_limit": 1024}
111
+ )
112
+ assert loader._is_distributed is True
113
+ assert os.environ["RUNAI_STREAMER_CONCURRENCY"] == "16"
114
+ assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == "1024"
115
+
116
+
117
+def test_runai_invalid_extra_config_leaves_environ_untouched():
118
+ # A later invalid key must not leave an earlier valid key applied to
119
+ # os.environ (all values are validated before any global mutation).
120
+ with patch.dict(os.environ, {}, clear=False):
121
+ os.environ.pop("RUNAI_STREAMER_CONCURRENCY", None)
122
+ with pytest.raises(ValueError, match="memory_limit must be a positive integer"):
123
+ _runai_loader({"concurrency": 16, "memory_limit": -5})
124
+ assert "RUNAI_STREAMER_CONCURRENCY" not in os.environ
vllm/model_executor/model_loader/runai_streamer_loader.py
+31
-5
@@ -31,12 +31,38 @@ class RunaiModelStreamerLoader(BaseModelLoader):
31
if load_config.model_loader_extra_config:
32
extra_config = load_config.model_loader_extra_config
33
34
- if isinstance(distributed := extra_config.get("distributed"), bool):
34
+ allowed_keys = {"distributed", "concurrency", "memory_limit"}
35
+ if unexpected_keys := set(extra_config) - allowed_keys:
36
+ raise ValueError(
37
+ "Unexpected extra config keys for runai_streamer: "
38
+ f"{unexpected_keys}"
39
+ )
40
+
41
+ if "distributed" in extra_config:
42
+ distributed = extra_config["distributed"]
43
+ if not isinstance(distributed, bool):
44
+ raise ValueError(f"distributed must be a bool, got {distributed!r}")
45
self._is_distributed = distributed
36
- if isinstance(concurrency := extra_config.get("concurrency"), int):
37
- os.environ["RUNAI_STREAMER_CONCURRENCY"] = str(concurrency)
38
- if isinstance(memory_limit := extra_config.get("memory_limit"), int):
39
- os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] = str(memory_limit)
46
+
47
+ # Validate every value before mutating os.environ, so a later
48
+ # invalid key cannot leave an earlier one partially applied.
49
+ env_updates: dict[str, str] = {}
50
+ for key, env_var in (
51
+ ("concurrency", "RUNAI_STREAMER_CONCURRENCY"),
52
+ ("memory_limit", "RUNAI_STREAMER_MEMORY_LIMIT"),
53
+ ):
54
+ if key in extra_config:
55
+ value = extra_config[key]
56
+ if (
57
+ isinstance(value, bool)
58
+ or not isinstance(value, int)
59
+ or value <= 0
60
+ ):
61
+ raise ValueError(
62
+ f"{key} must be a positive integer, got {value!r}"
63
+ )
64
+ env_updates[env_var] = str(value)
65
+ os.environ.update(env_updates)
66
67
runai_streamer_s3_endpoint = os.getenv("RUNAI_STREAMER_S3_ENDPOINT")
68
aws_endpoint_url = os.getenv("AWS_ENDPOINT_URL")