From c8c92c6b57cabf4c21e165c6b9b24d8a88ae4277 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 7 Oct 2026 14:30:48 +0200 Subject: [PATCH 1/2] Fix gateway recovery after high request load --- pyproject.toml | 2 + src/dstack/_internal/proxy/gateway/auth.py | 5 +- .../gateway/resources/licenses/httpcore.txt | 27 ++ .../resources/systemd/dstack.gateway.service | 1 + .../proxy/gateway/services/server_client.py | 156 ++++++- .../gateway/services/test_server_client.py | 440 ++++++++++++++++++ .../services/gateways/test_gateway_update.py | 149 ++++++ 7 files changed, 773 insertions(+), 7 deletions(-) create mode 100644 src/dstack/_internal/proxy/gateway/resources/licenses/httpcore.txt create mode 100644 src/tests/_internal/proxy/gateway/services/test_server_client.py create mode 100644 src/tests/_internal/server/services/gateways/test_gateway_update.py diff --git a/pyproject.toml b/pyproject.toml index bba8433701..6b1c852aa8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -220,6 +220,8 @@ gateway = [ "aiorwlock", "aiocache", "httpx>=0.28.0", + # Gateway pool workarounds use private HTTPcore APIs; review before upgrading. + "httpcore>=1.0.5,<1.1", "jinja2", ] server = [ diff --git a/src/dstack/_internal/proxy/gateway/auth.py b/src/dstack/_internal/proxy/gateway/auth.py index 2523e9ff59..fff1d2f313 100644 --- a/src/dstack/_internal/proxy/gateway/auth.py +++ b/src/dstack/_internal/proxy/gateway/auth.py @@ -4,12 +4,15 @@ from dstack._internal.proxy.lib.auth import BaseProxyAuthProvider from dstack._internal.proxy.lib.errors import UnexpectedProxyError +_AUTH_CACHE_TTL_SECONDS = 60 +"""Seconds to cache allowed or denied access for each project/token pair.""" + class GatewayProxyAuthProvider(BaseProxyAuthProvider): def __init__(self, server_client: httpx.AsyncClient) -> None: self._server_client = server_client - @cached(ttl=60, noself=True, skip_cache_func=lambda r: r is None) + @cached(ttl=_AUTH_CACHE_TTL_SECONDS, noself=True, skip_cache_func=lambda r: r is None) async def is_project_member(self, project_name: str, token: str) -> bool: try: resp = await self._server_client.post( diff --git a/src/dstack/_internal/proxy/gateway/resources/licenses/httpcore.txt b/src/dstack/_internal/proxy/gateway/resources/licenses/httpcore.txt new file mode 100644 index 0000000000..311b2b56c5 --- /dev/null +++ b/src/dstack/_internal/proxy/gateway/resources/licenses/httpcore.txt @@ -0,0 +1,27 @@ +Copyright © 2020, [Encode OSS Ltd](https://www.encode.io/). +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/src/dstack/_internal/proxy/gateway/resources/systemd/dstack.gateway.service b/src/dstack/_internal/proxy/gateway/resources/systemd/dstack.gateway.service index 409e51fe3f..a8027f451a 100644 --- a/src/dstack/_internal/proxy/gateway/resources/systemd/dstack.gateway.service +++ b/src/dstack/_internal/proxy/gateway/resources/systemd/dstack.gateway.service @@ -9,6 +9,7 @@ User=ubuntu Group=ubuntu Restart=always RestartSec=5 +LimitNOFILE=65535 [Install] WantedBy=default.target diff --git a/src/dstack/_internal/proxy/gateway/services/server_client.py b/src/dstack/_internal/proxy/gateway/services/server_client.py index cdcf3aa046..6300e8bb2b 100644 --- a/src/dstack/_internal/proxy/gateway/services/server_client.py +++ b/src/dstack/_internal/proxy/gateway/services/server_client.py @@ -2,9 +2,11 @@ import logging import random from dataclasses import dataclass, field +from itertools import chain from pathlib import Path from typing import Container, Dict, Generator, List +import httpcore import httpx logger = logging.getLogger(__name__) @@ -13,22 +15,25 @@ @dataclass class CachedClientInfo: + """HTTP client and connection-failure history for one server, owned by HTTPMultiClient.""" + client: httpx.AsyncClient socket: Path connect_errors: List[datetime.datetime] = field(default_factory=lambda: []) def seems_disconnected(self) -> bool: + # Only ConnectError is recorded; a returned HTTP response clears this history. if len(self.connect_errors) < 2: return False return self.connect_errors[-1] - self.connect_errors[0] >= datetime.timedelta(minutes=2) class HTTPMultiClient(httpx.AsyncClient): - """ - An HTTP client that sends requests to randomly chosen Unix sockets from a specified - directory. This allows to balance the load between multiple HTTP server replicas. - Automatically deletes sockets that stop responding. - Used for requesting random dstack-server replicas from the gateway. + """Gateway HTTP client for uncached project-access checks against dstack-server. + + GatewayProxyAuthProvider uses this client and owns authorization decisions and + their cache. Return the server's HTTP response, including error statuses. Raise + httpx.RequestError when no connected server can complete the request. """ def __init__(self, sockets_dir: Path): @@ -40,6 +45,7 @@ async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response: errors: List[httpx.RequestError] = [] clients_count = 0 + # Try another replica after request errors; HTTP error responses are returned. for clients_count, client in enumerate(self._iter_clients_rand(), start=1): try: resp = await client.client.send(request, *args, **kwargs) @@ -69,6 +75,7 @@ async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response: def _iter_clients_rand(self) -> Generator[CachedClientInfo, None, None]: sockets = list(self._sockets_dir.glob("*.sock")) self._evict_clients(stems_to_keep={s.stem for s in sockets}) + # Each socket forwards to a server replica; shuffle to distribute auth checks. random.shuffle(sockets) for socket in sockets: @@ -81,7 +88,7 @@ def _iter_clients_rand(self) -> Generator[CachedClientInfo, None, None]: @staticmethod def _make_client(socket: Path) -> CachedClientInfo: client = httpx.AsyncClient( - transport=httpx.AsyncHTTPTransport(uds=str(socket.absolute())), + transport=_ServerTransport(uds=str(socket.absolute())), base_url=BASE_URL, ) return CachedClientInfo( @@ -93,3 +100,140 @@ def _evict_clients(self, stems_to_keep: Container[str]) -> None: self._clients_cache = { stem: client for stem, client in self._clients_cache.items() if stem in stems_to_keep } + + +class _ServerTransport(httpx.AsyncHTTPTransport): + """HTTPMultiClient's transport for one dstack-server Unix socket. + + Preserve HTTPX's connection limits and use the request's HTTPX timeouts. + _ServerConnectionPool provides recovery after overload. + """ + + def __init__(self, uds: str): + super().__init__(uds=uds) + # Preserve HTTPX's existing default limits. + self._pool = _ServerConnectionPool( + uds=uds, + max_connections=100, + max_keepalive_connections=20, + keepalive_expiry=5.0, + ) + + +class _ServerConnectionPool(httpcore.AsyncConnectionPool): + """Authorization connection pool used by _ServerTransport, with overload recovery. + + Preserve HTTPcore's connection limits, reuse, ordering and timeout behavior. + Failed or cancelled requests must not permanently consume capacity or prevent + later requests from proceeding once connections are available. + See https://github.com/dstackai/dstack/issues/4338. + """ + + async def handle_async_request(self, request: httpcore.Request) -> httpcore.Response: + try: + return await super().handle_async_request(request) + except BaseException: + # HTTPcore removes the failed/cancelled request, but can leave its newly + # assigned connection behind before it starts. That slot is never reused. + # Remove this orphan cleanup once https://github.com/encode/httpcore/pull/1099 ships. + with self._optional_thread_lock: + owners = {id(request.connection) for request in self._requests} + orphaned = [ + connection + for connection in self._connections + if id(connection) not in owners + and not connection.is_closed() + and not connection.is_available() + and not connection.is_idle() + ] + if not orphaned: + raise + for connection in orphaned: + self._connections.remove(connection) + # Wake queued requests now that slots are free. Response streams still + # retain their owners, and HTTPcore shields connection closing below. + closing = orphaned + self._assign_requests_to_connections() + await self._close_connections(closing) + raise + + def _assign_requests_to_connections(self) -> List[httpcore.AsyncConnectionInterface]: + # Adapted from HTTPcore; see resources/licenses/httpcore.txt. + if ( + len(self._requests) > len(self._connections) + and len(self._connections) >= self._max_connections + and all( + not connection.is_closed() + and not connection.has_expired() + and not connection.is_idle() + and not connection.is_available() + for connection in self._connections + ) + ): + # A full, busy pool cannot assign or close anything. Scanning every queued + # request here makes a backlog of timeouts quadratic to drain. + return [] + closing = [] + for connection in list(self._connections): + if connection.is_closed(): + self._connections.remove(connection) + elif connection.has_expired(): + self._connections.remove(connection) + closing.append(connection) + elif connection.is_idle() and len(self._connections) > self._max_keepalive_connections: + # Preserve HTTPcore's existing total-connection keepalive check. + self._connections.remove(connection) + closing.append(connection) + + if not self._requests: + return closing + if len(self._requests) == 1: + request = self._requests[0] + if not request.is_queued(): + return closing + # Reuse a connection without building lists for a single request. + origin = request.request.url.origin + for connection in self._connections: + if connection.can_handle_request(origin) and connection.is_available(): + request.assign_to_connection(connection) + return closing + + # Skip assignment scans when no request is waiting for a connection. + queued_requests = (request for request in self._requests if request.is_queued()) + first_request = next(queued_requests, None) + if first_request is None: + return closing + + # No connection state changes while this synchronous pass runs. Inspect + # availability once, rather than scanning every connection for every waiter. + available = [connection for connection in self._connections if connection.is_available()] + idle = [connection for connection in self._connections if connection.is_idle()] + + def add_connection(origin: httpcore.Origin) -> httpcore.AsyncConnectionInterface: + connection = self.create_connection(origin) + self._connections.append(connection) + if connection.is_available(): + available.append(connection) + if connection.is_idle(): + idle.append(connection) + return connection + + for request in chain((first_request,), queued_requests): + if len(self._connections) >= self._max_connections and not available and not idle: + break + origin = request.request.url.origin + connection = next( + (connection for connection in available if connection.can_handle_request(origin)), + None, + ) + if connection is not None: + request.assign_to_connection(connection) + elif len(self._connections) < self._max_connections: + request.assign_to_connection(add_connection(origin)) + elif idle: + connection = idle.pop(0) + self._connections.remove(connection) + if connection in available: + available.remove(connection) + closing.append(connection) + request.assign_to_connection(add_connection(origin)) + return closing diff --git a/src/tests/_internal/proxy/gateway/services/test_server_client.py b/src/tests/_internal/proxy/gateway/services/test_server_client.py new file mode 100644 index 0000000000..adbcd2fad7 --- /dev/null +++ b/src/tests/_internal/proxy/gateway/services/test_server_client.py @@ -0,0 +1,440 @@ +import asyncio +from pathlib import Path + +import httpcore +import pytest +from httpcore._async.connection_pool import AsyncPoolRequest + +from dstack._internal.proxy.gateway.services.server_client import ( + HTTPMultiClient, + _ServerConnectionPool, + _ServerTransport, +) + +RESPONSE = b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok" + + +class TestServerConnectionPool: + @pytest.mark.asyncio + @pytest.mark.parametrize("failure", ["cancel", "timeout"]) + async def test_reclaims_connection_assigned_to_cancelled_waiter(self, monkeypatch, failure): + backend = httpcore.AsyncMockBackend([RESPONSE]) + queued = asyncio.Event() + wait_for_connection = AsyncPoolRequest.wait_for_connection + + async def wait(self, timeout=None): + was_queued = self.connection is None + if was_queued: + queued.set() + connection = await wait_for_connection(self, timeout) + if was_queued and failure == "timeout": + # The deadline fires while a connection is assigned, before the + # waiter resumes. Inject the outcome without waiting on a clock. + raise httpcore.PoolTimeout() + return connection + + monkeypatch.setattr(AsyncPoolRequest, "wait_for_connection", wait) + async with _ServerConnectionPool(max_connections=1, network_backend=backend) as pool: + first = await pool.handle_async_request( + httpcore.Request("GET", "http://first/", headers={"Host": "first"}) + ) + second = asyncio.create_task(pool.request("GET", "http://second/")) + await queued.wait() + # Releasing the only slot assigns a new, unstarted connection to second. + await first.aclose() + assert len(pool.connections) == 1 + assert pool.connections[0].info() == "CONNECTING" + if failure == "cancel": + second.cancel() + with pytest.raises( + asyncio.CancelledError if failure == "cancel" else httpcore.PoolTimeout + ): + await second + + assert not pool._requests + assert pool.connections == [] + response = await pool.request("GET", "http://next/") + assert response.status == 200 + assert response.content == b"ok" + + @pytest.mark.asyncio + async def test_assigns_reclaimed_slot_to_queued_waiter(self, monkeypatch): + backend = httpcore.AsyncMockBackend([RESPONSE]) + queued = {host: asyncio.Event() for host in (b"cancelled", b"survivor")} + waiting = {} + wait_for_connection = AsyncPoolRequest.wait_for_connection + + async def wait(self, timeout=None): + host = self.request.url.host + if self.connection is None and host in queued: + waiting[host] = self + queued[host].set() + return await wait_for_connection(self, timeout) + + monkeypatch.setattr(AsyncPoolRequest, "wait_for_connection", wait) + async with _ServerConnectionPool(max_connections=1, network_backend=backend) as pool: + first = await pool.handle_async_request( + httpcore.Request("GET", "http://first/", headers={"Host": "first"}) + ) + cancelled = asyncio.create_task(pool.request("GET", "http://cancelled/")) + survivor = asyncio.create_task(pool.request("GET", "http://survivor/")) + try: + await queued[b"cancelled"].wait() + await queued[b"survivor"].wait() + await first.aclose() + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + # Reclaiming the abandoned slot must wake the existing waiter without + # requiring another incoming request to trigger pool maintenance. + assert waiting[b"survivor"].connection is not None + response = await survivor + assert response.status == 200 + assert response.content == b"ok" + finally: + cancelled.cancel() + survivor.cancel() + await asyncio.gather(cancelled, survivor, return_exceptions=True) + + @pytest.mark.asyncio + async def test_preserves_connections_being_established(self, monkeypatch): + backend = httpcore.AsyncMockBackend([RESPONSE]) + connecting = asyncio.Event() + release_connection = asyncio.Event() + connect_tcp = backend.connect_tcp + + async def connect(host, *args, **kwargs): + if host == "first": + connecting.set() + await release_connection.wait() + else: + raise httpcore.ConnectError("simulated connection failure") + return await connect_tcp(host, *args, **kwargs) + + monkeypatch.setattr(backend, "connect_tcp", connect) + async with _ServerConnectionPool(max_connections=2, network_backend=backend) as pool: + first = asyncio.create_task(pool.request("GET", "http://first/")) + try: + await connecting.wait() + connection = pool.connections[0] + assert connection.info() == "CONNECTING" + # A failed request triggers cleanup while the first still owns its slot. + with pytest.raises(httpcore.ConnectError): + await pool.request("GET", "http://second/") + assert connection in pool.connections + assert not first.done() + release_connection.set() + response = await first + assert response.status == 200 + assert response.content == b"ok" + finally: + first.cancel() + await asyncio.gather(first, return_exceptions=True) + + @pytest.mark.asyncio + async def test_preserves_connections_with_unread_response_streams(self, monkeypatch): + backend = httpcore.AsyncMockBackend([RESPONSE]) + connect_tcp = backend.connect_tcp + + async def connect(host, *args, **kwargs): + if host == "second": + raise httpcore.ConnectError("simulated connection failure") + return await connect_tcp(host, *args, **kwargs) + + monkeypatch.setattr(backend, "connect_tcp", connect) + async with _ServerConnectionPool(max_connections=2, network_backend=backend) as pool: + first = await pool.handle_async_request( + httpcore.Request("GET", "http://first/", headers={"Host": "first"}) + ) + with pytest.raises(httpcore.ConnectError): + await pool.request("GET", "http://second/") + # Failed-request cleanup must not close the first response's connection. + assert await first.aread() == b"ok" + await first.aclose() + + @pytest.mark.asyncio + async def test_reuses_idle_connections(self): + backend = httpcore.AsyncMockBackend([RESPONSE, RESPONSE]) + async with _ServerConnectionPool(network_backend=backend) as pool: + await pool.request("GET", "http://dstack/") + connection = pool.connections[0] + response = await pool.request("GET", "http://dstack/") + assert response.content == b"ok" + assert pool.connections == [connection] + + @pytest.mark.asyncio + @pytest.mark.parametrize("http2", [False, True]) + async def test_single_waiter_reuses_first_match_without_scanning_later_connections( + self, monkeypatch, http2 + ): + async with _ServerConnectionPool( + max_connections=3, + max_keepalive_connections=3, + http2=http2, + network_backend=httpcore.AsyncMockBackend([RESPONSE]), + ) as pool: + scheme = "https" if http2 else "http" + requests = [ + httpcore.Request("GET", f"{scheme}://{host}/", headers={"Host": host}) + for host in ("other", "dstack", "dstack") + ] + if http2: + # HTTPS HTTP/2 connections can be shared while still connecting; + # reusing them must not depend on the connection being idle. + pool._connections = [ + pool.create_connection(request.url.origin) for request in requests + ] + assert all(not connection.is_idle() for connection in pool.connections) + else: + responses = [await pool.handle_async_request(request) for request in requests] + for response in responses: + assert await response.aread() == b"ok" + await response.aclose() + assert all(connection.is_idle() for connection in pool.connections) + different_origin, first_match, later_match = pool.connections + assert not different_origin.can_handle_request(requests[1].url.origin) + + def unexpected_scan(*args): + pytest.fail("A single waiter must stop scanning after its first reusable match") + + monkeypatch.setattr(later_match, "can_handle_request", unexpected_scan) + monkeypatch.setattr(later_match, "is_available", unexpected_scan) + waiter = AsyncPoolRequest(requests[1]) + pool._requests.append(waiter) + + assert pool._assign_requests_to_connections() == [] + assert waiter.connection is first_match + assert pool.connections == [different_origin, first_match, later_match] + + @pytest.mark.asyncio + @pytest.mark.parametrize("active_request", [False, True]) + @pytest.mark.parametrize("state", ["idle", "excess_idle", "expired", "closed"]) + async def test_response_completion_without_waiters_skips_availability_scan( + self, monkeypatch, active_request, state + ): + async with _ServerConnectionPool( + max_connections=2, + max_keepalive_connections=0 if state == "excess_idle" else 2, + keepalive_expiry=-float("inf") if state == "expired" else 5.0, + network_backend=httpcore.AsyncMockBackend([RESPONSE]), + ) as pool: + responses = [ + await pool.handle_async_request( + httpcore.Request("GET", "http://dstack/", headers={"Host": "dstack"}) + ) + for _ in range(1 + active_request) + ] + released, *busy = pool.connections + + def is_available(): + pytest.fail( + "Response completion must not scan availability without queued requests" + ) + + with monkeypatch.context() as patch: + for connection in pool.connections: + patch.setattr(connection, "is_available", is_available) + # Closing an unread HTTP/1.1 response closes its connection. Reading + # it first returns it idle, allowing expiry/keepalive cleanup instead. + if state != "closed": + assert await responses[0].aread() == b"ok" + await responses[0].aclose() + + if state == "idle": + assert pool.connections == [released, *busy] + assert released.is_idle() + else: + assert pool.connections == busy + assert released.is_closed() + assert len(pool._requests) == active_request + if active_request: + assert pool._requests[0].connection is busy[0] + assert not pool._requests[0].is_queued() + assert not busy[0].is_closed() + assert await responses[1].aread() == b"ok" + await responses[1].aclose() + + @pytest.mark.asyncio + async def test_saturated_pool_waits_without_scanning_unavailable_connections( + self, monkeypatch + ): + backend = httpcore.AsyncMockBackend([RESPONSE, RESPONSE]) + queued = asyncio.Event() + wait_for_connection = AsyncPoolRequest.wait_for_connection + + async def wait(self, timeout=None): + if self.connection is None: + queued.set() + return await wait_for_connection(self, timeout) + + monkeypatch.setattr(AsyncPoolRequest, "wait_for_connection", wait) + async with _ServerConnectionPool(max_connections=1, network_backend=backend) as pool: + first = await pool.handle_async_request( + httpcore.Request("GET", "http://dstack/", headers={"Host": "dstack"}) + ) + connection = pool.connections[0] + can_handle_request = connection.can_handle_request + scans = 0 + + def can_handle(origin): + nonlocal scans + scans += 1 + return can_handle_request(origin) + + monkeypatch.setattr(connection, "can_handle_request", can_handle) + second = asyncio.create_task(pool.request("GET", "http://dstack/")) + try: + await queued.wait() + assert scans == 0 + assert not second.done() + await first.aread() + await first.aclose() + response = await second + assert response.status == 200 + assert response.content == b"ok" + finally: + second.cancel() + await asyncio.gather(second, return_exceptions=True) + + @pytest.mark.asyncio + @pytest.mark.parametrize("reuse", [True, False]) + async def test_response_completion_does_not_rescan_busy_connections(self, monkeypatch, reuse): + async with _ServerConnectionPool( + max_connections=3, + max_keepalive_connections=3 if reuse else 1, + network_backend=httpcore.AsyncMockBackend([RESPONSE]), + ) as pool: + responses = [ + await pool.handle_async_request( + httpcore.Request("GET", "http://dstack/", headers={"Host": "dstack"}) + ) + for _ in range(3) + ] + released, *busy = pool.connections + scans = 0 + + def can_handle(origin): + nonlocal scans + scans += 1 + return True + + for connection in busy: + monkeypatch.setattr(connection, "can_handle_request", can_handle) + queued = [ + AsyncPoolRequest(httpcore.Request("GET", "http://dstack/")) for _ in range(100) + ] + pool._requests.extend(queued) + + await responses[0].aread() + await responses[0].aclose() + + # A completed response bypasses the all-busy shortcut. Pool maintenance + # must not rescan busy connections for each of the remaining waiters. + assert scans <= len(busy) + if reuse: + assert all(request.connection is released for request in queued) + else: + assert released.is_closed() + assert queued[0].connection is not None + assert queued[0].connection is not released + assert all(request.connection is None for request in queued[1:]) + assert len(pool.connections) == 3 + + @pytest.mark.asyncio + async def test_evicted_connection_is_not_reused_by_later_waiter(self): + async with _ServerConnectionPool( + max_connections=1, + network_backend=httpcore.AsyncMockBackend([RESPONSE]), + ) as pool: + await pool.request("GET", "http://first/") + original = pool.connections[0] + queued = [ + AsyncPoolRequest(httpcore.Request("GET", f"http://{host}/")) + for host in ("second", "first") + ] + pool._requests.extend(queued) + + closing = pool._assign_requests_to_connections() + + assert closing == [original] + assert queued[0].connection is pool.connections[0] + assert queued[0].connection is not original + assert queued[1].connection is None + await pool._close_connections(closing) + + @pytest.mark.asyncio + @pytest.mark.parametrize("state", ["spare_capacity", "idle", "closed", "expired"]) + async def test_assigns_when_capacity_or_connections_are_available(self, state): + async with _ServerConnectionPool( + max_connections=2 if state == "spare_capacity" else 1, + max_keepalive_connections=1, + network_backend=httpcore.AsyncMockBackend([RESPONSE]), + ) as pool: + response = await pool.handle_async_request( + httpcore.Request("GET", "http://dstack/", headers={"Host": "dstack"}) + ) + connection = pool.connections[0] + assert isinstance(connection, httpcore.AsyncHTTPConnection) + if state != "spare_capacity": + await response.aread() + await response.aclose() + if state == "closed": + await connection.aclose() + elif state == "expired": + assert isinstance(connection._connection, httpcore.AsyncHTTP11Connection) + connection._connection._expire_at = -float("inf") + queued = [ + AsyncPoolRequest(httpcore.Request("GET", "http://dstack/")) for _ in range(2) + ] + pool._requests.extend(queued) + closing = pool._assign_requests_to_connections() + assert queued[0].connection is not None + if state == "idle": + assert queued[0].connection is connection + else: + assert queued[0].connection is not connection + if state == "expired": + assert connection in closing + await pool._close_connections(closing) + + @pytest.mark.asyncio + @pytest.mark.parametrize("existing_connection", [True, False]) + async def test_assigns_available_http2_connection_to_multiple_requests( + self, existing_connection + ): + async with _ServerConnectionPool(max_connections=1, http2=True) as pool: + request = httpcore.Request("GET", "https://dstack/") + if existing_connection: + connection = pool.create_connection(request.url.origin) + pool._connections.append(connection) + assert not connection.is_idle() + assert connection.is_available() + queued = [AsyncPoolRequest(request) for _ in range(2)] + pool._requests.extend(queued) + assert pool._assign_requests_to_connections() == [] + assert len(pool.connections) == 1 + assert all(waiter.connection is pool.connections[0] for waiter in queued) + + +class TestHTTPMultiClient: + @pytest.mark.asyncio + async def test_server_transport_uses_fixed_pool(self, tmp_path: Path): + socket = tmp_path / "server.sock" + socket.touch() + client = HTTPMultiClient(tmp_path) + cached = next(client._iter_clients_rand()) + transport = cached.client._transport + assert isinstance(transport, _ServerTransport) + pool = transport._pool + assert isinstance(pool, _ServerConnectionPool) + assert pool._uds == str(socket) + assert (pool._max_connections, pool._max_keepalive_connections) == (100, 20) + assert pool._keepalive_expiry == 5.0 + pool._network_backend = httpcore.AsyncMockBackend([RESPONSE]) + try: + response = await client.post("/api/projects/test/get") + assert response.status_code == 200 + assert response.text == "ok" + finally: + await cached.client.aclose() + await client.aclose() diff --git a/src/tests/_internal/server/services/gateways/test_gateway_update.py b/src/tests/_internal/server/services/gateways/test_gateway_update.py new file mode 100644 index 0000000000..2a15bf3615 --- /dev/null +++ b/src/tests/_internal/server/services/gateways/test_gateway_update.py @@ -0,0 +1,149 @@ +import shlex +import subprocess +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from unittest.mock import AsyncMock, Mock + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal import settings as core_settings +from dstack._internal.server.services import gateways +from dstack._internal.server.testing.common import ( + create_backend, + create_gateway_replica, + create_project, +) + + +@pytest.mark.asyncio +class TestInitGateways: + @pytest.mark.parametrize("target_version", ["0.22.2", "0.22.3"]) + async def test_release_version_controls_existing_gateway_update( + self, + session: AsyncSession, + monkeypatch: pytest.MonkeyPatch, + gateway_update_sandbox: "_GatewayUpdateSandbox", + target_version: str, + ): + sandbox = gateway_update_sandbox + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + previously_updated_at = now - timedelta(minutes=2) + monkeypatch.setattr(core_settings, "DSTACK_GATEWAY_PACKAGE_URL", None) + monkeypatch.setattr(core_settings, "DSTACK_VERSION", target_version) + monkeypatch.setattr(gateways.settings, "SKIP_GATEWAY_UPDATE", False) + monkeypatch.setattr(gateways, "get_current_datetime", lambda: now) + monkeypatch.setattr( + gateways.gateway_connections_pool, + "get_or_add", + AsyncMock(return_value=sandbox.connection), + ) + monkeypatch.setattr( + gateways.gateway_connections_pool, + "all", + AsyncMock(return_value=[sandbox.connection]), + ) + configure = AsyncMock() + monkeypatch.setattr(gateways, "configure_gateway_replica", configure) + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + replica = await create_gateway_replica(session=session, backend=backend) + replica.app_updated_at = previously_updated_at + await session.commit() + + await gateways.init_gateways(session) + + configure.assert_awaited_once_with(sandbox.connection, attempts=7) + events = sandbox.events.read_text().splitlines() + if target_version == "0.22.2": + assert events == ["blue pip show dstack"] + assert (sandbox.root / "version").read_text().strip() == "blue" + assert replica.app_updated_at == previously_updated_at + return + + assert (sandbox.root / "version").read_text().strip() == "green" + assert replica.app_updated_at == now + install_package = f"green pip install dstack[gateway]=={target_version}" + install_service = "green python -m dstack._internal.proxy.gateway.systemd install" + reload_service = "systemctl daemon-reload active=green" + restart_service = "systemctl restart dstack.gateway active=green" + assert ( + events.index(install_package) + < events.index(install_service) + < events.index(reload_service) + < events.index(restart_service) + ) + + +@dataclass +class _GatewayUpdateSandbox: + root: Path + events: Path + connection: Mock + + +@pytest.fixture +def gateway_update_sandbox(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + # Execute the real shell: mocking aexec would miss quoting, version gating, and the order + # of service installation / venv switching / restart. Only the remote commands are stubbed; + # no SSH, package installation, sudo, or real systemd operations are performed. + root = tmp_path / "gateway root" + root.mkdir() + (root / "version").write_text("blue\n") + events = tmp_path / "events" + events.touch() + for color in ("blue", "green"): + venv_bin = root / color / "bin" + venv_bin.mkdir(parents=True) + for command in ("pip", "python"): + _write_executable( + venv_bin / command, + f'printf "%s\\n" "{color} {command} $*" >> "$TEST_EVENTS"\n' + + ( + 'if [ "$1" = show ]; then echo "Version: 0.22.2"; fi\n' + if command == "pip" + else "" + ), + ) + command_bin = tmp_path / "commands" + command_bin.mkdir() + _write_executable( + command_bin / "sudo", + 'if [ "$1" = systemctl ]; then\n' + ' printf "%s\\n" "$* active=$(cat "$TEST_ROOT/version")" >> "$TEST_EVENTS"\n' + "else\n" + ' exec "$@"\n' + "fi\n", + ) + # The only script substitution redirects its hard-coded remote root to this temporary dir. + remote_root = "root=/home/ubuntu/dstack" + assert gateways._GATEWAY_UPDATE_SCRIPT.count(remote_root) == 1 + monkeypatch.setattr( + gateways, + "_GATEWAY_UPDATE_SCRIPT", + gateways._GATEWAY_UPDATE_SCRIPT.replace(remote_root, f"root={shlex.quote(str(root))}"), + ) + + def execute(command: str) -> str: + return subprocess.run( + shlex.split(command), + check=True, + capture_output=True, + text=True, + timeout=5, + env={ + "PATH": f"{command_bin}:/usr/bin:/bin", + "TEST_EVENTS": str(events), + "TEST_ROOT": str(root), + }, + ).stdout + + connection = Mock(ip_address="1.1.1.1") + connection.tunnel.aexec = AsyncMock(side_effect=execute) + return _GatewayUpdateSandbox(root, events, connection) + + +def _write_executable(path: Path, script: str): + path.write_text("#!/bin/sh\nset -e\n" + script) + path.chmod(0o755) From d1e7f9ab05125e3901e2f152e927639dddb8a34c Mon Sep 17 00:00:00 2001 From: Alvils Sture Date: Thu, 8 Oct 2026 05:48:58 +0300 Subject: [PATCH 2/2] Reopen gateway replica SSH tunnels whose ssh process exited (#4352) --- src/dstack/_internal/proxy/gateway/app.py | 8 +- .../proxy/lib/services/service_connection.py | 71 +++++++++++- .../_internal/proxy/lib/services/__init__.py | 0 .../lib/services/test_service_connection.py | 108 ++++++++++++++++++ 4 files changed, 182 insertions(+), 5 deletions(-) create mode 100644 src/tests/_internal/proxy/lib/services/__init__.py create mode 100644 src/tests/_internal/proxy/lib/services/test_service_connection.py diff --git a/src/dstack/_internal/proxy/gateway/app.py b/src/dstack/_internal/proxy/gateway/app.py index 6627e9e3f3..313a29edcd 100644 --- a/src/dstack/_internal/proxy/gateway/app.py +++ b/src/dstack/_internal/proxy/gateway/app.py @@ -1,6 +1,7 @@ """FastAPI app running on a gateway.""" -from contextlib import asynccontextmanager +import asyncio +from contextlib import asynccontextmanager, suppress from pathlib import Path from typing import Optional @@ -29,6 +30,7 @@ from dstack._internal.proxy.gateway.services.server_client import HTTPMultiClient from dstack._internal.proxy.gateway.services.stats import StatsCollector from dstack._internal.proxy.lib.routers.model_proxy import router as model_proxy_router +from dstack._internal.proxy.lib.services.service_connection import maintain_service_connections from dstack._internal.utils.common import run_async from dstack.version import __version__ @@ -45,9 +47,13 @@ async def lifespan(app: FastAPI): service_conn_pool = await injector.get_service_connection_pool() await run_async(nginx.write_global_conf) await apply_all(repo, nginx, service_conn_pool) + maintenance = asyncio.create_task(maintain_service_connections(service_conn_pool)) yield + maintenance.cancel() + with suppress(asyncio.CancelledError): + await maintenance await service_conn_pool.remove_all() diff --git a/src/dstack/_internal/proxy/lib/services/service_connection.py b/src/dstack/_internal/proxy/lib/services/service_connection.py index c8229ad53a..9b8c0160ae 100644 --- a/src/dstack/_internal/proxy/lib/services/service_connection.py +++ b/src/dstack/_internal/proxy/lib/services/service_connection.py @@ -18,12 +18,15 @@ from dstack._internal.proxy.lib.errors import UnexpectedProxyError from dstack._internal.proxy.lib.models import Project, Replica, Service from dstack._internal.proxy.lib.repo import BaseProxyRepo -from dstack._internal.utils.common import get_or_error +from dstack._internal.utils.common import get_or_error, run_async +from dstack._internal.utils.env import environ from dstack._internal.utils.logging import get_logger from dstack._internal.utils.path import FileContent logger = get_logger(__name__) OPEN_TUNNEL_TIMEOUT = 10 +TUNNEL_CHECK_INTERVAL = environ.get_int("DSTACK_PROXY_TUNNEL_CHECK_INTERVAL", default=15) +"""Seconds between checks that reopen replica SSH tunnels whose ssh process has exited.""" class ServiceClient(httpx.AsyncClient): @@ -75,6 +78,9 @@ def __init__(self, project: Project, service: Service, replica: Replica) -> None timeout=service.read_timeout, ) self._is_open = asyncio.locks.Event() + self._replica_id = replica.id + self._lock = asyncio.Lock() + self._closed = False @property def app_socket_path(self) -> Path: @@ -85,9 +91,33 @@ async def open(self) -> None: self._is_open.set() async def close(self) -> None: - self._is_open.clear() - await self._client.aclose() - await self._tunnel.aclose() + async with self._lock: + self._closed = True + self._is_open.clear() + await self._client.aclose() + await self._tunnel.aclose() + + async def reopen_if_exited(self) -> bool: + """ + Reopen the SSH tunnel if its ssh process has exited, e.g. after the replica's SSH server + dropped the connection because the gateway missed keepalives. The local app socket path + is kept, so nginx and the HTTP client keep working without reconfiguration. + + Returns `True` if the tunnel was reopened. + """ + if self._closed or not self._is_open.is_set(): + return False + if await self._tunnel.acheck(): + return False + async with self._lock: + if self._closed or await self._tunnel.acheck(): + return False + logger.warning("SSH tunnel to service replica %s exited, reopening", self._replica_id) + # The control socket of a killed ssh process stays on disk. A new ssh with + # `ControlMaster=auto` would then run without a control socket and fail every check. + await run_async(_remove_file, Path(self._tunnel.control_sock_path)) + await self._tunnel.aopen() + return True async def client(self) -> ServiceClient: await asyncio.wait_for(self._is_open.wait(), timeout=OPEN_TUNNEL_TIMEOUT) @@ -122,6 +152,20 @@ async def remove(self, replica_id: str) -> None: if connection is not None: await connection.close() + async def reopen_exited(self) -> None: + connections = list(self.connections.items()) + results = await asyncio.gather( + *(connection.reopen_if_exited() for _, connection in connections), + return_exceptions=True, + ) + for (replica_id, _), result in zip(connections, results): + if isinstance(result, BaseException): + logger.warning( + "Failed to reopen SSH tunnel to service replica %s: %r", replica_id, result + ) + elif result: + logger.info("Reopened SSH tunnel to service replica %s", replica_id) + async def remove_all(self) -> None: replica_ids = list(self.connections) results = await asyncio.gather( @@ -158,3 +202,22 @@ async def get_service_replica_client( ) connection = await service_conn_pool.get_or_add(project, service, replica) return await connection.client() + + +async def maintain_service_connections( + service_conn_pool: ServiceConnectionPool, interval: float = TUNNEL_CHECK_INTERVAL +) -> None: + """Periodically reopen replica SSH tunnels whose ssh process has exited.""" + while True: + await asyncio.sleep(interval) + try: + await service_conn_pool.reopen_exited() + except Exception: + logger.exception("Failed to check service replica SSH tunnels") + + +def _remove_file(path: Path) -> None: + try: + path.unlink() + except FileNotFoundError: + pass diff --git a/src/tests/_internal/proxy/lib/services/__init__.py b/src/tests/_internal/proxy/lib/services/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/_internal/proxy/lib/services/test_service_connection.py b/src/tests/_internal/proxy/lib/services/test_service_connection.py new file mode 100644 index 0000000000..7175896310 --- /dev/null +++ b/src/tests/_internal/proxy/lib/services/test_service_connection.py @@ -0,0 +1,108 @@ +import asyncio +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from dstack._internal.proxy.lib.services.service_connection import ( + ServiceConnection, + ServiceConnectionPool, + maintain_service_connections, +) +from dstack._internal.proxy.lib.testing.common import make_project, make_service + + +def make_connection(alive: bool = True) -> ServiceConnection: + service = make_service("test-proj", "test-run") + connection = ServiceConnection(make_project("test-proj"), service, service.replicas[0]) + connection._tunnel.aopen = AsyncMock() + connection._tunnel.aclose = AsyncMock() + connection._tunnel.acheck = AsyncMock(return_value=alive) + return connection + + +@pytest.mark.asyncio +async def test_reopen_if_exited_keeps_live_tunnel() -> None: + connection = make_connection(alive=True) + await connection.open() + connection._tunnel.aopen.reset_mock() + + assert not await connection.reopen_if_exited() + connection._tunnel.aopen.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_reopen_if_exited_reopens_dead_tunnel_and_removes_stale_control_socket() -> None: + connection = make_connection(alive=False) + await connection.open() + connection._tunnel.aopen.reset_mock() + control_sock = Path(connection._tunnel.control_sock_path) + control_sock.touch() + + assert await connection.reopen_if_exited() + connection._tunnel.aopen.assert_awaited_once() + assert not control_sock.exists() + assert (await connection.client()) is connection._client + + +@pytest.mark.asyncio +async def test_reopen_if_exited_skips_unopened_and_closed_connections() -> None: + unopened = make_connection(alive=False) + assert not await unopened.reopen_if_exited() + unopened._tunnel.acheck.assert_not_awaited() + + closed = make_connection(alive=False) + await closed.open() + await closed.close() + closed._tunnel.aopen.reset_mock() + assert not await closed.reopen_if_exited() + closed._tunnel.aopen.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_reopen_if_exited_does_not_reopen_tunnel_closed_meanwhile() -> None: + connection = make_connection(alive=False) + await connection.open() + connection._tunnel.aopen.reset_mock() + + async def close_during_check() -> bool: + connection._closed = True + return False + + connection._tunnel.acheck = AsyncMock(side_effect=close_during_check) + assert not await connection.reopen_if_exited() + connection._tunnel.aopen.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_pool_reopen_exited_continues_after_failure() -> None: + pool = ServiceConnectionPool() + failing = make_connection(alive=False) + dead = make_connection(alive=False) + live = make_connection(alive=True) + for connection in (failing, dead, live): + await connection.open() + connection._tunnel.aopen.reset_mock() + failing._tunnel.aopen.side_effect = RuntimeError("ssh failed") + pool.connections = {"failing": failing, "dead": dead, "live": live} + + await pool.reopen_exited() + + failing._tunnel.aopen.assert_awaited_once() + dead._tunnel.aopen.assert_awaited_once() + live._tunnel.aopen.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_maintain_service_connections_checks_periodically() -> None: + pool = ServiceConnectionPool() + pool.reopen_exited = AsyncMock(side_effect=[RuntimeError("check failed"), None, None]) + task = asyncio.create_task(maintain_service_connections(pool, interval=0.01)) + for _ in range(100): + if pool.reopen_exited.await_count >= 3: + break + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert pool.reopen_exited.await_count >= 3