Add pipelines optimizations (#3719)

* Skip sibling jobs loading for running jobs

* Use status-specific processing intervals for instances

* Optimize run_model loading for submitted jobs

* Introduce *_STATUSES_WITH_MIN_PROCESSING_INTERVAL

* Use state-specific processing intervals for terminating jobs

* Pass pipeline hinter to pipelines

Author: r4victor

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

@@ -27,16 +27,16 @@ def __init__(self) -> None:
         self._pipelines: list[Pipeline] = []
         self._hinter = PipelineHinter()
         for builtin_pipeline in [
-            ComputeGroupPipeline(),
-            FleetPipeline(),
-            GatewayPipeline(),
-            JobSubmittedPipeline(),
-            JobRunningPipeline(),
-            JobTerminatingPipeline(),
-            InstancePipeline(),
-            PlacementGroupPipeline(),
-            RunPipeline(),
-            VolumePipeline(),
+            ComputeGroupPipeline(pipeline_hinter=self._hinter),
+            FleetPipeline(pipeline_hinter=self._hinter),
+            GatewayPipeline(pipeline_hinter=self._hinter),
+            JobSubmittedPipeline(pipeline_hinter=self._hinter),
+            JobRunningPipeline(pipeline_hinter=self._hinter),
+            JobTerminatingPipeline(pipeline_hinter=self._hinter),
+            InstancePipeline(pipeline_hinter=self._hinter),
+            PlacementGroupPipeline(pipeline_hinter=self._hinter),
+            RunPipeline(pipeline_hinter=self._hinter),
+            VolumePipeline(pipeline_hinter=self._hinter),
         ]:
             self.register_pipeline(builtin_pipeline)
 

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

@@ -24,6 +24,7 @@
 from sqlalchemy.orm import Mapped
 
 from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.services.pipelines import PipelineHinterProtocol
 from dstack._internal.utils.common import get_current_datetime
 from dstack._internal.utils.logging import get_logger
 
@@ -338,9 +339,11 @@ def __init__(
         self,
         queue: asyncio.Queue[ItemT],
         heartbeater: Heartbeater[ItemT],
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         self._queue = queue
         self._heartbeater = heartbeater
+        self._pipeline_hinter = pipeline_hinter
         self._running = False
 
     async def start(self):

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

@@ -32,6 +32,7 @@
 from dstack._internal.server.services.compute_groups import compute_group_model_to_compute_group
 from dstack._internal.server.services.instances import emit_instance_status_change_event
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.pipelines import PipelineHinterProtocol
 from dstack._internal.server.utils import sentry_utils
 from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
@@ -51,6 +52,8 @@ def __init__(
         min_processing_interval: timedelta = timedelta(seconds=15),
         lock_timeout: timedelta = timedelta(seconds=30),
         heartbeat_trigger: timedelta = timedelta(seconds=15),
+        *,
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         super().__init__(
             workers_num=workers_num,
@@ -73,7 +76,11 @@ def __init__(
             heartbeater=self._heartbeater,
         )
         self.__workers = [
-            ComputeGroupWorker(queue=self._queue, heartbeater=self._heartbeater)
+            ComputeGroupWorker(
+                queue=self._queue,
+                heartbeater=self._heartbeater,
+                pipeline_hinter=pipeline_hinter,
+            )
             for _ in range(self._workers_num)
         ]
 
@@ -173,10 +180,12 @@ def __init__(
         self,
         queue: asyncio.Queue[PipelineItem],
         heartbeater: Heartbeater[PipelineItem],
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         super().__init__(
             queue=queue,
             heartbeater=heartbeater,
+            pipeline_hinter=pipeline_hinter,
         )
 
     @sentry_utils.instrument_named_task("pipeline_tasks.ComputeGroupWorker.process")

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

@@ -49,6 +49,7 @@
     is_fleet_in_use,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.pipelines import PipelineHinterProtocol
 from dstack._internal.server.utils import sentry_utils
 from dstack._internal.utils.common import get_current_datetime
 from dstack._internal.utils.logging import get_logger
@@ -65,6 +66,8 @@ def __init__(
         min_processing_interval: timedelta = timedelta(seconds=30),
         lock_timeout: timedelta = timedelta(seconds=20),
         heartbeat_trigger: timedelta = timedelta(seconds=10),
+        *,
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         super().__init__(
             workers_num=workers_num,
@@ -87,7 +90,11 @@ def __init__(
             heartbeater=self._heartbeater,
         )
         self.__workers = [
-            FleetWorker(queue=self._queue, heartbeater=self._heartbeater)
+            FleetWorker(
+                queue=self._queue,
+                heartbeater=self._heartbeater,
+                pipeline_hinter=pipeline_hinter,
+            )
             for _ in range(self._workers_num)
         ]
 
@@ -188,10 +195,12 @@ def __init__(
         self,
         queue: asyncio.Queue[PipelineItem],
         heartbeater: Heartbeater[PipelineItem],
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         super().__init__(
             queue=queue,
             heartbeater=heartbeater,
+            pipeline_hinter=pipeline_hinter,
         )
 
     @sentry_utils.instrument_named_task("pipeline_tasks.FleetWorker.process")

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

@@ -38,6 +38,7 @@
 from dstack._internal.server.services.gateways.pool import gateway_connections_pool
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
+from dstack._internal.server.services.pipelines import PipelineHinterProtocol
 from dstack._internal.server.utils import sentry_utils
 from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
@@ -60,6 +61,8 @@ def __init__(
         min_processing_interval: timedelta = timedelta(seconds=15),
         lock_timeout: timedelta = timedelta(seconds=30),
         heartbeat_trigger: timedelta = timedelta(seconds=15),
+        *,
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         super().__init__(
             workers_num=workers_num,
@@ -82,7 +85,11 @@ def __init__(
             heartbeater=self._heartbeater,
         )
         self.__workers = [
-            GatewayWorker(queue=self._queue, heartbeater=self._heartbeater)
+            GatewayWorker(
+                queue=self._queue,
+                heartbeater=self._heartbeater,
+                pipeline_hinter=pipeline_hinter,
+            )
             for _ in range(self._workers_num)
         ]
 
@@ -192,10 +199,12 @@ def __init__(
         self,
         queue: asyncio.Queue[GatewayPipelineItem],
         heartbeater: Heartbeater[GatewayPipelineItem],
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         super().__init__(
             queue=queue,
             heartbeater=heartbeater,
+            pipeline_hinter=pipeline_hinter,
         )
 
     @sentry_utils.instrument_named_task("pipeline_tasks.GatewayWorker.process")

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

@@ -52,6 +52,7 @@
     is_ssh_instance,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.pipelines import PipelineHinterProtocol
 from dstack._internal.server.services.placement import (
     schedule_fleet_placement_groups_deletion,
 )
@@ -62,6 +63,13 @@
 logger = get_logger(__name__)
 
 
+INSTANCE_STATUSES_WITH_MIN_PROCESSING_INTERVAL = [
+    InstanceStatus.PENDING,
+    InstanceStatus.PROVISIONING,
+    InstanceStatus.TERMINATING,
+]
+
+
 @dataclass
 class InstancePipelineItem(PipelineItem):
     status: InstanceStatus
@@ -73,9 +81,11 @@ def __init__(
         workers_num: int = 20,
         queue_lower_limit_factor: float = 0.5,
         queue_upper_limit_factor: float = 2.0,
-        min_processing_interval: timedelta = timedelta(seconds=15),
+        min_processing_interval: timedelta = timedelta(seconds=7),
         lock_timeout: timedelta = timedelta(seconds=30),
         heartbeat_trigger: timedelta = timedelta(seconds=15),
+        *,
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         super().__init__(
             workers_num=workers_num,
@@ -98,7 +108,11 @@ def __init__(
             heartbeater=self._heartbeater,
         )
         self.__workers = [
-            InstanceWorker(queue=self._queue, heartbeater=self._heartbeater)
+            InstanceWorker(
+                queue=self._queue,
+                heartbeater=self._heartbeater,
+                pipeline_hinter=pipeline_hinter,
+            )
             for _ in range(self._workers_num)
         ]
 
@@ -167,7 +181,24 @@ async def fetch(self, limit: int) -> list[InstancePipelineItem]:
                         ),
                         InstanceModel.deleted == False,
                         or_(
-                            InstanceModel.last_processed_at <= now - self._min_processing_interval,
+                            # Process fast-moving instances (pending, provisioning, terminating)
+                            # at base interval for low-latency state transitions.
+                            # Steady-state instances (idle, busy) use a longer interval
+                            # since they only need periodic health checks.
+                            and_(
+                                InstanceModel.status.in_(
+                                    INSTANCE_STATUSES_WITH_MIN_PROCESSING_INTERVAL
+                                ),
+                                InstanceModel.last_processed_at
+                                <= now - self._min_processing_interval,
+                            ),
+                            and_(
+                                InstanceModel.status.not_in(
+                                    INSTANCE_STATUSES_WITH_MIN_PROCESSING_INTERVAL
+                                ),
+                                InstanceModel.last_processed_at
+                                <= now - self._min_processing_interval * 2,
+                            ),
                             InstanceModel.last_processed_at == InstanceModel.created_at,
                         ),
                         or_(
@@ -228,10 +259,12 @@ def __init__(
         self,
         queue: asyncio.Queue[InstancePipelineItem],
         heartbeater: Heartbeater[InstancePipelineItem],
+        pipeline_hinter: PipelineHinterProtocol,
     ) -> None:
         super().__init__(
             queue=queue,
             heartbeater=heartbeater,
+            pipeline_hinter=pipeline_hinter,
         )
 
     @sentry_utils.instrument_named_task("pipeline_tasks.InstanceWorker.process")