From 06da0ac4561f4eff104783209cd868445bc65f9a Mon Sep 17 00:00:00 2001
From: Benedikt Volkel <benedikt.volkel@cern.ch>
Date: Wed, 2 Aug 2023 10:41:16 +0200
Subject: [PATCH] Consolidate user stages

* 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
---
 src/o2tuner/backends.py | 11 +++++++++++
 src/o2tuner/run.py      | 22 ++++++++++++++++++++--
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/o2tuner/backends.py b/src/o2tuner/backends.py
index 10a2bbe..290bcf3 100644
--- a/src/o2tuner/backends.py
+++ b/src/o2tuner/backends.py
@@ -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
@@ -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
@@ -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:
diff --git a/src/o2tuner/run.py b/src/o2tuner/run.py
index 8a78f96..2e67823 100644
--- a/src/o2tuner/run.py
+++ b/src/o2tuner/run.py
@@ -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
@@ -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
@@ -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