Add tests for TuningBudget

Author: stijnh

kernel_tuner/runners/parallel.py

@@ -286,9 +286,9 @@ def run(self, parameter_space, tuning_options) -> List[Optional[dict]]:
                 params.update(tuning_options.cache[key])
 
                 # Simulate compile, verification, and benchmark time
-                tuning_options.budget.add_time_spent(params["compile_time"])
-                tuning_options.budget.add_time_spent(params["verification_time"])
-                tuning_options.budget.add_time_spent(params["benchmark_time"])
+                tuning_options.budget.add_time(milliseconds=params["compile_time"])
+                tuning_options.budget.add_time(milliseconds=params["verification_time"])
+                tuning_options.budget.add_time(milliseconds=params["benchmark_time"])
                 results.append(params)
             else:
                 assert key not in key2index, "duplicate jobs submitted"

kernel_tuner/runners/sequential.py

@@ -5,7 +5,7 @@
 
 from kernel_tuner.core import DeviceInterface
 from kernel_tuner.runners.runner import Runner
-from kernel_tuner.util import ErrorConfig, print_config_output, process_metrics, store_cache, stop_criterion_reached
+from kernel_tuner.util import ErrorConfig, print_config_output, process_metrics, store_cache
 
 
 class SequentialRunner(Runner):
@@ -70,12 +70,15 @@ def run(self, parameter_space, tuning_options):
 
         # iterate over parameter space
         for element in parameter_space:
+            # If the time limit is exceeded, just skip this element. Add `None` to
+            # indicate to CostFunc that no result is available for this config.
+            if tuning_options.budget.is_done():
+                results.append(None)
+                continue
+            
             tuning_options.budget.add_evaluations(1)
             params = dict(zip(tuning_options.tune_params.keys(), element))
 
-            if stop_criterion_reached(tuning_options):
-                return results
-
             result = None
             warmup_time = 0
 
@@ -85,9 +88,9 @@ def run(self, parameter_space, tuning_options):
                 params.update(tuning_options.cache[x_int])
 
                 # Simulate compile, verification, and benchmark time
-                tuning_options.budget.add_time_spent(params["compile_time"])
-                tuning_options.budget.add_time_spent(params["verification_time"])
-                tuning_options.budget.add_time_spent(params["benchmark_time"])
+                tuning_options.budget.add_time(milliseconds=params["compile_time"])
+                tuning_options.budget.add_time(milliseconds=params["verification_time"])
+                tuning_options.budget.add_time(milliseconds=params["benchmark_time"])
             else:
                 # attempt to warmup the GPU by running the first config in the parameter space and ignoring the result
                 if not self.warmed_up:

kernel_tuner/runners/simulation.py

@@ -116,9 +116,9 @@ def run(self, parameter_space, tuning_options):
 
                 # Simulate the evaluation of this configuration
                 tuning_options.budget.add_evaluations(1)
-                tuning_options.budget.add_time_spent(result["compile_time"])
-                tuning_options.budget.add_time_spent(result["verification_time"])
-                tuning_options.budget.add_time_spent(result["benchmark_time"])
+                tuning_options.budget.add_time(milliseconds=result["compile_time"])
+                tuning_options.budget.add_time(milliseconds=result["verification_time"])
+                tuning_options.budget.add_time(milliseconds=result["benchmark_time"])
 
                 try:
                     self.total_simulated_time += result["compile_time"] + result["verification_time"] + result["benchmark_time"]

kernel_tuner/strategies/common.py

@@ -145,7 +145,6 @@ def _run_configs(self, xs, check_restrictions=True):
         batch_keys = [] # The keys of the configs to run
         pending_indices_by_key = dict()  # Maps key => where to store result in `final_results`
         final_results = []  # List returned to the user
-        benchmark_config = []
 
         # Loop over all configurations. For each configurations there are four cases:
         # 1. The configuration is valid, we can skip it
@@ -201,16 +200,11 @@ def _run_configs(self, xs, check_restrictions=True):
             self.tuning_options.unique_results[key] = result
             self.results.append(result)
 
-
-        # check again for stop condition
-        # this check is necessary because some strategies cannot handle partially completed requests
-        # for example when only half of the configs in a population have been evaluated
-        self.budget_spent_fraction = util.check_stop_criterion(self.tuning_options)
-
         # upon returning from this function control will be given back to the strategy, so reset the start time
         self.runner.last_strategy_start_time = perf_counter()
 
