Select only ids in locking select

Author: r4victor

src/dstack/_internal/server/background/__init__.py

@@ -82,7 +82,6 @@ def start_background_tasks() -> AsyncIOScheduler:
     _scheduler.add_job(
         process_fleets,
         IntervalTrigger(seconds=10, jitter=2),
-        kwargs={"batch_size": 10},
         max_instances=1,
     )
     for replica in range(settings.SERVER_BACKGROUND_PROCESSING_FACTOR):

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

@@ -25,10 +25,11 @@
 logger = get_logger(__name__)
 
 
+BATCH_SIZE = 10
 MIN_PROCESSING_INTERVAL = timedelta(seconds=30)
 
 
-async def process_fleets(batch_size: int = 1):
+async def process_fleets():
     lock, lockset = get_locker(get_db().dialect_name).get_lockset(FleetModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
@@ -42,7 +43,7 @@ async def process_fleets(batch_size: int = 1):
                 )
                 .options(load_only(FleetModel.id))
                 .order_by(FleetModel.last_processed_at.asc())
-                .limit(batch_size)
+                .limit(BATCH_SIZE)
                 .with_for_update(skip_locked=True, key_share=True)
             )
             fleet_models = list(res.scalars().all())

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

@@ -7,7 +7,7 @@
 
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import joinedload, load_only
 
 from dstack._internal import settings
 from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT, DSTACK_SHIM_HTTP_PORT
@@ -110,6 +110,7 @@ async def _process_next_running_job():
                     JobModel.last_processed_at
                     < common_utils.get_current_datetime() - MIN_PROCESSING_INTERVAL,
                 )
+                .options(load_only(JobModel.id))
                 .order_by(JobModel.last_processed_at.asc())
                 .limit(1)
                 .with_for_update(

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

@@ -4,7 +4,7 @@
 
 from sqlalchemy import and_, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, selectinload
+from sqlalchemy.orm import joinedload, load_only, selectinload
 
 import dstack._internal.server.services.services.autoscalers as autoscalers
 from dstack._internal.core.errors import ServerError
@@ -102,6 +102,7 @@ async def _process_next_run():
                     ),
                 )
                 .options(joinedload(RunModel.jobs).load_only(JobModel.id))
+                .options(load_only(RunModel.id))
                 .order_by(RunModel.last_processed_at.asc())
                 .limit(1)
                 .with_for_update(skip_locked=True, key_share=True, of=RunModel)
@@ -134,7 +135,6 @@ async def _process_next_run():
 
 
 async def _process_run(session: AsyncSession, run_model: RunModel):
-    logger.debug("%s: processing run", fmt(run_model))
     # Refetch to load related attributes.
     res = await session.execute(
         select(RunModel)
@@ -150,6 +150,7 @@ async def _process_run(session: AsyncSession, run_model: RunModel):
         .execution_options(populate_existing=True)
     )
     run_model = res.unique().scalar_one()
+    logger.debug("%s: processing run", fmt(run_model))
     try:
         if run_model.status == RunStatus.PENDING:
             await _process_pending_run(session, run_model)

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

@@ -5,7 +5,7 @@
 
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, lazyload, selectinload
+from sqlalchemy.orm import joinedload, lazyload, load_only, selectinload
 
 from dstack._internal.core.backends.base.backend import Backend
 from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport
@@ -120,6 +120,7 @@ async def _process_next_submitted_job():
                     JobModel.status == JobStatus.SUBMITTED,
                     JobModel.id.not_in(lockset),
                 )
+                .options(load_only(JobModel.id))
                 # Jobs are process in FIFO sorted by priority globally,
                 # thus runs from different projects can "overtake" each other by using higher priorities.
                 # That's not a big problem as long as projects do not compete for the same compute resources.
@@ -152,7 +153,6 @@ async def _process_next_submitted_job():
 
 
 async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
-    logger.debug("%s: provisioning has started", fmt(job_model))
     # Refetch to load related attributes.
     res = await session.execute(
         select(JobModel).where(JobModel.id == job_model.id).options(joinedload(JobModel.instance))
@@ -166,6 +166,8 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         .options(joinedload(RunModel.fleet).joinedload(FleetModel.instances))
     )
     run_model = res.unique().scalar_one()
+    logger.debug("%s: provisioning has started", fmt(job_model))
+
     project = run_model.project
     run = run_model_to_run(run_model)
     run_spec = run.run_spec