diff --git a/src/o2tuner/backends.py b/src/o2tuner/backends.py index b639273..1907497 100644 --- a/src/o2tuner/backends.py +++ b/src/o2tuner/backends.py @@ -3,11 +3,13 @@ """ import sys from time import time -from os import getcwd, chdir, remove +from os import getcwd, chdir from os.path import join from inspect import signature from copy import deepcopy import pickle +# import needed to ensure sanity of storage +from sqlalchemy.exc import NoSuchModuleError import optuna @@ -38,7 +40,7 @@ def make_trial_directory(trial): return cwd -def get_storage_identifier(storage_path): +def adjust_storage_url(storage_url, workdir="./"): """ From the full storage path, try to extract the identifier for what to be used """ @@ -46,43 +48,59 @@ def get_storage_identifier(storage_path): storage_prefixes = ["sqlite:///", "mysql:///"] for prefix in storage_prefixes: - if storage_path.find(prefix) == 0: - return prefix + if storage_url.find(prefix) == 0: + if prefix == "sqlite:///": + # in case of SQLite, the file could be requested to be stored under an absolute or relative path + path = storage_url[len(prefix):] + if path[0] != "/": + # if not an absolute path, put the working directory in between + storage_url = prefix + join(workdir, path) + return storage_url - return None + LOG.warning("Unknown storage identifier in URL %s, might fail", storage_url) + + return storage_url + + +def get_default_storage_url(study_name="o2tuner_study"): + """ + Construct a default storage path + """ + return f"sqlite:///{study_name}.db" -def adjust_storage_path(storage_path, workdir="./"): +def create_storage(storage, workdir="./"): """ Make sure the path is either absolute path or relative to the specified workdir. Take care of storage identifier. Right now check for MySQL and SQLite. """ - if not storage_path: + if not storage: # Empty path, cannot know how to deal with it, return None return None - # check if there is a known identifier - check_prefix = get_storage_identifier(storage_path) - - if not check_prefix: - # either no or unknown identifier - return storage_path + # default arguments which we will use + # for now, use a high timeout so we don't fail if another process is currently using the storage backend + engine_kwargs = {"connect_args": {"timeout": 100}} - path = storage_path[len(check_prefix):] - if path[0] == "/": - # Absolute path, just return - return storage_path + if isinstance(storage, str): + # simply treat this as the storage url + url = storage + else: + # first pop the url... + url = storage.pop("url", get_default_storage_url()) + if storage: + # ...then check, if there is more in the dictionary; if so, use it + engine_kwargs = storage - # re-assemble, put the working directory in between - return check_prefix + join(workdir, path) - - -def get_default_storage(study_name): - """ - Construct a default storage path - """ - return f"sqlite:///{study_name}.db" + # check if there is a known identifier + url = adjust_storage_url(url, workdir) + try: + storage = optuna.storages.RDBStorage(url=url, engine_kwargs=engine_kwargs) + return storage + except (ImportError, NoSuchModuleError) as import_error: + LOG.error(import_error) + return None def load_or_create_study_from_storage(study_name, storage, sampler=None, create_if_not_exists=True): @@ -91,17 +109,13 @@ def load_or_create_study_from_storage(study_name, storage, sampler=None, create_ """ try: study = optuna.load_study(study_name=study_name, storage=storage, sampler=sampler) - LOG.debug("Loading existing study %s from storage %s", study_name, storage) + LOG.debug("Loading existing study %s from storage %s", study_name, storage.url) return study except KeyError: if create_if_not_exists: study = optuna.create_study(study_name=study_name, storage=storage, sampler=sampler) - LOG.debug("Creating new study %s at storage %s", study_name, storage) + LOG.debug("Creating new study %s at storage %s", study_name, storage.url) return study - except ImportError as exc: - # Probably cannot import MySQL or SQLite stuff - LOG.warning("Probably cannot import what is needed for database access. Will try to attempt a serial run.") - LOG.warning(exc) return None @@ -140,7 +154,7 @@ def load_or_create_study(study_name=None, storage=None, sampler=None, workdir=". file in the given directory with .pkl. If found, tru to load. If also this does not exist, create a new in-memory study. """ - storage = adjust_storage_path(storage, workdir) + storage = create_storage(storage, workdir) if study_name and storage: # Although optuna would come up with a unique name when study_name is None, # we force a name to be given by the user for those cases @@ -174,25 +188,11 @@ def pickle_study(study, workdir="./"): return file_name -def can_do_storage(storage): +def can_do_storage(storage_url): """ Basically a dry run to try and create a study for given storage """ - identifier = get_storage_identifier(storage) - if not identifier: - LOG.error("Storage %s has unknown identifier, cannot create study.", storage) - return False - filepath = "/tmp/o2tuner_dry_run.db" - if exists_file(filepath): - remove(filepath) - storage = f"{identifier}{filepath}" - can_do, _ = load_or_create_study("o2tuner_dry_study", storage) - if exists_file(filepath): - # E.g. in case of SQLite, remove it - remove(filepath) - if not can_do: - LOG.error("Tested storage via %s, cannot create study at storage %s.", identifier, storage) - return can_do + return create_storage(storage_url) is not None class OptunaHandler: diff --git a/src/o2tuner/optimise.py b/src/o2tuner/optimise.py index fa799a9..9dafeca 100644 --- a/src/o2tuner/optimise.py +++ b/src/o2tuner/optimise.py @@ -9,7 +9,7 @@ import functools from o2tuner.io import make_dir, parse_yaml -from o2tuner.backends import OptunaHandler, can_do_storage, get_default_storage +from o2tuner.backends import OptunaHandler, can_do_storage, get_default_storage_url from o2tuner.sampler import construct_sampler from o2tuner.inspector import O2TunerInspector from o2tuner.exception import O2TunerStopOptimisation @@ -81,13 +81,13 @@ def prepare_optimisation(optuna_config, work_dir="o2tuner_optimise"): if not storage and not in_memory: # make a default storage, optimisation via storage should be the way to go - storage = get_default_storage(study_name) + storage = get_default_storage_url(study_name) if not in_memory and not can_do_storage(storage): # no worries - at this point - if optimisation via storage is not possible optuna_storage_config["storage"] = None if jobs > 1: - # however, if more than 1 one requested, abort the preparation here + # however, if more than 1 job requested, abort the preparation here LOG.error("Requested %d jobs but problem to set up storage %s", jobs, storage) return None, None, None else: