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

dstack/_internal/core/compatibility/gateways.py

@@ -1,34 +1,35 @@
-from typing import Dict
-
+from dstack._internal.core.models.common import IncludeExcludeDictType
 from dstack._internal.core.models.gateways import GatewayConfiguration, GatewaySpec
 
 
-def get_gateway_spec_excludes(gateway_spec: GatewaySpec) -> Dict:
+def get_gateway_spec_excludes(gateway_spec: GatewaySpec) -> IncludeExcludeDictType:
     """
     Returns `gateway_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes = {}
+    spec_excludes: IncludeExcludeDictType = {}
     spec_excludes["configuration"] = _get_gateway_configuration_excludes(
         gateway_spec.configuration
     )
     return spec_excludes
 
 
-def get_create_gateway_excludes(configuration: GatewayConfiguration) -> Dict:
+def get_create_gateway_excludes(configuration: GatewayConfiguration) -> IncludeExcludeDictType:
     """
     Returns an exclude mapping to exclude certain fields from the create gateway request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    create_gateway_excludes = {}
+    create_gateway_excludes: IncludeExcludeDictType = {}
     create_gateway_excludes["configuration"] = _get_gateway_configuration_excludes(configuration)
     return create_gateway_excludes
 
 
-def _get_gateway_configuration_excludes(configuration: GatewayConfiguration) -> Dict:
-    configuration_excludes = {}
+def _get_gateway_configuration_excludes(
+    configuration: GatewayConfiguration,
+) -> IncludeExcludeDictType:
+    configuration_excludes: IncludeExcludeDictType = {}
     if configuration.tags is None:
         configuration_excludes["tags"] = True
     return configuration_excludes

dstack/_internal/core/compatibility/logs.py

@@ -1,15 +1,16 @@
-from typing import Dict, Optional
+from typing import Optional
 
+from dstack._internal.core.models.common import IncludeExcludeDictType
 from dstack._internal.server.schemas.logs import PollLogsRequest
 
 
-def get_poll_logs_excludes(request: PollLogsRequest) -> Optional[Dict]:
+def get_poll_logs_excludes(request: PollLogsRequest) -> Optional[IncludeExcludeDictType]:
     """
     Returns exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    excludes = {}
+    excludes: IncludeExcludeDictType = {}
     if request.next_token is None:
         excludes["next_token"] = True
     return excludes if excludes else None

dstack/_internal/core/compatibility/runs.py

@@ -1,29 +1,30 @@
-from typing import Any, Dict, Optional
+from typing import Optional
 
+from dstack._internal.core.models.common import IncludeExcludeDictType, IncludeExcludeSetType
 from dstack._internal.core.models.configurations import ServiceConfiguration
 from dstack._internal.core.models.runs import ApplyRunPlanInput, JobSpec, JobSubmission, RunSpec
 from dstack._internal.server.schemas.runs import GetRunPlanRequest
 
 
-def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
+def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeDictType]:
     """
     Returns `plan` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    apply_plan_excludes = {}
+    apply_plan_excludes: IncludeExcludeDictType = {}
     run_spec_excludes = get_run_spec_excludes(plan.run_spec)
     if run_spec_excludes is not None:
         apply_plan_excludes["run_spec"] = run_spec_excludes
     current_resource = plan.current_resource
     if current_resource is not None:
-        current_resource_excludes = {}
+        current_resource_excludes: IncludeExcludeDictType = {}
         current_resource_excludes["status_message"] = True
         if current_resource.deployment_num == 0:
             current_resource_excludes["deployment_num"] = True
         apply_plan_excludes["current_resource"] = current_resource_excludes
         current_resource_excludes["run_spec"] = get_run_spec_excludes(current_resource.run_spec)
