Support nodes in-place update for cloud fleets

Author: r4victor

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

@@ -63,7 +63,7 @@ def __init__(
         workers_num: int = 10,
         queue_lower_limit_factor: float = 0.5,
         queue_upper_limit_factor: float = 2.0,
-        min_processing_interval: timedelta = timedelta(seconds=30),
+        min_processing_interval: timedelta = timedelta(seconds=15),
         lock_timeout: timedelta = timedelta(seconds=20),
         heartbeat_trigger: timedelta = timedelta(seconds=10),
         *,

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

@@ -1225,20 +1225,25 @@ def _check_can_update_fleet_spec(current: FleetSpec, new: FleetSpec, diff: Model
         _check_can_update_fleet_configuration(current.configuration, new.configuration)
 
 
-@_check_can_update("ssh_config")
-def _check_can_update_fleet_configuration(
-    current: FleetConfiguration, new: FleetConfiguration, diff: ModelDiff
-):
+def _check_can_update_fleet_configuration(current: FleetConfiguration, new: FleetConfiguration):
+    diff = diff_models(current, new)
+    current_ssh_config = current.ssh_config
+    new_ssh_config = new.ssh_config
+    if current_ssh_config is None:
+        if new_ssh_config is not None:
+            raise ServerClientError("Fleet type changed from Cloud to SSH, cannot update")
+        # TODO: Support best-effort `nodes.target` apply semantics:
+        # create missing instances and terminate extra idle instances.
+        # Current in-place update only persists `target`; FleetPipeline reconciles `min`/`max`.
+        _check_can_update_inner(current, new, ("nodes",))
+        return
+
+    if new_ssh_config is None:
+        raise ServerClientError("Fleet type changed from SSH to Cloud, cannot update")
+
+    _check_can_update_inner(current, new, ("ssh_config",))
     if "ssh_config" in diff:
-        current_ssh_config = current.ssh_config
-        new_ssh_config = new.ssh_config
-        if current_ssh_config is None:
-            if new_ssh_config is not None:
-                raise ServerClientError("Fleet type changed from Cloud to SSH, cannot update")
-        elif new_ssh_config is None:
-            raise ServerClientError("Fleet type changed from SSH to Cloud, cannot update")
-        else:
-            _check_can_update_ssh_config(current_ssh_config, new_ssh_config)
+        _check_can_update_ssh_config(current_ssh_config, new_ssh_config)
 
 
 @_check_can_update("hosts")

src/tests/_internal/server/routers/test_fleets.py

@@ -1436,6 +1436,108 @@ async def test_updates_ssh_fleet(self, test_db, session: AsyncSession, client: A
         assert instance.status == InstanceStatus.PENDING
         assert instance.remote_connection_info is not None
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_updates_cloud_fleet_nodes_in_place_when_fleet_in_use(
+        self, test_db, session: AsyncSession, client: AsyncClient
+    ):
+        user = await create_user(session, global_role=GlobalRole.USER)
+        project = await create_project(session)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        current_spec = get_fleet_spec(
+            conf=get_fleet_configuration(nodes=FleetNodesSpec(min=0, target=0, max=2))
+        )
+        fleet = await create_fleet(session=session, project=project, spec=current_spec)
+        repo = await create_repo(session=session, project_id=project.id)
+        run = await create_run(session=session, project=project, repo=repo, user=user, fleet=fleet)
+        job = await create_job(session=session, run=run, fleet=fleet)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            fleet=fleet,
+            job=job,
+            status=InstanceStatus.BUSY,
+            instance_num=0,
+        )
+        spec = current_spec.copy(deep=True)
+        spec.configuration.nodes = FleetNodesSpec(min=1, target=1, max=3)
+
+        response = await client.post(
+            f"/api/project/{project.name}/fleets/apply",
+            headers=get_auth_headers(user.token),
+            json={
+                "plan": {
+                    "spec": spec.dict(),
+                    "current_resource": _fleet_model_to_json_dict(fleet),
+                },
+                "force": False,
+            },
+        )
+
+        response_json = response.json()
+        assert response.status_code == 200, response_json
+        assert response_json["id"] == str(fleet.id)
+        assert response_json["spec"]["configuration"]["nodes"] == {"min": 1, "max": 3}
+
+        await session.refresh(fleet)
+        await session.refresh(instance)
+        assert json.loads(fleet.spec)["configuration"]["nodes"] == {"min": 1, "max": 3}
+        assert instance.status == InstanceStatus.BUSY
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_updates_cloud_fleet_nodes_target_without_changing_instance_count(
+        self, test_db, session: AsyncSession, client: AsyncClient
+    ):
+        user = await create_user(session, global_role=GlobalRole.USER)
+        project = await create_project(session)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        current_spec = get_fleet_spec(
+            conf=get_fleet_configuration(nodes=FleetNodesSpec(min=0, target=0, max=1))
+        )
+        fleet = await create_fleet(session=session, project=project, spec=current_spec)
+        spec = current_spec.copy(deep=True)
+        spec.configuration.nodes = FleetNodesSpec(min=0, target=1, max=1)
+
+        response = await client.post(
+            f"/api/project/{project.name}/fleets/apply",
+            headers=get_auth_headers(user.token),
+            json={
+                "plan": {
+                    "spec": spec.dict(),
+                    "current_resource": _fleet_model_to_json_dict(fleet),
+                },
+                "force": False,
+            },
+        )
+
+        response_json = response.json()
+        assert response.status_code == 200, response_json
+        assert response_json["id"] == str(fleet.id)
+        assert response_json["spec"]["configuration"]["nodes"] == {
+            "min": 0,
+            "target": 1,
+            "max": 1,
+        }
+
+        await session.refresh(fleet)
+        assert json.loads(fleet.spec)["configuration"]["nodes"] == {
+            "min": 0,
+            "target": 1,
+            "max": 1,
+        }
+        res = await session.execute(
+            select(InstanceModel).where(
+                InstanceModel.fleet_id == fleet.id,
+                InstanceModel.deleted == False,
+            )
+        )
+        assert list(res.scalars().all()) == []
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc))
@@ -2118,6 +2220,62 @@ async def test_returns_update_plan_for_existing_fleet(
             "action": "update",
         }
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_returns_update_plan_for_existing_cloud_fleet_nodes_update(
+        self, test_db, session: AsyncSession, client: AsyncClient
+    ):
+        user = await create_user(session=session, global_role=GlobalRole.USER)
+        project = await create_project(session=session, owner=user)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        current_spec = get_fleet_spec(
+            conf=get_fleet_configuration(nodes=FleetNodesSpec(min=0, target=0, max=1))
+        )
+        spec = current_spec.copy(deep=True)
+        spec.configuration.nodes = FleetNodesSpec(min=1, target=1, max=1)
+        fleet = await create_fleet(session=session, project=project, spec=current_spec)
+
+        response = await client.post(
+            f"/api/project/{project.name}/fleets/get_plan",
+            headers=get_auth_headers(user.token),
+            json={"spec": spec.dict()},
+        )
+
+        response_json = response.json()
+        assert response.status_code == 200, response_json
+        assert response_json["current_resource"]["id"] == str(fleet.id)
+        assert response_json["action"] == "update"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_returns_create_plan_for_existing_cloud_fleet_blocks_update(
+        self, test_db, session: AsyncSession, client: AsyncClient
+    ):
+        user = await create_user(session=session, global_role=GlobalRole.USER)
+        project = await create_project(session=session, owner=user)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        current_spec = get_fleet_spec(
+            conf=get_fleet_configuration(nodes=FleetNodesSpec(min=0, target=0, max=1))
+        )
+        spec = current_spec.copy(deep=True)
+        spec.configuration.blocks = 2
+        fleet = await create_fleet(session=session, project=project, spec=current_spec)
+
+        response = await client.post(
+            f"/api/project/{project.name}/fleets/get_plan",
+            headers=get_auth_headers(user.token),
+            json={"spec": spec.dict()},
+        )
+
+        response_json = response.json()
+        assert response.status_code == 200, response_json
+        assert response_json["current_resource"]["id"] == str(fleet.id)
+        assert response_json["action"] == "create"
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_returns_create_plan_for_existing_fleet(

src/tests/_internal/server/services/test_fleets.py

@@ -7,6 +7,7 @@
 from dstack._internal.core.backends.base.backend import Backend
 from dstack._internal.core.errors import ServerClientError
 from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import ApplyAction
 from dstack._internal.core.models.fleets import (
     FleetConfiguration,
     FleetNodesSpec,
@@ -203,3 +204,53 @@ async def test_returns_none_without_current_master_instance(
         )
 
         assert master_provisioning_data is None
+
+
+class TestGetPlanCloudFleetUpdate:
+    @pytest.fixture
+    def get_project_backends_mock(self, monkeypatch: pytest.MonkeyPatch) -> list[Backend]:
+        mock = Mock(spec_set=get_project_backends, return_value=[])
+        monkeypatch.setattr("dstack._internal.server.services.backends.get_project_backends", mock)
+        return mock
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db", "get_project_backends_mock")
+    async def test_ok_nodes_update(self, session: AsyncSession):
+        user = await create_user(session=session)
+        project = await create_project(session=session, owner=user)
+        current_spec = get_fleet_spec(
+            conf=FleetConfiguration(
+                name="my-fleet",
+                nodes=FleetNodesSpec(min=0, target=0, max=1),
+            )
+        )
+        await create_fleet(session=session, project=project, spec=current_spec)
+        new_spec = current_spec.copy(deep=True)
+        new_spec.configuration.nodes = FleetNodesSpec(min=0, target=1, max=1)
+
+        plan = await get_plan(session=session, project=project, user=user, spec=new_spec)
+
+        assert plan.current_resource is not None
+        assert plan.action == ApplyAction.UPDATE
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db", "get_project_backends_mock")
+    async def test_placement_update_requires_recreate(self, session: AsyncSession):
+        user = await create_user(session=session)
+        project = await create_project(session=session, owner=user)
+        current_spec = get_fleet_spec(
+            conf=FleetConfiguration(
+                name="my-fleet",
+                nodes=FleetNodesSpec(min=0, target=0, max=1),
+            )
+        )
+        await create_fleet(session=session, project=project, spec=current_spec)
+        new_spec = current_spec.copy(deep=True)
+        new_spec.configuration.placement = InstanceGroupPlacement.CLUSTER
+
+        plan = await get_plan(session=session, project=project, user=user, spec=new_spec)
+
+        assert plan.current_resource is not None
+        assert plan.action == ApplyAction.CREATE