-        # Check the tuning budget again
+        # this check is necessary because some strategies cannot handle partially completed requests
+        # for example when only half of the configs in a population have been evaluated
         self.tuning_options.budget.raise_exception_if_done()
         self.budget_spent_fraction = self.tuning_options.budget.get_fraction_consumed()
 

kernel_tuner/util.py

@@ -190,19 +190,13 @@ def check_argument_list(kernel_name, kernel_string, args):
 
 class TuningBudget:
     def __init__(self, time_limit=None, max_fevals=None):
-        if max_fevals is None:
-            max_fevals = float("inf")
-            
-        if time_limit is None:
-            time_limit = timedelta.max
-
-        if not isinstance(time_limit, timedelta):
+        if time_limit is not None and not isinstance(time_limit, timedelta):
             time_limit = timedelta(seconds=time_limit)
 
-        if max_fevals <= 0:
+        if max_fevals is not None and max_fevals <= 0:
             raise ValueError("max_fevals must be greater than zero")
 
-        if time_limit <= timedelta(seconds=0):
+        if time_limit is not None and time_limit <= timedelta(seconds=0):
             raise ValueError("time_limit must be greater than zero")
 
         self.start_time_seconds = time.perf_counter()
@@ -214,50 +208,55 @@ def __init__(self, time_limit=None, max_fevals=None):
     def add_evaluations(self, n=1):
         self.num_fevals += n
 
-    def add_time_spent(self, delta):
-        if not isinstance(delta, timedelta):
-            delta = timedelta(seconds=delta)
-        self.time_spent_extra += delta
+    def add_time(self, seconds=0, milliseconds=0):
+        self.time_spent_extra += timedelta(seconds=seconds, milliseconds=milliseconds)
 
     def get_time_spent(self) -> timedelta:
         seconds_passed = time.perf_counter() - self.start_time_seconds
         return timedelta(seconds=seconds_passed) + self.time_spent_extra
 
     def get_time_remaining(self) -> timedelta:
-        return max(self.time_limit - self.get_time_spent(), timedelta(seconds=0))
+        if self.time_limit is not None:
+            return max(self.time_limit - self.get_time_spent(), timedelta(seconds=0))
+        else:
+            return timedelta.max
 
     def get_evaluations_spent(self) -> int:
-        return max(self.max_fevals - self.num_fevals, 0)
+        return self.num_fevals
 
     def get_evaluations_remaining(self) -> int:
-        return max(self.max_fevals - self.num_fevals, 0)
+        if self.max_fevals is not None:
+            return max(self.max_fevals - self.num_fevals, 0)
+        else:
+            return float("inf")
 
     def is_done(self) -> bool:
-        if self.num_fevals >= self.max_fevals:
+        if self.max_fevals is not None and self.num_fevals >= self.max_fevals:
             return True
 
-        if self.get_time_spent() > self.time_limit:
+        if self.time_limit is not None and self.get_time_spent() > self.time_limit:
             return True
 
         return False   
 
     def raise_exception_if_done(self):
-        if self.num_fevals >= self.max_fevals:
+        if self.max_fevals is not None and self.num_fevals >= self.max_fevals:
             raise StopCriterionReached(f"max_fevals ({self.max_fevals}) reached")
 
-        if self.get_time_spent() > self.time_limit:
+        if self.time_limit is not None and self.get_time_spent() > self.time_limit:
             raise StopCriterionReached("time limit exceeded")
 
     def get_fraction_consumed(self) -> float:
-        if self.num_fevals >= self.max_fevals:
-            return 1.0
-        
-        time_spent = self.get_time_spent()
+        if self.max_fevals is not None and self.time_limit is not None:
+            time_spent = self.get_time_spent()
+            return min(1.0, time_spent / self.time_limit, self.num_fevals / self.max_fevals)
+        elif self.max_fevals is not None:
+            return min(1.0, self.num_fevals / self.max_fevals)
+        elif self.time_limit is not None:
+            return min(1.0, self.get_time_spent() / self.time_limit)
+        else:
+            return 0.0
 
-        if time_spent > self.time_limit:
-            return 1.0
-        
-        return max(time_spent / self.time_limit, self.num_fevals / self.max_fevals)
 
 
 

