Implement /api/project/{project_name}/fleets/apply (#2577)

Author: r4victor

src/dstack/_internal/cli/services/configurators/fleet.py

@@ -17,7 +17,13 @@
 )
 from dstack._internal.cli.utils.fleet import get_fleets_table
 from dstack._internal.cli.utils.rich import MultiItemStatus
-from dstack._internal.core.errors import ConfigurationError, ResourceNotExistsError
+from dstack._internal.core.errors import (
+    CLIError,
+    ConfigurationError,
+    ResourceNotExistsError,
+    ServerClientError,
+    URLNotFoundError,
+)
 from dstack._internal.core.models.configurations import ApplyConfigurationType
 from dstack._internal.core.models.fleets import (
     Fleet,
@@ -31,6 +37,7 @@
 from dstack._internal.utils.common import local_time
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.ssh import convert_ssh_key_to_pem, generate_public_key, pkey_from_str
+from dstack.api._public import Client
 from dstack.api.utils import load_profile
 
 logger = get_logger(__name__)
@@ -109,11 +116,11 @@ def apply_configuration(
                     else:
                         time.sleep(1)
 
-        with console.status("Creating fleet..."):
-            fleet = self.api.client.fleets.create(
-                project_name=self.api.project,
-                spec=spec,
-            )
+        try:
+            with console.status("Applying plan..."):
+                fleet = _apply_plan(self.api, plan)
+        except ServerClientError as e:
+            raise CLIError(e.msg)
         if command_args.detach:
             console.print("Fleet configuration submitted. Exiting...")
             return
@@ -350,3 +357,17 @@ def _failed_provisioning(fleet: Fleet) -> bool:
         if instance.status == InstanceStatus.TERMINATED:
             return True
     return False
+
+
+def _apply_plan(api: Client, plan: FleetPlan) -> Fleet:
+    try:
+        return api.client.fleets.apply_plan(
+            project_name=api.project,
+            plan=plan,
+        )
+    except URLNotFoundError:
+        # TODO: Remove in 0.20
+        return api.client.fleets.create(
+            project_name=api.project,
+            spec=plan.spec,
+        )

src/dstack/_internal/core/errors.py

@@ -18,6 +18,10 @@ class ClientError(DstackError):
     pass
 
 
+class URLNotFoundError(ClientError):
+    pass
+
+
 class ServerClientErrorCode(str, enum.Enum):
     UNSPECIFIED_ERROR = "error"
     RESOURCE_EXISTS = "resource_exists"

src/dstack/_internal/core/models/fleets.py

@@ -312,7 +312,20 @@ class FleetPlan(CoreModel):
     project_name: str
     user: str
     spec: FleetSpec
-    current_resource: Optional[Fleet]
+    current_resource: Optional[Fleet] = None
     offers: List[InstanceOfferWithAvailability]
     total_offers: int
-    max_offer_price: Optional[float]
+    max_offer_price: Optional[float] = None
+
+
+class ApplyFleetPlanInput(CoreModel):
+    spec: FleetSpec
+    current_resource: Annotated[
+        Optional[Fleet],
+        Field(
+            description=(
+                "The expected current resource."
+                " If the resource has changed, the apply fails unless `force: true`."
+            )
+        ),
+    ] = None

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

@@ -9,6 +9,7 @@
 from dstack._internal.server.db import get_session
 from dstack._internal.server.models import ProjectModel, UserModel
 from dstack._internal.server.schemas.fleets import (
+    ApplyFleetPlanRequest,
     CreateFleetRequest,
     DeleteFleetInstancesRequest,
     DeleteFleetsRequest,
@@ -107,6 +108,27 @@ async def get_plan(
     return plan
 
 
+@project_router.post("/apply")
+async def apply_plan(
+    body: ApplyFleetPlanRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
+) -> Fleet:
+    """
+    Creates a new fleet or updates an existing fleet.
+    Errors if the expected current resource from the plan does not match the current resource.
+    Use `force: true` to apply even if the current resource does not match.
+    """
+    user, project = user_project
+    return await fleets_services.apply_plan(
+        session=session,
+        user=user,
+        project=project,
+        plan=body.plan,
+        force=body.force,
+    )
+
+
 @project_router.post("/create")
 async def create_fleet(
     body: CreateFleetRequest,

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

@@ -1,11 +1,11 @@
 from datetime import datetime
-from typing import List, Optional
+from typing import Annotated, List, Optional
 from uuid import UUID
 
 from pydantic import Field
 
 from dstack._internal.core.models.common import CoreModel
-from dstack._internal.core.models.fleets import FleetSpec
+from dstack._internal.core.models.fleets import ApplyFleetPlanInput, FleetSpec
 
 
 class ListFleetsRequest(CoreModel):
@@ -26,6 +26,16 @@ class GetFleetPlanRequest(CoreModel):
     spec: FleetSpec
 
 
+class ApplyFleetPlanRequest(CoreModel):
+    plan: ApplyFleetPlanInput
+    force: Annotated[
+        bool,
+        Field(
+            description="Use `force: true` to apply even if the expected resource does not match."
+        ),
+    ]
+
+
 class CreateFleetRequest(CoreModel):
     spec: FleetSpec
 

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

@@ -17,6 +17,7 @@
 )
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.fleets import (
+    ApplyFleetPlanInput,
     Fleet,
     FleetPlan,
     FleetSpec,
@@ -306,6 +307,21 @@ async def get_create_instance_offers(
     return offers
 
 
+async def apply_plan(
+    session: AsyncSession,
+    user: UserModel,
+    project: ProjectModel,
+    plan: ApplyFleetPlanInput,
+    force: bool,
+) -> Fleet:
+    return await create_fleet(
+        session=session,
+        project=project,
+        user=user,
+        spec=plan.spec,
+    )
+
+
 async def create_fleet(
     session: AsyncSession,
     project: ProjectModel,

src/dstack/api/server/__init__.py

@@ -6,7 +6,7 @@
 import requests
 
 from dstack import version
-from dstack._internal.core.errors import ClientError, ServerClientError
+from dstack._internal.core.errors import ClientError, ServerClientError, URLNotFoundError
 from dstack._internal.utils.logging import get_logger
 from dstack.api.server._backends import BackendsAPIClient
 from dstack.api.server._fleets import FleetsAPIClient
@@ -154,6 +154,8 @@ def _request(
                 raise ClientError(
                     f"Access to {resp.request.url} is denied. Please check your access token"
                 )
+            if resp.status_code == 404:
+                raise URLNotFoundError(f"Status code 404 when requesting {resp.request.url}")
             if 400 <= resp.status_code < 600:
                 raise ClientError(
                     f"Unexpected error: status code {resp.status_code}"

src/dstack/api/server/_fleets.py

@@ -1,9 +1,10 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import parse_obj_as
 
-from dstack._internal.core.models.fleets import Fleet, FleetPlan, FleetSpec
+from dstack._internal.core.models.fleets import ApplyFleetPlanInput, Fleet, FleetPlan, FleetSpec
 from dstack._internal.server.schemas.fleets import (
+    ApplyFleetPlanRequest,
     CreateFleetRequest,
     DeleteFleetInstancesRequest,
     DeleteFleetsRequest,
@@ -32,18 +33,20 @@ def get_plan(
         spec: FleetSpec,
     ) -> FleetPlan:
         body = GetFleetPlanRequest(spec=spec)
-        body_json = body.json(exclude=_get_fleet_spec_excludes(spec))
+        body_json = body.json(exclude=_get_get_plan_excludes(spec))
         resp = self._request(f"/api/project/{project_name}/fleets/get_plan", body=body_json)
         return parse_obj_as(FleetPlan.__response__, resp.json())
 
-    def create(
+    def apply_plan(
         self,
         project_name: str,
-        spec: FleetSpec,
+        plan: Union[FleetPlan, ApplyFleetPlanInput],
+        force: bool = False,
     ) -> Fleet:
-        body = CreateFleetRequest(spec=spec)
-        body_json = body.json(exclude=_get_fleet_spec_excludes(spec))
-        resp = self._request(f"/api/project/{project_name}/fleets/create", body=body_json)
+        plan_input = ApplyFleetPlanInput.__response__.parse_obj(plan)
+        body = ApplyFleetPlanRequest(plan=plan_input, force=force)
+        body_json = body.json(exclude=_get_apply_plan_excludes(plan_input))
+        resp = self._request(f"/api/project/{project_name}/fleets/apply", body=body_json)
         return parse_obj_as(Fleet.__response__, resp.json())
 
     def delete(self, project_name: str, names: List[str]) -> None:
@@ -54,6 +57,42 @@ def delete_instances(self, project_name: str, name: str, instance_nums: List[int
         body = DeleteFleetInstancesRequest(name=name, instance_nums=instance_nums)
         self._request(f"/api/project/{project_name}/fleets/delete_instances", body=body.json())
 
+    # Deprecated
+    # TODO: Remove in 0.20
+    def create(
+        self,
+        project_name: str,
+        spec: FleetSpec,
+    ) -> Fleet:
+        body = CreateFleetRequest(spec=spec)
+        body_json = body.json(exclude=_get_create_fleet_excludes(spec))
+        resp = self._request(f"/api/project/{project_name}/fleets/create", body=body_json)
+        return parse_obj_as(Fleet.__response__, resp.json())
+
+
+def _get_get_plan_excludes(fleet_spec: FleetSpec) -> Dict:
+    get_plan_excludes = {}
+    spec_excludes = _get_fleet_spec_excludes(fleet_spec)
+    if spec_excludes:
+        get_plan_excludes["spec"] = spec_excludes
+    return get_plan_excludes
+
+
+def _get_apply_plan_excludes(plan_input: ApplyFleetPlanInput) -> Dict:
+    apply_plan_excludes = {}
+    spec_excludes = _get_fleet_spec_excludes(plan_input.spec)
+    if spec_excludes:
+        apply_plan_excludes["spec"] = apply_plan_excludes
+    return {"plan": apply_plan_excludes}
+
+
+def _get_create_fleet_excludes(fleet_spec: FleetSpec) -> Dict:
+    create_fleet_excludes = {}
+    spec_excludes = _get_fleet_spec_excludes(fleet_spec)
+    if spec_excludes:
+        create_fleet_excludes["spec"] = spec_excludes
+    return create_fleet_excludes
+
 
 def _get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[Dict]:
     """
@@ -76,5 +115,5 @@ def _get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[Dict]:
     if profile_excludes:
         spec_excludes["profile"] = profile_excludes
     if spec_excludes:
-        return {"spec": spec_excludes}
+        return spec_excludes
     return None

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

@@ -292,13 +292,13 @@ async def test_not_returns_by_name_if_fleet_does_not_exist(
         assert response.status_code == 400
 
 
-class TestCreateFleet:
+class TestApplyFleetPlan:
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_returns_40x_if_not_authenticated(
         self, test_db, session: AsyncSession, client: AsyncClient
     ):
-        response = await client.post("/api/project/main/fleets/create")
+        response = await client.post("/api/project/main/fleets/apply")
         assert response.status_code == 403
 
     @pytest.mark.asyncio
@@ -314,9 +314,9 @@ async def test_creates_fleet(self, test_db, session: AsyncSession, client: Async
         with patch("uuid.uuid4") as m:
             m.return_value = UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e")
             response = await client.post(
-                f"/api/project/{project.name}/fleets/create",
+                f"/api/project/{project.name}/fleets/apply",
                 headers=get_auth_headers(user.token),
-                json={"spec": spec.dict()},
+                json={"plan": {"spec": spec.dict()}, "force": False},
             )
         assert response.status_code == 200
         assert response.json() == {
@@ -427,9 +427,9 @@ async def test_creates_ssh_fleet(self, test_db, session: AsyncSession, client: A
         with patch("uuid.uuid4") as m:
             m.return_value = UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e")
             response = await client.post(
-                f"/api/project/{project.name}/fleets/create",
+                f"/api/project/{project.name}/fleets/apply",
                 headers=get_auth_headers(user.token),
-                json={"spec": spec.dict()},
+                json={"plan": {"spec": spec.dict()}, "force": False},
             )
         assert response.status_code == 200, response.json()
         assert response.json() == {
@@ -559,9 +559,9 @@ async def test_errors_if_ssh_key_is_bad(
             )
         )
         response = await client.post(
-            f"/api/project/{project.name}/fleets/create",
+            f"/api/project/{project.name}/fleets/apply",
             headers=get_auth_headers(user.token),
-            json={"spec": spec.dict()},
+            json={"plan": {"spec": spec.dict()}, "force": False},
         )
         assert response.status_code == 400
 
@@ -590,9 +590,9 @@ async def test_forbids_if_no_permission_to_manage_ssh_fleets(
             DefaultPermissions(allow_non_admins_manage_ssh_fleets=False)
         ):
             response = await client.post(
-                f"/api/project/{project.name}/fleets/create",
+                f"/api/project/{project.name}/fleets/apply",
                 headers=get_auth_headers(user.token),
-                json={"spec": spec.dict()},
+                json={"plan": {"spec": spec.dict()}, "force": False},
             )
         assert response.status_code in [401, 403]