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

dstack/_internal/core/models/services.py

@@ -0,0 +1,76 @@
+"""
+Data structures related to `type: service` runs.
+"""
+
+from typing import Optional, Union
+
+from pydantic import Field
+from typing_extensions import Annotated, Literal
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class BaseChatModel(CoreModel):
+    type: Annotated[Literal["chat"], Field(description="The type of the model")] = "chat"
+    name: Annotated[str, Field(description="The name of the model")]
+    format: Annotated[
+        str, Field(description="The serving format. Supported values include `openai` and `tgi`")
+    ]
+
+
+class TGIChatModel(BaseChatModel):
+    """
+    Mapping of the model for the OpenAI-compatible endpoint.
+
+    Attributes:
+        type (str): The type of the model, e.g. "chat"
+        name (str): The name of the model. This name will be used both to load model configuration from the HuggingFace Hub and in the OpenAI-compatible endpoint.
+        format (str): The format of the model, e.g. "tgi" if the model is served with HuggingFace's Text Generation Inference.
+        chat_template (Optional[str]): The custom prompt template for the model. If not specified, the default prompt template from the HuggingFace Hub configuration will be used.
+        eos_token (Optional[str]): The custom end of sentence token. If not specified, the default end of sentence token from the HuggingFace Hub configuration will be used.
+    """
+
+    format: Annotated[
+        Literal["tgi"], Field(description="The serving format. Must be set to `tgi`")
+    ]
+    chat_template: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The custom prompt template for the model."
+                " If not specified, the default prompt template"
+                " from the HuggingFace Hub configuration will be used"
+            )
+        ),
+    ] = None  # will be set before registering the service
+    eos_token: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The custom end of sentence token."
+                " If not specified, the default end of sentence token"
+                " from the HuggingFace Hub configuration will be used"
+            )
+        ),
+    ] = None
+
+
+class OpenAIChatModel(BaseChatModel):
+    """
+    Mapping of the model for the OpenAI-compatible endpoint.
+
+    Attributes:
+        type (str): The type of the model, e.g. "chat"
+        name (str): The name of the model. This name will be used both to load model configuration from the HuggingFace Hub and in the OpenAI-compatible endpoint.
+        format (str): The format of the model, i.e. "openai".
+        prefix (str): The `base_url` prefix: `http://hostname/{prefix}/chat/completions`. Defaults to `/v1`.
+    """
+
+    format: Annotated[
+        Literal["openai"], Field(description="The serving format. Must be set to `openai`")
+    ]
+    prefix: Annotated[str, Field(description="The `base_url` prefix (after hostname)")] = "/v1"
+
+
+ChatModel = Annotated[Union[TGIChatModel, OpenAIChatModel], Field(discriminator="format")]
+AnyModel = Union[ChatModel]  # embeddings and etc.

dstack/_internal/core/models/unix.py

@@ -0,0 +1,53 @@
+from typing import Optional
+
+from pydantic import Field
+from typing_extensions import Annotated, Self
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class UnixUser(CoreModel):
+    uid: Annotated[Optional[int], Field(description="User ID", ge=0)] = None
+    gid: Annotated[Optional[int], Field(description="Group ID", ge=0)] = None
+    username: Annotated[Optional[str], Field(description="User name", min_length=1)] = None
+    groupname: Annotated[Optional[str], Field(description="Group name", min_length=1)] = None
+
+    @classmethod
+    def parse(cls, v: str) -> Self:
+        """
+        Parse `<user>[:<group>]` format used by Docker.
+        """
+        try:
+            return cls._parse(v)
+        except ValueError as e:
+            raise ValueError(f"invalid user format: {e}")
+
+    @classmethod
+    def _parse(cls, v: str) -> Self:
+        parts = v.split(":")
+        if len(parts) > 2:
+            raise ValueError("too many parts")
+        uid: Optional[int] = None
+        gid: Optional[int] = None
+        username: Optional[str] = None
+        groupname: Optional[str] = None
+        user_name_or_id = parts[0]
+        if not user_name_or_id:
+            raise ValueError("empty user name or id")
+        try:
+            uid = int(user_name_or_id)
+        except ValueError:
+            username = user_name_or_id
+        if uid is not None and uid < 0:
+            raise ValueError(f"negative uid {uid}")
+        if len(parts) == 2:
+            group_name_or_id = parts[1]
+            if not group_name_or_id:
+                raise ValueError("empty group name or id")
+            try:
+                gid = int(group_name_or_id)
+            except ValueError:
+                groupname = group_name_or_id
+            if gid is not None and gid < 0:
+                raise ValueError(f"negative gid {gid}")
+        return cls(uid=uid, gid=gid, username=username, groupname=groupname)

