Make brute-force strategies budget aware (#3805)

nni/retiarii/strategy/bruteforce.py

@@ -8,7 +8,7 @@ import random
 import time
 from typing import Any, Dict, List
 
-from .. import Sampler, submit_models, query_available_resources
+from .. import Sampler, submit_models, query_available_resources, budget_exhausted
 from .base import BaseStrategy
 from .utils import dry_run_for_search_space, get_targeted_model
 
@@ -63,6 +63,8 @@ class GridSearch(BaseStrategy):
         for sample in grid_generator(search_space, shuffle=self.shuffle):
             _logger.debug('New model created. Waiting for resource. %s', str(sample))
             while query_available_resources() <= 0:
+                if budget_exhausted():
+                    return
                 time.sleep(self._polling_interval)
             submit_models(get_targeted_model(base_model, applied_mutators, sample))
 
@@ -106,6 +108,8 @@ class Random(BaseStrategy):
                         model = mutator.apply(model)
                     _logger.debug('New model created. Applied mutators are: %s', str(applied_mutators))
                     submit_models(model)
+                elif budget_exhausted():
+                    break
                 else:
                     time.sleep(self._polling_interval)
         else:
@@ -114,5 +118,7 @@ class Random(BaseStrategy):
             for sample in random_generator(search_space, dedup=self.dedup):
                 _logger.debug('New model created. Waiting for resource. %s', str(sample))
                 while query_available_resources() <= 0:
+                    if budget_exhausted():
+                        return
                     time.sleep(self._polling_interval)
                 submit_models(get_targeted_model(base_model, applied_mutators, sample))