diff --git a/ax/core/base_trial.py b/ax/core/base_trial.py index 4d3589399b9..e3c5af698d4 100644 --- a/ax/core/base_trial.py +++ b/ax/core/base_trial.py @@ -83,11 +83,17 @@ def __init__( trial_type: str | None = None, ttl_seconds: int | None = None, index: int | None = None, + arm_name_prefix: str | None = None, ) -> None: """Initialize trial. Args: experiment: The experiment this trial belongs to. + arm_name_prefix: If provided, arms added to this trial will be + named ``{arm_name_prefix}_{trial_index}_{i}`` instead of the + default ``{trial_index}_{i}`` scheme. Retaining the trial and + arm indices keeps arm names unique even when the same prefix is + reused across multiple trials. """ self._experiment = experiment if ttl_seconds is not None and ttl_seconds <= 0: @@ -123,6 +129,7 @@ def __init__( # Counter to maintain how many arms have been named by this BatchTrial self._num_arms_created = 0 + self._arm_name_prefix: str | None = arm_name_prefix # NOTE: Please do not store any data related to trial deployment or data- # fetching in properties. It is intended to only store properties related @@ -476,6 +483,8 @@ def _check_existing_and_name_arm(self, arm: Arm) -> None: def _get_default_name(self, arm_index: int | None = None) -> str: if arm_index is None: arm_index = self._num_arms_created + if self._arm_name_prefix is not None: + return f"{self._arm_name_prefix}_{self.index}_{arm_index}" return f"{self.index}_{arm_index}" @property diff --git a/ax/core/batch_trial.py b/ax/core/batch_trial.py index 43d05839910..35f4d094d4f 100644 --- a/ax/core/batch_trial.py +++ b/ax/core/batch_trial.py @@ -115,6 +115,11 @@ class BatchTrial(BaseTrial): This should generally not be specified, as in the index will be automatically determined based on the number of existing trials. This is only used for the purpose of loading from storage. + arm_name_prefix: If provided, arms added to this trial will be + named ``{arm_name_prefix}_{trial_index}_{i}`` instead of the + default ``{trial_index}_{i}`` scheme. Retaining the trial and arm + indices keeps arm names unique even when the same prefix is reused + across multiple trials. """ def __init__( @@ -126,12 +131,14 @@ def __init__( should_add_status_quo_arm: bool | None = False, ttl_seconds: int | None = None, index: int | None = None, + arm_name_prefix: str | None = None, ) -> None: super().__init__( experiment=experiment, trial_type=trial_type, ttl_seconds=ttl_seconds, index=index, + arm_name_prefix=arm_name_prefix, ) self._arms_by_name: dict[str, Arm] = {} self._generator_runs: list[GeneratorRun] = [] diff --git a/ax/core/experiment.py b/ax/core/experiment.py index bc232d201a9..7f48f7fa2ea 100644 --- a/ax/core/experiment.py +++ b/ax/core/experiment.py @@ -1515,6 +1515,7 @@ def new_batch_trial( should_add_status_quo_arm: bool | None = False, trial_type: str | None = None, ttl_seconds: int | None = None, + arm_name_prefix: str | None = None, ) -> BatchTrial: """Create a new batch trial associated with this experiment. @@ -1537,6 +1538,11 @@ def new_batch_trial( 'dead' trials, for which the evaluation process might have crashed etc., and which should be considered stale after their 'time to live' has passed. + arm_name_prefix: If provided, arms added to this trial will be + named ``{arm_name_prefix}_{trial_index}_{i}`` instead of the + default ``{trial_index}_{i}`` scheme. Retaining the trial and + arm indices keeps arm names unique even when the same prefix is + reused across multiple trials. """ if ttl_seconds is not None: self._trials_have_ttl = True @@ -1547,6 +1553,7 @@ def new_batch_trial( generator_runs=generator_runs, should_add_status_quo_arm=should_add_status_quo_arm, ttl_seconds=ttl_seconds, + arm_name_prefix=arm_name_prefix, ) def get_batch_trial(self, trial_index: int) -> BatchTrial: diff --git a/ax/core/tests/test_batch_trial.py b/ax/core/tests/test_batch_trial.py index a83ee53dd41..b08589683b3 100644 --- a/ax/core/tests/test_batch_trial.py +++ b/ax/core/tests/test_batch_trial.py @@ -658,3 +658,43 @@ def test_attach_batch_trial_data(self) -> None: self.assertEqual(2.0, data[(arm1_name, "m2")]["mean"]) self.assertEqual(3.0, data[(arm2_name, "m1")]["mean"]) self.assertEqual(4.0, data[(arm2_name, "m2")]["mean"]) + + def test_arm_name_prefix(self) -> None: + """Arms added with arm_name_prefix are named + ``{prefix}_{trial_index}_{arm_index}``, retaining the trial and arm + indices so names stay unique when the prefix is reused.""" + experiment = get_experiment() + batch = experiment.new_batch_trial(arm_name_prefix="my_prefix") + arms = get_arms() + # Add unnamed arms via a generator run + gr = GeneratorRun(arms=arms[1:3], weights=[1.0, 1.0]) + batch.add_generator_run(gr) + arm_names = [arm.name for arm in batch.arms] + self.assertEqual( + arm_names, + [f"my_prefix_{batch.index}_0", f"my_prefix_{batch.index}_1"], + ) + + def test_arm_name_prefix_preserves_existing_names(self) -> None: + """Arms that already have names keep them even with arm_name_prefix.""" + experiment = get_experiment() + batch = experiment.new_batch_trial(arm_name_prefix="pfx") + named_arm = Arm( + parameters={"w": 0.5, "x": 5, "y": "bar", "z": False}, name="custom_name" + ) + gr = GeneratorRun(arms=[named_arm], weights=[1.0]) + batch.add_generator_run(gr) + arm_names = [arm.name for arm in batch.arms] + self.assertIn("custom_name", arm_names) + + def test_arm_name_prefix_default_uses_trial_index(self) -> None: + """Without arm_name_prefix the default {trial_index}_{i} scheme is used.""" + experiment = get_experiment() + batch = experiment.new_batch_trial() + gr = GeneratorRun( + arms=[Arm(parameters={"w": 0.1, "x": 1, "y": "foo", "z": True})], + weights=[1.0], + ) + batch.add_generator_run(gr) + arm_names = [arm.name for arm in batch.arms] + self.assertTrue(all(name.startswith(f"{batch.index}_") for name in arm_names)) diff --git a/ax/storage/sqa_store/utils.py b/ax/storage/sqa_store/utils.py index 6599af4a760..ac1278c3467 100644 --- a/ax/storage/sqa_store/utils.py +++ b/ax/storage/sqa_store/utils.py @@ -24,6 +24,7 @@ # of that data, so the recursion could fail somewhere) # * baseline_workflow_inputs (plain JSON dict with no _db_id fields) COPY_DB_IDS_ATTRS_TO_SKIP = { + "_arm_name_prefix", "baseline_workflow_inputs", "_best_arm_predictions", "_adapter_kwargs",