diff --git a/ax/telemetry/experiment.py b/ax/telemetry/experiment.py index 55d05e7de09..cddc2a0a002 100644 --- a/ax/telemetry/experiment.py +++ b/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 + ), ) diff --git a/ax/telemetry/optimization.py b/ax/telemetry/optimization.py index aa485e6486c..1c7cc37389d 100644 --- a/ax/telemetry/optimization.py +++ b/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, ) diff --git a/ax/telemetry/tests/test_experiment.py b/ax/telemetry/tests/test_experiment.py index 423e9c9ef76..06577d3ae93 100644 --- a/ax/telemetry/tests/test_experiment.py +++ b/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)