From a97f5b8211b71ee7803d8624455688ddac061a37 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Tue, 1 Sep 2026 20:51:03 +0800 Subject: [PATCH] feat(service): configure gateway proxy read timeout Signed-off-by: mikemikimike <13286568797@163.com> --- .../_internal/core/models/configurations.py | 17 ++++++++++++ .../gateway/resources/nginx/service.jinja2 | 4 +-- .../proxy/gateway/routers/registry.py | 1 + .../proxy/gateway/schemas/registry.py | 1 + .../_internal/proxy/gateway/services/nginx.py | 1 + .../proxy/gateway/services/registry.py | 7 ++++- src/dstack/_internal/proxy/lib/const.py | 3 +++ src/dstack/_internal/proxy/lib/models.py | 2 ++ .../pipeline_tasks/gateway_replicas.py | 5 ++++ .../server/services/gateways/client.py | 2 ++ .../_internal/server/services/proxy/repo.py | 10 ++++++- src/dstack/_internal/server/settings.py | 3 +++ .../core/models/test_configurations.py | 27 +++++++++++++++++++ .../proxy/gateway/routers/test_registry.py | 18 ++++++++++++- 14 files changed, 96 insertions(+), 5 deletions(-) diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 09fd5eec7..32f6a72be 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -1208,6 +1208,16 @@ class ServiceConfigurationParams(CoreModel): ] = None auth: Annotated[bool, Field(description="Enable the authorization")] = True + proxy_read_timeout: Annotated[ + Optional[Duration], + Field( + description=( + "Maximum interval between reads from the service upstream when using a gateway. " + "Defaults to 300 seconds." + ) + ), + ] = None + scaling: Annotated[ Optional[ScalingSpec], Field(description="The auto-scaling rules. Required if `replicas` is set to a range"), @@ -1222,6 +1232,13 @@ class ServiceConfigurationParams(CoreModel): ), ] = None # None = omitted (may get default when model is set); [] = explicit empty + @field_validator("proxy_read_timeout") + @classmethod + def validate_proxy_read_timeout(cls, v: Optional[Duration]) -> Optional[Duration]: + if v is not None and v < 1: + raise ValueError("Proxy read timeout must be at least 1 second") + return v + replicas: Annotated[ Optional[Range[int]], Field( diff --git a/src/dstack/_internal/proxy/gateway/resources/nginx/service.jinja2 b/src/dstack/_internal/proxy/gateway/resources/nginx/service.jinja2 index 7c8a5fcc9..d03fd72c2 100644 --- a/src/dstack/_internal/proxy/gateway/resources/nginx/service.jinja2 +++ b/src/dstack/_internal/proxy/gateway/resources/nginx/service.jinja2 @@ -68,7 +68,7 @@ server { proxy_http_version 1.1; proxy_set_header Upgrade $http_upgrade; proxy_set_header Connection "Upgrade"; - proxy_read_timeout 300s; + proxy_read_timeout {{ proxy_read_timeout }}s; {% else %} return 503; {% endif %} @@ -88,7 +88,7 @@ server { proxy_pass http://{{ domain }}.upstream; proxy_set_header X-Real-IP $remote_addr; proxy_set_header Host $host; - proxy_read_timeout 300s; + proxy_read_timeout {{ proxy_read_timeout }}s; {% else %} return 503; {% endif %} diff --git a/src/dstack/_internal/proxy/gateway/routers/registry.py b/src/dstack/_internal/proxy/gateway/routers/registry.py index c5fb7c0d6..a2d0de273 100644 --- a/src/dstack/_internal/proxy/gateway/routers/registry.py +++ b/src/dstack/_internal/proxy/gateway/routers/registry.py @@ -35,6 +35,7 @@ async def register_service( rate_limits=body.rate_limits, auth=body.auth, client_max_body_size=body.client_max_body_size, + proxy_read_timeout=body.proxy_read_timeout, model=body.options.openai.model if body.options.openai is not None else None, ssh_private_key=body.ssh_private_key, repo=repo, diff --git a/src/dstack/_internal/proxy/gateway/schemas/registry.py b/src/dstack/_internal/proxy/gateway/schemas/registry.py index 9a800c2ce..ae77ae958 100644 --- a/src/dstack/_internal/proxy/gateway/schemas/registry.py +++ b/src/dstack/_internal/proxy/gateway/schemas/registry.py @@ -43,6 +43,7 @@ class RegisterServiceRequest(BaseModel): https: bool auth: bool client_max_body_size: int + proxy_read_timeout: Optional[int] = Field(default=None, ge=1) options: Options ssh_private_key: str rate_limits: tuple[RateLimit, ...] = () diff --git a/src/dstack/_internal/proxy/gateway/services/nginx.py b/src/dstack/_internal/proxy/gateway/services/nginx.py index 29562c555..db25d14b7 100644 --- a/src/dstack/_internal/proxy/gateway/services/nginx.py +++ b/src/dstack/_internal/proxy/gateway/services/nginx.py @@ -62,6 +62,7 @@ class ServiceConfig(SiteConfig): project_name: str auth: bool client_max_body_size: int + proxy_read_timeout: int access_log_path: Path limit_req_zones: list[LimitReqZoneConfig] locations: list[LocationConfig] diff --git a/src/dstack/_internal/proxy/gateway/services/registry.py b/src/dstack/_internal/proxy/gateway/services/registry.py index 592a96492..e8d7d4ede 100644 --- a/src/dstack/_internal/proxy/gateway/services/registry.py +++ b/src/dstack/_internal/proxy/gateway/services/registry.py @@ -19,7 +19,7 @@ ServiceConfig, ) from dstack._internal.proxy.lib import models -from dstack._internal.proxy.lib.const import ROUTER_WHITELISTED_PATHS +from dstack._internal.proxy.lib.const import DEFAULT_PROXY_READ_TIMEOUT, ROUTER_WHITELISTED_PATHS from dstack._internal.proxy.lib.errors import ProxyError, UnexpectedProxyError from dstack._internal.proxy.lib.repo import BaseProxyRepo from dstack._internal.proxy.lib.services.service_connection import ( @@ -42,6 +42,7 @@ async def register_service( rate_limits: tuple[models.RateLimit, ...], auth: bool, client_max_body_size: int, + proxy_read_timeout: Optional[int], model: Optional[schemas.AnyModel], ssh_private_key: str, repo: GatewayProxyRepo, @@ -59,6 +60,9 @@ async def register_service( rate_limits=rate_limits, auth=auth, client_max_body_size=client_max_body_size, + proxy_read_timeout=( + proxy_read_timeout if proxy_read_timeout is not None else DEFAULT_PROXY_READ_TIMEOUT + ), replicas=(), has_router_replica=has_router_replica, cors_enabled=cors_enabled, @@ -401,6 +405,7 @@ async def get_nginx_service_config( project_name=service.project_name, auth=service.auth, client_max_body_size=service.client_max_body_size, + proxy_read_timeout=service.proxy_read_timeout, access_log_path=ACCESS_LOG_PATH, limit_req_zones=limit_req_zones, locations=locations, diff --git a/src/dstack/_internal/proxy/lib/const.py b/src/dstack/_internal/proxy/lib/const.py index 43ede03ac..082a19e92 100644 --- a/src/dstack/_internal/proxy/lib/const.py +++ b/src/dstack/_internal/proxy/lib/const.py @@ -2,6 +2,9 @@ Shared constants for proxy components (gateway + in-server proxy). """ +DEFAULT_PROXY_READ_TIMEOUT = 300 +"""Default maximum interval between reads from a service upstream, in seconds.""" + # Inference endpoints exposed by the in-replica HTTP router. Applies to both # SGLang's router and Dynamo's `dynamo.frontend` — they share the # OpenAI-compatible endpoint surface. diff --git a/src/dstack/_internal/proxy/lib/models.py b/src/dstack/_internal/proxy/lib/models.py index 53eb13e74..d6262f92a 100644 --- a/src/dstack/_internal/proxy/lib/models.py +++ b/src/dstack/_internal/proxy/lib/models.py @@ -7,6 +7,7 @@ from typing_extensions import Annotated from dstack._internal.core.models.instances import SSHConnectionParams +from dstack._internal.proxy.lib.const import DEFAULT_PROXY_READ_TIMEOUT from dstack._internal.proxy.lib.errors import UnexpectedProxyError @@ -59,6 +60,7 @@ class Service(ImmutableModel): rate_limits: tuple[RateLimit, ...] = () # only used on gateways auth: bool client_max_body_size: int # only enforced on gateways + proxy_read_timeout: int = DEFAULT_PROXY_READ_TIMEOUT # only enforced on gateways strip_prefix: bool = True # only used in-server replicas: tuple[Replica, ...] has_router_replica: bool = False diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py index 82891204a..2069806ca 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -1281,6 +1281,11 @@ async def _register_service( gateway_https=get_gateway_https(gateway_configuration), auth=run_spec.configuration.auth, client_max_body_size=settings.DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE, + proxy_read_timeout=( + run_spec.configuration.proxy_read_timeout + if run_spec.configuration.proxy_read_timeout is not None + else settings.DEFAULT_SERVICE_PROXY_READ_TIMEOUT + ), options=service_spec.options, rate_limits=run_spec.configuration.rate_limits, ssh_private_key=run_model.project.ssh_private_key, diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py index ed6c8d78b..fdee20ff0 100644 --- a/src/dstack/_internal/server/services/gateways/client.py +++ b/src/dstack/_internal/server/services/gateways/client.py @@ -45,6 +45,7 @@ async def register_service( gateway_https: bool, auth: bool, client_max_body_size: int, + proxy_read_timeout: Optional[int], options: dict, rate_limits: list[RateLimit], ssh_private_key: str, @@ -61,6 +62,7 @@ async def register_service( "https": service_https, "auth": auth, "client_max_body_size": client_max_body_size, + "proxy_read_timeout": proxy_read_timeout, "options": options, "rate_limits": [limit.model_dump() for limit in rate_limits], "ssh_private_key": ssh_private_key, diff --git a/src/dstack/_internal/server/services/proxy/repo.py b/src/dstack/_internal/server/services/proxy/repo.py index d0afd41de..3eb85d27e 100644 --- a/src/dstack/_internal/server/services/proxy/repo.py +++ b/src/dstack/_internal/server/services/proxy/repo.py @@ -32,7 +32,10 @@ from dstack._internal.server.services.instances import get_instance_remote_connection_info from dstack._internal.server.services.jobs import get_job_spec from dstack._internal.server.services.runs import get_run_spec -from dstack._internal.server.settings import DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE +from dstack._internal.server.settings import ( + DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE, + DEFAULT_SERVICE_PROXY_READ_TIMEOUT, +) from dstack._internal.utils.common import get_or_error _ANY_MODEL_ADAPTER = pydantic.TypeAdapter(AnyModel) @@ -133,6 +136,11 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic https=None, auth=run_spec.configuration.auth, client_max_body_size=DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE, + proxy_read_timeout=( + run_spec.configuration.proxy_read_timeout + if run_spec.configuration.proxy_read_timeout is not None + else DEFAULT_SERVICE_PROXY_READ_TIMEOUT + ), strip_prefix=run_spec.configuration.strip_prefix, replicas=tuple(replicas), has_router_replica=has_router_replica, diff --git a/src/dstack/_internal/server/settings.py b/src/dstack/_internal/server/settings.py index 360572535..7005464f6 100644 --- a/src/dstack/_internal/server/settings.py +++ b/src/dstack/_internal/server/settings.py @@ -174,6 +174,9 @@ def get_database_url() -> str: DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE = int( os.getenv("DSTACK_DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE", 64 * 1024 * 1024) ) +DEFAULT_SERVICE_PROXY_READ_TIMEOUT = int( + os.getenv("DSTACK_DEFAULT_SERVICE_PROXY_READ_TIMEOUT", 300) +) SERVER_DEFAULT_DOCKER_REGISTRY = os.getenv("DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY") or None SERVER_DEFAULT_DOCKER_REGISTRY_USERNAME = ( diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index e0d91c718..f5510bb73 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -74,6 +74,33 @@ def test_service_model_probes_none_when_omitted(self): assert isinstance(parsed, ServiceConfiguration) assert parsed.probes is None + @pytest.mark.windows + def test_service_proxy_read_timeout_accepts_duration(self): + parsed = parse_run_configuration( + { + "type": "service", + "commands": ["python3 -m http.server"], + "port": 8000, + "proxy_read_timeout": "15m", + } + ) + assert isinstance(parsed, ServiceConfiguration) + assert parsed.proxy_read_timeout == 15 * 60 + + @pytest.mark.windows + def test_service_proxy_read_timeout_rejects_zero(self): + with pytest.raises( + ConfigurationError, match="Proxy read timeout must be at least 1 second" + ): + parse_run_configuration( + { + "type": "service", + "commands": ["python3 -m http.server"], + "port": 8000, + "proxy_read_timeout": 0, + } + ) + def test_service_model_does_not_override_explicit_probes(self): conf = { "type": "service", diff --git a/src/tests/_internal/proxy/gateway/routers/test_registry.py b/src/tests/_internal/proxy/gateway/routers/test_registry.py index b7a5b2e2e..6594ed8fd 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_registry.py +++ b/src/tests/_internal/proxy/gateway/routers/test_registry.py @@ -11,7 +11,8 @@ from dstack._internal.core.errors import SSHError from dstack._internal.proxy.gateway.app import make_app from dstack._internal.proxy.gateway.repo.repo import GatewayProxyRepo -from dstack._internal.proxy.gateway.services.nginx import Nginx +from dstack._internal.proxy.gateway.services.nginx import Nginx, ReplicaConfig +from dstack._internal.proxy.gateway.services.registry import get_nginx_service_config from dstack._internal.proxy.gateway.testing.common import Mocks from dstack._internal.proxy.lib.models import ChatModel, OpenAIChatModelFormat from dstack._internal.proxy.lib.testing.common import make_project, make_service @@ -118,6 +119,21 @@ async def test_register(self, tmp_path: Path, system_mocks: Mocks) -> None: assert "upstream" not in conf assert "return 503;" in conf + @pytest.mark.windows + async def test_nginx_service_config_uses_proxy_read_timeout(self) -> None: + service = make_service( + "test-proj", "test-run", domain="test-run.gtw.test", https=False + ).model_copy(update={"proxy_read_timeout": 900}) + replicas = [ReplicaConfig(id="replica", socket=Path("/tmp/replica.sock"), port=80)] + config = await get_nginx_service_config(service, replicas) + assert "proxy_read_timeout 900s;" in config.render() + + default_config = await get_nginx_service_config( + make_service("test-proj", "default-run", domain="default.gtw.test", https=False), + replicas, + ) + assert "proxy_read_timeout 300s;" in default_config.render() + async def test_legacy_register_without_id(self, tmp_path: Path, system_mocks: Mocks) -> None: repo = GatewayProxyRepo() client = make_client(tmp_path, repo=repo)