Change rest-plugin to a builtin plugin

Author: Nadine-H

pyproject.toml

@@ -75,9 +75,6 @@ pattern = '<picture>\s*|<source[^>]*>\s*|\s*</picture>|<video[^>]*>\s*|</video>\
 replacement = ''
 ignore-case = true
 
-[tool.uv.workspace]
-members = ["src/plugins/rest_plugin"]
-
 [dependency-groups]
 dev = [
     "build>=1.2.2.post1",
@@ -179,6 +176,3 @@ nebius = [
 all = [
     "dstack[gateway,server,aws,azure,gcp,datacrunch,kubernetes,lambda,nebius,oci]",
 ]
-
-[project.entry-points."dstack.plugins"]
-rest_plugin = "plugins.rest_plugin.src.rest_plugin:RESTPlugin"

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

@@ -1,5 +1,6 @@
 import itertools
 from importlib import import_module
+from typing import Dict
 
 from backports.entry_points_selectable import entry_points  # backport for Python 3.9
 
@@ -12,50 +13,80 @@
 
 _PLUGINS: list[Plugin] = []
 
+_BUILTIN_PLUGINS: Dict[str, str] = {"rest_plugin": "dstack.plugins.builtin.rest_plugin:RESTPlugin"}
 
-def load_plugins(enabled_plugins: list[str]):
-    _PLUGINS.clear()
-    plugins_entrypoints = entry_points(group="dstack.plugins")
-    plugins_to_load = enabled_plugins.copy()
-    for entrypoint in plugins_entrypoints:
-        if entrypoint.name not in enabled_plugins:
-            logger.info(
-                ("Found not enabled plugin %s. Plugin will not be loaded."),
-                entrypoint.name,
-            )
-            continue
+
+class PluginEntrypoint:
+    def __init__(self, name: str, import_path: str, is_builtin: bool = False):
+        self.name = name
+        self.import_path = import_path
+        self.is_builtin = is_builtin
+
+    def load(self):
+        module_path, _, class_name = self.import_path.partition(":")
         try:
-            module_path, _, class_name = entrypoint.value.partition(":")
             module = import_module(module_path)
+            plugin_class = getattr(module, class_name, None)
+            if plugin_class is None:
+                logger.warning(
+                    ("Failed to load plugin %s: plugin class %s not found in module %s."),
+                    self.name,
+                    class_name,
+                    module_path,
+                )
+                return None
+            if not issubclass(plugin_class, Plugin):
+                logger.warning(
+                    ("Failed to load plugin %s: plugin class %s is not a subclass of Plugin."),
+                    self.name,
+                    class_name,
+                )
+                return None
+            return plugin_class()
         except ImportError:
             logger.warning(
                 (
                     "Failed to load plugin %s when importing %s."
                     " Ensure the module is on the import path."
                 ),
-                entrypoint.name,
-                entrypoint.value,
+                self.name,
+                self.import_path,
             )
-            continue
-        plugin_class = getattr(module, class_name, None)
-        if plugin_class is None:
-            logger.warning(
-                ("Failed to load plugin %s: plugin class %s not found in module %s."),
+            return None
+
+
+def load_plugins(enabled_plugins: list[str]):
+    _PLUGINS.clear()
+    entrypoints: dict[str, PluginEntrypoint] = {}
+    plugins_to_load = enabled_plugins.copy()
+    for entrypoint in entry_points(group="dstack.plugins"):
+        if entrypoint.name not in enabled_plugins:
+            logger.info(
+                ("Found not enabled plugin %s. Plugin will not be loaded."),
                 entrypoint.name,
-                class_name,
-                module_path,
             )
             continue
-        if not issubclass(plugin_class, Plugin):
-            logger.warning(
-                ("Failed to load plugin %s: plugin class %s is not a subclass of Plugin."),
-                entrypoint.name,
-                class_name,
+        else:
+            entrypoints[entrypoint.name] = PluginEntrypoint(
+                entrypoint.name, entrypoint.value, is_builtin=False
             )
-            continue
-        plugins_to_load.remove(entrypoint.name)
-        _PLUGINS.append(plugin_class())
-        logger.info("Loaded plugin %s", entrypoint.name)
+
+    for name, import_path in _BUILTIN_PLUGINS.items():
+        if name not in enabled_plugins:
+            logger.info(
+                ("Found not enabled builtin plugin %s. Plugin will not be loaded."),
+                name,
+            )
+        else:
+            entrypoints[name] = PluginEntrypoint(name, import_path, is_builtin=True)
+
+    for plugin_name, plugin_entrypoint in entrypoints.items():
+        plugin_instance = plugin_entrypoint.load()
+        if plugin_instance is not None:
+            _PLUGINS.append(plugin_instance)
+            plugins_to_load.remove(plugin_name)
+            logger.info("Loaded plugin %s", plugin_name)
+
     if plugins_to_load:
         logger.warning("Enabled plugins not found: %s", plugins_to_load)
 
@@ -65,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 ValueError as e:
+        except Exception as e:
             msg = None
             if len(e.args) > 0:
                 msg = e.args[0]

src/dstack/plugins/builtin/rest_plugin.py

@@ -0,0 +1,97 @@
+import json
+import os
+from typing import Generic, TypeVar
+
+import requests
+from pydantic import BaseModel, 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
+
+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):
+        self._plugin_service_uri = os.getenv(PLUGIN_SERVICE_URI_ENV_VAR_NAME)
+        logger.info(f"Found plugin service at {self._plugin_service_uri}")
+        if not self._plugin_service_uri:
+            logger.error(
+                f"Cannot create policy because {PLUGIN_SERVICE_URI_ENV_VAR_NAME} is not set"
+            )
+            raise ServerClientError(f"{PLUGIN_SERVICE_URI_ENV_VAR_NAME} is not set")
+
+    def _call_plugin_service(self, spec_request: SpecRequest, endpoint: str) -> ApplySpec:
+        response = None
+        try:
+            response = requests.post(
+                f"{self._plugin_service_uri}{endpoint}",
+                json=spec_request.dict(),
+                headers={"accept": "application/json", "Content-Type": "application/json"},
+                timeout=PLUGIN_REQUEST_TIMEOUT,
+            )
+            response.raise_for_status()
+            spec_json = json.loads(response.text)
+            return spec_json
+        except requests.exceptions.ConnectionError as e:
+            logger.error(
+                f"Could not connect to plugin service at {self._plugin_service_uri}: %s", e
+            )
+            raise e
+        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
+
+    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)
+
+    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)
+
+    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)
+
+    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)
+
+
+class RESTPlugin(Plugin):
+    def get_apply_policies(self) -> list[ApplyPolicy]:
+        return [CustomApplyPolicy()]