feat: Add support for conditional search spaces in the tuner (#172)

# Summary
Allows a custom function to be specified as the definition for the
`Tuner` search space. This allows for more complex search space
definitions, e.g. branching, looping. See
https://docs.ray.io/en/latest/tune/examples/optuna_example.html#conditional-search-spaces
for more information on how Ray provides this (via Optuna).

# Changes
* New `space` attribute on the `OptunaSpec` algorithm definition.
* Implementation and tests.

Author: toby-coleman

docs/examples/tutorials/tuning-a-process.md

@@ -73,3 +73,31 @@ Plugboard's YAML config supports an optional `tune` section, allowing you to def
 3.  Parameters need to reference a type, so that Plugboard knows the type of parameter to build.
 
 Now run `plugboard process tune model-with-tuner.yaml` to execute the optimisation job from the CLI.
+
+## Advanced usage: complex search spaces
+
+Occasionally you may need to define more complex search spaces, which go beyond what can be defined with a simple parameter configuration. For example:
+
+*  Conditional parameters, e.g. where parameter `a` must be greater than parameter `b`; or
+*  Looping, e.g. building up a list of tunable parameters that is of variable length.
+
+These conditional search space functions are supported by Ray Tune and can be defined as described in the [Ray documentation](https://docs.ray.io/en/latest/tune/examples/optuna_example.html#conditional-search-spaces). To use such a function you will need to:
+
+1. Setup the [`Tuner`][plugboard.tune.Tuner], defining your parameters as usual;
+2. Write a custom function to define the search space, where each tunable parameter has a name of the form `"{component_name.field_or_arg_name}"`; then
+3. Supply your custom function to the `OptunaSpec` algorithm configuration.
+
+For example, the following search space makes the velocity depend on the angle:
+```python
+--8<-- "examples/tutorials/006_optimisation/hello_tuner.py:custom_search_space"
+```
+
+Then use this configuration to point the  tuner to the `custom_space` function.
+```yaml
+--8<-- "examples/tutorials/006_optimisation/model-with-tuner-custom.yaml"
+```
+
+1. We can reference a [`Process`][plugboard.process.Process] from another YAML file here to avoid repetition.
+2. Add the algorithm configuration here.
+
+Then run using `plugboard process tune model-with-tuner-custom.yaml`.

examples/tutorials/006_optimisation/hello_tuner.py

@@ -2,7 +2,7 @@
 
 # fmt: off
 import typing as _t
-
+from optuna import Trial
 from plugboard.component import Component, IOController as IO
 from plugboard.process import ProcessBuilder
 from plugboard.schemas import ComponentArgsDict, ProcessSpec, ProcessArgsSpec, ObjectiveSpec
@@ -64,6 +64,15 @@ async def step(self) -> None:
 # --8<-- [end:components]
 
 
+# --8<-- [start:custom_search_space]
+def custom_space(trial: Trial) -> dict[str, _t.Any] | None:
+    """Defines a custom search space for Optuna."""
+    angle = trial.suggest_int("trajectory.angle", 0, 90)
+    # Make velocity depend on angle
+    trial.suggest_int("trajectory.velocity", angle, 100)
+# --8<-- [end:custom_search_space]
+
+
 if __name__ == "__main__":
     # --8<-- [start:define_process]
     process_spec = ProcessSpec(

examples/tutorials/006_optimisation/model-with-tuner-custom.yaml

@@ -0,0 +1,30 @@
+plugboard:
+  process: model-with-tuner.yaml  # (1)!
+  tune:
+    args:
+      objective:
+        object_name: max-height
+        field_type: field
+        field_name: max_y
+      parameters:
+      - type: ray.tune.uniform
+        object_type: component
+        object_name: trajectory
+        field_type: arg
+        field_name: angle
+        lower: 0
+        upper: 90
+      - type: ray.tune.uniform
+        object_type: component
+        object_name: trajectory
+        field_type: arg
+        field_name: velocity
+        lower: 0
+        upper: 100
+      num_samples: 40
+      mode: max
+      max_concurrent: 4
+      algorithm:  # (2)!
+        type: ray.tune.search.optuna.OptunaSearch
+        space: hello_tuner.custom_space
+  

plugboard/cli/process/__init__.py

@@ -129,7 +129,8 @@ def tune(
 ) -> None:
     """Optimise a Plugboard process by adjusting its tunable parameters."""
     config_spec = _read_yaml(config)
-    tuner = _build_tuner(config_spec)
+    with add_sys_path(config.parent):
+        tuner = _build_tuner(config_spec)
 
     with Progress(
         SpinnerColumn("arrow3"),

plugboard/schemas/tune.py

@@ -17,11 +17,14 @@ class OptunaSpec(PlugboardBaseModel):
 
     Attributes:
         type: The algorithm type to load.
+        space: Optional; A function defining the search space. Use this to define more complex
+            search spaces that cannot be represented using the built-in parameter types.
         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"
+    space: str | None = None
     study_name: str | None = None
     storage: str | None = None
 

plugboard/tune/tune.py

@@ -55,25 +55,19 @@ def __init__(
             algorithm: Configuration for the underlying Optuna algorithm used for optimisation.
         """
         self._logger = DI.logger.resolve_sync().bind(cls=self.__class__.__name__)
-        # Check that objective and mode are lists of the same length if multiple objectives are used
+        # Validate and normalize objective/mode
         self._check_objective(objective, mode)
-        self._objective = objective if isinstance(objective, list) else [objective]
-        self._mode = [str(m) for m in mode] if isinstance(mode, list) else str(mode)
-        self._metric = (
-            [obj.full_name for obj in self._objective]
-            if len(self._objective) > 1
-            else self._objective[0].full_name
+        self._objective, self._mode, self._metric = self._normalize_objective_and_mode(
+            objective, mode
         )
+        self._custom_space = bool(algorithm and algorithm.space)
 
-        self._parameters_dict = {p.full_name: p for p in parameters}
-        self._parameters = dict(self._build_parameter(p) for p in parameters)
-        _algo = self._build_algorithm(algorithm)
-        if max_concurrent is not None:
-            _algo = ray.tune.search.ConcurrencyLimiter(_algo, max_concurrent)
-        self._config = ray.tune.TuneConfig(
-            num_samples=num_samples,
-            search_alg=_algo,
-        )
+        # Prepare parameters and search algorithm
+        self._parameters_dict, self._parameters = self._prepare_parameters(parameters)
+        searcher = self._init_search_algorithm(algorithm, max_concurrent)
+
+        # Configure Ray Tune
+        self._config = ray.tune.TuneConfig(num_samples=num_samples, search_alg=searcher)
         self._result_grid: _t.Optional[ray.tune.ResultGrid] = None
         self._logger.info("Tuner created")
 
@@ -105,33 +99,58 @@ def _build_algorithm(
     ) -> ray.tune.search.Searcher:
         if algorithm is None:
             self._logger.info("Using default Optuna search algorithm")
-            return ray.tune.search.optuna.OptunaSearch(metric=self._metric, mode=self._mode)
-        _algo_kwargs = {
-            **algorithm.model_dump(exclude={"type"}),
-            "mode": self._mode,
-            "metric": self._metric,
-        }
-
-        # Convert storage URI string to optuna storage object if needed
-        # TODO: Make this more general to support other algorithms, e.g. use a builder class
-        if "storage" in _algo_kwargs and isinstance(_algo_kwargs["storage"], str):
-            _algo_kwargs["storage"] = optuna.storages.RDBStorage(url=_algo_kwargs["storage"])
-            self._logger.info(
-                "Converted storage URI to Optuna RDBStorage object",
-                storage_uri=algorithm.storage,
-            )
+            return self._default_searcher()
 
-        algo_cls: _t.Optional[_t.Any] = locate(algorithm.type)
-        if not algo_cls or not issubclass(algo_cls, ray.tune.search.searcher.Searcher):
-            raise ValueError(f"Could not locate `Searcher` class {algorithm.type}")
+        algo_kwargs = self._build_algo_kwargs(algorithm)
+        algo_cls = self._get_algo_class(algorithm.type)
         self._logger.info(
             "Using custom search algorithm",
             algorithm=algorithm.type,
-            params={
-                k: v if k != "storage" else f"<{type(v).__name__}>" for k, v in _algo_kwargs.items()
-            },
+            params={k: self._mask_param_value(k, v) for k, v in algo_kwargs.items()},
         )
-        return algo_cls(**_algo_kwargs)
+        return algo_cls(**algo_kwargs)
+
+    def _default_searcher(self) -> "ray.tune.search.Searcher":
+        return ray.tune.search.optuna.OptunaSearch(metric=self._metric, mode=self._mode)
+
+    def _build_algo_kwargs(self, algorithm: OptunaSpec) -> dict[str, _t.Any]:
+        """Prepare keyword args for the searcher, normalising storage/space."""
+        kwargs = algorithm.model_dump(exclude={"type"})
+        kwargs["mode"] = self._mode
+        kwargs["metric"] = self._metric
+
+        storage = kwargs.get("storage")
+        if isinstance(storage, str):
+            kwargs["storage"] = optuna.storages.RDBStorage(url=storage)
+            self._logger.info(
+                "Converted storage URI to Optuna RDBStorage object",
+                storage_uri=storage,
+            )
+
+        space = kwargs.get("space")
+        if space is not None:
+            kwargs["space"] = self._resolve_space_fn(space)
+
+        return kwargs
+
+    def _resolve_space_fn(self, space: str) -> _t.Callable:
+        space_fn = locate(space)
+        if not space_fn or not isfunction(space_fn):  # pragma: no cover
+            raise ValueError(f"Could not locate search space function {space}")
+        return space_fn
+
+    def _get_algo_class(self, type_path: str) -> _t.Type[ray.tune.search.searcher.Searcher]:
+        algo_cls: _t.Optional[_t.Any] = locate(type_path)
+        if not algo_cls or not issubclass(
+            algo_cls, ray.tune.search.searcher.Searcher
+        ):  # pragma: no cover
+            raise ValueError(f"Could not locate `Searcher` class {type_path}")
+        return algo_cls
+
+    def _mask_param_value(self, k: str, v: _t.Any) -> _t.Any:
+        if k == "storage" or (k == "space" and isfunction(v)):
+            return f"<{type(v).__name__}>"
+        return v
 
     def _build_parameter(
         self, parameter: ParameterSpec
@@ -203,12 +222,16 @@ def run(self, spec: ProcessSpec) -> ray.tune.Result | list[ray.tune.Result]:
             ),
         )
 
-        self._logger.info("Setting Tuner with parameters", params=list(self._parameters.keys()))
-        _tune = ray.tune.Tuner(
-            trainable_with_resources,
-            param_space=self._parameters,
-            tune_config=self._config,
-        )
+        tuner_kwargs: dict[str, _t.Any] = {
+            "tune_config": self._config,
+        }
+        if not self._custom_space:
+            self._logger.info("Setting Tuner with parameters", params=list(self._parameters.keys()))
+            tuner_kwargs["param_space"] = self._parameters
+        else:
+            self._logger.info("Setting Tuner with custom search space")
+
+        _tune = ray.tune.Tuner(trainable_with_resources, **tuner_kwargs)
         self._logger.info("Starting Tuner")
         self._result_grid = _tune.fit()
         self._logger.info("Tuner finished")
@@ -230,6 +253,10 @@ def fn(config: dict[str, _t.Any]) -> dict[str, _t.Any]:  # pragma: no cover
                 ComponentRegistry.add(cls, key=key)
 
             for name, value in config.items():
+                if name not in self._parameters_dict:
+                    # Custom search spaces may include intermediate parameters not in the Tuner
+                    self._logger.warning("Parameter from config not found in Tuner", param=name)
+                    continue
                 self._override_parameter(spec, self._parameters_dict[name], value)
 
             process = ProcessBuilder.build(spec)
@@ -253,3 +280,35 @@ def fn(config: dict[str, _t.Any]) -> dict[str, _t.Any]:  # pragma: no cover
             return result
 
         return fn
+
+    def _normalize_objective_and_mode(
+        self,
+        objective: ObjectiveSpec | list[ObjectiveSpec],
+        mode: Direction | list[Direction],
+    ) -> tuple[list[ObjectiveSpec], str | list[str], str | list[str]]:
+        """Return normalized objectives, modes and metric name(s)."""
+        objectives = objective if isinstance(objective, list) else [objective]
+        modes = [str(m) for m in mode] if isinstance(mode, list) else str(mode)
+        metric = (
+            [obj.full_name for obj in objectives]
+            if len(objectives) > 1
+            else objectives[0].full_name
+        )
+        return objectives, modes, metric
+
+    def _prepare_parameters(
+        self, parameters: list[ParameterSpec]
+    ) -> tuple[dict[str, ParameterSpec], dict[str, "ray.tune.search.sample.Sampler"]]:
+        """Build parameter lookup dict and Ray Tune parameter space."""
+        params_dict = {p.full_name: p for p in parameters}
+        params_space = dict(self._build_parameter(p) for p in parameters)
+        return params_dict, params_space
+
+    def _init_search_algorithm(
+        self, algorithm: _t.Optional[OptunaSpec], max_concurrent: _t.Optional[int]
+    ) -> "ray.tune.search.Searcher":
+        """Create the search algorithm and apply concurrency limits if requested."""
+        algo = self._build_algorithm(algorithm)
+        if max_concurrent is not None:
+            algo = ray.tune.search.ConcurrencyLimiter(algo, max_concurrent)
+        return algo

tests/data/dynamic-param-process.yaml

@@ -0,0 +1,21 @@
+plugboard:
+  process:
+    args:
+      components:
+      - type: tests.integration.test_process_with_components_run.A
+        args:
+          name: "a"
+          iters: 10
+      - type: tests.integration.test_tuner.DynamicListComponent
+        args:
+          name: "d"
+          list_param: [1.0, 2.0, 3.0]
+      - type: tests.integration.test_process_with_components_run.C
+        args:
+          name: "c"
+          path: "./c.txt"
+      connectors:
+      - source: "a.out_1"
+        target: "d.in_1"
+      - source: "d.out_1"
+        target: "c.in_1"