dstack/_internal/core/models/users.py

@@ -0,0 +1,60 @@
+import enum
+from datetime import datetime
+from typing import List, Optional, Union
+
+from pydantic import UUID4
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class ProjectRole(str, enum.Enum):
+    ADMIN = "admin"
+    MANAGER = "manager"
+    USER = "user"
+
+
+class GlobalRole(str, enum.Enum):
+    ADMIN = "admin"
+    USER = "user"
+
+
+class UserPermissions(CoreModel):
+    can_create_projects: bool
+
+
+class User(CoreModel):
+    id: UUID4
+    username: str
+    created_at: Optional[datetime] = None
+    global_role: GlobalRole
+    email: Optional[str]
+    active: bool
+    permissions: UserPermissions
+    ssh_public_key: Optional[str] = None
+
+
+class UserTokenCreds(CoreModel):
+    token: str
+
+
+class UserWithCreds(User):
+    creds: UserTokenCreds
+    ssh_private_key: Optional[str] = None
+
+
+class UsersInfoList(CoreModel):
+    total_count: Optional[int] = None
+    users: List[User]
+
+
+# For backward compatibility with 0.20 clients, endpoints return `List[User]` if `total_count` is None.
+# TODO: Replace with UsersInfoList in 0.21.
+UsersInfoListOrUsersList = Union[List[User], UsersInfoList]
+
+
+class UserHookConfig(CoreModel):
+    """
+    This class can be inherited to extend the user creation configuration passed to the hooks.
+    """
+
+    pass

dstack/_internal/core/models/volumes.py

