[tune] lazy trials (#10802)

* Lazily fill trial queue

* Update interface

* Update end to end reporter test

* Removed `next_trials()` method

* Lint

* Print total number of samples to be generated in progress reporter. Allow infinite samples.

* Nit check

python/ray/tune/suggest/basic_variant.py

@@ -1,14 +1,13 @@
 import itertools
 import os
-import random
 import uuid
 from typing import Dict, List, Union
 
 from ray.tune.error import TuneError
 from ray.tune.experiment import Experiment, convert_to_experiment_list
 from ray.tune.config_parser import make_parser, create_trial_from_spec
-from ray.tune.suggest.variant_generator import (generate_variants, format_vars,
-                                                flatten_resolved_vars)
+from ray.tune.suggest.variant_generator import (
+    count_variants, generate_variants, format_vars, flatten_resolved_vars)
 from ray.tune.suggest.search import SearchAlgorithm
 
 
@@ -17,10 +16,6 @@ class BasicVariantGenerator(SearchAlgorithm):
 
     See also: `ray.tune.suggest.variant_generator`.
 
-
-    Parameters:
-        shuffle (bool): Shuffles the generated list of configurations.
-
     User API:
 
     .. code-block:: python
@@ -39,19 +34,19 @@ class BasicVariantGenerator(SearchAlgorithm):
 
         searcher = BasicVariantGenerator()
         searcher.add_configurations({"experiment": { ... }})
-        list_of_trials = searcher.next_trials()
+        trial = searcher.next_trial()
         searcher.is_finished == True
     """
 
-    def __init__(self, shuffle: bool = False):
+    def __init__(self):
         """Initializes the Variant Generator.
 
         """
         self._parser = make_parser()
         self._trial_generator = []
+        self._trial_iter = None
         self._counter = 0
         self._finished = False
-        self._shuffle = shuffle
 
         # Unique prefix for all trials generated, e.g., trial ids start as
         # 2f1e_00001, 2f1ef_00002, 2f1ef_0003, etc. Overridable for testing.
@@ -61,6 +56,12 @@ class BasicVariantGenerator(SearchAlgorithm):
         else:
             self._uuid_prefix = str(uuid.uuid1().hex)[:5] + "_"
 
+        self._total_samples = 0
+
+    @property
+    def total_samples(self):
+        return self._total_samples
+
     def add_configurations(
             self,
             experiments: Union[Experiment, List[Experiment], Dict[str, Dict]]):
@@ -71,23 +72,28 @@ class BasicVariantGenerator(SearchAlgorithm):
         """
         experiment_list = convert_to_experiment_list(experiments)
         for experiment in experiment_list:
+            self._total_samples += count_variants(experiment.spec)
             self._trial_generator = itertools.chain(
                 self._trial_generator,
                 self._generate_trials(
                     experiment.spec.get("num_samples", 1), experiment.spec,
                     experiment.name))
 
-    def next_trials(self):
-        """Provides Trial objects to be queued into the TrialRunner.
+    def next_trial(self):
+        """Provides one Trial object to be queued into the TrialRunner.
 
         Returns:
-            trials (list): Returns a list of trials.
+            Trial: Returns a single trial.
         """
-        trials = list(self._trial_generator)
-        if self._shuffle:
-            random.shuffle(trials)
-        self.set_finished()
-        return trials
+        if not self._trial_iter:
+            self._trial_iter = iter(self._trial_generator)
+        try:
+            return next(self._trial_iter)
+        except StopIteration:
+            self._trial_generator = []
+            self._trial_iter = None
+            self.set_finished()
+            return None
 
     def _generate_trials(self, num_samples, unresolved_spec, output_path=""):
         """Generates Trial objects with the variant generation process.

python/ray/tune/suggest/search.py

@@ -1,7 +1,6 @@
 from typing import Dict, List, Optional, Union
 
 from ray.tune.experiment import Experiment
-from ray.tune.trial import Trial
 
 
 class SearchAlgorithm:
@@ -36,6 +35,11 @@ class SearchAlgorithm:
         """
         return True
 
+    @property
+    def total_samples(self):
+        """Get number of total trials to be generated"""
+        return 0
+
     def add_configurations(
             self,
             experiments: Union[Experiment, List[Experiment], Dict[str, Dict]]):
@@ -46,11 +50,11 @@ class SearchAlgorithm:
         """
         raise NotImplementedError
 
-    def next_trials(self) -> List[Trial]:
-        """Provides Trial objects to be queued into the TrialRunner.
+    def next_trial(self):
+        """Returns single Trial object to be queued into the TrialRunner.
 
         Returns:
-            trials (list): Returns a list of trials.
+            trial (Trial): Returns a Trial object.
         """
         raise NotImplementedError
 

python/ray/tune/suggest/search_generator.py

@@ -67,13 +67,17 @@ class SearchGenerator(SearchAlgorithm):
         self._parser = make_parser()
         self._experiment = None
         self._counter = 0  # Keeps track of number of trials created.
-        self._total_samples = None  # int: total samples to evaluate.
+        self._total_samples = 0  # int: total samples to evaluate.
         self._finished = False
 
     def set_search_properties(self, metric: Optional[str], mode: Optional[str],
                               config: Dict) -> bool:
         return self.searcher.set_search_properties(metric, mode, config)
 
+    @property
+    def total_samples(self):
+        return self._total_samples
+
     def add_configurations(
             self,
             experiments: Union[Experiment, List[Experiment], Dict[str, Dict]]):
@@ -95,20 +99,16 @@ class SearchGenerator(SearchAlgorithm):
         if "run" not in experiment_spec:
             raise TuneError("Must specify `run` in {}".format(experiment_spec))
 
-    def next_trials(self) -> List[Trial]:
-        """Provides a batch of Trial objects to be queued into the TrialRunner.
+    def next_trial(self):
+        """Provides one Trial object to be queued into the TrialRunner.
 
         Returns:
-            List[Trial]: A list of trials for the Runner to consume.
+            Trial: Returns a single trial.
         """
-        trials = []
-        while not self.is_finished():
-            trial = self.create_trial_if_possible(self._experiment.spec,
-                                                  self._experiment.name)
-            if trial is None:
-                break
-            trials.append(trial)
-        return trials
+        if not self.is_finished():
+            return self.create_trial_if_possible(self._experiment.spec,
+                                                 self._experiment.name)
+        return None
 
     def create_trial_if_possible(self, experiment_spec: Dict,
                                  output_path: str) -> Optional[Trial]:

python/ray/tune/suggest/variant_generator.py

@@ -138,6 +138,15 @@ def parse_spec_vars(spec: Dict) -> Tuple[List[Tuple[Tuple, Any]], List[Tuple[
     return resolved_vars, domain_vars, grid_vars
 
 
+def count_variants(spec: Dict) -> int:
+    spec = copy.deepcopy(spec)
+    _, domain_vars, grid_vars = parse_spec_vars(spec)
+    grid_count = 1
+    for path, domain in grid_vars:
+        grid_count *= len(domain.categories)
+    return spec.get("num_samples", 1) * grid_count
+
+
 def _generate_variants(spec: Dict) -> Tuple[Dict, Dict]:
     spec = copy.deepcopy(spec)
     _, domain_vars, grid_vars = parse_spec_vars(spec)