Exclude LILO labeling trials from trials_expecting_data (facebook#5144)

Summary:
Pull Request resolved: facebook#5144

LILO labeling trials have their pairwise preference data fetched inline
during the labeling loop (`_run_lilo_labeling_loop`), not via the normal
data refetch paths (orchestrator `poll_and_process_results`,
PTSClient `refetch_data`). Including them in `trials_expecting_data`
causes unnecessary data fetch attempts for metrics (e.g., Deltoid) that
don't exist on these trials, producing noisy errors and wasting time.

This filters LILO labeling trials (`trial_type == Keys.LILO_LABELING`)
from the `trials_expecting_data` property on `Experiment`, which is the
centralized source used by all downstream data refetch consumers.

Reviewed By: saitcakmak

Differential Revision: D99571562

fbshipit-source-id: 7064918189335ac2630fcd7ea57046ae296e437a

Author: ItsMrLin

ax/core/base_trial.py

@@ -198,6 +198,16 @@ def status(self) -> TrialStatus:
         self._mark_stale_if_past_TTL()
         return none_throws(self._status)
 
+    @property
+    def expecting_data(self) -> bool:
+        """Whether this trial expects data via the standard data-fetch pipeline.
+
+        Returns ``False`` for LILO labeling trials because their pairwise
+        preference data is fetched inline during the labeling loop and is
+        never refetched through the normal orchestration path.
+        """
+        return self.status.expecting_data and self.trial_type != Keys.LILO_LABELING
+
     @status.setter
     def status(self, status: TrialStatus) -> None:
         raise NotImplementedError("Use `trial.mark_*` methods to set trial status.")

ax/core/experiment.py

@@ -45,11 +45,7 @@
 from ax.core.runner import Runner
 from ax.core.search_space import SearchSpace
 from ax.core.trial import Trial
-from ax.core.trial_status import (
-    DEFAULT_STATUSES_TO_WARM_START,
-    STATUSES_EXPECTING_DATA,
-    TrialStatus,
-)
+from ax.core.trial_status import DEFAULT_STATUSES_TO_WARM_START, TrialStatus
 from ax.core.types import ComparisonOp, TParameterization
 from ax.exceptions.core import (
     AxError,
@@ -1406,10 +1402,10 @@ def trials_by_status(self) -> dict[TrialStatus, list[BaseTrial]]:
 
     @property
     def trials_expecting_data(self) -> list[BaseTrial]:
-        """list[BaseTrial]: the list of all trials for which data has arrived
-        or is expected to arrive.
+        """list[BaseTrial]: the list of all trials that expect data via the
+        standard data-fetch pipeline.
         """
-        return [trial for trial in self.trials.values() if trial.status.expecting_data]
+        return [trial for trial in self.trials.values() if trial.expecting_data]
 
     @property
     def completed_trials(self) -> list[BaseTrial]:
@@ -1433,15 +1429,10 @@ def running_trial_indices(self) -> set[int]:
 
     @property
     def trial_indices_expecting_data(self) -> set[int]:
-        """Set of indices of trials, statuses of which indicate that we expect
-        these trials to have data, either already or in the future.
+        """Set of indices of trials that expect data via the standard
+        data-fetch pipeline.
         """
-        return set.union(
-            *(
-                self.trial_indices_by_status[status]
-                for status in STATUSES_EXPECTING_DATA
-            )
-        )
+        return {trial.index for trial in self.trials.values() if trial.expecting_data}
 
     def trial_indices_with_data(
         self, critical_metrics_only: bool | None = True

ax/core/multi_type_experiment.py

@@ -275,7 +275,7 @@ def fetch_data(
             [
                 (
                     trial.fetch_data(**kwargs, metrics=metrics)
-                    if trial.status.expecting_data
+                    if trial.expecting_data
                     else Data()
                 )
                 for trial in self.trials.values()

ax/core/tests/test_experiment.py

@@ -881,14 +881,26 @@ def test_experiment_runner(self) -> None:
         candidate_batch.run()
         candidate_batch._status = TrialStatus.CANDIDATE
         self.assertEqual(self.experiment.trials_expecting_data, [batch])
+
+        # LILO labeling trials are excluded from trials_expecting_data
+        # (their data is fetched inline during the labeling loop).
+        lilo_batch = self.experiment.new_batch_trial(
+            trial_type=Keys.LILO_LABELING,
+        )
+        lilo_batch.run()
+        lilo_batch.mark_completed()
+        self.assertEqual(self.experiment.trials_expecting_data, [batch])
+
         tbs = self.experiment.trials_by_status  # All statuses should be present
         self.assertEqual(len(tbs), len(TrialStatus))
         self.assertEqual(tbs[TrialStatus.RUNNING], [batch])
         self.assertEqual(tbs[TrialStatus.CANDIDATE], [candidate_batch])
+        self.assertEqual(tbs[TrialStatus.COMPLETED], [lilo_batch])
         tibs = self.experiment.trial_indices_by_status
         self.assertEqual(len(tibs), len(TrialStatus))
         self.assertEqual(tibs[TrialStatus.RUNNING], {0})
         self.assertEqual(tibs[TrialStatus.CANDIDATE], {1})
+        self.assertEqual(tibs[TrialStatus.COMPLETED], {2})
 
         identifier = {"new_runner": True}
         # pyre-fixme[6]: For 1st param expected `Optional[str]` but got `Dict[str,
@@ -1727,6 +1739,12 @@ def test_trial_indices(self) -> None:
         )
         self.assertEqual(experiment.trial_indices_expecting_data, {2, 5})
 
+        # LILO labeling trials are excluded from trial_indices_expecting_data.
+        lilo_trial = experiment.new_batch_trial(trial_type=Keys.LILO_LABELING)
+        lilo_trial.mark_running(no_runner_required=True)
+        lilo_trial.mark_completed()
+        self.assertEqual(experiment.trial_indices_expecting_data, {2, 5})
+
     def test_trial_indices_with_data(self) -> None:
         exp = get_branin_experiment_with_multi_objective(
             with_status_quo=True,

ax/core/tests/test_trial.py

@@ -26,6 +26,7 @@
 from ax.exceptions.core import TrialMutationError, UnsupportedError, UserInputError
 from ax.metrics.branin import BraninMetric
 from ax.runners.synthetic import SyntheticRunner
+from ax.utils.common.constants import Keys
 from ax.utils.common.result import Ok
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
@@ -271,8 +272,22 @@ def test_mark_as(self) -> None:
                     TrialStatus.COMPLETED,
                 ]:
                     self.assertTrue(self.trial.status.expecting_data)
+                    # trial.expecting_data follows status for normal trials.
+                    self.assertTrue(self.trial.expecting_data)
                 else:
                     self.assertFalse(self.trial.status.expecting_data)
+                    self.assertFalse(self.trial.expecting_data)
+
+    def test_expecting_data_excludes_lilo(self) -> None:
+        """LILO labeling trials never expect data via the standard pipeline."""
+        self.trial._trial_type = Keys.LILO_LABELING
+        self.trial.mark_running(no_runner_required=True)
+        self.assertTrue(self.trial.status.expecting_data)
+        self.assertFalse(self.trial.expecting_data)
+
+        self.trial.mark_completed()
+        self.assertTrue(self.trial.status.expecting_data)
+        self.assertFalse(self.trial.expecting_data)
 
     def test_stop(self) -> None:
         # test bad old status

ax/orchestration/orchestrator.py

@@ -477,7 +477,7 @@ def trials_expecting_data(self) -> list[BaseTrial]:
         """
         trials = []
         for trial in self.experiment.trials.values():
-            if trial.status.expecting_data:
+            if trial.expecting_data:
                 if self.trial_type is None or trial.trial_type == self.trial_type:
                     trials.append(trial)
         return trials