Track all GeneratorRun-s on Experiment (+ a bit of cleanup) (#3327)

Summary: Pull Request resolved: #3327

Differential Revision: D69221377

ax/core/experiment.py

@@ -66,9 +66,7 @@
 )
 
 ROUND_FLOATS_IN_LOGS_TO_DECIMAL_PLACES: int = 6
-
-# pyre-fixme[5]: Global expression must be annotated.
-round_floats_for_logging = partial(
+round_floats_for_logging: partial[int] = partial(
     _round_floats_for_logging,
     decimal_places=ROUND_FLOATS_IN_LOGS_TO_DECIMAL_PLACES,
 )
@@ -82,7 +80,38 @@
 
 
 class Experiment(Base):
-    """Base class for defining an experiment."""
+    # I. Metadata:
+    _name: str | None
+    _is_test: bool  # pyre-ignore[13]: Initialized in a setter.
+    _experiment_type: str | None  # pyre-ignore[13]: Initialized in a setter.
+    _time_created: datetime
+    # NOTE: While ``_properties`` is an unstructured JSON blob, it's meant
+    # exclusively for storing a small set of primitives that pertain to the
+    # experiment state, and NOT deployment- or modeling-related information.
+    _properties: dict[str, Any]
+    _default_data_type: DataType
+
+    # II. Key components of experiment design and state:
+    _search_space: SearchSpace  # pyre-ignore[13]: Initialized in a setter.
+    # Experiment can be created without an Opt.Config, e.g. it can be
+    # configured after an exploratory trial.
+    _optimization_config: OptimizationConfig | None = None
+    _trials: dict[int, BaseTrial]
+    _data_by_trial: dict[int, OrderedDict[int, Data]]
+    _tracking_metrics: dict[str, Metric]
+    _status_quo: Arm | None = None
+    _generator_runs: dict[int, GeneratorRun]
+    # Arms are unique on the experiment: arms with the same underlying
+    # parameterization will always receive the same name.
+    _arms_by_signature: dict[str, Arm]
+    _arms_by_name: dict[str, Arm]
+    # Experiment caches trial indices by status for performant access.
+    _trial_indices_by_status: dict[TrialStatus, set[int]]
+    # Auxiliary experiments are sources of auxiliary data that can be
+    # used by the optimization methods for the ongoing experiment.
+    auxiliary_experiments_by_purpose: dict[
+        AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]
+    ]
 
     def __init__(
         self,
@@ -96,11 +125,13 @@ def __init__(
         is_test: bool = False,
         experiment_type: str | None = None,
         properties: dict[str, Any] | None = None,
-        default_data_type: DataType | None = None,
+        default_data_type: DataType = DataType.DATA,
         auxiliary_experiments_by_purpose: None
         | (dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]]) = None,
     ) -> None:
-        """Inits Experiment.
+        """Initializes an ``Experiment``: central object for tracking experiment
+        state in Ax. The only required argument is a ``search_space``, but
+        ``name`` and ``is_test`` are also recommended.
 
         Args:
             search_space: Search space of the experiment.
