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

dstack/_internal/server/services/secrets.py

@@ -0,0 +1,245 @@
+import re
+import uuid
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import Dict, List, Optional
+
+import sqlalchemy.exc
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.errors import (
+    ResourceExistsError,
+    ResourceNotExistsError,
+    ServerClientError,
+)
+from dstack._internal.core.models.secrets import Secret
+from dstack._internal.server.db import get_db
+from dstack._internal.server.models import DecryptedString, ProjectModel, SecretModel, UserModel
+from dstack._internal.server.services import events
+from dstack._internal.server.services.locking import get_locker
+
+_SECRET_NAME_REGEX = "^[A-Za-z0-9-_]{1,200}$"
+_SECRET_VALUE_MAX_LENGTH = 5000
+
+
+async def list_secrets(
+    session: AsyncSession,
+    project: ProjectModel,
+) -> List[Secret]:
+    secret_models = await list_project_secret_models(session=session, project=project)
+    return [secret_model_to_secret(s, include_value=False) for s in secret_models]
+
+
+async def get_project_secrets_mapping(
+    session: AsyncSession,
+    project: ProjectModel,
+) -> Dict[str, str]:
+    secret_models = await list_project_secret_models(session=session, project=project)
+    return {s.name: s.value.get_plaintext_or_error() for s in secret_models}
+
+
+async def get_secret(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+) -> Optional[Secret]:
+    secret_model = await get_project_secret_model_by_name(
+        session=session,
+        project=project,
+        name=name,
+    )
+    if secret_model is None:
+        return None
+    return secret_model_to_secret(secret_model, include_value=True)
+
+
+async def create_or_update_secret(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+    value: str,
+    user: UserModel,
+) -> Secret:
+    _validate_secret(name=name, value=value)
+    try:
+        secret_model = await create_secret(
+            session=session,
+            project=project,
+            name=name,
+            value=value,
+            user=user,
+        )
+    except ResourceExistsError:
+        secret_model = await update_secret(
+            session=session,
+            project=project,
+            name=name,
+            value=value,
+            user=user,
+        )
+    return secret_model_to_secret(secret_model, include_value=True)
+
+
+async def delete_secrets(
+    session: AsyncSession,
+    project: ProjectModel,
+    names: List[str],
+    user: UserModel,
+):
+    async with get_project_secret_models_by_name_for_update(
+        session=session, project=project, names=names
+    ) as secret_models:
+        existing_names = [s.name for s in secret_models]
+        missing_names = set(names) - set(existing_names)
+        if missing_names:
+            raise ResourceNotExistsError(f"Secrets not found: {', '.join(missing_names)}")
+        for secret_model in secret_models:
+            await session.delete(secret_model)
+            events.emit(
+                session,
+                "Secret deleted",
+                actor=events.UserActor.from_user(user),
+                targets=[events.Target.from_model(secret_model)],
+            )
+        await session.commit()
+
+
+def secret_model_to_secret(secret_model: SecretModel, include_value: bool = False) -> Secret:
+    value = None
+    if include_value:
+        value = secret_model.value.get_plaintext_or_error()
+    return Secret(
+        id=secret_model.id,
+        name=secret_model.name,
+        value=value,
+    )
+
+
+async def list_project_secret_models(
+    session: AsyncSession,
+    project: ProjectModel,
+) -> List[SecretModel]:
+    res = await session.execute(
+        select(SecretModel)
+        .where(
+            SecretModel.project_id == project.id,
+        )
+        .order_by(SecretModel.created_at.desc())
+    )
+    secret_models = list(res.scalars().all())
+    return secret_models
+
+
+async def get_project_secret_model_by_name(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+) -> Optional[SecretModel]:
+    res = await session.execute(
+        select(SecretModel).where(
+            SecretModel.project_id == project.id,
+            SecretModel.name == name,
+        )
+    )
+    return res.scalar_one_or_none()
+
+
+@asynccontextmanager
+async def get_project_secret_models_by_name_for_update(
+    session: AsyncSession, project: ProjectModel, names: list[str]
+) -> AsyncGenerator[list[SecretModel], None]:
+    """
+    Fetch secrets from the database and lock them for update.
+
+    **NOTE**: commit changes to the database before exiting from this context manager,
+              so that in-memory locks are only released after commit.
+    """
+    filters = [
+        SecretModel.project_id == project.id,
+        SecretModel.name.in_(names),
+    ]
+    res = await session.execute(select(SecretModel.id).where(*filters))
+    secret_ids = res.scalars().all()
+    if not secret_ids:
+        yield []
+    else:
+        async with get_locker(get_db().dialect_name).lock_ctx(
+            SecretModel.__tablename__, sorted(secret_ids)
+        ):
+            # Refetch after lock
+            res = await session.execute(
+                select(SecretModel)
+                .where(SecretModel.id.in_(secret_ids), *filters)
+                .with_for_update(key_share=True)
+                .order_by(SecretModel.id)  # take locks in order
+            )
+            yield list(res.scalars().all())
+
+
+async def create_secret(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+    value: str,
+    user: UserModel,
+) -> SecretModel:
+    secret_model = SecretModel(
+        id=uuid.uuid4(),
+        project_id=project.id,
+        name=name,
+        value=DecryptedString(plaintext=value),
+    )
+    try:
+        async with session.begin_nested():
+            session.add(secret_model)
+            events.emit(
+                session,
+                "Secret created",
+                actor=events.UserActor.from_user(user),
+                targets=[events.Target.from_model(secret_model)],
+            )
+    except sqlalchemy.exc.IntegrityError:
+        raise ResourceExistsError()
+    await session.commit()
+    return secret_model
+
+
+async def update_secret(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+    value: str,
+    user: UserModel,
+) -> SecretModel:
+    async with get_project_secret_models_by_name_for_update(
+        session=session, project=project, names=[name]
+    ) as secret_models:
+        if not secret_models:
+            raise ResourceNotExistsError()
+        secret_model = secret_models[0]
+        if secret_model.value.get_plaintext_or_error() != value:
+            secret_model.value = DecryptedString(plaintext=value)
+            events.emit(
+                session,
+                "Secret updated",
+                actor=events.UserActor.from_user(user),
+                targets=[events.Target.from_model(secret_model)],
+            )
+            await session.commit()
+    return secret_model
+
+
+def _validate_secret(name: str, value: str):
+    _validate_secret_name(name)
+    _validate_secret_value(value)
+
+
+def _validate_secret_name(name: str):
+    if re.match(_SECRET_NAME_REGEX, name) is None:
+        raise ServerClientError(f"Secret name should match regex '{_SECRET_NAME_REGEX}")
+
+
+def _validate_secret_value(value: str):
+    if len(value) > _SECRET_VALUE_MAX_LENGTH:
+        raise ServerClientError(f"Secret value length must not exceed {_SECRET_VALUE_MAX_LENGTH}")

