From d85c0ea974a1ff4f713da27097f4695f2b604366 Mon Sep 17 00:00:00 2001
From: Matthew Grange <mgrange@meta.com>
Date: Tue, 3 Sep 2024 07:35:32 -0700
Subject: [PATCH] Ensure new arm names do not match a different name on the
 experiment

Summary:
>
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

Differential Revision: D62130255
---
 ax/core/experiment.py            | 11 +++++++++
 ax/core/tests/test_experiment.py | 40 ++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/ax/core/experiment.py b/ax/core/experiment.py
index 0c48c6e5c94..6b3705d0cec 100644
--- a/ax/core/experiment.py
+++ b/ax/core/experiment.py
@@ -1362,6 +1362,17 @@ def _name_and_store_arm_if_not_exists(self, arm: Arm, proposed_name: str) -> Non
             proposed_name: The name to assign if it doesn't have one already.
         """
 
+        # 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:
diff --git a/ax/core/tests/test_experiment.py b/ax/core/tests/test_experiment.py
index cd285fe9260..f653f7e89ee 100644
--- a/ax/core/tests/test_experiment.py
+++ b/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"
+            )