feat: Optimisation

plugboard/cli/process/__init__.py

@@ -14,6 +14,7 @@
 from plugboard.diagram import MermaidDiagram
 from plugboard.process import Process, ProcessBuilder
 from plugboard.schemas import ConfigSpec
+from plugboard.tune import Tuner
 from plugboard.utils import add_sys_path
 
 
@@ -38,11 +39,39 @@
     return process
 
 
+def _build_tuner(config: ConfigSpec) -> Tuner:
+    tune_config = config.plugboard.tune
+    if tune_config is None:
+        stderr.print("[red]No tuning configuration found in YAML file[/red]")
+        raise typer.Exit(1)
+    return Tuner(
+        objective=tune_config.args.objective,
+        parameters=tune_config.args.parameters,
+        num_samples=tune_config.args.num_samples,
+        mode=tune_config.args.mode,
+        max_concurrent=tune_config.args.max_concurrent,
+        algorithm=tune_config.args.algorithm,
+    )
+
+
 async def _run_process(process: Process) -> None:
     async with process:
         await process.run()
 
 
+def _run_tune(tuner: Tuner, config_spec: ConfigSpec) -> None:
+    process_spec = config_spec.plugboard.process
+    # Build the process to import the component types
+    _build_process(config_spec)
+    result = tuner.run(spec=process_spec)
+    print("[green]Best parameters found:[/green]")
+    if isinstance(result, list):
+        for r in result:
+            print(f"Config: {r.config} - Metrics: {r.metrics}")
+    else:
+        print(f"Config: {result.config} - Metrics: {result.metrics}")
+
+
 @app.command()
 def run(
     config: Annotated[
@@ -73,6 +102,35 @@
         progress.update(task, description=f"[green]Process complete[/green]")
 
 
+@app.command()
+def tune(
+    config: Annotated[
+        Path,
+        typer.Argument(
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            writable=False,
+            readable=True,
+            resolve_path=True,
+            help="Path to the YAML configuration file.",
+        ),
+    ],
+) -> None:
+    """Optimise a Plugboard process by adjusting its tunable parameters."""
+    config_spec = _read_yaml(config)
+    tuner = _build_tuner(config_spec)
+
+    with Progress(
+        SpinnerColumn("arrow3"),
+        TextColumn("[progress.description]{task.description}"),
+    ) as progress:
+        task = progress.add_task(f"Running tune job from {config}", total=None)
+        with add_sys_path(config.parent):
+            _run_tune(tuner, config_spec)
+        progress.update(task, description=f"[green]Tune job complete[/green]")
+
+
 @app.command()
 def diagram(
     config: Annotated[

plugboard/connector/ray_channel.py

@@ -11,7 +11,7 @@
 
 try:
     import ray
-except ImportError:
+except ImportError:  # pragma: no cover
     pass
 
 _AsyncioChannelActor = build_actor_wrapper(AsyncioChannel)

plugboard/process/ray_process.py

@@ -8,12 +8,12 @@
 from plugboard.connector import Connector
 from plugboard.process.process import Process
 from plugboard.state import RayStateBackend, StateBackend
-from plugboard.utils import build_actor_wrapper, depends_on_optional, gather_except
+from plugboard.utils import build_actor_wrapper, depends_on_optional, gather_except, gen_rand_str
 
 
 try:
     import ray
-except ImportError:
+except ImportError:  # pragma: no cover
     pass
 
 
@@ -40,6 +40,8 @@ def __init__(
             parameters: Optional; Parameters for the `Process`.
             state: Optional; `StateBackend` for the `Process`.
         """
+        # TODO: Replace with a namespace based on the job ID or similar
+        self._namespace = f"plugboard-{gen_rand_str(16)}"
         self._component_actors = {
             # Recreate components on remote actors
             c.id: self._create_component_actor(c)
@@ -58,7 +60,9 @@ def _create_component_actor(self, component: Component) -> _t.Any:
         name = component.id
         args = component.export()["args"]
         actor_cls = build_actor_wrapper(component.__class__)
-        return ray.remote(num_cpus=0, name=name)(actor_cls).remote(**args)  # type: ignore
+        return ray.remote(num_cpus=0, name=name, namespace=self._namespace)(  # type: ignore
+            actor_cls
+        ).remote(**args)
 
     async def _update_component_attributes(self) -> None:
         """Updates attributes on local components from remote actors."""

plugboard/schemas/__init__.py

@@ -20,6 +20,15 @@
 from .io import IODirection
 from .process import ProcessArgsDict, ProcessArgsSpec, ProcessSpec
 from .state import StateBackendArgsDict, StateBackendArgsSpec, StateBackendSpec
+from .tune import (
+    Direction,
+    ObjectiveSpec,
+    OptunaSpec,
+    ParameterSpec,
+    TuneArgsDict,
+    TuneArgsSpec,
+    TuneSpec,
+)
 
 
 __all__ = [
@@ -33,13 +42,20 @@
     "ConnectorMode",
     "ConnectorSocket",
     "ConnectorSpec",
+    "Direction",
     "Entity",
     "IODirection",
+    "ObjectiveSpec",
+    "OptunaSpec",
+    "ParameterSpec",
     "ProcessConfigSpec",
     "ProcessSpec",
     "ProcessArgsDict",
     "ProcessArgsSpec",
     "StateBackendSpec",
     "StateBackendArgsDict",
     "StateBackendArgsSpec",
+    "TuneArgsDict",
+    "TuneArgsSpec",
+    "TuneSpec",
 ]

plugboard/schemas/_common.py

@@ -8,4 +8,6 @@
 class PlugboardBaseModel(BaseModel, ABC):
     """Custom base model for Plugboard schemas."""
 
-    model_config = ConfigDict(extra="forbid", populate_by_name=True, use_enum_values=True)
+    model_config = ConfigDict(
+        extra="forbid", populate_by_name=True, use_enum_values=True, validate_assignment=True
+    )

plugboard/schemas/config.py

@@ -1,17 +1,42 @@
 """Provides top-level `ConfigSpec` class for Plugboard configuration."""
 
+from pathlib import Path
+import typing as _t
+
+import msgspec
+from pydantic import field_validator
+
 from plugboard.schemas._common import PlugboardBaseModel
 from .process import ProcessSpec
+from .tune import TuneSpec
 
 
 class ProcessConfigSpec(PlugboardBaseModel):
     """A `ProcessSpec` within a Plugboard configuration.
 
     Attributes:
-        process: A `ProcessSpec` that specifies the process.
+        process: A `ProcessSpec` that specifies the process, or a path to a YAML file containing the
+            process specification.
+        tune: Optional; A `TuneSpec` that specifies an optimisation configuration.
     """
 
     process: ProcessSpec
+    tune: TuneSpec | None = None
+
+    @field_validator("process", mode="before")
+    @classmethod
+    def _auto_load_process(cls, value: _t.Any) -> _t.Any:
+        """Automatically loads the process specification from a YAML file if a path is provided."""
+        if isinstance(value, str) and Path(value).exists():
+            with open(value, "rb") as f:
+                other_config = msgspec.yaml.decode(f.read())
+            try:
+                return other_config["plugboard"]["process"]
+            except KeyError:
+                raise ValueError(
+                    "The provided YAML file does not contain a Plugboard process specification."
+                )
+        return value
 
 
 class ConfigSpec(PlugboardBaseModel):

plugboard/schemas/tune.py

@@ -0,0 +1,175 @@
+"""Provides the `TuneSpec` class for configuring optimisation jobs."""
+
+from abc import ABC
+import typing as _t
+
+from pydantic import Field, PositiveInt, model_validator
+
+from plugboard.schemas._common import PlugboardBaseModel
+
+
+class OptunaSpec(PlugboardBaseModel):
+    """Specification for the Optuna configuration.
+
+    See: https://docs.ray.io/en/latest/tune/api/doc/ray.tune.search.optuna.OptunaSearch.html
+    and https://optuna.readthedocs.io/en/stable/reference/index.html for more information on the
+    Optuna configuration.
+
+    Attributes:
+        type: The algorithm type to load.
+        study_name: Optional; The name of the study.
+        storage: Optional; The storage URI to save the optimisation results to.
+    """
+
+    type: _t.Literal["ray.tune.search.optuna.OptunaSearch"] = "ray.tune.search.optuna.OptunaSearch"
+    study_name: str | None = None
+    storage: str | None = None
+
+
+class BaseFieldSpec(PlugboardBaseModel, ABC):
+    """Base class for specifying fields within a Plugboard [`Process`][plugboard.process.Process].
+
+    These fields may be used as adjustable parameter inputs or as an optimisation objective.
+
+    Attributes:
+        object_type: The type of object on which the field is defined. Defaults to "component".
+        object_name: The name of the object on which the field is defined.
+        field_type: The type of field. This can be "arg", "initial_value", or "field".
+        field_name: The name of the field.
+    """
+
+    object_type: _t.Literal["component"] = Field("component", exclude=True)
+    object_name: str = Field(..., exclude=True)
+    field_type: _t.Literal["arg", "initial_value", "field"] = Field(..., exclude=True)
+    field_name: str = Field(..., exclude=True)
+
+    @property
+    def full_name(self) -> str:
+        """Returns the full name of the field, including the object name and field name."""
+        return f"{self.object_name}.{self.field_name}"
+
+
+class ObjectiveSpec(BaseFieldSpec):
+    """Specification for an objective field."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def _fill_defaults(
+        cls, data: dict[str, _t.Any] | list[dict[str, _t.Any]]
+    ) -> dict[str, _t.Any] | list[dict[str, _t.Any]]:
+        if isinstance(data, list):
+            # If the data is a list, skip because it is already a list of objectives
+            return data
+        if "field_type" not in data:
+            data["field_type"] = "field"
+        if data["field_type"] != "field":  # pragma: no cover
+            raise ValueError("The field type must be 'field' for an objective specification.")
+        return data
+
+
+class FloatParameterSpec(BaseFieldSpec):
+    """Specification for a uniform float parameter.
+
+    See: https://docs.ray.io/en/latest/tune/api/search_space.html.
+
+    Attributes:
+        type: The type of the parameter.
+        lower: The lower bound of the parameter.
+        upper: The upper bound of the parameter.
+    """
+
+    type: _t.Literal["ray.tune.uniform"] = "ray.tune.uniform"
+    lower: float
+    upper: float
+
+
+class IntParameterSpec(BaseFieldSpec):
+    """Specification for a uniform integer parameter.
+
+    See: https://docs.ray.io/en/latest/tune/api/search_space.html.
+
+    Attributes:
+        type: The type of the parameter.
+        lower: The lower bound of the parameter.
+        upper: The upper bound of the parameter.
+    """
+
+    type: _t.Literal["ray.tune.randint"] = "ray.tune.randint"
+    lower: int
+    upper: int
+
+
+class CategoricalParameterSpec(BaseFieldSpec):
+    """Specification for a categorical parameter.
+
+    See: https://docs.ray.io/en/latest/tune/api/search_space.html.
+
+    Attributes:
+        type: The type of the parameter.
+        categories: The categories of the parameter.
+    """
+
+    type: _t.Literal["ray.tune.choice"] = "ray.tune.choice"
+    categories: list[_t.Any]
+
+
+ParameterSpec = _t.Union[
+    FloatParameterSpec,
+    IntParameterSpec,
+    CategoricalParameterSpec,
+]
+
+Direction = _t.Literal["min", "max"]
+
+
+class TuneArgsDict(_t.TypedDict):
+    """`TypedDict` of the [`Tuner`][plugboard.tune.Tuner] constructor arguments."""
+
+    objective: str | list[str]
+    parameters: list[ParameterSpec]
+    num_samples: int
+    mode: _t.NotRequired[Direction | list[Direction]]
+    max_concurrent: _t.NotRequired[int | None]
+    algorithm: OptunaSpec
+
+
+class TuneArgsSpec(PlugboardBaseModel):
+    """Specification of the arguments for the `Tune` class.
+
+    Attributes:
+        objective: The location of the objective(s) to optimise for in the `Process`.
+        parameters: The parameters to optimise over.
+        num_samples: The number of samples to draw during the optimisation.
+        mode: The mode of optimisation. For multi-objective optimisation, this should be a list
+            containing a direction for each objective.
+        max_concurrent: The maximum number of concurrent trials.
+        algorithm: The algorithm to use for the optimisation.
+    """
+
+    objective: ObjectiveSpec | list[ObjectiveSpec]
+    parameters: list[ParameterSpec] = Field(min_length=1)
+    num_samples: PositiveInt
+    mode: Direction | list[Direction] = "max"
+    max_concurrent: PositiveInt | None = None
+    algorithm: _t.Union[OptunaSpec] = Field(OptunaSpec(), discriminator="type")
+
+    @model_validator(mode="after")
+    def _validate_model(self: _t.Self) -> _t.Self:
+        if isinstance(self.mode, list):
+            if not isinstance(self.objective, list):
+                raise ValueError(
+                    "In multi-objective optimisation, both `mode` and `objective` must be lists."
+                )
+            if len(self.mode) != len(self.objective):
+                raise ValueError("The length of `mode` must match the length of `objective`.")
+        return self
+
+
+class TuneSpec(PlugboardBaseModel):
+    """Configuration for an optimisation job.
+
+    Attributes:
+        args: The arguments for the `Tune` job.
+    """
+
+    args: TuneArgsSpec

plugboard/tune/__init__.py

@@ -0,0 +1,6 @@
+"""Tune submodule for configuring optimisation jobs."""
+
+from plugboard.tune.tune import Tuner
+
+
+__all__ = ["Tuner"]