@@ -0,0 +1,221 @@
+import uuid
+from datetime import datetime
+from enum import Enum
+from pathlib import PurePosixPath
+from typing import Dict, List, Literal, Optional, Tuple, Union
+
+from pydantic import Field, validator
+from typing_extensions import Annotated, Self
+
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.core.models.profiles import parse_idle_duration
+from dstack._internal.core.models.resources import Memory
+from dstack._internal.utils.common import get_or_error
+from dstack._internal.utils.tags import tags_validator
+
+
+class VolumeStatus(str, Enum):
+    SUBMITTED = "submitted"
+    PROVISIONING = "provisioning"
+    ACTIVE = "active"
+    FAILED = "failed"
+
+    def is_active(self) -> bool:
+        return self not in self.finished_statuses()
+
+    @classmethod
+    def finished_statuses(cls) -> List["VolumeStatus"]:
+        return [cls.FAILED]
+
+
+class VolumeConfiguration(CoreModel):
+    type: Literal["volume"] = "volume"
+    name: Annotated[Optional[str], Field(description="The volume name")] = None
+    backend: Annotated[BackendType, Field(description="The volume backend")]
+    region: Annotated[str, Field(description="The volume region")]
+    availability_zone: Annotated[
+        Optional[str], Field(description="The volume availability zone")
+    ] = None
+    size: Annotated[
+        Optional[Memory],
+        Field(description="The volume size. Must be specified when creating new volumes"),
+    ] = None
+    volume_id: Annotated[
+        Optional[str],
+        Field(description="The volume ID. Must be specified when registering external volumes"),
+    ] = None
+    auto_cleanup_duration: Annotated[
+        Optional[Union[str, int]],
+        Field(
+            description=(
+                "Time to wait after volume is no longer used by any job before deleting it. "
+                "Defaults to keep the volume indefinitely. "
+                "Use the value 'off' or -1 to disable auto-cleanup."
+            )
+        ),
+    ] = None
+    tags: Annotated[
+        Optional[Dict[str, str]],
+        Field(
+            description=(
+                "The custom tags to associate with the volume."
+                " The tags are also propagated to the underlying backend resources."
+                " If there is a conflict with backend-level tags, does not override them"
+            )
+        ),
+    ] = None
+
+    _validate_tags = validator("tags", pre=True, allow_reuse=True)(tags_validator)
+    _validate_auto_cleanup_duration = validator(
+        "auto_cleanup_duration", pre=True, allow_reuse=True
+    )(parse_idle_duration)
+
+    @property
+    def size_gb(self) -> int:
+        return int(get_or_error(self.size))
+
+
+class VolumeSpec(CoreModel):
+    configuration: VolumeConfiguration
+    configuration_path: Optional[str] = None
+
+
+class VolumeProvisioningData(CoreModel):
+    backend: Optional[BackendType] = None
+    volume_id: str
+    size_gb: int
+    availability_zone: Optional[str] = None
+    # price per month
+    price: Optional[float] = None
+    # should be manually attached/detached
+    attachable: bool = True
+    detachable: bool = True
+    backend_data: Optional[str] = None  # backend-specific data in json
+
+
+class VolumeAttachmentData(CoreModel):
+    device_name: Optional[str] = None
+
+
+class VolumeInstance(CoreModel):
+    name: str
+    fleet_name: Optional[str] = None
+    instance_num: int
+    instance_id: Optional[str] = None
+
+
+class VolumeAttachment(CoreModel):
+    instance: VolumeInstance
+    attachment_data: Optional[VolumeAttachmentData] = None
+
+
+class Volume(CoreModel):
+    id: uuid.UUID
+    name: str
+    user: str
+    project_name: str
+    configuration: VolumeConfiguration
+    external: bool
+    created_at: datetime
+    last_processed_at: datetime
+    status: VolumeStatus
+    status_message: Optional[str] = None
+    deleted: bool
+    deleted_at: Optional[datetime] = None
+    volume_id: Optional[str] = None  # id of the volume in the cloud
+    provisioning_data: Optional[VolumeProvisioningData] = None
+    cost: float = 0
+    attachments: Optional[List[VolumeAttachment]] = None
+    # attachment_data is deprecated in favor of attachments.
+    # It's only set for volumes that were attached before attachments.
+    attachment_data: Optional[VolumeAttachmentData] = None
+
+    def get_attachment_data_for_instance(self, instance_id: str) -> Optional[VolumeAttachmentData]:
+        if self.attachments is not None:
+            for attachment in self.attachments:
+                if attachment.instance.instance_id == instance_id:
+                    return attachment.attachment_data
+        # volume was attached before attachments were introduced
+        return self.attachment_data
+
+
+class VolumePlan(CoreModel):
+    project_name: str
+    user: str
+    spec: VolumeSpec
+    current_resource: Optional[Volume]
+
+
+def _split_mount_point(mount_point: str) -> Tuple[str, str]:
+    parts = mount_point.split(":")
+    if len(parts) != 2:
+        raise ValueError(f"invalid mount point format: {mount_point}")
+    src, dest = parts
+    return src, dest
+
+
+def _validate_mount_point_path(path: str) -> str:
+    if not path:
+        raise ValueError("empty path")
+    _path = PurePosixPath(path)
+    if not _path.is_absolute():
+        raise ValueError(f"path must be absolute: {path}")
+    if ".." in _path.parts:
+        raise ValueError(f".. are not allowed: {path}")
+    return str(_path)
+
+
+class VolumeMountPoint(CoreModel):
+    name: Annotated[
+        Union[str, List[str]],
+        Field(
+            description=(
+                "The network volume name or the list of network volume names to mount."
+                " If a list is specified, one of the volumes in the list will be mounted."
+                " Specify volumes from different backends/regions to increase availability"
+            )
+        ),
+    ]
+    path: Annotated[str, Field(description="The absolute container path to mount the volume at")]
+
+    _validate_path = validator("path", allow_reuse=True)(_validate_mount_point_path)
+
+    @classmethod
+    def parse(cls, v: str) -> Self:
+        name, path = _split_mount_point(v)
+        return cls(name=name, path=path)
+
+
+class InstanceMountPoint(CoreModel):
+    instance_path: Annotated[str, Field(description="The absolute path on the instance (host)")]
+    path: Annotated[str, Field(description="The absolute path in the container")]
+    optional: Annotated[
+        bool,
+        Field(
+            description=(
+                "Allow running without this volume"
+                " in backends that do not support instance volumes"
+            ),
+        ),
+    ] = False
+
+    _validate_instance_path = validator("instance_path", allow_reuse=True)(
+        _validate_mount_point_path
+    )
+    _validate_path = validator("path", allow_reuse=True)(_validate_mount_point_path)
+
+    @classmethod
+    def parse(cls, v: str) -> Self:
+        instance_path, path = _split_mount_point(v)
+        return cls(instance_path=instance_path, path=path)
+
+
+MountPoint = Union[VolumeMountPoint, InstanceMountPoint]
+
+
+def parse_mount_point(v: str) -> MountPoint:
+    src, dest = _split_mount_point(v)
+    if "/" in src:
+        return InstanceMountPoint(instance_path=src, path=dest)
+    return VolumeMountPoint(name=src, path=dest)

