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

dstack/_internal/server/services/config.py

@@ -0,0 +1,279 @@
+from typing import List, Optional
+
+import yaml
+from pydantic import Field, ValidationError
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import Annotated
+
+import dstack._internal.core.backends.configurators
+from dstack._internal.core.backends.models import (
+    AnyBackendConfigWithCreds,
+    AnyBackendFileConfigWithCreds,
+    BackendInfoYAML,
+)
+from dstack._internal.core.errors import (
+    BackendNotAvailable,
+    ResourceNotExistsError,
+    ServerClientError,
+)
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.server import settings
+from dstack._internal.server.models import ProjectModel, UserModel
+from dstack._internal.server.services import backends as backends_services
+from dstack._internal.server.services import encryption as encryption_services
+from dstack._internal.server.services import projects as projects_services
+from dstack._internal.server.services.backends.handlers import delete_backends_safe
+from dstack._internal.server.services.encryption import AnyEncryptionKeyConfig
+from dstack._internal.server.services.permissions import (
+    DefaultPermissions,
+    set_default_permissions,
+)
+from dstack._internal.server.services.plugins import load_plugins
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+# By default, PyYAML chooses the style of a collection depending on whether it has nested collections.
+# If a collection has nested collections, it will be assigned the block style. Otherwise it will have the flow style.
+#
+# We want mapping to always be displayed in block-style but lists without nested objects in flow-style.
+# So we define a custom representer.
+
+
+def seq_representer(dumper, sequence):
+    flow_style = len(sequence) == 0 or isinstance(sequence[0], str) or isinstance(sequence[0], int)
+    return dumper.represent_sequence("tag:yaml.org,2002:seq", sequence, flow_style)
+
+
+yaml.add_representer(list, seq_representer)
+
+
+BackendFileConfigWithCreds = Annotated[
+    AnyBackendFileConfigWithCreds, Field(..., discriminator="type")
+]
+
+
+class ProjectConfig(CoreModel):
+    name: Annotated[str, Field(description="The name of the project")]
+    backends: Annotated[
+        Optional[List[BackendFileConfigWithCreds]], Field(description="The list of backends")
+    ] = None
+
+
+EncryptionKeyConfig = Annotated[AnyEncryptionKeyConfig, Field(..., discriminator="type")]
+
+
+class EncryptionConfig(CoreModel):
+    keys: Annotated[List[EncryptionKeyConfig], Field(description="The encryption keys")]
+
+
+class ServerConfig(CoreModel):
+    projects: Annotated[List[ProjectConfig], Field(description="The list of projects")]
+    encryption: Annotated[
+        Optional[EncryptionConfig], Field(description="The encryption config")
+    ] = None
+    default_permissions: Annotated[
+        Optional[DefaultPermissions], Field(description="The default user permissions")
+    ] = None
+    plugins: Annotated[
+        Optional[List[str]], Field(description="The server-side plugins to enable")
+    ] = None
+
+
+class ServerConfigManager:
+    def load_config(self) -> bool:
+        self.config = self._load_config()
+        return self.config is not None
+
+    async def init_config(self, session: AsyncSession):
+        """
+        Initializes the default server/config.yml.
+        The default config is empty or contains an existing `main` project config.
+        """
+        self.config = await self._init_config(session)
+        if self.config is not None:
+            self._save_config(self.config)
+
+    async def sync_config(self, session: AsyncSession):
+        # Disable config.yml sync for https://github.com/dstackai/dstack/issues/815.
+        return
+
+    async def apply_encryption(self):
+        if self.config is None:
+            logger.info("No server/config.yml. Skipping encryption configuration.")
+            return
+        if self.config.encryption is not None:
+            encryption_services.init_encryption_keys(self.config.encryption.keys)
+
+    async def apply_config(self, session: AsyncSession, owner: UserModel):
+        if self.config is None:
+            raise ValueError("Config is not loaded")
+        if self.config.default_permissions is not None:
+            set_default_permissions(self.config.default_permissions)
+        for project_config in self.config.projects:
+            await self._apply_project_config(
+                session=session, owner=owner, project_config=project_config
+            )
+        load_plugins(enabled_plugins=self.config.plugins or [])
+
+    async def _apply_project_config(
+        self,
+        session: AsyncSession,
+        owner: UserModel,
+        project_config: ProjectConfig,
+    ):
+        project = await projects_services.get_project_model_by_name(
+            session=session,
+            project_name=project_config.name,
+        )
+        if not project:
+            await projects_services.create_project_model(
+                session=session, owner=owner, project_name=project_config.name
+            )
+            project = await projects_services.get_project_model_by_name_or_error(
+                session=session, project_name=project_config.name
+            )
+        backends_to_delete = set(
+            dstack._internal.core.backends.configurators.list_available_backend_types()
+        )
+        for backend_file_config in project_config.backends or []:
+            backend_config = file_config_to_config(backend_file_config)
+            backend_type = BackendType(backend_config.type)
+            backends_to_delete.difference_update([backend_type])
+            try:
+                current_backend_config = await backends_services.get_backend_config(
+                    project=project,
+                    backend_type=backend_type,
+                )
+            except BackendNotAvailable:
+                logger.warning(
+                    "Backend %s not available and won't be configured."
+                    " Check that backend dependencies are installed.",
+                    backend_type.value,
+                )
+                continue
+            if backend_config == current_backend_config:
+                continue
+            backend_exists = any(backend_type == b.type for b in project.backends)
+            try:
+                # current_backend_config may be None if backend exists
+                # but it's config is invalid (e.g. cannot be decrypted).
+                # Update backend in this case.
+                if current_backend_config is None and not backend_exists:
+                    await backends_services.create_backend(
+                        session=session, project=project, config=backend_config
+                    )
+                else:
+                    await backends_services.update_backend(
+                        session=session, project=project, config=backend_config
+                    )
+            except Exception as e:
+                logger.warning("Failed to configure backend %s: %s", backend_config.type, e)
+        await delete_backends_safe(
+            session=session,
+            project=project,
+            backends_types=list(backends_to_delete),
+            error=False,
+        )
+
+    async def _init_config(self, session: AsyncSession) -> Optional[ServerConfig]:
+        project = await projects_services.get_project_model_by_name(
+            session=session,
+            project_name=settings.DEFAULT_PROJECT_NAME,
+        )
+        if project is None:
+            return None
+        # Force project reload to reflect updates when syncing
+        await session.refresh(project)
+        backends = []
+        for (
+            backend_type
+        ) in dstack._internal.core.backends.configurators.list_available_backend_types():
+            backend_config = await backends_services.get_backend_config(
+                project=project, backend_type=backend_type
+            )
+            if backend_config is not None:
+                backends.append(backend_config)
+        return ServerConfig(
+            projects=[ProjectConfig(name=settings.DEFAULT_PROJECT_NAME, backends=backends)],
+            encryption=EncryptionConfig(keys=[]),
+            default_permissions=None,
+        )
+
+    def _load_config(self) -> Optional[ServerConfig]:
+        try:
+            with open(settings.SERVER_CONFIG_FILE_PATH) as f:
+                content = f.read()
+        except OSError:
+            return
+        config_dict = yaml.load(content, yaml.FullLoader)
+        return ServerConfig.parse_obj(config_dict)
+
+    def _save_config(self, config: ServerConfig):
+        with open(settings.SERVER_CONFIG_FILE_PATH, "w+") as f:
+            f.write(config_to_yaml(config))
+
+
+async def get_backend_config_yaml(
+    project: ProjectModel, backend_type: BackendType
+) -> BackendInfoYAML:
+    backend_config = await backends_services.get_backend_config(
+        project=project, backend_type=backend_type
+    )
+    if backend_config is None:
+        raise ResourceNotExistsError()
+    config_yaml = config_to_yaml(backend_config)
+    return BackendInfoYAML(
+        name=backend_type,
+        config_yaml=config_yaml,
+    )
+
+
+async def create_backend_config_yaml(
+    session: AsyncSession,
+    project: ProjectModel,
+    config_yaml: str,
+):
+    config = config_yaml_to_backend_config(config_yaml)
+    await backends_services.create_backend(session=session, project=project, config=config)
+
+
+async def update_backend_config_yaml(
+    session: AsyncSession,
+    project: ProjectModel,
+    config_yaml: str,
+):
+    config = config_yaml_to_backend_config(config_yaml)
+    await backends_services.update_backend(session=session, project=project, config=config)
+
+
+class _BackendConfigWithCreds(CoreModel):
+    """
+    Model for parsing API and file YAML configs.
+    """
+
+    __root__: Annotated[AnyBackendConfigWithCreds, Field(..., discriminator="type")]
+
+
+def config_yaml_to_backend_config(config_yaml: str) -> AnyBackendConfigWithCreds:
+    try:
+        config_dict = yaml.load(config_yaml, yaml.FullLoader)
+    except yaml.YAMLError:
+        raise ServerClientError("Error parsing YAML")
+    try:
+        backend_config = _BackendConfigWithCreds.parse_obj(config_dict).__root__
+    except ValidationError as e:
+        raise ServerClientError(str(e))
+    return backend_config
+
+
+def file_config_to_config(file_config: AnyBackendFileConfigWithCreds) -> AnyBackendConfigWithCreds:
+    backend_config_dict = file_config.dict()
+    backend_config = _BackendConfigWithCreds.parse_obj(backend_config_dict)
+    return backend_config.__root__
+
+
+def config_to_yaml(config: CoreModel) -> str:
+    return yaml.dump(config.dict(exclude_none=True), sort_keys=False)

