From 650609ceaeedfcf206c6661c68890a7cb8d3788d Mon Sep 17 00:00:00 2001 From: Grada Date: Tue, 8 Sep 2026 17:10:24 +0800 Subject: [PATCH] Disable gateway response buffering for model services --- .../gateway/resources/nginx/service.jinja2 | 3 + .../_internal/proxy/gateway/services/nginx.py | 1 + .../proxy/gateway/services/registry.py | 16 ++++++ src/dstack/_internal/proxy/lib/models.py | 1 + .../proxy/gateway/routers/test_registry.py | 28 ++++++++++ src/tests/_internal/proxy/gateway/test_app.py | 55 +++++++++++++++++++ 6 files changed, 104 insertions(+) diff --git a/src/dstack/_internal/proxy/gateway/resources/nginx/service.jinja2 b/src/dstack/_internal/proxy/gateway/resources/nginx/service.jinja2 index 7c8a5fcc9..380f1f4ea 100644 --- a/src/dstack/_internal/proxy/gateway/resources/nginx/service.jinja2 +++ b/src/dstack/_internal/proxy/gateway/resources/nginx/service.jinja2 @@ -76,6 +76,9 @@ server { location @ { set $dstack_replica_hit 1; {% if replicas %} + {% if not proxy_buffering %} + proxy_buffering off; + {% endif %} {% if cors_enabled %} proxy_hide_header 'Access-Control-Allow-Origin'; proxy_hide_header 'Access-Control-Allow-Methods'; diff --git a/src/dstack/_internal/proxy/gateway/services/nginx.py b/src/dstack/_internal/proxy/gateway/services/nginx.py index 74581c647..1486ea0c2 100644 --- a/src/dstack/_internal/proxy/gateway/services/nginx.py +++ b/src/dstack/_internal/proxy/gateway/services/nginx.py @@ -68,6 +68,7 @@ class ServiceConfig(SiteConfig): replicas: list[ReplicaConfig] has_router_replica: bool = False cors_enabled: bool = False + proxy_buffering: bool = True class ModelEntrypointConfig(SiteConfig): diff --git a/src/dstack/_internal/proxy/gateway/services/registry.py b/src/dstack/_internal/proxy/gateway/services/registry.py index 592a96492..bfa2eaecc 100644 --- a/src/dstack/_internal/proxy/gateway/services/registry.py +++ b/src/dstack/_internal/proxy/gateway/services/registry.py @@ -62,6 +62,7 @@ async def register_service( replicas=(), has_router_replica=has_router_replica, cors_enabled=cors_enabled, + proxy_buffering=model is None, ) async with lock: @@ -407,6 +408,7 @@ async def get_nginx_service_config( replicas=sorted(replicas, key=lambda r: r.id), # sort for reproducible configs has_router_replica=service.has_router_replica, cors_enabled=service.cors_enabled, + proxy_buffering=service.proxy_buffering, ) @@ -446,10 +448,24 @@ async def _migrate_cors_enabled(repo: GatewayProxyRepo) -> None: await repo.set_service(updated) +async def _migrate_proxy_buffering(repo: GatewayProxyRepo) -> None: + """Disable buffering for model services saved by older gateways.""" + services = await repo.list_services() + model_runs = { + (project_name, model.run_name) + for project_name in {service.project_name for service in services} + for model in await repo.list_models(project_name) + } + for service in services: + if service.proxy_buffering and (service.project_name, service.run_name) in model_runs: + await repo.set_service(service.model_copy(update={"proxy_buffering": False})) + + async def apply_all( repo: GatewayProxyRepo, nginx: Nginx, service_conn_pool: ServiceConnectionPool ) -> None: await _migrate_cors_enabled(repo) + await _migrate_proxy_buffering(repo) service_tasks = [ apply_service( service=service, diff --git a/src/dstack/_internal/proxy/lib/models.py b/src/dstack/_internal/proxy/lib/models.py index 53eb13e74..a05c521f1 100644 --- a/src/dstack/_internal/proxy/lib/models.py +++ b/src/dstack/_internal/proxy/lib/models.py @@ -63,6 +63,7 @@ class Service(ImmutableModel): replicas: tuple[Replica, ...] has_router_replica: bool = False cors_enabled: bool = False # only used on gateways; enabled for openai-format models + proxy_buffering: bool = True # only used on gateways; disabled for model services @model_validator(mode="before") @classmethod diff --git a/src/tests/_internal/proxy/gateway/routers/test_registry.py b/src/tests/_internal/proxy/gateway/routers/test_registry.py index 20a0e148c..a8f5e7f32 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_registry.py +++ b/src/tests/_internal/proxy/gateway/routers/test_registry.py @@ -88,6 +88,34 @@ def sample_model_options(name: str = "test-model") -> dict: @pytest.mark.asyncio class TestRegisterService: + @pytest.mark.parametrize("model_format", [None, "openai", "tgi"]) + async def test_proxy_buffering( + self, tmp_path: Path, system_mocks: Mocks, model_format: Optional[str] + ) -> None: + repo = GatewayProxyRepo() + client = make_client(tmp_path, repo=repo) + options = None + if model_format is not None: + model = {"type": "chat", "name": "test-model", "format": model_format} + if model_format == "openai": + model["prefix"] = "/v1" + else: + model.update(chat_template="{{ messages }}", eos_token="") + options = {"openai": {"model": model}} + response = await client.post( + "/api/registry/test-proj/services/register", + json=register_service_payload(options=options), + ) + assert response.status_code == 200 + response = await client.post( + "/api/registry/test-proj/services/test-run/replicas/register", + json=register_replica_payload(), + ) + assert response.status_code == 200 + conf = (tmp_path / "sites-enabled" / "443-test-run.gtw.test.conf").read_text() + service_location = conf.split("location @ {", 1)[1].split("}", 1)[0] + assert ("proxy_buffering off;" in service_location) == (model_format is not None) + async def test_register(self, tmp_path: Path, system_mocks: Mocks) -> None: client = make_client(tmp_path) resp = await client.post( diff --git a/src/tests/_internal/proxy/gateway/test_app.py b/src/tests/_internal/proxy/gateway/test_app.py index 335be84fa..efc434fbb 100644 --- a/src/tests/_internal/proxy/gateway/test_app.py +++ b/src/tests/_internal/proxy/gateway/test_app.py @@ -1,3 +1,5 @@ +import json +from datetime import datetime from pathlib import Path import pytest @@ -7,6 +9,11 @@ from dstack._internal.proxy.gateway.repo.repo import GatewayProxyRepo from dstack._internal.proxy.gateway.services.nginx import Nginx from dstack._internal.proxy.gateway.testing.common import Mocks +from dstack._internal.proxy.lib.models import ( + ChatModel, + OpenAIChatModelFormat, + TGIChatModelFormat, +) from dstack._internal.proxy.lib.testing.common import make_project, make_service @@ -32,3 +39,51 @@ async def test_lifespan(tmp_path: Path, system_mocks: Mocks) -> None: assert system_mocks.open_conn.call_count == 1 assert system_mocks.close_conn.call_count == 0 assert system_mocks.close_conn.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_format", ["openai", "tgi"]) +async def test_lifespan_migrates_proxy_buffering( + tmp_path: Path, system_mocks: Mocks, model_format: str +) -> None: + state_file = tmp_path / "state-v2.json" + repo = GatewayProxyRepo(file=state_file) + for project in ("model-proj", "web-proj"): + await repo.set_project(make_project(project)) + await repo.set_service( + make_service(project, "same-run", domain=f"{project}.gtw.test", https=False) + ) + await repo.set_model( + ChatModel( + project_name="model-proj", + name="test-model", + created_at=datetime(2026, 1, 1), + run_name="same-run", + format_spec=( + OpenAIChatModelFormat(prefix="/v1") + if model_format == "openai" + else TGIChatModelFormat(chat_template="{{ messages }}", eos_token="") + ), + ) + ) + # Old gateways saved services without a proxy_buffering field. + state = json.loads(state_file.read_text()) + for services in state["services"].values(): + for service in services.values(): + service.pop("proxy_buffering", None) + state_file.write_text(json.dumps(state)) + nginx_dir = tmp_path / "nginx" + conf_dir = nginx_dir / "sites-enabled" + conf_dir.mkdir(parents=True) + # Restart twice to cover both migration and subsequent persisted-state recovery. + for _ in range(2): + repo = GatewayProxyRepo.load(state_file) + app = make_app(repo=repo, nginx=Nginx(nginx_dir=nginx_dir)) + async with lifespan(app): + model_conf = (conf_dir / "443-model-proj.gtw.test.conf").read_text() + web_conf = (conf_dir / "443-web-proj.gtw.test.conf").read_text() + assert "proxy_buffering off;" in model_conf + assert "proxy_buffering off;" not in web_conf + persisted = json.loads(state_file.read_text()) + assert persisted["services"]["model-proj"]["same-run"]["proxy_buffering"] is False + assert persisted["services"]["web-proj"]["same-run"]["proxy_buffering"] is True