dstack/_internal/core/services/__init__.py

@@ -0,0 +1,16 @@
+import re
+
+from dstack._internal.core.errors import ServerClientError
+
+
+def validate_dstack_resource_name(resource_name: str):
+    if not is_valid_dstack_resource_name(resource_name):
+        raise ServerClientError("Resource name should match regex '^[a-z][a-z0-9-]{1,40}$'")
+
+
+def is_valid_dstack_resource_name(resource_name: str) -> bool:
+    return re.match("^[a-z][a-z0-9-]{1,40}$", resource_name) is not None
+
+
+def is_valid_replica_group_name(name: str) -> bool:
+    return re.match("^[a-z0-9][a-z0-9-]{0,39}$", name) is not None

dstack/_internal/core/services/api_client.py

@@ -0,0 +1,15 @@
+from typing import Optional, Tuple
+
+import dstack._internal.core.services.configs as configs
+from dstack._internal.core.errors import ConfigurationError
+from dstack.api.server import APIClient
+
+
+def get_api_client(project_name: Optional[str] = None) -> Tuple[APIClient, str]:
+    config = configs.ConfigManager()
+    project = config.get_project_config(project_name)
+    if project is None:
+        if project_name is not None:
+            raise ConfigurationError(f"Project {project_name} is not configured")
+        raise ConfigurationError("No default project, specify project name")
+    return APIClient(project.url, project.token), project.name

dstack/_internal/core/services/configs/__init__.py

