diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index c45a76f9df..fbae3a8e63 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -923,7 +923,9 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: action="store_true", default=False, help="""enable RL control plane (HTTP APIs, router rl_rpyc, model RlBackendOps). - When disabled (default), RL routes/services are not started.""", + When disabled (default), RL routes/services are not started. + When enabled, set LIGHTLLM_RL_API_KEY and send it as an Authorization: Bearer token + to the RL control-plane routes.""", ) parser.add_argument( "--enable_torch_memory_saver", diff --git a/lightllm/server/api_http.py b/lightllm/server/api_http.py index 8c8ed804ab..6b65e027d3 100755 --- a/lightllm/server/api_http.py +++ b/lightllm/server/api_http.py @@ -547,8 +547,15 @@ async def metrics() -> Response: # RL 控制面接口（abort / pause / flush / memory / weight update），见 api_http_rl.py from .api_http_rl import router as rl_router +from .rl_auth import RLStartArgs, mount_rl_router, validate_rl_auth + +_rl_router_mounted = False + + +def _mount_rl_router(args: RLStartArgs) -> None: + global _rl_router_mounted + _rl_router_mounted = mount_rl_router(app, rl_router, args, _rl_router_mounted) -app.include_router(rl_router) # PD 分离控制面接口（worker 发现、P/D 注册与 KV 状态上报），见 api_http_pd.py from .api_http_pd import router as pd_router @@ -595,7 +602,10 @@ async def shutdown(): async def startup_event(): logger.info("server start up") loop = asyncio.get_event_loop() - g_objs.set_args(get_env_start_args()) + args = get_env_start_args() + validate_rl_auth(args) + g_objs.set_args(args) + _mount_rl_router(args) loop.create_task(g_objs.httpserver_manager.handle_loop()) logger.info(f"server start up ok, loop use is {asyncio.get_event_loop()}") return diff --git a/lightllm/server/api_http_rl.py b/lightllm/server/api_http_rl.py index b73d76170e..621392d469 100644 --- a/lightllm/server/api_http_rl.py +++ b/lightllm/server/api_http_rl.py @@ -12,7 +12,7 @@ from http import HTTPStatus -from fastapi import APIRouter, Request +from fastapi import APIRouter, Depends, Request from fastapi.responses import JSONResponse, Response from lightllm.server.io_struct import ( @@ -30,10 +30,11 @@ from lightllm.utils.log_utils import init_logger from .api_errors import create_error_response +from .rl_auth import require_rl_auth logger = init_logger(__name__) -router = APIRouter() +router = APIRouter(dependencies=[Depends(require_rl_auth)]) async def handle_request_common(request_obj, handler): diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index c67ab9582d..2a11e5ba25 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -33,6 +33,7 @@ get_default_disable_chunked_mm_input, ) from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args +from .rl_auth import validate_rl_auth logger = init_logger(__name__) @@ -43,6 +44,7 @@ def _set_envs_and_config(args: StartArgs): def _launch_subprocesses(args: StartArgs): _validate_pd_features(args) + validate_rl_auth(args) _set_envs_and_config(args) if args.target_vocab_topk_sampling is not None: diff --git a/lightllm/server/rl_auth.py b/lightllm/server/rl_auth.py new file mode 100644 index 0000000000..62529f9285 --- /dev/null +++ b/lightllm/server/rl_auth.py @@ -0,0 +1,52 @@ +"""Authentication helpers for the RL control-plane HTTP routes.""" + +import os +import secrets +from typing import Protocol + +from fastapi import APIRouter, FastAPI, HTTPException, Request + +RL_API_KEY_ENV = "LIGHTLLM_RL_API_KEY" + + +class RLStartArgs(Protocol): + enable_rl: bool + + +def get_rl_api_key() -> str | None: + """Return the configured RL control-plane key without logging it.""" + key = os.getenv(RL_API_KEY_ENV) + return key if key else None + + +def validate_rl_auth(args: RLStartArgs) -> None: + """Reject an RL-enabled server that has no control-plane credential.""" + if args.enable_rl and get_rl_api_key() is None: + raise ValueError(f"{RL_API_KEY_ENV} must be set when --enable_rl is enabled") + + +def mount_rl_router(app: FastAPI, router: APIRouter, args: RLStartArgs, mounted: bool) -> bool: + """Mount the RL router once and only for an RL-enabled server.""" + if args.enable_rl and not mounted: + app.include_router(router) + return True + return mounted + + +def require_rl_auth(request: Request) -> None: + """Require a valid bearer token for an RL control-plane request.""" + configured_key = get_rl_api_key() + if configured_key is None: + raise HTTPException(status_code=503, detail="RL control-plane authentication is not configured") + + authorization_parts = request.headers.get("authorization", "").split() + if ( + len(authorization_parts) != 2 + or authorization_parts[0].lower() != "bearer" + or not secrets.compare_digest(authorization_parts[1].encode("utf-8"), configured_key.encode("utf-8")) + ): + raise HTTPException( + status_code=401, + detail="Invalid RL control-plane credentials", + headers={"WWW-Authenticate": "Bearer"}, + ) diff --git a/unit_tests/server/test_rl_api_auth.py b/unit_tests/server/test_rl_api_auth.py new file mode 100644 index 0000000000..5d798e0ad9 --- /dev/null +++ b/unit_tests/server/test_rl_api_auth.py @@ -0,0 +1,117 @@ +import sys +from types import ModuleType +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI, HTTPException +from fastapi.testclient import TestClient +from starlette.requests import Request + +from lightllm.server.api_http_rl import router as rl_router +from lightllm.server.rl_auth import ( + RL_API_KEY_ENV, + mount_rl_router, + require_rl_auth, + validate_rl_auth, +) + + +def _request(headers=None): + raw_headers = [(key.lower().encode(), value.encode()) for key, value in (headers or {}).items()] + return Request({"type": "http", "headers": raw_headers}) + + +def _route_paths(routes): + for route in routes: + path = getattr(route, "path", None) + if path is not None: + yield path + yield from _route_paths(getattr(route, "routes", [])) + original_router = getattr(route, "original_router", None) + if original_router is not None: + yield from _route_paths(original_router.routes) + + +def test_rl_routes_are_mounted_only_when_enabled(monkeypatch): + app = FastAPI() + + mounted = mount_rl_router(app, rl_router, SimpleNamespace(enable_rl=False), False) + assert "/pause_generation" not in set(_route_paths(app.routes)) + + monkeypatch.setenv(RL_API_KEY_ENV, "control-secret") + mounted = mount_rl_router(app, rl_router, SimpleNamespace(enable_rl=True), mounted) + assert "/pause_generation" in set(_route_paths(app.routes)) + assert mount_rl_router(app, rl_router, SimpleNamespace(enable_rl=True), mounted) is True + + +def test_mounted_rl_routes_reject_unauthenticated_requests(monkeypatch): + monkeypatch.setenv(RL_API_KEY_ENV, "control-secret") + app = FastAPI() + app.include_router(rl_router) + + response = TestClient(app).post("/pause_generation") + + assert response.status_code == 401 + assert response.headers["www-authenticate"] == "Bearer" + + +def test_every_rl_route_requires_auth(): + assert rl_router.routes + assert all( + any(dependency.call is require_rl_auth for dependency in route.dependant.dependencies) + for route in rl_router.routes + ) + + +def test_mounted_rl_route_dispatches_with_valid_bearer_token(monkeypatch): + monkeypatch.setenv(RL_API_KEY_ENV, "control-secret") + calls = [] + + async def pause_generation(): + calls.append("pause") + + api_http_module = ModuleType("lightllm.server.api_http") + api_http_module.g_objs = SimpleNamespace(httpserver_manager=SimpleNamespace(pause_generation=pause_generation)) + monkeypatch.setitem(sys.modules, "lightllm.server.api_http", api_http_module) + + app = FastAPI() + app.include_router(rl_router) + response = TestClient(app).post( + "/pause_generation", + headers={"Authorization": "Bearer control-secret"}, + ) + + assert response.status_code == 200 + assert calls == ["pause"] + + +def test_rl_auth_accepts_only_the_configured_bearer_token(monkeypatch): + monkeypatch.setenv(RL_API_KEY_ENV, "control-secret") + + assert require_rl_auth(_request({"authorization": "Bearer control-secret"})) is None + assert require_rl_auth(_request({"authorization": "bEaReR control-secret"})) is None + + with pytest.raises(HTTPException) as error: + require_rl_auth(_request({"authorization": "Bearer wrong-secret"})) + assert error.value.status_code == 401 + + with pytest.raises(HTTPException) as error: + require_rl_auth(_request({"authorization": "Bearer incorrect-" + chr(233)})) + assert error.value.status_code == 401 + + +def test_rl_auth_rejects_missing_configuration(monkeypatch): + monkeypatch.delenv(RL_API_KEY_ENV, raising=False) + + with pytest.raises(HTTPException) as error: + require_rl_auth(_request({"authorization": "Bearer control-secret"})) + assert error.value.status_code == 503 + + +def test_rl_startup_requires_a_control_plane_key(monkeypatch): + monkeypatch.delenv(RL_API_KEY_ENV, raising=False) + with pytest.raises(ValueError, match=RL_API_KEY_ENV): + validate_rl_auth(SimpleNamespace(enable_rl=True)) + + monkeypatch.setenv(RL_API_KEY_ENV, "control-secret") + validate_rl_auth(SimpleNamespace(enable_rl=True))