mirror of
https://github.com/clearml/clearml
synced 2025-01-31 17:17:00 +00:00
132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
import six
|
|
from pathlib2 import Path
|
|
|
|
from ..binding.frameworks import _patched_call, _Empty, WeightsFileHandler
|
|
from ..config import running_remotely
|
|
from ..debugging.log import LoggerRoot
|
|
from ..model import Framework
|
|
|
|
|
|
class PatchedJoblib(object):
|
|
_patched_joblib = False
|
|
_current_task = None
|
|
|
|
@staticmethod
|
|
def patch_joblib():
|
|
if PatchedJoblib._patched_joblib:
|
|
# We don't need to patch anything else, so we are done
|
|
return True
|
|
|
|
# whatever happens we should not retry to patch it
|
|
PatchedJoblib._patched_joblib = True
|
|
# noinspection PyBroadException
|
|
try:
|
|
try:
|
|
import joblib
|
|
except ImportError:
|
|
joblib = None
|
|
|
|
try:
|
|
from sklearn.externals import joblib as sk_joblib
|
|
except ImportError:
|
|
sk_joblib = None
|
|
|
|
if joblib:
|
|
joblib.dump = _patched_call(joblib.dump, PatchedJoblib._dump)
|
|
joblib.load = _patched_call(joblib.load, PatchedJoblib._load)
|
|
if sk_joblib:
|
|
sk_joblib.dump = _patched_call(sk_joblib.dump, PatchedJoblib._dump)
|
|
sk_joblib.load = _patched_call(sk_joblib.load, PatchedJoblib._load)
|
|
|
|
except Exception:
|
|
return False
|
|
return True
|
|
|
|
@staticmethod
|
|
def update_current_task(task):
|
|
if PatchedJoblib.patch_joblib():
|
|
PatchedJoblib._current_task = task
|
|
|
|
@staticmethod
|
|
def _dump(original_fn, obj, f, *args, **kwargs):
|
|
ret = original_fn(obj, f, *args, **kwargs)
|
|
if not PatchedJoblib._current_task:
|
|
return ret
|
|
|
|
if isinstance(f, six.string_types):
|
|
filename = f
|
|
elif hasattr(f, 'name'):
|
|
filename = f.name
|
|
# noinspection PyBroadException
|
|
try:
|
|
f.flush()
|
|
except Exception:
|
|
pass
|
|
else:
|
|
filename = None
|
|
|
|
# give the model a descriptive name based on the file name
|
|
# noinspection PyBroadException
|
|
try:
|
|
model_name = Path(filename).stem
|
|
except Exception:
|
|
model_name = None
|
|
current_framework = PatchedJoblib.get_model_framework(obj)
|
|
WeightsFileHandler.create_output_model(obj, filename, current_framework,
|
|
PatchedJoblib._current_task, singlefile=True, model_name=model_name)
|
|
return ret
|
|
|
|
@staticmethod
|
|
def _load(original_fn, f, *args, **kwargs):
|
|
if isinstance(f, six.string_types):
|
|
filename = f
|
|
elif hasattr(f, 'name'):
|
|
filename = f.name
|
|
else:
|
|
filename = None
|
|
|
|
if not PatchedJoblib._current_task:
|
|
return original_fn(f, *args, **kwargs)
|
|
|
|
# register input model
|
|
empty = _Empty()
|
|
# Hack: disabled
|
|
if False and running_remotely():
|
|
# we assume scikit-learn, for the time being
|
|
current_framework = Framework.scikitlearn
|
|
filename = WeightsFileHandler.restore_weights_file(empty, filename, current_framework,
|
|
PatchedJoblib._current_task)
|
|
model = original_fn(filename or f, *args, **kwargs)
|
|
else:
|
|
# try to load model before registering, in case we fail
|
|
model = original_fn(f, *args, **kwargs)
|
|
current_framework = PatchedJoblib.get_model_framework(model)
|
|
WeightsFileHandler.restore_weights_file(empty, filename, current_framework,
|
|
PatchedJoblib._current_task)
|
|
|
|
if empty.trains_in_model:
|
|
# noinspection PyBroadException
|
|
try:
|
|
model.trains_in_model = empty.trains_in_model
|
|
except Exception:
|
|
pass
|
|
return model
|
|
|
|
@staticmethod
|
|
def get_model_framework(obj):
|
|
object_orig_module = obj.__module__
|
|
framework = Framework.scikitlearn
|
|
try:
|
|
model = object_orig_module.partition(".")[0]
|
|
if model == 'sklearn':
|
|
framework = Framework.scikitlearn
|
|
elif model == 'xgboost':
|
|
framework = Framework.xgboost
|
|
else:
|
|
framework = Framework.scikitlearn
|
|
except Exception as _:
|
|
LoggerRoot.get_base_logger().debug(
|
|
"Can't get model framework {}, model framework will be: {} ".format(object_orig_module, framework))
|
|
finally:
|
|
return framework
|