@@ -0,0 +1,116 @@
+import json
+import os.path
+from pathlib import Path
+from typing import Optional
+
+import yaml
+from pydantic import ValidationError
+
+from dstack._internal.cli.utils.common import confirm_ask
+from dstack._internal.core.models.config import GlobalConfig, ProjectConfig
+from dstack._internal.utils.common import get_dstack_dir
+from dstack._internal.utils.logging import get_logger
+from dstack._internal.utils.path import PathLike
+
+logger = get_logger(__name__)
+
+
+class ConfigManager:
+    config: GlobalConfig
+
+    def __init__(self, dstack_dir: Optional[PathLike] = None):
+        self.dstack_dir = Path(dstack_dir) if dstack_dir else get_dstack_dir()
+        self.config_filepath = self.dstack_dir / "config.yml"
+        self.dstack_ssh_dir.mkdir(parents=True, exist_ok=True)
+        self.load()
+
+    def save(self):
+        self.config_filepath.parent.mkdir(parents=True, exist_ok=True)
+        with self.config_filepath.open("w") as f:
+            # hack to convert enums to strings, etc.
+            yaml.dump(json.loads(self.config.json()), f)
+
+    def load(self):
+        try:
+            with open(self.config_filepath, "r") as f:
+                config = yaml.safe_load(f)
+            self.config = GlobalConfig.parse_obj(config)
+        except FileNotFoundError:
+            self.config = GlobalConfig()
+        except ValidationError:
+            logger.error(f"Error in `{self.config_filepath}`", exc_info=True)
+            self.config = GlobalConfig()
+
+    def get_project_config(self, name: Optional[str] = None) -> Optional[ProjectConfig]:
+        for project in self.config.projects:
+            if name is None and project.default:
+                return project
+            if project.name == name:
+                return project
+        return None
+
+    def configure_project(self, name: str, url: str, token: str, default: bool):
+        if default:
+            for project in self.config.projects:
+                project.default = False
+        for project in self.config.projects:
+            if project.name == name:
+                project.url = url
+                project.token = token
+                project.default = default or project.default
+                return
+        self.config.projects.append(
+            ProjectConfig(name=name, url=url, token=token, default=default)
+        )
+        if len(self.config.projects) == 1:
+            self.config.projects[0].default = True
+
+    def list_project_configs(self) -> list[ProjectConfig]:
+        return self.config.projects
+
+    def delete_project(self, name: str):
+        self.config.projects = [p for p in self.config.projects if p.name != name]
+
+    @property
+    def dstack_ssh_dir(self) -> Path:
+        return self.dstack_dir / "ssh"
+
+    @property
+    def dstack_key_path(self) -> Path:
+        # TODO: Remove since 0.19.40
+        return self.dstack_ssh_dir / "id_rsa"
+
+    @property
+    def dstack_ssh_config_path(self) -> Path:
+        return self.dstack_ssh_dir / "config"
+
+
+def update_default_project(project_name: str, url: str, token: str, yes: bool, no: bool):
+    config_manager = ConfigManager()
+    default_project = config_manager.get_project_config()
+    config_dir = str(config_manager.config_filepath).replace(os.path.expanduser("~"), "~", 1)
+    if (
+        default_project is None
+        or default_project.name != project_name
+        or default_project.url != url
+        or default_project.token != token
+    ):
+        set_it_as_default = (
+            (
+                default_project is None
+                or yes
+                or confirm_ask(
+                    f"Update the [code]{project_name}[/] project in [code]{config_dir}[/]?"
+                )
+            )
+            if not no
+            else False
+        )
+        if set_it_as_default:
+            config_manager.configure_project(
+                name=project_name, url=url, token=token, default=set_it_as_default
+            )
+            config_manager.save()
+            logger.info(
+                f"Configured the {project_name} project in {config_dir}", {"show_path": False}
+            )

dstack/_internal/core/services/diff.py

