Add Keys.LILO_LABELING trial type for Language-in-the-Loop labeling trials

ax/core/experiment.py

@@ -1880,11 +1880,9 @@ def supports_trial_type(self, trial_type: str | None) -> bool:
         """
         return (
             trial_type is None
-            # We temporarily allow "short run" and "long run" trial
-            # types in single-type experiments during development of
-            # a new ``GenerationStrategy`` that needs them.
             or trial_type == Keys.SHORT_RUN
             or trial_type == Keys.LONG_RUN
+            or trial_type == Keys.LILO_LABELING
         )
 
     def attach_trial(

ax/core/tests/test_experiment.py

@@ -207,6 +207,18 @@ def test_basic_batch_creation(self) -> None:
             new_exp = get_experiment()
             new_exp._attach_trial(batch)
 
+    def test_supports_trial_type(self) -> None:
+        exp = get_experiment()
+        self.assertTrue(exp.supports_trial_type(None))
+        self.assertTrue(exp.supports_trial_type(Keys.SHORT_RUN))
+        self.assertTrue(exp.supports_trial_type(Keys.LONG_RUN))
+        self.assertTrue(exp.supports_trial_type(Keys.LILO_LABELING))
+        self.assertFalse(exp.supports_trial_type("unsupported_type"))
+
+        # Verify LILO_LABELING trial type works with new_batch_trial
+        batch = exp.new_batch_trial(trial_type=Keys.LILO_LABELING)
+        self.assertEqual(batch.trial_type, Keys.LILO_LABELING)
+
     def test_repr(self) -> None:
         self.assertEqual("Experiment(test)", str(self.experiment))
 

ax/generation_strategy/generation_node.py

@@ -179,12 +179,14 @@ def __init__(
             )
         if len(generator_specs) > 1 and best_model_selector is None:
             raise UserInputError(MISSING_MODEL_SELECTOR_MESSAGE)
-        if trial_type is not None and (
-            trial_type != Keys.SHORT_RUN and trial_type != Keys.LONG_RUN
+        if trial_type is not None and trial_type not in (
+            Keys.SHORT_RUN,
+            Keys.LONG_RUN,
+            Keys.LILO_LABELING,
         ):
             raise NotImplementedError(
-                f"Trial type must be either {Keys.SHORT_RUN} or {Keys.LONG_RUN},"
-                f" got {trial_type}."
+                f"Trial type must be one of {Keys.SHORT_RUN}, {Keys.LONG_RUN},"
+                f" or {Keys.LILO_LABELING}, got {trial_type}."
             )
         # If possible, assign `_generator_spec_to_gen_from` right away, for use in
         # `__repr__`

ax/generation_strategy/tests/test_generation_node.py

@@ -134,7 +134,7 @@ def test_input_constructor_none(self) -> None:
         self.assertEqual(self.sobol_generation_node.input_constructors, {})
 
     def test_incorrect_trial_type(self) -> None:
-        with self.assertRaisesRegex(NotImplementedError, "Trial type must be either"):
+        with self.assertRaisesRegex(NotImplementedError, "Trial type must be one of"):
             GenerationNode(
                 name="test",
                 generator_specs=[self.sobol_generator_spec],
@@ -147,12 +147,18 @@ def test_init_with_trial_type(self) -> None:
             generator_specs=[self.sobol_generator_spec],
             trial_type=Keys.LONG_RUN,
         )
+        node_lilo = GenerationNode(
+            name="test",
+            generator_specs=[self.sobol_generator_spec],
+            trial_type=Keys.LILO_LABELING,
+        )
         node_default = GenerationNode(
             name="test",
             generator_specs=[self.sobol_generator_spec],
         )
         self.assertEqual(self.node_short._trial_type, Keys.SHORT_RUN)
         self.assertEqual(node_long._trial_type, Keys.LONG_RUN)
+        self.assertEqual(node_lilo._trial_type, Keys.LILO_LABELING)
         self.assertIsNone(node_default._trial_type)
 
     def test_input_constructor(self) -> None:

ax/utils/common/constants.py

@@ -64,6 +64,7 @@ class Keys(StrEnum):
     FRAC_RANDOM = "frac_random"
     FULL_PARAMETERIZATION = "full_parameterization"
     IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF = "immutable_search_space_and_opt_config"
+    LILO_LABELING = "lilo_labeling"
     LLM_MESSAGES = "llm_messages"
     LONG_RUN = "long_run"
     MAXIMIZE = "maximize"