Fix tasks selects

Author: r4victor

src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py

@@ -462,6 +462,9 @@ async def _refetch_locked_job_model(
 
 
 async def _fetch_run_model(session: AsyncSession, run_id: uuid.UUID) -> RunModel:
+    # FIXME: Selecting all run's jobs on every processing iteration is highly inefficient:
+    # it's quadratic w.r.t. the number jobs within a run.
+    # Avoid selecting other jobs as much as possible.
     latest_submissions_sq = (
         select(
             JobModel.run_id.label("run_id"),

src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py

@@ -234,6 +234,11 @@ async def fetch(self, limit: int) -> list[JobSubmittedPipelineItem]:
                     .where(
                         JobModel.status == JobStatus.SUBMITTED,
                         JobModel.waiting_master_job.is_not(True),
+                        or_(
+                            # Non-master jobs must wait for the run to have the fleet assigned.
+                            JobModel.job_num == 0,
+                            RunModel.fleet_id.is_not(None),
+                        ),
                         or_(
                             JobModel.last_processed_at <= now - self._min_processing_interval,
                             JobModel.last_processed_at == JobModel.submitted_at,

src/dstack/_internal/server/background/scheduled_tasks/runs.py

@@ -114,8 +114,7 @@ async def _process_next_run():
                         ),
                         # Scaled-to-zero runs:
                         # Such runs cannot be scheduled, thus we check next_triggered_at.
-                        # If we allow scheduled services with downscaling to zero
-                        # This check won't pass.
+                        # If we allow scheduled services with downscaling to zero, this check won't pass.
                         and_(
                             RunModel.status == RunStatus.PENDING,
                             RunModel.resubmission_attempt == 0,
@@ -144,19 +143,10 @@ async def _process_next_run():
                 .where(
                     JobModel.run_id == run_model.id,
                     JobModel.id.not_in(job_lockset),
-                    or_(
-                        JobModel.lock_expires_at.is_(None),
-                        JobModel.lock_expires_at < now,
-                    ),
-                )
-                .options(
-                    load_only(JobModel.id),
-                    with_loader_criteria(
-                        JobModel,
-                        JobModel.status.not_in(JOB_STATUSES_EXCLUDED_FOR_LOCKING),
-                        include_aliases=True,
-                    ),
+                    JobModel.status.not_in(JOB_STATUSES_EXCLUDED_FOR_LOCKING),
+                    JobModel.lock_expires_at.is_(None),
                 )
+                .options(load_only(JobModel.id))
                 .order_by(JobModel.id)  # take locks in order
                 .with_for_update(skip_locked=True, key_share=True)
             )