@@ -0,0 +1,71 @@
+from typing import Any, Optional, TypedDict, TypeVar
+
+from pydantic import BaseModel
+
+from dstack._internal.core.models.common import IncludeExcludeType
+
+
+class ModelFieldDiff(TypedDict):
+    old: Any
+    new: Any
+
+
+ModelDiff = dict[str, ModelFieldDiff]
+
+
+# TODO: calculate nested diffs
+def diff_models(
+    old: BaseModel, new: BaseModel, reset: Optional[IncludeExcludeType] = None
+) -> ModelDiff:
+    """
+    Returns a diff of model instances fields.
+
+    The fields specified in the `reset` option are reset to their default values, effectively
+    excluding them from comparison (assuming that the default value is equal to itself, e.g,
+    `None == None`, `"task" == "task"`, but `math.nan != math.nan`).
+
+    Args:
+        old: The "old" model instance.
+        new: The "new" model instance.
+        reset: Fields to reset to their default values before comparison.
+
+    Returns:
+        A dict of changed fields in the form of
+        `{<field_name>: {"old": old_value, "new": new_value}}`
+    """
+    if type(old) is not type(new):
+        raise TypeError("Both instances must be of the same Pydantic model class.")
+
+    if reset is not None:
+        old = copy_model(old, reset=reset)
+        new = copy_model(new, reset=reset)
+
+    changes: ModelDiff = {}
+    for field in old.__fields__:
+        old_value = getattr(old, field)
+        new_value = getattr(new, field)
+        if old_value != new_value:
+            changes[field] = {"old": old_value, "new": new_value}
+
+    return changes
+
+
+M = TypeVar("M", bound=BaseModel)
+
+
+def copy_model(model: M, reset: Optional[IncludeExcludeType] = None) -> M:
+    """
+    Returns a deep copy of the model instance.
+
+    Implemented as `BaseModel.parse_obj(BaseModel.dict())`, thus,
+    unlike `BaseModel.copy(deep=True)`, runs all validations.
+
+    The fields specified in the `reset` option are reset to their default values.
+
+    Args:
+        reset: Fields to reset to their default values.
+
+    Returns:
+        A deep copy of the model instance.
+    """
+    return type(model).parse_obj(model.dict(exclude=reset))

dstack/_internal/core/services/logs.py

@@ -0,0 +1,58 @@
+import re
+import urllib.parse
+from typing import Dict, List, Optional
+
+from dstack._internal.core.models.runs import AppSpec
+from dstack._internal.utils.common import concat_url_path
+
+
+class URLReplacer:
+    def __init__(
+        self,
+        app_specs: List[AppSpec],
+        ports: Dict[int, int],
+        hostname: str,
+        secure: bool,
+        path_prefix: str = "",
+        ip_address: Optional[str] = None,
+    ):
+        self.app_specs = {app_spec.port: app_spec for app_spec in app_specs}
+        self.ports = ports
+        self.hostname = hostname
+        self.secure = secure
+        self.path_prefix = path_prefix.encode()
+
+        hosts = ["localhost", "0.0.0.0", "127.0.0.1"]
+        if ip_address and ip_address not in hosts:
+            hosts.append(ip_address)
+        hosts_re = "|".join(hosts)
+        self._url_re = re.compile(
+            rf"http://(?:{hosts_re}):(\d+)\b([-a-zA-Z0-9()@:%_\+.~#?&//=]*)".encode()
+        )
+
+    def _replace_url(self, match: re.Match) -> bytes:
+        remote_port = int(match.group(1))
+        if remote_port not in self.ports:
+            return match.group(0)
+        local_port = self.ports[remote_port]
+        omit_port = (not self.secure and local_port == 80) or (self.secure and local_port == 443)
+
+        app_spec = self.app_specs.get(remote_port)
+        url = urllib.parse.urlparse(match.group(0))
+        qs = {k: v[0] for k, v in urllib.parse.parse_qs(url.query).items()}
+        if app_spec and app_spec.url_query_params is not None:
+            qs.update({k.encode(): v.encode() for k, v in app_spec.url_query_params.items()})
+        path = url.path
+        if not path.startswith(self.path_prefix.removesuffix(b"/")):
+            path = concat_url_path(self.path_prefix, path)
+
+        url = url._replace(
+            scheme=("https" if self.secure else "http").encode(),
+            netloc=(self.hostname if omit_port else f"{self.hostname}:{local_port}").encode(),
+            path=path,
+            query=urllib.parse.urlencode(qs).encode(),
+        )
+        return url.geturl()
+
+    def __call__(self, entry: bytes) -> bytes:
+        return self._url_re.sub(self._replace_url, entry)