mirror of
https://github.com/clearml/clearml
synced 2025-01-31 17:17:00 +00:00
118 lines
5.3 KiB
Python
118 lines
5.3 KiB
Python
import io
|
|
import sys
|
|
|
|
from ..config import running_remotely
|
|
from ..debugging.log import LoggerRoot
|
|
|
|
|
|
class PatchHydra(object):
|
|
_original_run_job = None
|
|
_current_task = None
|
|
_config_section = 'OmegaConf'
|
|
_parameter_section = 'Hydra'
|
|
_parameter_allow_full_edit = '_allow_omegaconf_edit_'
|
|
|
|
@classmethod
|
|
def patch_hydra(cls):
|
|
# noinspection PyBroadException
|
|
try:
|
|
# only once
|
|
if cls._original_run_job:
|
|
return True
|
|
# if hydra is not loaded, do not patch anything
|
|
if not sys.modules.get('hydra'):
|
|
return False
|
|
|
|
from hydra.core import utils # noqa
|
|
from hydra._internal import hydra as internal_hydra # noqa
|
|
|
|
# check if hydra is already initialized
|
|
if utils.HydraConfig.initialized():
|
|
LoggerRoot.get_base_logger().warning(
|
|
"Hydra is already loaded storing read-only OmegaConf, "
|
|
"For full support call Task.init(...) before the Hydra App")
|
|
# noinspection PyBroadException
|
|
try:
|
|
# noinspection PyProtectedMember,PyUnresolvedReferences
|
|
PatchHydra._register_omegaconf(utils.HydraConfig.get()._get_root())
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
cls._original_run_job = internal_hydra.Hydra.run
|
|
internal_hydra.Hydra.run = cls._patched_run_job
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
@staticmethod
|
|
def update_current_task(task):
|
|
# set current Task before patching
|
|
PatchHydra._current_task = task
|
|
if not PatchHydra.patch_hydra():
|
|
# if patching failed set it to None
|
|
PatchHydra._current_task = None
|
|
|
|
@staticmethod
|
|
def _patched_run_job(self, config_name, task_function, overrides, *args, **kwargs):
|
|
if not PatchHydra._current_task:
|
|
return PatchHydra._original_run_job(self, config_name, task_function, overrides, *args, **kwargs)
|
|
allow_omegaconf_edit = False
|
|
|
|
def patched_task_function(a_config, *a_args, **a_kwargs):
|
|
from omegaconf import OmegaConf # noqa
|
|
if not running_remotely() or not allow_omegaconf_edit:
|
|
PatchHydra._register_omegaconf(a_config)
|
|
else:
|
|
# noinspection PyProtectedMember
|
|
omega_yaml = PatchHydra._current_task._get_configuration_text(PatchHydra._config_section)
|
|
loaded_config = OmegaConf.load(io.StringIO(omega_yaml))
|
|
a_config = OmegaConf.merge(a_config, loaded_config)
|
|
PatchHydra._register_omegaconf(a_config, is_read_only=False)
|
|
return task_function(a_config, *a_args, **a_kwargs)
|
|
|
|
# store the config
|
|
# noinspection PyBroadException
|
|
try:
|
|
if running_remotely():
|
|
# get the _parameter_allow_full_edit casted back to boolean
|
|
connected_config = dict()
|
|
connected_config[PatchHydra._parameter_allow_full_edit] = False
|
|
PatchHydra._current_task.connect(connected_config, name=PatchHydra._parameter_section)
|
|
allow_omegaconf_edit = connected_config.pop(PatchHydra._parameter_allow_full_edit, None)
|
|
# get all the overrides
|
|
full_parameters = PatchHydra._current_task.get_parameters(backwards_compatibility=False)
|
|
stored_config = {k[len(PatchHydra._parameter_section)+1:]: v for k, v in full_parameters.items()
|
|
if k.startswith(PatchHydra._parameter_section+'/')}
|
|
stored_config.pop(PatchHydra._parameter_allow_full_edit, None)
|
|
overrides = ['{}={}'.format(k, v) for k, v in stored_config.items()]
|
|
else:
|
|
stored_config = dict(arg.split('=', 1) for arg in overrides)
|
|
stored_config[PatchHydra._parameter_allow_full_edit] = False
|
|
PatchHydra._current_task.connect(stored_config, name=PatchHydra._parameter_section)
|
|
# todo: remove the overrides section from the Args (we have it here)
|
|
# PatchHydra._current_task.delete_parameter('Args/overrides')
|
|
except Exception:
|
|
pass
|
|
return PatchHydra._original_run_job(self, config_name, patched_task_function, overrides, *args, **kwargs)
|
|
|
|
@staticmethod
|
|
def _register_omegaconf(config, is_read_only=True):
|
|
from omegaconf import OmegaConf # noqa
|
|
|
|
if is_read_only:
|
|
description = 'Full OmegaConf YAML configuration. ' \
|
|
'This is a read-only section, unless \'{}/{}\' is set to True'.format(
|
|
PatchHydra._parameter_section, PatchHydra._parameter_allow_full_edit)
|
|
else:
|
|
description = 'Full OmegaConf YAML configuration overridden! ({}/{}=True)'.format(
|
|
PatchHydra._parameter_section, PatchHydra._parameter_allow_full_edit)
|
|
|
|
# noinspection PyProtectedMember
|
|
PatchHydra._current_task._set_configuration(
|
|
name=PatchHydra._config_section,
|
|
description=description,
|
|
config_type='OmegaConf YAML',
|
|
config_text=OmegaConf.to_yaml({k: v for k, v in config.items() if k not in ('hydra', )})
|
|
)
|