From c21b9e8de3d472299491c7dcc57081438a349aff Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Tue, 5 Mar 2024 12:28:53 +0100
Subject: [PATCH 01/17] Rework store and registry routes

---
 gateway/src/dstack/gateway/auth/routes.py     |   5 +-
 .../gateway/{services => core}/__init__.py    |   0
 .../dstack/gateway/{services => core}/auth.py |   0
 .../gateway/{services => core}/nginx.py       |   0
 .../gateway/{services => core}/persistent.py  |   0
 gateway/src/dstack/gateway/core/store.py      | 307 ++++++++++++++++++
 .../gateway/{services => core}/tunnel.py      |   0
 gateway/src/dstack/gateway/main.py            |  10 +-
 gateway/src/dstack/gateway/openai/routes.py   |   2 +-
 gateway/src/dstack/gateway/openai/store.py    |   4 +-
 gateway/src/dstack/gateway/registry/routes.py | 103 +++---
 .../src/dstack/gateway/registry/schemas.py    |  29 +-
 gateway/src/dstack/gateway/schemas.py         |  15 -
 gateway/src/dstack/gateway/services/store.py  | 178 ----------
 gateway/src/tests/services/test_store.py      |   4 +-
 15 files changed, 402 insertions(+), 255 deletions(-)
 rename gateway/src/dstack/gateway/{services => core}/__init__.py (100%)
 rename gateway/src/dstack/gateway/{services => core}/auth.py (100%)
 rename gateway/src/dstack/gateway/{services => core}/nginx.py (100%)
 rename gateway/src/dstack/gateway/{services => core}/persistent.py (100%)
 create mode 100644 gateway/src/dstack/gateway/core/store.py
 rename gateway/src/dstack/gateway/{services => core}/tunnel.py (100%)
 delete mode 100644 gateway/src/dstack/gateway/schemas.py
 delete mode 100644 gateway/src/dstack/gateway/services/store.py

diff --git a/gateway/src/dstack/gateway/auth/routes.py b/gateway/src/dstack/gateway/auth/routes.py
index e89b4c5fe1..5cd8e7bb63 100644
--- a/gateway/src/dstack/gateway/auth/routes.py
+++ b/gateway/src/dstack/gateway/auth/routes.py
@@ -1,11 +1,14 @@
 from fastapi import APIRouter, Depends, HTTPException, Security
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
-from dstack.gateway.services.auth import AuthProvider, get_auth
+from dstack.gateway.core.auth import AuthProvider, get_auth
 
 router = APIRouter()
 
 
