Add experiment_type to ExperimentConfig (#3201)

Summary: Pull Request resolved: #3201 Reviewed By: saitcakmak Differential Revision: D67454875 fbshipit-source-id: a8ef6638159cdac707cab8043f1b4605cfea1c89
facebook · Dec 19, 2024 · b538e50 · b538e50
1 parent 782b188
commit b538e50
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 0 deletions.
diff --git a/ax/preview/api/configs.py b/ax/preview/api/configs.py
@@ -78,6 +78,7 @@ class ExperimentConfig:
     parameter_constraints: list[str] = field(default_factory=list)
 
     description: str | None = None
+    experiment_type: str | None = None
     owner: str | None = None
 
 

diff --git a/ax/preview/api/utils/instantiation/from_config.py b/ax/preview/api/utils/instantiation/from_config.py
@@ -134,6 +134,7 @@ def experiment_from_config(config: ExperimentConfig) -> Experiment:
         search_space=search_space,
         name=config.name,
         description=config.description,
+        experiment_type=config.experiment_type,
         properties={"owners": [config.owner]},
         default_data_type=DataType.MAP_DATA,
     )

diff --git a/ax/preview/api/utils/instantiation/tests/test_from_config.py b/ax/preview/api/utils/instantiation/tests/test_from_config.py
@@ -216,6 +216,7 @@ def test_experiment_from_config(self) -> None:
             parameters=[float_parameter, int_parameter, choice_parameter],
             parameter_constraints=["int_param <= float_param"],
             description="test description",
+            experiment_type="TEST",
             owner="miles",
         )
 
@@ -252,6 +253,7 @@ def test_experiment_from_config(self) -> None:
                 ),
                 name="test_experiment",
                 description="test description",
+                experiment_type="TEST",
                 properties={"owners": ["miles"]},
                 default_data_type=DataType.MAP_DATA,
             ),