Add REST plugin for user-defined policies

pyproject.toml

@@ -75,6 +75,9 @@ 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",
@@ -176,3 +179,6 @@ 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/plugins/rest_plugin/README.md

@@ -0,0 +1 @@
+[TODO]

src/plugins/rest_plugin/pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "rest-plugin"
+version = "0.1.0"
+description = "A dstack plugin that enables validation and mutation of run specifications via REST API"
+readme = "README.md"
+requires-python = ">=3.9"
+dependencies = []
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src"]

src/plugins/rest_plugin/src/rest_plugin.py

@@ -0,0 +1,63 @@
+import json
+import os
+import pydantic
+import requests
+from dstack._internal.core.errors import ServerError
+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"
+
+class PreApplyPolicy(ApplyPolicy):
+    def __init__(self):
+        self._plugin_service_uri = os.getenv(PLUGIN_SERVICE_URI_ENV_VAR_NAME)
+        if not self._plugin_service_uri:
+            logger.error(f"Cannot create policy as {PLUGIN_SERVICE_URI_ENV_VAR_NAME} is not set")
+            raise ServerError(f"{PLUGIN_SERVICE_URI_ENV_VAR_NAME} is not set")
+
+    def _call_plugin_service(self, user: str, project: str, spec: ApplySpec, endpoint: str) -> ApplySpec:
+        # Make request to plugin service with run params
+        params = {
+            "user": user,
+            "project": project,
+            "spec": spec.json()
+        }
+        response = None
+        try:
+            response = requests.post(f"{self._plugin_service_uri}/{endpoint}", json=json.dumps(params))
+            response.raise_for_status()
+            spec_json = json.loads(response.text)
+            spec = RunSpec(**spec_json)
+        except requests.RequestException as e:
+            logger.error("Failed to call plugin service: %s", e)
+            if response:
+                logger.error(f"Error response from plugin service:\n{response.text}")
+            logger.info("Returning original run spec")
+            return spec
+        except pydantic.ValidationError as e:
+            # TODO: check response error status and report if plugin service rejected request as invalid
+            logger.exception(f"Plugin service returned invalid response:\n{response.text if response else None}")
+            logger.info("Returning original run spec")
+            return spec
+        logger.info(f"Using RunSpec from plugin service:\n{spec}")
+        return spec
+    def on_run_apply(self, user: str, project: str, spec: RunSpec) -> RunSpec:
+        return self._call_plugin_service(user, project, spec, '/runs/pre_apply')
+
+    def on_fleet_apply(self, user: str, project: str, spec: FleetSpec) -> FleetSpec:
+        return self._call_plugin_service(user, project, spec, '/fleets/pre_apply')
+
+    def on_volume_apply(self, user: str, project: str, spec: VolumeSpec) -> VolumeSpec:
+        return self._call_plugin_service(user, project, spec, '/volumes/pre_apply')
+
+    def on_gateway_apply(self, user: str, project: str, spec: GatewaySpec) -> GatewaySpec:
+        return self._call_plugin_service(user, project, spec, '/gateways/pre_apply')
+
+class RESTPlugin(Plugin):
+    def get_apply_policies(self) -> list[ApplyPolicy]:
+        return [PreApplyPolicy()]

src/tests/plugins/test_rest_plugin.py

@@ -0,0 +1,102 @@
+from dstack._internal.core.errors import ServerError
+from dstack._internal.server.models import ProjectModel, UserModel
+from plugins.rest_plugin.src.rest_plugin import PreApplyPolicy, PLUGIN_SERVICE_URI_ENV_VAR_NAME
+import pytest
+from sqlalchemy.ext.asyncio import AsyncSession
+from pydantic import parse_obj_as
+import os
+import json
+import requests
+from unittest.mock import Mock
+
+from dstack._internal.core.models.runs import RunSpec
+from dstack._internal.core.models.configurations import ServiceConfiguration
+from dstack._internal.core.models.profiles import Profile
+from dstack._internal.core.models.resources import Range
+from dstack._internal.server.testing.common import (
+    create_project,
+    create_user,
+    create_repo,
+    get_run_spec,
+)
+from dstack._internal.server.testing.conf import session, test_db  # noqa: F401
+from dstack._internal.server.services import encryption as encryption  # import for side-effect
+import pytest_asyncio
+from unittest import mock
+
+
+async def create_run_spec(
+    session: AsyncSession,
+    project: ProjectModel,
+    replicas: str = 1,
+) -> RunSpec:
+    repo = await create_repo(session=session, project_id=project.id)
+    run_name = "test-run"
+    profile = Profile(name="test-profile")
+    spec = get_run_spec(
+        repo_id=repo.name,
+        run_name=run_name,
+        profile=profile,
+        configuration=ServiceConfiguration(
+            commands=["echo hello"],
+            port=8000,
+            replicas=parse_obj_as(Range[int], replicas)
+        ),
+    )
+    return spec
+
+@pytest_asyncio.fixture
+async def project(session):
+    return await create_project(session=session)
+
+@pytest_asyncio.fixture
+async def user(session):
+    return await create_user(session=session)
+
+@pytest_asyncio.fixture
+async def run_spec(session, project):
+    return await create_run_spec(session=session, project=project)
+
+
+class TestRESTPlugin:
+    @pytest.mark.asyncio
+    async def test_on_run_apply_plugin_service_uri_not_set(self):
+        with pytest.raises(ServerError):
+            policy = PreApplyPolicy()
+
+    @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)
+    async def test_on_run_apply_plugin_service_returns_mutated_spec(self, test_db, user, project, run_spec):
+        policy = PreApplyPolicy()
+        mock_response = Mock()
+        run_spec_dict = run_spec.dict()
+        run_spec_dict["profile"]["tags"] = {"env": "test", "team": "qa"}
+        mock_response.text = json.dumps(run_spec_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=run_spec)
+            assert result == RunSpec(**run_spec_dict)
+
+    @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)
+    async def test_on_run_apply_plugin_service_call_fails(self, test_db, user, project, run_spec):
+        policy = PreApplyPolicy()
+        with mock.patch("requests.post", side_effect=requests.RequestException("fail")):
+            result = policy.on_apply(user=user.name, project=project.name, spec=run_spec)
+            assert result == run_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)
+    async def test_on_run_apply_plugin_service_returns_invalid_spec(self, test_db, user, project, run_spec):
+        policy = PreApplyPolicy()
+        mock_response = Mock()
+        mock_response.text = json.dumps({"invalid-key": "abc"})
+        mock_response.raise_for_status = Mock()
+        with mock.patch("requests.post", return_value=mock_response):
+            result = policy.on_apply(user.name, project=project.name, spec=run_spec)
+            # return original run spec
+            assert result == run_spec
+