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

dstack/_internal/server/services/locking.py

@@ -1,8 +1,10 @@
 import asyncio
+import collections.abc
 import hashlib
+from abc import abstractmethod
 from asyncio import Lock
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, Dict, List, Set, Tuple, TypeVar, Union
+from typing import AsyncGenerator, Iterable, Iterator, Protocol, TypeVar, Union
 
 from sqlalchemy import func, select
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
@@ -10,23 +12,54 @@ from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 KeyT = TypeVar("KeyT")
 
 
-class ResourceLocker:
-    def __init__(self):
-        self.namespace_to_locks_map: Dict[str, Tuple[Lock, set]] = {}
+class LocksetLock(Protocol):
+    async def acquire(self) -> bool: ...
+    def release(self) -> None: ...
+    async def __aenter__(self): ...
+    async def __aexit__(self, exc_type, exc, tb): ...
+
+
+T = TypeVar("T")
+
 
-    def get_lockset(self, namespace: str) -> Tuple[Lock, set]:
+class Lockset(Protocol[T]):
+    def __contains__(self, item: T) -> bool: ...
+    def __iter__(self) -> Iterator[T]: ...
+    def __len__(self) -> int: ...
+    def add(self, item: T) -> None: ...
+    def discard(self, item: T) -> None: ...
+    def update(self, other: Iterable[T]) -> None: ...
+    def difference_update(self, other: Iterable[T]) -> None: ...
+
+
+class ResourceLocker:
+    @abstractmethod
+    def get_lockset(self, namespace: str) -> tuple[LocksetLock, Lockset]:
         """
         Returns a lockset containing locked resources for in-memory locking.
         Also returns a lock that guards the lockset.
         """
-        return self.namespace_to_locks_map.setdefault(namespace, (Lock(), set()))
+        pass
 
+    @abstractmethod
     @asynccontextmanager
-    async def lock_ctx(self, namespace: str, keys: List[KeyT]):
+    async def lock_ctx(self, namespace: str, keys: list[KeyT]):
         """
         Acquires locks for all keys in namespace.
         The keys must be sorted to prevent deadlock.
         """
+        yield
+
+
+class InMemoryResourceLocker(ResourceLocker):
+    def __init__(self):
+        self.namespace_to_locks_map: dict[str, tuple[Lock, set]] = {}
+
+    def get_lockset(self, namespace: str) -> tuple[Lock, set]:
+        return self.namespace_to_locks_map.setdefault(namespace, (Lock(), set()))
+
+    @asynccontextmanager
+    async def lock_ctx(self, namespace: str, keys: list[KeyT]):
         lock, lockset = self.get_lockset(namespace)
         try:
             await _wait_to_lock_many(lock, lockset, keys)
@@ -35,6 +68,56 @@ class ResourceLocker:
             lockset.difference_update(keys)
 
 
+class DummyAsyncLock:
+    async def __aenter__(self):
+        pass
+
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
+    async def acquire(self):
+        return True
+
+    def release(self):
+        pass
+
+
+class DummySet(collections.abc.MutableSet):
+    def __contains__(self, item):
+        return False
+
+    def __iter__(self):
+        return iter(())
+
+    def __len__(self):
+        return 0
+
+    def add(self, value):
+        pass
+
+    def discard(self, value):
+        pass
+
+    def update(self, other):
+        pass
+
+    def difference_update(self, other):
+        pass
+
+
+class DummyResourceLocker(ResourceLocker):
+    def __init__(self):
+        self.lock = DummyAsyncLock()
+        self.lockset = DummySet()
+
+    def get_lockset(self, namespace: str) -> tuple[DummyAsyncLock, DummySet]:
+        return self.lock, self.lockset
+
+    @asynccontextmanager
+    async def lock_ctx(self, namespace: str, keys: list[KeyT]):
+        yield
+
+
 def string_to_lock_id(s: str) -> int:
     return int(hashlib.sha256(s.encode()).hexdigest(), 16) % (2**63)
 
