Make JobSubmittedPipeline to wait for master election

Author: r4victor

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

@@ -319,6 +319,10 @@ async def process(self, item: JobSubmittedPipelineItem):
         if context.job_model.instance_assigned:
             logger.debug("%s: provisioning has started", fmt(context.job_model))
             provisioning = await _process_provisioning(item=item, context=context)
+            _hint_pipelines_fetch(
+                pipeline_hinter=self._pipeline_hinter,
+                result=provisioning,
+            )
             await _apply_provisioning_result(
                 item=item,
                 provisioning=provisioning,
@@ -327,6 +331,10 @@ async def process(self, item: JobSubmittedPipelineItem):
 
         logger.debug("%s: assignment has started", fmt(context.job_model))
         assignment = await _process_assignment(context=context)
+        _hint_pipelines_fetch(
+            pipeline_hinter=self._pipeline_hinter,
+            result=assignment,
+        )
         await _apply_assignment_result(
             item=item,
             context=context,
@@ -365,6 +373,7 @@ class _DeferSubmittedJobResult:
     """The job is not ready yet, so apply should just mark it processed and unlock it."""
 
     log_message: str
+    hint_fleet_pipeline: bool = False
 
 
 @dataclass
@@ -1179,6 +1188,17 @@ async def _process_new_capacity_provisioning(
             job=context.job,
         )
     )
+    if (
+        is_master_job(context.job)
+        and fleet_model is not None
+        and _get_cluster_fleet_spec(fleet_model) is not None
+        and any(not instance.deleted for instance in fleet_model.instances)
+        and master_provisioning_data is None
+    ):
+        return _DeferSubmittedJobResult(
+            log_message="waiting for fleet master instance election",
+            hint_fleet_pipeline=True,
+        )
     provision_new_capacity_result = await _provision_new_capacity(
         project=context.project,
         fleet_model=fleet_model,
@@ -1834,6 +1854,17 @@ def _get_fleet_master_provisioning_data(
     )
 
 
+def _hint_pipelines_fetch(
+    pipeline_hinter: PipelineHinterProtocol,
+    result: Union[_AssignmentResult, _ProvisioningResult],
+) -> None:
+    if not isinstance(result, _DeferSubmittedJobResult):
+        return
+
+    if result.hint_fleet_pipeline:
+        pipeline_hinter.hint_fetch(FleetModel.__name__)
+
+
 def _select_jobs_to_provision(job: Job, replica_jobs: list[Job], job_model: JobModel) -> list[Job]:
     jobs_to_provision = [job]
     if is_multinode_job(job) and is_master_job(job) and job_model.waiting_master_job is not None:

src/dstack/_internal/server/services/fleets.py

@@ -961,12 +961,6 @@ def get_fleet_master_instance_provisioning_data(
                     instance_model.job_provisioning_data
                 )
 
-    # TODO: Drop the legacy instance-list fallback after scheduled tasks stop
-    # inferring cluster masters from loaded fleet instances.
-    for instance_model in fleet_model.instances:
-        if not instance_model.deleted and instance_model.job_provisioning_data is not None:
-            return JobProvisioningData.__response__.parse_raw(instance_model.job_provisioning_data)
-
     return None
 
 

src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py

@@ -31,6 +31,7 @@
 )
 from dstack._internal.server.models import (
     ComputeGroupModel,
+    FleetModel,
     InstanceModel,
     JobModel,
     PlacementGroupModel,
@@ -431,14 +432,15 @@ async def test_provisioning_master_job_respects_cluster_placement_in_non_empty_f
         fleet_spec.configuration.placement = InstanceGroupPlacement.CLUSTER
         fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=0, max=None)
         fleet = await create_fleet(session=session, project=project, spec=fleet_spec)
-        await create_instance(
+        instance = await create_instance(
             session=session,
             project=project,
             fleet=fleet,
             status=InstanceStatus.BUSY,
             backend=BackendType.AWS,
             job_provisioning_data=get_job_provisioning_data(region="eu-west-1"),
         )
+        fleet.current_master_instance_id = instance.id
         configuration = TaskConfiguration(image="debian", nodes=2)
         run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration)
         run = await create_run(
@@ -487,6 +489,57 @@ async def test_provisioning_master_job_respects_cluster_placement_in_non_empty_f
         selected_offer = backend_mock.compute.return_value.run_job.call_args[0][2]
         assert selected_offer.region == "eu-west-1"
 
+    async def test_defers_new_capacity_provisioning_until_fleet_master_is_elected(
+        self, test_db, session: AsyncSession, worker: JobSubmittedWorker
+    ):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        repo = await create_repo(session=session, project_id=project.id)
+        fleet_spec = get_fleet_spec()
+        fleet_spec.configuration.placement = InstanceGroupPlacement.CLUSTER
+        fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=0, max=None)
+        fleet = await create_fleet(session=session, project=project, spec=fleet_spec)
+        await create_instance(
+            session=session,
+            project=project,
+            fleet=fleet,
+            status=InstanceStatus.BUSY,
+            backend=BackendType.AWS,
+            job_provisioning_data=get_job_provisioning_data(region="eu-west-1"),
+        )
+        configuration = TaskConfiguration(image="debian", nodes=2)
+        run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration)
+        run = await create_run(
+            session=session,
+            run_name="run",
+            project=project,
+            repo=repo,
+            user=user,
+            run_spec=run_spec,
+            fleet=fleet,
+        )
+        job = await create_job(
+            session=session,
+            run=run,
+            instance_assigned=True,
+            waiting_master_job=False,
+        )
+        previous_last_processed_at = job.last_processed_at
+
+        with patch("dstack._internal.server.services.backends.get_project_backends") as m:
+            await _process_job(session=session, worker=worker, job_model=job)
+            m.assert_not_called()
+
+        await session.refresh(job)
+        assert job.status == JobStatus.SUBMITTED
+        assert job.instance_assigned
+        assert job.instance is None
+        assert job.last_processed_at > previous_last_processed_at
+        assert job.lock_owner is None
+        assert job.lock_token is None
+        assert job.lock_expires_at is None
+        worker._pipeline_hinter.hint_fetch.assert_called_once_with(FleetModel.__name__)
+
     async def test_provisioning_non_master_job_ignores_cluster_master_fleet_lock(
         self, test_db, session: AsyncSession, worker: JobSubmittedWorker
     ):

