Source code for mlrun.frameworks.auto_mlrun.auto_mlrun

# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
from typing import Callable, Dict, List, Tuple, Type, Union

import mlrun
from mlrun.artifacts import get_model

from .._common import CommonTypes, ModelHandler


[docs]def get_framework_by_instance(model: CommonTypes.ModelType) -> str: """ Get the framework name of the given model by its instance. :param model: The model to get his framework. :return: The model's framework. :raise MLRunInvalidArgumentError: If the given model type is not supported by AutoMLRun or not recognized. """ # TF.Keras: try: from tensorflow.keras import Model from mlrun.frameworks.tf_keras import TFKerasModelHandler if isinstance(model, Model): return TFKerasModelHandler.FRAMEWORK_NAME except ModuleNotFoundError: pass # PyTorch: try: from torch.nn import Module from mlrun.frameworks.pytorch import PyTorchModelHandler if isinstance(model, Module): return PyTorchModelHandler.FRAMEWORK_NAME except ModuleNotFoundError: pass # XGBoost: try: from xgboost import Booster, XGBModel from mlrun.frameworks.xgboost import XGBoostModelHandler if isinstance(model, (XGBModel, Booster)): return XGBoostModelHandler.FRAMEWORK_NAME except ModuleNotFoundError: pass # LightGBM: try: from lightgbm import Booster, LGBMModel from mlrun.frameworks.lgbm import LGBMModelHandler if isinstance(model, (LGBMModel, Booster)): return LGBMModelHandler.FRAMEWORK_NAME except ModuleNotFoundError: pass # ONNX: try: from onnx import ModelProto from mlrun.frameworks.onnx import ONNXModelHandler if isinstance(model, ModelProto): return ONNXModelHandler.FRAMEWORK_NAME except ModuleNotFoundError: pass # SciKit-Learn: (As SKLearn's API is commonly used in other frameworks, its important to check it last) try: import inspect import sklearn.base as sklearn_base from mlrun.frameworks.sklearn import SKLearnModelHandler if isinstance( model, tuple( [ class_tuple[1] for class_tuple in inspect.getmembers(sklearn_base, inspect.isclass) if "sklearn" in str(class_tuple[1]) ] ), ): return SKLearnModelHandler.FRAMEWORK_NAME except ModuleNotFoundError: pass # Unrecognized: raise mlrun.errors.MLRunInvalidArgumentError( f"The type of the model: {type(model)} was not recognized to be of any of the supported frameworks in " f"MLRun or perhaps the required package of the framework could not be imported. If the model's framework is " f"one of: TF.Keras, PyTorch, XGBoost, LightGBM, SciKit-Learn, please provide the 'framework' parameter to the " f"function." )
[docs]def get_framework_by_class_name(model: CommonTypes.ModelType) -> str: """ Get the framework name of the given model by its class name. :param model: The model to get its framework. :return: The model's framework. :raise MLRunInvalidArgumentError: If the given model's class name is not supported by AutoMLRun or not recognized. """ # Read the class name: class_name = str(model.__class__) # Look for the correct framework: if "tensorflow" in class_name: from mlrun.frameworks.tf_keras import TFKerasModelHandler return TFKerasModelHandler.FRAMEWORK_NAME if "torch" in class_name: from mlrun.frameworks.pytorch import PyTorchModelHandler return PyTorchModelHandler.FRAMEWORK_NAME if "xgboost" in class_name: from mlrun.frameworks.xgboost import XGBoostModelHandler return XGBoostModelHandler.FRAMEWORK_NAME if "lightgbm" in class_name: from mlrun.frameworks.lgbm import LGBMModelHandler return LGBMModelHandler.FRAMEWORK_NAME if "onnx" in class_name: from mlrun.frameworks.onnx import ONNXModelHandler return ONNXModelHandler.FRAMEWORK_NAME # As SKLearn's API is commonly used in other frameworks, its important to check it last: if "sklearn" in class_name: from mlrun.frameworks.sklearn import SKLearnModelHandler return SKLearnModelHandler.FRAMEWORK_NAME # Framework was not recognized: raise mlrun.errors.MLRunInvalidArgumentError( f"The model's class name: {class_name} was not recognized to be of any of the supported frameworks in " f"MLRun. If the model's framework is one of: TF.Keras, PyTorch, XGBoost, LightGBM, SciKit-Learn, " f"please provide the 'framework' parameter to the function." )
[docs]def framework_to_model_handler(framework: str) -> Type[ModelHandler]: """ Get the ModelHandler class of the given framework's name. :param framework: The framework's name. :return: The framework's ModelHandler class. :raise MLRunInvalidArgumentError: If the given framework is not supported by AutoMLRun. """ # Match the framework: if framework == "tensorflow.keras": from mlrun.frameworks.tf_keras import TFKerasModelHandler return TFKerasModelHandler if framework == "torch": from mlrun.frameworks.pytorch import PyTorchModelHandler return PyTorchModelHandler if framework == "sklearn": from mlrun.frameworks.sklearn import SKLearnModelHandler return SKLearnModelHandler if framework == "xgboost": from mlrun.frameworks.xgboost import XGBoostModelHandler return XGBoostModelHandler if framework == "lightgbm": from mlrun.frameworks.lgbm import LGBMModelHandler return LGBMModelHandler if framework == "onnx": from mlrun.frameworks.onnx import ONNXModelHandler return ONNXModelHandler # Framework was not recognized: raise mlrun.errors.MLRunInvalidArgumentError( f"The framework {framework} is not supported yet by AutoMLRun. " f"Please use the ModelHandler class of the framework from mlrun.frameworks.{framework}" )
[docs]def framework_to_apply_mlrun(framework: str) -> Callable[..., ModelHandler]: """ Get the 'apply_mlrun' shortcut function of the given framework's name. :param framework: The framework's name. :return: The framework's 'apply_mlrun' shortcut function. :raise MLRunInvalidArgumentError: If the given framework is not supported by AutoMLRun or if it does not have an 'apply_mlrun' yet. """ # Match the framework: if framework == "tensorflow.keras": from mlrun.frameworks.tf_keras import apply_mlrun return apply_mlrun if framework == "torch": raise mlrun.errors.MLRunInvalidArgumentError( "PyTorch has no 'apply_mlrun' shortcut yet. Please use the PyTorchMLRunInterface " "from mlrun.frameworks.pytorch instead." ) if framework == "sklearn": from mlrun.frameworks.sklearn import apply_mlrun return apply_mlrun if framework == "xgboost": from mlrun.frameworks.xgboost import apply_mlrun return apply_mlrun if framework == "lightgbm": from mlrun.frameworks.lgbm import apply_mlrun return apply_mlrun if framework == "onnx": raise mlrun.errors.MLRunInvalidArgumentError( "ONNX has no 'apply_mlrun' shortcut yet. Please use the ONNXMLRunInterface from mlrun.frameworks.onnx " "instead." ) # Framework was not recognized: raise mlrun.errors.MLRunInvalidArgumentError( f"The framework {framework} is not supported yet by AutoMLRun. " f"Please use the ModelHandler class of the framework from mlrun.frameworks.{framework}" )
[docs]class AutoMLRun: """ A library of automatic functions for managing models using MLRun's frameworks package. """ @staticmethod def _get_framework( model: CommonTypes.ModelType = None, model_path: str = None ) -> Union[Tuple[str, dict]]: """ Try to get the framework from the model or model path provided. The framework can be read from the model path only if the model path is of a logged model artifact (store object uri). :param model: The model instance to get its framework. :param model_path: The store object uri of a model artifact. :return: A tuple of: [0] = The framework. [1] = Collected data (model_file, model_artifact, extra_data) if read from model path or an empty dictionary if read from model. :raise MLRunInvalidArgumentError: If both values were None, the framework is not recognized / supported or the the given model path is not of a store object. """ # If the model is provided: if model is not None: try: # Try to get the framework by the class name: return get_framework_by_class_name(model=model), {} except mlrun.errors.MLRunInvalidArgumentError: # Try to get the framework by the instance's type: return get_framework_by_instance(model=model), {} # If the model path is provided: if model_path is not None: # Get the model primary file, extra data and artifact: model_file, model_artifact, extra_data = get_model(model_path) # If the model path is not of a store object, `model_artifact` will be None: if model_artifact is None: raise mlrun.errors.MLRunInvalidArgumentError( f"The model path provided: '{model_path}' is not of a model artifact (store uri) so the framework " f"attribute must be specified." ) # Return the framework and the collected files and artifacts: return ( model_artifact.labels["framework"], { "model_file": model_file, "model_artifact": model_artifact, "extra_data": extra_data, }, ) # Both values were None: raise mlrun.errors.MLRunInvalidArgumentError( "At least 'mode' or 'model_path' must be given in order to use AutoMLRun." )
[docs] @staticmethod def load_model( model_path: str, model_name: str = None, context: mlrun.MLClientCtx = None, modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None, custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None, custom_objects_directory: str = None, framework: str = None, **kwargs, ) -> ModelHandler: """ Load a model using MLRun's ModelHandler. The loaded model can be accessed from the model handler returned via model_handler.model. If the model is a store object uri (it is logged in MLRun) then the framework will be read automatically, otherwise (for local path and urls) it must be given. The other parameters will be automatically read in case its a logged model in MLRun. :param model_path: A store object path of a logged model object in MLRun. :param model_name: The model name to use for storing the model artifact. If not given will have a default name according to the framework. :param modules_map: A dictionary of all the modules required for loading the model. Each key is a path to a module and its value is the object name to import from it. All the modules will be imported globally. If multiple objects needed to be imported from the same module a list can be given. The map can be passed as a path to a json file as well. For example: .. code-block:: python { "module1": None, # import module1 "module2": ["func1", "func2"], # from module2 import func1, func2 "module3.sub_module": "func3", # from module3.sub_module import func3 } If the model path given is of a store object, the modules map will be read from the logged modules map artifact of the model. :param custom_objects_map: A dictionary of all the custom objects required for loading the model. Each key is a path to a python file and its value is the custom object name to import from it. If multiple objects needed to be imported from the same py file a list can be given. The map can be passed as a path to a json file as well. For example: .. code-block:: python { "/.../custom_model.py": "MyModel", "/.../custom_objects.py": ["object1", "object2"] } All the paths will be accessed from the given 'custom_objects_directory', meaning each py file will be read from 'custom_objects_directory/<MAP VALUE>'. If the model path given is of a store object, the custom objects map will be read from the logged custom object map artifact of the model. Notice: The custom objects will be imported in the order they came in this dictionary (or json). If a custom object is depended on another, make sure to put it below the one it relies on. :param custom_objects_directory: Path to the directory with all the python files required for the custom objects. Can be passed as a zip file as well (will be extracted during the run before loading the model). If the model path given is of a store object, the custom objects files will be read from the logged custom object artifact of the model. :param context: A MLRun context. :param framework: The model's framework. It must be provided for local paths or urls. If None, AutoMLRun will assume the model path is of a store uri model artifact and try to get the framework from it. Defaulted to None. :param kwargs: Additional parameters for the specific framework's ModelHandler class. :return: The model inside a MLRun model handler. :raise MLRunInvalidArgumentError: In case the framework is incorrect or missing. """ # Get the model's framework (if needed): collected_model_file_and_artifact = {} if framework is None: framework, collected_model_file_and_artifact = AutoMLRun._get_framework( model_path=model_path ) # Get the ModelHandler according to the framework: model_handler_class = framework_to_model_handler(framework=framework) # Initialize the model handler: model_handler = model_handler_class( model_path=model_path, model_name=model_name, modules_map=modules_map, custom_objects_map=custom_objects_map, custom_objects_directory=custom_objects_directory, context=context, **collected_model_file_and_artifact, **kwargs, ) # Load the model and return the handler: model_handler.load() return model_handler
[docs] @staticmethod def apply_mlrun( model: CommonTypes.ModelType = None, model_name: str = None, tag: str = "", model_path: str = None, modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None, custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None, custom_objects_directory: str = None, context: mlrun.MLClientCtx = None, framework: str = None, auto_log: bool = True, **kwargs, ) -> ModelHandler: """ Use MLRun's 'apply_mlrun' of the detected given model's framework to wrap the framework relevant methods and gain the framework's features in MLRun. A ModelHandler initialized with the model will be returned. :param model: The model to wrap. Can be loaded from the model path given as well. :param model_name: The model name to use for storing the model artifact. If not given will have a default name according to the framework. :param tag: The model's tag to log with. :param model_path: The model's store object path. Mandatory for evaluation (to know which model to update). If model is not provided, it will be loaded from this path. :param modules_map: A dictionary of all the modules required for loading the model. Each key is a path to a module and its value is the object name to import from it. All the modules will be imported globally. If multiple objects needed to be imported from the same module a list can be given. The map can be passed as a path to a json file as well. For example: .. code-block:: python { "module1": None, # import module1 "module2": ["func1", "func2"], # from module2 import func1, func2 "module3.sub_module": "func3", # from module3.sub_module import func3 } If the model path given is of a store object, the modules map will be read from the logged modules map artifact of the model. :param custom_objects_map: A dictionary of all the custom objects required for loading the model. Each key is a path to a python file and its value is the custom object name to import from it. If multiple objects needed to be imported from the same py file a list can be given. The map can be passed as a path to a json file as well. For example: .. code-block:: python { "/.../custom_model.py": "MyModel", "/.../custom_objects.py": ["object1", "object2"] } All the paths will be accessed from the given 'custom_objects_directory', meaning each py file will be read from 'custom_objects_directory/<MAP VALUE>'. If the model path given is of a store object, the custom objects map will be read from the logged custom object map artifact of the model. Notice: The custom objects will be imported in the order they came in this dictionary (or json). If a custom object is depended on another, make sure to put it below the one it relies on. :param custom_objects_directory: Path to the directory with all the python files required for the custom objects. Can be passed as a zip file as well (will be extracted during the run before loading the model). If the model path given is of a store object, the custom objects files will be read from the logged custom object artifact of the model. :param context: A MLRun context. :param auto_log: Whether to enable auto-logging capabilities of MLRun or not. Auto logging will add default artifacts and metrics besides the one you can pass here. :param framework: The model's framework. If None, AutoMLRun will try to figure out the framework. From the provided model or model path. Defaulted to None. :param kwargs: Additional parameters for the specific framework's 'apply_mlrun' function like metrics, callbacks and more (read the docs of the required framework to know more). :return: The framework's model handler initialized with the given model. """ # Get the model's framework: collected_model_file_and_artifact = {} if framework is None: framework, collected_model_file_and_artifact = AutoMLRun._get_framework( model=model, model_path=model_path ) # Get the framework's 'apply_mlrun': apply_mlrun = framework_to_apply_mlrun(framework=framework) # Initialize the model handler: return apply_mlrun( model=model, model_name=model_name, tag=tag, model_path=model_path, modules_map=modules_map, custom_objects_map=custom_objects_map, custom_objects_directory=custom_objects_directory, context=context, auto_log=auto_log, model_handler_kwargs=collected_model_file_and_artifact, **kwargs, )