Store ExperimentStatus on Experiment class (facebook#4738)

Summary:

# AOSC note

Land xdb schema changes first D88096914

# Summary

Store ExperimentStatus on Experiment class as `status` and introduce setters/getters, update SQA classes+encoder/decoder, and add tests.

Reviewed By: mgarrard

Differential Revision: D90089265

Author: CristianLara

ax/core/experiment.py

@@ -29,6 +29,7 @@
 from ax.core.base_trial import BaseTrial
 from ax.core.batch_trial import BatchTrial
 from ax.core.data import combine_data_rows_favoring_recent, Data
+from ax.core.experiment_status import ExperimentStatus
 from ax.core.generator_run import GeneratorRun
 from ax.core.metric import Metric, MetricFetchE, MetricFetchResult
 from ax.core.objective import MultiObjective
@@ -146,6 +147,7 @@ def __init__(
         self._optimization_config: OptimizationConfig | None = None
         self._tracking_metrics: dict[str, Metric] = {}
         self._time_created: datetime = datetime.now()
+        self._status: ExperimentStatus | None = None
         self._trials: dict[int, BaseTrial] = {}
         self._properties: dict[str, Any] = properties or {}
 
@@ -231,6 +233,27 @@ def experiment_type(self, experiment_type: str | None) -> None:
         """Set the type of the experiment."""
         self._experiment_type = experiment_type
 
+    @property
+    def status(self) -> ExperimentStatus | None:
+        """The current status of the experiment.
+
+        Status tracks the high-level lifecycle phase of the experiment:
+        DRAFT, INITIALIZATION, OPTIMIZATION, COMPLETED.
+
+        For new experiments, status defaults to DRAFT. For legacy experiments
+        that were created before the status field was added, status may be None.
+        """
+        return self._status
+
+    @status.setter
+    def status(self, status: ExperimentStatus | None) -> None:
+        """Set the status of the experiment.
+
+        Args:
+            status: The new status for the experiment.
+        """
+        self._status = status
+
     @property
     def search_space(self) -> SearchSpace:
         """The search space for this experiment.

ax/core/tests/test_experiment.py

@@ -17,6 +17,7 @@
 from ax.core.base_trial import BaseTrial, TrialStatus
 from ax.core.data import Data, sort_by_trial_index_and_arm_name
 from ax.core.evaluations_to_data import raw_evaluations_to_data
+from ax.core.experiment_status import ExperimentStatus
 from ax.core.map_metric import MapMetric
 from ax.core.metric import Metric
 from ax.core.objective import MultiObjective, Objective
@@ -1853,6 +1854,15 @@ def test_to_df_with_relativize(self) -> None:
                 "relativized value",
             )
 
+    def test_experiment_status_default(self) -> None:
+        """Test that new experiments have None status for backward compatibility."""
+        self.assertIsNone(self.experiment.status)
+
+    def test_experiment_status_property(self) -> None:
+        """Test the experiment status property getter and setter."""
+        self.experiment.status = ExperimentStatus.DRAFT
+        self.assertEqual(self.experiment.status, ExperimentStatus.DRAFT)
+
 
 class ExperimentWithMapDataTest(TestCase):
     def setUp(self) -> None:

ax/storage/sqa_store/decoder.py

@@ -409,6 +409,7 @@ def experiment_from_sqa(
             _cast_arm_parameters(sq, experiment.search_space)
             experiment._register_arm(sq)
         experiment._time_created = experiment_sqa.time_created
+        experiment._status = experiment_sqa.status
         experiment._experiment_type = self.get_enum_name(
             value=experiment_sqa.experiment_type, enum=self.config.experiment_type_enum
         )

ax/storage/sqa_store/encoder.py

@@ -247,6 +247,7 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment:
             status_quo_name=status_quo_name,
             status_quo_parameters=status_quo_parameters,
             time_created=experiment.time_created,
+            status=experiment.status,
             experiment_type=experiment_type,
             metrics=optimization_metrics + tracking_metrics,
             parameters=parameters,

ax/storage/sqa_store/sqa_classes.py

@@ -13,6 +13,7 @@
 from typing import Any
 
 from ax.core.evaluations_to_data import DataType
+from ax.core.experiment_status import ExperimentStatus
 from ax.core.parameter import ParameterType
 from ax.core.trial_status import TrialStatus
 from ax.core.types import (
@@ -375,6 +376,9 @@ class SQAExperiment(Base):
         JSONEncodedTextDict
     )
     time_created: Column[datetime] = Column(IntTimestamp, nullable=False)
+    status: Column[ExperimentStatus | None] = Column(
+        IntEnum(ExperimentStatus), nullable=True
+    )
     default_trial_type: Column[str | None] = Column(String(NAME_OR_TYPE_FIELD_LENGTH))
     default_data_type: Column[DataType] = Column(IntEnum(DataType), nullable=True)
     # pyre-fixme[8]: Incompatible attribute type [8]: Attribute

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -32,6 +32,7 @@
     TransferLearningMetadata,
 )
 from ax.core.experiment import Experiment
+from ax.core.experiment_status import ExperimentStatus
 from ax.core.generator_run import GeneratorRun
 from ax.core.metric import Metric
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
@@ -614,6 +615,27 @@ def test_saving_an_experiment_with_type_errors_with_missing_enum_value(
                 config=SQAConfig(experiment_type_enum=MockExperimentTypeEnum),
             )
 
+    def test_experiment_status_save_load(self) -> None:
+        """Test that experiment status is correctly saved and loaded."""
+        # Test None status (backward compatibility)
+        with self.subTest(status=None):
+            exp = get_experiment()
+            exp._name = "test_exp_status_none"
+            exp.status = None
+            save_experiment(exp)
+            loaded_exp = load_experiment(exp.name)
+            self.assertEqual(loaded_exp.status, None)
+
+        # Test all ExperimentStatus enum values
+        for status in ExperimentStatus:
+            with self.subTest(status=status):
+                exp = get_experiment()
+                exp._name = f"test_exp_status_{status.name.lower()}"
+                exp.status = status
+                save_experiment(exp)
+                loaded_exp = load_experiment(exp.name)
+                self.assertEqual(loaded_exp.status, status)
+
     def test_load_experiment_trials_in_batches(self) -> None:
         for _ in range(4):
             self.experiment.new_trial()