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.

I've gone ahead and also defined methods in `storage/sqa_store/` to enable efficiently updating the property. Not currently used but will be used later in the stack by the orchestrator to set this status.

Differential Revision:
D90089265

Privacy Context Container: L1307644

Author: CristianLara

ax/core/experiment.py

@@ -34,6 +34,7 @@
     Data,
     sort_by_trial_index_and_arm_name,
 )
+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
@@ -151,6 +152,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 {}
 
@@ -236,6 +238,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.
@@ -620,6 +643,18 @@ def remove_tracking_metric(self, metric_name: str) -> Experiment:
         del self._tracking_metrics[metric_name]
         return self
 
+    def set_status(self, status: ExperimentStatus) -> Experiment:
+        """Set the status of the experiment.
+
+        Args:
+            status: The new status of the experiment.
+
+        Returns:
+            The experiment instance.
+        """
+        self._status = status
+        return self
+
     @property
     def metrics(self) -> dict[str, Metric]:
         """The metrics attached to the 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
@@ -1841,6 +1842,19 @@ 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."""
+        exp = get_experiment()
+        self.assertIsNone(exp.status)
+
+    def test_experiment_status_property(self) -> None:
+        """Test the experiment status property getter and setter."""
+        exp = get_experiment()
+
+        # Set and get status
+        exp.status = ExperimentStatus.DRAFT
+        self.assertEqual(exp.status, ExperimentStatus.DRAFT)
+
 
 class ExperimentWithMapDataTest(TestCase):
     def setUp(self) -> None:

ax/storage/sqa_store/decoder.py

@@ -384,6 +384,7 @@ def experiment_from_sqa(
             sq = none_throws(experiment.status_quo)
             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

@@ -246,6 +246,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/save.py

@@ -536,6 +536,42 @@ def update_properties_on_experiment(
         )
 
 
+def update_experiment_status(
+    experiment_with_updated_status: Experiment,
+    config: SQAConfig | None = None,
+) -> None:
+    """Update experiment status in the database.
+
+    This function provides an efficient way to update only the experiment's status
+    field without re-saving the entire experiment. Use this when you need to persist
+    status changes immediately after calling status transition methods
+    (e.g., mark_initialization(), mark_optimization()).
+
+    Note: save_experiment() already handles status updates, so this function is
+    optional. Use it when you need status-only updates for efficiency.
+
+    Args:
+        experiment_with_updated_status: Experiment with updated status.
+        config: SQAConfig to use for database operations.
+
+    Raises:
+        ValueError: If experiment has not been saved to the database yet.
+    """
+    config = SQAConfig() if config is None else config
+    exp_sqa_class = config.class_to_sqa_class[Experiment]
+
+    exp_id = experiment_with_updated_status.db_id
+    if exp_id is None:
+        raise ValueError("Experiment must be saved before being updated.")
+
+    with session_scope() as session:
+        session.query(exp_sqa_class).filter_by(id=exp_id).update(
+            {
+                "status": experiment_with_updated_status.status,
+            }
+        )
+
+
 def update_properties_on_trial(
     trial_with_updated_properties: BaseTrial,
     config: SQAConfig | None = None,

ax/storage/sqa_store/sqa_classes.py

@@ -14,6 +14,7 @@
 
 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
@@ -379,6 +380,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

@@ -527,6 +527,29 @@ def test_saving_an_experiment_with_type_errors_with_missing_enum_value(
                 config=SQAConfig(experiment_type_enum=TestExperimentTypeEnum),
             )
 
+    def test_experiment_status_save_load(self) -> None:
+        """Test that experiment status is correctly saved and loaded."""
+        from ax.core.experiment_status import ExperimentStatus
+
+        # 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_LoadExperimentTrialsInBatches(self) -> None:
         for _ in range(4):
             self.experiment.new_trial()

ax/storage/sqa_store/with_db_settings_base.py

@@ -66,6 +66,7 @@
         _save_or_update_trials,
         _update_generation_strategy,
         save_analysis_card,
+        update_experiment_status,
         update_properties_on_experiment,
         update_runner_on_experiment,
     )
@@ -468,6 +469,27 @@ def _update_experiment_properties_in_db(
             return True
         return False
 
+    def _update_experiment_status_in_db_if_possible(
+        self,
+        experiment_with_updated_status: Experiment,
+    ) -> bool:
+        """Update experiment status in the database if DB settings are configured.
+
+        Args:
+            experiment_with_updated_status: Experiment with updated status.
+
+        Returns:
+            True if the update was performed, False if DB settings are not configured.
+        """
+        if self.db_settings_set:
+            _update_experiment_status_in_db(
+                experiment_with_updated_status=experiment_with_updated_status,
+                sqa_config=self.db_settings.encoder.config,
+                suppress_all_errors=self._suppress_all_errors,
+            )
+            return True
+        return False
+
     def _save_analysis_card_to_db_if_possible(
         self,
         experiment: Experiment,
@@ -625,6 +647,22 @@ def _update_experiment_properties_in_db(
     )
 
 
+@retry_on_exception(
+    retries=3,
+    default_return_on_suppression=False,
+    exception_types=RETRY_EXCEPTION_TYPES,
+)
+def _update_experiment_status_in_db(
+    experiment_with_updated_status: Experiment,
+    sqa_config: SQAConfig,
+    suppress_all_errors: bool,  # Used by the decorator.
+) -> None:
+    update_experiment_status(
+        experiment_with_updated_status=experiment_with_updated_status,
+        config=sqa_config,
+    )
+
+
 @retry_on_exception(
     retries=3,
     default_return_on_suppression=False,