-        job_submissions_excludes = {}
+        job_submissions_excludes: IncludeExcludeDictType = {}
         current_resource_excludes["jobs"] = {
             "__all__": {
                 "job_spec": get_job_spec_excludes([job.job_spec for job in current_resource.jobs]),
@@ -45,7 +46,7 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
             job_submissions_excludes["deployment_num"] = True
         latest_job_submission = current_resource.latest_job_submission
         if latest_job_submission is not None:
-            latest_job_submission_excludes = {}
+            latest_job_submission_excludes: IncludeExcludeDictType = {}
             current_resource_excludes["latest_job_submission"] = latest_job_submission_excludes
             if _should_exclude_job_submission_jpd_cpu_arch(latest_job_submission):
                 latest_job_submission_excludes["job_provisioning_data"] = {
@@ -62,12 +63,12 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
     return {"plan": apply_plan_excludes}
 
 
-def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[Dict]:
+def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[IncludeExcludeDictType]:
     """
     Excludes new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    get_plan_excludes = {}
+    get_plan_excludes: IncludeExcludeDictType = {}
     run_spec_excludes = get_run_spec_excludes(request.run_spec)
     if run_spec_excludes is not None:
         get_plan_excludes["run_spec"] = run_spec_excludes
@@ -76,15 +77,15 @@ def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[Dict]:
     return get_plan_excludes
 
 
-def get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
+def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType:
     """
     Returns `run_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes: dict[str, Any] = {}
-    configuration_excludes: dict[str, Any] = {}
-    profile_excludes: set[str] = set()
+    spec_excludes: IncludeExcludeDictType = {}
+    configuration_excludes: IncludeExcludeDictType = {}
+    profile_excludes: IncludeExcludeSetType = set()
     configuration = run_spec.configuration
     profile = run_spec.profile
 
@@ -121,18 +122,16 @@ def get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
         spec_excludes["configuration"] = configuration_excludes
     if profile_excludes:
         spec_excludes["profile"] = profile_excludes
-    if spec_excludes:
-        return spec_excludes
-    return None
+    return spec_excludes
 
 
-def get_job_spec_excludes(job_specs: list[JobSpec]) -> Optional[dict]:
+def get_job_spec_excludes(job_specs: list[JobSpec]) -> IncludeExcludeDictType:
     """
     Returns `job_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes: dict[str, Any] = {}
+    spec_excludes: IncludeExcludeDictType = {}
 
     if all(s.repo_code_hash is None for s in job_specs):
         spec_excludes["repo_code_hash"] = True
@@ -141,9 +140,7 @@ def get_job_spec_excludes(job_specs: list[JobSpec]) -> Optional[dict]:
     if all(not s.file_archives for s in job_specs):
         spec_excludes["file_archives"] = True
 
-    if spec_excludes:
-        return spec_excludes
-    return None
+    return spec_excludes
 
 
 def _should_exclude_job_submission_jpd_cpu_arch(job_submission: JobSubmission) -> bool:

dstack/_internal/core/compatibility/volumes.py

@@ -1,32 +1,33 @@
-from typing import Dict
-
+from dstack._internal.core.models.common import IncludeExcludeDictType
 from dstack._internal.core.models.volumes import VolumeConfiguration, VolumeSpec
 
 
-def get_volume_spec_excludes(volume_spec: VolumeSpec) -> Dict:
+def get_volume_spec_excludes(volume_spec: VolumeSpec) -> IncludeExcludeDictType:
     """
     Returns `volume_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes = {}
+    spec_excludes: IncludeExcludeDictType = {}
     spec_excludes["configuration"] = _get_volume_configuration_excludes(volume_spec.configuration)
     return spec_excludes
 
 
-def get_create_volume_excludes(configuration: VolumeConfiguration) -> Dict:
+def get_create_volume_excludes(configuration: VolumeConfiguration) -> IncludeExcludeDictType:
     """
     Returns an exclude mapping to exclude certain fields from the create volume request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    create_volume_excludes = {}
+    create_volume_excludes: IncludeExcludeDictType = {}
     create_volume_excludes["configuration"] = _get_volume_configuration_excludes(configuration)
     return create_volume_excludes
 
 
-def _get_volume_configuration_excludes(configuration: VolumeConfiguration) -> Dict:
-    configuration_excludes = {}
+def _get_volume_configuration_excludes(
+    configuration: VolumeConfiguration,
+) -> IncludeExcludeDictType:
+    configuration_excludes: IncludeExcludeDictType = {}
     if configuration.tags is None:
         configuration_excludes["tags"] = True
     return configuration_excludes

dstack/_internal/core/errors.py

@@ -110,6 +110,10 @@ class PlacementGroupInUseError(ComputeError):
     pass
 
 
+class PlacementGroupNotSupportedError(ComputeError):
+    pass
+
+
 class NotYetTerminated(ComputeError):
     """
     Used by Compute.terminate_instance to signal that instance termination is not complete

dstack/_internal/core/models/common.py

@@ -6,6 +6,13 @@ from pydantic import Field
 from pydantic_duality import DualBaseModel
 from typing_extensions import Annotated
 
+IncludeExcludeFieldType = Union[int, str]
+IncludeExcludeSetType = set[IncludeExcludeFieldType]
+IncludeExcludeDictType = dict[
+    IncludeExcludeFieldType, Union[bool, IncludeExcludeSetType, "IncludeExcludeDictType"]
+]
+IncludeExcludeType = Union[IncludeExcludeSetType, IncludeExcludeDictType]
+
 
 # DualBaseModel creates two classes for the model:
 # one with extra = "forbid" (CoreModel/CoreModel.__request__),

dstack/_internal/core/services/diff.py

@@ -1,14 +1,47 @@
-from typing import Any, Dict
+from typing import Any, Optional, TypedDict
 
 from pydantic import BaseModel
 
+from dstack._internal.core.models.common import IncludeExcludeType
+
+
+class ModelFieldDiff(TypedDict):
+    old: Any
+    new: Any
+
+
+ModelDiff = dict[str, ModelFieldDiff]
+
 
 # TODO: calculate nested diffs
-def diff_models(old: BaseModel, new: BaseModel) -> Dict[str, Any]:
+def diff_models(
+    old: BaseModel, new: BaseModel, ignore: Optional[IncludeExcludeType] = None
+) -> ModelDiff:
+    """
+    Returns a diff of model instances fields.
+
+    NOTE: `ignore` is implemented as `BaseModel.parse_obj(BaseModel.dict(exclude=ignore))`,
+    that is, the "ignored" fields are actually not ignored but reset to the default values
+    before comparison, meaning that 1) any field in `ignore` must have a default value,
+    2) the default value must be equal to itself (e.g. `math.nan` != `math.nan`).
+
+    Args:
+        old: The "old" model instance.
+        new: The "new" model instance.
+        ignore: Optional fields to ignore.
+
+    Returns:
+        A dict of changed fields in the form of
+        `{<field_name>: {"old": old_value, "new": new_value}}`
+    """
     if type(old) is not type(new):
         raise TypeError("Both instances must be of the same Pydantic model class.")
 
-    changes = {}
+    if ignore is not None:
+        old = type(old).parse_obj(old.dict(exclude=ignore))
+        new = type(new).parse_obj(new.dict(exclude=ignore))
+
+    changes: ModelDiff = {}
     for field in old.__fields__:
         old_value = getattr(old, field)
         new_value = getattr(new, field)

dstack/_internal/server/app.py

@@ -2,6 +2,7 @@ import asyncio
 import importlib.resources
 import os
 import time
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Awaitable, Callable, List
@@ -97,6 +98,8 @@ def create_app() -> FastAPI:
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     configure_logging()
+    server_executor = ThreadPoolExecutor(max_workers=settings.SERVER_EXECUTOR_MAX_WORKERS)
+    asyncio.get_running_loop().set_default_executor(server_executor)
     await migrate()
     _print_dstack_logo()
     if not check_required_ssh_version():
@@ -242,6 +245,23 @@ def register_routes(app: FastAPI, ui: bool = True):
         )
         return response
 
+    if settings.SERVER_PROFILING_ENABLED:
+        from pyinstrument import Profiler
+
+        @app.middleware("http")
+        async def profile_request(request: Request, call_next):
+            profiling = request.query_params.get("profile", False)
+            if profiling:
+                profiler = Profiler()
+                profiler.start()
+                respone = await call_next(request)
+                profiler.stop()
+                with open("profiling_results.html", "w+") as f:
+                    f.write(profiler.output_html())
+                return respone
+            else:
+                return await call_next(request)
+
     # this middleware must be defined after the log_request middleware
     @app.middleware("http")
     async def log_http_metrics(request: Request, call_next):

dstack/_internal/server/background/__init__.py

@@ -37,15 +37,31 @@ def get_scheduler() -> AsyncIOScheduler:
 
 
 def start_background_tasks() -> AsyncIOScheduler:
+    # We try to process as many resources as possible without exhausting DB connections.
+    #
+    # Quick tasks can process multiple resources per transaction.
+    # Potentially long tasks process one resource per transaction
+    # to avoid holding locks for all the resources if one is slow to process.
+    # Still, the next batch won't be processed unless all resources are processed,
+    # so larger batches do not increase processing rate linearly.
+    #
+    # The interval, batch_size, and max_instances determine background tasks processing rates.
+    # By default, one server replica can handle:
+    #
+    # * 150 active jobs with 2 minutes processing latency
+    # * 150 active runs with 2 minutes processing latency
+    # * 150 active instances with 2 minutes processing latency
+    #
+    # These latency numbers do not account for provisioning time,
+    # so it may be slower if a backend is slow to provision.
+    #
+    # Users can set SERVER_BACKGROUND_PROCESSING_FACTOR to process more resources per replica.
+    # They also need to increase max db connections on the client side and db side.
+    #
     # In-memory locking via locksets does not guarantee
     # that the first waiting for the lock will acquire it.
     # The jitter is needed to give all tasks a chance to acquire locks.
 
-    # The batch_size and interval determine background tasks processing rates.
-    # Currently one server replica can handle:
-    # * 150 active jobs with up to 2 minutes processing latency
-    # * 150 active runs with up to 2 minutes processing latency
-    # * 150 active instances with up to 2 minutes processing latency
     _scheduler.add_job(collect_metrics, IntervalTrigger(seconds=10), max_instances=1)
     _scheduler.add_job(delete_metrics, IntervalTrigger(minutes=5), max_instances=1)
     if settings.ENABLE_PROMETHEUS_METRICS:
@@ -53,38 +69,6 @@ def start_background_tasks() -> AsyncIOScheduler:
             collect_prometheus_metrics, IntervalTrigger(seconds=10), max_instances=1
         )
         _scheduler.add_job(delete_prometheus_metrics, IntervalTrigger(minutes=5), max_instances=1)
-    # process_submitted_jobs and process_instances max processing rate is 75 jobs(instances) per minute.
-    _scheduler.add_job(
-        process_submitted_jobs,
-        IntervalTrigger(seconds=4, jitter=2),
-        kwargs={"batch_size": 5},
-        max_instances=2,
-    )
-    _scheduler.add_job(
-        process_running_jobs,
-        IntervalTrigger(seconds=4, jitter=2),
-        kwargs={"batch_size": 5},
-        max_instances=2,
-    )
-    _scheduler.add_job(
-        process_terminating_jobs,
-        IntervalTrigger(seconds=4, jitter=2),
-        kwargs={"batch_size": 5},
-        max_instances=2,
-    )
-    _scheduler.add_job(
-        process_runs,
-        IntervalTrigger(seconds=2, jitter=1),
-        kwargs={"batch_size": 5},
-        max_instances=2,
-    )
-    _scheduler.add_job(
-        process_instances,
-        IntervalTrigger(seconds=4, jitter=2),
-        kwargs={"batch_size": 5},
-        max_instances=2,
-    )
-    _scheduler.add_job(process_fleets, IntervalTrigger(seconds=10, jitter=2))
     _scheduler.add_job(process_gateways_connections, IntervalTrigger(seconds=15))
     _scheduler.add_job(
         process_submitted_gateways, IntervalTrigger(seconds=10, jitter=2), max_instances=5
@@ -93,5 +77,45 @@ def start_background_tasks() -> AsyncIOScheduler:
         process_submitted_volumes, IntervalTrigger(seconds=10, jitter=2), max_instances=5
     )
     _scheduler.add_job(process_placement_groups, IntervalTrigger(seconds=30, jitter=5))
+    for replica in range(settings.SERVER_BACKGROUND_PROCESSING_FACTOR):
+        # Add multiple copies of tasks if requested.
+        # max_instances=1 for additional copies to avoid running too many tasks.
+        # Move other tasks here when they need per-replica scaling.
+        _scheduler.add_job(
+            process_submitted_jobs,
+            IntervalTrigger(seconds=4, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=4 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_running_jobs,
+            IntervalTrigger(seconds=4, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_terminating_jobs,
+            IntervalTrigger(seconds=4, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_runs,
+            IntervalTrigger(seconds=2, jitter=1),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_instances,
+            IntervalTrigger(seconds=4, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
+        _scheduler.add_job(
+            process_fleets,
+            IntervalTrigger(seconds=10, jitter=2),
+            kwargs={"batch_size": 5},
+            max_instances=2 if replica == 0 else 1,
+        )
     _scheduler.start()
     return _scheduler

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

@@ -1,9 +1,12 @@
+import asyncio
+from datetime import timedelta
+
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.models.fleets import FleetStatus
-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
 from dstack._internal.server.services.fleets import (
     is_fleet_empty,
@@ -17,8 +20,18 @@ from dstack._internal.utils.logging import get_logger
 logger = get_logger(__name__)
 
 
-async def process_fleets():
-    lock, lockset = get_locker().get_lockset(FleetModel.__tablename__)
+MIN_PROCESSING_INTERVAL = timedelta(seconds=30)
+
+
+async def process_fleets(batch_size: int = 1):
+    tasks = []
+    for _ in range(batch_size):
+        tasks.append(_process_next_fleet())
+    await asyncio.gather(*tasks)
+
+
+async def _process_next_fleet():
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(FleetModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
             res = await session.execute(
@@ -26,6 +39,8 @@ async def process_fleets():
                 .where(
                     FleetModel.deleted == False,
                     FleetModel.id.not_in(lockset),
+                    FleetModel.last_processed_at
+                    < get_current_datetime().replace(tzinfo=None) - MIN_PROCESSING_INTERVAL,
                 )
                 .order_by(FleetModel.last_processed_at.asc())
                 .limit(1)
@@ -43,6 +58,7 @@ async def process_fleets():
 
 
 async def _process_fleet(session: AsyncSession, fleet_model: FleetModel):
+    logger.debug("Processing fleet %s", fleet_model.name)
     # Refetch to load related attributes.
     # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
     res = await session.execute(

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

@@ -28,7 +28,7 @@ async def process_gateways_connections():
 
 
 async def process_submitted_gateways():
-    lock, lockset = get_locker().get_lockset(GatewayModel.__tablename__)
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(GatewayModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
             res = await session.execute(

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

@@ -34,10 +34,11 @@ 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,
@@ -79,6 +80,7 @@ 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.
@@ -93,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: