dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

dstack/_internal/server/app.py

@@ -0,0 +1,432 @@
+import asyncio
+import importlib.resources
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Annotated, Awaitable, Callable, List, Optional
+
+import sentry_sdk
+from fastapi import Depends, FastAPI, Request, Response, status
+from fastapi.datastructures import URL
+from fastapi.responses import HTMLResponse, RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from packaging.version import Version
+from prometheus_client import Counter, Histogram
+from sentry_sdk.types import SamplingContext
+
+from dstack._internal import settings as core_settings
+from dstack._internal.cli.utils.common import console
+from dstack._internal.core.errors import ForbiddenError, ServerClientError
+from dstack._internal.core.services.configs import update_default_project
+from dstack._internal.proxy.lib.deps import get_injector_from_app
+from dstack._internal.proxy.lib.routers import model_proxy
+from dstack._internal.server import settings
+from dstack._internal.server.background import start_background_tasks
+from dstack._internal.server.background.tasks.process_probes import PROBES_SCHEDULER
+from dstack._internal.server.db import get_db, get_session_ctx, migrate
+from dstack._internal.server.routers import (
+    auth,
+    backends,
+    events,
+    files,
+    fleets,
+    gateways,
+    gpus,
+    instances,
+    logs,
+    metrics,
+    projects,
+    prometheus,
+    repos,
+    runs,
+    secrets,
+    server,
+    users,
+    volumes,
+)
+from dstack._internal.server.services.config import ServerConfigManager
+from dstack._internal.server.services.gateways import gateway_connections_pool, init_gateways
+from dstack._internal.server.services.locking import advisory_lock_ctx
+from dstack._internal.server.services.projects import get_or_create_default_project
+from dstack._internal.server.services.proxy.deps import ServerProxyDependencyInjector
+from dstack._internal.server.services.proxy.routers import service_proxy
+from dstack._internal.server.services.storage import init_default_storage
+from dstack._internal.server.services.users import get_or_create_admin_user
+from dstack._internal.server.settings import (
+    DEFAULT_PROJECT_NAME,
+    DO_NOT_UPDATE_DEFAULT_PROJECT,
+    SERVER_CONFIG_FILE_PATH,
+    SERVER_URL,
+    UPDATE_DEFAULT_PROJECT,
+)
+from dstack._internal.server.utils import sentry_utils
+from dstack._internal.server.utils.logging import configure_logging
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    check_client_server_compatibility,
+    error_detail,
+    get_client_version,
+    get_server_client_error_details,
+)
+from dstack._internal.utils.logging import get_logger
+from dstack._internal.utils.ssh import check_required_ssh_version
+
+logger = get_logger(__name__)
+
+# Server HTTP metrics
+REQUESTS_TOTAL = Counter(
+    "dstack_server_requests_total",
+    "Total number of HTTP requests",
+    ["method", "endpoint", "http_status", "project_name"],
+)
+REQUEST_DURATION = Histogram(
+    "dstack_server_request_duration_seconds",
+    "HTTP request duration in seconds",
+    ["method", "endpoint", "http_status", "project_name"],
+)
+
+
+def create_app() -> FastAPI:
+    app = FastAPI(
+        docs_url="/api/docs",
+        lifespan=lifespan,
+        dependencies=[
+            Depends(_check_client_version),
+        ],
+    )
+    app.state.proxy_dependency_injector = ServerProxyDependencyInjector()
+    return app
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    configure_logging()
+    if settings.SENTRY_DSN is not None:
+        sentry_sdk.init(
+            dsn=settings.SENTRY_DSN,
+            release=core_settings.DSTACK_VERSION,
+            environment=settings.SERVER_ENVIRONMENT,
+            enable_tracing=True,
+            traces_sampler=_sentry_traces_sampler,
+            profiles_sample_rate=settings.SENTRY_PROFILES_SAMPLE_RATE,
+            before_send=sentry_utils.AsyncioCancelledErrorFilterEventProcessor(),
+        )
+    server_executor = ThreadPoolExecutor(max_workers=settings.SERVER_EXECUTOR_MAX_WORKERS)
+    asyncio.get_running_loop().set_default_executor(server_executor)
+    await migrate()
+    _print_dstack_logo()
+    if not check_required_ssh_version():
+        logger.warning("OpenSSH 8.4+ is required. The dstack server may not work properly")
+    server_config_manager = None
+    server_config_loaded = False
+    if settings.SERVER_CONFIG_ENABLED:
+        server_config_manager = ServerConfigManager()
+        server_config_loaded = server_config_manager.load_config()
+        # Encryption has to be configured before working with users and projects
+        await server_config_manager.apply_encryption()
+    async with get_session_ctx() as session:
+        async with advisory_lock_ctx(
+            bind=session,
+            dialect_name=get_db().dialect_name,
+            resource="server_init",
+        ):
+            admin, _ = await get_or_create_admin_user(session=session)
+            await get_or_create_default_project(
+                session=session,
+                user=admin,
+            )
+            if server_config_manager is not None:
+                server_config_dir = _get_server_config_dir()
+                if not server_config_loaded:
+                    logger.info("Initializing the default configuration...", {"show_path": False})
+                    await server_config_manager.init_config(session=session)
+                    logger.info(
+                        f"Initialized the default configuration at [link=file://{SERVER_CONFIG_FILE_PATH}]{server_config_dir}[/link]",
+                        {"show_path": False},
+                    )
+                else:
+                    logger.info(
+                        f"Applying [link=file://{SERVER_CONFIG_FILE_PATH}]{server_config_dir}[/link]...",
+                        {"show_path": False},
+                    )
+                    await server_config_manager.apply_config(session=session, owner=admin)
+            await init_gateways(session=session)
+    update_default_project(
+        project_name=DEFAULT_PROJECT_NAME,
+        url=SERVER_URL,
+        token=admin.token.get_plaintext_or_error(),
+        yes=UPDATE_DEFAULT_PROJECT,
+        no=DO_NOT_UPDATE_DEFAULT_PROJECT,
+    )
+    if settings.SERVER_S3_BUCKET is not None or settings.SERVER_GCS_BUCKET is not None:
+        init_default_storage()
+    scheduler = None
+    if settings.SERVER_BACKGROUND_PROCESSING_ENABLED:
+        scheduler = start_background_tasks()
+    else:
+        logger.info("Background processing is disabled")
+    PROBES_SCHEDULER.start()
+    dstack_version = (
+        core_settings.DSTACK_VERSION if core_settings.DSTACK_VERSION else "(no version)"
+    )
+    job_network_mode_log = (
+        logger.info
+        if settings.JOB_NETWORK_MODE != settings.DEFAULT_JOB_NETWORK_MODE
+        else logger.debug
+    )
+    job_network_mode_log(
+        "Job network mode: %s (%d)",
+        settings.JOB_NETWORK_MODE.name,
+        settings.JOB_NETWORK_MODE.value,
+    )
+    logger.info(f"The admin token is {admin.token.get_plaintext_or_error()}", {"show_path": False})
+    logger.info(
+        f"The dstack server {dstack_version} is running at {SERVER_URL}",
+        {"show_path": False},
+    )
+    for func in _ON_STARTUP_HOOKS:
+        await func(app)
+    yield
+    if scheduler is not None:
+        scheduler.shutdown()
+    PROBES_SCHEDULER.shutdown(wait=False)
+    await gateway_connections_pool.remove_all()
+    service_conn_pool = await get_injector_from_app(app).get_service_connection_pool()
+    await service_conn_pool.remove_all()
+    await get_db().engine.dispose()
+    # Let checked-out DB connections close as dispose() only closes checked-in connections
+    await asyncio.sleep(3)
+
+
+_ON_STARTUP_HOOKS = []
+
+
+def register_on_startup_hook(func: Callable[[FastAPI], Awaitable[None]]):
+    _ON_STARTUP_HOOKS.append(func)
+
+
+_NO_API_VERSION_CHECK_ROUTES = ["/api/docs"]
+
+
+def add_no_api_version_check_routes(paths: List[str]):
+    _NO_API_VERSION_CHECK_ROUTES.extend(paths)
+
+
+def register_routes(app: FastAPI, ui: bool = True):
+    app.include_router(server.router)
+    app.include_router(users.router)
+    app.include_router(auth.router)
+    app.include_router(projects.router)
+    app.include_router(backends.root_router)
+    app.include_router(backends.project_router)
+    app.include_router(fleets.root_router)
+    app.include_router(fleets.project_router)
+    app.include_router(instances.root_router)
+    app.include_router(instances.project_router)
+    app.include_router(repos.router)
+    app.include_router(runs.root_router)
+    app.include_router(runs.project_router)
+    app.include_router(gpus.project_router)
+    app.include_router(metrics.router)
+    app.include_router(logs.router)
+    app.include_router(secrets.router)
+    app.include_router(gateways.router)
+    app.include_router(volumes.root_router)
+    app.include_router(volumes.project_router)
+    app.include_router(service_proxy.router, prefix="/proxy/services", tags=["service-proxy"])
+    app.include_router(model_proxy.router, prefix="/proxy/models", tags=["model-proxy"])
+    app.include_router(prometheus.router)
+    app.include_router(files.router)
+    app.include_router(events.root_router)
+
+    @app.exception_handler(ForbiddenError)
+    async def forbidden_error_handler(request: Request, exc: ForbiddenError):
+        msg = "Access denied"
+        if len(exc.args) > 0:
+            msg = exc.args[0]
+        return CustomORJSONResponse(
+            status_code=status.HTTP_403_FORBIDDEN,
+            content=error_detail(msg),
+        )
+
+    @app.exception_handler(ServerClientError)
+    async def server_client_error_handler(request: Request, exc: ServerClientError):
+        return CustomORJSONResponse(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            content={"detail": get_server_client_error_details(exc)},
+        )
+
+    @app.exception_handler(OSError)
+    async def os_error_handler(request, exc: OSError):
+        if exc.errno in [36, 63]:
+            return CustomORJSONResponse(
+                {"detail": "Filename too long"},
+                status_code=status.HTTP_400_BAD_REQUEST,
+            )
+        raise exc
+
+    @app.middleware("http")
+    async def log_request(request: Request, call_next):
+        start_time = time.time()
+        response: Response = await call_next(request)
+        process_time = time.time() - start_time
+        # log process_time to be used in the log_http_metrics middleware
+        request.state.process_time = process_time
+        logger.debug(
+            "Processed request %s %s in %s. Status: %s",
+            request.method,
+            request.url,
+            f"{process_time:0.6f}s",
+            response.status_code,
+        )
+        return response
+
+    if settings.SERVER_PROFILING_ENABLED:
+        from pyinstrument import Profiler
+
+        @app.middleware("http")
+        async def profile_request(request: Request, call_next):
+            profiling = request.query_params.get("profile", False)
+            if profiling:
+                profiler = Profiler()
+                profiler.start()
+                respone = await call_next(request)
+                profiler.stop()
+                with open("profiling_results.html", "w+") as f:
+                    f.write(profiler.output_html())
+                return respone
+            else:
+                return await call_next(request)
+
+    # this middleware must be defined after the log_request middleware
+    @app.middleware("http")
+    async def log_http_metrics(request: Request, call_next):
+        def _extract_project_name(request: Request):
+            project_name = None
+            prefix = "/api/project/"
+            if request.url.path.startswith(prefix):
+                rest = request.url.path[len(prefix) :]
+                project_name = rest.split("/", 1)[0] if rest else None
+
+            return project_name
+
+        def _extract_endpoint_label(request: Request, response: Response) -> str:
+            route = request.scope.get("route")
+            route_path = getattr(route, "path", None)
+            if route_path:
+                return route_path
+            if not request.url.path.startswith("/api/"):
+                return "__non_api__"
+            if response.status_code == status.HTTP_404_NOT_FOUND:
+                return "__not_found__"
+            return "__unmatched__"
+
+        project_name = _extract_project_name(request)
+        response: Response = await call_next(request)
+        endpoint_label = _extract_endpoint_label(request, response)
+
+        REQUEST_DURATION.labels(
+            method=request.method,
+            endpoint=endpoint_label,
+            http_status=response.status_code,
+            project_name=project_name,
+        ).observe(request.state.process_time)
+
+        REQUESTS_TOTAL.labels(
+            method=request.method,
+            endpoint=endpoint_label,
+            http_status=response.status_code,
+            project_name=project_name,
+        ).inc()
+        return response
+
+    @app.get("/healthcheck")
+    async def healthcheck():
+        return CustomORJSONResponse(content={"status": "running"})
+
+    if ui and Path(__file__).parent.joinpath("statics").exists():
+        app.mount(
+            "/", StaticFiles(packages=["dstack._internal.server"], html=True), name="statics"
+        )
+
+        @app.exception_handler(404)
+        async def custom_http_exception_handler(request, exc):
+            if (
+                request.url.path.startswith("/api")
+                or _is_proxy_request(request)
+                or _is_prometheus_request(request)
+            ):
+                return CustomORJSONResponse(
+                    {"detail": exc.detail},
+                    status_code=status.HTTP_404_NOT_FOUND,
+                )
+            else:
+                return HTMLResponse(
+                    importlib.resources.files("dstack._internal.server")
+                    .joinpath("statics/index.html")
+                    .read_text()
+                )
+
+    else:
+
+        @app.get("/")
+        async def index():
+            return RedirectResponse("/api/docs")
+
+
+def _check_client_version(
+    request: Request, client_version: Annotated[Optional[Version], Depends(get_client_version)]
+) -> None:
+    if (
+        request.url.path.startswith("/api/")
+        and request.url.path not in _NO_API_VERSION_CHECK_ROUTES
+    ):
+        check_client_server_compatibility(
+            client_version=client_version,
+            server_version=core_settings.DSTACK_VERSION,
+        )
+
+
+def _is_proxy_request(request: Request) -> bool:
+    if request.url.path.startswith("/proxy"):
+        return True
+    # Attempt detecting requests originating from services proxied by dstack-proxy.
+    # Such requests can "leak" to dstack server paths if the service does not support
+    # running under a path prefix properly.
+    referrer = URL(request.headers.get("Referer", ""))
+    return (
+        referrer.netloc == "" or referrer.netloc == request.url.netloc
+    ) and referrer.path.startswith("/proxy")
+
+
+def _is_prometheus_request(request: Request) -> bool:
+    return request.url.path.startswith("/metrics")
+
+
+def _sentry_traces_sampler(sampling_context: SamplingContext) -> float:
+    parent_sampling_decision = sampling_context["parent_sampled"]
+    if parent_sampling_decision is not None:
+        return float(parent_sampling_decision)
+    transaction_context = sampling_context["transaction_context"]
+    name = transaction_context.get("name")
+    if name is not None:
+        if name.startswith("background."):
+            return settings.SENTRY_TRACES_BACKGROUND_SAMPLE_RATE
+    return settings.SENTRY_TRACES_SAMPLE_RATE
+
+
+def _print_dstack_logo():
+    console.print(
+        r"""[purple]     _     _             _
+  __| |___| |_ __ _  ___| | __  ___  ___ _ ____   _____ _ __
+ / _` / __| __/ _` |/ __| |/ / / __|/ _ \ '__\ \ / / _ \ '__|
+| (_| \__ \ || (_| | (__|   <  \__ \  __/ |   \ V /  __/ |
+ \__,_|___/\__\__,_|\___|_|\_\ |___/\___|_|    \_/ \___|_|
+[/]"""
+    )
+
+
+def _get_server_config_dir() -> str:
+    return str(SERVER_CONFIG_FILE_PATH).replace(os.path.expanduser("~"), "~", 1)

dstack/_internal/server/background/__init__.py

@@ -0,0 +1,142 @@
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+from apscheduler.triggers.interval import IntervalTrigger
+
+from dstack._internal.server import settings
+from dstack._internal.server.background.tasks.process_compute_groups import process_compute_groups
+from dstack._internal.server.background.tasks.process_events import delete_events
+from dstack._internal.server.background.tasks.process_fleets import process_fleets
+from dstack._internal.server.background.tasks.process_gateways import (
+    process_gateways,
+    process_gateways_connections,
+)
+from dstack._internal.server.background.tasks.process_idle_volumes import process_idle_volumes
+from dstack._internal.server.background.tasks.process_instances import (
+    delete_instance_health_checks,
+    process_instances,
+)
+from dstack._internal.server.background.tasks.process_metrics import (
+    collect_metrics,
+    delete_metrics,
+)
+from dstack._internal.server.background.tasks.process_placement_groups import (
+    process_placement_groups,
+)
+from dstack._internal.server.background.tasks.process_probes import process_probes
+from dstack._internal.server.background.tasks.process_prometheus_metrics import (
+    collect_prometheus_metrics,
+    delete_prometheus_metrics,
+)
+from dstack._internal.server.background.tasks.process_running_jobs import process_running_jobs
+from dstack._internal.server.background.tasks.process_runs import process_runs
+from dstack._internal.server.background.tasks.process_submitted_jobs import process_submitted_jobs
+from dstack._internal.server.background.tasks.process_terminating_jobs import (
+    process_terminating_jobs,
+)
+from dstack._internal.server.background.tasks.process_volumes import process_submitted_volumes
+
+_scheduler = AsyncIOScheduler()
+
+
+def get_scheduler() -> AsyncIOScheduler:
+    return _scheduler
+
+
+def start_background_tasks() -> AsyncIOScheduler:
+    # Background processing is implemented via in-memory locks on SQLite
+    # and SELECT FOR UPDATE on Postgres. Locks may be held for a long time.
+    # This is currently the main bottleneck for scaling dstack processing
+    # as processing more resources requires more DB connections.
+    # TODO: Make background processing efficient by committing locks to DB
+    # and processing outside of DB transactions.
+    #
+    # Now we just try to process as many resources as possible without exhausting DB connections.
+    #
+    # Quick tasks can process multiple resources per transaction.
+    # Potentially long tasks process one resource per transaction
+    # to avoid holding locks for all the resources if one is slow to process.
+    # Still, the next batch won't be processed unless all resources are processed,
+    # so larger batches do not increase processing rate linearly.
+    #
+    # The interval, batch_size, and max_instances determine background tasks processing rates.
+    # By default, one server replica can handle:
+    #
+    # * 150 active jobs with 2 minutes processing latency
+    # * 150 active runs with 2 minutes processing latency
+    # * 150 active instances with 2 minutes processing latency
+    #
+    # These latency numbers do not account for provisioning time,
+    # so it may be slower if a backend is slow to provision.
+    #
+    # Users can set SERVER_BACKGROUND_PROCESSING_FACTOR to process more resources per replica.
+    # They also need to increase max db connections on the client side and db side.
+    #
+    # In-memory locking via locksets does not guarantee
+    # that the first waiting for the lock will acquire it.
+    # The jitter is needed to give all tasks a chance to acquire locks.
+
+    _scheduler.add_job(process_probes, IntervalTrigger(seconds=3, jitter=1))
+    _scheduler.add_job(collect_metrics, IntervalTrigger(seconds=10), max_instances=1)
+    _scheduler.add_job(delete_metrics, IntervalTrigger(minutes=5), max_instances=1)
+    _scheduler.add_job(delete_events, IntervalTrigger(minutes=7), max_instances=1)
+    if settings.ENABLE_PROMETHEUS_METRICS:
+        _scheduler.add_job(
+            collect_prometheus_metrics, IntervalTrigger(seconds=10), max_instances=1
+        )
+        _scheduler.add_job(delete_prometheus_metrics, IntervalTrigger(minutes=5), max_instances=1)
+    _scheduler.add_job(process_gateways_connections, IntervalTrigger(seconds=15))
+    _scheduler.add_job(process_gateways, IntervalTrigger(seconds=10, jitter=2), max_instances=5)
+    _scheduler.add_job(
+        process_submitted_volumes, IntervalTrigger(seconds=10, jitter=2), max_instances=5
+    )
+    _scheduler.add_job(
+        process_idle_volumes, IntervalTrigger(seconds=60, jitter=10), max_instances=1
+    )
+    _scheduler.add_job(process_placement_groups, IntervalTrigger(seconds=30, jitter=5))
+    _scheduler.add_job(
+        process_fleets,
+        IntervalTrigger(seconds=10, jitter=2),
+        max_instances=1,
+    )
+    _scheduler.add_job(delete_instance_health_checks, IntervalTrigger(minutes=5), max_instances=1)
+    for replica in range(settings.SERVER_BACKGROUND_PROCESSING_FACTOR):
+        # Add multiple copies of tasks if requested.
+        # max_instances=1 for additional copies to avoid running too many tasks.
+        # Move other tasks here when they need per-replica scaling.
+        _scheduler.add_job(
+            process_submitted_jobs,
+            IntervalTrigger(seconds=4, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=4 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_running_jobs,
+            IntervalTrigger(seconds=4, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_terminating_jobs,
+            IntervalTrigger(seconds=4, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_runs,
+            IntervalTrigger(seconds=2, jitter=1),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_instances,
+            IntervalTrigger(seconds=4, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_compute_groups,
+            IntervalTrigger(seconds=15, jitter=2),
+            kwargs={"batch_size": 1},
+            max_instances=2 if replica == 0 else 1,
+        )
+    _scheduler.start()
+    return _scheduler

dstack/_internal/server/background/tasks/common.py

@@ -0,0 +1,24 @@
+from datetime import timedelta
+
+from dstack._internal.core.models.backends.base import BackendType
+
+
+def get_provisioning_timeout(backend_type: BackendType, instance_type_name: str) -> timedelta:
+    """
+    This timeout is used in a few places, but roughly refers to the max time between
+    requesting instance creation and the instance becoming ready to accept jobs.
+    For container-based backends, this also includes the image pulling time.
+    """
+    if backend_type == BackendType.LAMBDA:
+        return timedelta(minutes=30)
+    if backend_type == BackendType.RUNPOD:
+        return timedelta(minutes=20)
+    if backend_type == BackendType.KUBERNETES:
+        return timedelta(minutes=20)
+    if backend_type == BackendType.OCI and instance_type_name.startswith("BM."):
+        return timedelta(minutes=20)
+    if backend_type == BackendType.VULTR and instance_type_name.startswith("vbm"):
+        return timedelta(minutes=55)
+    if backend_type == BackendType.GCP and instance_type_name == "a4-highgpu-8g":
+        return timedelta(minutes=16)
+    return timedelta(minutes=10)