File size: 2,870 Bytes
c34ff1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | import os
from scripts.launch_vllm import (
DEFAULT_RENDERER_NUM_WORKERS,
_enable_scale_out_endpoints,
_preprocessing_workers,
_set_render_thread_defaults,
_with_render_defaults,
render_throughput_defaults,
)
from speculators.data_generation.preprocessing import default_preprocessing_workers
def test_defaults_added_when_absent():
args = _with_render_defaults(["--port", "8000"])
api_servers, renderer_workers = render_throughput_defaults()
assert args == [
"--api-server-count",
str(api_servers),
"--renderer-num-workers",
str(renderer_workers),
"--port",
"8000",
]
def test_explicit_flag_follows_default():
args = _with_render_defaults(["--api-server-count", "1"])
assert args[-2:] == ["--api-server-count", "1"]
def test_headless_does_not_get_api_server_defaults():
args = _with_render_defaults(["--headless"])
assert "--api-server-count" not in args
assert "--renderer-num-workers" not in args
def test_render_thread_defaults_are_bounded_and_overrideable(monkeypatch):
for name in (
"OMP_NUM_THREADS",
"OPENBLAS_NUM_THREADS",
"MKL_NUM_THREADS",
"RAYON_NUM_THREADS",
):
monkeypatch.delenv(name, raising=False)
_set_render_thread_defaults()
assert os.environ["OMP_NUM_THREADS"] == "1"
assert os.environ["OPENBLAS_NUM_THREADS"] == "1"
assert os.environ["MKL_NUM_THREADS"] == "1"
assert os.environ["RAYON_NUM_THREADS"] == "2"
monkeypatch.setenv("RAYON_NUM_THREADS", "8")
_set_render_thread_defaults()
assert os.environ["RAYON_NUM_THREADS"] == "8"
def test_scale_out_endpoints_enabled_by_default(monkeypatch):
monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False)
_enable_scale_out_endpoints()
assert os.environ["VLLM_ENABLE_SCALE_OUT_ENDPOINTS"] == "1"
def test_explicit_scale_out_endpoint_setting_is_preserved(monkeypatch):
monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0")
_enable_scale_out_endpoints()
assert os.environ["VLLM_ENABLE_SCALE_OUT_ENDPOINTS"] == "0"
def test_sizing_respects_the_combined_budget():
assert default_preprocessing_workers(384) == 72
assert _preprocessing_workers(384) == 72
assert render_throughput_defaults(384) == (18, DEFAULT_RENDERER_NUM_WORKERS)
def test_sizing_scales_down_on_small_hosts():
assert default_preprocessing_workers(16) == 3
assert _preprocessing_workers(16) == 3
assert render_throughput_defaults(16) == (1, DEFAULT_RENDERER_NUM_WORKERS)
def test_sizing_uses_one_combined_budget():
preprocessing_workers = default_preprocessing_workers(160)
api_servers, _ = render_throughput_defaults(160)
assert preprocessing_workers == 30
assert api_servers == 7
assert preprocessing_workers * 3 + api_servers * 4 <= int(160 * 0.75)
|