dstack 0.19.15rc1__py3-none-any.whl → 0.19.17__py3-none-any.whl

dstack/_internal/server/routers/projects.py

@@ -7,13 +7,19 @@ from dstack._internal.core.models.projects import Project
 from dstack._internal.server.db import get_session
 from dstack._internal.server.models import ProjectModel, UserModel
 from dstack._internal.server.schemas.projects import (
+    AddProjectMemberRequest,
     CreateProjectRequest,
     DeleteProjectsRequest,
+    RemoveProjectMemberRequest,
     SetProjectMembersRequest,
+    UpdateProjectRequest,
 )
 from dstack._internal.server.security.permissions import (
     Authenticated,
+    ProjectAdmin,
     ProjectManager,
+    ProjectManagerOrPublicProject,
+    ProjectManagerOrSelfLeave,
     ProjectMemberOrPublicAccess,
 )
 from dstack._internal.server.services import projects
@@ -92,3 +98,60 @@ async def set_project_members(
     )
     await session.refresh(project)
     return projects.project_model_to_project(project)
+
+
+@router.post(
+    "/{project_name}/add_members",
+)
+async def add_project_members(
+    body: AddProjectMemberRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManagerOrPublicProject()),
+) -> Project:
+    user, project = user_project
+    await projects.add_project_members(
+        session=session,
+        user=user,
+        project=project,
+        members=body.members,
+    )
+    await session.refresh(project)
+    return projects.project_model_to_project(project)
+
+
+@router.post(
+    "/{project_name}/remove_members",
+)
+async def remove_project_members(
+    body: RemoveProjectMemberRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManagerOrSelfLeave()),
+) -> Project:
+    user, project = user_project
+    await projects.remove_project_members(
+        session=session,
+        user=user,
+        project=project,
+        usernames=body.usernames,
+    )
+    await session.refresh(project)
+    return projects.project_model_to_project(project)
+
+
+@router.post(
+    "/{project_name}/update",
+)
+async def update_project(
+    body: UpdateProjectRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
+) -> Project:
+    user, project = user_project
+    await projects.update_project(
+        session=session,
+        user=user,
+        project=project,
+        is_public=body.is_public,
+    )
+    await session.refresh(project)
+    return projects.project_model_to_project(project)

dstack/_internal/server/routers/prometheus.py

@@ -1,15 +1,15 @@
 import os
 from typing import Annotated
 
+import prometheus_client
 from fastapi import APIRouter, Depends
 from fastapi.responses import PlainTextResponse
-from prometheus_client import generate_latest
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.server import settings
 from dstack._internal.server.db import get_session
 from dstack._internal.server.security.permissions import OptionalServiceAccount
-from dstack._internal.server.services import prometheus
+from dstack._internal.server.services.prometheus import custom_metrics
 from dstack._internal.server.utils.routers import error_not_found
 
 _auth = OptionalServiceAccount(os.getenv("DSTACK_PROMETHEUS_AUTH_TOKEN"))
@@ -27,6 +27,6 @@ async def get_prometheus_metrics(
 ) -> str:
     if not settings.ENABLE_PROMETHEUS_METRICS:
         raise error_not_found()
-    custom_metrics = await prometheus.get_metrics(session=session)
-    prometheus_metrics = generate_latest()
-    return custom_metrics + prometheus_metrics.decode()
+    custom_metrics_ = await custom_metrics.get_metrics(session=session)
+    client_metrics = prometheus_client.generate_latest().decode()
+    return custom_metrics_ + client_metrics

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/logs.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from typing import Optional
 
-from pydantic import UUID4, Field
+from pydantic import UUID4, Field, validator
 
 from dstack._internal.core.models.common import CoreModel
 
@@ -12,5 +12,14 @@ class PollLogsRequest(CoreModel):
     start_time: Optional[datetime]
     end_time: Optional[datetime]
     descending: bool = False
+    next_token: Optional[str] = None
     limit: int = Field(100, ge=0, le=1000)
     diagnose: bool = False
+
+    @validator("descending")
+    @classmethod
+    def validate_descending(cls, v):
+        # Descending is not supported until we migrate from base64-encoded logs to plain text logs.
+        if v is True:
+            raise ValueError("descending: true is not supported")
+        return v

dstack/_internal/server/schemas/projects.py

@@ -11,6 +11,10 @@ class CreateProjectRequest(CoreModel):
     is_public: bool = False
 
 
+class UpdateProjectRequest(CoreModel):
+    is_public: bool
+
+
 class DeleteProjectsRequest(CoreModel):
     projects_names: List[str]
 
@@ -25,3 +29,11 @@ class MemberSetting(CoreModel):
 
 class SetProjectMembersRequest(CoreModel):
     members: List[MemberSetting]
