From 43d34a3680f8d141cf7cf3c84116bedd25c656d3 Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Sun, 29 Jun 2025 19:35:11 +0100 Subject: [PATCH 01/12] Ad ConstraintError --- plugboard/exceptions/__init__.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/plugboard/exceptions/__init__.py b/plugboard/exceptions/__init__.py index 75fccd74..5dd488da 100644 --- a/plugboard/exceptions/__init__.py +++ b/plugboard/exceptions/__init__.py @@ -89,3 +89,9 @@ class ValidationError(Exception): """Raised when an invalid `Process` or `Component` is encountered.""" pass + + +class ConstraintError(Exception): + """Raised when a constraint is violated.""" + + pass From ce2b9ddc605e713acfc7ea70d599f861d2b2a337 Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 9 Jul 2025 20:40:19 +0100 Subject: [PATCH 02/12] Implement early stopping --- plugboard/tune/tune.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/plugboard/tune/tune.py b/plugboard/tune/tune.py index 6e0fd090..729eb514 100644 --- a/plugboard/tune/tune.py +++ b/plugboard/tune/tune.py @@ -7,6 +7,7 @@ import ray.tune.search.optuna from plugboard.component.component import Component, ComponentRegistry +from plugboard.exceptions import ConstraintError from plugboard.process import Process, ProcessBuilder from plugboard.schemas import ( Direction, @@ -172,6 +173,7 @@ def run(self, spec: ProcessSpec) -> ray.tune.Result | list[ray.tune.Result]: _tune = ray.tune.Tuner( trainable_with_resources, param_space=self._parameters, + run_config=ray.tune.RunConfig(stop={"constraint_hit": True}), tune_config=self._config, ) self._logger.info("Starting Tuner") @@ -197,8 +199,17 @@ def fn(config: dict[str, _t.Any]) -> _t.Any: # pragma: no-cover for name, value in config.items(): self._override_parameter(spec, self._parameters_dict[name], value) + ray.tune.report({"constraint_hit": False}) process = ProcessBuilder.build(spec) - run_coro_sync(self._run_process(process)) + try: + run_coro_sync(self._run_process(process)) + except ConstraintError as e: + # If a constraint is violated, we report it and return the objectives as None + self._logger.warning( + "Constraint violated during optimisation", + constraint_error=str(e), + ) + ray.tune.report({"constraint_hit": True}) return {obj.full_name: self._get_objective(process, obj) for obj in self._objective} From e721788fab57b59e889cc6fbb89c5108c2eee752 Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Tue, 15 Jul 2025 20:27:12 +0100 Subject: [PATCH 03/12] Implement constraints using +/- inf objective --- plugboard/tune/tune.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/plugboard/tune/tune.py b/plugboard/tune/tune.py index 729eb514..67c9d651 100644 --- a/plugboard/tune/tune.py +++ b/plugboard/tune/tune.py @@ -1,6 +1,7 @@ """Provides `Tuner` class for optimising Plugboard processes.""" from inspect import isfunction +import math from pydoc import locate import typing as _t @@ -191,7 +192,7 @@ def run(self, spec: ProcessSpec) -> ray.tune.Result | list[ray.tune.Result]: def _build_objective( self, component_classes: dict[str, type[Component]], spec: ProcessSpec ) -> _t.Callable: - def fn(config: dict[str, _t.Any]) -> _t.Any: # pragma: no-cover + def fn(config: dict[str, _t.Any]) -> dict[str, _t.Any]: # pragma: no-cover # Recreate the ComponentRegistry in the Ray worker for key, cls in component_classes.items(): ComponentRegistry.add(cls, key=key) @@ -199,17 +200,19 @@ def fn(config: dict[str, _t.Any]) -> _t.Any: # pragma: no-cover for name, value in config.items(): self._override_parameter(spec, self._parameters_dict[name], value) - ray.tune.report({"constraint_hit": False}) process = ProcessBuilder.build(spec) try: run_coro_sync(self._run_process(process)) except ConstraintError as e: - # If a constraint is violated, we report it and return the objectives as None + modes = self._mode if isinstance(self._mode, list) else [self._mode] self._logger.warning( - "Constraint violated during optimisation", + "Constraint violated during optimisation, stopping early", constraint_error=str(e), ) - ray.tune.report({"constraint_hit": True}) + return { + obj.full_name: math.inf if mode == "min" else -math.inf + for obj, mode in zip(self._objective, modes) + } return {obj.full_name: self._get_objective(process, obj) for obj in self._objective} From 8999c61f207ca312b34627efd4125d01a59934ea Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 19:45:26 +0100 Subject: [PATCH 04/12] Add test for constrained optimisation --- tests/integration/test_tuner.py | 51 +++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tests/integration/test_tuner.py b/tests/integration/test_tuner.py index 9dd3189e..d477a546 100644 --- a/tests/integration/test_tuner.py +++ b/tests/integration/test_tuner.py @@ -3,12 +3,23 @@ import msgspec import pytest +from plugboard.exceptions import ConstraintError from plugboard.schemas import ConfigSpec, ConnectorBuilderSpec, ObjectiveSpec from plugboard.schemas.tune import CategoricalParameterSpec, IntParameterSpec, OptunaSpec from plugboard.tune import Tuner from tests.integration.test_process_with_components_run import A, B, C # noqa: F401 +class ConstrainedB(B): + """Component with a constraint on the output value.""" + + async def step(self) -> None: + """Override step to apply a constraint.""" + if self.in_1 > 10: + raise ConstraintError("Input must not be greater than 10") + await super().step() + + @pytest.fixture def config() -> dict: """Loads the YAML config.""" @@ -120,3 +131,43 @@ async def test_multi_objective_tune(config: dict, ray_ctx: None) -> None: assert best_result[1].config["b.factor"] == -1 assert best_result[0].metrics["c.in_1"] == 1 assert best_result[1].metrics["b.out_1"] == -1 + + +@pytest.mark.tuner +@pytest.mark.asyncio +async def test_tune_with_constraint(config: dict, ray_ctx: None) -> None: + """Tests running of optimisation jobs with a constraint.""" + spec = ConfigSpec.model_validate(config) + process_spec = spec.plugboard.process + # Replace component B with a constrained version + process_spec.args.components[1].type = "tests.integration.test_tuner.ConstrainedB" + tuner = Tuner( + objective=ObjectiveSpec( + object_type="component", + object_name="c", + field_type="field", + field_name="in_1", + ), + parameters=[ + IntParameterSpec( + object_type="component", + object_name="a", + field_type="arg", + field_name="iters", + lower=5, + upper=15, + ) + ], + num_samples=20, + mode="max", + max_concurrent=2, + algorithm=OptunaSpec(), + ) + best_result = tuner.run( + spec=process_spec, + ) + result = tuner.result_grid + # There must be no failed trials with output less than or equal to 10 + assert all(t.metrics["c.in_1"] <= 10 for t in result if not t.error) + # Optimum must be less than or equal to 10 + assert best_result.metrics["c.in_1"] <= 10 From ab69b5d994df1192e3d9d9d434f85e7da7389f6b Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 19:45:51 +0100 Subject: [PATCH 05/12] Increase tolerance --- tests/integration/test_tuner.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_tuner.py b/tests/integration/test_tuner.py index d477a546..ec4218ab 100644 --- a/tests/integration/test_tuner.py +++ b/tests/integration/test_tuner.py @@ -11,7 +11,7 @@ class ConstrainedB(B): - """Component with a constraint on the output value.""" + """Component with a constraint.""" async def step(self) -> None: """Override step to apply a constraint.""" @@ -70,11 +70,11 @@ async def test_tune(config: dict, mode: str, process_type: str, ray_ctx: None) - assert not [t for t in result if t.error] # Correct optimimum must be found (within tolerance) if mode == "min": - assert best_result.config["a.iters"] <= tuner._parameters["a.iters"].lower + 1 - assert best_result.metrics["c.in_1"] == best_result.config["a.iters"] - 1 + assert best_result.config["a.iters"] <= tuner._parameters["a.iters"].lower + 2 + assert best_result.metrics["c.in_1"] == best_result.config["a.iters"] - 2 else: - assert best_result.config["a.iters"] >= tuner._parameters["a.iters"].upper - 1 - assert best_result.metrics["c.in_1"] == best_result.config["a.iters"] - 1 + assert best_result.config["a.iters"] >= tuner._parameters["a.iters"].upper - 2 + assert best_result.metrics["c.in_1"] == best_result.config["a.iters"] - 2 @pytest.mark.tuner @@ -158,7 +158,7 @@ async def test_tune_with_constraint(config: dict, ray_ctx: None) -> None: upper=15, ) ], - num_samples=20, + num_samples=12, mode="max", max_concurrent=2, algorithm=OptunaSpec(), From 368f41db55114af5439409f3c4c5980063d333f1 Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 19:52:42 +0100 Subject: [PATCH 06/12] Fixup --- tests/integration/test_tuner.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_tuner.py b/tests/integration/test_tuner.py index ec4218ab..38c5e510 100644 --- a/tests/integration/test_tuner.py +++ b/tests/integration/test_tuner.py @@ -71,10 +71,10 @@ async def test_tune(config: dict, mode: str, process_type: str, ray_ctx: None) - # Correct optimimum must be found (within tolerance) if mode == "min": assert best_result.config["a.iters"] <= tuner._parameters["a.iters"].lower + 2 - assert best_result.metrics["c.in_1"] == best_result.config["a.iters"] - 2 + assert best_result.metrics["c.in_1"] == best_result.config["a.iters"] - 1 else: assert best_result.config["a.iters"] >= tuner._parameters["a.iters"].upper - 2 - assert best_result.metrics["c.in_1"] == best_result.config["a.iters"] - 2 + assert best_result.metrics["c.in_1"] == best_result.config["a.iters"] - 1 @pytest.mark.tuner From 3eb253de4637987ad37b96adeb4f0d32821ff14b Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 20:01:41 +0100 Subject: [PATCH 07/12] Use constraint group --- plugboard/tune/tune.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/plugboard/tune/tune.py b/plugboard/tune/tune.py index 67c9d651..e35cdd2f 100644 --- a/plugboard/tune/tune.py +++ b/plugboard/tune/tune.py @@ -201,19 +201,23 @@ def fn(config: dict[str, _t.Any]) -> dict[str, _t.Any]: # pragma: no-cover self._override_parameter(spec, self._parameters_dict[name], value) process = ProcessBuilder.build(spec) + result = {} try: run_coro_sync(self._run_process(process)) - except ConstraintError as e: + result = { + obj.full_name: self._get_objective(process, obj) for obj in self._objective + } + except* ConstraintError as e: modes = self._mode if isinstance(self._mode, list) else [self._mode] self._logger.warning( "Constraint violated during optimisation, stopping early", constraint_error=str(e), ) - return { + result = { obj.full_name: math.inf if mode == "min" else -math.inf for obj, mode in zip(self._objective, modes) } - return {obj.full_name: self._get_objective(process, obj) for obj in self._objective} + return result return fn From adbefdaa171fcb247b0642b3ede810e2f5c797ba Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 20:01:49 +0100 Subject: [PATCH 08/12] Update test --- tests/integration/test_tuner.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_tuner.py b/tests/integration/test_tuner.py index 38c5e510..d77dfdc4 100644 --- a/tests/integration/test_tuner.py +++ b/tests/integration/test_tuner.py @@ -1,5 +1,7 @@ """Integration tests for the `Tuner` class.""" +import math + import msgspec import pytest @@ -167,7 +169,11 @@ async def test_tune_with_constraint(config: dict, ray_ctx: None) -> None: spec=process_spec, ) result = tuner.result_grid - # There must be no failed trials with output less than or equal to 10 - assert all(t.metrics["c.in_1"] <= 10 for t in result if not t.error) + # There must be no failed trials + assert not [t for t in result if t.error] + # Constraint must be respected + assert all(t.metrics["c.in_1"] <= 10 for t in result) # Optimum must be less than or equal to 10 assert best_result.metrics["c.in_1"] <= 10 + # If a.iters is greater than 11, the constraint will be violated + assert all(t.metrics["c.in_1"] == -math.inf for t in result if t.config["a.iters"] > 11) From 02123e818e4831313d26b6115f25af92e472fc50 Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 20:18:08 +0100 Subject: [PATCH 09/12] Better error handling for multi-objective --- plugboard/tune/tune.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/plugboard/tune/tune.py b/plugboard/tune/tune.py index e35cdd2f..c2b1a8bc 100644 --- a/plugboard/tune/tune.py +++ b/plugboard/tune/tune.py @@ -54,6 +54,8 @@ 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 + 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 = ( @@ -81,6 +83,22 @@ def result_grid(self) -> ray.tune.ResultGrid: raise ValueError("No result grid available. Run the optimisation job first.") return self._result_grid + @classmethod + def _check_objective( + cls, objective: ObjectiveSpec | list[ObjectiveSpec], mode: Direction | list[Direction] + ) -> None: + """Check that the objective and mode are valid.""" + if isinstance(objective, list): + if not isinstance(mode, list): + raise ValueError("If using multiple objectives, `mode` must also be a list.") + if len(objective) != len(mode): + raise ValueError( + "If using multiple objectives, `mode` and `objective` must be the same length." + ) + else: + if isinstance(mode, list): + raise ValueError("If using a single objective, `mode` must not be a list.") + def _build_algorithm( self, algorithm: _t.Optional[OptunaSpec] = None ) -> ray.tune.search.Searcher: From 65c8f88c9448fe32345a69378ca5fded5786a961 Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 20:19:36 +0100 Subject: [PATCH 10/12] Remove unused config from early experiment --- plugboard/tune/tune.py | 1 - 1 file changed, 1 deletion(-) diff --git a/plugboard/tune/tune.py b/plugboard/tune/tune.py index c2b1a8bc..c1ebcda4 100644 --- a/plugboard/tune/tune.py +++ b/plugboard/tune/tune.py @@ -192,7 +192,6 @@ def run(self, spec: ProcessSpec) -> ray.tune.Result | list[ray.tune.Result]: _tune = ray.tune.Tuner( trainable_with_resources, param_space=self._parameters, - run_config=ray.tune.RunConfig(stop={"constraint_hit": True}), tune_config=self._config, ) self._logger.info("Starting Tuner") From 817c62e3e40c845b9c2f261fb76a0378ebd39c70 Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 20:29:12 +0100 Subject: [PATCH 11/12] Update docs on constraints --- docs/examples/tutorials/tuning-a-process.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/examples/tutorials/tuning-a-process.md b/docs/examples/tutorials/tuning-a-process.md index 071f76d9..4c8d3651 100644 --- a/docs/examples/tutorials/tuning-a-process.md +++ b/docs/examples/tutorials/tuning-a-process.md @@ -58,6 +58,9 @@ Running this code will execute an optimisation job and print out information on !!! tip Since [Optuna](https://optuna.org/) is used under the hood, you can configure the optional `algorithm` argument on the `Tuner` with additional configuration defined in [`OptunaSpec`][plugboard.schemas.OptunaSpec]. For example, the [`storage`](https://optuna.readthedocs.io/en/stable/reference/storages.html) argument allows you to save the optimisation results to a database or SQLite file. You can then use a tool like [Optuna Dashboard](https://optuna-dashboard.readthedocs.io/en/stable/getting-started.html) to study the optimisation output in more detail. +!!! tip + You can impose arbitary constraints on variables within a `Process`. In your `step` method you can raise a [`ConstraintError`][plugboard.exceptions.ConstraintError] to indicate to the `Tuner` that a constraint has been breached. This will cause the trial to be stopped, and the optimisation will continue trying to find parameters that don't cause the constraint violation. + ## Using YAML config Plugboard's YAML config supports an optional `tune` section, allowing you to define optimisation jobs alongside your model configuration: From c7327b400b2a6e254b9983be22917b87eb97daea Mon Sep 17 00:00:00 2001 From: Toby Coleman Date: Wed, 23 Jul 2025 20:37:27 +0100 Subject: [PATCH 12/12] Fixup --- plugboard/tune/tune.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugboard/tune/tune.py b/plugboard/tune/tune.py index c1ebcda4..aa678fc4 100644 --- a/plugboard/tune/tune.py +++ b/plugboard/tune/tune.py @@ -209,7 +209,7 @@ def run(self, spec: ProcessSpec) -> ray.tune.Result | list[ray.tune.Result]: def _build_objective( self, component_classes: dict[str, type[Component]], spec: ProcessSpec ) -> _t.Callable: - def fn(config: dict[str, _t.Any]) -> dict[str, _t.Any]: # pragma: no-cover + def fn(config: dict[str, _t.Any]) -> dict[str, _t.Any]: # pragma: no cover # Recreate the ComponentRegistry in the Ray worker for key, cls in component_classes.items(): ComponentRegistry.add(cls, key=key)