test/test_util_functions.py

@@ -3,6 +3,7 @@
 import json
 import os
 import warnings
+import datetime
 
 import numpy as np
 import pytest
@@ -429,6 +430,92 @@ def test_check_argument_list7():
     assert_user_warning(check_argument_list, [kernel_name, kernel_string, args])
 
 
+def test_tuning_budget1():
+    budget = TuningBudget()
+    assert budget.get_evaluations_spent() == 0
+    assert budget.get_evaluations_remaining() == float("inf")
+    assert not budget.is_done()
+    budget.raise_exception_if_done() # Should not raise
+    assert budget.get_fraction_consumed() == 0.0
+
+    budget.add_evaluations(9000)
+    assert budget.get_evaluations_spent() == 9000
+    assert budget.get_evaluations_remaining() == float("inf")
+    assert not budget.is_done()
+    budget.raise_exception_if_done() # Should not raise
+    assert budget.get_fraction_consumed() == 0.0
+
+    budget.add_time(seconds=9000)
+    assert budget.get_evaluations_spent() == 9000
+    assert budget.get_evaluations_remaining() == float("inf")
+    assert not budget.is_done()
+    budget.raise_exception_if_done() # Should not raise
+    assert budget.get_fraction_consumed() == 0.0
+
+def test_tuning_budget2():
+    budget = TuningBudget(max_fevals=5)
+    assert budget.get_evaluations_spent() == 0
+    assert budget.get_evaluations_remaining() == 5
+    assert not budget.is_done()
+    budget.raise_exception_if_done() # Should not raise
+    assert budget.get_fraction_consumed() == 0.0
+
+    budget.add_evaluations(4)
+    assert budget.get_evaluations_spent() == 4
+    assert budget.get_evaluations_remaining() == 1
+    assert not budget.is_done()
+    budget.raise_exception_if_done() # Should not raise
+    assert budget.get_fraction_consumed() == 4/5
+
+    budget.add_evaluations(1)
+    assert budget.get_evaluations_spent() == 5
+    assert budget.get_evaluations_remaining() == 0
+    assert budget.is_done()
+    assert pytest.raises(StopCriterionReached, budget.raise_exception_if_done)
+    assert budget.get_fraction_consumed() == 1.0
+
+
+def test_tuning_budget3():
+    # Two values are similar if they are within 0.01
+    approx = lambda x: pytest.approx(x, abs=0.01)
+
+    budget = TuningBudget(time_limit=5)
+    assert budget.get_time_spent().total_seconds() == approx(0)
+    assert budget.get_time_remaining().total_seconds() == approx(5)
+    assert budget.get_evaluations_spent() == 0
+    assert budget.get_evaluations_remaining() == float("inf")
+    assert not budget.is_done()
+    budget.raise_exception_if_done() # Should not raise
+    assert budget.get_fraction_consumed() == approx(0.0)
+
+    budget.add_evaluations(1)
+    assert budget.get_time_spent().total_seconds() == approx(0)
+    assert budget.get_time_remaining().total_seconds() == approx(5)
+    assert budget.get_evaluations_spent() == 1
+    assert budget.get_evaluations_remaining() == float("inf")
+    assert not budget.is_done()
+    budget.raise_exception_if_done() # Should not raise
+    assert budget.get_fraction_consumed() == approx(0.0)
+
+    budget.add_time(seconds=2)
+    assert budget.get_time_spent().total_seconds() == approx(2)
+    assert budget.get_time_remaining().total_seconds() == approx(3)
+    assert budget.get_evaluations_spent() == 1
+    assert budget.get_evaluations_remaining() == float("inf")
+    assert not budget.is_done()
+    budget.raise_exception_if_done() # Should not raise
+    assert budget.get_fraction_consumed() == approx(2/5)
+
+    budget.add_time(seconds=4)
+    assert budget.get_time_spent().total_seconds() == approx(6)
+    assert budget.get_time_remaining().total_seconds() == approx(0)
+    assert budget.get_evaluations_spent() == 1
+    assert budget.get_evaluations_remaining() == float("inf")
+    assert budget.is_done()
+    assert pytest.raises(StopCriterionReached, budget.raise_exception_if_done)
+    assert budget.get_fraction_consumed() == 1.0
+
+
 def test_check_tune_params_list():
     tune_params = dict(
         zip(