Add joblib (equivalent to scikit) to Task.init() auto_connect_frameworks argument

This commit is contained in:
allegroai 2021-09-25 23:05:12 +03:00
parent 5f390f3a54
commit 5a5b9488e8

View File

@ -357,7 +357,7 @@ class Task(_Task):
auto_connect_frameworks={
'matplotlib': True, 'tensorflow': True, 'tensorboard': True, 'pytorch': True,
'xgboost': True, 'scikit': True, 'fastai': True, 'lightgbm': True,
'hydra': True, 'detect_repository': True, 'tfdefines': True,
'hydra': True, 'detect_repository': True, 'tfdefines': True, 'joblib': True,
}
:param bool auto_resource_monitoring: Automatically create machine resource monitoring plots
@ -545,7 +545,9 @@ class Task(_Task):
is_auto_connect_frameworks_bool = not isinstance(auto_connect_frameworks, dict)
if is_auto_connect_frameworks_bool or auto_connect_frameworks.get('hydra', True):
PatchHydra.update_current_task(task)
if is_auto_connect_frameworks_bool or auto_connect_frameworks.get('scikit', True):
if is_auto_connect_frameworks_bool or (
auto_connect_frameworks.get('scikit', True) and
auto_connect_frameworks.get('joblib', True)):
PatchedJoblib.update_current_task(task)
if is_auto_connect_frameworks_bool or auto_connect_frameworks.get('matplotlib', True):
PatchedMatplotlib.update_current_task(Task.__main_task)