dstack/_internal/server/services/docker.py

@@ -0,0 +1,162 @@
+import re
+from dataclasses import dataclass
+from typing import List, Optional
+
+import requests
+from dxf import DXF
+from dxf.exceptions import DXFError
+from pydantic import Field, ValidationError, validator
+from typing_extensions import Annotated
+
+from dstack._internal.core.errors import DockerRegistryError
+from dstack._internal.core.models.common import (
+    CoreModel,
+    FrozenCoreModel,
+    RegistryAuth,
+)
+from dstack._internal.server.utils.common import join_byte_stream_checked
+
+DEFAULT_PLATFORM = "linux/amd64"
+DEFAULT_REGISTRY = "index.docker.io"
+MAX_CONFIG_OBJECT_SIZE = 2**22  # 4 MiB
+REGISTRY_REQUEST_TIMEOUT = 20
+
+
+@dataclass
+class DXFAuthAdapter:
+    registry_auth: Optional[RegistryAuth]
+
+    def __call__(self, dxf: DXF, response: requests.Response) -> None:
+        dxf.authenticate(
+            username=self.registry_auth.username if self.registry_auth else None,
+            password=self.registry_auth.password if self.registry_auth else None,
+            response=response,
+        )
+
+
+class DockerImage(FrozenCoreModel):
+    image: str
+    registry: Optional[str] = None
+    repo: str
+    tag: str
+    digest: Optional[str] = None
+
+
+class ImageConfig(CoreModel):
+    user: Annotated[Optional[str], Field(alias="User")] = None
+    entrypoint: Annotated[Optional[List[str]], Field(alias="Entrypoint")] = None
+    cmd: Annotated[Optional[List[str]], Field(alias="Cmd")] = None
+
+    @validator("user")
+    def normalize_user(cls, v: Optional[str]) -> Optional[str]:
+        # If USER is not set, the corresponding field may be missing or set to an empty string
+        if v == "":
+            return None
+        return v
+
+
+class ImageConfigObject(CoreModel):
+    config: ImageConfig = ImageConfig()
+
+    @validator("config", pre=True)
+    def config_set_default_if_null(cls, value):
+        return ImageConfig() if value is None else value
+
+
+class ImageManifestConfigField(CoreModel):
+    digest: str
+
+
+class ImageManifest(CoreModel):
+    config: ImageManifestConfigField
+
+
+def get_image_config(image_name: str, registry_auth: Optional[RegistryAuth]) -> ImageConfigObject:
+    image = parse_image_name(image_name)
+
+    registry_client = DXF(
+        host=image.registry or DEFAULT_REGISTRY,
+        repo=image.repo,
+        auth=DXFAuthAdapter(registry_auth),  # type: ignore[assignment]
+        timeout=REGISTRY_REQUEST_TIMEOUT,
+    )
+
+    with registry_client:
+        try:
+            manifest_resp = registry_client.get_manifest(
+                alias=image.digest or image.tag, platform=DEFAULT_PLATFORM
+            )
+            manifest = ImageManifest.__response__.parse_raw(manifest_resp)
+            config_stream = registry_client.pull_blob(manifest.config.digest)
+            config_resp = join_byte_stream_checked(config_stream, MAX_CONFIG_OBJECT_SIZE)  # type: ignore[arg-type]
+            if config_resp is None:
+                raise DockerRegistryError(
+                    f"Image config object exceeds the size limit of {MAX_CONFIG_OBJECT_SIZE} bytes"
+                )
+            return ImageConfigObject.__response__.parse_raw(config_resp)
+
+        except (DXFError, requests.RequestException, ValidationError) as e:
+            raise DockerRegistryError(e)
+
+
+def parse_image_name(image: str) -> DockerImage:
+    """
+    :param image: docker image name
+    :return: registry host, repo, tag, digest
+
+    >>> parse_image_name("ubuntu:22.04")
+    DockerImage(registry=None, repo='library/ubuntu', tag='22.04', digest=None)
+    >>> parse_image_name("dstackai/miniforge:py3.9-0.2")
+    DockerImage(registry=None, repo='dstackai/miniforge', tag='py3.9-0.2', digest=None)
+    >>> parse_image_name("ghcr.io/dstackai/miniforge")
+    DockerImage(registry='ghcr.io', repo='dstackai/miniforge', tag='latest', digest=None)
+    >>> parse_image_name("dstackai/miniforge@sha256:a4ba18a847a172a248d68faf6689e69fae4779b90b250211b79a26d21ddd6a15")
+    DockerImage(registry=None, repo='dstackai/miniforge', tag='latest', digest='sha256:a4ba18a847a172a248d68faf6689e69fae4779b90b250211b79a26d21ddd6a15')
+    """
+
+    digest = None
+    if "@" in image.split("/")[-1]:
+        image, digest = image.rsplit("@", maxsplit=1)
+
+    tag = "latest"
+    if ":" in image.split("/")[-1]:  # avoid detecting port as a tag
+        image, tag = image.rsplit(":", maxsplit=1)
+
+    registry = None
+    components = image.split("/")
+    if len(components) == 1:  # default registry, official image
+        repo = "library/" + components[0]
+    elif not is_host(components[0]):  # default registry, custom image
+        repo = "/".join(components)
+    else:  # custom registry
+        registry = components[0]
+        repo = "/".join(components[1:])
+
+    return DockerImage(image=image, registry=registry, repo=repo, tag=tag, digest=digest)
+
+
+def is_host(s: str) -> bool:
+    """
+    >>> is_host("localhost")
+    True
+    >>> is_host("localhost:5000")
+    True
+    >>> is_host("ghcr.io")
+    True
+    >>> is_host("127.0.0.1")
+    True
+    >>> is_host("dstackai")
+    False
+    """
+    return s == "localhost" or ":" in s or "." in s
+
+
+DOCKER_TARGET_PATH_PATTERN = re.compile(r"^(/[^/\0]*)+/?$")
+
+
+def is_valid_docker_volume_target(path: str) -> bool:
+    if not path.startswith("/"):
+        return False
+    if path.endswith("/") and path != "/":
+        return False
+    return DOCKER_TARGET_PATH_PATTERN.match(path) is not None

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