+# TODO(egor-s): support Authorization header alternative for web browsers
+
+
 @router.get("/{project}")
 async def get_auth(
     project: str,
diff --git a/gateway/src/dstack/gateway/services/__init__.py b/gateway/src/dstack/gateway/core/__init__.py
similarity index 100%
rename from gateway/src/dstack/gateway/services/__init__.py
rename to gateway/src/dstack/gateway/core/__init__.py
diff --git a/gateway/src/dstack/gateway/services/auth.py b/gateway/src/dstack/gateway/core/auth.py
similarity index 100%
rename from gateway/src/dstack/gateway/services/auth.py
rename to gateway/src/dstack/gateway/core/auth.py
diff --git a/gateway/src/dstack/gateway/services/nginx.py b/gateway/src/dstack/gateway/core/nginx.py
similarity index 100%
rename from gateway/src/dstack/gateway/services/nginx.py
rename to gateway/src/dstack/gateway/core/nginx.py
diff --git a/gateway/src/dstack/gateway/services/persistent.py b/gateway/src/dstack/gateway/core/persistent.py
similarity index 100%
rename from gateway/src/dstack/gateway/services/persistent.py
rename to gateway/src/dstack/gateway/core/persistent.py
diff --git a/gateway/src/dstack/gateway/core/store.py b/gateway/src/dstack/gateway/core/store.py
new file mode 100644
index 0000000000..ccffb4d209
--- /dev/null
+++ b/gateway/src/dstack/gateway/core/store.py
@@ -0,0 +1,307 @@
+import asyncio
+import concurrent
+import functools
+import logging
+import os
+from abc import ABC, abstractmethod
+from asyncio import CancelledError, Lock
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import AsyncExitStack
+from functools import lru_cache
+from pathlib import Path
+from typing import DefaultDict, Dict, List, Optional, Set, Tuple
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+from dstack.gateway.common import run_async
+from dstack.gateway.core.nginx import Nginx
+from dstack.gateway.core.persistent import get_persistent_state
+from dstack.gateway.core.tunnel import SSHTunnel
+from dstack.gateway.errors import GatewayError
+
+logger = logging.getLogger(__name__)
+
+
+class Replica(BaseModel):
+    id: str
+    app_port: int
+    ssh_host: str
+    ssh_port: int
+    ssh_jump_host: Optional[str]
+    ssh_jump_port: Optional[int]
+    ssh_tunnel: Optional[SSHTunnel] = None
+
+
+class Service(BaseModel):
+    id: str
+    domain: str
+    auth: bool
+    options: dict
+    replicas: List[Replica] = []
+
+
+class Store(BaseModel):
+    """
+    Store is a central place to register and unregister services.
+    Other components can subscribe to updates.
+    Its internal state could be serialized to a file and restored from it using pydantic.
+    """
+
+    services: Dict[str, Service] = {}
+    projects: DefaultDict[str, Set[str]] = defaultdict(set)
+    entrypoints: Dict[str, Tuple[str, str]] = {}
+    nginx: Nginx = Field(default_factory=Nginx)
+    _lock: Lock = Lock()
+    _subscribers: List["StoreSubscriber"] = []
+    _ssh_keys_dir = PrivateAttr(
+        default_factory=lambda: Path("~/.ssh/projects").expanduser().resolve()
+    )
+
+    async def register_service(self, project: str, service: Service, ssh_private_key: str):
+        async with self._lock:
+            if service.id in self.services:
+                raise GatewayError(f"Service ID {service.id!r} is already registered")
+            if service.replicas:
+                raise GatewayError("Not implemented: replicas should be registered separately")
+
+            logger.debug("%s: registering service %s (%s)", project, service.id, service.domain)
+
+            # Save project SSH key
+            self._ssh_keys_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
+            ssh_key_path = self._ssh_keys_dir / project
+            if (
+                ssh_key_path.exists()
+                and ssh_key_path.read_text().strip() != ssh_private_key.strip()
+            ):
+                logger.warning(
+                    "%s: SSH key for service %s (%s) is different from the previous one",
+                    project,
+                    service.id,
+                    service.domain,
+                )
+            with open(
+                ssh_key_path, "w", opener=lambda path, flags: os.open(path, flags, 0o600)
+            ) as f:
+                f.write(ssh_private_key)
+
+            async with AsyncExitStack() as stack:
+                # Configure nginx and issue SSL cert
+                await self.nginx.register_service(
+                    project, service.id, service.domain, service.auth
+                )
+                stack.push_async_callback(
+                    supress_exc_async(self.nginx.unregister_domain, service.domain)
+                )
+
+                # Notify subscribers
+                for subscriber in self._subscribers:
+                    await subscriber.on_register(project, service)
+                    stack.push_async_callback(
+                        supress_exc_async(subscriber.on_unregister, project, service.id)
+                    )
+
+                # All fine, remove rollbacks
+                stack.pop_all()
+
+            self.services[service.id] = service
+            self.projects[project].add(service.id)
+
+        logger.info("%s: service %s (%s) is registered now", project, service.id, service.domain)
+
+    async def unregister_service(self, project: str, service_id: str):
+        async with self._lock:
+            if service_id not in self.projects[project]:
+                raise GatewayError(
+                    f"Service ID {service_id!r} is not registered in project {project!r}"
+                )
+            service = self.services[service_id]
+
+            logger.debug("%s: unregistering service %s (%s)", project, service_id, service.domain)
+
+            await asyncio.gather(
+                # Terminate all SSH tunnels
+                *(
+                    run_async(replica.ssh_tunnel.stop)
+                    for replica in service.replicas
+                    if replica.ssh_tunnel is not None
+                ),
+                # Unregister from nginx
+                self.nginx.unregister_domain(service.domain),
+                # Notify subscribers
+                *(
+                    subscriber.on_unregister(project, service.id)
+                    for subscriber in self._subscribers
+                ),
+                return_exceptions=True,
+            )
+
+            self.projects[project].remove(service_id)
+            self.services.pop(service_id)
+
+        logger.info("%s: service %s (%s) is unregistered now", project, service_id, service.domain)
+
+    async def register_replica(self, project: str, service_id: str, replica: Replica):
+        async with self._lock:
+            if service_id not in self.projects[project]:
+                raise GatewayError(
+                    f"Service ID {service_id!r} is not registered in project {project!r}"
+                )
+            if replica.ssh_tunnel:
+                raise GatewayError("Not implemented: replica should not have a tunnel yet")
+            service = self.services[service_id]
+
+            logger.debug(
+                "%s: registering replica %s for service %s (%s)",
+                project,
+                replica.id,
+                service_id,
+                service.domain,
+            )
+
+            with AsyncExitStack() as stack:
+                # Start SSH tunnel
+                ssh_tunnel = SSHTunnel.create(
+                    host=replica.ssh_host,
+                    port=replica.ssh_port,
+                    app_port=replica.app_port,
+                    id_rsa_path=(self._ssh_keys_dir / project).as_posix(),
+                    docker_host=replica.ssh_jump_host,
+                    docker_port=replica.ssh_jump_port,
+                )
+                await run_async(ssh_tunnel.start)
+                stack.push_async_callback(supress_exc_async(run_async, ssh_tunnel.stop))
+
+                # Add to nginx
+                await self.nginx.add_upstream(service.domain, replica.id, ssh_tunnel.sock_path)
+                stack.push_async_callback(
+                    supress_exc_async(self.nginx.remove_upstream, service.domain, replica.id)
+                )
+
+                # All fine, remove rollbacks
+                stack.pop_all()
+
+            replica.ssh_tunnel = ssh_tunnel
+            service.replicas.append(replica)
+
+        logger.info(
+            "%s: replica %s for service %s (%s) is registered now",
+            project,
+            replica.id,
+            service_id,
+            service.domain,
+        )
+
+    async def unregister_replica(self, project: str, service_id: str, replica_id: str):
+        async with self._lock:
+            if service_id not in self.projects[project]:
+                raise GatewayError(
+                    f"Service ID {service_id!r} is not registered in project {project!r}"
+                )
+            service = self.services[service_id]
+
+            for replica in service.replicas:
+                if replica.id == replica_id:
+                    break
+            else:
+                raise GatewayError(
+                    f"Replica ID {replica_id!r} is not registered in service {service_id!r}"
+                )
+
+            logger.debug(
+                "%s: unregistering replica %s for service %s (%s)",
+                project,
+                replica_id,
+                service_id,
+                service.domain,
+            )
+
+            await asyncio.gather(
+                # Terminate SSH tunnel
+                run_async(replica.ssh_tunnel.stop),
+                # Remove from nginx
+                self.nginx.remove_upstream(service.domain, replica.id),
+                return_exceptions=True,
+            )
+
+            service.replicas.remove(replica)
+
+        logger.info(
+            "%s: replica %s for service %s (%s) is unregistered now",
+            project,
+            replica_id,
+            service_id,
+            service.domain,
+        )
+
+    async def register_entrypoint(self, project: str, domain: str, module: str):
+        async with self._lock:
+            if domain in self.entrypoints:
+                if self.entrypoints[domain] == (project, module):
+                    return
+                raise GatewayError(
+                    f"Domain {domain} is already registered as {self.entrypoints[domain]}"
+                )
+
+            logger.debug("%s: registering entrypoint %s for module %s", project, domain, module)
+
+            await self.nginx.register_entrypoint(domain, f"/api/{module}/{project}")
+            self.entrypoints[domain] = (project, module)
+
+        logger.info("%s: entrypoint %s is now registered", project, domain)
+
+    async def subscribe(self, subscriber: "StoreSubscriber"):
+        async with self._lock:
+            self._subscribers.append(subscriber)
+
+    def start_tunnels(self):
+        with ThreadPoolExecutor(max_workers=8) as executor:
+            futures = [
+                executor.submit(supress_exc(replica.ssh_tunnel.start))
+                for service_id, service in self.services.items()
+                for replica in service.replicas
+                if replica.ssh_tunnel is not None
+            ]
+            concurrent.futures.wait(futures)
+
+
+class StoreSubscriber(ABC):
+    @abstractmethod
+    async def on_register(self, project: str, service: Service):
+        ...
+
+    @abstractmethod
+    async def on_unregister(self, project: str, service_id: str):
+        ...
+
+
+def supress_exc_async(func, *args, **kwargs):
+    @functools.wraps(func)
+    async def wrapper():
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            if isinstance(e, CancelledError):
+                raise
+
+    return wrapper
+
+
+def supress_exc(func, *args, **kwargs):
+    @functools.wraps(func)
+    def wrapper():
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            if isinstance(e, CancelledError):
+                raise
+
+    return wrapper
+
+
+@lru_cache()
+def get_store() -> Store:
+    # TODO(egor-s): keep going if failed to restore the state
+    store = Store.model_validate(get_persistent_state().get("store", {}))
+    store.start_tunnels()  # start tunnels after restoring the state
+    return store
diff --git a/gateway/src/dstack/gateway/services/tunnel.py b/gateway/src/dstack/gateway/core/tunnel.py
similarity index 100%
rename from gateway/src/dstack/gateway/services/tunnel.py
rename to gateway/src/dstack/gateway/core/tunnel.py
diff --git a/gateway/src/dstack/gateway/main.py b/gateway/src/dstack/gateway/main.py
index 626cf3ca88..9ab670ea37 100644
--- a/gateway/src/dstack/gateway/main.py
+++ b/gateway/src/dstack/gateway/main.py
@@ -7,11 +7,12 @@
 import dstack.gateway.openai.store as openai_store
 import dstack.gateway.version
 from dstack.gateway.auth.routes import router as auth_router
+from dstack.gateway.core.persistent import save_persistent_state
+from dstack.gateway.core.store import get_store
+from dstack.gateway.errors import GatewayError
 from dstack.gateway.logging import configure_logging
 from dstack.gateway.openai.routes import router as openai_router
 from dstack.gateway.registry.routes import router as registry_router
-from dstack.gateway.services.persistent import save_persistent_state
-from dstack.gateway.services.store import get_store
 
 
 @asynccontextmanager
@@ -43,3 +44,8 @@ async def lifespan(app: FastAPI):
 @app.get("/")
 def get_info():
     return {"version": dstack.gateway.version.__version__}
+
+
+@app.exception_handler(GatewayError)
+async def gateway_error_handler(request, exc: GatewayError):
+    return exc.http()
diff --git a/gateway/src/dstack/gateway/openai/routes.py b/gateway/src/dstack/gateway/openai/routes.py
index 14338de50c..02cb0a8789 100644
--- a/gateway/src/dstack/gateway/openai/routes.py
+++ b/gateway/src/dstack/gateway/openai/routes.py
@@ -4,6 +4,7 @@
 from fastapi.responses import StreamingResponse
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
+from dstack.gateway.core.auth import AuthProvider, get_auth
 from dstack.gateway.errors import GatewayError
 from dstack.gateway.openai.schemas import (
     ChatCompletionsChunk,
@@ -12,7 +13,6 @@
     ModelsResponse,
 )
 from dstack.gateway.openai.store import OpenAIStore, get_store
-from dstack.gateway.services.auth import AuthProvider, get_auth
 
 
 async def auth_required(
diff --git a/gateway/src/dstack/gateway/openai/store.py b/gateway/src/dstack/gateway/openai/store.py
index c364be3501..87b5875872 100644
--- a/gateway/src/dstack/gateway/openai/store.py
+++ b/gateway/src/dstack/gateway/openai/store.py
@@ -5,6 +5,8 @@
 
 from pydantic import BaseModel, ValidationError
 
+from dstack.gateway.core.persistent import get_persistent_state
+from dstack.gateway.core.store import StoreSubscriber
 from dstack.gateway.errors import GatewayError, NotFoundError
 from dstack.gateway.openai.clients import ChatCompletionsClient
 from dstack.gateway.openai.clients.openai import OpenAIChatCompletions
@@ -12,8 +14,6 @@
 from dstack.gateway.openai.models import OpenAIOptions, ServiceModel
 from dstack.gateway.openai.schemas import Model
 from dstack.gateway.schemas import Service
-from dstack.gateway.services.persistent import get_persistent_state
-from dstack.gateway.services.store import StoreSubscriber
 
 
 class OpenAIStore(BaseModel, StoreSubscriber):
diff --git a/gateway/src/dstack/gateway/registry/routes.py b/gateway/src/dstack/gateway/registry/routes.py
index 72a454ed78..cbffffc70e 100644
--- a/gateway/src/dstack/gateway/registry/routes.py
+++ b/gateway/src/dstack/gateway/registry/routes.py
@@ -2,60 +2,77 @@
 
 from fastapi import APIRouter, Depends
 
-from dstack.gateway.errors import GatewayError
+from dstack.gateway.core.store import Replica, Service, Store, get_store
 from dstack.gateway.registry.schemas import (
-    PreflightRequest,
+    OkResponse,
     RegisterEntrypointRequest,
-    RegisterRequest,
-    UnregisterRequest,
+    RegisterReplicaRequest,
+    RegisterServiceRequest,
 )
-from dstack.gateway.services.store import Store, get_store
 
-router = APIRouter()
+router = APIRouter(prefix="/{project}")
 
 
-@router.post("/{project}/register")
-async def post_register(
-    project: str, body: RegisterRequest, store: Annotated[Store, Depends(get_store)]
-):
-    try:
-        await store.register(project, body)
-    except GatewayError as e:
-        raise e.http()
-    return "ok"
+@router.post("/services/register")
+async def post_register_service(
+    project: str, body: RegisterServiceRequest, store: Annotated[Store, Depends(get_store)]
+) -> OkResponse:
+    await store.register_service(
+        project,
+        Service(
+            id=body.run_id,
+            domain=body.domain,
+            auth=body.auth,
+            options=body.options,
+        ),
+        body.ssh_private_key,
+    )
+    return OkResponse()
 
 
-@router.post("/{project}/unregister")
-async def post_unregister(
-    project: str, body: UnregisterRequest, store: Annotated[Store, Depends(get_store)]
-):
-    try:
-        await store.unregister(project, body.public_domain)
-    except GatewayError as e:
-        raise e.http()
-    return "ok"
+@router.post("/services/{run_id}/unregister")
+async def post_unregister_services(
+    project: str, run_id: str, store: Annotated[Store, Depends(get_store)]
+) -> OkResponse:
+    await store.unregister_service(project, run_id)
+    return OkResponse()
 
 
-@router.post("/{project}/{module}/register")
+@router.post("/services/{run_id}/replicas/register")
+async def post_register_replica(
+    project: str,
+    run_id: str,
+    body: RegisterReplicaRequest,
+    store: Annotated[Store, Depends(get_store)],
+) -> OkResponse:
+    await store.register_replica(
+        project,
+        run_id,
+        Replica(
+            id=body.job_id,
+            app_port=body.app_port,
+            ssh_host=body.ssh_host,
+            ssh_port=body.ssh_port,
+            ssh_jump_host=body.ssh_jump_host,
+            ssh_jump_port=body.ssh_jump_port,
+        ),
+    )
+    return OkResponse()
+
+
+@router.post("/services/{run_id}/replicas/{job_id}/unregister")
+async def post_unregister_replica(
+    project: str, run_id: str, job_id: str, store: Annotated[Store, Depends(get_store)]
+) -> OkResponse:
+    await store.unregister_replica(project, run_id, job_id)
+    return OkResponse()
+
+
+@router.post("/entrypoints/register")
 async def post_register_entrypoint(
     project: str,
-    module: str,
     body: RegisterEntrypointRequest,
     store: Annotated[Store, Depends(get_store)],
-):
-    try:
-        await store.register_entrypoint(project, body.domain, module)
-    except GatewayError as e:
-        raise e.http()
-    return "ok"
-
-
-@router.post("/{project}/preflight")
-async def post_preflight(
-    project: str, body: PreflightRequest, store: Annotated[Store, Depends(get_store)]
-):
-    try:
-        await store.preflight(project, body.public_domain, body.ssh_private_key)
-    except GatewayError as e:
-        raise e.http()
-    return "ok"
+) -> OkResponse:
+    await store.register_entrypoint(project, body.domain, body.module)
+    return OkResponse()
diff --git a/gateway/src/dstack/gateway/registry/schemas.py b/gateway/src/dstack/gateway/registry/schemas.py
index 9ff3a11aa1..ef564e9b9d 100644
--- a/gateway/src/dstack/gateway/registry/schemas.py
+++ b/gateway/src/dstack/gateway/registry/schemas.py
@@ -1,22 +1,29 @@
-from pydantic import BaseModel
+from typing import Literal, Optional
 
-import dstack.gateway.schemas
+from pydantic import BaseModel
 
 
-class RegisterRequest(dstack.gateway.schemas.Service):
-    pass  # TODO(egor-s): adapters and auth requirements
+class RegisterServiceRequest(BaseModel):
+    run_id: str
+    domain: str
+    auth: bool = True
+    options: dict = {}
+    ssh_private_key: str
 
 
-class UnregisterRequest(BaseModel):
-    public_domain: str
+class RegisterReplicaRequest(BaseModel):
+    job_id: str
+    app_port: int
+    ssh_host: str
+    ssh_port: int
+    ssh_jump_host: Optional[str] = None
+    ssh_jump_port: Optional[int] = None
 
 
 class RegisterEntrypointRequest(BaseModel):
+    module: Literal["openai"]
     domain: str
 
 
-class PreflightRequest(BaseModel):
-    public_domain: str
-    ssh_private_key: str
-
-    options: dict = {}
+class OkResponse(BaseModel):
+    status: str = "ok"
diff --git a/gateway/src/dstack/gateway/schemas.py b/gateway/src/dstack/gateway/schemas.py
deleted file mode 100644
index 8c2e43c343..0000000000
--- a/gateway/src/dstack/gateway/schemas.py
+++ /dev/null
@@ -1,15 +0,0 @@
-from typing import Optional
-
-from pydantic import BaseModel
-
-
-class Service(BaseModel):
-    public_domain: str  # only https & 443 port
-    app_port: int
-    ssh_host: str  # user@hostname
-    ssh_port: int
-    docker_ssh_host: Optional[str] = None
-    docker_ssh_port: Optional[int] = None
-
-    auth: bool = True
-    options: dict = {}
diff --git a/gateway/src/dstack/gateway/services/store.py b/gateway/src/dstack/gateway/services/store.py
deleted file mode 100644
index 03366abc94..0000000000
--- a/gateway/src/dstack/gateway/services/store.py
+++ /dev/null
@@ -1,178 +0,0 @@
-import asyncio
-import concurrent
-import functools
-import logging
-import os
-from abc import ABC, abstractmethod
-from asyncio import Lock
-from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor
-from contextlib import AsyncExitStack
-from functools import lru_cache
-from pathlib import Path
-from typing import DefaultDict, Dict, List, Set, Tuple
-
-from pydantic import BaseModel, Field, PrivateAttr
-
-from dstack.gateway.common import run_async
-from dstack.gateway.errors import GatewayError
-from dstack.gateway.schemas import Service
-from dstack.gateway.services.nginx import Nginx
-from dstack.gateway.services.persistent import get_persistent_state
-from dstack.gateway.services.tunnel import SSHTunnel
-
-logger = logging.getLogger(__name__)
-
-
-class Store(BaseModel):
-    """
-    Store is a central place to register and unregister services.
-    Other components can subscribe to updates.
-    Its internal state could be serialized to a file and restored from it using pydantic.
-    """
-
-    services: Dict[str, Tuple[Service, SSHTunnel]] = {}
-    projects: DefaultDict[str, Set[str]] = defaultdict(set)
-    entrypoints: Dict[str, Tuple[str, str]] = {}
-    nginx: Nginx = Field(default_factory=Nginx)
-    _lock: Lock = Lock()
-    _subscribers: List["StoreSubscriber"] = []
-    _ssh_keys_dir = PrivateAttr(
-        default_factory=lambda: Path("~/.ssh/projects").expanduser().resolve()
-    )
-
-    async def register(self, project: str, service: Service):
-        async with self._lock:
-            if service.public_domain in self.services:
-                raise GatewayError(f"Domain {service.public_domain} is already registered")
-            logger.info("%s: registering service %s", project, service.public_domain)
-
-            tunnel = SSHTunnel.create(
-                host=service.ssh_host,
-                port=service.ssh_port,
-                app_port=service.app_port,
-                id_rsa_path=(self._ssh_keys_dir / project).as_posix(),
-                docker_host=service.docker_ssh_host,
-                docker_port=service.docker_ssh_port,
-            )
-            async with AsyncExitStack() as stack:
-                await run_async(tunnel.start)
-                stack.push_async_callback(supress_exc_async(run_async, tunnel.stop))
-
-                await self.nginx.register_service(
-                    project, service.public_domain, tunnel.sock_path, auth=service.auth
-                )
-                stack.push_async_callback(
-                    supress_exc_async(self.nginx.unregister_domain, service.public_domain)
-                )
-
-                for subscriber in self._subscribers:
-                    await subscriber.on_register(project, service)
-                    stack.push_async_callback(
-                        supress_exc_async(subscriber.on_unregister, project, service.public_domain)
-                    )
-
-                stack.pop_all()  # no need to rollback
-            self.projects[project].add(service.public_domain)
-            self.services[service.public_domain] = (service, tunnel)
-
-    async def register_entrypoint(self, project: str, domain: str, module: str):
-        async with self._lock:
-            if domain in self.entrypoints:
-                if self.entrypoints[domain] == (project, module):
-                    return
-                raise GatewayError(
-                    f"Domain {domain} is already registered for {self.entrypoints[domain]}"
-                )
-
-            logger.info("%s: registering entrypoint %s", project, domain)
-            await self.nginx.register_entrypoint(domain, f"/api/{module}/{project}")
-            self.entrypoints[domain] = (project, module)
-
-    async def unregister(self, project: str, domain: str):
-        async with self._lock:
-            if domain not in self.services:
-                raise GatewayError(f"Domain {domain} is not registered")
-            if domain not in self.projects[project]:
-                raise GatewayError(f"Domain {domain} is not registered in project {project}")
-            logger.info("%s: unregistering service %s", project, domain)
-
-            self.projects[project].remove(domain)
-            service, tunnel = self.services.pop(domain)
-            await asyncio.gather(
-                run_async(tunnel.stop),
-                self.nginx.unregister_domain(domain),
-                *(subscriber.on_unregister(project, domain) for subscriber in self._subscribers),
-                return_exceptions=True,
-            )
-
-    async def subscribe(self, subscriber: "StoreSubscriber"):
-        async with self._lock:
-            self._subscribers.append(subscriber)
-
-    async def preflight(self, project: str, domain: str, ssh_private_key: str):
-        async with self._lock:
-            if domain in self.services:
-                raise GatewayError(f"Domain {domain} is already registered")
-            logger.info("%s: preflighting service %s", project, domain)
-
-            await run_async(self.nginx.run_certbot, domain)
-
-            self._ssh_keys_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
-            ssh_key_path = self._ssh_keys_dir / project
-            if (
-                ssh_key_path.exists()
-                and ssh_key_path.read_text().strip() != ssh_private_key.strip()
-            ):
-                logger.warning("%s: SSH key for project %s is different", domain, project)
-            with open(
-                ssh_key_path, "w", opener=lambda path, flags: os.open(path, flags, 0o600)
-            ) as f:
-                f.write(ssh_private_key)
-
-    def start_tunnels(self):
-        with ThreadPoolExecutor(max_workers=8) as executor:
-            futures = [
-                executor.submit(supress_exc(tunnel.start))
-                for domain, (service, tunnel) in self.services.items()
-            ]
-            concurrent.futures.wait(futures)
-
-
-class StoreSubscriber(ABC):
-    @abstractmethod
-    async def on_register(self, project: str, service: Service):
-        ...
-
-    @abstractmethod
-    async def on_unregister(self, project: str, domain: str):
-        ...
-
-
-def supress_exc_async(func, *args, **kwargs):
-    @functools.wraps(func)
-    async def wrapper():
-        try:
-            return await func(*args, **kwargs)
-        except Exception:
-            pass
-
-    return wrapper
-
-
-def supress_exc(func, *args, **kwargs):
-    @functools.wraps(func)
-    def wrapper():
-        try:
-            return func(*args, **kwargs)
-        except Exception:
-            pass
-
-    return wrapper
-
-
-@lru_cache()
-def get_store() -> Store:
-    store = Store.model_validate(get_persistent_state().get("store", {}))
-    store.start_tunnels()  # start tunnels after restoring the state
-    return store
diff --git a/gateway/src/tests/services/test_store.py b/gateway/src/tests/services/test_store.py
index e53bfbb89c..be7834ae47 100644
--- a/gateway/src/tests/services/test_store.py
+++ b/gateway/src/tests/services/test_store.py
@@ -3,10 +3,10 @@
 
 import pytest
 
+from dstack.gateway.core.nginx import Nginx
+from dstack.gateway.core.store import Store
 from dstack.gateway.openai.store import OpenAIStore
 from dstack.gateway.schemas import Service
-from dstack.gateway.services.nginx import Nginx
-from dstack.gateway.services.store import Store
 
 
 @pytest.fixture()

From c1764467b4c6c5dac293f5222aba7b387401bf94 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Tue, 5 Mar 2024 14:14:23 +0100
Subject: [PATCH 02/17] Refactor nginx service

---
 gateway/src/dstack/gateway/core/nginx.py      | 89 ++++++++++++++++---
 gateway/src/dstack/gateway/core/store.py      |  8 +-
 gateway/src/dstack/gateway/main.py            |  6 +-
 .../src/dstack/gateway/monitoring/__init__.py |  0
 .../src/dstack/gateway/monitoring/routes.py   | 13 +++
 .../gateway/resources/nginx/service.jinja2    |  3 +-
 6 files changed, 102 insertions(+), 17 deletions(-)
 create mode 100644 gateway/src/dstack/gateway/monitoring/__init__.py
 create mode 100644 gateway/src/dstack/gateway/monitoring/routes.py

diff --git a/gateway/src/dstack/gateway/core/nginx.py b/gateway/src/dstack/gateway/core/nginx.py
index 8156025dd7..1cc5c9f3b0 100644
--- a/gateway/src/dstack/gateway/core/nginx.py
+++ b/gateway/src/dstack/gateway/core/nginx.py
@@ -5,7 +5,7 @@
 import tempfile
 from asyncio import Lock
 from pathlib import Path
-from typing import Optional, Set
+from typing import Set
 
 import jinja2
 from pydantic import BaseModel
@@ -27,23 +27,32 @@ class Nginx(BaseModel):
     domains: Set[str] = set()
     _lock: Lock = Lock()
 
-    async def register_service(self, project: str, domain: str, sock_path: str, auth: bool):
-        logger.info("Registering service %s", domain)
+    async def register_service(
+        self, project: str, service_id: str, domain: str, auth: bool, fallback: str
+    ):
         async with self._lock:
             if domain in self.domains:
-                raise GatewayError("Domain is already registered")
+                raise GatewayError(f"Domain {domain} is already registered")
+
+            logger.debug("Registering service domain %s", domain)
+
+            await run_async(self.run_certbot, domain)
             self.write_conf(
-                self.get_service_conf(project, domain, f"unix:{sock_path}", auth),
+                self.get_service_conf(project, domain, auth, service_id, fallback),
                 CONFIGS_DIR / f"443-{domain}.conf",
             )
             self.domains.add(domain)
             await run_async(self.reload)
 
+        logger.info("Service domain %s is registered now", domain)
+
     async def register_entrypoint(self, domain: str, prefix: str):
-        logger.info("Registering entrypoint %s", domain)
         async with self._lock:
             if domain in self.domains:
-                raise GatewayError("Domain is already registered")
+                raise GatewayError(f"Domain {domain} is already registered")
+
+            logger.debug("Registering entrypoint domain %s", domain)
+
             await run_async(self.run_certbot, domain)
             self.write_conf(
                 self.get_entrypoint_conf(domain, prefix),
@@ -52,11 +61,15 @@ async def register_entrypoint(self, domain: str, prefix: str):
             self.domains.add(domain)
             await run_async(self.reload)
 
+        logger.info("Entrypoint domain %s is registered now", domain)
+
     async def unregister_domain(self, domain: str):
-        logger.info("Unregistering domain %s", domain)
         async with self._lock:
             if domain not in self.domains:
                 raise GatewayError("Domain is not registered")
+
+            logger.debug("Unregistering domain %s", domain)
+
             conf_path = CONFIGS_DIR / f"443-{domain}.conf"
             r = subprocess.run(["sudo", "rm", conf_path])
             if r.returncode != 0:
@@ -64,18 +77,54 @@ async def unregister_domain(self, domain: str):
             self.domains.remove(domain)
             await run_async(self.reload)
 
+        logger.info("Domain %s is unregistered now", domain)
+
+    async def add_upstream(self, domain: str, server: str, replica_id: str):
+        async with self._lock:
+            if domain not in self.domains:
+                raise GatewayError("Domain is not registered")
+
+            logger.debug("Adding upstream %s to domain %s", server, domain)
+
+            config_path = CONFIGS_DIR / f"443-{domain}.conf"
+            with open(config_path, "r") as f:
+                conf = f.read()
+
+            conf = self.add_upstream_to_conf(conf, server, replica_id)
+
+            self.write_conf(conf, config_path)
+            await run_async(self.reload)
+
+        logger.debug("Upstream %s is added to domain %s", server, domain)
+
+    async def remove_upstream(self, domain: str, replica_id: str):
+        async with self._lock:
+            if domain not in self.domains:
+                raise GatewayError("Domain is not registered")
+
+            logger.debug("Removing upstream %s from domain %s", replica_id, domain)
+
+            config_path = CONFIGS_DIR / f"443-{domain}.conf"
+            with open(config_path, "r") as f:
+                conf = f.read()
+
+            conf = self.remove_upstream_from_conf(conf, replica_id)
+
+            self.write_conf(conf, config_path)
+            await run_async(self.reload)
+
+        logger.debug("Upstream %s is removed from domain %s", replica_id, domain)
+
     @classmethod
     def get_service_conf(
-        cls, project: str, domain: str, server: str, auth: bool, upstream: Optional[str] = None
+        cls, project: str, domain: str, auth: bool, upstream: str, fallback: str
     ) -> str:
-        if upstream is None:
-            upstream = re.sub(r"[^a-z0-9_.\-]", "_", server, flags=re.IGNORECASE)
         template = importlib.resources.read_text(
             "dstack.gateway.resources.nginx", "service.jinja2"
         )
         return jinja2.Template(template).render(
             upstream=upstream,
-            server=server,
+            fallback=fallback,
             domain=domain,
             auth=auth,
             port=GATEWAY_PORT,
@@ -93,6 +142,22 @@ def get_entrypoint_conf(cls, domain: str, prefix: str) -> str:
             prefix=prefix,
         )
 
+    @classmethod
+    def add_upstream_to_conf(cls, conf: str, server: str, replica_id: str) -> str:
+        return re.sub(
+            r"(upstream +[^ ]+ *\{)",
+            f"\\1\n  server {server}; # REPLICA:{replica_id}",
+            conf,
+        )
+
+    @classmethod
+    def remove_upstream_from_conf(cls, conf: str, replica_id: str) -> str:
+        return re.sub(
+            f" *server [^;]+; # REPLICA:{replica_id}\n",
+            "",
+            conf,
+        )
+
     @staticmethod
     def reload():
         cmd = ["sudo", "systemctl", "reload", "nginx.service"]
diff --git a/gateway/src/dstack/gateway/core/store.py b/gateway/src/dstack/gateway/core/store.py
index ccffb4d209..6c8afd4822 100644
--- a/gateway/src/dstack/gateway/core/store.py
+++ b/gateway/src/dstack/gateway/core/store.py
@@ -88,7 +88,11 @@ async def register_service(self, project: str, service: Service, ssh_private_key
             async with AsyncExitStack() as stack:
                 # Configure nginx and issue SSL cert
                 await self.nginx.register_service(
-                    project, service.id, service.domain, service.auth
+                    project,
+                    service.id,
+                    service.domain,
+                    service.auth,
+                    fallback=f"/api/monitoring/{project}/services/{service.id}/fallback",
                 )
                 stack.push_async_callback(
                     supress_exc_async(self.nginx.unregister_domain, service.domain)
@@ -173,7 +177,7 @@ async def register_replica(self, project: str, service_id: str, replica: Replica
                 stack.push_async_callback(supress_exc_async(run_async, ssh_tunnel.stop))
 
                 # Add to nginx
-                await self.nginx.add_upstream(service.domain, replica.id, ssh_tunnel.sock_path)
+                await self.nginx.add_upstream(service.domain, ssh_tunnel.sock_path, replica.id)
                 stack.push_async_callback(
                     supress_exc_async(self.nginx.remove_upstream, service.domain, replica.id)
                 )
diff --git a/gateway/src/dstack/gateway/main.py b/gateway/src/dstack/gateway/main.py
index 9ab670ea37..3eb7935f24 100644
--- a/gateway/src/dstack/gateway/main.py
+++ b/gateway/src/dstack/gateway/main.py
@@ -11,6 +11,7 @@
 from dstack.gateway.core.store import get_store
 from dstack.gateway.errors import GatewayError
 from dstack.gateway.logging import configure_logging
+from dstack.gateway.monitoring.routes import router as monitoring_router
 from dstack.gateway.openai.routes import router as openai_router
 from dstack.gateway.registry.routes import router as registry_router
 
@@ -36,9 +37,10 @@ async def lifespan(app: FastAPI):
 
 configure_logging(logging.DEBUG)
 app = FastAPI(lifespan=lifespan)
-app.include_router(registry_router, prefix="/api/registry")
-app.include_router(openai_router, prefix="/api/openai")
 app.include_router(auth_router, prefix="/auth")
+app.include_router(monitoring_router, prefix="/api/monitoring")
+app.include_router(openai_router, prefix="/api/openai")
+app.include_router(registry_router, prefix="/api/registry")
 
 
 @app.get("/")
diff --git a/gateway/src/dstack/gateway/monitoring/__init__.py b/gateway/src/dstack/gateway/monitoring/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/gateway/src/dstack/gateway/monitoring/routes.py b/gateway/src/dstack/gateway/monitoring/routes.py
new file mode 100644
index 0000000000..019a1c6907
--- /dev/null
+++ b/gateway/src/dstack/gateway/monitoring/routes.py
@@ -0,0 +1,13 @@
+from fastapi import APIRouter
+from starlette.responses import JSONResponse
+
+router = APIRouter()
+
+
+@router.api_route("/{project}/services/{run_id}/fallback")
+async def service_upstream_fallback(project: str, run_id: str):
+    # TODO(egor-s): store event
+    return JSONResponse(
+        status_code=503,
+        content={"message": "Service temporarily unavailable, try again later"},
+    )
diff --git a/gateway/src/dstack/gateway/resources/nginx/service.jinja2 b/gateway/src/dstack/gateway/resources/nginx/service.jinja2
index 33b22e616f..f933c94a9b 100644
--- a/gateway/src/dstack/gateway/resources/nginx/service.jinja2
+++ b/gateway/src/dstack/gateway/resources/nginx/service.jinja2
@@ -1,5 +1,6 @@
 upstream {{ upstream }} {
-  server {{ server }};
+
+  server http://localhost:{{ port }}/{{ fallback.strip('/') }} backup;
 }
 server {
   server_name {{ domain }};

From a1e560e86c209da4950efdc0467710d76a6112a8 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Tue, 5 Mar 2024 14:46:09 +0100
Subject: [PATCH 03/17] Refactor SSH tunnel, add nginx tests

---
 gateway/src/dstack/gateway/core/store.py      | 20 +++++++-----
 gateway/src/dstack/gateway/core/tunnel.py     | 15 ++++-----
 .../src/tests/{services => core}/__init__.py  |  0
 gateway/src/tests/core/test_nginx.py          | 31 +++++++++++++++++++
 .../tests/{services => core}/test_store.py    |  0
 5 files changed, 50 insertions(+), 16 deletions(-)
 rename gateway/src/tests/{services => core}/__init__.py (100%)
 create mode 100644 gateway/src/tests/core/test_nginx.py
 rename gateway/src/tests/{services => core}/test_store.py (100%)

diff --git a/gateway/src/dstack/gateway/core/store.py b/gateway/src/dstack/gateway/core/store.py
index 6c8afd4822..568890d8ce 100644
--- a/gateway/src/dstack/gateway/core/store.py
+++ b/gateway/src/dstack/gateway/core/store.py
@@ -12,7 +12,7 @@
 from pathlib import Path
 from typing import DefaultDict, Dict, List, Optional, Set, Tuple
 
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, Field, PrivateAttr, ValidationError
 
 from dstack.gateway.common import run_async
 from dstack.gateway.core.nginx import Nginx
@@ -170,14 +170,16 @@ async def register_replica(self, project: str, service_id: str, replica: Replica
                     port=replica.ssh_port,
                     app_port=replica.app_port,
                     id_rsa_path=(self._ssh_keys_dir / project).as_posix(),
-                    docker_host=replica.ssh_jump_host,
-                    docker_port=replica.ssh_jump_port,
+                    jump_host=replica.ssh_jump_host,
+                    jump_port=replica.ssh_jump_port,
                 )
                 await run_async(ssh_tunnel.start)
                 stack.push_async_callback(supress_exc_async(run_async, ssh_tunnel.stop))
 
                 # Add to nginx
-                await self.nginx.add_upstream(service.domain, ssh_tunnel.sock_path, replica.id)
+                await self.nginx.add_upstream(
+                    service.domain, f"unix:{ssh_tunnel.sock_path}", replica.id
+                )
                 stack.push_async_callback(
                     supress_exc_async(self.nginx.remove_upstream, service.domain, replica.id)
                 )
@@ -305,7 +307,11 @@ def wrapper():
 
 @lru_cache()
 def get_store() -> Store:
-    # TODO(egor-s): keep going if failed to restore the state
-    store = Store.model_validate(get_persistent_state().get("store", {}))
-    store.start_tunnels()  # start tunnels after restoring the state
+    try:
+        store = Store.model_validate(get_persistent_state().get("store", {}))
+    except ValidationError as e:
+        logger.warning("Failed to load store state: %s", e)
+        store = Store()
+    # Start tunnels after restoring the state
+    store.start_tunnels()
     return store
diff --git a/gateway/src/dstack/gateway/core/tunnel.py b/gateway/src/dstack/gateway/core/tunnel.py
index a760a37f27..91207e4509 100644
--- a/gateway/src/dstack/gateway/core/tunnel.py
+++ b/gateway/src/dstack/gateway/core/tunnel.py
@@ -32,8 +32,8 @@ def create(
         app_port: int,
         *,
         id_rsa_path: str = "~/.ssh/id_rsa",
-        docker_host: Optional[str] = None,
-        docker_port: Optional[int] = None,
+        jump_host: Optional[str] = None,
+        jump_port: Optional[int] = None,
     ) -> "SSHTunnel":
         temp_dir = tempfile.mkdtemp()
         os.chmod(temp_dir, 0o755)  # grant any user read access
@@ -44,17 +44,13 @@ def create(
         cmd += ["-o", "StreamLocalBindMask=0111", "-o", "StreamLocalBindUnlink=yes"]
         cmd += ["-o", "ServerAliveInterval=60"]
         cmd += ["-f", "-N", "-L", f"{_sock_path(temp_dir)}:localhost:{app_port}"]
-        if docker_host is not None:
+        if jump_host is not None:
             # use `host` as a jump host
             proxy = ["ssh", "-F", "none", "-i", id_rsa_path, "-W", "%h:%p"]
             proxy += ["-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null"]
-            proxy += ["-p", str(port), host]
+            proxy += ["-p", str(jump_port), jump_host]
             cmd += ["-o", f"ProxyCommand={shlex.join(proxy)}"]
-            # connect to `docker_host`
-            cmd += ["-p", str(docker_port), docker_host]
-        else:
-            # connect to `host` directly
-            cmd += ["-p", str(port), host]
+        cmd += ["-p", str(port), host]
 
         start_cmd = cmd
         exit_cmd = ["ssh", "-S", control_path, "-O", "exit"]
@@ -65,6 +61,7 @@ def create(
     def sock_path(self):
         return _sock_path(self.temp_dir)
 
+    # TODO(egor-s): make it async
     def start(self):
         logger.info("Starting SSH tunnel for %s", self.sock_path)
         logger.debug("Executing %s", shlex.join(self.start_cmd))
diff --git a/gateway/src/tests/services/__init__.py b/gateway/src/tests/core/__init__.py
similarity index 100%
rename from gateway/src/tests/services/__init__.py
rename to gateway/src/tests/core/__init__.py
diff --git a/gateway/src/tests/core/test_nginx.py b/gateway/src/tests/core/test_nginx.py
new file mode 100644
index 0000000000..373b2b6d89
--- /dev/null
+++ b/gateway/src/tests/core/test_nginx.py
@@ -0,0 +1,31 @@
+from dstack.gateway.core.nginx import Nginx
+
+
+class TestNginxConf:
+    def test_add_upstream(self):
+        conf = "upstream hex123 {\n" "  server http://fallback backup;\n" "}"
+        expected = (
+            "upstream hex123 {\n"
+            "  server unix:/tmp/job1.sock; # REPLICA:12345\n"
+            "  server http://fallback backup;\n"
+            "}"
+        )
+        conf_out = Nginx.add_upstream_to_conf(conf, "unix:/tmp/job1.sock", "12345")
+        assert conf_out == expected
+
+    def test_remove_upstream(self):
+        conf = (
+            "upstream hex123 {\n"
+            "  server unix:/tmp/job1.sock; # REPLICA:12345\n"
+            "  server unix:/tmp/job2.sock; # REPLICA:67890\n"
+            "  server http://fallback backup;\n"
+            "}"
+        )
+        expected = (
+            "upstream hex123 {\n"
+            "  server unix:/tmp/job1.sock; # REPLICA:12345\n"
+            "  server http://fallback backup;\n"
+            "}"
+        )
+        conf_out = Nginx.remove_upstream_from_conf(conf, "67890")
+        assert conf_out == expected
diff --git a/gateway/src/tests/services/test_store.py b/gateway/src/tests/core/test_store.py
similarity index 100%
rename from gateway/src/tests/services/test_store.py
rename to gateway/src/tests/core/test_store.py

From 99f779fe4d106b69d0a339e2685a1c956719c0f3 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Tue, 5 Mar 2024 15:09:28 +0100
Subject: [PATCH 04/17] Fix openai store, update readme

---
 gateway/README.md                          | 40 +++++++++++++---------
 gateway/src/dstack/gateway/openai/store.py | 21 ++++++------
 2 files changed, 35 insertions(+), 26 deletions(-)

diff --git a/gateway/README.md b/gateway/README.md
index 87b35e4e71..7da860767f 100644
--- a/gateway/README.md
+++ b/gateway/README.md
@@ -10,20 +10,28 @@
 
 ## Development
 
-1. Build the wheel:
-    ```
-    python -m build .
-    ```
-2. Upload the wheel:
-    ```shell
-    scp dist/dstack_gateway-0.0.0-py3-none-any.whl ubuntu@${GATEWAY}:/tmp/
-    ```
-3. Install the wheel:
-    ```
-    ssh ubuntu@${GATEWAY} "pip install --force-reinstall /tmp/dstack_gateway-0.0.0-py3-none-any.whl"
-    ```
-4. Run the tunnel and the gateway:
-    ```
-    ssh -L 9001:localhost:8000 -t ubuntu@${GATEWAY} "uvicorn dstack.gateway.main:app"
-    ```
+1. Provision a gateway through dstack:
+   ```shell
+   dstack gateway create --backend aws --region us-east-1 --domain my.wildcard.domain.com
+   ```
+2. Extract the project key from the sqlite to the file
+3. Build gateway locally and deploy it:
+   ```shell
+   HOST=ubuntu@x.my.wildcard.domain.com
+   ID_RSA=/path/to/the/project/key
+   WHEEL=dstack_gateway-0.0.0-py3-none-any.whl
+   
+   python -m build .
+   scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -i "${ID_RSA}" "./dist/${WHEEL}" "${HOST}":/tmp/
+   ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -i "${ID_RSA}" "${HOST}" "/bin/sh /home/ubuntu/dstack/update.sh /tmp/${WHEEL} dev"
+   ```
+4. Open SSH tunnel to the gateway:
+   ```shell
+   ssh -L 9001:localhost:8000 -i "${ID_RSA}" "${HOST}"
+   ```
 5. Visit the gateway docs page at http://localhost:9001/docs
+
+To follow logs, use the command:
+```shell
+journalctl -u dstack.gateway.service -f
+```
diff --git a/gateway/src/dstack/gateway/openai/store.py b/gateway/src/dstack/gateway/openai/store.py
index 87b5875872..330d466b8b 100644
--- a/gateway/src/dstack/gateway/openai/store.py
+++ b/gateway/src/dstack/gateway/openai/store.py
@@ -6,14 +6,13 @@
 from pydantic import BaseModel, ValidationError
 
 from dstack.gateway.core.persistent import get_persistent_state
-from dstack.gateway.core.store import StoreSubscriber
+from dstack.gateway.core.store import Service, StoreSubscriber
 from dstack.gateway.errors import GatewayError, NotFoundError
 from dstack.gateway.openai.clients import ChatCompletionsClient
 from dstack.gateway.openai.clients.openai import OpenAIChatCompletions
 from dstack.gateway.openai.clients.tgi import TGIChatCompletions
 from dstack.gateway.openai.models import OpenAIOptions, ServiceModel
 from dstack.gateway.openai.schemas import Model
-from dstack.gateway.schemas import Service
 
 
 class OpenAIStore(BaseModel, StoreSubscriber):
@@ -23,7 +22,7 @@ class OpenAIStore(BaseModel, StoreSubscriber):
     """
 
     index: Dict[str, Dict[str, Dict[str, ServiceModel]]] = {}
-    domain_index: Dict[str, Tuple[str, str, str]] = {}
+    services_index: Dict[str, Tuple[str, str, str]] = {}
     _lock: asyncio.Lock = asyncio.Lock()
 
     async def on_register(self, project: str, service: Service):
@@ -41,18 +40,20 @@ async def on_register(self, project: str, service: Service):
                 self.index[project][model.type] = {}
             self.index[project][model.type][model.name] = ServiceModel(
                 model=model,
-                domain=service.public_domain,
+                domain=service.domain,
                 created=int(datetime.datetime.utcnow().timestamp()),
             )
-            self.domain_index[service.public_domain] = (project, model.type, model.name)
+            self.services_index[service.id] = (project, model.type, model.name)
 
-    async def on_unregister(self, project: str, domain: str):
+    async def on_unregister(self, project: str, service_id: str):
         async with self._lock:
-            if domain not in self.domain_index or self.domain_index[domain][0] != project:
+            if (
+                service_id not in self.services_index
+                or self.services_index[service_id][0] != project
+            ):
                 return
-            project, model_type, model_name = self.domain_index[domain]
-            del self.domain_index[domain]
-            del self.index[project][model_type][model_name]
+            project, model_type, model_name = self.services_index.pop(service_id)
+            self.index[project][model_type].pop(model_name)
 
     async def list_models(self, project: str) -> List[Model]:
         models = []

From 9c3a85ff5d5944ec176d494a20493c82fb4db803 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Tue, 5 Mar 2024 17:22:32 +0100
Subject: [PATCH 05/17] Start client refactoring, nginx config is broken

---
 gateway/src/dstack/gateway/core/nginx.py      |  1 +
 gateway/src/dstack/gateway/errors.py          |  9 +--
 gateway/src/dstack/gateway/main.py            |  2 +-
 gateway/src/dstack/gateway/openai/routes.py   | 22 +++-----
 .../gateway/resources/nginx/service.jinja2    |  2 +-
 .../server/services/gateways/__init__.py      |  2 +-
 .../server/services/gateways/client.py        | 55 +++++++++----------
 7 files changed, 44 insertions(+), 49 deletions(-)

diff --git a/gateway/src/dstack/gateway/core/nginx.py b/gateway/src/dstack/gateway/core/nginx.py
index 1cc5c9f3b0..f01ae5cc95 100644
--- a/gateway/src/dstack/gateway/core/nginx.py
+++ b/gateway/src/dstack/gateway/core/nginx.py
@@ -164,6 +164,7 @@ def reload():
         r = subprocess.run(cmd)
         if r.returncode != 0:
             raise GatewayError("Failed to reload nginx")
+        # TODO(egor-s) implement nginx config rollback
 
     @classmethod
     def write_conf(cls, conf: str, conf_path: Path):
diff --git a/gateway/src/dstack/gateway/errors.py b/gateway/src/dstack/gateway/errors.py
index ff3f481205..d5c394f392 100644
--- a/gateway/src/dstack/gateway/errors.py
+++ b/gateway/src/dstack/gateway/errors.py
@@ -1,11 +1,12 @@
 from fastapi import HTTPException
+from starlette.responses import JSONResponse
 
 
 class GatewayError(Exception):
-    def http(self, code: int = 400, **kwargs) -> HTTPException:
-        return HTTPException(
-            code,
-            {
+    def to_response(self, code: int = 400, **kwargs) -> JSONResponse:
+        return JSONResponse(
+            status_code=code,
+            content={
                 "error": self.__class__.__name__,
                 "message": str(self),
                 **kwargs,
diff --git a/gateway/src/dstack/gateway/main.py b/gateway/src/dstack/gateway/main.py
index 3eb7935f24..b903e51b9b 100644
--- a/gateway/src/dstack/gateway/main.py
+++ b/gateway/src/dstack/gateway/main.py
@@ -50,4 +50,4 @@ def get_info():
 
 @app.exception_handler(GatewayError)
 async def gateway_error_handler(request, exc: GatewayError):
-    return exc.http()
+    return exc.to_response()
diff --git a/gateway/src/dstack/gateway/openai/routes.py b/gateway/src/dstack/gateway/openai/routes.py
index 02cb0a8789..5674ed2246 100644
--- a/gateway/src/dstack/gateway/openai/routes.py
+++ b/gateway/src/dstack/gateway/openai/routes.py
@@ -5,7 +5,6 @@
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 from dstack.gateway.core.auth import AuthProvider, get_auth
-from dstack.gateway.errors import GatewayError
 from dstack.gateway.openai.schemas import (
     ChatCompletionsChunk,
     ChatCompletionsRequest,
@@ -38,18 +37,15 @@ async def get_models(
 async def post_chat_completions(
     project: str, body: ChatCompletionsRequest, store: Annotated[OpenAIStore, Depends(get_store)]
 ):
-    try:
-        client = await store.get_chat_client(project, body.model)
-        if not body.stream:
-            return await client.generate(body)
-        else:
-            return StreamingResponse(
-                stream_chunks(client.stream(body)),
-                media_type="text/event-stream",
-                headers={"X-Accel-Buffering": "no"},
-            )
-    except GatewayError as e:
-        raise e.http()
+    client = await store.get_chat_client(project, body.model)
+    if not body.stream:
+        return await client.generate(body)
+    else:
+        return StreamingResponse(
+            stream_chunks(client.stream(body)),
+            media_type="text/event-stream",
+            headers={"X-Accel-Buffering": "no"},
+        )
 
 
 async def stream_chunks(chunks: AsyncIterator[ChatCompletionsChunk]) -> AsyncIterator[bytes]:
diff --git a/gateway/src/dstack/gateway/resources/nginx/service.jinja2 b/gateway/src/dstack/gateway/resources/nginx/service.jinja2
index f933c94a9b..c742f8467f 100644
--- a/gateway/src/dstack/gateway/resources/nginx/service.jinja2
+++ b/gateway/src/dstack/gateway/resources/nginx/service.jinja2
@@ -1,6 +1,6 @@
 upstream {{ upstream }} {
 
-  server http://localhost:{{ port }}/{{ fallback.strip('/') }} backup;
+{#  server localhost:{{ port }}/{{ fallback.strip('/') }} backup;#}
 }
 server {
   server_name {{ domain }};
diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py
index 4f7f4713bf..e60dc9e5b8 100644
--- a/src/dstack/_internal/server/services/gateways/__init__.py
+++ b/src/dstack/_internal/server/services/gateways/__init__.py
@@ -334,7 +334,7 @@ async def register_replica(
 
     try:
         await run_async(
-            conn.client.register_service,
+            conn.client.register_replica,
             run,
             job_provisioning_data,
         )
diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py
index 786ef16251..8cfe2bc546 100644
--- a/src/dstack/_internal/server/services/gateways/client.py
+++ b/src/dstack/_internal/server/services/gateways/client.py
@@ -1,10 +1,11 @@
+import uuid
 from typing import Optional
 from urllib.parse import urlparse
 
 import httpx
 
 from dstack._internal.core.errors import GatewayError
-from dstack._internal.core.models.runs import JobProvisioningData, Run
+from dstack._internal.core.models.runs import Run
 
 GATEWAY_MANAGEMENT_PORT = 8000
 
@@ -19,46 +20,43 @@ def __init__(self, uds: Optional[str] = None, port: Optional[int] = None):
         self.base_url = "http://gateway" if uds else f"http://localhost:{port}"
         self.s = httpx.Client(transport=httpx.HTTPTransport(uds=uds) if uds else None, timeout=30)
 
-    def register_service(self, run: Run, job_provisioning_data: JobProvisioningData):
-        conf = run.run_spec.configuration
+    def register_service(self, run: Run, ssh_private_key: str):
         payload = {
-            "public_domain": urlparse(run.service.url).hostname,
-            "app_port": conf.port.container_port,
-            "auth": conf.auth,
+            "run_id": run.id.hex,
+            "domain": urlparse(run.service.url).hostname,
+            "auth": run.run_spec.configuration.auth,
             "options": run.service.options,
+            "ssh_private_key": ssh_private_key,
         }
-        ssh_proxy = job_provisioning_data.ssh_proxy
-        if ssh_proxy is None:
-            payload[
-                "ssh_host"
-            ] = f"{job_provisioning_data.username}@{job_provisioning_data.hostname}"
-            payload["ssh_port"] = job_provisioning_data.ssh_port
-            if job_provisioning_data.dockerized:
-                payload["docker_ssh_host"] = "root@localhost"
-                payload["docker_ssh_port"] = 10022
-        else:
-            payload["ssh_host"] = f"{ssh_proxy.username}@{ssh_proxy.hostname}"
-            payload["ssh_port"] = ssh_proxy.port
-            payload[
-                "docker_ssh_host"
-            ] = f"{job_provisioning_data.username}@{job_provisioning_data.hostname}"
-            payload["docker_ssh_port"] = job_provisioning_data.ssh_port
-        resp = self.s.post(self._url(f"/api/registry/{run.project_name}/register"), json=payload)
+        resp = self.s.post(
+            self._url(f"/api/registry/{run.project_name}/services/register"), json=payload
+        )
         if resp.status_code == 400:
             raise gateway_error(resp.json())
         resp.raise_for_status()
 
-    def register_openai_entrypoint(self, project: str, domain: str):
+    def unregister_service(self, project: str, run_id: uuid.UUID):
+        resp = self.s.post(self._url(f"/api/registry/{project}/services/{run_id.hex}/unregister"))
+        if resp.status_code == 400:
+            raise gateway_error(resp.json())
+        resp.raise_for_status()
+
+    def register_replica(self):
+        pass  # TODO(egor-s): implement
+
+    def unregister_replica(self, project: str, run_id: uuid.UUID, job_id: uuid.UUID):
         resp = self.s.post(
-            self._url(f"/api/registry/{project}/openai/register"), json={"domain": domain}
+            self._url(
+                f"/api/registry/{project}/services/{run_id.hex}/replicas/{job_id.hex}/unregister"
+            )
         )
         if resp.status_code == 400:
             raise gateway_error(resp.json())
         resp.raise_for_status()
 
-    def unregister_service(self, project: str, public_domain: str):
+    def register_openai_entrypoint(self, project: str, domain: str):
         resp = self.s.post(
-            self._url(f"/api/registry/{project}/unregister"), json={"public_domain": public_domain}
+            self._url(f"/api/registry/{project}/openai/register"), json={"domain": domain}
         )
         if resp.status_code == 400:
             raise gateway_error(resp.json())
@@ -89,5 +87,4 @@ def _url(self, path: str) -> str:
 
 
 def gateway_error(data: dict) -> GatewayError:
-    detail = data["detail"]
-    return GatewayError(msg=f"{detail['error']}: {detail['message']}")
+    return GatewayError(msg=f"{data['error']}: {data['message']}")

From 8ea12f7765e824583c54d979825ac074e3f05d78 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Wed, 6 Mar 2024 10:27:25 +0100
Subject: [PATCH 06/17] Drop nginx fallback

---
 gateway/src/dstack/gateway/core/nginx.py            | 11 +++--------
 gateway/src/dstack/gateway/core/store.py            |  1 -
 gateway/src/dstack/gateway/main.py                  |  2 --
 gateway/src/dstack/gateway/monitoring/__init__.py   |  0
 gateway/src/dstack/gateway/monitoring/routes.py     | 13 -------------
 .../dstack/gateway/resources/nginx/service.jinja2   |  1 -
 6 files changed, 3 insertions(+), 25 deletions(-)
 delete mode 100644 gateway/src/dstack/gateway/monitoring/__init__.py
 delete mode 100644 gateway/src/dstack/gateway/monitoring/routes.py

diff --git a/gateway/src/dstack/gateway/core/nginx.py b/gateway/src/dstack/gateway/core/nginx.py
index f01ae5cc95..a1c33badd2 100644
--- a/gateway/src/dstack/gateway/core/nginx.py
+++ b/gateway/src/dstack/gateway/core/nginx.py
@@ -27,9 +27,7 @@ class Nginx(BaseModel):
     domains: Set[str] = set()
     _lock: Lock = Lock()
 
-    async def register_service(
-        self, project: str, service_id: str, domain: str, auth: bool, fallback: str
-    ):
+    async def register_service(self, project: str, service_id: str, domain: str, auth: bool):
         async with self._lock:
             if domain in self.domains:
                 raise GatewayError(f"Domain {domain} is already registered")
@@ -38,7 +36,7 @@ async def register_service(
 
             await run_async(self.run_certbot, domain)
             self.write_conf(
-                self.get_service_conf(project, domain, auth, service_id, fallback),
+                self.get_service_conf(project, domain, auth, service_id),
                 CONFIGS_DIR / f"443-{domain}.conf",
             )
             self.domains.add(domain)
@@ -116,15 +114,12 @@ async def remove_upstream(self, domain: str, replica_id: str):
         logger.debug("Upstream %s is removed from domain %s", replica_id, domain)
 
     @classmethod
-    def get_service_conf(
-        cls, project: str, domain: str, auth: bool, upstream: str, fallback: str
-    ) -> str:
+    def get_service_conf(cls, project: str, domain: str, auth: bool, upstream: str) -> str:
         template = importlib.resources.read_text(
             "dstack.gateway.resources.nginx", "service.jinja2"
         )
         return jinja2.Template(template).render(
             upstream=upstream,
-            fallback=fallback,
             domain=domain,
             auth=auth,
             port=GATEWAY_PORT,
diff --git a/gateway/src/dstack/gateway/core/store.py b/gateway/src/dstack/gateway/core/store.py
index 568890d8ce..5b37a2969d 100644
--- a/gateway/src/dstack/gateway/core/store.py
+++ b/gateway/src/dstack/gateway/core/store.py
@@ -92,7 +92,6 @@ async def register_service(self, project: str, service: Service, ssh_private_key
                     service.id,
                     service.domain,
                     service.auth,
-                    fallback=f"/api/monitoring/{project}/services/{service.id}/fallback",
                 )
                 stack.push_async_callback(
                     supress_exc_async(self.nginx.unregister_domain, service.domain)
diff --git a/gateway/src/dstack/gateway/main.py b/gateway/src/dstack/gateway/main.py
index b903e51b9b..0b7e649af0 100644
--- a/gateway/src/dstack/gateway/main.py
+++ b/gateway/src/dstack/gateway/main.py
@@ -11,7 +11,6 @@
 from dstack.gateway.core.store import get_store
 from dstack.gateway.errors import GatewayError
 from dstack.gateway.logging import configure_logging
-from dstack.gateway.monitoring.routes import router as monitoring_router
 from dstack.gateway.openai.routes import router as openai_router
 from dstack.gateway.registry.routes import router as registry_router
 
@@ -38,7 +37,6 @@ async def lifespan(app: FastAPI):
 configure_logging(logging.DEBUG)
 app = FastAPI(lifespan=lifespan)
 app.include_router(auth_router, prefix="/auth")
-app.include_router(monitoring_router, prefix="/api/monitoring")
 app.include_router(openai_router, prefix="/api/openai")
 app.include_router(registry_router, prefix="/api/registry")
 
diff --git a/gateway/src/dstack/gateway/monitoring/__init__.py b/gateway/src/dstack/gateway/monitoring/__init__.py
deleted file mode 100644
index e69de29bb2..0000000000
diff --git a/gateway/src/dstack/gateway/monitoring/routes.py b/gateway/src/dstack/gateway/monitoring/routes.py
deleted file mode 100644
index 019a1c6907..0000000000
--- a/gateway/src/dstack/gateway/monitoring/routes.py
+++ /dev/null
@@ -1,13 +0,0 @@
-from fastapi import APIRouter
-from starlette.responses import JSONResponse
-
-router = APIRouter()
-
-
-@router.api_route("/{project}/services/{run_id}/fallback")
-async def service_upstream_fallback(project: str, run_id: str):
-    # TODO(egor-s): store event
-    return JSONResponse(
-        status_code=503,
-        content={"message": "Service temporarily unavailable, try again later"},
-    )
diff --git a/gateway/src/dstack/gateway/resources/nginx/service.jinja2 b/gateway/src/dstack/gateway/resources/nginx/service.jinja2
index c742f8467f..da506fd94d 100644
--- a/gateway/src/dstack/gateway/resources/nginx/service.jinja2
+++ b/gateway/src/dstack/gateway/resources/nginx/service.jinja2
@@ -1,6 +1,5 @@
 upstream {{ upstream }} {
 
-{#  server localhost:{{ port }}/{{ fallback.strip('/') }} backup;#}
 }
 server {
   server_name {{ domain }};

From 1d51b8f92088dea00374d843c8bc460cb70c8b06 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Wed, 6 Mar 2024 12:56:49 +0100
Subject: [PATCH 07/17] Refactor nginx to return 503 without replicas

---
 gateway/src/dstack/gateway/core/nginx.py      | 203 ++++++++++--------
 .../gateway/resources/nginx/entrypoint.jinja2 |  46 ++--
 .../gateway/resources/nginx/service.jinja2    | 115 +++++-----
 gateway/src/tests/core/test_nginx.py          |  31 ---
 gateway/src/tests/core/test_store.py          |  87 --------
 5 files changed, 199 insertions(+), 283 deletions(-)
 delete mode 100644 gateway/src/tests/core/test_nginx.py
 delete mode 100644 gateway/src/tests/core/test_store.py

diff --git a/gateway/src/dstack/gateway/core/nginx.py b/gateway/src/dstack/gateway/core/nginx.py
index a1c33badd2..664f2b2cb5 100644
--- a/gateway/src/dstack/gateway/core/nginx.py
+++ b/gateway/src/dstack/gateway/core/nginx.py
@@ -1,14 +1,13 @@
 import importlib.resources
 import logging
-import re
 import subprocess
 import tempfile
 from asyncio import Lock
 from pathlib import Path
-from typing import Set
+from typing import Annotated, Dict, Literal, Union
 
 import jinja2
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from dstack.gateway.common import run_async
 from dstack.gateway.errors import GatewayError
@@ -18,158 +17,156 @@
 logger = logging.getLogger(__name__)
 
 
+class SiteConfig(BaseModel):
+    type: str
+    domain: str
+
+    def render(self) -> str:
+        template = importlib.resources.read_text(
+            "dstack.gateway.resources.nginx", f"{self.type}.jinja2"
+        )
+        return jinja2.Template(template).render(
+            **self.model_dump(),
+            gateway_port=GATEWAY_PORT,
+        )
+
+
+class ServiceConfig(SiteConfig):
+    type: Literal["service"] = "service"
+    project: str
+    service_id: str
+    auth: bool
+    servers: Dict[str, str] = {}
+
+
+class EntrypointConfig(SiteConfig):
+    type: Literal["entrypoint"] = "entrypoint"
+    proxy_path: str
+
+
 class Nginx(BaseModel):
     """
     Nginx keeps track of registered domains, updates nginx config and issues SSL certificates.
     Its internal state could be serialized to a file and restored from it using pydantic.
     """
 
-    domains: Set[str] = set()
+    configs: Dict[
+        str, Annotated[Union[ServiceConfig, EntrypointConfig], Field(discriminator="type")]
+    ] = {}
     _lock: Lock = Lock()
 
     async def register_service(self, project: str, service_id: str, domain: str, auth: bool):
+        config_name = self.get_config_name(domain)
+        conf = ServiceConfig(
+            project=project,
+            service_id=service_id,
+            domain=domain,
+            auth=auth,
+        )
+
         async with self._lock:
-            if domain in self.domains:
+            if config_name in self.configs:
                 raise GatewayError(f"Domain {domain} is already registered")
 
             logger.debug("Registering service domain %s", domain)
 
             await run_async(self.run_certbot, domain)
-            self.write_conf(
-                self.get_service_conf(project, domain, auth, service_id),
-                CONFIGS_DIR / f"443-{domain}.conf",
-            )
-            self.domains.add(domain)
-            await run_async(self.reload)
+            await run_async(self.write_conf, conf.render(), config_name)
+            self.configs[config_name] = conf
 
         logger.info("Service domain %s is registered now", domain)
 
     async def register_entrypoint(self, domain: str, prefix: str):
+        config_name = self.get_config_name(domain)
+        conf = EntrypointConfig(
+            domain=domain,
+            proxy_path=prefix,
+        )
+
         async with self._lock:
-            if domain in self.domains:
+            if config_name in self.configs:
                 raise GatewayError(f"Domain {domain} is already registered")
 
             logger.debug("Registering entrypoint domain %s", domain)
 
             await run_async(self.run_certbot, domain)
-            self.write_conf(
-                self.get_entrypoint_conf(domain, prefix),
-                CONFIGS_DIR / f"443-{domain}.conf",
-            )
-            self.domains.add(domain)
-            await run_async(self.reload)
+            await run_async(self.write_conf, conf.render(), config_name)
+            self.configs[config_name] = conf
 
         logger.info("Entrypoint domain %s is registered now", domain)
 
     async def unregister_domain(self, domain: str):
+        config_name = self.get_config_name(domain)
+
         async with self._lock:
-            if domain not in self.domains:
+            if config_name not in self.configs:
                 raise GatewayError("Domain is not registered")
 
             logger.debug("Unregistering domain %s", domain)
 
-            conf_path = CONFIGS_DIR / f"443-{domain}.conf"
-            r = subprocess.run(["sudo", "rm", conf_path])
-            if r.returncode != 0:
-                raise GatewayError("Failed to remove nginx config")
-            self.domains.remove(domain)
+            await run_async(sudo_rm, CONFIGS_DIR / config_name)
             await run_async(self.reload)
+            self.configs.pop(config_name)
 
         logger.info("Domain %s is unregistered now", domain)
 
     async def add_upstream(self, domain: str, server: str, replica_id: str):
+        config_name = self.get_config_name(domain)
+
         async with self._lock:
-            if domain not in self.domains:
-                raise GatewayError("Domain is not registered")
+            if config_name not in self.configs:
+                raise GatewayError(f"Domain {domain} is not registered")
 
             logger.debug("Adding upstream %s to domain %s", server, domain)
 
-            config_path = CONFIGS_DIR / f"443-{domain}.conf"
-            with open(config_path, "r") as f:
-                conf = f.read()
-
-            conf = self.add_upstream_to_conf(conf, server, replica_id)
-
-            self.write_conf(conf, config_path)
-            await run_async(self.reload)
+            conf = self.configs[config_name].model_copy(deep=True)
+            conf.servers[replica_id] = server
+            await run_async(self.write_conf, conf.render(), config_name)
+            self.configs[config_name] = conf
 
         logger.debug("Upstream %s is added to domain %s", server, domain)
 
     async def remove_upstream(self, domain: str, replica_id: str):
+        config_name = self.get_config_name(domain)
+
         async with self._lock:
-            if domain not in self.domains:
-                raise GatewayError("Domain is not registered")
+            if config_name not in self.configs:
+                raise GatewayError(f"Domain {domain} is not registered")
+            if replica_id not in self.configs[config_name].servers:
+                raise GatewayError(f"Upstream {replica_id} is not registered")
 
             logger.debug("Removing upstream %s from domain %s", replica_id, domain)
 
-            config_path = CONFIGS_DIR / f"443-{domain}.conf"
-            with open(config_path, "r") as f:
-                conf = f.read()
-
-            conf = self.remove_upstream_from_conf(conf, replica_id)
-
-            self.write_conf(conf, config_path)
-            await run_async(self.reload)
+            conf = self.configs[config_name].model_copy(deep=True)
+            conf.servers.pop(replica_id)
+            await run_async(self.write_conf, conf.render(), config_name)
+            self.configs[config_name] = conf
 
         logger.debug("Upstream %s is removed from domain %s", replica_id, domain)
 
-    @classmethod
-    def get_service_conf(cls, project: str, domain: str, auth: bool, upstream: str) -> str:
-        template = importlib.resources.read_text(
-            "dstack.gateway.resources.nginx", "service.jinja2"
-        )
-        return jinja2.Template(template).render(
-            upstream=upstream,
-            domain=domain,
-            auth=auth,
-            port=GATEWAY_PORT,
-            project=project,
-        )
-
-    @classmethod
-    def get_entrypoint_conf(cls, domain: str, prefix: str) -> str:
-        template = importlib.resources.read_text(
-            "dstack.gateway.resources.nginx", "entrypoint.jinja2"
-        )
-        return jinja2.Template(template).render(
-            domain=domain,
-            port=GATEWAY_PORT,
-            prefix=prefix,
-        )
-
-    @classmethod
-    def add_upstream_to_conf(cls, conf: str, server: str, replica_id: str) -> str:
-        return re.sub(
-            r"(upstream +[^ ]+ *\{)",
-            f"\\1\n  server {server}; # REPLICA:{replica_id}",
-            conf,
-        )
-
-    @classmethod
-    def remove_upstream_from_conf(cls, conf: str, replica_id: str) -> str:
-        return re.sub(
-            f" *server [^;]+; # REPLICA:{replica_id}\n",
-            "",
-            conf,
-        )
-
     @staticmethod
     def reload():
         cmd = ["sudo", "systemctl", "reload", "nginx.service"]
         r = subprocess.run(cmd)
         if r.returncode != 0:
             raise GatewayError("Failed to reload nginx")
-        # TODO(egor-s) implement nginx config rollback
 
     @classmethod
-    def write_conf(cls, conf: str, conf_path: Path):
-        temp = tempfile.NamedTemporaryFile("w")
-        temp.write(conf)
-        temp.flush()
-        temp.seek(0)
-        r = subprocess.run(["sudo", "cp", temp.name, conf_path])
-        if r.returncode != 0:
-            raise GatewayError("Failed to write nginx config")
+    def write_conf(cls, conf: str, conf_name: str):
+        """Update config and reload nginx. Rollback changes on error."""
+        conf_path = CONFIGS_DIR / conf_name
+        old_conf = conf_path.read_text() if conf_path.exists() else None
+
+        sudo_write(conf_path, conf)
+        try:
+            cls.reload()
+        except GatewayError:
+            # rollback changes
+            if old_conf is not None:
+                sudo_write(conf_path, old_conf)
+            else:
+                sudo_rm(conf_path)
+            raise
 
     @staticmethod
     def run_certbot(domain: str):
@@ -180,3 +177,23 @@ def run_certbot(domain: str):
         r = subprocess.run(cmd, capture_output=True)
         if r.returncode != 0:
             raise GatewayError(f"Certbot failed:\n{r.stderr.decode()}")
+
+    @staticmethod
+    def get_config_name(domain: str) -> str:
+        return f"443-{domain}.conf"
+
+
+def sudo_write(path: Path, content: str):
+    with tempfile.NamedTemporaryFile("w") as temp:
+        temp.write(content)
+        temp.flush()
+        temp.seek(0)
+        r = subprocess.run(["sudo", "cp", "-p", temp.name, path])
+        if r.returncode != 0:
+            raise GatewayError("Failed to copy file as sudo")
+
+
+def sudo_rm(path: Path):
+    r = subprocess.run(["sudo", "rm", path])
+    if r.returncode != 0:
+        raise GatewayError("Failed to remove file as sudo")
diff --git a/gateway/src/dstack/gateway/resources/nginx/entrypoint.jinja2 b/gateway/src/dstack/gateway/resources/nginx/entrypoint.jinja2
index ee785bd44d..2fcdfb2a0c 100644
--- a/gateway/src/dstack/gateway/resources/nginx/entrypoint.jinja2
+++ b/gateway/src/dstack/gateway/resources/nginx/entrypoint.jinja2
@@ -1,25 +1,25 @@
 server {
-  server_name {{ domain }};
-  location / {
-    proxy_pass http://localhost:{{ port }}/{{ prefix.strip('/') }}/;
-    proxy_set_header X-Real-IP $remote_addr;
-    proxy_set_header Host $host;
-    proxy_read_timeout 300s;
-  }
-  listen 80;
-  listen 443 ssl;
-  ssl_certificate /etc/letsencrypt/live/{{ domain }}/fullchain.pem;
-  ssl_certificate_key /etc/letsencrypt/live/{{ domain }}/privkey.pem;
-  include /etc/letsencrypt/options-ssl-nginx.conf;
-  ssl_dhparam /etc/letsencrypt/ssl-dhparams.pem;
-  set $force_https 1;
-  if ($scheme = "https") {
-    set $force_https 0;
-  }
-  if ($remote_addr = 127.0.0.1) {
-    set $force_https 0;
-  }
-  if ($force_https) {
-    return 301 https://$host$request_uri;
-  }
+    server_name {{ domain }};
+    location / {
+        proxy_pass http://localhost:{{ gateway_port }}/{{ proxy_path.strip('/') }}/;
+        proxy_set_header X-Real-IP $remote_addr;
+        proxy_set_header Host $host;
+        proxy_read_timeout 300s;
+    }
+    listen 80;
+    listen 443 ssl;
+    ssl_certificate /etc/letsencrypt/live/{{ domain }}/fullchain.pem;
+    ssl_certificate_key /etc/letsencrypt/live/{{ domain }}/privkey.pem;
+    include /etc/letsencrypt/options-ssl-nginx.conf;
+    ssl_dhparam /etc/letsencrypt/ssl-dhparams.pem;
+    set $force_https 1;
+    if ($scheme = "https") {
+        set $force_https 0;
+    }
+    if ($remote_addr = 127.0.0.1) {
+        set $force_https 0;
+    }
+    if ($force_https) {
+        return 301 https://$host$request_uri;
+    }
 }
diff --git a/gateway/src/dstack/gateway/resources/nginx/service.jinja2 b/gateway/src/dstack/gateway/resources/nginx/service.jinja2
index da506fd94d..8867039699 100644
--- a/gateway/src/dstack/gateway/resources/nginx/service.jinja2
+++ b/gateway/src/dstack/gateway/resources/nginx/service.jinja2
@@ -1,56 +1,73 @@
-upstream {{ upstream }} {
-
+{% if servers %}
+upstream {{ service_id }} {
+    {% for replica_id, server in servers.items() %}
+    server {{ server }}; # REPLICA:{{ replica_id }}
+    {% endfor %}
 }
+{% else %}
+
+{% endif %}
 server {
-  server_name {{ domain }};
-  location / {
+    server_name {{ domain }};
+
+    location / {
+        {% if auth %}
+        auth_request /auth;
+        {% endif %}
+
+        {% if servers %}
+        try_files /nonexistent @$http_upgrade;
+        {% else %}
+        return 503;
+        {% endif %}
+    }
+
+    {% if servers %}
+    location @websocket {
+        proxy_pass http://{{ service_id }};
+        proxy_set_header X-Real-IP $remote_addr;
+        proxy_set_header Host $host;
+        proxy_http_version 1.1;
+        proxy_set_header Upgrade $http_upgrade;
+        proxy_set_header Connection "Upgrade";
+        proxy_read_timeout 300s;
+    }
+    location @ {
+        proxy_pass http://{{ service_id }};
+        proxy_set_header X-Real-IP $remote_addr;
+        proxy_set_header Host $host;
+        proxy_read_timeout 300s;
+    }
+    {% endif %}
+
     {% if auth %}
-    auth_request /auth;
+    location = /auth {
+        internal;
+        if ($remote_addr = 127.0.0.1) {
+            return 200;
+        }
+        proxy_pass http://localhost:{{ gateway_port }}/auth/{{ project }};
+        proxy_pass_request_body off;
+        proxy_set_header Content-Length "";
+        proxy_set_header X-Original-URI $request_uri;
+        proxy_set_header Authorization $http_authorization;
+    }
     {% endif %}
-    try_files /nonexistent @$http_upgrade;
-  }
-  location @websocket {
-    proxy_pass http://{{ upstream }};
-    proxy_set_header X-Real-IP $remote_addr;
-    proxy_set_header Host $host;
-    proxy_http_version 1.1;
-    proxy_set_header Upgrade $http_upgrade;
-    proxy_set_header Connection "Upgrade";
-    proxy_read_timeout 300s;
-  }
-  location @ {
-    proxy_pass http://{{ upstream }};
-    proxy_set_header X-Real-IP $remote_addr;
-    proxy_set_header Host $host;
-    proxy_read_timeout 300s;
-  }
-  {% if auth %}
-  location = /auth {
-    internal;
+
+    listen 80;
+    listen 443 ssl;
+    ssl_certificate /etc/letsencrypt/live/{{ domain }}/fullchain.pem;
+    ssl_certificate_key /etc/letsencrypt/live/{{ domain }}/privkey.pem;
+    include /etc/letsencrypt/options-ssl-nginx.conf;
+    ssl_dhparam /etc/letsencrypt/ssl-dhparams.pem;
+    set $force_https 1;
+    if ($scheme = "https") {
+        set $force_https 0;
+    }
     if ($remote_addr = 127.0.0.1) {
-      return 200;
+        set $force_https 0;
+    }
+    if ($force_https) {
+        return 301 https://$host$request_uri;
     }
-    proxy_pass http://localhost:{{ port }}/auth/{{ project }};
-    proxy_pass_request_body off;
-    proxy_set_header Content-Length "";
-    proxy_set_header X-Original-URI $request_uri;
-    proxy_set_header Authorization $http_authorization;
-  }
-  {% endif %}
-  listen 80;
-  listen 443 ssl;
-  ssl_certificate /etc/letsencrypt/live/{{ domain }}/fullchain.pem;
-  ssl_certificate_key /etc/letsencrypt/live/{{ domain }}/privkey.pem;
-  include /etc/letsencrypt/options-ssl-nginx.conf;
-  ssl_dhparam /etc/letsencrypt/ssl-dhparams.pem;
-  set $force_https 1;
-  if ($scheme = "https") {
-    set $force_https 0;
-  }
-  if ($remote_addr = 127.0.0.1) {
-    set $force_https 0;
-  }
-  if ($force_https) {
-    return 301 https://$host$request_uri;
-  }
 }
diff --git a/gateway/src/tests/core/test_nginx.py b/gateway/src/tests/core/test_nginx.py
deleted file mode 100644
index 373b2b6d89..0000000000
--- a/gateway/src/tests/core/test_nginx.py
+++ /dev/null
@@ -1,31 +0,0 @@
-from dstack.gateway.core.nginx import Nginx
-
-
-class TestNginxConf:
-    def test_add_upstream(self):
-        conf = "upstream hex123 {\n" "  server http://fallback backup;\n" "}"
-        expected = (
-            "upstream hex123 {\n"
-            "  server unix:/tmp/job1.sock; # REPLICA:12345\n"
-            "  server http://fallback backup;\n"
-            "}"
-        )
-        conf_out = Nginx.add_upstream_to_conf(conf, "unix:/tmp/job1.sock", "12345")
-        assert conf_out == expected
-
-    def test_remove_upstream(self):
-        conf = (
-            "upstream hex123 {\n"
-            "  server unix:/tmp/job1.sock; # REPLICA:12345\n"
-            "  server unix:/tmp/job2.sock; # REPLICA:67890\n"
-            "  server http://fallback backup;\n"
-            "}"
-        )
-        expected = (
-            "upstream hex123 {\n"
-            "  server unix:/tmp/job1.sock; # REPLICA:12345\n"
-            "  server http://fallback backup;\n"
-            "}"
-        )
-        conf_out = Nginx.remove_upstream_from_conf(conf, "67890")
-        assert conf_out == expected
diff --git a/gateway/src/tests/core/test_store.py b/gateway/src/tests/core/test_store.py
deleted file mode 100644
index be7834ae47..0000000000
--- a/gateway/src/tests/core/test_store.py
+++ /dev/null
@@ -1,87 +0,0 @@
-from typing import Dict, Optional
-from unittest.mock import Mock, patch
-
-import pytest
-
-from dstack.gateway.core.nginx import Nginx
-from dstack.gateway.core.store import Store
-from dstack.gateway.openai.store import OpenAIStore
-from dstack.gateway.schemas import Service
-
-
-@pytest.fixture()
-def ssh_tunnel():
-    with patch("dstack.gateway.services.store.SSHTunnel.create") as mock:
-        yield mock.return_value
-
-
-@pytest.fixture()
-def nginx():
-    yield Mock(Nginx)
-
-
-class TestRegister:
-    @pytest.mark.asyncio
-    async def test_fail_tunnel(self, ssh_tunnel, nginx):
-        store = Store(nginx=nginx)
-        ssh_tunnel.start.side_effect = FooException()
-        with pytest.raises(FooException):
-            await store.register("project", get_service("domain.com"))
-        ssh_tunnel.start.assert_called_once()
-        assert not ssh_tunnel.stop.called
-        assert not nginx.register_service.called
-        assert not nginx.unregister_domain.called
-
-    @pytest.mark.asyncio
-    async def test_fail_nginx(self, ssh_tunnel, nginx):
-        store = Store(nginx=nginx)
-        nginx.register_service.side_effect = FooException()
-        with pytest.raises(FooException):
-            await store.register("project", get_service("domain.com"))
-        ssh_tunnel.start.assert_called_once()
-        ssh_tunnel.stop.assert_called_once()
-        nginx.register_service.assert_called_once()
-        assert not nginx.unregister_domain.called
-
-    @pytest.mark.asyncio
-    async def test_fail_rollback(self, ssh_tunnel, nginx):
-        store = Store(nginx=nginx)
-        nginx.register_service.side_effect = FooException()
-        ssh_tunnel.stop.side_effect = BarException()
-        with pytest.raises(FooException):
-            await store.register("project", get_service("domain.com"))
-        ssh_tunnel.start.assert_called_once()
-        ssh_tunnel.stop.assert_called_once()
-        nginx.register_service.assert_called_once()
-        assert not nginx.unregister_domain.called
-
-    @pytest.mark.asyncio
-    async def test_fail_subscriber(self, ssh_tunnel, nginx):
-        store = Store(nginx=nginx)
-        openai_store = Mock(OpenAIStore)
-        await store.subscribe(openai_store)
-        openai_store.on_register.side_effect = FooException()
-        with pytest.raises(FooException):
-            await store.register("project", get_service("domain.com", {"openai": {}}))
-        ssh_tunnel.start.assert_called_once()
-        ssh_tunnel.stop.assert_called_once()
-        nginx.register_service.assert_called_once()
-        nginx.unregister_domain.assert_called_once()
-
-
-def get_service(domain: str, options: Optional[Dict] = None) -> Service:
-    return Service(
-        public_domain=domain,
-        app_port=8000,
-        ssh_host="user@host",
-        ssh_port=22,
-        options=options or {},
-    )
-
-
-class FooException(Exception):
-    pass
-
-
-class BarException(Exception):
-    pass

From ef67b6e131a0e5b6552b5c7e00df902614b8b341 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Wed, 6 Mar 2024 15:48:39 +0100
Subject: [PATCH 08/17] Normalize domain name

---
 gateway/src/dstack/gateway/registry/routes.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gateway/src/dstack/gateway/registry/routes.py b/gateway/src/dstack/gateway/registry/routes.py
index cbffffc70e..7d4cc0d40d 100644
--- a/gateway/src/dstack/gateway/registry/routes.py
+++ b/gateway/src/dstack/gateway/registry/routes.py
@@ -21,7 +21,7 @@ async def post_register_service(
         project,
         Service(
             id=body.run_id,
-            domain=body.domain,
+            domain=body.domain.lower(),
             auth=body.auth,
             options=body.options,
         ),
@@ -74,5 +74,5 @@ async def post_register_entrypoint(
     body: RegisterEntrypointRequest,
     store: Annotated[Store, Depends(get_store)],
 ) -> OkResponse:
-    await store.register_entrypoint(project, body.domain, body.module)
+    await store.register_entrypoint(project, body.domain.lower(), body.module)
     return OkResponse()

From c4f0e7f70be2085ed5121959e40e0829a459318d Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Wed, 6 Mar 2024 16:45:20 +0100
Subject: [PATCH 09/17] Test gateway client with replicas

---
 gateway/src/dstack/gateway/core/store.py      |  2 +-
 .../server/services/gateways/client.py        | 62 ++++++++++++++-----
 2 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/gateway/src/dstack/gateway/core/store.py b/gateway/src/dstack/gateway/core/store.py
index 5b37a2969d..0e585aed29 100644
--- a/gateway/src/dstack/gateway/core/store.py
+++ b/gateway/src/dstack/gateway/core/store.py
@@ -162,7 +162,7 @@ async def register_replica(self, project: str, service_id: str, replica: Replica
                 service.domain,
             )
 
-            with AsyncExitStack() as stack:
+            async with AsyncExitStack() as stack:
                 # Start SSH tunnel
                 ssh_tunnel = SSHTunnel.create(
                     host=replica.ssh_host,
diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py
index 8cfe2bc546..7052f8e697 100644
--- a/src/dstack/_internal/server/services/gateways/client.py
+++ b/src/dstack/_internal/server/services/gateways/client.py
@@ -5,7 +5,7 @@
 import httpx
 
 from dstack._internal.core.errors import GatewayError
-from dstack._internal.core.models.runs import Run
+from dstack._internal.core.models.runs import JobSubmission, Run
 
 GATEWAY_MANAGEMENT_PORT = 8000
 
@@ -21,9 +21,15 @@ def __init__(self, uds: Optional[str] = None, port: Optional[int] = None):
         self.s = httpx.Client(transport=httpx.HTTPTransport(uds=uds) if uds else None, timeout=30)
 
     def register_service(self, run: Run, ssh_private_key: str):
+        domain = urlparse(run.service.url).hostname
+
+        if "openai" in run.service.options:
+            entrypoint = f"gateway.{domain.split('.', maxsplit=1)[1]}"
+            self.register_openai_entrypoint(run.project_name, entrypoint)
+
         payload = {
             "run_id": run.id.hex,
-            "domain": urlparse(run.service.url).hostname,
+            "domain": domain,
             "auth": run.run_spec.configuration.auth,
             "options": run.service.options,
             "ssh_private_key": ssh_private_key,
@@ -41,35 +47,57 @@ def unregister_service(self, project: str, run_id: uuid.UUID):
             raise gateway_error(resp.json())
         resp.raise_for_status()
 
-    def register_replica(self):
-        pass  # TODO(egor-s): implement
+    def register_replica(self, run: Run, job_submission: JobSubmission):
+        payload = {
+            "job_id": job_submission.id.hex,
+            "app_port": run.run_spec.configuration.port.container_port,
+        }
+        jpd = job_submission.job_provisioning_data
+        if not jpd.dockerized:
+            payload.update(
+                {
+                    "ssh_port": jpd.ssh_port,
+                    "ssh_host": f"{jpd.username}@{jpd.hostname}",
+                }
+            )
+            if jpd.ssh_proxy is not None:
+                payload.update(
+                    {
+                        "ssh_jump_port": jpd.ssh_proxy.port,
+                        "ssh_jump_host": f"{jpd.ssh_proxy.username}@{jpd.ssh_proxy.hostname}",
+                    }
+                )
+        else:
+            payload.update(
+                {
+                    "ssh_port": 10022,
+                    "ssh_host": "root@localhost",
+                    "ssh_jump_port": jpd.ssh_port,
+                    "ssh_jump_host": f"{jpd.username}@{jpd.hostname}",
+                }
+            )
 
-    def unregister_replica(self, project: str, run_id: uuid.UUID, job_id: uuid.UUID):
         resp = self.s.post(
-            self._url(
-                f"/api/registry/{project}/services/{run_id.hex}/replicas/{job_id.hex}/unregister"
-            )
+            self._url(f"/api/registry/{run.project_name}/services/{run.id.hex}/replicas/register"),
+            json=payload,
         )
         if resp.status_code == 400:
             raise gateway_error(resp.json())
         resp.raise_for_status()
 
-    def register_openai_entrypoint(self, project: str, domain: str):
+    def unregister_replica(self, project: str, run_id: uuid.UUID, job_id: uuid.UUID):
         resp = self.s.post(
-            self._url(f"/api/registry/{project}/openai/register"), json={"domain": domain}
+            self._url(
+                f"/api/registry/{project}/services/{run_id.hex}/replicas/{job_id.hex}/unregister"
+            )
         )
         if resp.status_code == 400:
             raise gateway_error(resp.json())
         resp.raise_for_status()
 
-    def preflight(self, project: str, domain: str, private_ssh_key: str, options: dict):
-        if "openai" in options:
-            # TODO(egor-s): custom entrypoint domain
-            entrypoint = f"gateway.{domain.split('.', maxsplit=1)[1]}"
-            self.register_openai_entrypoint(project, entrypoint)
+    def register_openai_entrypoint(self, project: str, domain: str):
         resp = self.s.post(
-            self._url(f"/api/registry/{project}/preflight"),
-            json={"public_domain": domain, "ssh_private_key": private_ssh_key, "options": options},
+            self._url(f"/api/registry/{project}/openai/register"), json={"domain": domain}
         )
         if resp.status_code == 400:
             raise gateway_error(resp.json())

From a70ddcab0a703160d60f9be95541508da3078ab5 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Wed, 6 Mar 2024 17:20:01 +0100
Subject: [PATCH 10/17] Run services end to end

---
 .../background/tasks/process_running_jobs.py  |  5 +-
 .../server/services/gateways/__init__.py      | 83 ++++++++++++-------
 .../server/services/gateways/client.py        | 27 +++---
 .../server/services/jobs/__init__.py          |  3 +-
 4 files changed, 72 insertions(+), 46 deletions(-)

diff --git a/src/dstack/_internal/server/background/tasks/process_running_jobs.py b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
index 1f87dee7ab..a3f783be5b 100644
--- a/src/dstack/_internal/server/background/tasks/process_running_jobs.py
+++ b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
@@ -22,7 +22,6 @@
 )
 from dstack._internal.server.db import get_session_ctx
 from dstack._internal.server.models import (
-    GatewayModel,
     JobModel,
     ProjectModel,
     RepoModel,
@@ -222,10 +221,8 @@ async def _process_job(job_id: UUID):
             and job_model.job_num == 0  # gateway connects only to the first node
             and run.run_spec.configuration.type == "service"
         ):
-            # TODO(egor-s): move code to gateways module
-            gateway = await session.get(GatewayModel, run_model.gateway_id)
             try:
-                await gateways.register_replica(gateway, run, job_provisioning_data)
+                await gateways.register_replica(session, run_model.gateway_id, run, job_submission)
                 logger.debug(
                     "%s: service replica is registered: %s, age=%s",
                     fmt(job_model),
diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py
index e60dc9e5b8..7bdfa3dece 100644
--- a/src/dstack/_internal/server/services/gateways/__init__.py
+++ b/src/dstack/_internal/server/services/gateways/__init__.py
@@ -1,9 +1,11 @@
 import asyncio
+import uuid
 from datetime import timezone
 from typing import List, Optional, Sequence
 from urllib.parse import urlparse
 
 import httpx
+import sqlalchemy.orm as sa_orm
 from sqlalchemy import delete, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -21,7 +23,7 @@
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.gateways import Gateway
 from dstack._internal.core.models.runs import (
-    JobProvisioningData,
+    JobSubmission,
     Run,
     RunSpec,
     ServiceModelSpec,
@@ -31,6 +33,7 @@
 from dstack._internal.server.models import (
     GatewayComputeModel,
     GatewayModel,
+    JobModel,
     ProjectModel,
     RunModel,
 )
@@ -312,11 +315,13 @@ async def register_service(session: AsyncSession, run_model: RunModel):
     try:
         logger.debug("%s: registering service as %s", fmt(run_model), service_spec.url)
         await run_async(
-            conn.client.preflight,
-            run_model.project.name,
-            urlparse(service_spec.url).hostname,
-            run_model.project.ssh_private_key,
-            service_spec.options,
+            conn.client.register_service,
+            project=run_model.project.name,
+            run_id=run_model.id,
+            domain=urlparse(service_spec.url).hostname,
+            auth=run_spec.configuration.auth,
+            options=service_spec.options,
+            ssh_private_key=run_model.project.ssh_private_key,
         )
     except SSHError:
         raise ServerClientError("Gateway tunnel is not working")
@@ -325,48 +330,68 @@ async def register_service(session: AsyncSession, run_model: RunModel):
 
 
 async def register_replica(
-    gateway: Optional[GatewayModel], run: Run, job_provisioning_data: JobProvisioningData
+    session: AsyncSession, gateway_id: uuid.UUID, run: Run, job_submission: JobSubmission
 ):
-    if gateway is None:
-        raise GatewayError("Gateway is not found")
-    if (conn := await gateway_connections_pool.get(gateway.gateway_compute.ip_address)) is None:
-        raise GatewayError("Gateway is not connected")
-
+    conn = await get_gateway_connection(session, gateway_id)
     try:
         await run_async(
             conn.client.register_replica,
-            run,
-            job_provisioning_data,
+            run=run,
+            job_submission=job_submission,
         )
     except (httpx.RequestError, SSHError) as e:
         raise GatewayError(str(e))
 
 
 async def unregister_service(session: AsyncSession, run_model: RunModel):
-    gateway = await session.get(GatewayModel, run_model.gateway_id)
-    if gateway.gateway_compute is None:
-        raise GatewayError("Gateway is broken")
-    conn = await gateway_connections_pool.get(gateway.gateway_compute.ip_address)
-    if conn is None:
-        raise GatewayError("Gateway is not connected")
-
-    # TODO(egor-s): unregister by run_id
-    wildcard_domain = gateway.wildcard_domain.lstrip("*.") if gateway.wildcard_domain else None
-    if wildcard_domain is None:
-        raise ServerClientError("Domain is required for gateway")
-    service_address = f"{run_model.run_name}.{wildcard_domain}"
-
+    conn = await get_gateway_connection(session, run_model.gateway_id)
     project = await session.get(ProjectModel, run_model.project_id)
     try:
         await run_async(
             conn.client.unregister_service,
-            project.name,
-            service_address,
+            project=project.name,
+            run_id=run_model.id,
         )
     except (httpx.RequestError, SSHError) as e:
         raise GatewayError(str(e))
 
 
+async def unregister_replica(session: AsyncSession, job_model: JobModel):
+    res = await session.execute(
+        select(RunModel)
+        .where(RunModel.id == job_model.run_id)
+        .options(sa_orm.joinedload(RunModel.project))
+    )
+    run_model = res.scalar()
+    if run_model.gateway_id is None:
+        return
+
+    conn = await get_gateway_connection(session, run_model.gateway_id)
+    try:
+        await run_async(
+            conn.client.unregister_replica,
+            project=run_model.project.name,
+            run_id=run_model.id,
+            job_id=job_model.id,
+        )
+    except (httpx.RequestError, SSHError) as e:
+        raise GatewayError(str(e))
+
+
+async def get_gateway_connection(
+    session: AsyncSession, gateway_id: uuid.UUID
+) -> GatewayConnection:
+    gateway = await session.get(GatewayModel, gateway_id)
+    if gateway is None:
+        raise GatewayError("Gateway doesn't exist")
+    if gateway.gateway_compute is None:
+        raise GatewayError("Gateway is broken, no compute")
+    conn = await gateway_connections_pool.get(gateway.gateway_compute.ip_address)
+    if conn is None:
+        raise GatewayError("Gateway is not connected")
+    return conn
+
+
 async def init_gateways(session: AsyncSession):
     res = await session.execute(
         select(GatewayComputeModel).where(GatewayComputeModel.deleted == False)
diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py
index 7052f8e697..df835e4591 100644
--- a/src/dstack/_internal/server/services/gateways/client.py
+++ b/src/dstack/_internal/server/services/gateways/client.py
@@ -1,6 +1,5 @@
 import uuid
 from typing import Optional
-from urllib.parse import urlparse
 
 import httpx
 
@@ -20,23 +19,27 @@ def __init__(self, uds: Optional[str] = None, port: Optional[int] = None):
         self.base_url = "http://gateway" if uds else f"http://localhost:{port}"
         self.s = httpx.Client(transport=httpx.HTTPTransport(uds=uds) if uds else None, timeout=30)
 
-    def register_service(self, run: Run, ssh_private_key: str):
-        domain = urlparse(run.service.url).hostname
-
-        if "openai" in run.service.options:
+    def register_service(
+        self,
+        project: str,
+        run_id: uuid.UUID,
+        domain: str,
+        auth: bool,
+        options: dict,
+        ssh_private_key: str,
+    ):
+        if "openai" in options:
             entrypoint = f"gateway.{domain.split('.', maxsplit=1)[1]}"
-            self.register_openai_entrypoint(run.project_name, entrypoint)
+            self.register_openai_entrypoint(project, entrypoint)
 
         payload = {
-            "run_id": run.id.hex,
+            "run_id": run_id.hex,
             "domain": domain,
-            "auth": run.run_spec.configuration.auth,
-            "options": run.service.options,
+            "auth": auth,
+            "options": options,
             "ssh_private_key": ssh_private_key,
         }
-        resp = self.s.post(
-            self._url(f"/api/registry/{run.project_name}/services/register"), json=payload
-        )
+        resp = self.s.post(self._url(f"/api/registry/{project}/services/register"), json=payload)
         if resp.status_code == 400:
             raise gateway_error(resp.json())
         resp.raise_for_status()
diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py
index 5a4ff97b52..4834431534 100644
--- a/src/dstack/_internal/server/services/jobs/__init__.py
+++ b/src/dstack/_internal/server/services/jobs/__init__.py
@@ -8,6 +8,7 @@
 import sqlalchemy.orm as sa_orm
 from sqlalchemy.ext.asyncio import AsyncSession
 
+import dstack._internal.server.services.gateways as gateways
 from dstack._internal.core.errors import SSHError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.configurations import ConfigurationType
@@ -206,7 +207,7 @@ async def process_terminating_job(session: AsyncSession, job_model: JobModel):
                 instance.name,
                 instance.status.name,
             )
-            # TODO(egor-s): unregister service replica
+            await gateways.unregister_replica(session, job_model)
         finally:
             PROCESSING_POOL_IDS.remove(instance.id)
 

From 1c366e8ec9302899d2a830c92e3fed91f6cbc433 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Wed, 6 Mar 2024 17:45:39 +0100
Subject: [PATCH 11/17] Improve services logging

---
 .../background/tasks/process_running_jobs.py     |  8 +-------
 .../server/services/gateways/__init__.py         | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/dstack/_internal/server/background/tasks/process_running_jobs.py b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
index a3f783be5b..eee79c3a48 100644
--- a/src/dstack/_internal/server/background/tasks/process_running_jobs.py
+++ b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
@@ -222,13 +222,7 @@ async def _process_job(job_id: UUID):
             and run.run_spec.configuration.type == "service"
         ):
             try:
-                await gateways.register_replica(session, run_model.gateway_id, run, job_submission)
-                logger.debug(
-                    "%s: service replica is registered: %s, age=%s",
-                    fmt(job_model),
-                    run.service.url,
-                    job_submission.age,
-                )
+                await gateways.register_replica(session, run_model.gateway_id, run, job_model)
             except GatewayError as e:
                 logger.warning(
                     "%s: failed to register service replica: %s, age=%s",
diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py
index 7bdfa3dece..2309981c9d 100644
--- a/src/dstack/_internal/server/services/gateways/__init__.py
+++ b/src/dstack/_internal/server/services/gateways/__init__.py
@@ -9,6 +9,7 @@
 from sqlalchemy import delete, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 
+import dstack._internal.server.services.jobs as jobs_services
 import dstack._internal.utils.random_names as random_names
 from dstack._internal.core.backends.base.compute import (
     get_dstack_gateway_wheel,
@@ -23,7 +24,6 @@
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.gateways import Gateway
 from dstack._internal.core.models.runs import (
-    JobSubmission,
     Run,
     RunSpec,
     ServiceModelSpec,
@@ -323,6 +323,7 @@ async def register_service(session: AsyncSession, run_model: RunModel):
             options=service_spec.options,
             ssh_private_key=run_model.project.ssh_private_key,
         )
+        logger.info("%s: service is registered as %s", fmt(run_model), service_spec.url)
     except SSHError:
         raise ServerClientError("Gateway tunnel is not working")
     except httpx.RequestError as e:
@@ -330,15 +331,18 @@ async def register_service(session: AsyncSession, run_model: RunModel):
 
 
 async def register_replica(
-    session: AsyncSession, gateway_id: uuid.UUID, run: Run, job_submission: JobSubmission
+    session: AsyncSession, gateway_id: uuid.UUID, run: Run, job_model: JobModel
 ):
     conn = await get_gateway_connection(session, gateway_id)
+    job_submission = jobs_services.job_model_to_job_submission(job_model)
     try:
+        logger.debug("%s: registering replica for service %s", fmt(job_model), run.id.hex)
         await run_async(
             conn.client.register_replica,
             run=run,
             job_submission=job_submission,
         )
+        logger.info("%s: replica is registered for service %s", fmt(job_model), run.id.hex)
     except (httpx.RequestError, SSHError) as e:
         raise GatewayError(str(e))
 
@@ -347,11 +351,13 @@ async def unregister_service(session: AsyncSession, run_model: RunModel):
     conn = await get_gateway_connection(session, run_model.gateway_id)
     project = await session.get(ProjectModel, run_model.project_id)
     try:
+        logger.debug("%s: unregistering service", fmt(run_model))
         await run_async(
             conn.client.unregister_service,
             project=project.name,
             run_id=run_model.id,
         )
+        logger.debug("%s: service is unregistered", fmt(run_model))
     except (httpx.RequestError, SSHError) as e:
         raise GatewayError(str(e))
 
@@ -368,12 +374,18 @@ async def unregister_replica(session: AsyncSession, job_model: JobModel):
 
     conn = await get_gateway_connection(session, run_model.gateway_id)
     try:
+        logger.debug(
+            "%s: unregistering replica from service %s", fmt(job_model), job_model.run_id.hex
+        )
         await run_async(
             conn.client.unregister_replica,
             project=run_model.project.name,
             run_id=run_model.id,
             job_id=job_model.id,
         )
+        logger.info(
+            "%s: replica is unregistered from service %s", fmt(job_model), job_model.run_id.hex
+        )
     except (httpx.RequestError, SSHError) as e:
         raise GatewayError(str(e))
 

From 9096336187d336ff8b06e020529239c19cb86cc2 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Thu, 7 Mar 2024 12:19:06 +0100
Subject: [PATCH 12/17] Add replicas to ServiceConfiguration, add replica_num
 to JobSpec

---
 .../_internal/core/models/configurations.py   |  4 +-
 src/dstack/_internal/core/models/runs.py      |  1 +
 .../server/services/jobs/__init__.py          |  8 ++--
 .../services/jobs/configurators/base.py       |  6 ++-
 src/dstack/_internal/server/services/runs.py  | 42 ++++++++++++-------
 src/dstack/_internal/server/testing/common.py |  4 +-
 .../_internal/server/routers/test_runs.py     |  6 ++-
 7 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py
index a95105503c..e123b12b85 100644
--- a/src/dstack/_internal/core/models/configurations.py
+++ b/src/dstack/_internal/core/models/configurations.py
@@ -10,7 +10,7 @@
 from dstack._internal.core.models.gateways import AnyModel
 from dstack._internal.core.models.repos.base import Repo
 from dstack._internal.core.models.repos.virtual import VirtualRepo
-from dstack._internal.core.models.resources import ResourcesSpec
+from dstack._internal.core.models.resources import Range, ResourcesSpec
 
 CommandsList = List[str]
 ValidPort = conint(gt=0, le=65536)
@@ -178,6 +178,7 @@ class ServiceConfiguration(BaseConfiguration):
         resources (Optional[ResourcesSpec]): The requirements to run the configuration.
         model (Optional[ModelMapping]): Mapping of the model for the OpenAI-compatible endpoint.
         auth (bool): Enable the authorization. Defaults to `True`.
+        replicas Range[int]: The range of the number of replicas. Defaults to `1`.
     """
 
     type: Literal["service"] = "service"
@@ -191,6 +192,7 @@ class ServiceConfiguration(BaseConfiguration):
         Field(description="Mapping of the model for the OpenAI-compatible endpoint"),
     ] = None
     auth: Annotated[bool, Field(description="Enable the authorization")] = True
+    replicas: Annotated[Range[int], Field(description="The range ")] = Range[int](min=1, max=1)
 
     @validator("port")
     def convert_port(cls, v) -> PortMapping:
diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py
index c0991fab84..0c389d265d 100644
--- a/src/dstack/_internal/core/models/runs.py
+++ b/src/dstack/_internal/core/models/runs.py
@@ -119,6 +119,7 @@ class Gateway(BaseModel):
 
 
 class JobSpec(BaseModel):
+    replica_num: int = 0  # default value for backward compatibility
     job_num: int
     job_name: str
     app_specs: Optional[List[AppSpec]]
diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py
index 4834431534..d77f057360 100644
--- a/src/dstack/_internal/server/services/jobs/__init__.py
+++ b/src/dstack/_internal/server/services/jobs/__init__.py
@@ -52,15 +52,15 @@
 TERMINATING_PROCESSING_JOBS_IDS = set()
 
 
-def get_jobs_from_run_spec(run_spec: RunSpec) -> List[Job]:
+def get_jobs_from_run_spec(run_spec: RunSpec, replica_num: int = 0) -> List[Job]:
     job_configurator = _get_job_configurator(run_spec)
-    job_specs = job_configurator.get_job_specs()
+    job_specs = job_configurator.get_job_specs(replica_num=replica_num)
     return [Job(job_spec=s, job_submissions=[]) for s in job_specs]
 
 
-def get_job_specs_from_run_spec(run_spec: RunSpec) -> List[JobSpec]:
+def get_job_specs_from_run_spec(run_spec: RunSpec, replica_num: int) -> List[JobSpec]:
     job_configurator = _get_job_configurator(run_spec)
-    job_specs = job_configurator.get_job_specs()
+    job_specs = job_configurator.get_job_specs(replica_num=replica_num)
     return job_specs
 
 
diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py
index 1bc22fcb32..9048c477f1 100644
--- a/src/dstack/_internal/server/services/jobs/configurators/base.py
+++ b/src/dstack/_internal/server/services/jobs/configurators/base.py
@@ -37,10 +37,12 @@ class JobConfigurator(ABC):
     def __init__(self, run_spec: RunSpec):
         self.run_spec = run_spec
 
-    def get_job_specs(self) -> List[JobSpec]:
+    def get_job_specs(self, replica_num: int) -> List[JobSpec]:
         job_spec = JobSpec(
+            replica_num=replica_num,  # TODO(egor-s): add to env variables in the runner
             job_num=0,
-            job_name=self.run_spec.run_name + "-0",
+            job_name=self.run_spec.run_name
+            + f"-0-{replica_num}",  # TODO(egor-s): use actual job_num
             app_specs=self._app_specs(),
             commands=self._commands(),
             env=self._env(),
diff --git a/src/dstack/_internal/server/services/runs.py b/src/dstack/_internal/server/services/runs.py
index b175b4e833..9efd5ac025 100644
--- a/src/dstack/_internal/server/services/runs.py
+++ b/src/dstack/_internal/server/services/runs.py
@@ -342,9 +342,17 @@ async def submit_run(
     )
     session.add(run_model)
 
+    replicas = 1
     if run_spec.configuration.type == "service":
+        replicas = run_spec.configuration.replicas.min
+        if replicas < 1:
+            raise ServerClientError("Replicas count should be at least 1")
+        if replicas != run_spec.configuration.replicas.max:
+            raise ServerClientError("Auto-scaling is not supported yet")
+
         await gateways.register_service(session, run_model)
 
+    # TODO(egor-s) spawn enough replicas
     jobs = get_jobs_from_run_spec(run_spec)
     for job in jobs:
         job.job_spec.pool_name = pool.name
@@ -373,8 +381,8 @@ def create_job_model_for_new_submission(
         run_id=run_model.id,
         run_name=run_model.run_name,
         job_num=job.job_spec.job_num,
-        job_name=job.job_spec.job_name,
-        replica_num=0,  # TODO(egor-s): replace with actual replica number
+        job_name=f"{job.job_spec.job_name}",
+        replica_num=job.job_spec.replica_num,
         submission_num=len(job.job_submissions),
         submitted_at=now,
         last_processed_at=now,
@@ -599,19 +607,22 @@ async def create_instance(
 
 def run_model_to_run(run_model: RunModel, include_job_submissions: bool = True) -> Run:
     jobs: List[Job] = []
-    # JobSpec from JobConfigurator doesn't have gateway information for `service` type
-    # TODO(egor-s): consider replicas
-    run_jobs = sorted(run_model.jobs, key=lambda j: (j.job_num, j.submission_num))
-    for job_num, job_submissions in itertools.groupby(run_jobs):
-        job_spec = None
-        submissions = []
-        for job_model in job_submissions:
-            if job_spec is None:
-                job_spec = JobSpec.parse_raw(job_model.job_spec_data)
-            if include_job_submissions:
-                submissions.append(job_model_to_job_submission(job_model))
-        if job_spec is not None:
-            jobs.append(Job(job_spec=job_spec, job_submissions=submissions))
+    run_jobs = sorted(run_model.jobs, key=lambda j: (j.replica_num, j.job_num, j.submission_num))
+    for replica_num, replica_submissions in itertools.groupby(
+        run_jobs, key=lambda j: j.replica_num
+    ):
+        for job_num, job_submissions in itertools.groupby(
+            replica_submissions, key=lambda j: j.job_num
+        ):
+            job_spec = None
+            submissions = []
+            for job_model in job_submissions:
+                if job_spec is None:
+                    job_spec = JobSpec.parse_raw(job_model.job_spec_data)
+                if include_job_submissions:
+                    submissions.append(job_model_to_job_submission(job_model))
+            if job_spec is not None:
+                jobs.append(Job(job_spec=job_spec, job_submissions=submissions))
 
     run_spec = RunSpec.parse_raw(run_model.run_spec)
 
@@ -634,7 +645,6 @@ def run_model_to_run(run_model: RunModel, include_job_submissions: bool = True)
         latest_job_submission=latest_job_submission,
         service=service_spec,
     )
-    # TODO(egor-s): add replicas support
     run.cost = _get_run_cost(run)
     return run
 
diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py
index 1e55ec7eec..36f0d89eb2 100644
--- a/src/dstack/_internal/server/testing/common.py
+++ b/src/dstack/_internal/server/testing/common.py
@@ -210,13 +210,13 @@ async def create_job(
     replica_num: int = 0,
 ) -> JobModel:
     run_spec = RunSpec.parse_raw(run.run_spec)
-    job_spec = get_job_specs_from_run_spec(run_spec)[0]
+    job_spec = get_job_specs_from_run_spec(run_spec, replica_num=0)[0]
     job = JobModel(
         project_id=run.project_id,
         run_id=run.id,
         run_name=run.run_name,
         job_num=job_num,
-        job_name=run.run_name + "-0",
+        job_name=run.run_name + "-0-0",
         replica_num=replica_num,
         submission_num=submission_num,
         submitted_at=submitted_at,
diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py
index 58ad7b8bad..5b1e31cdf1 100644
--- a/src/tests/_internal/server/routers/test_runs.py
+++ b/src/tests/_internal/server/routers/test_runs.py
@@ -128,7 +128,8 @@ def get_dev_env_run_plan_dict(
                     "env": {},
                     "home_dir": "/root",
                     "image_name": "dstackai/base:py3.8-0.4rc4-cuda-12.1",
-                    "job_name": f"{run_name}-0",
+                    "job_name": f"{run_name}-0-0",
+                    "replica_num": 0,
                     "job_num": 0,
                     "max_duration": None,
                     "pool_name": DEFAULT_POOL_NAME,
@@ -240,7 +241,8 @@ def get_dev_env_run_dict(
                     "env": {},
                     "home_dir": "/root",
                     "image_name": "dstackai/base:py3.8-0.4rc4-cuda-12.1",
-                    "job_name": f"{run_name}-0",
+                    "job_name": f"{run_name}-0-0",
+                    "replica_num": 0,
                     "job_num": 0,
                     "max_duration": None,
                     "pool_name": DEFAULT_POOL_NAME,

From f1b6bb762b2ecb371a63886b8a3c9b222d0cac9e Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Thu, 7 Mar 2024 13:41:32 +0100
Subject: [PATCH 13/17] Drop JobSpec.pool_name, spawn replicas on submit

---
 src/dstack/_internal/core/models/runs.py      |  1 -
 .../server/background/tasks/process_runs.py   | 31 ++++++++++++++++---
 .../server/services/jobs/__init__.py          |  9 +++---
 .../services/jobs/configurators/base.py       |  4 ---
 src/dstack/_internal/server/services/runs.py  | 24 +++++++-------
 src/dstack/_internal/server/testing/common.py |  4 +--
 .../_internal/server/routers/test_runs.py     |  2 --
 7 files changed, 44 insertions(+), 31 deletions(-)

diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py
index 0c389d265d..680a78baa4 100644
--- a/src/dstack/_internal/core/models/runs.py
+++ b/src/dstack/_internal/core/models/runs.py
@@ -132,7 +132,6 @@ class JobSpec(BaseModel):
     requirements: Requirements
     retry_policy: RetryPolicy
     working_dir: str
-    pool_name: Optional[str]  # TODO: remove pool_name from JobSpec
 
 
 class JobProvisioningData(BaseModel):
diff --git a/src/dstack/_internal/server/background/tasks/process_runs.py b/src/dstack/_internal/server/background/tasks/process_runs.py
index 100c8cd707..b66b61d48a 100644
--- a/src/dstack/_internal/server/background/tasks/process_runs.py
+++ b/src/dstack/_internal/server/background/tasks/process_runs.py
@@ -26,6 +26,7 @@
     SUBMITTED_PROCESSING_JOBS_LOCK,
     TERMINATING_PROCESSING_JOBS_IDS,
     TERMINATING_PROCESSING_JOBS_LOCK,
+    get_jobs_from_run_spec,
 )
 from dstack._internal.server.services.runs import (
     PROCESSING_RUNS_IDS,
@@ -106,7 +107,6 @@ async def process_pending_run(session: AsyncSession, run_model: RunModel):
     """Jobs are not created yet"""
 
     # TODO(egor-s): consider retry delay
-    # TODO(egor-s): respect min_replicas and auto-scaling
 
     await session.execute(
         sa.select(RunModel)
@@ -118,15 +118,36 @@ async def process_pending_run(session: AsyncSession, run_model: RunModel):
     )
     run = run_model_to_run(run_model)
 
-    for replica_num, job_models in group_jobs_by_replica_latest(run_model.jobs):
-        for job_model in job_models:
+    replicas = 1
+    if run.run_spec.configuration.type == "service":
+        # TODO(egor-s): consider auto-scaling
+        replicas = run.run_spec.configuration.replicas.min
+
+    scheduled_replicas = 0
+    # Resubmit existing replicas
+    for replica_num, replica_jobs in itertools.groupby(
+        run.jobs, key=lambda j: j.job_spec.replica_num
+    ):
+        if scheduled_replicas >= replicas:
+            break
+        scheduled_replicas += 1
+        for job in replica_jobs:
             new_job_model = create_job_model_for_new_submission(
                 run_model=run_model,
-                job=run.jobs[job_model.job_num],
+                job=job,
                 status=JobStatus.SUBMITTED,
             )
             session.add(new_job_model)
-        break  # TODO(egor-s): add replicas support
+    # Create missing replicas
+    for replica_num in range(scheduled_replicas, replicas):
+        jobs = get_jobs_from_run_spec(run.run_spec, replica_num=replica_num)
+        for job in jobs:
+            job_model = create_job_model_for_new_submission(
+                run_model=run_model,
+                job=job,
+                status=JobStatus.SUBMITTED,
+            )
+            session.add(job_model)
 
     run_model.status = RunStatus.SUBMITTED
     logger.info("%s: run status has changed PENDING -> SUBMITTED", fmt(run_model))
diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py
index d77f057360..7baedff34f 100644
--- a/src/dstack/_internal/server/services/jobs/__init__.py
+++ b/src/dstack/_internal/server/services/jobs/__init__.py
@@ -52,10 +52,11 @@
 TERMINATING_PROCESSING_JOBS_IDS = set()
 
 
-def get_jobs_from_run_spec(run_spec: RunSpec, replica_num: int = 0) -> List[Job]:
-    job_configurator = _get_job_configurator(run_spec)
-    job_specs = job_configurator.get_job_specs(replica_num=replica_num)
-    return [Job(job_spec=s, job_submissions=[]) for s in job_specs]
+def get_jobs_from_run_spec(run_spec: RunSpec, replica_num: int) -> List[Job]:
+    return [
+        Job(job_spec=s, job_submissions=[])
+        for s in get_job_specs_from_run_spec(run_spec, replica_num)
+    ]
 
 
 def get_job_specs_from_run_spec(run_spec: RunSpec, replica_num: int) -> List[JobSpec]:
diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py
index 9048c477f1..5830ca8b22 100644
--- a/src/dstack/_internal/server/services/jobs/configurators/base.py
+++ b/src/dstack/_internal/server/services/jobs/configurators/base.py
@@ -53,7 +53,6 @@ def get_job_specs(self, replica_num: int) -> List[JobSpec]:
             requirements=self._requirements(),
             retry_policy=self._retry_policy(),
             working_dir=self._working_dir(),
-            pool_name=self._pool_name(),
         )
         return [job_spec]
 
@@ -149,9 +148,6 @@ def _python(self) -> str:
             return self.run_spec.configuration.python.value
         return get_default_python_verison()
 
-    def _pool_name(self):
-        return self.run_spec.profile.pool_name
-
 
 def _join_shell_commands(commands: List[str], env: Optional[Dict[str, str]] = None) -> str:
     if env is None:
diff --git a/src/dstack/_internal/server/services/runs.py b/src/dstack/_internal/server/services/runs.py
index 9efd5ac025..672e934f3b 100644
--- a/src/dstack/_internal/server/services/runs.py
+++ b/src/dstack/_internal/server/services/runs.py
@@ -226,7 +226,8 @@ async def get_run_plan(
 
     run_name = run_spec.run_name  # preserve run_name
     run_spec.run_name = "dry-run"  # will regenerate jobs on submission
-    jobs = get_jobs_from_run_spec(run_spec)
+    # TODO(egor-s): do we need to generate all replicas here?
+    jobs = get_jobs_from_run_spec(run_spec, replica_num=0)
     job_plans = []
 
     for job in jobs:
@@ -325,8 +326,6 @@ async def submit_run(
         _validate_run_name(run_spec.run_name)
         await delete_runs(session=session, project=project, runs_names=[run_spec.run_name])
 
-    pool = await get_or_create_pool_by_name(session, project, run_spec.profile.pool_name)
-
     submitted_at = common_utils.get_current_datetime()
     run_model = RunModel(
         id=uuid.uuid4(),
@@ -352,16 +351,15 @@ async def submit_run(
 
         await gateways.register_service(session, run_model)
 
-    # TODO(egor-s) spawn enough replicas
-    jobs = get_jobs_from_run_spec(run_spec)
-    for job in jobs:
-        job.job_spec.pool_name = pool.name
-        job_model = create_job_model_for_new_submission(
-            run_model=run_model,
-            job=job,
-            status=JobStatus.SUBMITTED,
-        )
-        session.add(job_model)
+    for replica_num in range(replicas):
+        jobs = get_jobs_from_run_spec(run_spec, replica_num=replica_num)
+        for job in jobs:
+            job_model = create_job_model_for_new_submission(
+                run_model=run_model,
+                job=job,
+                status=JobStatus.SUBMITTED,
+            )
+            session.add(job_model)
     await session.commit()
     await session.refresh(run_model)
 
diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py
index 36f0d89eb2..9e47e38394 100644
--- a/src/dstack/_internal/server/testing/common.py
+++ b/src/dstack/_internal/server/testing/common.py
@@ -210,13 +210,13 @@ async def create_job(
     replica_num: int = 0,
 ) -> JobModel:
     run_spec = RunSpec.parse_raw(run.run_spec)
-    job_spec = get_job_specs_from_run_spec(run_spec, replica_num=0)[0]
+    job_spec = get_job_specs_from_run_spec(run_spec, replica_num=replica_num)[0]
     job = JobModel(
         project_id=run.project_id,
         run_id=run.id,
         run_name=run.run_name,
         job_num=job_num,
-        job_name=run.run_name + "-0-0",
+        job_name=run.run_name + f"-0-{replica_num}",
         replica_num=replica_num,
         submission_num=submission_num,
         submitted_at=submitted_at,
diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py
index 5b1e31cdf1..6faa14966b 100644
--- a/src/tests/_internal/server/routers/test_runs.py
+++ b/src/tests/_internal/server/routers/test_runs.py
@@ -132,7 +132,6 @@ def get_dev_env_run_plan_dict(
                     "replica_num": 0,
                     "job_num": 0,
                     "max_duration": None,
-                    "pool_name": DEFAULT_POOL_NAME,
                     "registry_auth": None,
                     "requirements": {
                         "resources": {
@@ -245,7 +244,6 @@ def get_dev_env_run_dict(
                     "replica_num": 0,
                     "job_num": 0,
                     "max_duration": None,
-                    "pool_name": DEFAULT_POOL_NAME,
                     "registry_auth": None,
                     "requirements": {
                         "resources": {

From 0b455fb5ce1ed7c529dafc4652b9eae7e630bc38 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Thu, 7 Mar 2024 15:36:16 +0100
Subject: [PATCH 14/17] Show provisioning status per job, show replicas in run
 table

---
 src/dstack/_internal/cli/commands/run.py      |  6 +-
 src/dstack/_internal/cli/utils/common.py      | 17 ++++-
 src/dstack/_internal/cli/utils/run.py         | 73 ++++++++++---------
 src/dstack/_internal/core/models/runs.py      |  1 +
 src/dstack/_internal/server/services/runs.py  |  1 +
 src/dstack/api/_public/runs.py                |  5 +-
 .../_internal/server/routers/test_runs.py     |  2 +
 7 files changed, 68 insertions(+), 37 deletions(-)

diff --git a/src/dstack/_internal/cli/commands/run.py b/src/dstack/_internal/cli/commands/run.py
index 7bb9c12e75..8575b5e76e 100644
--- a/src/dstack/_internal/cli/commands/run.py
+++ b/src/dstack/_internal/cli/commands/run.py
@@ -163,8 +163,12 @@ def _command(self, args: argparse.Namespace):
                     RunStatus.PENDING,
                     RunStatus.PROVISIONING,
                 ):
+                    job_statuses = "\n".join(
+                        f"  - {job.job_spec.job_name} [secondary]({job.job_submissions[-1].status.value})[/]"
+                        for job in run._run.jobs
+                    )
                     status.update(
-                        f"Launching [code]{run.name}[/] [secondary]({run.status.value})[/]"
+                        f"Launching [code]{run.name}[/] [secondary]({run.status.value})[/]\n{job_statuses}"
                     )
                     time.sleep(5)
                     run.refresh()
diff --git a/src/dstack/_internal/cli/utils/common.py b/src/dstack/_internal/cli/utils/common.py
index dff26f27c7..00d3155761 100644
--- a/src/dstack/_internal/cli/utils/common.py
+++ b/src/dstack/_internal/cli/utils/common.py
@@ -1,11 +1,12 @@
 import logging
 import os
 from datetime import datetime
-from typing import Optional
+from typing import Any, Dict, Optional, Union
 
 from rich.console import Console, ConsoleRenderable
 from rich.logging import RichHandler
 from rich.prompt import Confirm
+from rich.table import Table
 from rich.theme import Theme
 from rich.traceback import Traceback
 
@@ -63,3 +64,17 @@ def render(
 def confirm_ask(prompt, **kwargs) -> bool:
     kwargs["console"] = console
     return Confirm.ask(prompt=prompt, **kwargs)
+
+
+def add_row_from_dict(table: Table, data: Dict[Union[str, int], Any], **kwargs):
+    """Maps dict keys to a table columns. `data` key is a column name or index. Missing keys are ignored."""
+    row = []
+    for i, col in enumerate(table.columns):
+        # TODO(egor-s): clear header style
+        if col.header in data:
+            row.append(data[col.header])
+        elif i in data:
+            row.append(data[i])
+        else:
+            row.append("")
+    table.add_row(*row, **kwargs)
diff --git a/src/dstack/_internal/cli/utils/run.py b/src/dstack/_internal/cli/utils/run.py
index 1865faea16..fd48ca6d52 100644
--- a/src/dstack/_internal/cli/utils/run.py
+++ b/src/dstack/_internal/cli/utils/run.py
@@ -1,9 +1,9 @@
-from typing import List, Optional
+from typing import List
 
 from rich.table import Table
 
-from dstack._internal.cli.utils.common import console
-from dstack._internal.core.models.instances import InstanceAvailability, InstanceType
+from dstack._internal.cli.utils.common import add_row_from_dict, console
+from dstack._internal.core.models.instances import InstanceAvailability
 from dstack._internal.core.models.profiles import TerminationPolicy
 from dstack._internal.core.models.runs import RunPlan
 from dstack._internal.utils.common import pretty_date
@@ -112,7 +112,7 @@ def generate_runs_table(
     runs: List[Run], include_configuration: bool = False, verbose: bool = False
 ) -> Table:
     table = Table(box=None)
-    table.add_column("RUN", style="bold", no_wrap=True)
+    table.add_column("NAME", style="bold", no_wrap=True)
     if include_configuration:
         table.add_column("CONFIGURATION", style="grey58")
     table.add_column("BACKEND", style="grey58", no_wrap=True, max_width=16)
@@ -128,33 +128,40 @@ def generate_runs_table(
         table.add_column("ERROR", no_wrap=True)
 
     for run in runs:
-        run = run._run  # TODO
-        job = run.jobs[0]  # TODO
-        provisioning = job.job_submissions[-1].job_provisioning_data  # TODO
-
-        renderables = [run.run_spec.run_name]
-        if include_configuration:
-            renderables.append(run.run_spec.configuration_path)
-        renderables += [
-            provisioning.backend.value if provisioning else "",
-            provisioning.region if provisioning else "",
-            *_render_instance_and_resources(
-                provisioning.instance_type if provisioning else None, verbose
-            ),
-            ("yes" if provisioning.instance_type.resources.spot else "no") if provisioning else "",
-            f"${provisioning.price:.4}" if provisioning else "",
-            run.status,
-            pretty_date(run.submitted_at),
-        ]
-        if verbose:
-            renderables.append("-")  # TODO
-        table.add_row(*renderables)
-    return table
-
+        run = run._run  # TODO(egor-s): make public attribute
+
+        run_row = {
+            "NAME": run.run_spec.run_name,
+            "CONFIGURATION": run.run_spec.configuration_path,
+            "STATUS": run.status,
+            "SUBMITTED": pretty_date(run.submitted_at),
+            "ERROR": run.termination_reason,
+        }
+        if len(run.jobs) != 1:
+            add_row_from_dict(table, run_row)
+
+        for job in run.jobs:
+            job_row = {
+                "NAME": f"  replica {job.job_spec.replica_num}",  # TODO(egor-s): show job_num
+                "STATUS": job.job_submissions[-1].status,
+                "SUBMITTED": pretty_date(job.job_submissions[-1].submitted_at),
+                "ERROR": job.job_submissions[-1].termination_reason,
+            }
+            jpd = job.job_submissions[-1].job_provisioning_data
+            if jpd is not None:
+                job_row.update(
+                    {
+                        "BACKEND": jpd.backend.value,
+                        "REGION": jpd.region,
+                        "INSTANCE": jpd.instance_type.name,
+                        "RESOURCES": jpd.instance_type.resources.pretty_format(),
+                        "SPOT": "yes" if jpd.instance_type.resources.spot else "no",
+                        "PRICE": f"${jpd.price:.4}",
+                    }
+                )
+            if len(run.jobs) == 1:
+                # merge rows
+                job_row.update(run_row)
+            add_row_from_dict(table, job_row, style="secondary" if len(run.jobs) != 1 else None)
 
-def _render_instance_and_resources(instance: Optional[InstanceType], verbose: bool) -> List[str]:
-    if not instance:
-        return [""] if not verbose else ["", ""]
-    rows = [] if not verbose else [instance.name]
-    rows.append(instance.resources.pretty_format())
-    return rows
+    return table
diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py
index 680a78baa4..56bbe17ce0 100644
--- a/src/dstack/_internal/core/models/runs.py
+++ b/src/dstack/_internal/core/models/runs.py
@@ -222,6 +222,7 @@ class Run(BaseModel):
     user: str
     submitted_at: datetime
     status: RunStatus
+    termination_reason: Optional[RunTerminationReason]
     run_spec: RunSpec
     jobs: List[Job]
     latest_job_submission: Optional[JobSubmission]
diff --git a/src/dstack/_internal/server/services/runs.py b/src/dstack/_internal/server/services/runs.py
index 672e934f3b..494afd0a27 100644
--- a/src/dstack/_internal/server/services/runs.py
+++ b/src/dstack/_internal/server/services/runs.py
@@ -638,6 +638,7 @@ def run_model_to_run(run_model: RunModel, include_job_submissions: bool = True)
         user=run_model.user.name,
         submitted_at=run_model.submitted_at.replace(tzinfo=timezone.utc),
         status=run_model.status,
+        termination_reason=run_model.termination_reason,
         run_spec=run_spec,
         jobs=jobs,
         latest_job_submission=latest_job_submission,
diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py
index 2b968604ab..ccbdc567ef 100644
--- a/src/dstack/api/_public/runs.py
+++ b/src/dstack/api/_public/runs.py
@@ -239,7 +239,8 @@ def attach(
             if self.status.is_finished() and self.status != RunStatus.DONE:
                 return False
 
-            provisioning_data = self._run.jobs[0].job_submissions[-1].job_provisioning_data
+            job = self._run.jobs[0]  # TODO(egor-s): pull logs from all replicas?
+            provisioning_data = job.job_submissions[-1].job_provisioning_data
 
             control_sock_path_and_port_locks = SSHAttach.reuse_control_sock_path_and_port_locks(
                 run_name=self.name
@@ -247,7 +248,7 @@ def attach(
 
             if control_sock_path_and_port_locks is None:
                 if self._ports_lock is None:
-                    self._ports_lock = _reserve_ports(self._run.jobs[0].job_spec)
+                    self._ports_lock = _reserve_ports(job.job_spec)
 
                 logger.debug(
                     "Attaching to %s (%s: %s)",
diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py
index 6faa14966b..efb8759c64 100644
--- a/src/tests/_internal/server/routers/test_runs.py
+++ b/src/tests/_internal/server/routers/test_runs.py
@@ -283,6 +283,7 @@ def get_dev_env_run_dict(
         },
         "cost": 0.0,
         "service": None,
+        "termination_reason": None,
     }
 
 
@@ -360,6 +361,7 @@ async def test_lists_runs(self, test_db, session: AsyncSession):
                 },
                 "cost": 0,
                 "service": None,
+                "termination_reason": None,
             }
         ]
 

From 7070ca5b3d8242487dceb2ae0db5f593ee3ba707 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Fri, 8 Mar 2024 11:57:39 +0100
Subject: [PATCH 15/17] Test and debug replicas replacement

---
 .../background/tasks/process_running_jobs.py  | 26 +++++------
 .../server/background/tasks/process_runs.py   | 45 +++++++++++++++----
 .../tasks/process_submitted_jobs.py           | 14 +++++-
 .../server/services/gateways/__init__.py      |  6 +++
 .../server/services/jobs/__init__.py          | 11 ++++-
 5 files changed, 77 insertions(+), 25 deletions(-)

diff --git a/src/dstack/_internal/server/background/tasks/process_running_jobs.py b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
index eee79c3a48..cac76818d8 100644
--- a/src/dstack/_internal/server/background/tasks/process_running_jobs.py
+++ b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
@@ -13,7 +13,6 @@
 from dstack._internal.core.models.configurations import RegistryAuth
 from dstack._internal.core.models.repos import RemoteRepoCreds
 from dstack._internal.core.models.runs import (
-    InstanceStatus,
     Job,
     JobSpec,
     JobStatus,
@@ -31,6 +30,7 @@
 from dstack._internal.server.services.jobs import (
     RUNNING_PROCESSING_JOBS_IDS,
     RUNNING_PROCESSING_JOBS_LOCK,
+    find_job,
     job_model_to_job_submission,
 )
 from dstack._internal.server.services.logging import fmt
@@ -97,9 +97,17 @@ async def _process_job(job_id: UUID):
         repo_model = run_model.repo
         project = run_model.project
         run = run_model_to_run(run_model)
-        job = run.jobs[job_model.job_num]
         job_submission = job_model_to_job_submission(job_model)
         job_provisioning_data = job_submission.job_provisioning_data
+        job = find_job(run.jobs, job_model.replica_num, job_model.job_num)
+        if job is None:
+            # Impossible scenario
+            logger.warning("%s: job not found in run, age=%s", fmt(job_model), job_submission.age)
+            job_model.status = JobStatus.TERMINATING
+            job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+            job_model.last_processed_at = common_utils.get_current_datetime()
+            await session.commit()
+            return
 
         server_ssh_private_key = project.ssh_private_key
         secrets = {}  # TODO secrets
@@ -146,10 +154,6 @@ async def _process_job(job_id: UUID):
                     secrets,
                     repo_creds,
                 )
-                if job_model.instance is not None:
-                    job_model.used_instance_id = job_model.instance.id
-                    if success:
-                        job_model.instance.status = InstanceStatus.BUSY
 
             if not success:
                 # check timeout
@@ -165,9 +169,7 @@ async def _process_job(job_id: UUID):
                     job_model.termination_reason = (
                         JobTerminationReason.WAITING_RUNNER_LIMIT_EXCEEDED
                     )
-                    job_model.used_instance_id = job_model.instance.id
-                    job_model.instance.last_job_processed_at = common_utils.get_current_datetime()
-                    job_model.instance = None
+                    # instance will be emptied by process_terminating_jobs
 
         else:  # fails are not acceptable
             if initial_status == JobStatus.PULLING:
@@ -209,11 +211,7 @@ async def _process_job(job_id: UUID):
                 )
                 job_model.status = JobStatus.TERMINATING
                 job_model.termination_reason = JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY
-                job_model.used_instance_id = job_model.instance.id
-                job_model.instance.last_job_processed_at = common_utils.get_current_datetime()
-                job_model.instance = None
-
-                # job will be terminated by process_finished_jobs
+                # job will be terminated and instance will be emptied by process_terminating_jobs
 
         if (
             initial_status != job_model.status
diff --git a/src/dstack/_internal/server/background/tasks/process_runs.py b/src/dstack/_internal/server/background/tasks/process_runs.py
index b66b61d48a..fbe3f6ed87 100644
--- a/src/dstack/_internal/server/background/tasks/process_runs.py
+++ b/src/dstack/_internal/server/background/tasks/process_runs.py
@@ -11,6 +11,8 @@
 from dstack._internal.core.models.instances import InstanceOffer
 from dstack._internal.core.models.profiles import ProfileRetryPolicy
 from dstack._internal.core.models.runs import (
+    Job,
+    JobSpec,
     JobStatus,
     JobTerminationReason,
     RunSpec,
@@ -42,7 +44,7 @@
 
 logger = get_logger(__name__)
 PROCESSING_INTERVAL = datetime.timedelta(seconds=2)
-JOB_ERROR_CODES_TO_RETRY = {
+JOB_TERMINATION_REASONS_TO_RETRY = {
     JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY,
     JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY,
 }
@@ -160,6 +162,7 @@ async def process_active_run(session: AsyncSession, run_model: RunModel):
     """
     run_spec = RunSpec.parse_raw(run_model.run_spec)
     retry_policy = run_spec.profile.retry_policy or ProfileRetryPolicy()
+    retry_single_job = can_retry_single_job(run_spec)
 
     run_statuses: Set[RunStatus] = set()
     run_termination_reasons: Set[RunTerminationReason] = set()
@@ -214,7 +217,7 @@ async def process_active_run(session: AsyncSession, run_model: RunModel):
         else:
             if replica_needs_retry:
                 replicas_to_retry.append((replica_num, jobs))
-            if not replica_needs_retry or can_retry_single_job(run_spec):
+            if not replica_needs_retry or retry_single_job:
                 run_statuses.update(replica_statuses)
 
     termination_reason: Optional[RunTerminationReason] = None
@@ -241,7 +244,10 @@ async def process_active_run(session: AsyncSession, run_model: RunModel):
     if new_status not in {RunStatus.TERMINATING, RunStatus.PENDING}:
         # No need to retry if the run is terminating,
         # pending run will retry replicas in `process_pending_run`
-        pass  # TODO(egor-s): retry replicas
+        for _, replica_jobs in replicas_to_retry:
+            await retry_replica_jobs(
+                session, run_model, replica_jobs, only_failed=retry_single_job
+            )
 
     if run_model.status != new_status:
         logger.info(
@@ -279,7 +285,7 @@ async def is_retry_enabled(
     session: AsyncSession, job: JobModel, retry_policy: ProfileRetryPolicy
 ) -> bool:
     # retry for spot instances only
-    if retry_policy.retry and job.termination_reason in JOB_ERROR_CODES_TO_RETRY:
+    if retry_policy.retry and job.termination_reason in JOB_TERMINATION_REASONS_TO_RETRY:
         instance = await session.get(InstanceModel, job.used_instance_id)
         instance_offer = InstanceOffer.parse_raw(instance.offer)
         if instance_offer.instance.resources.spot:
@@ -291,11 +297,9 @@ async def is_retry_enabled(
 async def is_retry_limit_exceeded(
     session: AsyncSession, job: JobModel, retry_policy: ProfileRetryPolicy
 ) -> bool:
-    if (
-        retry_policy.limit is not None
-        and get_current_datetime() - job.submitted_at
-        > datetime.timedelta(seconds=retry_policy.limit)
-    ):
+    if retry_policy.limit is not None and get_current_datetime() - job.submitted_at.replace(
+        tzinfo=datetime.timezone.utc
+    ) > datetime.timedelta(seconds=retry_policy.limit):
         return True
     return False
 
@@ -303,3 +307,26 @@ async def is_retry_limit_exceeded(
 def can_retry_single_job(run_spec: RunSpec) -> bool:
     # TODO(egor-s): handle independent and interconnected clusters
     return False
+
+
+async def retry_replica_jobs(
+    session: AsyncSession, run_model: RunModel, latest_jobs: List[JobModel], *, only_failed: bool
+):
+    for job_model in latest_jobs:
+        if job_model.termination_reason not in JOB_TERMINATION_REASONS_TO_RETRY:
+            if only_failed:
+                # No need to resubmit, skip
+                continue
+            if not (job_model.status.is_finished() or job_model.status == JobStatus.TERMINATING):
+                # The job is not finished, but we have to retry all jobs. Terminate it
+                job_model.status = JobStatus.TERMINATING
+                job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+
+        new_job_model = create_job_model_for_new_submission(
+            run_model=run_model,
+            job=Job(job_spec=JobSpec.parse_raw(job_model.job_spec_data), job_submissions=[]),
+            status=JobStatus.SUBMITTED,
+        )
+        # dirty hack to avoid passing all job submissions
+        new_job_model.submission_num = job_model.submission_num + 1
+        session.add(new_job_model)
diff --git a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
index ae7e18e86f..a7fed688c7 100644
--- a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
+++ b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
@@ -31,6 +31,7 @@
     PROCESSING_POOL_LOCK,
     SUBMITTED_PROCESSING_JOBS_IDS,
     SUBMITTED_PROCESSING_JOBS_LOCK,
+    find_job,
 )
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.pools import (
@@ -106,7 +107,15 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
     )
 
     run = run_model_to_run(run_model)
-    job = run.jobs[job_model.job_num]
+    job = find_job(run.jobs, job_model.replica_num, job_model.job_num)
+    if job is None:
+        # Impossible scenario
+        logger.warning("%s: job not found in run", fmt(job_model))
+        job_model.status = JobStatus.TERMINATING
+        job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+        job_model.last_processed_at = common_utils.get_current_datetime()
+        await session.commit()
+        return
 
     async with PROCESSING_POOL_LOCK:
         pool_instances = get_pool_instances(pool)
@@ -130,6 +139,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
 
             logger.info("%s: now is provisioning on '%s'", fmt(job_model), instance.name)
             job_model.job_provisioning_data = instance.job_provisioning_data
+            job_model.used_instance_id = instance.id
             job_model.status = JobStatus.PROVISIONING
             job_model.last_processed_at = common_utils.get_current_datetime()
             await session.commit()
@@ -190,6 +200,8 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
             region=offer.region,
         )
         session.add(im)
+        await session.flush()  # to get im.id
+        job_model.used_instance_id = im.id
     job_model.last_processed_at = common_utils.get_current_datetime()
     await session.commit()
 
diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py
index 2309981c9d..fbe7564f98 100644
--- a/src/dstack/_internal/server/services/gateways/__init__.py
+++ b/src/dstack/_internal/server/services/gateways/__init__.py
@@ -358,6 +358,9 @@ async def unregister_service(session: AsyncSession, run_model: RunModel):
             run_id=run_model.id,
         )
         logger.debug("%s: service is unregistered", fmt(run_model))
+    except GatewayError as e:
+        # ignore if service is not registered
+        logger.warning("%s: unregistering service: %s", fmt(run_model), e)
     except (httpx.RequestError, SSHError) as e:
         raise GatewayError(str(e))
 
@@ -386,6 +389,9 @@ async def unregister_replica(session: AsyncSession, job_model: JobModel):
         logger.info(
             "%s: replica is unregistered from service %s", fmt(job_model), job_model.run_id.hex
         )
+    except GatewayError as e:
+        # ignore if replica is not registered
+        logger.warning("%s: unregistering replica from service: %s", fmt(job_model), e)
     except (httpx.RequestError, SSHError) as e:
         raise GatewayError(str(e))
 
diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py
index 7baedff34f..1d260680c3 100644
--- a/src/dstack/_internal/server/services/jobs/__init__.py
+++ b/src/dstack/_internal/server/services/jobs/__init__.py
@@ -93,6 +93,13 @@ def job_model_to_job_submission(job_model: JobModel) -> JobSubmission:
     )
 
 
+def find_job(jobs: List[Job], replica_num: int, job_num: int) -> Optional[Job]:
+    for job in jobs:
+        if job.job_spec.replica_num == replica_num and job.job_spec.job_num == job_num:
+            return job
+    return None
+
+
 async def terminate_job_provisioning_data_instance(
     project: ProjectModel, job_provisioning_data: JobProvisioningData
 ):
@@ -208,7 +215,9 @@ async def process_terminating_job(session: AsyncSession, job_model: JobModel):
                 instance.name,
                 instance.status.name,
             )
-            await gateways.unregister_replica(session, job_model)
+            await gateways.unregister_replica(
+                session, job_model
+            )  # TODO(egor-s) ensure always runs
         finally:
             PROCESSING_POOL_IDS.remove(instance.id)
 

From ee30ca6bce12a5c251ca2ca04d5b88c365496281 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Fri, 8 Mar 2024 15:59:43 +0100
Subject: [PATCH 16/17] Add tests for replicas

---
 src/dstack/_internal/server/testing/common.py |   8 +-
 .../background/tasks/test_process_runs.py     | 166 +++++++++++++++---
 2 files changed, 151 insertions(+), 23 deletions(-)

diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py
index 9e47e38394..671256689b 100644
--- a/src/dstack/_internal/server/testing/common.py
+++ b/src/dstack/_internal/server/testing/common.py
@@ -7,7 +7,10 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.models.backends.base import BackendType
-from dstack._internal.core.models.configurations import DevEnvironmentConfiguration
+from dstack._internal.core.models.configurations import (
+    AnyRunConfiguration,
+    DevEnvironmentConfiguration,
+)
 from dstack._internal.core.models.instances import InstanceType, Resources
 from dstack._internal.core.models.profiles import (
     DEFAULT_POOL_NAME,
@@ -150,6 +153,7 @@ def get_run_spec(
     run_name: str,
     repo_id: str,
     profile: Optional[Profile] = None,
+    configuration: Optional[AnyRunConfiguration] = None,
 ) -> RunSpec:
     if profile is None:
         profile = Profile(name="default")
@@ -160,7 +164,7 @@ def get_run_spec(
         repo_code_hash=None,
         working_dir=".",
         configuration_path="dstack.yaml",
-        configuration=DevEnvironmentConfiguration(ide="vscode"),
+        configuration=configuration or DevEnvironmentConfiguration(ide="vscode"),
         profile=profile,
         ssh_key_pub="",
     )
diff --git a/src/tests/_internal/server/background/tasks/test_process_runs.py b/src/tests/_internal/server/background/tasks/test_process_runs.py
index cbbc706d32..c564479502 100644
--- a/src/tests/_internal/server/background/tasks/test_process_runs.py
+++ b/src/tests/_internal/server/background/tasks/test_process_runs.py
@@ -1,12 +1,15 @@
 import datetime
+from typing import Union
 from unittest.mock import patch
 
 import pytest
-import pytest_asyncio
+from pydantic import parse_obj_as
 from sqlalchemy.ext.asyncio import AsyncSession
 
 import dstack._internal.server.background.tasks.process_runs as process_runs
+from dstack._internal.core.models.configurations import ServiceConfiguration
 from dstack._internal.core.models.profiles import Profile, ProfileRetryPolicy
+from dstack._internal.core.models.resources import Range
 from dstack._internal.core.models.runs import (
     JobStatus,
     JobTerminationReason,
@@ -27,8 +30,9 @@
 )
 
 
-@pytest_asyncio.fixture
-async def run(session: AsyncSession) -> RunModel:
+async def make_run(
+    session: AsyncSession, status: RunStatus = RunStatus.SUBMITTED, replicas: Union[str, int] = 1
+) -> RunModel:
     project = await create_project(session=session)
     user = await create_user(session=session)
     repo = await create_repo(
@@ -43,7 +47,16 @@ async def run(session: AsyncSession) -> RunModel:
         name="test-profile",
         retry_policy=ProfileRetryPolicy(retry=True),
     )
-    run_spec = get_run_spec(repo_id=repo.name, run_name=run_name, profile=profile)
+    run_spec = get_run_spec(
+        repo_id=repo.name,
+        run_name=run_name,
+        profile=profile,
+        configuration=ServiceConfiguration(
+            commands=["echo hello"],
+            port=8000,
+            replicas=parse_obj_as(Range[int], replicas),
+        ),
+    )
     return await create_run(
         session=session,
         project=project,
@@ -51,13 +64,14 @@ async def run(session: AsyncSession) -> RunModel:
         user=user,
         run_name=run_name,
         run_spec=run_spec,
+        status=status,
     )
 
 
 class TestProcessRuns:
     @pytest.mark.asyncio
-    async def test_submitted_to_provisioning(self, test_db, session: AsyncSession, run: RunModel):
-        run.status = RunStatus.SUBMITTED
+    async def test_submitted_to_provisioning(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.SUBMITTED)
         await create_job(session=session, run=run, status=JobStatus.PROVISIONING)
 
         await process_runs.process_single_run(run.id, [])
@@ -65,8 +79,8 @@ async def test_submitted_to_provisioning(self, test_db, session: AsyncSession, r
         assert run.status == RunStatus.PROVISIONING
 
     @pytest.mark.asyncio
-    async def test_provisioning_to_running(self, test_db, session: AsyncSession, run: RunModel):
-        run.status = RunStatus.PROVISIONING
+    async def test_provisioning_to_running(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.PROVISIONING)
         await create_job(session=session, run=run, status=JobStatus.RUNNING)
 
         await process_runs.process_single_run(run.id, [])
@@ -74,8 +88,8 @@ async def test_provisioning_to_running(self, test_db, session: AsyncSession, run
         assert run.status == RunStatus.RUNNING
 
     @pytest.mark.asyncio
-    async def test_keep_provisioning(self, test_db, session: AsyncSession, run: RunModel):
-        run.status = RunStatus.PROVISIONING
+    async def test_keep_provisioning(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.PROVISIONING)
         await create_job(session=session, run=run, status=JobStatus.PULLING)
 
         await process_runs.process_single_run(run.id, [])
@@ -83,8 +97,8 @@ async def test_keep_provisioning(self, test_db, session: AsyncSession, run: RunM
         assert run.status == RunStatus.PROVISIONING
 
     @pytest.mark.asyncio
-    async def test_running_to_done(self, test_db, session: AsyncSession, run: RunModel):
-        run.status = RunStatus.RUNNING
+    async def test_running_to_done(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.RUNNING)
         await create_job(session=session, run=run, status=JobStatus.DONE)
 
         await process_runs.process_single_run(run.id, [])
@@ -93,8 +107,8 @@ async def test_running_to_done(self, test_db, session: AsyncSession, run: RunMod
         assert run.termination_reason == RunTerminationReason.ALL_JOBS_DONE
 
     @pytest.mark.asyncio
-    async def test_terminate_run_jobs(self, test_db, session: AsyncSession, run: RunModel):
-        run.status = RunStatus.TERMINATING
+    async def test_terminate_run_jobs(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.TERMINATING)
         run.termination_reason = RunTerminationReason.JOB_FAILED
         job = await create_job(
             session=session,
@@ -113,11 +127,11 @@ async def test_terminate_run_jobs(self, test_db, session: AsyncSession, run: Run
         assert run.status == RunStatus.TERMINATING
 
     @pytest.mark.asyncio
-    async def test_retry_running_to_pending(self, test_db, session: AsyncSession, run: RunModel):
+    async def test_retry_running_to_pending(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.RUNNING)
         instance = await create_instance(
             session, project=run.project, pool=run.project.default_pool, spot=True
         )
-        run.status = RunStatus.RUNNING
         await create_job(
             session=session,
             run=run,
@@ -133,11 +147,11 @@ async def test_retry_running_to_pending(self, test_db, session: AsyncSession, ru
         assert run.status == RunStatus.PENDING
 
     @pytest.mark.asyncio
-    async def test_retry_running_to_failed(self, test_db, session: AsyncSession, run: RunModel):
+    async def test_retry_running_to_failed(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.RUNNING)
         instance = await create_instance(
             session, project=run.project, pool=run.project.default_pool, spot=True
         )
-        run.status = RunStatus.RUNNING
         # job exited with non-zero code
         await create_job(
             session=session,
@@ -155,8 +169,8 @@ async def test_retry_running_to_failed(self, test_db, session: AsyncSession, run
         assert run.termination_reason == RunTerminationReason.JOB_FAILED
 
     @pytest.mark.asyncio
-    async def test_pending_to_submitted(self, test_db, session: AsyncSession, run: RunModel):
-        run.status = RunStatus.PENDING
+    async def test_pending_to_submitted(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.PENDING)
         await create_job(session=session, run=run, status=JobStatus.FAILED)
 
         await process_runs.process_single_run(run.id, [])
@@ -167,6 +181,116 @@ async def test_pending_to_submitted(self, test_db, session: AsyncSession, run: R
         assert run.jobs[1].status == JobStatus.SUBMITTED
 
 
+class TestProcessRunsReplicas:
+    @pytest.mark.asyncio
+    async def test_submitted_to_provisioning_if_any(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.SUBMITTED, replicas=2)
+        await create_job(session=session, run=run, status=JobStatus.SUBMITTED, replica_num=0)
+        await create_job(session=session, run=run, status=JobStatus.PROVISIONING, replica_num=1)
+
+        await process_runs.process_single_run(run.id, [])
+        await session.refresh(run)
+        assert run.status == RunStatus.PROVISIONING
+
+    @pytest.mark.asyncio
+    async def test_provisioning_to_running_if_any(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.PROVISIONING, replicas=2)
+        await create_job(session=session, run=run, status=JobStatus.RUNNING, replica_num=0)
+        await create_job(session=session, run=run, status=JobStatus.PROVISIONING, replica_num=1)
+
+        await process_runs.process_single_run(run.id, [])
+        await session.refresh(run)
+        assert run.status == RunStatus.RUNNING
+
+    @pytest.mark.asyncio
+    async def test_all_no_capacity_to_pending(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.RUNNING, replicas=2)
+        await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.TERMINATING,
+            termination_reason=JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY,
+            replica_num=0,
+            instance=await create_instance(
+                session, project=run.project, pool=run.project.default_pool, spot=True
+            ),
+        )
+        await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.TERMINATING,
+            termination_reason=JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY,
+            replica_num=1,
+            instance=await create_instance(
+                session, project=run.project, pool=run.project.default_pool, spot=True
+            ),
+        )
+
+        with patch("dstack._internal.utils.common.get_current_datetime") as datetime_mock:
+            datetime_mock.return_value = run.submitted_at + datetime.timedelta(minutes=3)
+            await process_runs.process_single_run(run.id, [])
+        await session.refresh(run)
+        assert run.status == RunStatus.PENDING
+
+    @pytest.mark.asyncio
+    async def test_some_no_capacity_keep_running(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.RUNNING, replicas=2)
+        await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.TERMINATING,
+            termination_reason=JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY,
+            replica_num=0,
+            instance=await create_instance(
+                session, project=run.project, pool=run.project.default_pool, spot=True
+            ),
+        )
+        await create_job(session=session, run=run, status=JobStatus.RUNNING, replica_num=1)
+
+        await process_runs.process_single_run(run.id, [])
+        await session.refresh(run)
+        assert run.status == RunStatus.RUNNING
+        assert len(run.jobs) == 3
+        assert run.jobs[2].status == JobStatus.SUBMITTED
+        assert run.jobs[2].replica_num == 0
+
+    @pytest.mark.asyncio
+    async def test_some_failed_to_terminating(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.RUNNING, replicas=2)
+        await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.FAILED,
+            termination_reason=JobTerminationReason.CONTAINER_EXITED_WITH_ERROR,
+            replica_num=0,
+        )
+        await create_job(session=session, run=run, status=JobStatus.RUNNING, replica_num=1)
+
+        await process_runs.process_single_run(run.id, [])
+        await session.refresh(run)
+        assert run.status == RunStatus.TERMINATING
+        assert run.termination_reason == RunTerminationReason.JOB_FAILED
+
+    @pytest.mark.asyncio
+    async def test_pending_to_submitted_adds_replicas(self, test_db, session: AsyncSession):
+        run = await make_run(session, status=RunStatus.PENDING, replicas=2)
+        await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.FAILED,
+            termination_reason=JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY,
+            replica_num=0,
+        )
+
+        await process_runs.process_single_run(run.id, [])
+        await session.refresh(run)
+        assert run.status == RunStatus.SUBMITTED
+        assert len(run.jobs) == 3
+        assert run.jobs[1].status == JobStatus.SUBMITTED
+        assert run.jobs[1].replica_num == 0
+        assert run.jobs[2].status == JobStatus.SUBMITTED
+        assert run.jobs[2].replica_num == 1
+
+
 # TODO(egor-s): TestProcessRunsMultiNode
-# TODO(egor-s): TestProcessRunsReplicas
 # TODO(egor-s): TestProcessRunsAutoScaling

From 52cb2a9ad7a775d29d3f2262ebe704f1278384c8 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Mon, 11 Mar 2024 13:10:51 +0100
Subject: [PATCH 17/17] Address review comments

---
 gateway/src/dstack/gateway/core/store.py      | 22 +++++++++++++++++--
 gateway/src/dstack/gateway/registry/routes.py | 10 ++++-----
 .../background/tasks/process_running_jobs.py  |  8 -------
 .../server/background/tasks/process_runs.py   |  4 ++--
 .../tasks/process_submitted_jobs.py           |  8 -------
 .../server/services/gateways/client.py        |  6 ++++-
 .../server/services/jobs/__init__.py          |  8 ++++---
 src/dstack/_internal/server/services/runs.py  |  2 +-
 8 files changed, 38 insertions(+), 30 deletions(-)

diff --git a/gateway/src/dstack/gateway/core/store.py b/gateway/src/dstack/gateway/core/store.py
index 0e585aed29..5d3eeb6eb6 100644
--- a/gateway/src/dstack/gateway/core/store.py
+++ b/gateway/src/dstack/gateway/core/store.py
@@ -46,6 +46,8 @@ class Store(BaseModel):
     Store is a central place to register and unregister services.
     Other components can subscribe to updates.
     Its internal state could be serialized to a file and restored from it using pydantic.
+
+    Domains and project names must be lowercased.
     """
 
     services: Dict[str, Service] = {}
@@ -122,7 +124,7 @@ async def unregister_service(self, project: str, service_id: str):
 
             logger.debug("%s: unregistering service %s (%s)", project, service_id, service.domain)
 
-            await asyncio.gather(
+            results = await asyncio.gather(
                 # Terminate all SSH tunnels
                 *(
                     run_async(replica.ssh_tunnel.stop)
@@ -138,6 +140,14 @@ async def unregister_service(self, project: str, service_id: str):
                 ),
                 return_exceptions=True,
             )
+            for exc in results:
+                if isinstance(exc, Exception):
+                    logger.error(
+                        "%s: exception during unregistering service %s: %s",
+                        project,
+                        service_id,
+                        exc,
+                    )
 
             self.projects[project].remove(service_id)
             self.services.pop(service_id)
@@ -221,13 +231,21 @@ async def unregister_replica(self, project: str, service_id: str, replica_id: st
                 service.domain,
             )
 
-            await asyncio.gather(
+            results = await asyncio.gather(
                 # Terminate SSH tunnel
                 run_async(replica.ssh_tunnel.stop),
                 # Remove from nginx
                 self.nginx.remove_upstream(service.domain, replica.id),
                 return_exceptions=True,
             )
+            for exc in results:
+                if isinstance(exc, Exception):
+                    logger.error(
+                        "%s: exception during unregistering replica %s: %s",
+                        project,
+                        replica_id,
+                        exc,
+                    )
 
             service.replicas.remove(replica)
 
diff --git a/gateway/src/dstack/gateway/registry/routes.py b/gateway/src/dstack/gateway/registry/routes.py
index 7d4cc0d40d..facb5c01dd 100644
--- a/gateway/src/dstack/gateway/registry/routes.py
+++ b/gateway/src/dstack/gateway/registry/routes.py
@@ -18,7 +18,7 @@ async def post_register_service(
     project: str, body: RegisterServiceRequest, store: Annotated[Store, Depends(get_store)]
 ) -> OkResponse:
     await store.register_service(
-        project,
+        project.lower(),
         Service(
             id=body.run_id,
             domain=body.domain.lower(),
@@ -34,7 +34,7 @@ async def post_register_service(
 async def post_unregister_services(
     project: str, run_id: str, store: Annotated[Store, Depends(get_store)]
 ) -> OkResponse:
-    await store.unregister_service(project, run_id)
+    await store.unregister_service(project.lower(), run_id)
     return OkResponse()
 
 
@@ -46,7 +46,7 @@ async def post_register_replica(
     store: Annotated[Store, Depends(get_store)],
 ) -> OkResponse:
     await store.register_replica(
-        project,
+        project.lower(),
         run_id,
         Replica(
             id=body.job_id,
@@ -64,7 +64,7 @@ async def post_register_replica(
 async def post_unregister_replica(
     project: str, run_id: str, job_id: str, store: Annotated[Store, Depends(get_store)]
 ) -> OkResponse:
-    await store.unregister_replica(project, run_id, job_id)
+    await store.unregister_replica(project.lower(), run_id, job_id)
     return OkResponse()
 
 
@@ -74,5 +74,5 @@ async def post_register_entrypoint(
     body: RegisterEntrypointRequest,
     store: Annotated[Store, Depends(get_store)],
 ) -> OkResponse:
-    await store.register_entrypoint(project, body.domain.lower(), body.module)
+    await store.register_entrypoint(project.lower(), body.domain.lower(), body.module)
     return OkResponse()
diff --git a/src/dstack/_internal/server/background/tasks/process_running_jobs.py b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
index cac76818d8..aca4f231c5 100644
--- a/src/dstack/_internal/server/background/tasks/process_running_jobs.py
+++ b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
@@ -100,14 +100,6 @@ async def _process_job(job_id: UUID):
         job_submission = job_model_to_job_submission(job_model)
         job_provisioning_data = job_submission.job_provisioning_data
         job = find_job(run.jobs, job_model.replica_num, job_model.job_num)
-        if job is None:
-            # Impossible scenario
-            logger.warning("%s: job not found in run, age=%s", fmt(job_model), job_submission.age)
-            job_model.status = JobStatus.TERMINATING
-            job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
-            job_model.last_processed_at = common_utils.get_current_datetime()
-            await session.commit()
-            return
 
         server_ssh_private_key = project.ssh_private_key
         secrets = {}  # TODO secrets
diff --git a/src/dstack/_internal/server/background/tasks/process_runs.py b/src/dstack/_internal/server/background/tasks/process_runs.py
index fbe3f6ed87..9f36ee12d1 100644
--- a/src/dstack/_internal/server/background/tasks/process_runs.py
+++ b/src/dstack/_internal/server/background/tasks/process_runs.py
@@ -122,8 +122,8 @@ async def process_pending_run(session: AsyncSession, run_model: RunModel):
 
     replicas = 1
     if run.run_spec.configuration.type == "service":
-        # TODO(egor-s): consider auto-scaling
-        replicas = run.run_spec.configuration.replicas.min
+        # TODO(egor-s): consider max for auto-scaling
+        replicas = run.run_spec.configuration.replicas.min or 0
 
     scheduled_replicas = 0
     # Resubmit existing replicas
diff --git a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
index a7fed688c7..f2c72cca09 100644
--- a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
+++ b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
@@ -108,14 +108,6 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
 
     run = run_model_to_run(run_model)
     job = find_job(run.jobs, job_model.replica_num, job_model.job_num)
-    if job is None:
-        # Impossible scenario
-        logger.warning("%s: job not found in run", fmt(job_model))
-        job_model.status = JobStatus.TERMINATING
-        job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
-        job_model.last_processed_at = common_utils.get_current_datetime()
-        await session.commit()
-        return
 
     async with PROCESSING_POOL_LOCK:
         pool_instances = get_pool_instances(pool)
diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py
index df835e4591..c3778e441e 100644
--- a/src/dstack/_internal/server/services/gateways/client.py
+++ b/src/dstack/_internal/server/services/gateways/client.py
@@ -100,7 +100,11 @@ def unregister_replica(self, project: str, run_id: uuid.UUID, job_id: uuid.UUID)
 
     def register_openai_entrypoint(self, project: str, domain: str):
         resp = self.s.post(
-            self._url(f"/api/registry/{project}/openai/register"), json={"domain": domain}
+            self._url(f"/api/registry/{project}/entrypoints/register"),
+            json={
+                "module": "openai",
+                "domain": domain,
+            },
         )
         if resp.status_code == 400:
             raise gateway_error(resp.json())
diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py
index 1d260680c3..a78be713db 100644
--- a/src/dstack/_internal/server/services/jobs/__init__.py
+++ b/src/dstack/_internal/server/services/jobs/__init__.py
@@ -9,7 +9,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 import dstack._internal.server.services.gateways as gateways
-from dstack._internal.core.errors import SSHError
+from dstack._internal.core.errors import ResourceNotFoundError, SSHError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.configurations import ConfigurationType
 from dstack._internal.core.models.runs import (
@@ -93,11 +93,13 @@ def job_model_to_job_submission(job_model: JobModel) -> JobSubmission:
     )
 
 
-def find_job(jobs: List[Job], replica_num: int, job_num: int) -> Optional[Job]:
+def find_job(jobs: List[Job], replica_num: int, job_num: int) -> Job:
     for job in jobs:
         if job.job_spec.replica_num == replica_num and job.job_spec.job_num == job_num:
             return job
-    return None
+    raise ResourceNotFoundError(
+        f"Job with replica_num={replica_num} and job_num={job_num} not found"
+    )
 
 
 async def terminate_job_provisioning_data_instance(
diff --git a/src/dstack/_internal/server/services/runs.py b/src/dstack/_internal/server/services/runs.py
index 494afd0a27..a3c6e8e027 100644
--- a/src/dstack/_internal/server/services/runs.py
+++ b/src/dstack/_internal/server/services/runs.py
@@ -344,7 +344,7 @@ async def submit_run(
     replicas = 1
     if run_spec.configuration.type == "service":
         replicas = run_spec.configuration.replicas.min
-        if replicas < 1:
+        if replicas is None or replicas < 1:
             raise ServerClientError("Replicas count should be at least 1")
         if replicas != run_spec.configuration.replicas.max:
             raise ServerClientError("Auto-scaling is not supported yet")