| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from argparse import Namespace |
| 5 | |
| 6 | from fastapi import FastAPI |
| 7 | |
| 8 | from vllm.entrypoints.launchers.api_server.routers import register_api_routers |
| 9 | from vllm.tasks import SupportedTask |
| 10 | |
| 11 | RENDER_PATHS = { |
| 12 | "/v1/chat/completions/render", |
| 13 | "/v1/chat/completions/derender", |
| 14 | "/v1/completions/render", |
| 15 | "/v1/completions/derender", |
| 16 | "/v1/messages/render", |
| 17 | "/v1/responses/render", |
| 18 | } |
| 19 | GENERATE_PATH = "/inference/v1/generate" |
| 20 | SCALE_OUT_PATHS = RENDER_PATHS | {GENERATE_PATH} |
| 21 | |
| 22 | |
| 23 | def registered_paths( |
| 24 | supported_tasks: tuple[SupportedTask, ...], |
| 25 | *, |
| 26 | tokens_only: bool = False, |
| 27 | enable_scale_out: bool = False, |
| 28 | ) -> set[str]: |
| 29 | app = FastAPI() |
| 30 | args = Namespace( |
| 31 | tokens_only=tokens_only, |
| 32 | enable_scale_out=enable_scale_out, |
| 33 | enable_fault_tolerance=False, |
| 34 | ) |
| 35 | app.state.args = args |
| 36 | register_api_routers(args, app, supported_tasks) |
| 37 | return {route.path for route in app.routes} |
| 38 | |
| 39 | |
| 40 | def test_scale_out_routes_are_disabled_by_default(): |
| 41 | paths = registered_paths(("generate",)) |
| 42 | |
| 43 | assert SCALE_OUT_PATHS.isdisjoint(paths) |
| 44 | |
| 45 | |
| 46 | def test_disabled_scale_out_routes_are_logged(caplog): |
| 47 | with caplog.at_level("INFO", logger="vllm.entrypoints.scale_out.factories"): |
| 48 | registered_paths(("generate",)) |
| 49 | |
| 50 | assert any("--enable-scale-out" in record.message for record in caplog.records) |
| 51 | |
| 52 | |
| 53 | def test_scale_out_routes_can_be_enabled_for_generate_server(): |
| 54 | paths = registered_paths(("generate",), enable_scale_out=True) |
| 55 | |
| 56 | assert paths >= SCALE_OUT_PATHS |
| 57 | |
| 58 | |
| 59 | def test_render_routes_remain_enabled_for_render_server(): |
| 60 | paths = registered_paths(("render",)) |
| 61 | |
| 62 | assert paths >= RENDER_PATHS |
| 63 | assert GENERATE_PATH not in paths |
| 64 | |
| 65 | |
| 66 | def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset(): |
| 67 | paths = registered_paths(("generate",), tokens_only=True) |
| 68 | |
| 69 | assert paths >= {GENERATE_PATH, "/abort_requests"} |