@@ -0,0 +1,102 @@
+from contextlib import contextmanager
+from typing import List, Tuple, Union
+
+from dstack._internal.core.errors import DstackError
+from dstack._internal.server.models import EncryptedString
+from dstack._internal.server.services.encryption.keys.aes import (
+    AESEncryptionKey,
+    AESEncryptionKeyConfig,
+)
+from dstack._internal.server.services.encryption.keys.base import EncryptionKey
+from dstack._internal.server.services.encryption.keys.identity import (
+    IdentityEncryptionKey,
+    IdentityEncryptionKeyConfig,
+)
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class EncryptionError(DstackError):
+    pass
+
+
+AnyEncryptionKeyConfig = Union[
+    AESEncryptionKeyConfig,
+    IdentityEncryptionKeyConfig,
+]
+
+
+_ENCRYPTION_KEY_CLASSES = [
+    IdentityEncryptionKey,
+    AESEncryptionKey,
+]
+_ENCRYPTION_KEY_TYPE_TO_ENCRYPTION_KEY_CLASS = {c.TYPE: c for c in _ENCRYPTION_KEY_CLASSES}
+
+
+# TODO: Introduce EncryptionKeyConfigurator to support external providers
+def get_encryption_key(config: AnyEncryptionKeyConfig) -> EncryptionKey:
+    return _ENCRYPTION_KEY_TYPE_TO_ENCRYPTION_KEY_CLASS[config.type](config)
+
+
+def get_identity_encryption_key() -> IdentityEncryptionKey:
+    return IdentityEncryptionKey(IdentityEncryptionKeyConfig())
+
+
+_encryption_keys = [get_identity_encryption_key()]
+
+
+def init_encryption_keys(encryption_key_configs: List[AnyEncryptionKeyConfig]):
+    global _encryption_keys
+    _encryption_keys = [get_encryption_key(c) for c in encryption_key_configs]
+    if not any(isinstance(key, IdentityEncryptionKey) for key in _encryption_keys):
+        _encryption_keys.append(get_identity_encryption_key())
+
+
+@contextmanager
+def encryption_keys_context(encryption_keys: List[EncryptionKey]):
+    """
+    A helper context manager to be used in tests. It's not concurrency-safe.
+    """
+    global _encryption_keys
+    prev_encryption_keys = _encryption_keys
+    _encryption_keys = encryption_keys
+    try:
+        yield
+    finally:
+        _encryption_keys = prev_encryption_keys
+
+
+def encrypt(plaintext: str) -> str:
+    key = _encryption_keys[0]
+    ciphertext = key.encrypt(plaintext)
+    packed_ciphertext = _pack_ciphertext(ciphertext, key_type=key.TYPE, key_name=key.name)
+    return packed_ciphertext
+
+
+def decrypt(ciphertext: str) -> str:
+    key_type, _, ciphertext = _unpack_ciphertext(ciphertext)
+    # Ignore key_name when decrypting
+    for i, key in enumerate(_encryption_keys):
+        if key.TYPE != key_type:
+            continue
+        try:
+            return key.decrypt(ciphertext)
+        except Exception:
+            logger.debug(f"Attempt to decrypt ciphertext with key #{i} failed")
+    raise EncryptionError("All keys failed to decrypt ciphertext")
+
+
+def _pack_ciphertext(ciphertext: str, key_type: str, key_name: str) -> str:
+    return f"enc:{key_type}:{key_name}:{ciphertext}"
+
+
+def _unpack_ciphertext(packed_ciphertext: str) -> Tuple[str, str, str]:
+    _, key_type, key_name, ciphertext = packed_ciphertext.split(":", maxsplit=3)
+    return key_type, key_name, ciphertext
+
+
+EncryptedString.set_encrypt_decrypt(
+    encrypt_func=encrypt,
+    decrypt_func=decrypt,
+)

