Ensure new arm names do not match a different name on the experiment (f…

…acebook#2732)

Summary:
Pull Request resolved: facebook#2732

>
In Experiment._name_and_store_arm_if_not_exists (pointer)`, ensure that no other arm with a different signature (but same name) exists on the experiment before adding.

This change checks for signature conflict by arm name in Experiment._name_and_store_arm_if_not_exists.

A "replace" flag is also added, to maintain the existing functionality where a conflicting arm by the same name existed.
=> In "clone_with", desired functionality is to replace existing status quo arm with new status quo arm.

Differential Revision: D62130255

ax/core/batch_trial.py

@@ -366,7 +366,9 @@ def set_status_quo_with_weight(
                 status_quo.parameters, raise_error=True
             )
             self.experiment._name_and_store_arm_if_not_exists(
-                arm=status_quo, proposed_name="status_quo_" + str(self.index)
+                arm=status_quo,
+                proposed_name="status_quo_" + str(self.index),
+                replace=True,
             )
         self._status_quo = status_quo.clone() if status_quo is not None else None
         self._status_quo_weight_override = weight

ax/core/experiment.py

@@ -1350,7 +1350,9 @@ def warm_start_from_old_experiment(
 
         return copied_trials
 
-    def _name_and_store_arm_if_not_exists(self, arm: Arm, proposed_name: str) -> None:
+    def _name_and_store_arm_if_not_exists(
+        self, arm: Arm, proposed_name: str, replace: bool = False
+    ) -> None:
         """Tries to lookup arm with same signature, otherwise names and stores it.
 
         - Looks up if arm already exists on experiment
@@ -1360,8 +1362,22 @@ def _name_and_store_arm_if_not_exists(self, arm: Arm, proposed_name: str) -> Non
         Args:
             arm: The arm object to name.
             proposed_name: The name to assign if it doesn't have one already.
+            replace: If true, override arm w/ same name and different signature.
+                If false, raise an error if this conflict occurs.
         """
 
+        if not replace:
+            # Check for signature conflict by arm name/proposed name
+            search_name = proposed_name if not arm.has_name else arm.name
+            if (
+                search_name in self.arms_by_name
+                and arm.signature != self.arms_by_name[search_name].signature
+            ):
+                raise ValueError(
+                    f"Arm with name {search_name} already exists on experiment "
+                    f"with different signature."
+                )
+
         # If arm is identical to an existing arm, return that
         # so that the names match.
         if arm.signature in self.arms_by_signature:

ax/core/tests/test_experiment.py

@@ -1541,3 +1541,43 @@ class TestAuxiliaryExperimentPurpose(AuxiliaryExperimentPurpose):
                     ],
                 },
             )
+
+    def test_name_and_store_arm_if_not_exists_same_name_different_signature(
+        self,
+    ) -> None:
+        experiment = self.experiment
+        shared_name = "shared_name"
+
+        arm_1 = Arm({"x1": -1.0, "x2": 1.0}, name=shared_name)
+        arm_2 = Arm({"x1": -1.7, "x2": 0.2, "x3": 1})
+        self.assertNotEqual(arm_1.signature, arm_2.signature)
+
+        experiment._register_arm(arm=arm_1)
+        with self.assertRaisesRegex(
+            ValueError,
+            f"Arm with name {shared_name} already exists on experiment "
+            f"with different signature.",
+        ):
+            experiment._name_and_store_arm_if_not_exists(
+                arm=arm_2, proposed_name=shared_name
+            )
+
+    def test_name_and_store_arm_if_not_exists_same_proposed_name_different_signature(
+        self,
+    ) -> None:
+        experiment = self.experiment
+        shared_name = "shared_name"
+
+        arm_1 = Arm({"x1": -1.0, "x2": 1.0}, name=shared_name)
+        arm_2 = Arm({"x1": -1.7, "x2": 0.2, "x3": 1}, name=shared_name)
+        self.assertNotEqual(arm_1.signature, arm_2.signature)
+
+        experiment._register_arm(arm=arm_1)
+        with self.assertRaisesRegex(
+            ValueError,
+            f"Arm with name {shared_name} already exists on experiment "
+            f"with different signature.",
+        ):
+            experiment._name_and_store_arm_if_not_exists(
+                arm=arm_2, proposed_name="different proposed name"
+            )