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

dstack/_internal/server/background/tasks/process_instances.py

@@ -45,6 +45,7 @@ from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT
 from dstack._internal.core.errors import (
     BackendError,
     NotYetTerminated,
+    PlacementGroupNotSupportedError,
     ProvisioningError,
 )
 from dstack._internal.core.models.backends.base import BackendType
@@ -73,7 +74,7 @@ from dstack._internal.core.models.runs import (
 from dstack._internal.core.services.profiles import get_retry
 from dstack._internal.server import settings as server_settings
 from dstack._internal.server.background.tasks.common import get_provisioning_timeout
-from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     FleetModel,
     InstanceModel,
@@ -110,6 +111,8 @@ from dstack._internal.utils.ssh import (
     pkey_from_str,
 )
 
+MIN_PROCESSING_INTERVAL = timedelta(seconds=10)
+
 PENDING_JOB_RETRY_INTERVAL = timedelta(seconds=60)
 
 TERMINATION_DEADLINE_OFFSET = timedelta(minutes=20)
@@ -129,7 +132,7 @@ async def process_instances(batch_size: int = 1):
 
 
 async def _process_next_instance():
-    lock, lockset = get_locker().get_lockset(InstanceModel.__tablename__)
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(InstanceModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
             res = await session.execute(
@@ -145,6 +148,8 @@ async def _process_next_instance():
                         ]
                     ),
                     InstanceModel.id.not_in(lockset),
+                    InstanceModel.last_processed_at
+                    < get_current_datetime().replace(tzinfo=None) - MIN_PROCESSING_INTERVAL,
                 )
                 .options(lazyload(InstanceModel.jobs))
                 .order_by(InstanceModel.last_processed_at.asc())
@@ -1063,6 +1068,12 @@ async def _create_placement_group(
             placement_group_model_to_placement_group(placement_group_model),
             master_instance_offer,
         )
