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

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/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
 

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

dstack/_internal/server/services/proxy/repo.py

@@ -7,6 +7,7 @@ from sqlalchemy.orm import joinedload
 
 import dstack._internal.server.services.jobs as jobs_services
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
+from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.configurations import ServiceConfiguration
 from dstack._internal.core.models.instances import RemoteConnectionInfo, SSHConnectionParams
 from dstack._internal.core.models.runs import (
@@ -86,6 +87,8 @@ class ServerProxyRepo(BaseProxyRepo):
                     username=jpd.username,
                     port=jpd.ssh_port,
                 )
+                if jpd.backend == BackendType.LOCAL:
+                    ssh_proxy = None
             ssh_head_proxy: Optional[SSHConnectionParams] = None
             ssh_head_proxy_private_key: Optional[str] = None
             instance = get_or_error(job.instance)

dstack/_internal/server/services/runner/client.py

@@ -109,6 +109,14 @@ class RunnerClient:
         )
         resp.raise_for_status()
 
+    def upload_archive(self, id: uuid.UUID, file: Union[BinaryIO, bytes]):
+        resp = requests.post(
+            self._url("/api/upload_archive"),
+            files={"archive": (str(id), file)},
+            timeout=UPLOAD_CODE_REQUEST_TIMEOUT,
+        )
+        resp.raise_for_status()
+
     def upload_code(self, file: Union[BinaryIO, bytes]):
         resp = requests.post(
             self._url("/api/upload_code"), data=file, timeout=UPLOAD_CODE_REQUEST_TIMEOUT

dstack/_internal/server/services/runs.py

@@ -82,6 +82,7 @@ from dstack._internal.server.services.offers import get_offers_by_requirements
 from dstack._internal.server.services.plugins import apply_plugin_policies
 from dstack._internal.server.services.projects import list_project_models, list_user_project_models
 from dstack._internal.server.services.resources import set_resources_defaults
+from dstack._internal.server.services.secrets import get_project_secrets_mapping
 from dstack._internal.server.services.users import get_user_model_by_name
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.random_names import generate_name
@@ -311,7 +312,12 @@ async def get_plan(
             ):
                 action = ApplyAction.UPDATE
 
-    jobs = await get_jobs_from_run_spec(effective_run_spec, replica_num=0)
+    secrets = await get_project_secrets_mapping(session=session, project=project)
+    jobs = await get_jobs_from_run_spec(
+        run_spec=effective_run_spec,
+        secrets=secrets,
+        replica_num=0,
+    )
 
     volumes = await get_job_configured_volumes(
         session=session,
@@ -462,6 +468,10 @@ async def submit_run(
         project=project,
         run_spec=run_spec,
     )
+    secrets = await get_project_secrets_mapping(
+        session=session,
+        project=project,
+    )
 
     lock_namespace = f"run_names_{project.name}"
     if get_db().dialect_name == "sqlite":
@@ -513,7 +523,11 @@ async def submit_run(
             await services.register_service(session, run_model, run_spec)
 
         for replica_num in range(replicas):
-            jobs = await get_jobs_from_run_spec(run_spec, replica_num=replica_num)
+            jobs = await get_jobs_from_run_spec(
+                run_spec=run_spec,
+                secrets=secrets,
+                replica_num=replica_num,
+            )
             for job in jobs:
                 job_model = create_job_model_for_new_submission(
                     run_model=run_model,
@@ -898,7 +912,16 @@ def _validate_run_spec_and_set_defaults(run_spec: RunSpec):
     set_resources_defaults(run_spec.configuration.resources)
 
 
-_UPDATABLE_SPEC_FIELDS = ["repo_code_hash", "configuration"]
+_UPDATABLE_SPEC_FIELDS = ["configuration_path", "configuration"]
+_TYPE_SPECIFIC_UPDATABLE_SPEC_FIELDS = {
+    "service": [
+        # rolling deployment
+        "repo_data",
+        "repo_code_hash",
+        "file_archives",
+        "working_dir",
+    ],
+}
 _CONF_UPDATABLE_FIELDS = ["priority"]
 _TYPE_SPECIFIC_CONF_UPDATABLE_FIELDS = {
     "dev-environment": ["inactivity_duration"],
@@ -909,10 +932,13 @@ _TYPE_SPECIFIC_CONF_UPDATABLE_FIELDS = {
         # rolling deployment
         "resources",
         "volumes",
+        "docker",
+        "files",
         "image",
         "user",
         "privileged",
         "entrypoint",
+        "working_dir",
         "python",
         "nvcc",
         "single_branch",
@@ -935,11 +961,14 @@ def _can_update_run_spec(current_run_spec: RunSpec, new_run_spec: RunSpec) -> bo
 def _check_can_update_run_spec(current_run_spec: RunSpec, new_run_spec: RunSpec):
     spec_diff = diff_models(current_run_spec, new_run_spec)
     changed_spec_fields = list(spec_diff.keys())
+    updatable_spec_fields = _UPDATABLE_SPEC_FIELDS + _TYPE_SPECIFIC_UPDATABLE_SPEC_FIELDS.get(
+        new_run_spec.configuration.type, []
+    )
     for key in changed_spec_fields:
-        if key not in _UPDATABLE_SPEC_FIELDS:
+        if key not in updatable_spec_fields:
             raise ServerClientError(
                 f"Failed to update fields {changed_spec_fields}."
-                f" Can only update {_UPDATABLE_SPEC_FIELDS}."
+                f" Can only update {updatable_spec_fields}."
             )
     _check_can_update_configuration(current_run_spec.configuration, new_run_spec.configuration)
 
@@ -1068,10 +1097,20 @@ async def scale_run_replicas(session: AsyncSession, run_model: RunModel, replica
             await retry_run_replica_jobs(session, run_model, replica_jobs, only_failed=False)
             scheduled_replicas += 1
 
+        secrets = await get_project_secrets_mapping(
+            session=session,
+            project=run_model.project,
+        )
+
         for replica_num in range(
             len(active_replicas) + scheduled_replicas, len(active_replicas) + replicas_diff
         ):
-            jobs = await get_jobs_from_run_spec(run_spec, replica_num=replica_num)
+            # FIXME: Handle getting image configuration errors or skip it.
+            jobs = await get_jobs_from_run_spec(
+                run_spec=run_spec,
+                secrets=secrets,
+                replica_num=replica_num,
+            )
             for job in jobs:
                 job_model = create_job_model_for_new_submission(
                     run_model=run_model,
@@ -1084,8 +1123,14 @@ async def scale_run_replicas(session: AsyncSession, run_model: RunModel, replica
 async def retry_run_replica_jobs(
     session: AsyncSession, run_model: RunModel, latest_jobs: List[JobModel], *, only_failed: bool
 ):
+    # FIXME: Handle getting image configuration errors or skip it.
+    secrets = await get_project_secrets_mapping(
+        session=session,
+        project=run_model.project,
+    )
     new_jobs = await get_jobs_from_run_spec(
-        RunSpec.__response__.parse_raw(run_model.run_spec),
+        run_spec=RunSpec.__response__.parse_raw(run_model.run_spec),
+        secrets=secrets,
         replica_num=latest_jobs[0].replica_num,
     )
     assert len(new_jobs) == len(latest_jobs), (

dstack/_internal/server/services/secrets.py

@@ -0,0 +1,204 @@
+import re
+from typing import Dict, List, Optional
+
+import sqlalchemy.exc
+from sqlalchemy import delete, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.errors import (
+    ResourceExistsError,
+    ResourceNotExistsError,
+    ServerClientError,
+)
+from dstack._internal.core.models.secrets import Secret
+from dstack._internal.server.models import DecryptedString, ProjectModel, SecretModel
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+_SECRET_NAME_REGEX = "^[A-Za-z0-9-_]{1,200}$"
+_SECRET_VALUE_MAX_LENGTH = 2000
+
+
+async def list_secrets(
+    session: AsyncSession,
+    project: ProjectModel,
+) -> List[Secret]:
+    secret_models = await list_project_secret_models(session=session, project=project)
+    return [secret_model_to_secret(s, include_value=False) for s in secret_models]
+
+
+async def get_project_secrets_mapping(
+    session: AsyncSession,
+    project: ProjectModel,
+) -> Dict[str, str]:
+    secret_models = await list_project_secret_models(session=session, project=project)
+    return {s.name: s.value.get_plaintext_or_error() for s in secret_models}
+
+
+async def get_secret(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+) -> Optional[Secret]:
+    secret_model = await get_project_secret_model_by_name(
+        session=session,
+        project=project,
+        name=name,
+    )
+    if secret_model is None:
+        return None
+    return secret_model_to_secret(secret_model, include_value=True)
+
+
+async def create_or_update_secret(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+    value: str,
+) -> Secret:
+    _validate_secret(name=name, value=value)
+    try:
+        secret_model = await create_secret(
+            session=session,
+            project=project,
+            name=name,
+            value=value,
+        )
+    except ResourceExistsError:
+        secret_model = await update_secret(
+            session=session,
+            project=project,
+            name=name,
+            value=value,
+        )
+    return secret_model_to_secret(secret_model, include_value=True)
+
+
+async def delete_secrets(
+    session: AsyncSession,
+    project: ProjectModel,
+    names: List[str],
+):
+    existing_secrets_query = await session.execute(
+        select(SecretModel).where(
+            SecretModel.project_id == project.id,
+            SecretModel.name.in_(names),
+        )
+    )
+    existing_names = [s.name for s in existing_secrets_query.scalars().all()]
+    missing_names = set(names) - set(existing_names)
+    if missing_names:
+        raise ResourceNotExistsError(f"Secrets not found: {', '.join(missing_names)}")
+
+    await session.execute(
+        delete(SecretModel).where(
+            SecretModel.project_id == project.id,
+            SecretModel.name.in_(names),
+        )
+    )
+    await session.commit()
+    logger.info("Deleted secrets %s in project %s", names, project.name)
+
+
+def secret_model_to_secret(secret_model: SecretModel, include_value: bool = False) -> Secret:
+    value = None
+    if include_value:
+        value = secret_model.value.get_plaintext_or_error()
+    return Secret(
+        id=secret_model.id,
+        name=secret_model.name,
+        value=value,
+    )
+
+
+async def list_project_secret_models(
+    session: AsyncSession,
+    project: ProjectModel,
+) -> List[SecretModel]:
+    res = await session.execute(
+        select(SecretModel)
+        .where(
+            SecretModel.project_id == project.id,
+        )
+        .order_by(SecretModel.created_at.desc())
+    )
+    secret_models = list(res.scalars().all())
+    return secret_models
+
+
+async def get_project_secret_model_by_name(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+) -> Optional[SecretModel]:
+    res = await session.execute(
+        select(SecretModel).where(
+            SecretModel.project_id == project.id,
+            SecretModel.name == name,
+        )
+    )
+    return res.scalar_one_or_none()
+
+
+async def create_secret(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+    value: str,
+) -> SecretModel:
+    secret_model = SecretModel(
+        project_id=project.id,
+        name=name,
+        value=DecryptedString(plaintext=value),
+    )
+    try:
+        async with session.begin_nested():
+            session.add(secret_model)
+    except sqlalchemy.exc.IntegrityError:
+        raise ResourceExistsError()
+    await session.commit()
+    return secret_model
+
+
+async def update_secret(
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+    value: str,
+) -> SecretModel:
+    await session.execute(
+        update(SecretModel)
+        .where(
+            SecretModel.project_id == project.id,
+            SecretModel.name == name,
+        )
+        .values(
+            value=DecryptedString(plaintext=value),
+        )
+    )
+    await session.commit()
+    secret_model = await get_project_secret_model_by_name(
+        session=session,
+        project=project,
+        name=name,
+    )
+    if secret_model is None:
+        raise ResourceNotExistsError()
+    return secret_model
+
+
+def _validate_secret(name: str, value: str):
+    _validate_secret_name(name)
+    _validate_secret_value(value)
+
+
+def _validate_secret_name(name: str):
+    if re.match(_SECRET_NAME_REGEX, name) is None:
+        raise ServerClientError(f"Secret name should match regex '{_SECRET_NAME_REGEX}")
+
+
+def _validate_secret_value(value: str):
+    if len(value) > _SECRET_VALUE_MAX_LENGTH:
+        raise ServerClientError(f"Secret value length must not exceed {_SECRET_VALUE_MAX_LENGTH}")

dstack/_internal/server/services/storage/base.py

@@ -22,6 +22,27 @@ class BaseStorage(ABC):
     ) -> Optional[bytes]:
         pass
 
+    @abstractmethod
+    def upload_archive(
+        self,
+        user_id: str,
+        archive_hash: str,
+        blob: bytes,
+    ):
+        pass
+
+    @abstractmethod
+    def get_archive(
+        self,
+        user_id: str,
+        archive_hash: str,
+    ) -> Optional[bytes]:
+        pass
+
     @staticmethod
     def _get_code_key(project_id: str, repo_id: str, code_hash: str) -> str:
         return f"data/projects/{project_id}/codes/{repo_id}/{code_hash}"
+
+    @staticmethod
+    def _get_archive_key(user_id: str, archive_hash: str) -> str:
+        return f"data/users/{user_id}/file_archives/{archive_hash}"

dstack/_internal/server/services/storage/gcs.py

@@ -25,9 +25,8 @@ class GCSStorage(BaseStorage):
         code_hash: str,
         blob: bytes,
     ):
-        blob_name = self._get_code_key(project_id, repo_id, code_hash)
-        blob_obj = self._bucket.blob(blob_name)
-        blob_obj.upload_from_string(blob)
+        key = self._get_code_key(project_id, repo_id, code_hash)
+        self._upload(key, blob)
 
     def get_code(
         self,
@@ -35,10 +34,33 @@ class GCSStorage(BaseStorage):
         repo_id: str,
         code_hash: str,
     ) -> Optional[bytes]:
+        key = self._get_code_key(project_id, repo_id, code_hash)
+        return self._get(key)
+
+    def upload_archive(
+        self,
+        user_id: str,
+        archive_hash: str,
+        blob: bytes,
+    ):
+        key = self._get_archive_key(user_id, archive_hash)
+        self._upload(key, blob)
+
+    def get_archive(
+        self,
+        user_id: str,
+        archive_hash: str,
+    ) -> Optional[bytes]:
+        key = self._get_archive_key(user_id, archive_hash)
+        return self._get(key)
+
+    def _upload(self, key: str, blob: bytes):
+        blob_obj = self._bucket.blob(key)
+        blob_obj.upload_from_string(blob)
+
+    def _get(self, key: str) -> Optional[bytes]:
         try:
-            blob_name = self._get_code_key(project_id, repo_id, code_hash)
-            blob = self._bucket.blob(blob_name)
+            blob = self._bucket.blob(key)
         except NotFound:
             return None
-
         return blob.download_as_bytes()