@@ -118,52 +149,39 @@ def __init__(
                 should be stored elsewhere, e.g. in ``run_metadata`` of the trials.
             default_data_type: Enum representing the data type this experiment uses.
             auxiliary_experiments_by_purpose: Dictionary of auxiliary experiments
-                for different purposes (e.g., transfer learning).
+                (sources of additional data for the optimization methodology leveraged
+                for the current expeirment), grouped by different purposes
+                (e.g., historical experiments for the purpose of Transfer Learning).
         """
-        # appease pyre
-        # pyre-fixme[13]: Attribute `_search_space` is never initialized.
-        self._search_space: SearchSpace
-        self._status_quo: Arm | None = None
-        # pyre-fixme[13]: Attribute `_is_test` is never initialized.
-        self._is_test: bool
-
         self._name = name
         self.description = description
         self.runner = runner
-        self.is_test = is_test
 
-        self._data_by_trial: dict[int, OrderedDict[int, Data]] = {}
-        self._experiment_type: str | None = experiment_type
-        # pyre-fixme[4]: Attribute must be annotated.
-        self._optimization_config = None
-        self._tracking_metrics: dict[str, Metric] = {}
-        self._time_created: datetime = datetime.now()
-        self._trials: dict[int, BaseTrial] = {}
-        self._properties: dict[str, Any] = properties or {}
-        # pyre-fixme[4]: Attribute must be annotated.
-        self._default_data_type = default_data_type or DataType.DATA
+        self._data_by_trial = {}
+        self._tracking_metrics = {}
+        self._time_created = datetime.now()
+        self._trials = {}
+        self._generator_runs = {}
+        self._properties = properties or {}
+        self._default_data_type = default_data_type
         # Used to keep track of whether any trials on the experiment
         # specify a TTL. Since trials need to be checked for their TTL's
         # expiration often, having this attribute helps avoid unnecessary
         # TTL checks for experiments that do not use TTL.
         self._trials_have_ttl = False
         # Make sure all statuses appear in this dict, to avoid key errors.
-        self._trial_indices_by_status: dict[TrialStatus, set[int]] = {
-            status: set() for status in TrialStatus
-        }
-        self._arms_by_signature: dict[str, Arm] = {}
-        self._arms_by_name: dict[str, Arm] = {}
-
-        self.auxiliary_experiments_by_purpose: dict[
-            AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]
-        ] = auxiliary_experiments_by_purpose or {}
-
-        self.add_tracking_metrics(tracking_metrics or [])
+        self._trial_indices_by_status = {status: set() for status in TrialStatus}
+        self._arms_by_signature = {}
+        self._arms_by_name = {}
+        self.auxiliary_experiments_by_purpose = auxiliary_experiments_by_purpose or {}
 
-        # call setters defined below
+        # Call setters to set the respective protected attributes.
         self.search_space = search_space
         self.status_quo = status_quo
-        if optimization_config is not None:
+        self.is_test = is_test
+        self.experiment_type = experiment_type
+        self.add_tracking_metrics(tracking_metrics or [])
+        if optimization_config:
             self.optimization_config = optimization_config
 
     @property
@@ -1040,6 +1058,15 @@ def trials(self) -> dict[int, BaseTrial]:
         self._check_TTL_on_running_trials()
         return self._trials
 
+    @property
+    def generator_runs(self) -> dict[int, GeneratorRun]:
+        """The generator_runs associated with the experiment.
+
+        NOTE: Generator runs are associated with an experiment when arms from
+        them are added to a trial on the expeirment.
+        """
+        return self._generator_runs
+
     @property
     def trials_by_status(self) -> dict[TrialStatus, list[BaseTrial]]:
         """Trials associated with the experiment, grouped by trial status."""

ax/core/multi_type_experiment.py

@@ -27,20 +27,15 @@
 
 
 class MultiTypeExperiment(Experiment):
-    """Class for experiment with multiple trial types.
-
-    A canonical use case for this is tuning a large production system
-    with limited evaluation budget and a simulator which approximates
-    evaluations on the main system. Trial deployment and data fetching
-    is separate for the two systems, but the final data is combined and
-    fed into multi-task models.
-
-    See the Multi-Task Modeling tutorial for more details.
-
-    Attributes:
-        name: Name of the experiment.
-        description: Description of the experiment.
-    """
+    _default_trial_type: str
+    # Map from trial type to default runner of that type
+    _trial_type_to_runner: dict[str, Runner | None]
+    # Specifies which trial type each metric belongs to
+    _metric_to_trial_type: dict[str, str]
+    # Maps certain metric names to a canonical name. Useful for ancillary trial
+    # types' metrics, to specify which primary metrics they correspond to
+    # (e.g. 'comment_prediction' => 'comment')
+    _metric_to_canonical_name: dict[str, str]
 
     def __init__(
         self,
@@ -55,43 +50,50 @@ def __init__(
         is_test: bool = False,
         experiment_type: str | None = None,
         properties: dict[str, Any] | None = None,
-        default_data_type: DataType | None = None,
+        default_data_type: DataType = DataType.DATA,
     ) -> None:
-        """Inits Experiment.
+        """Initializes a ``MultiTypeExperiment``: special type of an experiment, in
+        which trials will have different "types", with each type determining how the
+        trial might be evaluated/deployed, as well as how it will be treated by the
+        optimization methodologies.
+
+        A canonical use case for this is tuning a large production system
+        with limited evaluation budget and a simulator which approximates
+        evaluations on the main system. Trial deployment and data fetching
+        is separate for the two systems, but the final data is combined and
+        fed into multi-task models.
+
+        See the Multi-Task Modeling tutorial for more details.
 
         Args:
-            name: Name of the experiment.
             search_space: Search space of the experiment.
-            default_trial_type: Default type for trials on this experiment.
-            default_runner: Default runner for trials of the default type.
+            default_trial_type: Type to assign to trials unless otherwise specified.
+            default_runner: Runner to use for the trial of the default trial type on
+                this experiment.
+            name: Name of the experiment.
             optimization_config: Optimization config of the experiment.
             tracking_metrics: Additional tracking metrics not used for optimization.
-                These are associated with the default trial type.
             runner: Default runner used for trials on this experiment.
             status_quo: Arm representing existing "control" arm.
             description: Description of the experiment.
             is_test: Convenience metadata tracker for the user to mark test experiments.
             experiment_type: The class of experiments this one belongs to.
-            properties: Dictionary of this experiment's properties.
+            properties: Dictionary of this experiment's properties.  It is meant to
+                only store primitives that pertain to Ax experiment state. Any trial
+                deployment-related information and modeling-layer configuration
+                should be stored elsewhere, e.g. in ``run_metadata`` of the trials.
             default_data_type: Enum representing the data type this experiment uses.
+            auxiliary_experiments_by_purpose: Dictionary of auxiliary experiments
+                (sources of additional data for the optimization methodology leveraged
+                for the current expeirment), grouped by different purposes
+                (e.g., historical experiments for the purpose of Transfer Learning).
         """
-
         self._default_trial_type = default_trial_type
+        self._trial_type_to_runner = {default_trial_type: default_runner}
+        self._metric_to_trial_type = {}
+        self._metric_to_canonical_name = {}
 
-        # Map from trial type to default runner of that type
-        self._trial_type_to_runner: dict[str, Runner | None] = {
-            default_trial_type: default_runner
-        }
-
-        # Specifies which trial type each metric belongs to
-        self._metric_to_trial_type: dict[str, str] = {}
-
-        # Maps certain metric names to a canonical name. Useful for ancillary trial
-        # types' metrics, to specify which primary metrics they correspond to
-        # (e.g. 'comment_prediction' => 'comment')
-        self._metric_to_canonical_name: dict[str, str] = {}
-
-        # call super.__init__() after defining fields above, because we need
+        # call ``super.__init__()`` after defining fields above, because we need
         # them to be populated before optimization config is set
         super().__init__(
             name=name,

ax/core/tests/test_experiment.py

@@ -118,6 +118,7 @@ def test_ExperimentInit(self) -> None:
         self.assertIsNotNone(self.experiment.time_created)
         self.assertEqual(self.experiment.experiment_type, None)
         self.assertEqual(self.experiment.num_abandoned_arms, 0)
+        self.assertEqual(self.experiment.generator_runs, {})
 
     def test_ExperimentName(self) -> None:
         self.assertTrue(self.experiment.has_name)