Allow updating of opt config on completion events

Summary: This will allow for accuracy in human in the loop experiments where metrics change through the lifecycle of experiments.

Differential Revision: D61996011

ax/telemetry/experiment.py

@@ -263,6 +263,12 @@ class ExperimentCompletedRecord:
     total_fit_time: int
     total_gen_time: int
 
+    # OptimizationConfig info which might be updated for human in the
+    # loop experiments
+    num_objectives: int
+    num_tracking_metrics: int
+    num_outcome_constraints: int  # Includes ObjectiveThresholds in MOO
+
     @classmethod
     def from_experiment(cls, experiment: Experiment) -> ExperimentCompletedRecord:
         trial_count_by_status = {
@@ -308,4 +314,15 @@ def from_experiment(cls, experiment: Experiment) -> ExperimentCompletedRecord:
             num_early_stopped_trials=trial_count_by_status[TrialStatus.EARLY_STOPPED],
             total_fit_time=int(fit_time),
             total_gen_time=int(gen_time),
+            num_objectives=(
+                len(experiment.optimization_config.objective.metrics)
+                if experiment.optimization_config is not None
+                else 0
+            ),
+            num_tracking_metrics=len(experiment.tracking_metrics),
+            num_outcome_constraints=(
+                len(experiment.optimization_config.outcome_constraints)
+                if experiment.optimization_config is not None
+                else 0
+            ),
         )

ax/telemetry/optimization.py

@@ -458,6 +458,12 @@ class OptimizationCompletedRecord:
     estimated_early_stopping_savings: float
     estimated_global_stopping_savings: float
 
+    # OptimizationConfig info which might be updated for human in the
+    # loop experiments
+    num_objectives: int
+    num_tracking_metrics: int
+    num_outcome_constraints: int  # Includes ObjectiveThresholds in MOO
+
     @classmethod
     def from_scheduler(
         cls,
@@ -503,6 +509,9 @@ def from_scheduler(
             deployed_job_id=deployed_job_id,
             estimated_early_stopping_savings=estimated_early_stopping_savings,
             estimated_global_stopping_savings=estimated_global_stopping_savings,
+            num_objectives=experiment_completed_record.num_objectives,
+            num_tracking_metrics=experiment_completed_record.num_tracking_metrics,
+            num_outcome_constraints=experiment_completed_record.num_outcome_constraints,
         )
 
     @classmethod
@@ -545,6 +554,9 @@ def from_ax_client(
             improvement_over_baseline=float("nan"),
             num_metric_fetch_e_encountered=-1,
             num_trials_bad_due_to_err=-1,
+            num_objectives=experiment_completed_record.num_objectives,
+            num_tracking_metrics=experiment_completed_record.num_tracking_metrics,
+            num_outcome_constraints=experiment_completed_record.num_outcome_constraints,
         )
 
 

ax/telemetry/tests/test_experiment.py

@@ -66,6 +66,9 @@ def test_experiment_completed_record_from_experiment(self) -> None:
             num_early_stopped_trials=0,
             total_fit_time=int(fit_time),
             total_gen_time=int(gen_time),
+            num_objectives=1,
+            num_outcome_constraints=1,
+            num_tracking_metrics=1,
         )
         self.assertEqual(record, expected)