Refactor rest-plugin + polish models

Author: Nadine-H

examples/plugins/example_plugin_server/app/main.py

@@ -2,8 +2,17 @@
 
 from fastapi import FastAPI
 
-from app.models import FleetSpecRequest, GatewaySpecRequest, RunSpecRequest, VolumeSpecRequest
 from app.utils import configure_logging
+from dstack.plugins.builtin.models import (
+    FleetSpecRequest,
+    FleetSpecResponse,
+    GatewaySpecRequest,
+    GatewaySpecResponse,
+    RunSpecRequest,
+    RunSpecResponse,
+    VolumeSpecRequest,
+    VolumeSpecResponse,
+)
 
 configure_logging()
 logger = logging.getLogger(__name__)
@@ -12,32 +21,36 @@
 
 
 @app.post("/apply_policies/on_run_apply")
-async def on_run_apply(request: RunSpecRequest):
+async def on_run_apply(request: RunSpecRequest) -> RunSpecResponse:
     logger.info(
         f"Received run spec request from user {request.user} and project {request.project}"
     )
-    return request.spec
+    response = RunSpecResponse(spec=request.spec, error=None)
+    return response
 
 
 @app.post("/apply_policies/on_fleet_apply")
-async def on_fleet_apply(request: FleetSpecRequest):
+async def on_fleet_apply(request: FleetSpecRequest) -> FleetSpecResponse:
     logger.info(
         f"Received fleet spec request from user {request.user} and project {request.project}"
     )
-    return request.spec
+    response = FleetSpecResponse(request.spec, error=None)
+    return response
 
 
 @app.post("/apply_policies/on_volume_apply")
-async def on_volume_apply(request: VolumeSpecRequest):
+async def on_volume_apply(request: VolumeSpecRequest) -> VolumeSpecResponse:
     logger.info(
         f"Received volume spec request from user {request.user} and project {request.project}"
     )
-    return request.spec
+    response = VolumeSpecResponse(request.spec, error=None)
+    return response
 
 
 @app.post("/apply_policies/on_gateway_apply")
-async def on_gateway_apply(request: GatewaySpecRequest):
+async def on_gateway_apply(request: GatewaySpecRequest) -> GatewaySpecResponse:
     logger.info(
         f"Received gateway spec request from user {request.user} and project {request.project}"
     )
-    return request.spec
+    response = GatewaySpecResponse(request.spec, error=None)
+    return response

examples/plugins/example_plugin_server/app/models.py

@@ -96,7 +96,7 @@ def apply_plugin_policies(user: str, project: str, spec: ApplySpec) -> ApplySpec
     for policy in policies:
         try:
             spec = policy.on_apply(user=user, project=project, spec=spec)
-        except Exception as e:
+        except ValueError as e:
             msg = None
             if len(e.args) > 0:
                 msg = e.args[0]

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

@@ -0,0 +1,38 @@
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel
+
+from dstack._internal.core.models.fleets import FleetSpec
+from dstack._internal.core.models.gateways import GatewaySpec
+from dstack._internal.core.models.runs import RunSpec
+from dstack._internal.core.models.volumes import VolumeSpec
+
+SpecType = TypeVar("SpecType", RunSpec, FleetSpec, VolumeSpec, GatewaySpec)
+
+
+class SpecApplyRequest(BaseModel, Generic[SpecType]):
+    user: str
+    project: str
+    spec: SpecType
+
+    def dict(self, *args, **kwargs):
+        d = super().dict(*args, **kwargs)
+        d.pop("__orig_class__", None)
+        return d
+
+
+RunSpecRequest = SpecApplyRequest[RunSpec]
+FleetSpecRequest = SpecApplyRequest[FleetSpec]
+VolumeSpecRequest = SpecApplyRequest[VolumeSpec]
+GatewaySpecRequest = SpecApplyRequest[GatewaySpec]
+
+
+class SpecApplyResponse(BaseModel, Generic[SpecType]):
+    spec: SpecType
+    error: str | None = None
+
+
+RunSpecResponse = SpecApplyResponse[RunSpec]
+FleetSpecResponse = SpecApplyResponse[FleetSpec]
+VolumeSpecResponse = SpecApplyResponse[VolumeSpec]
+GatewaySpecResponse = SpecApplyResponse[GatewaySpec]

src/dstack/plugins/builtin/__init__.py

@@ -1,36 +1,33 @@
 import json
 import os
-from typing import Generic, TypeVar
 
 import requests
-from pydantic import BaseModel, ValidationError
+from pydantic import ValidationError
 
 from dstack._internal.core.errors import ServerClientError
 from dstack._internal.core.models.fleets import FleetSpec
 from dstack._internal.core.models.gateways import GatewaySpec
 from dstack._internal.core.models.volumes import VolumeSpec
 from dstack.plugins import ApplyPolicy, Plugin, RunSpec, get_plugin_logger
 from dstack.plugins._models import ApplySpec
