Skip to content

Commit

Permalink
Consolidate user stages
Browse files Browse the repository at this point in the history
* also no optimisation are fine

* accept different function signatures of user python functions

* if optuna.samplers.GridSampler is used, attach the search space to the
  user config with key O2TUNER_GridSampler_SearchSpace
  • Loading branch information
Benedikt Volkel committed Aug 4, 2023
1 parent 92b519c commit 06da0ac
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 2 deletions.
11 changes: 11 additions & 0 deletions src/o2tuner/backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from os import getcwd, chdir, remove
from os.path import join
from inspect import signature
from copy import deepcopy
import pickle

import optuna
Expand Down Expand Up @@ -224,6 +225,14 @@ def __init__(self, db_study_name=None, db_storage=None, workdir=None, user_confi
# Flag to indicate if this is an in-memory run
self.in_memory = in_memory

def finalise_configuration(self):
"""
Pass in some O2Tuner related configurations that might be interesting for the user
"""
sampler = self._study.sampler
if isinstance(sampler, optuna.sampler.GridSampler):
self.user_config["O2TUNER_GridSampler_SearchSpace"] = deepcopy(sampler._search_space)

def objective_cwd_wrapper(self, trial):
"""
If this trial needs a dedicated cwd, create it, change into it, run the objective and go back
Expand Down Expand Up @@ -257,6 +266,8 @@ def optimise(self):
if not self._n_trials or self._n_trials < 0 or not self._objective:
LOG.error("Not initialised: Number of trials and objective function need to be set")
return
# add some special configuration
self.finalise_configuration()
try:
self._study.optimize(self.objective_cwd_wrapper, n_trials=self._n_trials)
except O2TunerStopOptimisation:
Expand Down
22 changes: 20 additions & 2 deletions src/o2tuner/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
"""
from os.path import join, abspath, expanduser, dirname
from os import getcwd, chdir
from inspect import signature

from o2tuner.system import run_command, import_function_from_file, get_signal_handler
from o2tuner.optimise import optimise
Expand Down Expand Up @@ -35,7 +36,7 @@ def run_cmd_or_python(cwd, name, config, stages_optimisation):
func = import_function_from_file(config["python"]["file"], config["python"]["entrypoint"])
# see if we need to pass in any inspectors with loaded optimisations
inspectors = []
for optimisation in config["optimisations"]:
for optimisation in config.get("optimisations", []):
if optimisation not in stages_optimisation:
LOG.warning("Optimisation stage %s not defined, cannot construct inspector for that. Skip...", optimisation)
continue
Expand All @@ -50,7 +51,24 @@ def run_cmd_or_python(cwd, name, config, stages_optimisation):
# change to this cwd and afterwards back
chdir(cwd)
pass_config = config.get("config", None)
ret = func(inspectors, pass_config)

# now we check the signature of the function
sig = signature(func)
n_params = len(sig.parameters)
if n_params == 1:
if inspectors:
LOG.error("Your function signature has no place to pass in the optimisations. Following optimisations were given")
for optimisation in config.get("optimisations", []):
LOG.append_log(optimisation)
ret = 1
ret = func(pass_config)
elif n_params == 2:
if not not inspectors:
LOG.warning("Given the function signature, your user stage %s might expect inspectors to be passed in", name)
ret = func(inspectors, pass_config)
else:
LOG.error("The function signature of user stage %s is wrong. Need either 1 (config) or 2 (inspectors, config) arguments.")
ret = 1
chdir(this_dir)
return ret

Expand Down

0 comments on commit 06da0ac

Please sign in to comment.