dstack/_internal/server/services/services/__init__.py

@@ -0,0 +1,345 @@
+"""
+Application logic related to `type: service` runs.
+"""
+
+import json
+import uuid
+from datetime import datetime
+from typing import Optional
+
+import httpx
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
+
+import dstack._internal.server.services.jobs as jobs_services
+from dstack._internal.core.errors import (
+    GatewayError,
+    ResourceNotExistsError,
+    ServerClientError,
+    SSHError,
+)
+from dstack._internal.core.models.configurations import (
+    DEFAULT_REPLICA_GROUP_NAME,
+    SERVICE_HTTPS_DEFAULT,
+    ServiceConfiguration,
+)
+from dstack._internal.core.models.gateways import GatewayConfiguration, GatewayStatus
+from dstack._internal.core.models.instances import SSHConnectionParams
+from dstack._internal.core.models.runs import JobSpec, Run, RunSpec, ServiceModelSpec, ServiceSpec
+from dstack._internal.server import settings
+from dstack._internal.server.models import GatewayModel, JobModel, ProjectModel, RunModel
+from dstack._internal.server.services.gateways import (
+    get_gateway_configuration,
+    get_or_add_gateway_connection,
+    get_project_default_gateway_model,
+    get_project_gateway_model_by_name,
+)
+from dstack._internal.server.services.logging import fmt
+from dstack._internal.server.services.services.autoscalers import get_service_scaler
+from dstack._internal.server.services.services.options import get_service_options
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+async def register_service(session: AsyncSession, run_model: RunModel, run_spec: RunSpec):
+    assert isinstance(run_spec.configuration, ServiceConfiguration)
+
+    if isinstance(run_spec.configuration.gateway, str):
+        gateway = await get_project_gateway_model_by_name(
+            session=session, project=run_model.project, name=run_spec.configuration.gateway
+        )
+        if gateway is None:
+            raise ResourceNotExistsError(
+                f"Gateway {run_spec.configuration.gateway} does not exist"
+            )
+    elif run_spec.configuration.gateway == False:
+        gateway = None
+    else:
+        gateway = await get_project_default_gateway_model(
+            session=session, project=run_model.project
+        )
+        if gateway is None and run_spec.configuration.gateway == True:
+            raise ResourceNotExistsError(
+                "The service requires a gateway, but there is no default gateway in the project"
+            )
+
+    if gateway is not None:
+        service_spec = await _register_service_in_gateway(session, run_model, run_spec, gateway)
+        run_model.gateway = gateway
+    elif not settings.FORBID_SERVICES_WITHOUT_GATEWAY:
+        service_spec = _register_service_in_server(run_model, run_spec)
+    else:
+        raise ResourceNotExistsError(
+            "This dstack-server installation forbids services without a gateway."
+            " Please configure a gateway."
+        )
+    run_model.service_spec = service_spec.json()
+
+
+async def _register_service_in_gateway(
+    session: AsyncSession, run_model: RunModel, run_spec: RunSpec, gateway: GatewayModel
+) -> ServiceSpec:
+    assert run_spec.configuration.type == "service"
+
+    if gateway.gateway_compute is None:
+        raise ServerClientError("Gateway has no instance associated with it")
+
+    if gateway.status != GatewayStatus.RUNNING:
+        raise ServerClientError("Gateway status is not running")
+
+    gateway_configuration = get_gateway_configuration(gateway)
+    service_https = _get_service_https(run_spec, gateway_configuration)
+    router = gateway_configuration.router
+    service_protocol = "https" if service_https else "http"
+
+    if service_https and gateway_configuration.certificate is None:
+        raise ServerClientError(
+            "Cannot run HTTPS service on gateway with no SSL certificates configured"
+        )
+
+    gateway_https = _get_gateway_https(gateway_configuration)
+    gateway_protocol = "https" if gateway_https else "http"
+
+    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_spec = get_service_spec(
+        configuration=run_spec.configuration,
+        service_url=f"{service_protocol}://{run_model.run_name}.{wildcard_domain}",
+        model_url=f"{gateway_protocol}://gateway.{wildcard_domain}",
+    )
+
+    domain = service_spec.get_domain()
+    assert domain is not None
+
+    conn = await get_or_add_gateway_connection(session, gateway.id)
+    try:
+        logger.debug("%s: registering service as %s", fmt(run_model), service_spec.url)
+        async with conn.client() as client:
+            await client.register_service(
+                project=run_model.project.name,
+                run_name=run_model.run_name,
+                domain=domain,
+                service_https=service_https,
+                gateway_https=gateway_https,
+                auth=run_spec.configuration.auth,
+                client_max_body_size=settings.DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE,
+                options=service_spec.options,
+                rate_limits=run_spec.configuration.rate_limits,
+                ssh_private_key=run_model.project.ssh_private_key,
+                router=router,
+            )
+        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:
+        logger.debug("Gateway request failed", exc_info=True)
+        raise GatewayError(f"Gateway is not working: {e!r}")
+
+    return service_spec
+
+
+def _register_service_in_server(run_model: RunModel, run_spec: RunSpec) -> ServiceSpec:
+    assert run_spec.configuration.type == "service"
+    if run_spec.configuration.https != SERVICE_HTTPS_DEFAULT:
+        # Note: if the user sets `https: <default-value>`, it will be ignored silently
+        # TODO: in 0.19, make `https` Optional to be able to tell if it was set or omitted
+        raise ServerClientError(
+            "The `https` configuration property is not applicable when running services without a gateway."
+            " Please configure a gateway or remove the `https` property from the service configuration"
+        )
+    # Check if any group has autoscaling (min != max)
+    has_autoscaling = any(
+        group.count.min != group.count.max for group in run_spec.configuration.replica_groups
+    )
+    if has_autoscaling:
+        raise ServerClientError(
+            "Auto-scaling is not supported when running services without a gateway."
+            " Please configure a gateway or set `replicas` to a fixed value in the service configuration"
+        )
+    if run_spec.configuration.rate_limits:
+        raise ServerClientError(
+            "Rate limits are not supported when running services without a gateway."
+            " Please configure a gateway or remove `rate_limits` from the service configuration"
+        )
+    return get_service_spec(
+        configuration=run_spec.configuration,
+        service_url=f"/proxy/services/{run_model.project.name}/{run_model.run_name}/",
+        model_url=f"/proxy/models/{run_model.project.name}/",
+    )
+
+
+def get_service_spec(
+    configuration: ServiceConfiguration, service_url: str, model_url: str
+) -> ServiceSpec:
+    service_spec = ServiceSpec(url=service_url)
+    if configuration.model is not None:
+        service_spec.model = ServiceModelSpec(
+            name=configuration.model.name,
+            base_url=model_url,
+            type=configuration.model.type,
+        )
+        service_spec.options = get_service_options(configuration)
+    return service_spec
+
+
+async def register_replica(
+    session: AsyncSession,
+    gateway_id: Optional[uuid.UUID],
+    run: Run,
+    job_model: JobModel,
+    ssh_head_proxy: Optional[SSHConnectionParams],
+    ssh_head_proxy_private_key: Optional[str],
+):
+    if gateway_id is not None:
+        conn = await get_or_add_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)
+            async with conn.client() as client:
+                await client.register_replica(
+                    run=run,
+                    job_spec=JobSpec.__response__.parse_raw(job_model.job_spec_data),
+                    job_submission=job_submission,
+                    ssh_head_proxy=ssh_head_proxy,
+                    ssh_head_proxy_private_key=ssh_head_proxy_private_key,
+                )
+        except (httpx.RequestError, SSHError) as e:
+            logger.debug("Gateway request failed", exc_info=True)
+            raise GatewayError(repr(e))
+        except GatewayError as e:
+            if "already exists in service" in e.msg:
+                # Pre-0.19.25 servers never mark the job as `registered`, so 0.19.25+ servers may
+                # attempt to register a replica that is already registered on the gateway.
+                logger.warning(
+                    (
+                        "%s: could not register replica in gateway: %s."
+                        " NOTE: if you just updated dstack from pre-0.19.25 to 0.19.25+,"
+                        " expect to see this warning once for every running service replica"
+                    ),
+                    fmt(job_model),
+                    e.msg,
+                )
+            else:
+                raise
+    job_model.registered = True
+    logger.info(
+        "%s: service replica registered to receive requests, gateway=%s",
+        fmt(job_model),
+        gateway_id is not None,
+    )
+
+
+async def unregister_service(session: AsyncSession, run_model: RunModel):
+    if run_model.gateway_id is None:  # in-server proxy
+        return
+    conn = await get_or_add_gateway_connection(session, run_model.gateway_id)
+    res = await session.execute(
+        select(ProjectModel).where(ProjectModel.id == run_model.project_id)
+    )
+    project = res.scalar_one()
+    try:
+        logger.debug("%s: unregistering service", fmt(run_model))
+        async with conn.client() as client:
+            await client.unregister_service(
+                project=project.name,
+                run_name=run_model.run_name,
+            )
+        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:
+        logger.debug("Gateway request failed", exc_info=True)
+        raise GatewayError(repr(e))
+
+
+async def unregister_replica(session: AsyncSession, job_model: JobModel):
+    res = await session.execute(
+        select(RunModel)
+        .where(RunModel.id == job_model.run_id)
+        .options(joinedload(RunModel.project).joinedload(ProjectModel.backends))
+    )
+    run_model = res.unique().scalar_one()
+    if run_model.gateway_id is not None:
+        conn = await get_or_add_gateway_connection(session, run_model.gateway_id)
+        try:
+            logger.debug(
+                "%s: unregistering replica from service %s", fmt(job_model), job_model.run_id.hex
+            )
+            async with conn.client() as client:
+                await client.unregister_replica(
+                    project=run_model.project.name,
+                    run_name=run_model.run_name,
+                    job_id=job_model.id,
+                )
+        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:
+            logger.debug("Gateway request failed", exc_info=True)
+            raise GatewayError(repr(e))
+    job_model.registered = False
+    logger.info(
+        "%s: service replica unregistered from receiving requests, gateway=%s",
+        fmt(job_model),
+        run_model.gateway_id is not None,
+    )
+
+
+def _get_service_https(run_spec: RunSpec, configuration: GatewayConfiguration) -> bool:
+    assert run_spec.configuration.type == "service"
+    if not run_spec.configuration.https:
+        return False
+    if configuration.certificate is not None and configuration.certificate.type == "acm":
+        return False
+    return True
+
+
+def _get_gateway_https(configuration: GatewayConfiguration) -> bool:
+    if configuration.certificate is not None and configuration.certificate.type == "acm":
+        return False
+    if configuration.certificate is not None and configuration.certificate.type == "lets-encrypt":
+        return True
+    return False
+
+
+async def update_service_desired_replica_count(
+    session: AsyncSession,
+    run_model: RunModel,
+    configuration: ServiceConfiguration,
+    last_scaled_at: Optional[datetime],
+) -> None:
+    stats = None
+    if run_model.gateway_id is not None:
+        conn = await get_or_add_gateway_connection(session, run_model.gateway_id)
+        stats = await conn.get_stats(run_model.project.name, run_model.run_name)
+    replica_groups = configuration.replica_groups
+    desired_replica_counts = {}
+    total = 0
+    prev_counts = (
+        json.loads(run_model.desired_replica_counts) if run_model.desired_replica_counts else {}
+    )
+    if (
+        prev_counts == {}
+        and len(replica_groups) == 1
+        and replica_groups[0].name == DEFAULT_REPLICA_GROUP_NAME
+    ):
+        # Special case to avoid dropping the replica count to group.count.min
+        # when a 0.20.7+ server first processes a service created by a pre-0.20.7 server.
+        # TODO: remove once most users upgrade to 0.20.7+.
+        prev_counts = {DEFAULT_REPLICA_GROUP_NAME: run_model.desired_replica_count}
+    for group in replica_groups:
+        scaler = get_service_scaler(group.count, group.scaling)
+        assert group.name is not None, "Group name is always set"
+        group_desired = scaler.get_desired_count(
+            current_desired_count=prev_counts.get(group.name, group.count.min or 0),
+            stats=stats,
+            last_scaled_at=last_scaled_at,
+        )
+        desired_replica_counts[group.name] = group_desired
+        total += group_desired
+    run_model.desired_replica_counts = json.dumps(desired_replica_counts)
+    run_model.desired_replica_count = total