Add arm_name_prefix parameter to make_candidate_trials

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

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] = []

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:

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))

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",