From 19c5868381e2f8cb5bcb43aa2ec16a22abfa6cd1 Mon Sep 17 00:00:00 2001 From: josephinecowley Date: Thu, 20 Feb 2025 19:25:10 -0800 Subject: [PATCH 1/7] Allow passing raise_parameter_error from ax_client.attach_trial through to search_space.check_membership to allow using non-validated parameters in trials --- ax/core/experiment.py | 5 ++++- ax/core/search_space.py | 6 ++++-- ax/service/ax_client.py | 2 ++ 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/ax/core/experiment.py b/ax/core/experiment.py index 25b522fbc6d..a2cc53e22a2 100644 --- a/ax/core/experiment.py +++ b/ax/core/experiment.py @@ -1535,6 +1535,7 @@ def attach_trial( ttl_seconds: int | None = None, run_metadata: dict[str, Any] | None = None, optimize_for_power: bool = False, + raise_parameter_error: bool = True, ) -> tuple[dict[str, TParameterization], int]: """Attach a new trial with the given parameterization to the experiment. @@ -1564,7 +1565,9 @@ def attach_trial( # Validate search space membership for all parameterizations for parameterization in parameterizations: - self.search_space.validate_membership(parameters=parameterization) + self.search_space.validate_membership( + parameters=parameterization, raise_error=raise_parameter_error + ) # Validate number of arm names if any arm names are provided. named_arms = False diff --git a/ax/core/search_space.py b/ax/core/search_space.py index b75cf380596..a998779034e 100644 --- a/ax/core/search_space.py +++ b/ax/core/search_space.py @@ -378,8 +378,10 @@ def _validate_parameter_constraints( f"`{parameter_name}` does not exist in search space." ) - def validate_membership(self, parameters: TParameterization) -> None: - self.check_membership(parameterization=parameters, raise_error=True) + def validate_membership( + self, parameters: TParameterization, raise_error: bool = True + ) -> None: + self.check_membership(parameterization=parameters, raise_error=raise_error) # `check_membership` uses int and float interchangeably, which we don't # want here. for p_name, parameter in self.parameters.items(): diff --git a/ax/service/ax_client.py b/ax/service/ax_client.py index 01a9af29c62..cbcea382e33 100644 --- a/ax/service/ax_client.py +++ b/ax/service/ax_client.py @@ -881,6 +881,7 @@ def attach_trial( ttl_seconds: int | None = None, run_metadata: dict[str, Any] | None = None, arm_name: str | None = None, + raise_parameter_error: bool = True, ) -> tuple[TParameterization, int]: """Attach a new trial with the given parameterization to the experiment. @@ -899,6 +900,7 @@ def attach_trial( arm_names=[arm_name] if arm_name else None, ttl_seconds=ttl_seconds, run_metadata=run_metadata, + raise_parameter_error=raise_parameter_error, ) self._save_or_update_trial_in_db_if_possible( experiment=self.experiment, From a17c0f8b8ce529196526e2f0d5d4de3727742697 Mon Sep 17 00:00:00 2001 From: josephinecowley Date: Thu, 20 Feb 2025 19:35:38 -0800 Subject: [PATCH 2/7] Created test for raise_parameter_error --- ax/service/tests/test_ax_client.py | 33 ++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/ax/service/tests/test_ax_client.py b/ax/service/tests/test_ax_client.py index 7988563b6b8..4760ab53601 100644 --- a/ax/service/tests/test_ax_client.py +++ b/ax/service/tests/test_ax_client.py @@ -1940,6 +1940,39 @@ def test_attach_trial_and_get_trial_parameters(self) -> None: ) # No trial #10 in experiment. with self.assertRaisesRegex(UnsupportedError, ".* is of type"): ax_client.attach_trial({"x": 1, "y": 2}) + + def test_attach_trial_invalid_parameters(self) -> None: + ax_client = AxClient() + ax_client.create_experiment( + parameters=[ + {"name": "x", "type": "range", "bounds": [-5.0, 10.0]}, + {"name": "y", "type": "range", "bounds": [0.0, 15.0]}, + ], + ) + # Test attaching trial with parameter outside bounds fails by default + with self.assertRaisesRegex( + ValueError, "20.0 is not a valid value for parameter RangeParameter\(name='x'" + ): + ax_client.attach_trial( + parameters={"x": 20.0, "y": 1.0} + ) + + # Test attaching trial with parameter outside bounds fails with raise_parameter_error=True + with self.assertRaisesRegex( + ValueError, "20.0 is not a valid value for parameter RangeParameter\(name='x'" + ): + ax_client.attach_trial( + parameters={"x": 20.0, "y": 1.0}, + raise_parameter_error=True + ) + + # Test attaching trial with parameter outside bounds succeeds with raise_parameter_error=False + _, idx = ax_client.attach_trial( + parameters={"x": 20.0, "y": 1.0}, + raise_parameter_error=False + ) + ax_client.complete_trial(trial_index=idx, raw_data=5) + self.assertEqual(ax_client.get_trial_parameters(trial_index=idx), {"x": 20.0, "y": 1.0}) def test_attach_trial_ttl_seconds(self) -> None: ax_client = AxClient() From d38ad617a958121d3ac13a4b244e2fa301eacf8c Mon Sep 17 00:00:00 2001 From: josephinecowley Date: Thu, 20 Feb 2025 19:45:02 -0800 Subject: [PATCH 3/7] reformat test_ax_client.py with ruff --- ax/service/tests/test_ax_client.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/ax/service/tests/test_ax_client.py b/ax/service/tests/test_ax_client.py index 4760ab53601..088e2ecec05 100644 --- a/ax/service/tests/test_ax_client.py +++ b/ax/service/tests/test_ax_client.py @@ -1940,7 +1940,7 @@ def test_attach_trial_and_get_trial_parameters(self) -> None: ) # No trial #10 in experiment. with self.assertRaisesRegex(UnsupportedError, ".* is of type"): ax_client.attach_trial({"x": 1, "y": 2}) - + def test_attach_trial_invalid_parameters(self) -> None: ax_client = AxClient() ax_client.create_experiment( @@ -1951,28 +1951,28 @@ def test_attach_trial_invalid_parameters(self) -> None: ) # Test attaching trial with parameter outside bounds fails by default with self.assertRaisesRegex( - ValueError, "20.0 is not a valid value for parameter RangeParameter\(name='x'" + ValueError, + "20.0 is not a valid value for parameter RangeParameter\(name='x'", ): - ax_client.attach_trial( - parameters={"x": 20.0, "y": 1.0} - ) - + ax_client.attach_trial(parameters={"x": 20.0, "y": 1.0}) + # Test attaching trial with parameter outside bounds fails with raise_parameter_error=True with self.assertRaisesRegex( - ValueError, "20.0 is not a valid value for parameter RangeParameter\(name='x'" + ValueError, + "20.0 is not a valid value for parameter RangeParameter\(name='x'", ): ax_client.attach_trial( - parameters={"x": 20.0, "y": 1.0}, - raise_parameter_error=True + parameters={"x": 20.0, "y": 1.0}, raise_parameter_error=True ) # Test attaching trial with parameter outside bounds succeeds with raise_parameter_error=False _, idx = ax_client.attach_trial( - parameters={"x": 20.0, "y": 1.0}, - raise_parameter_error=False + parameters={"x": 20.0, "y": 1.0}, raise_parameter_error=False ) ax_client.complete_trial(trial_index=idx, raw_data=5) - self.assertEqual(ax_client.get_trial_parameters(trial_index=idx), {"x": 20.0, "y": 1.0}) + self.assertEqual( + ax_client.get_trial_parameters(trial_index=idx), {"x": 20.0, "y": 1.0} + ) def test_attach_trial_ttl_seconds(self) -> None: ax_client = AxClient() From ac78135a8545644a09c34b1e842c29dfc18148b8 Mon Sep 17 00:00:00 2001 From: josephinecowley Date: Thu, 20 Feb 2025 19:49:14 -0800 Subject: [PATCH 4/7] add docstrings for raise_parameter_error --- ax/core/experiment.py | 3 +++ ax/service/ax_client.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/ax/core/experiment.py b/ax/core/experiment.py index a2cc53e22a2..a6f5f6afda8 100644 --- a/ax/core/experiment.py +++ b/ax/core/experiment.py @@ -1553,6 +1553,9 @@ def attach_trial( trial such that the experiment's power to detect effects of certain size is as high as possible. Refer to documentation of `BatchTrial.set_status_quo_and_optimize_power` for more detail. + raise_parameter_error: If True, raise an error if validating membership + of the parameterization in the search space fails. If False, do not + raise an error. Returns: Tuple of arm name to parameterization dict, and trial index from diff --git a/ax/service/ax_client.py b/ax/service/ax_client.py index cbcea382e33..b7a63b0cb11 100644 --- a/ax/service/ax_client.py +++ b/ax/service/ax_client.py @@ -890,6 +890,9 @@ def attach_trial( ttl_seconds: If specified, will consider the trial failed after this many seconds. Used to detect dead trials that were not marked failed properly. + raise_parameter_error: If True, raise an error if validating membership + of the parameterization in the search space fails. If False, do not + raise an error. Returns: Tuple of parameterization and trial index from newly created trial. From df16ec985fb51cfc797bcea6e92e9ca216ea7d0d Mon Sep 17 00:00:00 2001 From: josephinecowley Date: Thu, 20 Feb 2025 19:54:14 -0800 Subject: [PATCH 5/7] add partial docstring for validate_membership --- ax/core/search_space.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/ax/core/search_space.py b/ax/core/search_space.py index a998779034e..2f1824db704 100644 --- a/ax/core/search_space.py +++ b/ax/core/search_space.py @@ -381,6 +381,13 @@ def _validate_parameter_constraints( def validate_membership( self, parameters: TParameterization, raise_error: bool = True ) -> None: + """Validates if a parameterization belongs in the search space. + + Args: + raise_error: If True, raise an error if validating membership + of the parameterization in the search space fails. If False, do not + raise an error. + """ self.check_membership(parameterization=parameters, raise_error=raise_error) # `check_membership` uses int and float interchangeably, which we don't # want here. From 87616ceda26a210c462d5dcae032914327b0b74a Mon Sep 17 00:00:00 2001 From: josephinecowley Date: Thu, 20 Feb 2025 19:56:38 -0800 Subject: [PATCH 6/7] fix flake8 warning for test_ax_client --- ax/service/tests/test_ax_client.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ax/service/tests/test_ax_client.py b/ax/service/tests/test_ax_client.py index 088e2ecec05..04fa5ab0fc5 100644 --- a/ax/service/tests/test_ax_client.py +++ b/ax/service/tests/test_ax_client.py @@ -1949,23 +1949,23 @@ def test_attach_trial_invalid_parameters(self) -> None: {"name": "y", "type": "range", "bounds": [0.0, 15.0]}, ], ) - # Test attaching trial with parameter outside bounds fails by default + # Test parameter outside bounds fails by default with self.assertRaisesRegex( ValueError, - "20.0 is not a valid value for parameter RangeParameter\(name='x'", + "20.0 is not a valid value for parameter RangeParameter\\(name='x'", ): ax_client.attach_trial(parameters={"x": 20.0, "y": 1.0}) - # Test attaching trial with parameter outside bounds fails with raise_parameter_error=True + # Test parameter outside bounds fails with raise_parameter_error=True with self.assertRaisesRegex( ValueError, - "20.0 is not a valid value for parameter RangeParameter\(name='x'", + "20.0 is not a valid value for parameter RangeParameter\\(name='x'", ): ax_client.attach_trial( parameters={"x": 20.0, "y": 1.0}, raise_parameter_error=True ) - # Test attaching trial with parameter outside bounds succeeds with raise_parameter_error=False + # Test parameter outside bounds succeeds with raise_parameter_error=False _, idx = ax_client.attach_trial( parameters={"x": 20.0, "y": 1.0}, raise_parameter_error=False ) From f2c7341dd1d2b8ff6ed0572ac7dee94cbdc8f920 Mon Sep 17 00:00:00 2001 From: josephinecowley Date: Thu, 20 Feb 2025 20:00:54 -0800 Subject: [PATCH 7/7] escape . in regex --- ax/service/tests/test_ax_client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ax/service/tests/test_ax_client.py b/ax/service/tests/test_ax_client.py index 04fa5ab0fc5..26c30bdd8ce 100644 --- a/ax/service/tests/test_ax_client.py +++ b/ax/service/tests/test_ax_client.py @@ -1952,14 +1952,14 @@ def test_attach_trial_invalid_parameters(self) -> None: # Test parameter outside bounds fails by default with self.assertRaisesRegex( ValueError, - "20.0 is not a valid value for parameter RangeParameter\\(name='x'", + "20\\.0 is not a valid value for parameter RangeParameter\\(name='x'", ): ax_client.attach_trial(parameters={"x": 20.0, "y": 1.0}) # Test parameter outside bounds fails with raise_parameter_error=True with self.assertRaisesRegex( ValueError, - "20.0 is not a valid value for parameter RangeParameter\\(name='x'", + "20\\.0 is not a valid value for parameter RangeParameter\\(name='x'", ): ax_client.attach_trial( parameters={"x": 20.0, "y": 1.0}, raise_parameter_error=True