+    except PlacementGroupNotSupportedError:
+        logger.debug(
+            "Skipping offer %s because placement group not supported",
+            master_instance_offer.instance.name,
+        )
+        return None
     except BackendError as e:
         logger.warning(
             "Failed to create placement group %s in %s/%s: %r",

dstack/_internal/server/background/tasks/process_placement_groups.py

@@ -7,7 +7,7 @@ from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport
 from dstack._internal.core.errors import PlacementGroupInUseError
-from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import PlacementGroupModel, ProjectModel
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services.locking import get_locker
@@ -19,7 +19,9 @@ logger = get_logger(__name__)
 
 
 async def process_placement_groups():
-    lock, lockset = get_locker().get_lockset(PlacementGroupModel.__tablename__)
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(
+        PlacementGroupModel.__tablename__
+    )
     async with get_session_ctx() as session:
         async with lock:
             res = await session.execute(

dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -1,5 +1,6 @@
 import asyncio
 import re
+import uuid
 from collections.abc import Iterable
 from datetime import timedelta, timezone
 from typing import Dict, List, Optional
@@ -14,6 +15,7 @@ from dstack._internal.core.errors import GatewayError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import NetworkMode, RegistryAuth
 from dstack._internal.core.models.configurations import DevEnvironmentConfiguration
+from dstack._internal.core.models.files import FileArchiveMapping
 from dstack._internal.core.models.instances import (
     InstanceStatus,
     RemoteConnectionInfo,
@@ -32,18 +34,21 @@ from dstack._internal.core.models.runs import (
     JobTerminationReason,
     Run,
     RunSpec,
+    RunStatus,
 )
 from dstack._internal.core.models.volumes import InstanceMountPoint, Volume, VolumeMountPoint
 from dstack._internal.server.background.tasks.common import get_provisioning_timeout
-from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     InstanceModel,
     JobModel,
     ProjectModel,
     RepoModel,
     RunModel,
+    UserModel,
 )
 from dstack._internal.server.schemas.runner import GPUDevice, TaskStatus
+from dstack._internal.server.services import files as files_services
 from dstack._internal.server.services import logs as logs_services
 from dstack._internal.server.services import services
 from dstack._internal.server.services.instances import get_instance_ssh_private_keys
@@ -66,14 +71,16 @@ from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
 from dstack._internal.server.services.runs import (
     run_model_to_run,
 )
+from dstack._internal.server.services.secrets import get_project_secrets_mapping
 from dstack._internal.server.services.storage import get_default_storage
 from dstack._internal.utils import common as common_utils
-from dstack._internal.utils.interpolator import VariablesInterpolator
+from dstack._internal.utils.interpolator import InterpolatorError, VariablesInterpolator
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
 
 
+MIN_PROCESSING_INTERVAL = timedelta(seconds=10)
 # Minimum time before terminating active job in case of connectivity issues.
 # Should be sufficient to survive most problems caused by
 # the server network flickering and providers' glitches.
@@ -88,20 +95,29 @@ async def process_running_jobs(batch_size: int = 1):
 
 
 async def _process_next_running_job():
-    lock, lockset = get_locker().get_lockset(JobModel.__tablename__)
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
             res = await session.execute(
                 select(JobModel)
+                .join(JobModel.run)
                 .where(
                     JobModel.status.in_(
                         [JobStatus.PROVISIONING, JobStatus.PULLING, JobStatus.RUNNING]
                     ),
+                    RunModel.status.not_in([RunStatus.TERMINATING]),
                     JobModel.id.not_in(lockset),
+                    JobModel.last_processed_at
+                    < common_utils.get_current_datetime().replace(tzinfo=None)
+                    - MIN_PROCESSING_INTERVAL,
                 )
                 .order_by(JobModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True, key_share=True)
+                .with_for_update(
+                    skip_locked=True,
+                    key_share=True,
+                    of=JobModel,
+                )
             )
             job_model = res.unique().scalar()
             if job_model is None:
@@ -177,7 +193,17 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
         common_utils.get_or_error(job_model.instance)
     )
 
-    secrets = {}  # TODO secrets
+    secrets = await get_project_secrets_mapping(session=session, project=project)
+
+    try:
+        _interpolate_secrets(secrets, job.job_spec)
+    except InterpolatorError as e:
+        logger.info("%s: terminating due to secrets interpolation error", fmt(job_model))
+        job_model.status = JobStatus.TERMINATING
+        job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+        job_model.termination_reason_message = e.args[0]
+        job_model.last_processed_at = common_utils.get_current_datetime()
+        return
 
     repo_creds_model = await get_repo_creds(session=session, repo=repo_model, user=run_model.user)
     repo_creds = repo_model_to_repo_head_with_creds(repo_model, repo_creds_model).repo_creds
@@ -214,7 +240,6 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                     job_model,
                     job_provisioning_data,
                     volumes,
-                    secrets,
                     job.job_spec.registry_auth,
                     public_keys,
                     ssh_user,
@@ -226,12 +251,20 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                     fmt(job_model),
                     job_submission.age,
                 )
+                # FIXME: downloading file archives and code here is a waste of time if
+                # the runner is not ready yet
+                file_archives = await _get_job_file_archives(
+                    session=session,
+                    archive_mappings=job.job_spec.file_archives,
+                    user=run_model.user,
+                )
                 code = await _get_job_code(
                     session=session,
                     project=project,
                     repo=repo_model,
-                    code_hash=run.run_spec.repo_code_hash,
+                    code_hash=_get_repo_code_hash(run, job),
                 )
