dstack 0.19.16__py3-none-any.whl → 0.19.18__py3-none-any.whl

dstack/_internal/server/migrations/versions/644b8a114187_add_secretmodel.py

@@ -0,0 +1,49 @@
+"""Add SecretModel
+
+Revision ID: 644b8a114187
+Revises: 5f1707c525d2
+Create Date: 2025-06-30 11:00:04.326290
+
+"""
+
+import sqlalchemy as sa
+import sqlalchemy_utils
+from alembic import op
+
+import dstack._internal.server.models
+
+# revision identifiers, used by Alembic.
+revision = "644b8a114187"
+down_revision = "5f1707c525d2"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "secrets",
+        sa.Column("id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False),
+        sa.Column(
+            "project_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False
+        ),
+        sa.Column("created_at", dstack._internal.server.models.NaiveDateTime(), nullable=False),
+        sa.Column("updated_at", dstack._internal.server.models.NaiveDateTime(), nullable=False),
+        sa.Column("name", sa.String(length=200), nullable=False),
+        sa.Column("value", dstack._internal.server.models.EncryptedString(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["project_id"],
+            ["projects.id"],
+            name=op.f("fk_secrets_project_id_projects"),
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_secrets")),
+        sa.UniqueConstraint("project_id", "name", name="uq_secrets_project_id_name"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("secrets")
+    # ### end Alembic commands ###

dstack/_internal/server/models.py

@@ -315,6 +315,21 @@ class CodeModel(BaseModel):
     blob: Mapped[Optional[bytes]] = mapped_column(LargeBinary)  # None means blob is stored on s3
 
 
+class FileArchiveModel(BaseModel):
+    __tablename__ = "file_archives"
+    __table_args__ = (
+        UniqueConstraint("user_id", "blob_hash", name="uq_file_archives_user_id_blob_hash"),
+    )
+
+    id: Mapped[uuid.UUID] = mapped_column(
+        UUIDType(binary=False), primary_key=True, default=uuid.uuid4
+    )
+    user_id: Mapped["UserModel"] = mapped_column(ForeignKey("users.id", ondelete="CASCADE"))
+    user: Mapped["UserModel"] = relationship()
+    blob_hash: Mapped[str] = mapped_column(Text)
+    blob: Mapped[Optional[bytes]] = mapped_column(LargeBinary)  # None means blob is stored on s3
+
+
 class RunModel(BaseModel):
     __tablename__ = "runs"
 
@@ -711,3 +726,21 @@ class JobPrometheusMetrics(BaseModel):
     collected_at: Mapped[datetime] = mapped_column(NaiveDateTime)
     # Raw Prometheus text response
     text: Mapped[str] = mapped_column(Text)
+
+
+class SecretModel(BaseModel):
+    __tablename__ = "secrets"
+    __table_args__ = (UniqueConstraint("project_id", "name", name="uq_secrets_project_id_name"),)
+
+    id: Mapped[uuid.UUID] = mapped_column(
+        UUIDType(binary=False), primary_key=True, default=uuid.uuid4
+    )
+
+    project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE"))
+    project: Mapped["ProjectModel"] = relationship()
+
+    created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime)
+    updated_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime)
+
+    name: Mapped[str] = mapped_column(String(200))
+    value: Mapped[DecryptedString] = mapped_column(EncryptedString())

dstack/_internal/server/routers/files.py

@@ -0,0 +1,67 @@
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, Request, UploadFile
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.errors import ResourceNotExistsError, ServerClientError
+from dstack._internal.core.models.files import FileArchive
+from dstack._internal.server.db import get_session
+from dstack._internal.server.models import UserModel
+from dstack._internal.server.schemas.files import GetFileArchiveByHashRequest
+from dstack._internal.server.security.permissions import Authenticated
+from dstack._internal.server.services import files
+from dstack._internal.server.settings import SERVER_CODE_UPLOAD_LIMIT
+from dstack._internal.server.utils.routers import (
+    get_base_api_additional_responses,
+    get_request_size,
+)
+from dstack._internal.utils.common import sizeof_fmt
+
+router = APIRouter(
+    prefix="/api/files",
+    tags=["files"],
+    responses=get_base_api_additional_responses(),
+)
+
+
+@router.post("/get_archive_by_hash")
+async def get_archive_by_hash(
+    body: GetFileArchiveByHashRequest,
+    session: Annotated[AsyncSession, Depends(get_session)],
+    user: Annotated[UserModel, Depends(Authenticated())],
+) -> FileArchive:
+    archive = await files.get_archive_by_hash(
+        session=session,
+        user=user,
+        hash=body.hash,
+    )
+    if archive is None:
+        raise ResourceNotExistsError()
+    return archive
+
+
+@router.post("/upload_archive")
+async def upload_archive(
+    request: Request,
+    file: UploadFile,
+    session: Annotated[AsyncSession, Depends(get_session)],
+    user: Annotated[UserModel, Depends(Authenticated())],
+) -> FileArchive:
+    request_size = get_request_size(request)
+    if SERVER_CODE_UPLOAD_LIMIT > 0 and request_size > SERVER_CODE_UPLOAD_LIMIT:
+        diff_size_fmt = sizeof_fmt(request_size)
+        limit_fmt = sizeof_fmt(SERVER_CODE_UPLOAD_LIMIT)
+        if diff_size_fmt == limit_fmt:
+            diff_size_fmt = f"{request_size}B"
+            limit_fmt = f"{SERVER_CODE_UPLOAD_LIMIT}B"
+        raise ServerClientError(
+            f"Archive size is {diff_size_fmt}, which exceeds the limit of {limit_fmt}."
+            " Use .gitignore/.dstackignore to exclude large files."
+            " This limit can be modified by setting the DSTACK_SERVER_CODE_UPLOAD_LIMIT environment variable."
+        )
+    archive = await files.upload_archive(
+        session=session,
+        user=user,
+        file=file,
+    )
+    return archive

dstack/_internal/server/routers/secrets.py

@@ -1,15 +1,19 @@
-from typing import List
+from typing import List, Tuple
 
-from fastapi import APIRouter
+from fastapi import APIRouter, Depends
+from sqlalchemy.ext.asyncio import AsyncSession
 
-from dstack._internal.core.models.runs import Run
+from dstack._internal.core.errors import ResourceNotExistsError
 from dstack._internal.core.models.secrets import Secret
+from dstack._internal.server.db import get_session
+from dstack._internal.server.models import ProjectModel, UserModel
 from dstack._internal.server.schemas.secrets import (
-    AddSecretRequest,
+    CreateOrUpdateSecretRequest,
     DeleteSecretsRequest,
-    GetSecretsRequest,
-    ListSecretsRequest,
+    GetSecretRequest,
 )
+from dstack._internal.server.security.permissions import ProjectAdmin
+from dstack._internal.server.services import secrets as secrets_services
 
 router = APIRouter(
     prefix="/api/project/{project_name}/secrets",
@@ -18,20 +22,58 @@ router = APIRouter(
 
 
 @router.post("/list")
-async def list_secrets(project_name: str, body: ListSecretsRequest) -> List[Run]:
-    pass
+async def list_secrets(
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
+) -> List[Secret]:
+    _, project = user_project
+    return await secrets_services.list_secrets(
+        session=session,
+        project=project,
+    )
 
 
 @router.post("/get")
-async def get_secret(project_name: str, body: GetSecretsRequest) -> Secret:
-    pass
+async def get_secret(
+    body: GetSecretRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
+) -> Secret:
+    _, project = user_project
+    secret = await secrets_services.get_secret(
+        session=session,
+        project=project,
+        name=body.name,
+    )
+    if secret is None:
+        raise ResourceNotExistsError()
+    return secret
 
 
-@router.post("/add")
-async def add_or_update_secret(project_name: str, body: AddSecretRequest) -> Secret:
-    pass
+@router.post("/create_or_update")
+async def create_or_update_secret(
+    body: CreateOrUpdateSecretRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
+) -> Secret:
+    _, project = user_project
+    return await secrets_services.create_or_update_secret(
+        session=session,
+        project=project,
+        name=body.name,
+        value=body.value,
+    )
 
 
 @router.post("/delete")
-async def delete_secrets(project_name: str, body: DeleteSecretsRequest):
-    pass
+async def delete_secrets(
+    body: DeleteSecretsRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
+):
+    _, project = user_project
+    await secrets_services.delete_secrets(
+        session=session,
+        project=project,
+        names=body.secrets_names,
+    )

dstack/_internal/server/schemas/files.py

@@ -0,0 +1,5 @@
+from dstack._internal.core.models.common import CoreModel
+
+
+class GetFileArchiveByHashRequest(CoreModel):
+    hash: str

dstack/_internal/server/schemas/runner.py

@@ -77,6 +77,8 @@ class SubmitBody(CoreModel):
                 "max_duration",
                 "ssh_key",
                 "working_dir",
+                "repo_data",
+                "file_archives",
             }
         ),
     ]

dstack/_internal/server/schemas/secrets.py

@@ -1,20 +1,16 @@
 from typing import List
 
-from dstack._internal.core.models.secrets import Secret
-from dstack._internal.server.schemas.common import RepoRequest
+from dstack._internal.core.models.common import CoreModel
 
 
-class ListSecretsRequest(RepoRequest):
-    pass
+class GetSecretRequest(CoreModel):
+    name: str
 
 
-class GetSecretsRequest(RepoRequest):
-    pass
+class CreateOrUpdateSecretRequest(CoreModel):
+    name: str
+    value: str
 
 
-class AddSecretRequest(RepoRequest):
-    secret: Secret
-
-
-class DeleteSecretsRequest(RepoRequest):
+class DeleteSecretsRequest(CoreModel):
     secrets_names: List[str]

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

@@ -35,7 +35,7 @@ from dstack._internal.core.models.instances import (
 from dstack._internal.core.models.runs import Requirements
 from dstack._internal.server import settings
 from dstack._internal.server.models import BackendModel, DecryptedString, ProjectModel
-from dstack._internal.server.settings import LOCAL_BACKEND_ENABLED
+from dstack._internal.settings import LOCAL_BACKEND_ENABLED
 from dstack._internal.utils.common import run_async
 from dstack._internal.utils.logging import get_logger
 

dstack/_internal/server/services/files.py

@@ -0,0 +1,91 @@
+import uuid
+from typing import Optional
+
+from fastapi import UploadFile
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.errors import ServerClientError
+from dstack._internal.core.models.files import FileArchive
+from dstack._internal.server.models import FileArchiveModel, UserModel
+from dstack._internal.server.services.storage import get_default_storage
+from dstack._internal.utils.common import run_async
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+async def get_archive_model(
+    session: AsyncSession,
+    id: uuid.UUID,
+    user: Optional[UserModel] = None,
+) -> Optional[FileArchiveModel]:
+    stmt = select(FileArchiveModel).where(FileArchiveModel.id == id)
+    if user is not None:
+        stmt = stmt.where(FileArchiveModel.user_id == user.id)
+    res = await session.execute(stmt)
+    return res.scalar()
+
+
+async def get_archive_model_by_hash(
+    session: AsyncSession,
+    user: UserModel,
+    hash: str,
+) -> Optional[FileArchiveModel]:
+    res = await session.execute(
+        select(FileArchiveModel).where(
+            FileArchiveModel.user_id == user.id,
+            FileArchiveModel.blob_hash == hash,
+        )
+    )
+    return res.scalar()
+
+
+async def get_archive_by_hash(
+    session: AsyncSession,
+    user: UserModel,
+    hash: str,
+) -> Optional[FileArchive]:
+    archive_model = await get_archive_model_by_hash(
+        session=session,
+        user=user,
+        hash=hash,
+    )
+    if archive_model is None:
+        return None
+    return archive_model_to_archive(archive_model)
+
+
+async def upload_archive(
+    session: AsyncSession,
+    user: UserModel,
+    file: UploadFile,
+) -> FileArchive:
+    if file.filename is None:
+        raise ServerClientError("filename not specified")
+    archive_hash = file.filename
+    archive_model = await get_archive_model_by_hash(
+        session=session,
+        user=user,
+        hash=archive_hash,
+    )
+    if archive_model is not None:
+        logger.debug("File archive (user_id=%s, hash=%s) already uploaded", user.id, archive_hash)
+        return archive_model_to_archive(archive_model)
+    blob = await file.read()
+    storage = get_default_storage()
+    if storage is not None:
+        await run_async(storage.upload_archive, str(user.id), archive_hash, blob)
+    archive_model = FileArchiveModel(
+        user_id=user.id,
+        blob_hash=archive_hash,
+        blob=blob if storage is None else None,
+    )
+    session.add(archive_model)
+    await session.commit()
+    logger.debug("File archive (user_id=%s, hash=%s) has been uploaded", user.id, archive_hash)
+    return archive_model_to_archive(archive_model)
+
+
+def archive_model_to_archive(archive_model: FileArchiveModel) -> FileArchive:
+    return FileArchive(id=archive_model.id, hash=archive_model.blob_hash)

dstack/_internal/server/services/fleets.py

@@ -362,7 +362,7 @@ async def create_fleet(
             select(func.pg_advisory_xact_lock(string_to_lock_id(lock_namespace)))
         )
 
-    lock, _ = get_locker().get_lockset(lock_namespace)
+    lock, _ = get_locker(get_db().dialect_name).get_lockset(lock_namespace)
     async with lock:
         if spec.configuration.name is not None:
             fleet_model = await get_project_fleet_model_by_name(
@@ -516,11 +516,12 @@ async def delete_fleets(
     await session.commit()
     logger.info("Deleting fleets: %s", [v.name for v in fleet_models])
     async with (
-        get_locker().lock_ctx(FleetModel.__tablename__, fleets_ids),
-        get_locker().lock_ctx(InstanceModel.__tablename__, instances_ids),
+        get_locker(get_db().dialect_name).lock_ctx(FleetModel.__tablename__, fleets_ids),
+        get_locker(get_db().dialect_name).lock_ctx(InstanceModel.__tablename__, instances_ids),
     ):
         # Refetch after lock
-        # TODO lock instances with FOR UPDATE?
+        # TODO: Lock instances with FOR UPDATE?
+        # TODO: Do not lock fleet when deleting only instances
         res = await session.execute(
             select(FleetModel)
             .where(

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

@@ -162,7 +162,7 @@ async def create_gateway(
             select(func.pg_advisory_xact_lock(string_to_lock_id(lock_namespace)))
         )
 
-    lock, _ = get_locker().get_lockset(lock_namespace)
+    lock, _ = get_locker(get_db().dialect_name).get_lockset(lock_namespace)
     async with lock:
         if configuration.name is None:
             configuration.name = await generate_gateway_name(session=session, project=project)
@@ -229,7 +229,9 @@ async def delete_gateways(
     gateways_ids = sorted([g.id for g in gateway_models])
     await session.commit()
     logger.info("Deleting gateways: %s", [g.name for g in gateway_models])
-    async with get_locker().lock_ctx(GatewayModel.__tablename__, gateways_ids):
+    async with get_locker(get_db().dialect_name).lock_ctx(
+        GatewayModel.__tablename__, gateways_ids
+    ):
         # Refetch after lock
         res = await session.execute(
             select(GatewayModel)

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

@@ -33,6 +33,7 @@ from dstack._internal.core.models.runs import (
     RunSpec,
 )
 from dstack._internal.core.models.volumes import Volume, VolumeMountPoint, VolumeStatus
+from dstack._internal.server import settings
 from dstack._internal.server.models import (
     InstanceModel,
     JobModel,
@@ -64,15 +65,23 @@ from dstack._internal.utils.logging import get_logger
 logger = get_logger(__name__)
 
 
-async def get_jobs_from_run_spec(run_spec: RunSpec, replica_num: int) -> List[Job]:
+async def get_jobs_from_run_spec(
+    run_spec: RunSpec, secrets: Dict[str, str], replica_num: int
+) -> List[Job]:
     return [
         Job(job_spec=s, job_submissions=[])
-        for s in await get_job_specs_from_run_spec(run_spec, replica_num)
+        for s in await get_job_specs_from_run_spec(
+            run_spec=run_spec,
+            secrets=secrets,
+            replica_num=replica_num,
+        )
     ]
 
 
-async def get_job_specs_from_run_spec(run_spec: RunSpec, replica_num: int) -> List[JobSpec]:
-    job_configurator = _get_job_configurator(run_spec)
+async def get_job_specs_from_run_spec(
+    run_spec: RunSpec, secrets: Dict[str, str], replica_num: int
+) -> List[JobSpec]:
+    job_configurator = _get_job_configurator(run_spec=run_spec, secrets=secrets)
     job_specs = await job_configurator.get_job_specs(replica_num=replica_num)
     return job_specs
 
@@ -158,10 +167,10 @@ def delay_job_instance_termination(job_model: JobModel):
     job_model.remove_at = common.get_current_datetime() + timedelta(seconds=15)
 
 
-def _get_job_configurator(run_spec: RunSpec) -> JobConfigurator:
+def _get_job_configurator(run_spec: RunSpec, secrets: Dict[str, str]) -> JobConfigurator:
     configuration_type = RunConfigurationType(run_spec.configuration.type)
     configurator_class = _configuration_type_to_configurator_class_map[configuration_type]
-    return configurator_class(run_spec)
+    return configurator_class(run_spec=run_spec, secrets=secrets)
 
 
 _job_configurator_classes = [
@@ -380,8 +389,10 @@ def _shim_submit_stop(ports: Dict[int, int], job_model: JobModel):
             message=job_model.termination_reason_message,
             timeout=0,
         )
-        # maybe somehow postpone removing old tasks to allow inspecting failed jobs?
-        shim_client.remove_task(task_id=job_model.id)
+        # maybe somehow postpone removing old tasks to allow inspecting failed jobs without
+        # the following setting?
+        if not settings.SERVER_KEEP_SHIM_TASKS:
+            shim_client.remove_task(task_id=job_model.id)
     else:
         shim_client.stop(force=True)
 

dstack/_internal/server/services/jobs/configurators/base.py

@@ -1,5 +1,6 @@
 import shlex
 import sys
+import threading
 from abc import ABC, abstractmethod
 from pathlib import PurePosixPath
 from typing import Dict, List, Optional, Union
@@ -68,8 +69,13 @@ class JobConfigurator(ABC):
     # JobSSHKey should be shared for all jobs in a replica for inter-node communication.
     _job_ssh_key: Optional[JobSSHKey] = None
 
-    def __init__(self, run_spec: RunSpec):
+    def __init__(
+        self,
+        run_spec: RunSpec,
+        secrets: Optional[Dict[str, str]] = None,
+    ):
         self.run_spec = run_spec
+        self.secrets = secrets or {}
 
     async def get_job_specs(self, replica_num: int) -> List[JobSpec]:
         job_spec = await self._get_job_spec(replica_num=replica_num, job_num=0, jobs_per_replica=1)
@@ -98,10 +104,20 @@ class JobConfigurator(ABC):
     async def _get_image_config(self) -> ImageConfig:
         if self._image_config is not None:
             return self._image_config
+        interpolate = VariablesInterpolator({"secrets": self.secrets}).interpolate_or_error
+        registry_auth = self.run_spec.configuration.registry_auth
+        if registry_auth is not None:
+            try:
+                registry_auth = RegistryAuth(
+                    username=interpolate(registry_auth.username),
+                    password=interpolate(registry_auth.password),
+                )
+            except InterpolatorError as e:
+                raise ServerClientError(e.args[0])
         image_config = await run_async(
             _get_image_config,
             self._image_name(),
-            self.run_spec.configuration.registry_auth,
+            registry_auth,
         )
         self._image_config = image_config
         return image_config
@@ -134,6 +150,9 @@ class JobConfigurator(ABC):
             working_dir=self._working_dir(),
             volumes=self._volumes(job_num),
             ssh_key=self._ssh_key(jobs_per_replica),
+            repo_data=self.run_spec.repo_data,
+            repo_code_hash=self.run_spec.repo_code_hash,
+            file_archives=self.run_spec.file_archives,
         )
         return job_spec
 
@@ -336,7 +355,10 @@ def _join_shell_commands(commands: List[str]) -> str:
     return " && ".join(commands)
 
 
-@cached(TTLCache(maxsize=2048, ttl=80))
+@cached(
+    cache=TTLCache(maxsize=2048, ttl=80),
+    lock=threading.Lock(),
+)
 def _get_image_config(image: str, registry_auth: Optional[RegistryAuth]) -> ImageConfig:
     try:
         return get_image_config(image, registry_auth).config

dstack/_internal/server/services/jobs/configurators/dev.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from dstack._internal.core.errors import ServerClientError
 from dstack._internal.core.models.configurations import PortMapping, RunConfigurationType
@@ -17,7 +17,7 @@ INSTALL_IPYKERNEL = (
 class DevEnvironmentJobConfigurator(JobConfigurator):
     TYPE: RunConfigurationType = RunConfigurationType.DEV_ENVIRONMENT
 
-    def __init__(self, run_spec: RunSpec):
+    def __init__(self, run_spec: RunSpec, secrets: Dict[str, str]):
         if run_spec.configuration.ide == "vscode":
             __class = VSCodeDesktop
         elif run_spec.configuration.ide == "cursor":
@@ -29,7 +29,7 @@ class DevEnvironmentJobConfigurator(JobConfigurator):
             version=run_spec.configuration.version,
             extensions=["ms-python.python", "ms-toolsai.jupyter"],
         )
-        super().__init__(run_spec)
+        super().__init__(run_spec=run_spec, secrets=secrets)
 
     def _shell_commands(self) -> List[str]:
         commands = self.ide.get_install_commands()