Restructure rest_plugin modules

Author: Nadine-H

examples/plugins/example_plugin_server/app/main.py

@@ -3,7 +3,7 @@
 from fastapi import FastAPI
 
 from app.utils import configure_logging
-from dstack.plugins.builtin.models import (
+from dstack.plugins.builtin.rest_plugin import (
     FleetSpecRequest,
     FleetSpecResponse,
     GatewaySpecRequest,

src/dstack/plugins/builtin/rest_plugin/__init__.py

@@ -0,0 +1,18 @@
+# ruff: noqa: F401
+from dstack.plugins.builtin.rest_plugin._models import (
+    FleetSpecRequest,
+    FleetSpecResponse,
+    GatewaySpecRequest,
+    GatewaySpecResponse,
+    RunSpecRequest,
+    RunSpecResponse,
+    SpecApplyRequest,
+    SpecApplyResponse,
+    VolumeSpecRequest,
+    VolumeSpecResponse,
+)
+from dstack.plugins.builtin.rest_plugin._plugin import (
+    PLUGIN_SERVICE_URI_ENV_VAR_NAME,
+    CustomApplyPolicy,
+    RESTPlugin,
+)

src/dstack/plugins/builtin/models.py renamed to src/dstack/plugins/builtin/rest_plugin/_models.py

@@ -15,6 +15,8 @@ class SpecApplyRequest(BaseModel, Generic[SpecType]):
     project: str
     spec: SpecType
 
+    # Override dict() to remove __orig_class__ attribute and avoid "TypeError: Object of type _GenericAlias is not JSON serializable"
+    # error. This issue doesn't happen though when running the code in pytest, only when running the server.
     def dict(self, *args, **kwargs):
         d = super().dict(*args, **kwargs)
         d.pop("__orig_class__", None)

src/dstack/plugins/builtin/rest_plugin.py renamed to src/dstack/plugins/builtin/rest_plugin/_plugin.py

@@ -8,9 +8,8 @@
 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 (
+from dstack.plugins import ApplyPolicy, ApplySpec, Plugin, RunSpec, get_plugin_logger
+from dstack.plugins.builtin.rest_plugin import (
     FleetSpecRequest,
     FleetSpecResponse,
     GatewaySpecRequest,

src/tests/plugins/test_rest_plugin.py

@@ -3,7 +3,6 @@
 from unittest import mock
 from unittest.mock import Mock
 
-import pydantic
 import pytest
 import pytest_asyncio
 import requests
@@ -164,5 +163,5 @@ async def test_on_apply_plugin_service_returns_invalid_spec(
         mock_response.text = json.dumps({"invalid-key": "abc"})
         mock_response.raise_for_status = Mock()
         with mock.patch("requests.post", return_value=mock_response):
-            with pytest.raises(pydantic.ValidationError):
+            with pytest.raises(ServerClientError):
                 policy.on_apply(user.name, project=project.name, spec=spec)