diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index 54b742a2..c1d72608 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -780,7 +780,7 @@ def tune_kernel( def preprocess_cache(filepath): if isinstance(filepath, Path): filepath = str(filepath.resolve()) - if filepath[-5:] != ".json": + if filepath[-5:] != ".json" and filepath[-8:] != ".json.gz": filepath += ".json" return filepath @@ -877,11 +877,26 @@ def preprocess_cache(filepath): def tune_cache( - cache_path, + cachefile, restrictions=None, **kwargs, ): - cache = util.read_cache(cache_path, open_cache=False) + """ Simulate a tuning session based on a Kernel Tuner cache file + + See tune_kernel for full documentation of options. + + :param cachefile: filename or path to the cachefile + :type cachefile: string or Path + + :param restrictions: search space restrictions / constraints, if not + passed these are simply inferred from the cachefile, which may be + much less efficient. Default: None + :type restrictions: list of string expressions or lambdas + + """ + if isinstance(cachefile, Path): + cachefile = str(cachefile.resolve()) + cache = util.read_cache(cachefile, open_cache=False) tune_args = util.infer_args_from_cache(cache) _restrictions = [util.infer_restrictions_from_cache(cache)] @@ -894,7 +909,7 @@ def tune_cache( tune_args.update(kwargs) - return tune_kernel(**tune_args, cache=cache_path, restrictions=_restrictions, simulation_mode=True) + return tune_kernel(**tune_args, cache=cachefile, restrictions=_restrictions, simulation_mode=True) _run_kernel_docstring = ( diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index 5ceb3cce..a917b68f 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -2,6 +2,7 @@ import ast import errno +import gzip import json import logging import os @@ -1380,6 +1381,13 @@ def check_matching_problem_size(cached_problem_size, problem_size): ) +def open_cachefile(cachefile, mode): + """ Open a cachefile for reading or writing (depending on mode) """ + if cachefile[-3:] == ".gz": + return gzip.open(cachefile, mode) + return open(cachefile, mode) + + def process_cache(cachefile, kernel_options, tuning_options, runner): """Cache file for storing tuned configurations. @@ -1391,7 +1399,8 @@ def process_cache(cachefile, kernel_options, tuning_options, runner): kernel_name: "name of kernel" problem_size: (int, int, int) tune_params_keys: list - tune_params: + tune_params: dict + objective: string cache: { "x1,x2,..xN": {"block_size_x": x1, ..., time=0.234342}, "y1,y2,..yN": {"block_size_x": y1, ..., time=0.134233}, @@ -1427,7 +1436,7 @@ def process_cache(cachefile, kernel_options, tuning_options, runner): contents = json.dumps(c, cls=NpEncoder, indent="")[:-3] # except the last "}\n}" # write the header to the cachefile - with open(cachefile, "w") as f: + with open_cachefile(cachefile, "wt") as f: f.write(contents) return {} @@ -1471,7 +1480,7 @@ def process_cache(cachefile, kernel_options, tuning_options, runner): def correct_open_cache(cachefile, open_cache=True): """If cache file was not properly closed, pretend it was properly closed.""" - with open(cachefile, "r") as f: + with open_cachefile(cachefile, "rt") as f: filestr = f.read().strip() # if file was not properly closed, pretend it was properly closed @@ -1484,7 +1493,7 @@ def correct_open_cache(cachefile, open_cache=True): else: if open_cache: # if it was properly closed, open it for appending new entries - with open(cachefile, "w") as f: + with open_cachefile(cachefile, "wt") as f: f.write(filestr[:-3] + ",") return filestr @@ -1519,12 +1528,12 @@ def close_cache(cachefile): if not os.path.isfile(cachefile): raise ValueError("close_cache expects cache file to exist") - with open(cachefile, "r") as fh: + with open_cachefile(cachefile, "rt") as fh: contents = fh.read() # close to file to make sure it can be read by JSON parsers if contents[-1] == ",": - with open(cachefile, "w") as fh: + with open_cachefile(cachefile, "wt") as fh: fh.write(contents[:-1] + "}\n}") @@ -1542,14 +1551,14 @@ def store_cache(key, params, cachefile, cache): output_params[k] = str(v) if cachefile: - with open(cachefile, "a") as f: + with open_cachefile(cachefile, "at") as f: f.write("\n" + json.dumps({key: output_params}, cls=NpEncoder)[1:-1] + ",") def dump_cache(obj: str, tuning_options): """Dumps a string in the cache, this omits the several checks of store_cache() to speed up the process - with great power comes great responsibility!""" if isinstance(tuning_options.cache, dict) and tuning_options.cachefile: - with open(tuning_options.cachefile, "a") as cachefile: + with open(tuning_options.cachefile, "at") as cachefile: cachefile.write(obj) diff --git a/test/test_cache_file.json.gz b/test/test_cache_file.json.gz new file mode 100644 index 00000000..80dc1bbf Binary files /dev/null and b/test/test_cache_file.json.gz differ diff --git a/test/test_runners.py b/test/test_runners.py index a88ee6e6..daadc743 100644 --- a/test/test_runners.py +++ b/test/test_runners.py @@ -13,6 +13,7 @@ cache_filename = os.path.dirname( os.path.realpath(__file__)) + "/test_cache_file.json" +cache_filename_zipped = cache_filename + ".gz" @pytest.fixture def env(): @@ -135,6 +136,14 @@ def test_tune_cache(env): assert len(results) > 10 +def test_zipped_cache_file(env): + result, _ = tune_kernel(*env, + verbose=True, + cache=cache_filename_zipped, + simulation_mode=True) + assert len(result) > 0 + + def test_constraint_aware_GA(env): options = dict(method="uniform", constraint_aware=True,