Implement plugins (#2581)

* Implement plugins prototype

* Add per-configuration ApplyPolicy methods

* Add get_plugin_logger()

* feat: Add tests for loading dstack plugins

* Document get_plugin_logger()

* Backport entry_points for Python 3.9

* Add plugin example

* Document Concepts->Plugins

* Add example on setting service-specific properties in plugins

Author: r4victor

docs/docs/concepts/plugins.md

@@ -0,0 +1,116 @@
+# Plugins
+
+The `dstack` plugin system allows extending `dstack` server functionality using external Python packages.
+
+!!! info "Experimental"
+    Plugins are currently an _experimental_ feature.
+    Backward compatibility is not guaranteed across releases.
+
+## Enable plugins
+
+To enable a plugin, list it under `plugins` in [`server/config.yml`](../reference/server/config.yml.md):
+
+<div editor-title="server/config.yml"> 
+
+```yaml
+plugins:
+  - my_dstack_plugin
+  - some_other_plugin
+projects:
+- name: main
+```
+
+</div>
+
+On the next server restart, you should see a log message indicating that the plugin is loaded.
+
+## Create plugins
+
+To create a plugin, create a Python package that implements a subclass of
+`dstack.plugins.Plugin` and exports this subclass as a "dstack.plugins" entry point.
+
+1. Init the plugin package:
+
+    <div class="termy">
+
+    ```shell
+    $ uv init --library
+    ```
+
+    </div>
+
+2. Define `ApplyPolicy` and `Plugin` subclasses:
+
+    <div editor-title="src/example_plugin/__init__.py"> 
+
+    ```python
+    from dstack.plugins import ApplyPolicy, Plugin, RunSpec, get_plugin_logger
+
+    logger = get_plugin_logger(__name__)
+
+    class ExamplePolicy(ApplyPolicy):
+        def on_run_apply(self, user: str, project: str, spec: RunSpec) -> RunSpec:
+            # ...
+            return spec
+    
+    class ExamplePlugin(Plugin):
+        
+        def get_apply_policies(self) -> list[ApplyPolicy]:
+            return [ExamplePolicy()]
+    ```
+
+    </div>
+
+3. Specify a "dstack.plugins" entry point in `pyproject.toml`:
+
+    <div editor-title="pyproject.toml"> 
+
+    ```toml
+    [project.entry-points."dstack.plugins"]
+    example_plugin = "example_plugin:ExamplePlugin"
+    ```
+
+    </div>
+
+Then you can install the plugin package into your Python environment and enable it via `server/config.yml`.
+
+??? info "Plugins in Docker"
+    If you deploy `dstack` using a Docker image you can add plugins either
+    by including them in your custom image built upon the `dstack` server image,
+    or by mounting installed plugins as volumes.
+
+## Apply policies
+
+Currently the only plugin functionality is apply policies.
+Apply policies allow modifying specs of runs, fleets, volumes, and gateways submitted on `dstack apply`.
+Subclass `dstack.plugins.ApplyPolicy` to implement them.
+
+Here's an example of how to enforce certain rules using apply policies:
+
+<div editor-title="src/example_plugin/__init__.py"> 
+
+```python
+class ExamplePolicy(ApplyPolicy):
+    def on_run_apply(self, user: str, project: str, spec: RunSpec) -> RunSpec:
+        # Forcing some limits
+        spec.configuration.max_price = 2.0
+        spec.configuration.max_duration = "1d"
+        # Setting some extra tags
+        if spec.configuration.tags is None:
+            spec.configuration.tags = {}
+        spec.configuration.tags |= {
+            "team": "my_team",
+        }
+        # Forbid something
+        if spec.configuration.privileged:
+            logger.warning("User %s tries to run privileged containers", user)
+            raise ValueError("Running privileged containers is forbidden")
+        # Set some service-specific properties
+        if isinstance(spec.configuration, Service):  
+            spec.configuration.https = True
+        return spec
+```
+
+</div>
+
+For more information on the plugin development, see the [plugin example](https://github.com/dstackai/dstack/tree/master/examples/plugins/example_plugin).

examples/plugins/example_plugin/.python-version

@@ -0,0 +1 @@
+3.11

examples/plugins/example_plugin/README.md

@@ -0,0 +1,52 @@
+## Overview
+
+This is a basic `dstack` plugin example.
+You can use it as a reference point when implementing new `dstack` plugins.
+
+## Steps
+
+1. Init the plugin package:
+
+    ```
+    uv init --library
+    ```
+
+2. Define `ApplyPolicy` and `Plugin` subclasses:
+
+    ```python
+    from dstack.plugins import ApplyPolicy, Plugin, RunSpec, get_plugin_logger
+
+
+    logger = get_plugin_logger(__name__)
+
+
+    class ExamplePolicy(ApplyPolicy):
+
+        def on_run_apply(self, user: str, project: str, spec: RunSpec) -> RunSpec:
+            # ...
+            return spec
+
+
+    class ExamplePlugin(Plugin):
+        
+        def get_apply_policies(self) -> list[ApplyPolicy]:
+            return [ExamplePolicy()]
+
+    ```
+
+3. Specify a "dstack.plugins" entry point in `pyproject.toml`:
+
+    ```toml
+    [project.entry-points."dstack.plugins"]
+    example_plugin = "example_plugin:ExamplePlugin"
+    ```
+
+4. Make sure to install the plugin and enable it in the `server/config.yml`:
+
+    ```yaml
+    plugins:
+    - example_plugin
+    projects:
+    - name: main
+      # ...
+    ```

examples/plugins/example_plugin/pyproject.toml

@@ -0,0 +1,17 @@
+[project]
+name = "example-plugin"
+version = "0.1.0"
+description = "A dstack plugin example"
+readme = "README.md"
+authors = [
+    { name = "Victor Skvortsov", email = "victor@dstack.ai" }
+]
+requires-python = ">=3.9"
+dependencies = []
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project.entry-points."dstack.plugins"]
+example_plugin = "example_plugin:ExamplePlugin"

examples/plugins/example_plugin/src/example_plugin/__init__.py

@@ -0,0 +1,34 @@
+from dstack.api import Service
+from dstack.plugins import ApplyPolicy, GatewaySpec, Plugin, RunSpec, get_plugin_logger
+
+logger = get_plugin_logger(__name__)
+
+
+class ExamplePolicy(ApplyPolicy):
+    def on_run_apply(self, user: str, project: str, spec: RunSpec) -> RunSpec:
+        # Forcing some limits
+        spec.configuration.max_price = 2.0
+        spec.configuration.max_duration = "1d"
+        # Setting some extra tags
+        if spec.configuration.tags is None:
+            spec.configuration.tags = {}
+        spec.configuration.tags |= {
+            "team": "my_team",
+        }
+        # Forbid something
+        if spec.configuration.privileged:
+            logger.warning("User %s tries to run privileged containers", user)
+            raise ValueError("Running privileged containers is forbidden")
+        # Set some service-specific properties
+        if isinstance(spec.configuration, Service):
+            spec.configuration.https = True
+        return spec
+
+    def on_gateway_apply(self, user: str, project: str, spec: GatewaySpec) -> GatewaySpec:
+        # Forbid creating new gateways altogether
+        raise ValueError("Creating gateways is forbidden")
+
+
+class ExamplePlugin(Plugin):
+    def get_apply_policies(self) -> list[ApplyPolicy]:
+        return [ExamplePolicy()]

examples/plugins/example_plugin/src/example_plugin/py.typed

@@ -206,6 +206,7 @@ nav:
           - Backends: docs/concepts/backends.md
           - Projects: docs/concepts/projects.md
           - Gateways: docs/concepts/gateways.md
+          - Plugins: docs/concepts/plugins.md
       - Guides:
           - Protips: docs/guides/protips.md
           - Server deployment: docs/guides/server-deployment.md

mkdocs.yml

@@ -125,6 +125,7 @@ server = [
     "python-json-logger>=3.1.0",
     "prometheus-client",
     "grpcio>=1.50",
+    "backports.entry-points-selectable",
 ]
 aws = [
     "boto3",

pyproject.toml

@@ -269,6 +269,8 @@ class FleetSpec(CoreModel):
     configuration_path: Optional[str] = None
     profile: Profile
     autocreated: bool = False
+    # merged_profile stores profile parameters merged from profile and configuration.
+    # Read profile parameters from merged_profile instead of profile directly.
     # TODO: make merged_profile a computed field after migrating to pydanticV2
     merged_profile: Annotated[Profile, Field(exclude=True)] = None
 

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

@@ -357,6 +357,8 @@ class RunSpec(CoreModel):
             description="The contents of the SSH public key that will be used to connect to the run."
         ),
     ]
+    # merged_profile stores profile parameters merged from profile and configuration.
+    # Read profile parameters from merged_profile instead of profile directly.
     # TODO: make merged_profile a computed field after migrating to pydanticV2
     merged_profile: Annotated[Profile, Field(exclude=True)] = None