@@ -67,15 +150,21 @@ async def try_advisory_lock_ctx(
             await bind.execute(select(func.pg_advisory_unlock(string_to_lock_id(resource))))
 
 
-_locker = ResourceLocker()
+_in_memory_locker = InMemoryResourceLocker()
+_dummy_locker = DummyResourceLocker()
 
 
-def get_locker() -> ResourceLocker:
-    return _locker
+def get_locker(dialect_name: str) -> ResourceLocker:
+    if dialect_name == "sqlite":
+        return _in_memory_locker
+    # We could use an in-memory locker on Postgres
+    # but it can lead to unnecessary lock contention,
+    # so we use a dummy locker that does not take any locks.
+    return _dummy_locker
 
 
 async def _wait_to_lock_many(
-    lock: asyncio.Lock, locked: Set[KeyT], keys: List[KeyT], *, delay: float = 0.1
+    lock: asyncio.Lock, locked: set[KeyT], keys: list[KeyT], *, delay: float = 0.1
 ):
     """
     Retry locking until all the keys are locked.
@@ -88,7 +177,7 @@ async def _wait_to_lock_many(
             locked_now_num = 0
             for key in left_to_lock:
                 if key in locked:
-                    # Someone already aquired the lock, wait
+                    # Someone already acquired the lock, wait
                     break
                 locked.add(key)
                 locked_now_num += 1

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":
@@ -472,8 +482,9 @@ async def submit_run(
             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:
+        # FIXME: delete_runs commits, so Postgres lock is released too early.
         if run_spec.run_name is None:
             run_spec.run_name = await _generate_run_name(
                 session=session,
@@ -513,7 +524,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,
@@ -572,46 +587,29 @@ async def stop_runs(
     )
     run_models = res.scalars().all()
     run_ids = sorted([r.id for r in run_models])
-    res = await session.execute(select(JobModel).where(JobModel.run_id.in_(run_ids)))
-    job_models = res.scalars().all()
-    job_ids = sorted([j.id for j in job_models])
     await session.commit()
-    async with (
-        get_locker().lock_ctx(RunModel.__tablename__, run_ids),
-        get_locker().lock_ctx(JobModel.__tablename__, job_ids),
-    ):
+    async with get_locker(get_db().dialect_name).lock_ctx(RunModel.__tablename__, run_ids):
+        res = await session.execute(
+            select(RunModel)
+            .where(RunModel.id.in_(run_ids))
+            .order_by(RunModel.id)  # take locks in order
+            .with_for_update(key_share=True)
+            .execution_options(populate_existing=True)
+        )
+        run_models = res.scalars().all()
+        now = common_utils.get_current_datetime()
         for run_model in run_models:
-            await stop_run(session=session, run_model=run_model, abort=abort)
-
-
-async def stop_run(session: AsyncSession, run_model: RunModel, abort: bool):
-    res = await session.execute(
-        select(RunModel)
-        .where(RunModel.id == run_model.id)
-        .order_by(RunModel.id)  # take locks in order
-        .with_for_update(key_share=True)
-        .execution_options(populate_existing=True)
-    )
-    run_model = res.scalar_one()
-    await session.execute(
-        select(JobModel)
-        .where(JobModel.run_id == run_model.id)
-        .order_by(JobModel.id)  # take locks in order
-        .with_for_update(key_share=True)
-        .execution_options(populate_existing=True)
-    )
-    if run_model.status.is_finished():
-        return
-    run_model.status = RunStatus.TERMINATING
-    if abort:
-        run_model.termination_reason = RunTerminationReason.ABORTED_BY_USER
-    else:
-        run_model.termination_reason = RunTerminationReason.STOPPED_BY_USER
-    # process the run out of turn
-    logger.debug("%s: terminating because %s", fmt(run_model), run_model.termination_reason.name)
-    await process_terminating_run(session, run_model)
-    run_model.last_processed_at = common_utils.get_current_datetime()
-    await session.commit()
+            if run_model.status.is_finished():
+                continue
+            run_model.status = RunStatus.TERMINATING
+            if abort:
+                run_model.termination_reason = RunTerminationReason.ABORTED_BY_USER
+            else:
+                run_model.termination_reason = RunTerminationReason.STOPPED_BY_USER
+            run_model.last_processed_at = now
+            # The run will be terminated by process_runs.
+            # Terminating synchronously is problematic since it may take a long time.
+        await session.commit()
 
 
 async def delete_runs(
@@ -628,7 +626,7 @@ async def delete_runs(
     run_models = res.scalars().all()
     run_ids = sorted([r.id for r in run_models])
     await session.commit()
-    async with get_locker().lock_ctx(RunModel.__tablename__, run_ids):
+    async with get_locker(get_db().dialect_name).lock_ctx(RunModel.__tablename__, run_ids):
         res = await session.execute(
             select(RunModel)
             .where(RunModel.id.in_(run_ids))
@@ -898,7 +896,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 +916,13 @@ _TYPE_SPECIFIC_CONF_UPDATABLE_FIELDS = {
         # rolling deployment
         "resources",
         "volumes",
+        "docker",
+        "files",
         "image",
         "user",
         "privileged",
         "entrypoint",
+        "working_dir",
         "python",
         "nvcc",
         "single_branch",
@@ -935,11 +945,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 +1081,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 +1107,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()