# Copyright 2023 Iguazio
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# flake8: noqa - this is until we take care of the F401 violations with respect to __all__ & sphinx
from typing import Any, Union
from tensorflow import keras
import mlrun
from .callbacks import MLRunLoggingCallback, TensorboardLoggingCallback
from .mlrun_interface import TFKerasMLRunInterface
from .model_handler import TFKerasModelHandler
from .model_server import TFKerasModelServer
from .utils import TFKerasTypes, TFKerasUtils
[docs]def apply_mlrun(
model: keras.Model = None,
model_name: str = None,
tag: str = "",
model_path: str = None,
model_format: str = TFKerasModelHandler.ModelFormats.SAVED_MODEL,
save_traces: bool = False,
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,
auto_log: bool = True,
tensorboard_directory: str = None,
mlrun_callback_kwargs: dict[str, Any] = None,
tensorboard_callback_kwargs: dict[str, Any] = None,
use_horovod: bool = None,
**kwargs,
) -> TFKerasModelHandler:
"""
Wrap the given model with MLRun's interface providing it with mlrun's additional features.
: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, the
tf.keras.Model.name will be used.
: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 model_format: The format to use for saving and loading the model. Should be passed as a
member of the class 'ModelFormats'. Default: 'ModelFormats.SAVED_MODEL'.
:param save_traces: Whether or not to use functions saving (only available for the 'SavedModel'
format) for loading the model later without the custom objects dictionary. Only
from tensorflow version >= 2.4.0. Using this setting will increase the model
saving size.
: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_optimizer.py": "optimizer",
"/.../custom_layers.py": ["layer1", "layer2"],
}
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: MLRun context to work with. If no context is given it will be retrieved via
'mlrun.get_or_create_ctx(None)'
:param auto_log: Whether or not to apply MLRun's auto logging on the model. Default: True.
:param tensorboard_directory: If context is not given, or if wished to set the directory even with context,
this will be the output for the event logs of tensorboard. If not given, the
'tensorboard_dir' parameter will be tried to be taken from the provided context.
If not found in the context, the default tensorboard output directory will be:
/User/.tensorboard/<PROJECT_NAME> or if working on local, the set artifacts
path.
:param mlrun_callback_kwargs: Key word arguments for the MLRun callback. For further information see the
documentation of the class 'MLRunLoggingCallback'. Note that both 'context'
and 'auto_log' parameters are already given here.
:param tensorboard_callback_kwargs: Key word arguments for the tensorboard callback. For further information see
the documentation of the class 'TensorboardLoggingCallback'. Note that both
'context' and 'auto_log' parameters are already given here.
:param use_horovod: Whether or not to use horovod - a distributed training framework. Default:
None, meaning it will be read from context if available and if not - False.
:return: The model with MLRun's interface.
"""
# Get parameters defaults:
# # Context:
if context is None:
context = mlrun.get_or_create_ctx(TFKerasMLRunInterface.DEFAULT_CONTEXT_NAME)
# # Use horovod:
if use_horovod is None:
use_horovod = (
context.labels.get("kind", "") == "mpijob" if context is not None else False
)
# Create a model handler:
model_handler_kwargs = (
kwargs.pop("model_handler_kwargs") if "model_handler_kwargs" in kwargs else {}
)
handler = TFKerasModelHandler(
model_name=model_name,
model_path=model_path,
model=model,
model_format=model_format,
save_traces=save_traces,
context=context,
modules_map=modules_map,
custom_objects_map=custom_objects_map,
custom_objects_directory=custom_objects_directory,
**model_handler_kwargs,
)
# Load the model if it was not provided:
if model is None:
handler.load()
model = handler.model
# Add MLRun's interface to the model:
TFKerasMLRunInterface.add_interface(obj=model)
# Initialize horovod if needed:
if use_horovod is True:
model.use_horovod()
# Add auto-logging if needed:
if auto_log:
# Set the kwargs dictionaries defaults:
mlrun_callback_kwargs = (
{} if mlrun_callback_kwargs is None else mlrun_callback_kwargs
)
tensorboard_callback_kwargs = (
{} if tensorboard_callback_kwargs is None else tensorboard_callback_kwargs
)
# Add the logging callbacks with the provided parameters:
model.add_logging_callback(
logging_callback=MLRunLoggingCallback(
context=context,
model_handler=handler,
log_model_tag=tag,
auto_log=auto_log,
**mlrun_callback_kwargs,
)
)
model.add_logging_callback(
logging_callback=TensorboardLoggingCallback(
context=context,
tensorboard_directory=tensorboard_directory,
auto_log=auto_log,
**tensorboard_callback_kwargs,
)
)
return handler
