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/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/core/nginx.py b/gateway/src/dstack/gateway/core/nginx.py new file mode 100644 index 0000000000..664f2b2cb5 --- /dev/null +++ b/gateway/src/dstack/gateway/core/nginx.py @@ -0,0 +1,199 @@ +import importlib.resources +import logging +import subprocess +import tempfile +from asyncio import Lock +from pathlib import Path +from typing import Annotated, Dict, Literal, Union + +import jinja2 +from pydantic import BaseModel, Field + +from dstack.gateway.common import run_async +from dstack.gateway.errors import GatewayError + +CONFIGS_DIR = Path("/etc/nginx/sites-enabled") +GATEWAY_PORT = 8000 +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. + """ + + 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 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) + 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 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) + 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 config_name not in self.configs: + raise GatewayError("Domain is not registered") + + logger.debug("Unregistering domain %s", 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 config_name not in self.configs: + raise GatewayError(f"Domain {domain} is not registered") + + logger.debug("Adding upstream %s to domain %s", server, domain) + + 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 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) + + 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) + + @staticmethod + def reload(): + cmd = ["sudo", "systemctl", "reload", "nginx.service"] + r = subprocess.run(cmd) + if r.returncode != 0: + raise GatewayError("Failed to reload nginx") + + @classmethod + 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): + logger.info("Running certbot for %s", domain) + cmd = ["sudo", "certbot", "certonly"] + cmd += ["--non-interactive", "--agree-tos", "--register-unsafely-without-email"] + cmd += ["--nginx", "--domain", domain] + 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/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..5d3eeb6eb6 --- /dev/null +++ b/gateway/src/dstack/gateway/core/store.py @@ -0,0 +1,334 @@ +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, ValidationError + +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. + + Domains and project names must be lowercased. + """ + + 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) + + results = 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, + ) + 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) + + 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, + ) + + async 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(), + 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, f"unix:{ssh_tunnel.sock_path}", replica.id + ) + 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, + ) + + 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) + + 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: + 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/services/tunnel.py b/gateway/src/dstack/gateway/core/tunnel.py similarity index 88% rename from gateway/src/dstack/gateway/services/tunnel.py rename to gateway/src/dstack/gateway/core/tunnel.py index a760a37f27..91207e4509 100644 --- a/gateway/src/dstack/gateway/services/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/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 626cf3ca88..0b7e649af0 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 @@ -35,11 +36,16 @@ 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(openai_router, prefix="/api/openai") +app.include_router(registry_router, prefix="/api/registry") @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.to_response() diff --git a/gateway/src/dstack/gateway/openai/routes.py b/gateway/src/dstack/gateway/openai/routes.py index 14338de50c..5674ed2246 100644 --- a/gateway/src/dstack/gateway/openai/routes.py +++ b/gateway/src/dstack/gateway/openai/routes.py @@ -4,7 +4,7 @@ from fastapi.responses import StreamingResponse from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer -from dstack.gateway.errors import GatewayError +from dstack.gateway.core.auth import AuthProvider, get_auth from dstack.gateway.openai.schemas import ( ChatCompletionsChunk, ChatCompletionsRequest, @@ -12,7 +12,6 @@ ModelsResponse, ) from dstack.gateway.openai.store import OpenAIStore, get_store -from dstack.gateway.services.auth import AuthProvider, get_auth async def auth_required( @@ -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/openai/store.py b/gateway/src/dstack/gateway/openai/store.py index c364be3501..330d466b8b 100644 --- a/gateway/src/dstack/gateway/openai/store.py +++ b/gateway/src/dstack/gateway/openai/store.py @@ -5,15 +5,14 @@ from pydantic import BaseModel, ValidationError +from dstack.gateway.core.persistent import get_persistent_state +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 -from dstack.gateway.services.persistent import get_persistent_state -from dstack.gateway.services.store import StoreSubscriber 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 = [] diff --git a/gateway/src/dstack/gateway/registry/routes.py b/gateway/src/dstack/gateway/registry/routes.py index 72a454ed78..facb5c01dd 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.lower(), + Service( + id=body.run_id, + domain=body.domain.lower(), + 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.lower(), 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.lower(), + 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.lower(), 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.lower(), body.domain.lower(), 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/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 33b22e616f..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 }} { - server {{ server }}; +{% 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/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/nginx.py b/gateway/src/dstack/gateway/services/nginx.py deleted file mode 100644 index 8156025dd7..0000000000 --- a/gateway/src/dstack/gateway/services/nginx.py +++ /dev/null @@ -1,121 +0,0 @@ -import importlib.resources -import logging -import re -import subprocess -import tempfile -from asyncio import Lock -from pathlib import Path -from typing import Optional, Set - -import jinja2 -from pydantic import BaseModel - -from dstack.gateway.common import run_async -from dstack.gateway.errors import GatewayError - -CONFIGS_DIR = Path("/etc/nginx/sites-enabled") -GATEWAY_PORT = 8000 -logger = logging.getLogger(__name__) - - -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() - _lock: Lock = Lock() - - async def register_service(self, project: str, domain: str, sock_path: str, auth: bool): - logger.info("Registering service %s", domain) - async with self._lock: - if domain in self.domains: - raise GatewayError("Domain is already registered") - self.write_conf( - self.get_service_conf(project, domain, f"unix:{sock_path}", auth), - CONFIGS_DIR / f"443-{domain}.conf", - ) - self.domains.add(domain) - await run_async(self.reload) - - 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") - 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) - - 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") - 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(self.reload) - - @classmethod - def get_service_conf( - cls, project: str, domain: str, server: str, auth: bool, upstream: Optional[str] = None - ) -> 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, - 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, - ) - - @staticmethod - def reload(): - cmd = ["sudo", "systemctl", "reload", "nginx.service"] - r = subprocess.run(cmd) - if r.returncode != 0: - raise GatewayError("Failed to reload nginx") - - @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") - - @staticmethod - def run_certbot(domain: str): - logger.info("Running certbot for %s", domain) - cmd = ["sudo", "certbot", "certonly"] - cmd += ["--non-interactive", "--agree-tos", "--register-unsafely-without-email"] - cmd += ["--nginx", "--domain", domain] - r = subprocess.run(cmd, capture_output=True) - if r.returncode != 0: - raise GatewayError(f"Certbot failed:\n{r.stderr.decode()}") 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/__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/services/test_store.py b/gateway/src/tests/services/test_store.py deleted file mode 100644 index e53bfbb89c..0000000000 --- a/gateway/src/tests/services/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.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() -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 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/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..56bbe17ce0 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]] @@ -131,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): @@ -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/background/tasks/process_running_jobs.py b/src/dstack/_internal/server/background/tasks/process_running_jobs.py index 1f87dee7ab..aca4f231c5 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, @@ -22,7 +21,6 @@ ) from dstack._internal.server.db import get_session_ctx from dstack._internal.server.models import ( - GatewayModel, JobModel, ProjectModel, RepoModel, @@ -32,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 @@ -98,9 +97,9 @@ 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) server_ssh_private_key = project.ssh_private_key secrets = {} # TODO secrets @@ -147,10 +146,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 @@ -166,9 +161,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: @@ -210,11 +203,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 @@ -222,16 +211,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) - 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/background/tasks/process_runs.py b/src/dstack/_internal/server/background/tasks/process_runs.py index 100c8cd707..9f36ee12d1 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, @@ -26,6 +28,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, @@ -41,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, } @@ -106,7 +109,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 +120,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 max for auto-scaling + replicas = run.run_spec.configuration.replicas.min or 0 + + 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)) @@ -139,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() @@ -193,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 @@ -220,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( @@ -258,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: @@ -270,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 @@ -282,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..f2c72cca09 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,7 @@ 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) async with PROCESSING_POOL_LOCK: pool_instances = get_pool_instances(pool) @@ -130,6 +131,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 +192,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 4f7f4713bf..fbe7564f98 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -1,12 +1,15 @@ 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 +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, @@ -21,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 ( - JobProvisioningData, Run, RunSpec, ServiceModelSpec, @@ -31,6 +33,7 @@ from dstack._internal.server.models import ( GatewayComputeModel, GatewayModel, + JobModel, ProjectModel, RunModel, ) @@ -312,12 +315,15 @@ 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, ) + 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: @@ -325,48 +331,85 @@ 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_model: JobModel ): - 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) + 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_service, - run, - job_provisioning_data, + 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)) 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: + logger.debug("%s: unregistering service", fmt(run_model)) await run_async( conn.client.unregister_service, - project.name, - service_address, + project=project.name, + 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)) +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: + 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 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)) + + +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 786ef16251..c3778e441e 100644 --- a/src/dstack/_internal/server/services/gateways/client.py +++ b/src/dstack/_internal/server/services/gateways/client.py @@ -1,10 +1,10 @@ +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 JobSubmission, Run GATEWAY_MANAGEMENT_PORT = 8000 @@ -19,59 +19,92 @@ 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, + 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(project, entrypoint) + payload = { - "public_domain": urlparse(run.service.url).hostname, - "app_port": conf.port.container_port, - "auth": conf.auth, - "options": run.service.options, + "run_id": run_id.hex, + "domain": domain, + "auth": auth, + "options": 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/{project}/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, 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}", + } + ) + resp = self.s.post( - self._url(f"/api/registry/{project}/openai/register"), json={"domain": domain} + 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 unregister_service(self, project: str, public_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}/unregister"), json={"public_domain": public_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}/entrypoints/register"), + json={ + "module": "openai", + "domain": domain, + }, ) if resp.status_code == 400: raise gateway_error(resp.json()) @@ -89,5 +122,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']}") diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py index 5a4ff97b52..a78be713db 100644 --- a/src/dstack/_internal/server/services/jobs/__init__.py +++ b/src/dstack/_internal/server/services/jobs/__init__.py @@ -8,7 +8,8 @@ import sqlalchemy.orm as sa_orm from sqlalchemy.ext.asyncio import AsyncSession -from dstack._internal.core.errors import SSHError +import dstack._internal.server.services.gateways as gateways +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 ( @@ -51,15 +52,16 @@ TERMINATING_PROCESSING_JOBS_IDS = set() -def get_jobs_from_run_spec(run_spec: RunSpec) -> List[Job]: - job_configurator = _get_job_configurator(run_spec) - job_specs = job_configurator.get_job_specs() - 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) -> 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 @@ -91,6 +93,15 @@ def job_model_to_job_submission(job_model: JobModel) -> JobSubmission: ) +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 + raise ResourceNotFoundError( + f"Job with replica_num={replica_num} and job_num={job_num} not found" + ) + + async def terminate_job_provisioning_data_instance( project: ProjectModel, job_provisioning_data: JobProvisioningData ): @@ -206,7 +217,9 @@ 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 + ) # TODO(egor-s) ensure always runs finally: PROCESSING_POOL_IDS.remove(instance.id) diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index 1bc22fcb32..5830ca8b22 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(), @@ -51,7 +53,6 @@ def get_job_specs(self) -> List[JobSpec]: requirements=self._requirements(), retry_policy=self._retry_policy(), working_dir=self._working_dir(), - pool_name=self._pool_name(), ) return [job_spec] @@ -147,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 b175b4e833..a3c6e8e027 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(), @@ -342,18 +341,25 @@ async def submit_run( ) session.add(run_model) + replicas = 1 if run_spec.configuration.type == "service": + replicas = run_spec.configuration.replicas.min + 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") + await gateways.register_service(session, run_model) - 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) @@ -373,8 +379,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 +605,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) @@ -629,12 +638,12 @@ 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, 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..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="", ) @@ -210,13 +214,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=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", + 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/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/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 diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 58ad7b8bad..efb8759c64 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -128,10 +128,10 @@ 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, "registry_auth": None, "requirements": { "resources": { @@ -240,10 +240,10 @@ 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, "registry_auth": None, "requirements": { "resources": { @@ -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, } ]