src/tests/_internal/server/services/runs/__init__.py

@@ -0,0 +1,71 @@
+from unittest.mock import AsyncMock
+
+import pytest
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.models.configurations import TaskConfiguration
+from dstack._internal.core.models.fleets import FleetNodesSpec, InstanceGroupPlacement
+from dstack._internal.server.services.jobs import get_jobs_from_run_spec
+from dstack._internal.server.services.runs.plan import _get_backend_offers_in_fleet
+from dstack._internal.server.testing.common import (
+    create_fleet,
+    create_instance,
+    create_project,
+    create_repo,
+    create_user,
+    get_fleet_spec,
+    get_instance_offer_with_availability,
+    get_job_provisioning_data,
+    get_run_spec,
+)
+
+pytestmark = pytest.mark.usefixtures("image_config_mock")
+
+
+class TestGetBackendOffersInFleet:
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_keeps_unconstrained_offers_for_non_empty_cluster_fleet_without_elected_master(
+        self, test_db, session: AsyncSession, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        user = await create_user(session=session)
+        project = await create_project(session=session, owner=user)
+        repo = await create_repo(session=session, project_id=project.id)
+        fleet_spec = get_fleet_spec()
+        fleet_spec.configuration.placement = InstanceGroupPlacement.CLUSTER
+        fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=1, max=2)
+        fleet = await create_fleet(session=session, project=project, spec=fleet_spec)
+        await create_instance(
+            session=session,
+            project=project,
+            fleet=fleet,
+            job_provisioning_data=get_job_provisioning_data(region="eu-west-1"),
+        )
+        run_spec = get_run_spec(
+            repo_id=repo.name,
+            configuration=TaskConfiguration(image="debian", nodes=2),
+        )
+        jobs = await get_jobs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0)
+        get_offers_by_requirements_mock = AsyncMock()
+        monkeypatch.setattr(
+            "dstack._internal.server.services.runs.plan.get_offers_by_requirements",
+            get_offers_by_requirements_mock,
+        )
+        offer = get_instance_offer_with_availability()
+        backend = AsyncMock()
+        get_offers_by_requirements_mock.return_value = [(backend, offer)]
+
+        offers = await _get_backend_offers_in_fleet(
+            project=project,
+            fleet_model=fleet,
+            run_spec=run_spec,
+            job=jobs[0],
+            volumes=None,
+        )
+
+        assert offers == [(backend, offer)]
+        get_offers_by_requirements_mock.assert_awaited_once()
+        assert (
+            get_offers_by_requirements_mock.await_args.kwargs["master_job_provisioning_data"]
+            is None
+        )

src/tests/_internal/server/services/runs/test_plan.py

@@ -9,20 +9,26 @@
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.fleets import (
     FleetConfiguration,
+    FleetNodesSpec,
     FleetSpec,
+    InstanceGroupPlacement,
     SSHHostParams,
     SSHParams,
 )
 from dstack._internal.core.models.instances import RemoteConnectionInfo
 from dstack._internal.server.models import FleetModel, ProjectModel
 from dstack._internal.server.services.backends import get_project_backends
-from dstack._internal.server.services.fleets import get_plan
+from dstack._internal.server.services.fleets import (
+    get_fleet_master_instance_provisioning_data,
+    get_plan,
+)
 from dstack._internal.server.testing.common import (
     create_fleet,
     create_instance,
     create_project,
     create_user,
     get_fleet_spec,
+    get_job_provisioning_data,
     get_ssh_key,
 )
 
@@ -171,3 +177,29 @@ async def test_error_fleet_spec_without_name(self, session: AsyncSession):
             await get_plan(
                 session=session, project=project, user=user, spec=fleet_spec_without_name
             )
+
+
+class TestGetFleetMasterInstanceProvisioningData:
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_returns_none_without_current_master_instance(
+        self, test_db, session: AsyncSession
+    ) -> None:
+        project = await create_project(session=session)
+        fleet_spec = get_fleet_spec()
+        fleet_spec.configuration.placement = InstanceGroupPlacement.CLUSTER
+        fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=1, max=2)
+        fleet = await create_fleet(session=session, project=project, spec=fleet_spec)
+        await create_instance(
+            session=session,
+            project=project,
+            fleet=fleet,
+            job_provisioning_data=get_job_provisioning_data(region="eu-west-1"),
+        )
+
+        master_provisioning_data = get_fleet_master_instance_provisioning_data(
+            fleet_model=fleet,
+            fleet_spec=fleet_spec,
+        )
+
+        assert master_provisioning_data is None