| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Malformed bodies on RL dev routes are 400s that never reach the engine.""" |
| 4 | |
| 5 | import json |
| 6 | from types import SimpleNamespace |
| 7 | |
| 8 | import pytest |
| 9 | from fastapi import FastAPI |
| 10 | from fastapi.testclient import TestClient |
| 11 | |
| 12 | from vllm.entrypoints.serve.dev.rlhf.api_router import router as rlhf_router |
| 13 | from vllm.entrypoints.serve.dev.rpc.api_router import router as rpc_router |
| 14 | from vllm.entrypoints.serve.exception_handling.register import init_exception_handler |
| 15 | |
| 16 | pytestmark = pytest.mark.cpu_test |
| 17 | |
| 18 | |
| 19 | class Engine: |
| 20 | def __init__(self): |
| 21 | self.calls: list[str] = [] |
| 22 | |
| 23 | def __getattr__(self, name): |
| 24 | async def call(*args, **kwargs): |
| 25 | self.calls.append(name) |
| 26 | |
| 27 | return call |
| 28 | |
| 29 | |
| 30 | @pytest.fixture |
| 31 | def client(): |
| 32 | app = FastAPI() |
| 33 | app.include_router(rlhf_router) |
| 34 | app.include_router(rpc_router) |
| 35 | init_exception_handler(app) |
| 36 | app.state.engine_client = Engine() |
| 37 | app.state.args = SimpleNamespace(log_error_stack=False) |
| 38 | return TestClient(app) |
| 39 | |
| 40 | |
| 41 | def post(client, path, body): |
| 42 | return client.post( |
| 43 | path, content=json.dumps(body), headers={"content-type": "application/json"} |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | @pytest.mark.parametrize( |
| 48 | "path", |
| 49 | [ |
| 50 | "/abort_requests", |
| 51 | "/init_weight_transfer_engine", |
| 52 | "/update_weights", |
| 53 | "/collective_rpc", |
| 54 | ], |
| 55 | ) |
| 56 | @pytest.mark.parametrize("body", [[], 1, "x", None]) |
| 57 | def test_non_object_body_is_rejected(client, path, body): |
| 58 | assert post(client, path, body).status_code == 400 |
| 59 | assert client.app.state.engine_client.calls == [] |
| 60 | |
| 61 | |
| 62 | @pytest.mark.parametrize( |
| 63 | "path,body", |
| 64 | [ |
| 65 | ("/init_weight_transfer_engine", {}), |
| 66 | ("/init_weight_transfer_engine", {"init_info": []}), |
| 67 | ("/init_weight_transfer_engine", {"init_info": 1}), |
| 68 | ("/update_weights", {}), |
| 69 | ("/update_weights", {"update_info": 1}), |
| 70 | ("/update_weights", {"update_info": [1, 2]}), |
| 71 | ("/update_weights", {"update_info": "x"}), |
| 72 | ], |
| 73 | ) |
| 74 | def test_invalid_field_is_rejected_before_the_engine(client, path, body): |
| 75 | # Reaching the engine would also abort an in-progress weight update. |
| 76 | assert post(client, path, body).status_code == 400 |
| 77 | assert client.app.state.engine_client.calls == [] |
| 78 | |
| 79 | |
| 80 | @pytest.mark.parametrize( |
| 81 | "path,body,engine_call", |
| 82 | [ |
| 83 | ("/abort_requests", {"request_ids": ["a"]}, "abort"), |
| 84 | ( |
| 85 | "/init_weight_transfer_engine", |
| 86 | {"init_info": {}}, |
| 87 | "init_weight_transfer_engine", |
| 88 | ), |
| 89 | ("/update_weights", {"update_info": {"names": []}}, "update_weights"), |
| 90 | ("/update_weights", {"update_info": [{}, {}]}, "update_weights"), |
| 91 | ("/collective_rpc", {"method": "m"}, "collective_rpc"), |
| 92 | ], |
| 93 | ) |
| 94 | def test_valid_body_reaches_the_engine(client, path, body, engine_call): |
| 95 | assert post(client, path, body).status_code == 200 |
| 96 | assert client.app.state.engine_client.calls == [engine_call] |