+from dstack.plugins.builtin.models import (
+    FleetSpecRequest,
+    FleetSpecResponse,
+    GatewaySpecRequest,
+    GatewaySpecResponse,
+    RunSpecRequest,
+    RunSpecResponse,
+    SpecApplyRequest,
+    SpecApplyResponse,
+    VolumeSpecRequest,
+    VolumeSpecResponse,
+)
 
 logger = get_plugin_logger(__name__)
 
 PLUGIN_SERVICE_URI_ENV_VAR_NAME = "DSTACK_PLUGIN_SERVICE_URI"
 PLUGIN_REQUEST_TIMEOUT = 8  # in seconds
 
-SpecType = TypeVar("SpecType", RunSpec, FleetSpec, VolumeSpec, GatewaySpec)
-
-
-class SpecRequest(BaseModel, Generic[SpecType]):
-    user: str
-    project: str
-    spec: SpecType
-
-
-RunSpecRequest = SpecRequest[RunSpec]
-FleetSpecRequest = SpecRequest[FleetSpec]
-VolumeSpecRequest = SpecRequest[VolumeSpec]
-GatewaySpecRequest = SpecRequest[GatewaySpec]
-
 
 class CustomApplyPolicy(ApplyPolicy):
     def __init__(self):
@@ -42,7 +39,12 @@ def __init__(self):
             )
             raise ServerClientError(f"{PLUGIN_SERVICE_URI_ENV_VAR_NAME} is not set")
 