+
+
+class AddProjectMemberRequest(CoreModel):
+    members: List[MemberSetting]
+
+
+class RemoveProjectMemberRequest(CoreModel):
+    usernames: List[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/security/permissions.py

@@ -58,7 +58,7 @@ class ProjectAdmin:
             raise error_invalid_token()
         project = await get_project_model_by_name(session=session, project_name=project_name)
         if project is None:
-            raise error_forbidden()
+            raise error_not_found()
         if user.global_role == GlobalRole.ADMIN:
             return user, project
         project_role = get_user_project_role(user=user, project=project)
@@ -68,6 +68,10 @@ class ProjectAdmin:
 
 
 class ProjectManager:
+    """
+    Allows project admins and managers to manage projects.
+    """
+
     async def __call__(
         self,
         project_name: str,
@@ -79,12 +83,15 @@ class ProjectManager:
             raise error_invalid_token()
         project = await get_project_model_by_name(session=session, project_name=project_name)
         if project is None:
-            raise error_forbidden()
+            raise error_not_found()
+
         if user.global_role == GlobalRole.ADMIN:
             return user, project
+
         project_role = get_user_project_role(user=user, project=project)
         if project_role in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
             return user, project
+
         raise error_forbidden()
 
 
@@ -135,6 +142,72 @@ class ProjectMemberOrPublicAccess:
         raise error_forbidden()
 
 
+class ProjectManagerOrPublicProject:
+    """
+    Allows:
+    1. Project managers to perform member management operations
+    2. Access to public projects for any authenticated user
+    """
+
+    def __init__(self):
+        self.project_manager = ProjectManager()
+
+    async def __call__(
+        self,
+        project_name: str,
+        session: AsyncSession = Depends(get_session),
+        token: HTTPAuthorizationCredentials = Security(HTTPBearer()),
+    ) -> Tuple[UserModel, ProjectModel]:
+        user = await log_in_with_token(session=session, token=token.credentials)
+        if user is None:
+            raise error_invalid_token()
+        project = await get_project_model_by_name(session=session, project_name=project_name)
+        if project is None:
+            raise error_not_found()
+
+        if user.global_role == GlobalRole.ADMIN:
+            return user, project
+
+        project_role = get_user_project_role(user=user, project=project)
+        if project_role in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
+            return user, project
+
+        if project.is_public:
+            return user, project
+
+        raise error_forbidden()
+
+
+class ProjectManagerOrSelfLeave:
+    """
+    Allows:
+    1. Project managers to remove any members
+    2. Any project member to leave (remove themselves)
+    """
+
+    async def __call__(
+        self,
+        project_name: str,
+        session: AsyncSession = Depends(get_session),
+        token: HTTPAuthorizationCredentials = Security(HTTPBearer()),
+    ) -> Tuple[UserModel, ProjectModel]:
+        user = await log_in_with_token(session=session, token=token.credentials)
+        if user is None:
+            raise error_invalid_token()
+        project = await get_project_model_by_name(session=session, project_name=project_name)
+        if project is None:
+            raise error_not_found()
+
+        if user.global_role == GlobalRole.ADMIN:
+            return user, project
+
+        project_role = get_user_project_role(user=user, project=project)
+        if project_role is not None:
+            return user, project
+
+        raise error_forbidden()
+
+
 class OptionalServiceAccount:
     def __init__(self, token: Optional[str]) -> None:
         self._token = token

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

@@ -532,7 +532,7 @@ async def delete_fleets(
             .options(selectinload(FleetModel.runs))
             .execution_options(populate_existing=True)
             .order_by(FleetModel.id)  # take locks in order
-            .with_for_update()
+            .with_for_update(key_share=True)
         )
         fleet_models = res.scalars().unique().all()
         fleets = [fleet_model_to_fleet(m) for m in fleet_models]

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

@@ -240,7 +240,7 @@ async def delete_gateways(
             .options(selectinload(GatewayModel.gateway_compute))
             .execution_options(populate_existing=True)
             .order_by(GatewayModel.id)  # take locks in order
-            .with_for_update()
+            .with_for_update(key_share=True)
         )
         gateway_models = res.scalars().all()
         for gateway_model in gateway_models:

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

@@ -68,8 +68,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 +103,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 +149,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
 
@@ -171,6 +189,8 @@ class JobConfigurator(ABC):
         return result
 
     def _dstack_image_commands(self) -> List[str]:
+        if self.run_spec.configuration.docker is True:
+            return ["start-dockerd"]
         if (
             self.run_spec.configuration.image is not None
             or self.run_spec.configuration.entrypoint is not None
@@ -201,7 +221,9 @@ class JobConfigurator(ABC):
         return self.run_spec.configuration.home_dir
 
     def _image_name(self) -> str:
-        if self.run_spec.configuration.image is not None:
+        if self.run_spec.configuration.docker is True:
+            return settings.DSTACK_DIND_IMAGE
+        elif self.run_spec.configuration.image is not None:
             return self.run_spec.configuration.image
         return get_default_image(nvcc=bool(self.run_spec.configuration.nvcc))
 
@@ -215,6 +237,8 @@ class JobConfigurator(ABC):
         return UnixUser.parse(user)
 
     def _privileged(self) -> bool:
+        if self.run_spec.configuration.docker is True:
+            return True
         return self.run_spec.configuration.privileged
 
     def _single_branch(self) -> bool:

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()