Reuse gateway model proxy HTTP clients per service

src/dstack/_internal/proxy/gateway/services/registry.py

@@ -129,6 +129,7 @@ async def unregister_service(
         await nginx.unregister(service)
         await repo.delete_models_by_run(project_name, run_name)
         await repo.delete_service(project_name, run_name)
+        await service_conn_pool.remove_service_client(project_name, run_name)
 
     logger.info("Service %s is unregistered now", service.fmt())
 

src/dstack/_internal/proxy/lib/services/service_connection.py

@@ -202,6 +202,8 @@ class ServiceConnectionPool:
     def __init__(self) -> None:
         # TODO(#2238): remove connections to stopped replicas in-server
         self.connections: Dict[str, ServiceConnection] = {}
+        self._service_clients: Dict[tuple[str, str], ServiceClient] = {}
+        self._service_clients_lock = asyncio.Lock()
         self._reconnect_semaphore = asyncio.Semaphore(MAX_CONCURRENT_TUNNEL_RECONNECTS)
 
     async def get(self, replica_id: str) -> Optional[ServiceConnection]:
@@ -244,6 +246,28 @@ async def remove(self, replica_id: str) -> None:
         if connection is not None:
             await connection.close()
 
+    async def get_or_add_service_client(self, service: Service) -> ServiceClient:
+        if service.domain is None:
+            raise UnexpectedProxyError(f"Expected a domain for service {service.fmt()}")
+        key = (service.project_name, service.run_name)
+        async with self._service_clients_lock:
+            client = self._service_clients.get(key)
+            if client is None or client.is_closed:
+                client = ServiceClient(
+                    base_url="http://127.0.0.1",
+                    headers={"Host": service.domain},
+                    timeout=service.read_timeout,
+                )
+                self._service_clients[key] = client
+            return client
+
+    async def remove_service_client(self, project_name: str, run_name: str) -> None:
+        key = (project_name, run_name)
+        async with self._service_clients_lock:
+            client = self._service_clients.pop(key, None)
+        if client is not None:
+            await client.aclose()
+
     async def remove_all(self) -> None:
         replica_ids = list(self.connections)
         results = await asyncio.gather(
@@ -254,6 +278,15 @@ async def remove_all(self) -> None:
                 logger.error(
                     "Error removing connection to service replica %s: %s", replica_ids[i], exc
                 )
+        async with self._service_clients_lock:
+            clients = list(self._service_clients.values())
+            self._service_clients.clear()
+        results = await asyncio.gather(
+            *(client.aclose() for client in clients), return_exceptions=True
+        )
+        for exc in results:
+            if isinstance(exc, Exception):
+                logger.error("Error closing HTTP client for a service: %s", exc)
 
 
 async def get_service_replica_client(
@@ -264,11 +297,7 @@ async def get_service_replica_client(
     """
     if service.domain is not None:
         # Forward to Nginx so that requests are visible to StatsCollector in the access log
-        return httpx.AsyncClient(
-            base_url="http://127.0.0.1",
-            headers={"Host": service.domain},
-            timeout=service.read_timeout,
-        )
+        return await service_conn_pool.get_or_add_service_client(service)
     # Nginx not available, forward directly to the tunnel
     replica = random.choice(service.replicas)
     connection = await service_conn_pool.get(replica.id)

src/tests/_internal/proxy/lib/services/test_service_connection.py

@@ -191,6 +191,31 @@ async def close_tunnel():
 
 @pytest.mark.asyncio
 class TestServiceConnectionPool:
+    async def test_gateway_service_client_is_reused_and_closed_on_service_removal(
+        self, service, pool
+    ):
+        clients = await asyncio.gather(
+            *(pool.get_or_add_service_client(service) for _ in range(5))
+        )
+
+        assert all(client is clients[0] for client in clients)
+        assert clients[0].headers["Host"] == service.domain
+        assert not clients[0].is_closed
+
+        await pool.remove_service_client(service.project_name, service.run_name)
+
+        assert clients[0].is_closed
+        replacement = await pool.get_or_add_service_client(service)
+        assert replacement is not clients[0]
+
+    async def test_remove_all_closes_gateway_service_clients(self, service, pool):
+        client = await pool.get_or_add_service_client(service)
+
+        await pool.remove_all()
+
+        assert client.is_closed
+        assert not pool._service_clients
+
     @pytest.mark.parametrize("gateway", [False, True])
     async def test_proxy_jump_key_update_preserves_client_socket_and_reconnects_with_new_key(
         self, project, service, pool, tunnel_factory, gateway