-    def _call_plugin_service(self, spec_request: SpecRequest, endpoint: str) -> ApplySpec:
+    def _check_request_rejected(self, response: SpecApplyResponse):
+        if response.error is not None:
+            logger.error(f"Plugin service rejected apply request: {response.error}")
+            raise ServerClientError(f"Apply request rejected: {response.error}")
+
+    def _call_plugin_service(self, spec_request: SpecApplyRequest, endpoint: str) -> ApplySpec:
         response = None
         try:
             response = requests.post(
@@ -58,38 +60,58 @@ def _call_plugin_service(self, spec_request: SpecRequest, endpoint: str) -> Appl
             logger.error(
                 f"Could not connect to plugin service at {self._plugin_service_uri}: %s", e
             )
-            raise e
+            raise ServerClientError(
+                f"Could not connect to plugin service at {self._plugin_service_uri}"
+            )
         except requests.RequestException as e:
             logger.error("Request to the plugin service failed: %s", e)
-            if response:
-                logger.error(f"Error response from plugin service:\n{response.text}")
-            raise e
-        except ValidationError as e:
-            # Received 200 code but response body is invalid
-            logger.exception(
-                f"Plugin service returned invalid response:\n{response.text if response else None}"
-            )
-            raise e
+            raise ServerClientError("Request to the plugin service failed")
+
+    def _on_apply(self, request_cls, response_cls, endpoint, user, project, spec):
+        try:
+            spec_request = request_cls(user=user, project=project, spec=spec)
+            spec_json = self._call_plugin_service(spec_request, endpoint)
+            response = response_cls(**spec_json)
+            self._check_request_rejected(response)
+            return response.spec
+        except ValidationError:
+            logger.error(f"Plugin service returned invalid response:\n{spec_json}")
+            raise ServerClientError("Plugin service returned an invalid response")
 
     def on_run_apply(self, user: str, project: str, spec: RunSpec) -> RunSpec:
-        spec_request = RunSpecRequest(user=user, project=project, spec=spec)
-        spec_json = self._call_plugin_service(spec_request, "/apply_policies/on_run_apply")
-        return RunSpec(**spec_json)
+        return self._on_apply(
+            RunSpecRequest, RunSpecResponse, "/apply_policies/on_run_apply", user, project, spec
+        )
 
     def on_fleet_apply(self, user: str, project: str, spec: FleetSpec) -> FleetSpec:
-        spec_request = FleetSpecRequest(user=user, project=project, spec=spec)
-        spec_json = self._call_plugin_service(spec_request, "/apply_policies/on_fleet_apply")
-        return FleetSpec(**spec_json)
+        return self._on_apply(
+            FleetSpecRequest,
+            FleetSpecResponse,
+            "/apply_policies/on_fleet_apply",
+            user,
+            project,
+            spec,
+        )
 
     def on_volume_apply(self, user: str, project: str, spec: VolumeSpec) -> VolumeSpec:
-        spec_request = VolumeSpecRequest(user=user, project=project, spec=spec)
-        spec_json = self._call_plugin_service(spec_request, "/apply_policies/on_volume_apply")
-        return VolumeSpec(**spec_json)
+        return self._on_apply(
+            VolumeSpecRequest,
+            VolumeSpecResponse,
+            "/apply_policies/on_volume_apply",
+            user,
+            project,
+            spec,
+        )
 
     def on_gateway_apply(self, user: str, project: str, spec: GatewaySpec) -> GatewaySpec:
-        spec_request = GatewaySpecRequest(user=user, project=project, spec=spec)
-        spec_json = self._call_plugin_service(spec_request, "/apply_policies/on_gateway_apply")
-        return GatewaySpec(**spec_json)
+        return self._on_apply(
+            GatewaySpecRequest,
+            GatewaySpecResponse,
+            "/apply_policies/on_gateway_apply",
+            user,
+            project,
+            spec,
+        )
 
 
 class RESTPlugin(Plugin):

src/dstack/plugins/builtin/models.py

@@ -10,7 +10,7 @@
 from pydantic import parse_obj_as
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from dstack._internal.core.errors import ServerError
+from dstack._internal.core.errors import ServerClientError, ServerError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.configurations import ServiceConfiguration
 from dstack._internal.core.models.fleets import FleetConfiguration, FleetSpec
@@ -106,34 +106,34 @@ async def test_on_run_apply_plugin_service_uri_not_set(self):
     @pytest.mark.parametrize(
         "spec", ["run_spec", "fleet_spec", "volume_spec", "gateway_spec"], indirect=True
     )
-    async def test_on_run_apply_plugin_service_returns_mutated_spec(
+    async def test_on_apply_plugin_service_returns_mutated_spec(
         self, test_db, user, project, spec
     ):
         policy = CustomApplyPolicy()
         mock_response = Mock()
-        spec_dict = spec.dict()
+        response_dict = {"spec": spec.dict(), "error": None}
 
         if isinstance(spec, (RunSpec, FleetSpec)):
-            spec_dict["profile"]["tags"] = {"env": "test", "team": "qa"}
+            response_dict["spec"]["profile"]["tags"] = {"env": "test", "team": "qa"}
         else:
-            spec_dict["configuration_path"] = "/path/to/something"
+            response_dict["spec"]["configuration_path"] = "/path/to/something"
 
-        mock_response.text = json.dumps(spec_dict)
+        mock_response.text = json.dumps(response_dict)
         mock_response.raise_for_status = Mock()
         with mock.patch("requests.post", return_value=mock_response):
             result = policy.on_apply(user=user.name, project=project.name, spec=spec)
-            assert result == type(spec)(**spec_dict)
+            assert result == type(spec)(**response_dict["spec"])
 
     @pytest.mark.asyncio
     @mock.patch.dict(os.environ, {PLUGIN_SERVICE_URI_ENV_VAR_NAME: "http://mock"})
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     @pytest.mark.parametrize(
         "spec", ["run_spec", "fleet_spec", "volume_spec", "gateway_spec"], indirect=True
     )
-    async def test_on_run_apply_plugin_service_call_fails(self, test_db, user, project, spec):
+    async def test_on_apply_plugin_service_call_fails(self, test_db, user, project, spec):
         policy = CustomApplyPolicy()
         with mock.patch("requests.post", side_effect=requests.RequestException("fail")):
-            with pytest.raises(requests.RequestException):
+            with pytest.raises(ServerClientError):
                 policy.on_apply(user=user.name, project=project.name, spec=spec)
 
     @pytest.mark.asyncio
@@ -142,14 +142,12 @@ async def test_on_run_apply_plugin_service_call_fails(self, test_db, user, proje
     @pytest.mark.parametrize(
         "spec", ["run_spec", "fleet_spec", "volume_spec", "gateway_spec"], indirect=True
     )
-    async def test_on_run_apply_plugin_service_connection_fails(
-        self, test_db, user, project, spec
-    ):
+    async def test_on_apply_plugin_service_connection_fails(self, test_db, user, project, spec):
         policy = CustomApplyPolicy()
         with mock.patch(
             "requests.post", side_effect=requests.ConnectionError("Failed to connect")
         ):
-            with pytest.raises(requests.ConnectionError):
+            with pytest.raises(ServerClientError):
                 policy.on_apply(user=user.name, project=project.name, spec=spec)
 
     @pytest.mark.asyncio
@@ -158,7 +156,7 @@ async def test_on_run_apply_plugin_service_connection_fails(
     @pytest.mark.parametrize(
         "spec", ["run_spec", "fleet_spec", "volume_spec", "gateway_spec"], indirect=True
     )
-    async def test_on_run_apply_plugin_service_returns_invalid_spec(
+    async def test_on_apply_plugin_service_returns_invalid_spec(
         self, test_db, user, project, spec
     ):
         policy = CustomApplyPolicy()