dstack/_internal/server/services/encryption/keys/aes.py

@@ -0,0 +1,68 @@
+import os
+from base64 import b64decode, b64encode
+from typing import Literal
+
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from pydantic import Field, validator
+from typing_extensions import Annotated
+
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.server.services.encryption.keys.base import EncryptionKey
+
+
+class AESEncryptionKeyConfig(CoreModel):
+    type: Annotated[Literal["aes"], Field(description="The type of the key")] = "aes"
+    name: Annotated[str, Field(description="The key name for key identification")]
+    secret: Annotated[str, Field(description="Base64-encoded AES-256 key")]
+
+    @validator("name")
+    def validate_name(cls, v):
+        if not v.isalnum():
+            raise ValueError("Key name must be alphanumeric")
+        return v
+
+    @validator("secret")
+    def validate_secret(cls, v):
+        try:
+            key = b64decode(v, validate=True)
+        except Exception as e:
+            raise ValueError("Failed to decode secret from base64") from e
+        if len(key) != 32:
+            raise ValueError(f"AES key must be 32 bytes. Got {len(key)} bytes")
+        return v
+
+
+class AESEncryptionKey(EncryptionKey):
+    TYPE = "aes"
+
+    def __init__(self, config: AESEncryptionKeyConfig) -> None:
+        self.config = config
+        self.key = b64decode(config.secret)
+
+    @property
+    def name(self) -> str:
+        return self.config.name
+
+    def encrypt(self, plaintext: str) -> str:
+        # Generate a random 12-byte (96-bit) nonce (recommended size for GCM)
+        nonce = os.urandom(12)
+
+        # Create an AESGCM object and encrypt the plaintext
+        aesgcm = AESGCM(self.key)
+        ciphertext = aesgcm.encrypt(nonce, plaintext.encode("utf-8"), None)
+
+        # base64-encode the nonce and ciphertext for storage
+        return b64encode(nonce + ciphertext).decode("utf-8")
+
+    def decrypt(self, ciphertext: str) -> str:
+        data = b64decode(ciphertext)
+
+        # Extract the nonce and ciphertext
+        nonce = data[:12]
+        decoded_ciphertext = data[12:]
+
+        # Create an AESGCM object and decrypt the ciphertext
+        aesgcm = AESGCM(self.key)
+        plaintext = aesgcm.decrypt(nonce, decoded_ciphertext, None)
+
+        return plaintext.decode("utf-8")

dstack/_internal/server/services/encryption/keys/base.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+from typing import ClassVar
+
+
+class EncryptionKey(ABC):
+    TYPE: ClassVar[str]
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        pass
+
+    @abstractmethod
+    def encrypt(self, plaintext: str) -> str:
+        pass
+
+    @abstractmethod
+    def decrypt(self, ciphertext: str) -> str:
+        pass

dstack/_internal/server/services/encryption/keys/identity.py

@@ -0,0 +1,28 @@
+from typing import Literal
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.server.services.encryption.keys.base import EncryptionKey
+
+
+class IdentityEncryptionKeyConfig(CoreModel):
+    type: Annotated[Literal["identity"], Field(description="The type of the key")] = "identity"
+
+
+class IdentityEncryptionKey(EncryptionKey):
+    TYPE = "identity"
+
+    def __init__(self, config: IdentityEncryptionKeyConfig) -> None:
+        pass
+
+    @property
+    def name(self) -> str:
+        return "noname"
+
+    def encrypt(self, plaintext: str) -> str:
+        return plaintext
+
+    def decrypt(self, ciphertext: str) -> str:
+        return ciphertext