+
                 success = await common_utils.run_async(
                     _submit_job_to_runner,
                     server_ssh_private_keys,
@@ -242,6 +275,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                     job,
                     cluster_info,
                     code,
+                    file_archives,
                     secrets,
                     repo_creds,
                     success_if_not_available=False,
@@ -269,11 +303,18 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
             logger.debug(
                 "%s: process pulling job with shim, age=%s", fmt(job_model), job_submission.age
             )
+            # FIXME: downloading file archives and code here is a waste of time if
+            # the runner is not ready yet
+            file_archives = await _get_job_file_archives(
+                session=session,
+                archive_mappings=job.job_spec.file_archives,
+                user=run_model.user,
+            )
             code = await _get_job_code(
                 session=session,
                 project=project,
                 repo=repo_model,
-                code_hash=run.run_spec.repo_code_hash,
+                code_hash=_get_repo_code_hash(run, job),
             )
             success = await common_utils.run_async(
                 _process_pulling_with_shim,
@@ -285,6 +326,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 job,
                 cluster_info,
                 code,
+                file_archives,
                 secrets,
                 repo_creds,
                 server_ssh_private_keys,
@@ -306,8 +348,9 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
         else:
             if job_model.termination_reason:
                 logger.warning(
-                    "%s: failed because shim/runner returned an error, age=%s",
+                    "%s: failed due to %s, age=%s",
                     fmt(job_model),
+                    job_model.termination_reason.value,
                     job_submission.age,
                 )
                 job_model.status = JobStatus.TERMINATING
@@ -450,7 +493,6 @@ def _process_provisioning_with_shim(
     job_model: JobModel,
     job_provisioning_data: JobProvisioningData,
     volumes: List[Volume],
-    secrets: Dict[str, str],
     registry_auth: Optional[RegistryAuth],
     public_keys: List[str],
     ssh_user: str,
@@ -476,10 +518,8 @@ def _process_provisioning_with_shim(
     registry_username = ""
     registry_password = ""
     if registry_auth is not None:
-        logger.debug("%s: authenticating to the registry...", fmt(job_model))
-        interpolate = VariablesInterpolator({"secrets": secrets}).interpolate
-        registry_username = interpolate(registry_auth.username)
-        registry_password = interpolate(registry_auth.password)
+        registry_username = registry_auth.username
+        registry_password = registry_auth.password
 
     volume_mounts: List[VolumeMountPoint] = []
     instance_mounts: List[InstanceMountPoint] = []
@@ -588,6 +628,7 @@ def _process_pulling_with_shim(
     job: Job,
     cluster_info: ClusterInfo,
     code: bytes,
+    file_archives: Iterable[tuple[uuid.UUID, bytes]],
     secrets: Dict[str, str],
     repo_credentials: Optional[RemoteRepoCreds],
     server_ssh_private_keys: tuple[str, Optional[str]],
@@ -663,6 +704,7 @@ def _process_pulling_with_shim(
         job=job,
         cluster_info=cluster_info,
         code=code,
+        file_archives=file_archives,
         secrets=secrets,
         repo_credentials=repo_credentials,
         success_if_not_available=True,
@@ -826,6 +868,19 @@ def _get_cluster_info(
     return cluster_info
 
 
+def _get_repo_code_hash(run: Run, job: Job) -> Optional[str]:
+    # TODO: drop this function when supporting jobs submitted before 0.19.17 is no longer relevant.
+    if (
+        job.job_spec.repo_code_hash is None
+        and run.run_spec.repo_code_hash is not None
+        and job.job_submissions[-1].deployment_num == run.deployment_num
+    ):
+        # The job spec does not have `repo_code_hash`, because it was submitted before 0.19.17.
+        # Use `repo_code_hash` from the run.
+        return run.run_spec.repo_code_hash
+    return job.job_spec.repo_code_hash
+
+
 async def _get_job_code(
     session: AsyncSession, project: ProjectModel, repo: RepoModel, code_hash: Optional[str]
 ) -> bytes:
@@ -853,6 +908,43 @@ async def _get_job_code(
     return blob
 
 
+async def _get_job_file_archives(
+    session: AsyncSession,
+    archive_mappings: Iterable[FileArchiveMapping],
+    user: UserModel,
+) -> list[tuple[uuid.UUID, bytes]]:
+    archives: list[tuple[uuid.UUID, bytes]] = []
+    for archive_mapping in archive_mappings:
+        archive_id = archive_mapping.id
+        archive_blob = await _get_job_file_archive(
+            session=session, archive_id=archive_id, user=user
+        )
+        archives.append((archive_id, archive_blob))
+    return archives
+
+
+async def _get_job_file_archive(
+    session: AsyncSession, archive_id: uuid.UUID, user: UserModel
+) -> bytes:
+    archive_model = await files_services.get_archive_model(session, id=archive_id, user=user)
+    if archive_model is None:
+        return b""
+    if archive_model.blob is not None:
+        return archive_model.blob
+    storage = get_default_storage()
+    if storage is None:
+        return b""
+    blob = await common_utils.run_async(
+        storage.get_archive,
+        str(archive_model.user_id),
+        archive_model.blob_hash,
+    )
+    if blob is None:
+        logger.error("Failed to get file archive %s from storage", archive_id)
+        return b""
+    return blob
+
+
 @runner_ssh_tunnel(ports=[DSTACK_RUNNER_HTTP_PORT], retries=1)
 def _submit_job_to_runner(
     ports: Dict[int, int],
@@ -861,6 +953,7 @@ def _submit_job_to_runner(
     job: Job,
     cluster_info: ClusterInfo,
     code: bytes,
+    file_archives: Iterable[tuple[uuid.UUID, bytes]],
     secrets: Dict[str, str],
     repo_credentials: Optional[RemoteRepoCreds],
     success_if_not_available: bool,
@@ -896,10 +989,15 @@ def _submit_job_to_runner(
         run=run,
         job=job,
         cluster_info=cluster_info,
-        secrets=secrets,
+        # Do not send all the secrets since interpolation is already done by the server.
+        # TODO: Passing secrets may be necessary for filtering out secret values from logs.
+        secrets={},
         repo_credentials=repo_credentials,
         instance_env=instance_env,
     )
+    logger.debug("%s: uploading file archive(s)", fmt(job_model))
+    for archive_id, archive in file_archives:
+        runner_client.upload_archive(archive_id, archive)
     logger.debug("%s: uploading code", fmt(job_model))
     runner_client.upload_code(code)
     logger.debug("%s: starting job", fmt(job_model))
@@ -911,6 +1009,16 @@ def _submit_job_to_runner(
     return True
 
 
+def _interpolate_secrets(secrets: Dict[str, str], job_spec: JobSpec):
+    interpolate = VariablesInterpolator({"secrets": secrets}).interpolate_or_error
+    job_spec.env = {k: interpolate(v) for k, v in job_spec.env.items()}
+    if job_spec.registry_auth is not None:
+        job_spec.registry_auth = RegistryAuth(
+            username=interpolate(job_spec.registry_auth.username),
+            password=interpolate(job_spec.registry_auth.password),
+        )
+
+
 def _get_instance_specific_mounts(
     backend_type: BackendType, instance_type_name: str
 ) -> List[InstanceMountPoint]:

dstack/_internal/server/background/tasks/process_runs.py

@@ -19,7 +19,7 @@ from dstack._internal.core.models.runs import (
     RunStatus,
     RunTerminationReason,
 )
-from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import JobModel, ProjectModel, RunModel
 from dstack._internal.server.services.jobs import (
     find_job,
@@ -35,11 +35,14 @@ from dstack._internal.server.services.runs import (
     run_model_to_run,
     scale_run_replicas,
 )
+from dstack._internal.server.services.secrets import get_project_secrets_mapping
 from dstack._internal.server.services.services import update_service_desired_replica_count
 from dstack._internal.utils import common
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
+
+MIN_PROCESSING_INTERVAL = datetime.timedelta(seconds=5)
 ROLLING_DEPLOYMENT_MAX_SURGE = 1  # at most one extra replica during rolling deployment
 
 
@@ -51,8 +54,8 @@ async def process_runs(batch_size: int = 1):
 
 
 async def _process_next_run():
-    run_lock, run_lockset = get_locker().get_lockset(RunModel.__tablename__)
-    job_lock, job_lockset = get_locker().get_lockset(JobModel.__tablename__)
+    run_lock, run_lockset = get_locker(get_db().dialect_name).get_lockset(RunModel.__tablename__)
+    job_lock, job_lockset = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
     async with get_session_ctx() as session:
         async with run_lock, job_lock:
             res = await session.execute(
@@ -60,6 +63,8 @@ async def _process_next_run():
                 .where(
                     RunModel.status.not_in(RunStatus.finished_statuses()),
                     RunModel.id.not_in(run_lockset),
+                    RunModel.last_processed_at
+                    < common.get_current_datetime().replace(tzinfo=None) - MIN_PROCESSING_INTERVAL,
                 )
                 .order_by(RunModel.last_processed_at.asc())
                 .limit(1)
@@ -336,7 +341,7 @@ async def _process_active_run(session: AsyncSession, run_model: RunModel):
                 current_time - run_model.submitted_at.replace(tzinfo=datetime.timezone.utc)
             ).total_seconds()
             logger.info(
-                "%s: run took %.2f seconds from submision to provisioning.",
+                "%s: run took %.2f seconds from submission to provisioning.",
                 fmt(run_model),
                 submit_to_provision_duration,
             )
@@ -404,7 +409,11 @@ async def _handle_run_replicas(
         )
         return
 
-    await _update_jobs_to_new_deployment_in_place(run_model, run_spec)
+    await _update_jobs_to_new_deployment_in_place(
+        session=session,
+        run_model=run_model,
+        run_spec=run_spec,
+    )
     if _has_out_of_date_replicas(run_model):
         non_terminated_replica_count = len(
             {j.replica_num for j in run_model.jobs if not j.status.is_finished()}
@@ -444,18 +453,25 @@ async def _handle_run_replicas(
             )
 
 
-async def _update_jobs_to_new_deployment_in_place(run_model: RunModel, run_spec: RunSpec) -> None:
+async def _update_jobs_to_new_deployment_in_place(
+    session: AsyncSession, run_model: RunModel, run_spec: RunSpec
+) -> None:
     """
     Bump deployment_num for jobs that do not require redeployment.
     """
-
+    secrets = await get_project_secrets_mapping(
+        session=session,
+        project=run_model.project,
+    )
     for replica_num, job_models in group_jobs_by_replica_latest(run_model.jobs):
         if all(j.status.is_finished() for j in job_models):
             continue
         if all(j.deployment_num == run_model.deployment_num for j in job_models):
             continue
+        # FIXME: Handle getting image configuration errors or skip it.
         new_job_specs = await get_job_specs_from_run_spec(
             run_spec=run_spec,
+            secrets=secrets,
             replica_num=replica_num,
         )
         assert len(new_job_specs) == len(job_models), (

dstack/_internal/server/background/tasks/process_submitted_jobs.py

@@ -1,5 +1,6 @@
 import asyncio
 import uuid
+from datetime import datetime, timedelta
 from typing import List, Optional, Tuple
 
 from sqlalchemy import select
@@ -80,15 +81,35 @@ from dstack._internal.utils.logging import get_logger
 logger = get_logger(__name__)
 
 
+# Track when we last processed a job.
+# This is needed for a trick:
+# If no tasks were processed recently, we force batch_size 1.
+# If there are lots of runs/jobs with same offers submitted,
+# we warm up the cache instead of requesting the offers concurrently.
+# Mostly useful when runs are submitted via API without getting run plan first.
+BATCH_SIZE_RESET_TIMEOUT = timedelta(minutes=2)
+last_processed_at: Optional[datetime] = None
+
+
 async def process_submitted_jobs(batch_size: int = 1):
     tasks = []
-    for _ in range(batch_size):
+    effective_batch_size = _get_effective_batch_size(batch_size)
+    for _ in range(effective_batch_size):
         tasks.append(_process_next_submitted_job())
     await asyncio.gather(*tasks)
 
 
+def _get_effective_batch_size(batch_size: int) -> int:
+    if (
+        last_processed_at is None
+        or last_processed_at < common_utils.get_current_datetime() - BATCH_SIZE_RESET_TIMEOUT
+    ):
+        return 1
+    return batch_size
+
+
 async def _process_next_submitted_job():
-    lock, lockset = get_locker().get_lockset(JobModel.__tablename__)
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
             res = await session.execute(
@@ -125,6 +146,8 @@ async def _process_next_submitted_job():
             await _process_submitted_job(session=session, job_model=job_model)
         finally:
             lockset.difference_update([job_model_id])
+        global last_processed_at
+        last_processed_at = common_utils.get_current_datetime()
 
 
 async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
@@ -214,7 +237,9 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         if get_db().dialect_name == "sqlite":
             # Start new transaction to see committed changes after lock
             await session.commit()
-        async with get_locker().lock_ctx(InstanceModel.__tablename__, instances_ids):
+        async with get_locker(get_db().dialect_name).lock_ctx(
+            InstanceModel.__tablename__, instances_ids
+        ):
             # If another job freed the instance but is still trying to detach volumes,
             # do not provision on it to prevent attaching volumes that are currently detaching.
             detaching_instances_ids = await get_instances_ids_with_detaching_volumes(session)
@@ -243,8 +268,10 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
             )
             job_model.instance_assigned = True
             job_model.last_processed_at = common_utils.get_current_datetime()
-            await session.commit()
-            return
+            if len(pool_instances) > 0:
+                await session.commit()
+                return
+            # If no instances were locked, we can proceed in the same transaction.
 
     if job_model.instance is not None:
         res = await session.execute(
@@ -334,7 +361,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         .order_by(VolumeModel.id)  # take locks in order
         .with_for_update(key_share=True)
     )
-    async with get_locker().lock_ctx(VolumeModel.__tablename__, volumes_ids):
+    async with get_locker(get_db().dialect_name).lock_ctx(VolumeModel.__tablename__, volumes_ids):
         if len(volume_models) > 0:
             await _attach_volumes(
                 session=session,
@@ -527,7 +554,9 @@ async def _get_next_instance_num(session: AsyncSession, fleet_model: FleetModel)
     if len(fleet_model.instances) == 0:
         # No instances means the fleet is not in the db yet, so don't lock.
         return 0
-    async with get_locker().lock_ctx(FleetModel.__tablename__, [fleet_model.id]):
+    async with get_locker(get_db().dialect_name).lock_ctx(
+        FleetModel.__tablename__, [fleet_model.id]
+    ):
         fleet_model = (
             (
                 await session.execute(

dstack/_internal/server/background/tasks/process_terminating_jobs.py

@@ -5,7 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload, lazyload
 
 from dstack._internal.core.models.runs import JobStatus
-from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     InstanceModel,
     JobModel,
@@ -32,8 +32,10 @@ async def process_terminating_jobs(batch_size: int = 1):
 
 
 async def _process_next_terminating_job():
-    job_lock, job_lockset = get_locker().get_lockset(JobModel.__tablename__)
-    instance_lock, instance_lockset = get_locker().get_lockset(InstanceModel.__tablename__)
+    job_lock, job_lockset = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
+    instance_lock, instance_lockset = get_locker(get_db().dialect_name).get_lockset(
+        InstanceModel.__tablename__
+    )
     async with get_session_ctx() as session:
         async with job_lock, instance_lock:
             res = await session.execute(

dstack/_internal/server/background/tasks/process_volumes.py

@@ -5,7 +5,7 @@ from sqlalchemy.orm import joinedload
 from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport
 from dstack._internal.core.errors import BackendError, BackendNotAvailable
 from dstack._internal.core.models.volumes import VolumeStatus
-from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     InstanceModel,
     ProjectModel,
@@ -22,7 +22,7 @@ logger = get_logger(__name__)
 
 
 async def process_submitted_volumes():
-    lock, lockset = get_locker().get_lockset(VolumeModel.__tablename__)
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(VolumeModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
             res = await session.execute(

dstack/_internal/server/migrations/versions/5f1707c525d2_add_filearchivemodel.py

@@ -0,0 +1,39 @@
+"""Add FileArchiveModel
+
+Revision ID: 5f1707c525d2
+Revises: 35e90e1b0d3e
+Create Date: 2025-06-12 12:28:26.678380
+
+"""
+
+import sqlalchemy as sa
+import sqlalchemy_utils
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "5f1707c525d2"
+down_revision = "35e90e1b0d3e"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "file_archives",
+        sa.Column("id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False),
+        sa.Column("user_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False),
+        sa.Column("blob_hash", sa.Text(), nullable=False),
+        sa.Column("blob", sa.LargeBinary(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["users.id"],
+            name=op.f("fk_file_archives_user_id_users"),
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_file_archives")),
+        sa.UniqueConstraint("user_id", "blob_hash", name="uq_file_archives_user_id_blob_hash"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("file_archives")