Fix for ax batch_client

Differential Revision: D15203695

fbshipit-source-id: d05fe14d80c245b1854ad1a56ae68a7491beaa5f

ax/service/ax_client.py

@@ -320,6 +320,17 @@ def _save_experiment_if_possible(self) -> bool:
             )
         return False
 
+    def _get_new_data(self) -> Data:
+        """
+        Returns new data since the last run of the generator.
+
+        Returns:
+            Latest data.
+        """
+        return Data.from_multiple_data(
+            [self.experiment.lookup_data_for_trial(idx) for idx in self._updated_trials]
+        )
+
     def _suggest_new_trial(self) -> Trial:
         """
         Suggest new candidate for this experiment.
@@ -330,9 +341,7 @@ def _suggest_new_trial(self) -> Trial:
         Returns:
             Trial with candidate.
         """
-        new_data = Data.from_multiple_data(
-            [self.experiment.lookup_data_for_trial(idx) for idx in self._updated_trials]
-        )
+        new_data = self._get_new_data()
         generator_run = not_none(self.generation_strategy).gen(
             experiment=self.experiment, new_data=new_data
         )