def patch_class_tree(klass):
"""
Patches all subclasses that override any auto-loggable method via monkey patching using
the gorilla package, taking the argument as the tree root in the class hierarchy. Every
auto-loggable method found in any of the subclasses is replaced by the patched version.
:param klass: root in the class hierarchy to be analyzed and patched recursively
"""
# TODO: add more autologgable methods here (e.g. fit_regularized, from_formula, etc)
# See https://www.statsmodels.org/dev/api.html
autolog_supported_func = {"fit": wrapper_fit}
glob_subclasses = set(find_subclasses(klass))
# Create a patch for every method that needs to be patched, i.e. those
# which actually override an autologgable method
patches_list = [
# Link the patched function with the original via a local variable in the closure
# to allow invoking superclass methods in the context of the subclass, and not
# losing the trace of the true original method
(clazz, method_name, wrapper_func)
for clazz in glob_subclasses
for (method_name, wrapper_func) in autolog_supported_func.items()
if overrides(clazz, method_name)
]
for clazz, method_name, patch_impl in patches_list:
safe_patch(FLAVOR_NAME, clazz, method_name, patch_impl, manage_run=True)
def autolog(disable=False): # pylint: disable=unused-argument
"""
Enables (or disables) and configures logging of Spark datasource paths, versions
(if applicable), and formats when they are read. This method is not threadsafe and assumes a
`SparkSession
<https://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.SparkSession>`_
already exists with the
`mlflow-spark JAR
<http://mlflow.org/docs/latest/tracking.html#automatic-logging-from-spark-experimental>`_
attached. It should be called on the Spark driver, not on the executors (i.e. do not call
this method within a function parallelized by Spark). This API requires Spark 3.0 or above.
Datasource information is cached in memory and logged to all subsequent MLflow runs,
including the active MLflow run (if one exists when the data is read). Note that autologging of
Spark ML (MLlib) models is not currently supported via this API. Datasource autologging is
best-effort, meaning that if Spark is under heavy load or MLflow logging fails for any reason
(e.g., if the MLflow server is unavailable), logging may be dropped.
For any unexpected issues with autologging, check Spark driver and executor logs in addition
to stderr & stdout generated from your MLflow code - datasource information is pulled from
Spark, so logs relevant to debugging may show up amongst the Spark logs.
.. code-block:: python
:caption: Example
import mlflow.spark
import os
import shutil
from pyspark.sql import SparkSession
# Create and persist some dummy data
# Note: On environments like Databricks with pre-created SparkSessions,
# ensure the org.mlflow:mlflow-spark:1.11.0 is attached as a library to
# your cluster
spark = (SparkSession.builder
.config("spark.jars.packages", "org.mlflow:mlflow-spark:1.11.0")
.master("local[*]")
.getOrCreate())
df = spark.createDataFrame([
(4, "spark i j k"),
(5, "l m n"),
(6, "spark hadoop spark"),
(7, "apache hadoop")], ["id", "text"])
import tempfile
tempdir = tempfile.mkdtemp()
df.write.csv(os.path.join(tempdir, "my-data-path"), header=True)
# Enable Spark datasource autologging.
mlflow.spark.autolog()
loaded_df = spark.read.csv(os.path.join(tempdir, "my-data-path"),
header=True, inferSchema=True)
# Call toPandas() to trigger a read of the Spark datasource. Datasource info
# (path and format) is logged to the current active run, or the
# next-created MLflow run if no run is currently active
with mlflow.start_run() as active_run:
pandas_df = loaded_df.toPandas()
:param disable: If ``True``, disables the Spark datasource autologging integration.
If ``False``, enables the Spark datasource autologging integration.
"""
from mlflow.utils._spark_utils import _get_active_spark_session
from mlflow._spark_autologging import _listen_for_spark_activity
from pyspark.sql import SparkSession
from pyspark import SparkContext
def __init__(original, self, *args, **kwargs):
original(self, *args, **kwargs)
_listen_for_spark_activity(self._sc)
safe_patch(FLAVOR_NAME, SparkSession, "__init__", __init__, manage_run=False)
active_session = _get_active_spark_session()
if active_session is not None:
# We know SparkContext exists here already, so get it
sc = SparkContext.getOrCreate()
_listen_for_spark_activity(sc)
def autolog(log_models=True,
disable=False,
exclusive=False,
disable_for_unsupported_versions=False): # pylint: disable=unused-argument
"""
Enable automatic logging from Fastai to MLflow.
Logs loss and any other metrics specified in the fit
function, and optimizer data as parameters. Model checkpoints
are logged as artifacts to a 'models' directory.
MLflow will also log the parameters of the
`EarlyStoppingCallback <https://docs.fast.ai/callbacks.html#EarlyStoppingCallback>`_
and `OneCycleScheduler <https://docs.fast.ai/callbacks.html#OneCycleScheduler>`_ callbacks
:param log_models: If ``True``, trained models are logged as MLflow model artifacts.
If ``False``, trained models are not logged.
:param disable: If ``True``, disables the Fastai autologging integration. If ``False``,
enables the Fastai autologging integration.
:param exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
If ``False``, autologged content is logged to the active fluent run,
which may be user-created.
:param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
fastai that have not been tested against this version of the MLflow client
or are incompatible.
.. code-block:: python
:caption: Example
# This is a modified example from
# https://github.com/mlflow/mlflow/tree/master/examples/fastai
# demonstrating autolog capabilites.
import fastai.vision as vis
import mlflow.fastai
from mlflow.tracking import MlflowClient
def print_auto_logged_info(r):
tags = {k: v for k, v in r.data.tags.items() if not k.startswith("mlflow.")}
artifacts = [f.path for f in MlflowClient().list_artifacts(r.info.run_id, "model")]
print("run_id: {}".format(r.info.run_id))
print("artifacts: {}".format(artifacts))
print("params: {}".format(r.data.params))
print("metrics: {}".format(r.data.metrics))
print("tags: {}".format(tags))
def main(epochs=5, learning_rate=0.01):
# Download and untar the MNIST data set
path = vis.untar_data(vis.URLs.MNIST_SAMPLE)
# Prepare, transform, and normalize the data
data = vis.ImageDataBunch.from_folder(path, ds_tfms=(vis.rand_pad(2, 28), []), bs=64)
data.normalize(vis.imagenet_stats)
# Create CNN the Learner model
model = vis.cnn_learner(data, vis.models.resnet18, metrics=vis.accuracy)
# Enable auto logging
mlflow.fastai.autolog()
# Start MLflow session
with mlflow.start_run() as run:
model.fit(epochs, learning_rate)
# fetch the auto logged parameters, metrics, and artifacts
print_auto_logged_info(mlflow.get_run(run_id=run.info.run_id))
main()
.. code-block:: text
:caption: output
run_id: 5a23dcbcaa334637814dbce7a00b2f6a
artifacts: ['model/MLmodel', 'model/conda.yaml', 'model/model.fastai']
params: {'wd': 'None',
'bn_wd': 'True',
'opt_func': 'Adam',
'epochs': '5', '
train_bn': 'True',
'num_layers': '60',
'lr': '0.01',
'true_wd': 'True'}
metrics: {'train_loss': 0.024,
'accuracy': 0.99214,
'valid_loss': 0.021}
# Tags model summary omitted too long
tags: {...}
.. figure:: ../_static/images/fastai_autolog.png
Fastai autologged MLflow entities
"""
from fastai.basic_train import LearnerCallback, Learner
from fastai.callbacks.hooks import model_summary, layers_info
from fastai.callbacks import EarlyStoppingCallback, OneCycleScheduler
def getFastaiCallback(metrics_logger, learner):
class __MLflowFastaiCallback(LearnerCallback,
metaclass=ExceptionSafeClass):
"""
Callback for auto-logging metrics and parameters.
Records model structural information as params when training begins
"""
def __init__(self, learner):
super().__init__(learner)
self.learner = learner
self.opt = self.learn.opt
self.metrics_names = ["train_loss", "valid_loss"] + [
o.__name__ for o in learner.metrics
]
def on_epoch_end(self, **kwargs):
"""
Log loss and other metrics values after each epoch
"""
if kwargs["smooth_loss"] is None or kwargs[
"last_metrics"] is None:
return
epoch = kwargs["epoch"]
metrics = [kwargs["smooth_loss"]] + kwargs["last_metrics"]
metrics = map(float, metrics)
metrics = dict(zip(self.metrics_names, metrics))
metrics_logger.record_metrics(metrics, epoch)
def on_train_begin(self, **kwargs):
info = layers_info(self.learner)
try_mlflow_log(mlflow.log_param, "num_layers", len(info))
try_mlflow_log(mlflow.log_param, "opt_func",
self.opt_func.func.__name__)
if hasattr(self.opt, "true_wd"):
try_mlflow_log(mlflow.log_param, "true_wd",
self.opt.true_wd)
if hasattr(self.opt, "bn_wd"):
try_mlflow_log(mlflow.log_param, "bn_wd", self.opt.bn_wd)
if hasattr(self.opt, "train_bn"):
try_mlflow_log(mlflow.log_param, "train_bn", self.train_bn)
summary = model_summary(self.learner)
tempdir = tempfile.mkdtemp()
try:
summary_file = os.path.join(tempdir, "model_summary.txt")
with open(summary_file, "w") as f:
f.write(summary)
try_mlflow_log(mlflow.log_artifact,
local_path=summary_file)
finally:
shutil.rmtree(tempdir)
def on_train_end(self, **kwargs):
if log_models:
try_mlflow_log(log_model,
self.learner,
artifact_path="model")
return __MLflowFastaiCallback(learner)
def _find_callback_of_type(callback_type, callbacks):
for callback in callbacks:
if isinstance(callback, callback_type):
return callback
return None
def _log_early_stop_callback_params(callback):
if callback:
try:
earlystopping_params = {
"early_stop_monitor": callback.monitor,
"early_stop_min_delta": callback.min_delta,
"early_stop_patience": callback.patience,
"early_stop_mode": callback.mode,
}
try_mlflow_log(mlflow.log_params, earlystopping_params)
except Exception: # pylint: disable=W0703
return
def _log_one_cycle_callback_params(callback):
if callback:
try:
params = {
"lr_max": callback.lr_max,
"div_factor": callback.div_factor,
"pct_start": callback.pct_start,
"final_div": callback.final_div,
"tot_epochs": callback.tot_epochs,
"start_epoch": callback.start_epoch,
"moms": callback.moms,
}
try_mlflow_log(mlflow.log_params, params)
except Exception: # pylint: disable=W0703
return
def _run_and_log_function(self, original, args, kwargs, unlogged_params,
callback_arg_index):
log_fn_args_as_params(original, list(args), kwargs, unlogged_params)
callbacks = [cb(self)
for cb in self.callback_fns] + (self.callbacks or [])
run_id = mlflow.active_run().info.run_id
with batch_metrics_logger(run_id) as metrics_logger:
mlflowFastaiCallback = getFastaiCallback(metrics_logger, self)
# Checking if the 'callback' argument of the function is set
if len(args) > callback_arg_index:
tmp_list = list(args)
callbacks += list(args[callback_arg_index])
tmp_list[callback_arg_index] += [mlflowFastaiCallback]
args = tuple(tmp_list)
elif kwargs.get("callbacks"):
callbacks += list(kwargs["callbacks"])
kwargs["callbacks"] += [mlflowFastaiCallback]
else:
kwargs["callbacks"] = [mlflowFastaiCallback]
early_stop_callback = _find_callback_of_type(
EarlyStoppingCallback, callbacks)
one_cycle_callback = _find_callback_of_type(
OneCycleScheduler, callbacks)
_log_early_stop_callback_params(early_stop_callback)
_log_one_cycle_callback_params(one_cycle_callback)
result = original(self, *args, **kwargs)
return result
def fit(original, self, *args, **kwargs):
unlogged_params = ["self", "callbacks", "learner"]
return _run_and_log_function(self, original, args, kwargs,
unlogged_params, 3)
safe_patch(FLAVOR_NAME, Learner, "fit", fit, manage_run=True)
def autolog(
log_input_examples=False,
log_model_signatures=True,
log_models=True,
disable=False,
exclusive=False,
disable_for_unsupported_versions=False,
): # pylint: disable=unused-argument
"""
Enables (or disables) and configures autologging for scikit-learn estimators.
**When is autologging performed?**
Autologging is performed when you call:
- ``estimator.fit()``
- ``estimator.fit_predict()``
- ``estimator.fit_transform()``
**Logged information**
**Parameters**
- Parameters obtained by ``estimator.get_params(deep=True)``. Note that ``get_params``
is called with ``deep=True``. This means when you fit a meta estimator that chains
a series of estimators, the parameters of these child estimators are also logged.
**Metrics**
- A training score obtained by ``estimator.score``. Note that the training score is
computed using parameters given to ``fit()``.
- Common metrics for classifier:
- `precision score`_
.. _precision score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.precision_score.html
- `recall score`_
.. _recall score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.recall_score.html
- `f1 score`_
.. _f1 score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html
- `accuracy score`_
.. _accuracy score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.accuracy_score.html
If the classifier has method ``predict_proba``, we additionally log:
- `log loss`_
.. _log loss:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.log_loss.html
- `roc auc score`_
.. _roc auc score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.roc_auc_score.html
- Common metrics for regressor:
- `mean squared error`_
.. _mean squared error:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html
- root mean squared error
- `mean absolute error`_
.. _mean absolute error:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_absolute_error.html
- `r2 score`_
.. _r2 score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html
**Tags**
- An estimator class name (e.g. "LinearRegression").
- A fully qualified estimator class name
(e.g. "sklearn.linear_model._base.LinearRegression").
**Artifacts**
- An MLflow Model with the :py:mod:`mlflow.sklearn` flavor containing a fitted estimator
(logged by :py:func:`mlflow.sklearn.log_model()`). The Model also contains the
:py:mod:`mlflow.pyfunc` flavor when the scikit-learn estimator defines `predict()`.
**How does autologging work for meta estimators?**
When a meta estimator (e.g. `Pipeline`_, `GridSearchCV`_) calls ``fit()``, it internally calls
``fit()`` on its child estimators. Autologging does NOT perform logging on these constituent
``fit()`` calls.
**Parameter search**
In addition to recording the information discussed above, autologging for parameter
search meta estimators (`GridSearchCV`_ and `RandomizedSearchCV`_) records child runs
with metrics for each set of explored parameters, as well as artifacts and parameters
for the best model (if available).
**Supported estimators**
- All estimators obtained by `sklearn.utils.all_estimators`_ (including meta estimators).
- `Pipeline`_
- Parameter search estimators (`GridSearchCV`_ and `RandomizedSearchCV`_)
.. _sklearn.utils.all_estimators:
https://scikit-learn.org/stable/modules/generated/sklearn.utils.all_estimators.html
.. _Pipeline:
https://scikit-learn.org/stable/modules/generated/sklearn.pipeline.Pipeline.html
.. _GridSearchCV:
https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.GridSearchCV.html
.. _RandomizedSearchCV:
https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.RandomizedSearchCV.html
**Example**
`See more examples <https://github.com/mlflow/mlflow/blob/master/examples/sklearn_autolog>`_
.. code-block:: python
from pprint import pprint
import numpy as np
from sklearn.linear_model import LinearRegression
import mlflow
def fetch_logged_data(run_id):
client = mlflow.tracking.MlflowClient()
data = client.get_run(run_id).data
tags = {k: v for k, v in data.tags.items() if not k.startswith("mlflow.")}
artifacts = [f.path for f in client.list_artifacts(run_id, "model")]
return data.params, data.metrics, tags, artifacts
# enable autologging
mlflow.sklearn.autolog()
# prepare training data
X = np.array([[1, 1], [1, 2], [2, 2], [2, 3]])
y = np.dot(X, np.array([1, 2])) + 3
# train a model
model = LinearRegression()
with mlflow.start_run() as run:
model.fit(X, y)
# fetch logged data
params, metrics, tags, artifacts = fetch_logged_data(run.info.run_id)
pprint(params)
# {'copy_X': 'True',
# 'fit_intercept': 'True',
# 'n_jobs': 'None',
# 'normalize': 'False'}
pprint(metrics)
# {'training_score': 1.0,
'training_mae': 2.220446049250313e-16,
'training_mse': 1.9721522630525295e-31,
'training_r2_score': 1.0,
'training_rmse': 4.440892098500626e-16}
pprint(tags)
# {'estimator_class': 'sklearn.linear_model._base.LinearRegression',
# 'estimator_name': 'LinearRegression'}
pprint(artifacts)
# ['model/MLmodel', 'model/conda.yaml', 'model/model.pkl']
:param log_input_examples: If ``True``, input examples from training datasets are collected and
logged along with scikit-learn model artifacts during training. If
``False``, input examples are not logged.
Note: Input examples are MLflow model attributes
and are only collected if ``log_models`` is also ``True``.
:param log_model_signatures: If ``True``,
:py:class:`ModelSignatures <mlflow.models.ModelSignature>`
describing model inputs and outputs are collected and logged along
with scikit-learn model artifacts during training. If ``False``,
signatures are not logged.
Note: Model signatures are MLflow model attributes
and are only collected if ``log_models`` is also ``True``.
:param log_models: If ``True``, trained models are logged as MLflow model artifacts.
If ``False``, trained models are not logged.
Input examples and model signatures, which are attributes of MLflow models,
are also omitted when ``log_models`` is ``False``.
:param disable: If ``True``, disables the scikit-learn autologging integration. If ``False``,
enables the scikit-learn autologging integration.
:param exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
If ``False``, autologged content is logged to the active fluent run,
which may be user-created.
:param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
scikit-learn that have not been tested against this version of the MLflow
client or are incompatible.
"""
import pandas as pd
import sklearn
from mlflow.models import infer_signature
from mlflow.sklearn.utils import (
_MIN_SKLEARN_VERSION,
_is_supported_version,
_chunk_dict,
_get_args_for_score,
_log_specialized_estimator_content,
_get_Xy,
_all_estimators,
_truncate_dict,
_get_arg_names,
_get_estimator_info_tags,
_get_meta_estimators_for_autologging,
_is_parameter_search_estimator,
_log_parameter_search_results_as_artifact,
_create_child_runs_for_parameter_search,
)
from mlflow.tracking.context import registry as context_registry
from mlflow.utils.validation import (
MAX_PARAMS_TAGS_PER_BATCH,
MAX_PARAM_VAL_LENGTH,
MAX_ENTITY_KEY_LENGTH,
)
if not _is_supported_version():
warnings.warn(
"Autologging utilities may not work properly on scikit-learn < {} "
.format(_MIN_SKLEARN_VERSION) +
"(current version: {})".format(sklearn.__version__),
stacklevel=2,
)
def fit_mlflow(original, self, *args, **kwargs):
"""
Autologging function that performs model training by executing the training method
referred to be `func_name` on the instance of `clazz` referred to by `self` & records
MLflow parameters, metrics, tags, and artifacts to a corresponding MLflow Run.
"""
_log_pretraining_metadata(self, *args, **kwargs)
fit_output = original(self, *args, **kwargs)
_log_posttraining_metadata(self, *args, **kwargs)
return fit_output
def _log_pretraining_metadata(estimator, *args, **kwargs): # pylint: disable=unused-argument
"""
Records metadata (e.g., params and tags) for a scikit-learn estimator prior to training.
This is intended to be invoked within a patched scikit-learn training routine
(e.g., `fit()`, `fit_transform()`, ...) and assumes the existence of an active
MLflow run that can be referenced via the fluent Tracking API.
:param estimator: The scikit-learn estimator for which to log metadata.
:param args: The arguments passed to the scikit-learn training routine (e.g.,
`fit()`, `fit_transform()`, ...).
:param kwargs: The keyword arguments passed to the scikit-learn training routine.
"""
# Deep parameter logging includes parameters from children of a given
# estimator. For some meta estimators (e.g., pipelines), recording
# these parameters is desirable. For parameter search estimators,
# however, child estimators act as seeds for the parameter search
# process; accordingly, we avoid logging initial, untuned parameters
# for these seed estimators.
should_log_params_deeply = not _is_parameter_search_estimator(
estimator)
# Chunk model parameters to avoid hitting the log_batch API limit
for chunk in _chunk_dict(
estimator.get_params(deep=should_log_params_deeply),
chunk_size=MAX_PARAMS_TAGS_PER_BATCH,
):
truncated = _truncate_dict(chunk, MAX_ENTITY_KEY_LENGTH,
MAX_PARAM_VAL_LENGTH)
try_mlflow_log(mlflow.log_params, truncated)
try_mlflow_log(mlflow.set_tags, _get_estimator_info_tags(estimator))
def _log_posttraining_metadata(estimator, *args, **kwargs):
"""
Records metadata for a scikit-learn estimator after training has completed.
This is intended to be invoked within a patched scikit-learn training routine
(e.g., `fit()`, `fit_transform()`, ...) and assumes the existence of an active
MLflow run that can be referenced via the fluent Tracking API.
:param estimator: The scikit-learn estimator for which to log metadata.
:param args: The arguments passed to the scikit-learn training routine (e.g.,
`fit()`, `fit_transform()`, ...).
:param kwargs: The keyword arguments passed to the scikit-learn training routine.
"""
if hasattr(estimator, "score"):
try:
score_args = _get_args_for_score(estimator.score,
estimator.fit, args, kwargs)
training_score = estimator.score(*score_args)
except Exception as e:
msg = (
estimator.score.__qualname__ +
" failed. The 'training_score' metric will not be recorded. Scoring error: "
+ str(e))
_logger.warning(msg)
else:
try_mlflow_log(mlflow.log_metric, "training_score",
training_score)
# log common metrics and artifacts for estimators (classifier, regressor)
_log_specialized_estimator_content(estimator,
mlflow.active_run().info.run_id,
args, kwargs)
def get_input_example():
# Fetch an input example using the first several rows of the array-like
# training data supplied to the training routine (e.g., `fit()`)
fit_arg_names = _get_arg_names(estimator.fit)
X_var_name, y_var_name = fit_arg_names[:2]
input_example = _get_Xy(args, kwargs, X_var_name,
y_var_name)[0][:INPUT_EXAMPLE_SAMPLE_ROWS]
return input_example
def infer_model_signature(input_example):
if not hasattr(estimator, "predict"):
raise Exception(
"the trained model does not specify a `predict` function, "
+ "which is required in order to infer the signature")
return infer_signature(input_example,
estimator.predict(input_example))
if log_models:
# Will only resolve `input_example` and `signature` if `log_models` is `True`.
input_example, signature = resolve_input_example_and_signature(
get_input_example,
infer_model_signature,
log_input_examples,
log_model_signatures,
_logger,
)
try_mlflow_log(
log_model,
estimator,
artifact_path="model",
signature=signature,
input_example=input_example,
)
if _is_parameter_search_estimator(estimator):
if hasattr(estimator, "best_estimator_") and log_models:
try_mlflow_log(
log_model,
estimator.best_estimator_,
artifact_path="best_estimator",
signature=signature,
input_example=input_example,
)
if hasattr(estimator, "best_score_"):
try_mlflow_log(mlflow.log_metric, "best_cv_score",
estimator.best_score_)
if hasattr(estimator, "best_params_"):
best_params = {
"best_{param_name}".format(param_name=param_name):
param_value
for param_name, param_value in
estimator.best_params_.items()
}
try_mlflow_log(mlflow.log_params, best_params)
if hasattr(estimator, "cv_results_"):
try:
# Fetch environment-specific tags (e.g., user and source) to ensure that lineage
# information is consistent with the parent run
child_tags = context_registry.resolve_tags()
child_tags.update({MLFLOW_AUTOLOGGING: FLAVOR_NAME})
_create_child_runs_for_parameter_search(
cv_estimator=estimator,
parent_run=mlflow.active_run(),
child_tags=child_tags,
)
except Exception as e:
msg = (
"Encountered exception during creation of child runs for parameter search."
" Child runs may be missing. Exception: {}".format(
str(e)))
_logger.warning(msg)
try:
cv_results_df = pd.DataFrame.from_dict(
estimator.cv_results_)
_log_parameter_search_results_as_artifact(
cv_results_df,
mlflow.active_run().info.run_id)
except Exception as e:
msg = (
"Failed to log parameter search results as an artifact."
" Exception: {}".format(str(e)))
_logger.warning(msg)
def patched_fit(original, self, *args, **kwargs):
"""
Autologging patch function to be applied to a sklearn model class that defines a `fit`
method and inherits from `BaseEstimator` (thereby defining the `get_params()` method)
:param clazz: The scikit-learn model class to which this patch function is being applied for
autologging (e.g., `sklearn.linear_model.LogisticRegression`)
:param func_name: The function name on the specified `clazz` that this patch is overriding
for autologging (e.g., specify "fit" in order to indicate that
`sklearn.linear_model.LogisticRegression.fit()` is being patched)
"""
with _SklearnTrainingSession(clazz=self.__class__,
allow_children=False) as t:
if t.should_log():
return fit_mlflow(original, self, *args, **kwargs)
else:
return original(self, *args, **kwargs)
_, estimators_to_patch = zip(*_all_estimators())
# Ensure that relevant meta estimators (e.g. GridSearchCV, Pipeline) are selected
# for patching if they are not already included in the output of `all_estimators()`
estimators_to_patch = set(estimators_to_patch).union(
set(_get_meta_estimators_for_autologging()))
# Exclude certain preprocessing & feature manipulation estimators from patching. These
# estimators represent data manipulation routines (e.g., normalization, label encoding)
# rather than ML algorithms. Accordingly, we should not create MLflow runs and log
# parameters / metrics for these routines, unless they are captured as part of an ML pipeline
# (via `sklearn.pipeline.Pipeline`)
excluded_module_names = [
"sklearn.preprocessing",
"sklearn.impute",
"sklearn.feature_extraction",
"sklearn.feature_selection",
]
estimators_to_patch = [
estimator for estimator in estimators_to_patch if not any([
estimator.__module__.startswith(excluded_module_name)
for excluded_module_name in excluded_module_names
])
]
for class_def in estimators_to_patch:
for func_name in ["fit", "fit_transform", "fit_predict"]:
if hasattr(class_def, func_name):
original = getattr(class_def, func_name)
# A couple of estimators use property methods to return fitting functions,
# rather than defining the fitting functions on the estimator class directly.
#
# Example: https://github.com/scikit-learn/scikit-learn/blob/0.23.2/sklearn/neighbors/_lof.py#L183 # noqa
#
# We currently exclude these property fitting methods from patching because
# it's challenging to patch them correctly.
#
# Excluded fitting methods:
# - sklearn.cluster._agglomerative.FeatureAgglomeration.fit_predict
# - sklearn.neighbors._lof.LocalOutlierFactor.fit_predict
#
# You can list property fitting methods by inserting "print(class_def, func_name)"
# in the if clause below.
if isinstance(original, property):
continue
safe_patch(
FLAVOR_NAME,
class_def,
func_name,
patched_fit,
manage_run=True,
)
def autolog(
log_every_n_epoch=1,
log_models=True,
disable=False,
exclusive=False,
disable_for_unsupported_versions=False,
): # pylint: disable=unused-argument
"""
Enables (or disables) and configures autologging from `PyTorch Lightning
<https://pytorch-lightning.readthedocs.io/en/latest>`_ to MLflow.
Autologging is performed when you call the `fit` method of
`pytorch_lightning.Trainer() \
<https://pytorch-lightning.readthedocs.io/en/latest/trainer.html#>`_.
Explore the complete `PyTorch MNIST \
<https://github.com/mlflow/mlflow/tree/master/examples/pytorch/MNIST/example1>`_ for
an expansive example with implementation of additional lightening steps.
**Note**: Autologging is only supported for PyTorch Lightning models,
i.e., models that subclass
`pytorch_lightning.LightningModule \
<https://pytorch-lightning.readthedocs.io/en/latest/lightning_module.html>`_.
In particular, autologging support for vanilla PyTorch models that only subclass
`torch.nn.Module <https://pytorch.org/docs/stable/generated/torch.nn.Module.html>`_
is not yet available.
:param log_every_n_epoch: If specified, logs metrics once every `n` epochs. By default, metrics
are logged after every epoch.
:param log_models: If ``True``, trained models are logged as MLflow model artifacts.
If ``False``, trained models are not logged.
:param disable: If ``True``, disables the PyTorch Lightning autologging integration.
If ``False``, enables the PyTorch Lightning autologging integration.
:param exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
If ``False``, autologged content is logged to the active fluent run,
which may be user-created.
:param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
pytorch and pytorch-lightning that have not been tested against this version
of the MLflow client or are incompatible.
.. code-block:: python
:caption: Example
import os
import pytorch_lightning as pl
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
from pytorch_lightning.metrics.functional import accuracy
import mlflow.pytorch
from mlflow.tracking import MlflowClient
# For brevity, here is the simplest most minimal example with just a training
# loop step, (no validation, no testing). It illustrates how you can use MLflow
# to auto log parameters, metrics, and models.
class MNISTModel(pl.LightningModule):
def __init__(self):
super(MNISTModel, self).__init__()
self.l1 = torch.nn.Linear(28 * 28, 10)
def forward(self, x):
return torch.relu(self.l1(x.view(x.size(0), -1)))
def training_step(self, batch, batch_nb):
x, y = batch
loss = F.cross_entropy(self(x), y)
acc = accuracy(loss, y)
# Use the current of PyTorch logger
self.log("train_loss", loss, on_epoch=True)
self.log("acc", acc, on_epoch=True)
return loss
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=0.02)
def print_auto_logged_info(r):
tags = {k: v for k, v in r.data.tags.items() if not k.startswith("mlflow.")}
artifacts = [f.path for f in MlflowClient().list_artifacts(r.info.run_id, "model")]
print("run_id: {}".format(r.info.run_id))
print("artifacts: {}".format(artifacts))
print("params: {}".format(r.data.params))
print("metrics: {}".format(r.data.metrics))
print("tags: {}".format(tags))
# Initialize our model
mnist_model = MNISTModel()
# Initialize DataLoader from MNIST Dataset
train_ds = MNIST(os.getcwd(), train=True,
download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_ds, batch_size=32)
# Initialize a trainer
trainer = pl.Trainer(max_epochs=20, progress_bar_refresh_rate=20)
# Auto log all MLflow entities
mlflow.pytorch.autolog()
# Train the model
with mlflow.start_run() as run:
trainer.fit(mnist_model, train_loader)
# fetch the auto logged parameters and metrics
print_auto_logged_info(mlflow.get_run(run_id=run.info.run_id))
.. code-block:: text
:caption: Output
run_id: 42caa17b60cb489c8083900fb52506a7
artifacts: ['model/MLmodel', 'model/conda.yaml', 'model/data']
params: {'betas': '(0.9, 0.999)',
'weight_decay': '0',
'epochs': '20',
'eps': '1e-08',
'lr': '0.02',
'optimizer_name': 'Adam', '
amsgrad': 'False'}
metrics: {'acc_step': 0.0,
'train_loss_epoch': 1.0917967557907104,
'train_loss_step': 1.0794280767440796,
'train_loss': 1.0794280767440796,
'acc_epoch': 0.0033333334140479565,
'acc': 0.0}
tags: {'Mode': 'training'}
.. figure:: ../_static/images/pytorch_lightening_autolog.png
PyTorch autologged MLflow entities
"""
import pytorch_lightning as pl
from mlflow.pytorch._pytorch_autolog import _create_patch_fit
fit = _create_patch_fit(log_every_n_epoch=log_every_n_epoch,
log_models=log_models)
safe_patch(FLAVOR_NAME, pl.Trainer, "fit", fit, manage_run=True)
def autolog(
log_models=True, disable=False, exclusive=False, disable_for_unsupported_versions=False
): # pylint: disable=unused-argument
# pylint: disable=E0611
"""
Enables (or disables) and configures autologging from Keras to MLflow. Autologging captures
the following information:
**Metrics** and **Parameters**
- Training loss; validation loss; user-specified metrics
- Metrics associated with the ``EarlyStopping`` callbacks: ``stopped_epoch``,
``restored_epoch``, ``restore_best_weight``, ``last_epoch``, etc
- ``fit()`` or ``fit_generator()`` parameters; optimizer name; learning rate; epsilon
- ``fit()`` or ``fit_generator()`` parameters associated with ``EarlyStopping``: ``min_delta``,
``patience``, ``baseline``, ``restore_best_weights``, etc
**Artifacts**
- Model summary on training start
- `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (Keras model) on training end
.. code-block:: python
:caption: Example
import mlflow
import mlflow.keras
# Build, compile, enable autologging, and train your model
keras_model = ...
keras_model.compile(optimizer="rmsprop", loss="mse", metrics=["accuracy"])
# autolog your metrics, parameters, and model
mlflow.keras.autolog()
results = keras_model.fit(
x_train, y_train, epochs=20, batch_size=128, validation_data=(x_val, y_val))
``EarlyStopping Integration with Keras AutoLogging``
MLflow will detect if an ``EarlyStopping`` callback is used in a ``fit()`` or
``fit_generator()`` call, and if the ``restore_best_weights`` parameter is set to be ``True``,
then MLflow will log the metrics associated with the restored model as a final, extra step.
The epoch of the restored model will also be logged as the metric ``restored_epoch``.
This allows for easy comparison between the actual metrics of the restored model and
the metrics of other models.
If ``restore_best_weights`` is set to be ``False``, then MLflow will not log an additional step.
Regardless of ``restore_best_weights``, MLflow will also log ``stopped_epoch``,
which indicates the epoch at which training stopped due to early stopping.
If training does not end due to early stopping, then ``stopped_epoch`` will be logged as ``0``.
MLflow will also log the parameters of the ``EarlyStopping`` callback,
excluding ``mode`` and ``verbose``.
:param log_models: If ``True``, trained models are logged as MLflow model artifacts.
If ``False``, trained models are not logged.
:param disable: If ``True``, disables the Keras autologging integration. If ``False``,
enables the Keras autologging integration.
:param exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
If ``False``, autologged content is logged to the active fluent run,
which may be user-created.
:param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
keras that have not been tested against this version of the MLflow client
or are incompatible.
"""
import keras
def getKerasCallback(metrics_logger):
class __MLflowKerasCallback(keras.callbacks.Callback, metaclass=ExceptionSafeClass):
"""
Callback for auto-logging metrics and parameters.
Records available logs after each epoch.
Records model structural information as params when training begins
"""
def on_train_begin(self, logs=None): # pylint: disable=unused-argument
try_mlflow_log(mlflow.log_param, "num_layers", len(self.model.layers))
try_mlflow_log(
mlflow.log_param, "optimizer_name", type(self.model.optimizer).__name__
)
if hasattr(self.model.optimizer, "lr"):
lr = (
self.model.optimizer.lr
if type(self.model.optimizer.lr) is float
else keras.backend.eval(self.model.optimizer.lr)
)
try_mlflow_log(mlflow.log_param, "learning_rate", lr)
if hasattr(self.model.optimizer, "epsilon"):
epsilon = (
self.model.optimizer.epsilon
if type(self.model.optimizer.epsilon) is float
else keras.backend.eval(self.model.optimizer.epsilon)
)
try_mlflow_log(mlflow.log_param, "epsilon", epsilon)
sum_list = []
self.model.summary(print_fn=sum_list.append)
summary = "\n".join(sum_list)
tempdir = tempfile.mkdtemp()
try:
summary_file = os.path.join(tempdir, "model_summary.txt")
with open(summary_file, "w") as f:
f.write(summary)
try_mlflow_log(mlflow.log_artifact, local_path=summary_file)
finally:
shutil.rmtree(tempdir)
def on_epoch_end(self, epoch, logs=None):
if not logs:
return
metrics_logger.record_metrics(logs, epoch)
def on_train_end(self, logs=None):
if log_models:
try_mlflow_log(log_model, self.model, artifact_path="model")
# As of Keras 2.4.0, Keras Callback implementations must define the following
# methods indicating whether or not the callback overrides functions for
# batch training/testing/inference
def _implements_train_batch_hooks(self):
return False
def _implements_test_batch_hooks(self):
return False
def _implements_predict_batch_hooks(self):
return False
return __MLflowKerasCallback()
def _early_stop_check(callbacks):
if LooseVersion(keras.__version__) < LooseVersion("2.3.0") or LooseVersion(
keras.__version__
) >= LooseVersion("2.4.0"):
es_callback = keras.callbacks.EarlyStopping
else:
es_callback = keras.callbacks.callbacks.EarlyStopping
for callback in callbacks:
if isinstance(callback, es_callback):
return callback
return None
def _log_early_stop_callback_params(callback):
if callback:
earlystopping_params = {
"monitor": callback.monitor,
"min_delta": callback.min_delta,
"patience": callback.patience,
"baseline": callback.baseline,
"restore_best_weights": callback.restore_best_weights,
}
try_mlflow_log(mlflow.log_params, earlystopping_params)
def _get_early_stop_callback_attrs(callback):
try:
return callback.stopped_epoch, callback.restore_best_weights, callback.patience
except Exception: # pylint: disable=W0703
return None
def _log_early_stop_callback_metrics(callback, history, metrics_logger):
if callback:
callback_attrs = _get_early_stop_callback_attrs(callback)
if callback_attrs is None:
return
stopped_epoch, restore_best_weights, patience = callback_attrs
metrics_logger.record_metrics({"stopped_epoch": stopped_epoch})
# Weights are restored only if early stopping occurs
if stopped_epoch != 0 and restore_best_weights:
restored_epoch = stopped_epoch - max(1, patience)
metrics_logger.record_metrics({"restored_epoch": restored_epoch})
restored_index = history.epoch.index(restored_epoch)
restored_metrics = {
key: history.history[key][restored_index] for key in history.history.keys()
}
# Checking that a metric history exists
metric_key = next(iter(history.history), None)
if metric_key is not None:
last_epoch = len(history.history[metric_key])
metrics_logger.record_metrics(restored_metrics, last_epoch)
def _run_and_log_function(self, original, args, kwargs, unlogged_params, callback_arg_index):
log_fn_args_as_params(original, args, kwargs, unlogged_params)
early_stop_callback = None
# Checking if the 'callback' argument of the function is set
run_id = mlflow.active_run().info.run_id
with batch_metrics_logger(run_id) as metrics_logger:
mlflowKerasCallback = getKerasCallback(metrics_logger)
if len(args) > callback_arg_index:
tmp_list = list(args)
early_stop_callback = _early_stop_check(tmp_list[callback_arg_index])
tmp_list[callback_arg_index] += [mlflowKerasCallback]
args = tuple(tmp_list)
elif kwargs.get("callbacks"):
early_stop_callback = _early_stop_check(kwargs["callbacks"])
kwargs["callbacks"] += [mlflowKerasCallback]
else:
kwargs["callbacks"] = [mlflowKerasCallback]
try_mlflow_log(_log_early_stop_callback_params, early_stop_callback)
history = original(self, *args, **kwargs)
try_mlflow_log(
_log_early_stop_callback_metrics, early_stop_callback, history, metrics_logger
)
return history
def fit(original, self, *args, **kwargs):
unlogged_params = ["self", "x", "y", "callbacks", "validation_data", "verbose"]
return _run_and_log_function(self, original, args, kwargs, unlogged_params, 5)
def fit_generator(original, self, *args, **kwargs):
"""
NOTE: `fit_generator()` is deprecated in Keras >= 2.4.0 and simply wraps `fit()`.
To avoid unintentional creation of nested MLflow runs caused by a patched
`fit_generator()` method calling a patched `fit()` method, we only patch
`fit_generator()` in Keras < 2.4.0.
"""
unlogged_params = ["self", "generator", "callbacks", "validation_data", "verbose"]
return _run_and_log_function(self, original, args, kwargs, unlogged_params, 4)
safe_patch(FLAVOR_NAME, keras.Model, "fit", fit, manage_run=True)
# `fit_generator()` is deprecated in Keras >= 2.4.0 and simply wraps `fit()`.
# To avoid unintentional creation of nested MLflow runs caused by a patched
# `fit_generator()` method calling a patched `fit()` method, we only patch
# `fit_generator()` in Keras < 2.4.0.
if LooseVersion(keras.__version__) < LooseVersion("2.4.0"):
safe_patch(FLAVOR_NAME, keras.Model, "fit_generator", fit_generator, manage_run=True)
def autolog(
log_input_examples=False,
log_model_signatures=True,
log_models=True,
disable=False,
exclusive=False,
disable_for_unsupported_versions=False,
): # pylint: disable=unused-argument
"""
Enables (or disables) and configures autologging from LightGBM to MLflow. Logs the following:
- parameters specified in `lightgbm.train`_.
- metrics on each iteration (if ``valid_sets`` specified).
- metrics at the best iteration (if ``early_stopping_rounds`` specified).
- feature importance (both "split" and "gain") as JSON files and plots.
- trained model, including:
- an example of valid input.
- inferred signature of the inputs and outputs of the model.
Note that the `scikit-learn API`_ is not supported.
:param log_input_examples: If ``True``, input examples from training datasets are collected and
logged along with LightGBM model artifacts during training. If
``False``, input examples are not logged.
Note: Input examples are MLflow model attributes
and are only collected if ``log_models`` is also ``True``.
:param log_model_signatures: If ``True``,
:py:class:`ModelSignatures <mlflow.models.ModelSignature>`
describing model inputs and outputs are collected and logged along
with LightGBM model artifacts during training. If ``False``,
signatures are not logged.
Note: Model signatures are MLflow model attributes
and are only collected if ``log_models`` is also ``True``.
:param log_models: If ``True``, trained models are logged as MLflow model artifacts.
If ``False``, trained models are not logged.
Input examples and model signatures, which are attributes of MLflow models,
are also omitted when ``log_models`` is ``False``.
:param disable: If ``True``, disables the LightGBM autologging integration. If ``False``,
enables the LightGBM autologging integration.
:param exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
If ``False``, autologged content is logged to the active fluent run,
which may be user-created.
:param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
lightgbm that have not been tested against this version of the MLflow client
or are incompatible.
"""
import lightgbm
import numpy as np
# Patching this function so we can get a copy of the data given to Dataset.__init__
# to use as an input example and for inferring the model signature.
# (there is no way to get the data back from a Dataset object once it is consumed by train)
# We store it on the Dataset object so the train function is able to read it.
def __init__(original, self, *args, **kwargs):
data = args[0] if len(args) > 0 else kwargs.get("data")
if data is not None:
try:
if isinstance(data, str):
raise Exception(
"cannot gather example input when dataset is loaded from a file."
)
input_example_info = _InputExampleInfo(
input_example=deepcopy(data[:INPUT_EXAMPLE_SAMPLE_ROWS]))
except Exception as e:
input_example_info = _InputExampleInfo(error_msg=str(e))
setattr(self, "input_example_info", input_example_info)
original(self, *args, **kwargs)
def train(original, *args, **kwargs):
def record_eval_results(eval_results, metrics_logger):
"""
Create a callback function that records evaluation results.
"""
@exception_safe_function
def callback(env):
res = {}
for data_name, eval_name, value, _ in env.evaluation_result_list:
key = data_name + "-" + eval_name
res[key] = value
metrics_logger.record_metrics(res, env.iteration)
eval_results.append(res)
return callback
def log_feature_importance_plot(features, importance, importance_type):
"""
Log feature importance plot.
"""
import matplotlib.pyplot as plt
indices = np.argsort(importance)
features = np.array(features)[indices]
importance = importance[indices]
num_features = len(features)
# If num_features > 10, increase the figure height to prevent the plot
# from being too dense.
w, h = [6.4, 4.8] # matplotlib's default figure size
h = h + 0.1 * num_features if num_features > 10 else h
fig, ax = plt.subplots(figsize=(w, h))
yloc = np.arange(num_features)
ax.barh(yloc, importance, align="center", height=0.5)
ax.set_yticks(yloc)
ax.set_yticklabels(features)
ax.set_xlabel("Importance")
ax.set_title("Feature Importance ({})".format(importance_type))
fig.tight_layout()
tmpdir = tempfile.mkdtemp()
try:
# pylint: disable=undefined-loop-variable
filepath = os.path.join(
tmpdir, "feature_importance_{}.png".format(imp_type))
fig.savefig(filepath)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
plt.close(fig)
shutil.rmtree(tmpdir)
# logging booster params separately via mlflow.log_params to extract key/value pairs
# and make it easier to compare them across runs.
params = args[0] if len(args) > 0 else kwargs["params"]
try_mlflow_log(mlflow.log_params, params)
unlogged_params = [
"params",
"train_set",
"valid_sets",
"valid_names",
"fobj",
"feval",
"init_model",
"evals_result",
"learning_rates",
"callbacks",
]
log_fn_args_as_params(original, args, kwargs, unlogged_params)
all_arg_names = inspect.getargspec(original)[0] # pylint: disable=W1505
num_pos_args = len(args)
# adding a callback that records evaluation results.
eval_results = []
callbacks_index = all_arg_names.index("callbacks")
run_id = mlflow.active_run().info.run_id
with batch_metrics_logger(run_id) as metrics_logger:
callback = record_eval_results(eval_results, metrics_logger)
if num_pos_args >= callbacks_index + 1:
tmp_list = list(args)
tmp_list[callbacks_index] += [callback]
args = tuple(tmp_list)
elif "callbacks" in kwargs and kwargs["callbacks"] is not None:
kwargs["callbacks"] += [callback]
else:
kwargs["callbacks"] = [callback]
# training model
model = original(*args, **kwargs)
# If early_stopping_rounds is present, logging metrics at the best iteration
# as extra metrics with the max step + 1.
early_stopping_index = all_arg_names.index("early_stopping_rounds")
early_stopping = (num_pos_args >= early_stopping_index + 1
or "early_stopping_rounds" in kwargs)
if early_stopping:
extra_step = len(eval_results)
metrics_logger.record_metrics(
{"stopped_iteration": extra_step})
# best_iteration is set even if training does not stop early.
metrics_logger.record_metrics(
{"best_iteration": model.best_iteration})
# iteration starts from 1 in LightGBM.
results = eval_results[model.best_iteration - 1]
metrics_logger.record_metrics(results, step=extra_step)
# logging feature importance as artifacts.
for imp_type in ["split", "gain"]:
features = model.feature_name()
importance = model.feature_importance(importance_type=imp_type)
try:
log_feature_importance_plot(features, importance, imp_type)
except Exception:
_logger.exception(
"Failed to log feature importance plot. LightGBM autologging "
"will ignore the failure and continue. Exception: ")
imp = {ft: imp for ft, imp in zip(features, importance.tolist())}
tmpdir = tempfile.mkdtemp()
try:
filepath = os.path.join(
tmpdir, "feature_importance_{}.json".format(imp_type))
with open(filepath, "w") as f:
json.dump(imp, f, indent=2)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
shutil.rmtree(tmpdir)
# train_set must exist as the original train function already ran successfully
train_set = args[1] if len(args) > 1 else kwargs.get("train_set")
# it is possible that the dataset was constructed before the patched
# constructor was applied, so we cannot assume the input_example_info exists
input_example_info = getattr(train_set, "input_example_info", None)
def get_input_example():
if input_example_info is None:
raise Exception(ENSURE_AUTOLOGGING_ENABLED_TEXT)
if input_example_info.error_msg is not None:
raise Exception(input_example_info.error_msg)
return input_example_info.input_example
def infer_model_signature(input_example):
model_output = model.predict(input_example)
model_signature = infer_signature(input_example, model_output)
return model_signature
# Whether to automatically log the trained model based on boolean flag.
if log_models:
# Will only resolve `input_example` and `signature` if `log_models` is `True`.
input_example, signature = resolve_input_example_and_signature(
get_input_example,
infer_model_signature,
log_input_examples,
log_model_signatures,
_logger,
)
try_mlflow_log(
log_model,
model,
artifact_path="model",
signature=signature,
input_example=input_example,
)
return model
safe_patch(FLAVOR_NAME, lightgbm, "train", train, manage_run=True)
safe_patch(FLAVOR_NAME, lightgbm.Dataset, "__init__", __init__)
def autolog(
log_models=True, disable=False, exclusive=False, disable_for_unsupported_versions=False
): # pylint: disable=unused-argument
"""
Enables (or disables) and configures autologging from Gluon to MLflow.
Logs loss and any other metrics specified in the fit
function, and optimizer data as parameters. Model checkpoints
are logged as artifacts to a 'models' directory.
:param log_models: If ``True``, trained models are logged as MLflow model artifacts.
If ``False``, trained models are not logged.
:param disable: If ``True``, disables the MXNet Gluon autologging integration. If ``False``,
enables the MXNet Gluon autologging integration.
:param exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
If ``False``, autologged content is logged to the active fluent run,
which may be user-created.
:param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
gluon that have not been tested against this version of the MLflow client
or are incompatible.
"""
from mxnet.gluon.contrib.estimator import Estimator, EpochEnd, TrainBegin, TrainEnd
from mxnet.gluon.nn import HybridSequential
def getGluonCallback(metrics_logger):
class __MLflowGluonCallback(EpochEnd, TrainEnd, TrainBegin, metaclass=ExceptionSafeClass):
def __init__(self):
self.current_epoch = 0
def epoch_end(self, estimator, *args, **kwargs):
logs = {}
for metric in estimator.train_metrics:
metric_name, metric_val = metric.get()
logs[metric_name] = metric_val
for metric in estimator.val_metrics:
metric_name, metric_val = metric.get()
logs[metric_name] = metric_val
metrics_logger.record_metrics(logs, self.current_epoch)
self.current_epoch += 1
def train_begin(self, estimator, *args, **kwargs):
try_mlflow_log(mlflow.log_param, "num_layers", len(estimator.net))
if estimator.max_epoch is not None:
try_mlflow_log(mlflow.log_param, "epochs", estimator.max_epoch)
if estimator.max_batch is not None:
try_mlflow_log(mlflow.log_param, "batches", estimator.max_batch)
try_mlflow_log(
mlflow.log_param, "optimizer_name", type(estimator.trainer.optimizer).__name__
)
if hasattr(estimator.trainer.optimizer, "lr"):
try_mlflow_log(
mlflow.log_param, "learning_rate", estimator.trainer.optimizer.lr
)
if hasattr(estimator.trainer.optimizer, "epsilon"):
try_mlflow_log(mlflow.log_param, "epsilon", estimator.trainer.optimizer.epsilon)
def train_end(self, estimator, *args, **kwargs):
if isinstance(estimator.net, HybridSequential) and log_models:
try_mlflow_log(log_model, estimator.net, artifact_path="model")
return __MLflowGluonCallback()
def fit(original, self, *args, **kwargs):
# Wrap `fit` execution within a batch metrics logger context.
run_id = mlflow.active_run().info.run_id
with batch_metrics_logger(run_id) as metrics_logger:
mlflowGluonCallback = getGluonCallback(metrics_logger)
if len(args) >= 4:
l = list(args)
l[3] += [mlflowGluonCallback]
args = tuple(l)
elif "event_handlers" in kwargs:
kwargs["event_handlers"] += [mlflowGluonCallback]
else:
kwargs["event_handlers"] = [mlflowGluonCallback]
result = original(self, *args, **kwargs)
return result
safe_patch(FLAVOR_NAME, Estimator, "fit", fit, manage_run=True)
def autolog(
importance_types=None,
log_input_examples=False,
log_model_signatures=True,
log_models=True,
disable=False,
exclusive=False,
disable_for_unsupported_versions=False,
): # pylint: disable=W0102,unused-argument
"""
Enables (or disables) and configures autologging from XGBoost to MLflow. Logs the following:
- parameters specified in `xgboost.train`_.
- metrics on each iteration (if ``evals`` specified).
- metrics at the best iteration (if ``early_stopping_rounds`` specified).
- feature importance as JSON files and plots.
- trained model, including:
- an example of valid input.
- inferred signature of the inputs and outputs of the model.
Note that the `scikit-learn API`_ is not supported.
:param importance_types: Importance types to log. If unspecified, defaults to ``["weight"]``.
:param log_input_examples: If ``True``, input examples from training datasets are collected and
logged along with XGBoost model artifacts during training. If
``False``, input examples are not logged.
Note: Input examples are MLflow model attributes
and are only collected if ``log_models`` is also ``True``.
:param log_model_signatures: If ``True``,
:py:class:`ModelSignatures <mlflow.models.ModelSignature>`
describing model inputs and outputs are collected and logged along
with XGBoost model artifacts during training. If ``False``,
signatures are not logged.
Note: Model signatures are MLflow model attributes
and are only collected if ``log_models`` is also ``True``.
:param log_models: If ``True``, trained models are logged as MLflow model artifacts.
If ``False``, trained models are not logged.
Input examples and model signatures, which are attributes of MLflow models,
are also omitted when ``log_models`` is ``False``.
:param disable: If ``True``, disables the XGBoost autologging integration. If ``False``,
enables the XGBoost autologging integration.
:param exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
If ``False``, autologged content is logged to the active fluent run,
which may be user-created.
:param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
xgboost that have not been tested against this version of the MLflow client
or are incompatible.
"""
import xgboost
import numpy as np
if importance_types is None:
importance_types = ["weight"]
# Patching this function so we can get a copy of the data given to DMatrix.__init__
# to use as an input example and for inferring the model signature.
# (there is no way to get the data back from a DMatrix object)
# We store it on the DMatrix object so the train function is able to read it.
def __init__(original, self, *args, **kwargs):
data = args[0] if len(args) > 0 else kwargs.get("data")
if data is not None:
try:
if isinstance(data, str):
raise Exception(
"cannot gather example input when dataset is loaded from a file."
)
input_example_info = _InputExampleInfo(
input_example=deepcopy(data[:INPUT_EXAMPLE_SAMPLE_ROWS]))
except Exception as e:
input_example_info = _InputExampleInfo(error_msg=str(e))
setattr(self, "input_example_info", input_example_info)
original(self, *args, **kwargs)
def train(original, *args, **kwargs):
def record_eval_results(eval_results, metrics_logger):
"""
Create a callback function that records evaluation results.
"""
if LooseVersion(xgboost.__version__) >= LooseVersion("1.3.0"):
# In xgboost >= 1.3.0, user-defined callbacks should inherit
# `xgboost.callback.TrainingCallback`:
# https://xgboost.readthedocs.io/en/latest/python/callbacks.html#defining-your-own-callback # noqa
class Callback(
xgboost.callback.TrainingCallback,
metaclass=ExceptionSafeAbstractClass,
):
def after_iteration(self, model, epoch, evals_log):
"""
Run after each iteration. Return True when training should stop.
"""
# `evals_log` is a nested dict (type: Dict[str, Dict[str, List[float]]])
# that looks like this:
# {
# "train": {
# "auc": [0.5, 0.6, 0.7, ...],
# ...
# },
# ...
# }
evaluation_result_dict = {}
for data_name, metric_dict in evals_log.items():
for metric_name, metric_values_on_each_iter in metric_dict.items(
):
key = "{}-{}".format(data_name, metric_name)
# The last element in `metric_values_on_each_iter` corresponds to
# the meric on the current iteration
evaluation_result_dict[
key] = metric_values_on_each_iter[-1]
metrics_logger.record_metrics(evaluation_result_dict,
epoch)
eval_results.append(evaluation_result_dict)
# Return `False` to indicate training should not stop
return False
return Callback()
else:
@exception_safe_function
def callback(env):
metrics_logger.record_metrics(
dict(env.evaluation_result_list), env.iteration)
eval_results.append(dict(env.evaluation_result_list))
return callback
def log_feature_importance_plot(features, importance, importance_type):
"""
Log feature importance plot.
"""
import matplotlib.pyplot as plt
features = np.array(features)
importance = np.array(importance)
indices = np.argsort(importance)
features = features[indices]
importance = importance[indices]
num_features = len(features)
# If num_features > 10, increase the figure height to prevent the plot
# from being too dense.
w, h = [6.4, 4.8] # matplotlib's default figure size
h = h + 0.1 * num_features if num_features > 10 else h
fig, ax = plt.subplots(figsize=(w, h))
yloc = np.arange(num_features)
ax.barh(yloc, importance, align="center", height=0.5)
ax.set_yticks(yloc)
ax.set_yticklabels(features)
ax.set_xlabel("Importance")
ax.set_title("Feature Importance ({})".format(importance_type))
fig.tight_layout()
tmpdir = tempfile.mkdtemp()
try:
# pylint: disable=undefined-loop-variable
filepath = os.path.join(
tmpdir, "feature_importance_{}.png".format(imp_type))
fig.savefig(filepath)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
plt.close(fig)
shutil.rmtree(tmpdir)
# logging booster params separately via mlflow.log_params to extract key/value pairs
# and make it easier to compare them across runs.
params = args[0] if len(args) > 0 else kwargs["params"]
try_mlflow_log(mlflow.log_params, params)
unlogged_params = [
"params",
"dtrain",
"evals",
"obj",
"feval",
"evals_result",
"xgb_model",
"callbacks",
"learning_rates",
]
log_fn_args_as_params(original, args, kwargs, unlogged_params)
all_arg_names = inspect.getargspec(original)[0] # pylint: disable=W1505
num_pos_args = len(args)
# adding a callback that records evaluation results.
eval_results = []
callbacks_index = all_arg_names.index("callbacks")
run_id = mlflow.active_run().info.run_id
with batch_metrics_logger(run_id) as metrics_logger:
callback = record_eval_results(eval_results, metrics_logger)
if num_pos_args >= callbacks_index + 1:
tmp_list = list(args)
tmp_list[callbacks_index] += [callback]
args = tuple(tmp_list)
elif "callbacks" in kwargs and kwargs["callbacks"] is not None:
kwargs["callbacks"] += [callback]
else:
kwargs["callbacks"] = [callback]
# training model
model = original(*args, **kwargs)
# If early_stopping_rounds is present, logging metrics at the best iteration
# as extra metrics with the max step + 1.
early_stopping_index = all_arg_names.index("early_stopping_rounds")
early_stopping = (num_pos_args >= early_stopping_index + 1
or "early_stopping_rounds" in kwargs)
if early_stopping:
extra_step = len(eval_results)
metrics_logger.record_metrics(
{"stopped_iteration": extra_step - 1})
metrics_logger.record_metrics(
{"best_iteration": model.best_iteration})
metrics_logger.record_metrics(
eval_results[model.best_iteration], extra_step)
# logging feature importance as artifacts.
for imp_type in importance_types:
imp = None
try:
imp = model.get_score(importance_type=imp_type)
features, importance = zip(*imp.items())
log_feature_importance_plot(features, importance, imp_type)
except Exception:
_logger.exception(
"Failed to log feature importance plot. XGBoost autologging "
"will ignore the failure and continue. Exception: ")
if imp is not None:
tmpdir = tempfile.mkdtemp()
try:
filepath = os.path.join(
tmpdir, "feature_importance_{}.json".format(imp_type))
with open(filepath, "w") as f:
json.dump(imp, f)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
shutil.rmtree(tmpdir)
# dtrain must exist as the original train function already ran successfully
dtrain = args[1] if len(args) > 1 else kwargs.get("dtrain")
# it is possible that the dataset was constructed before the patched
# constructor was applied, so we cannot assume the input_example_info exists
input_example_info = getattr(dtrain, "input_example_info", None)
def get_input_example():
if input_example_info is None:
raise Exception(ENSURE_AUTOLOGGING_ENABLED_TEXT)
if input_example_info.error_msg is not None:
raise Exception(input_example_info.error_msg)
return input_example_info.input_example
def infer_model_signature(input_example):
model_output = model.predict(xgboost.DMatrix(input_example))
model_signature = infer_signature(input_example, model_output)
return model_signature
# Only log the model if the autolog() param log_models is set to True.
if log_models:
# Will only resolve `input_example` and `signature` if `log_models` is `True`.
input_example, signature = resolve_input_example_and_signature(
get_input_example,
infer_model_signature,
log_input_examples,
log_model_signatures,
_logger,
)
try_mlflow_log(
log_model,
model,
artifact_path="model",
signature=signature,
input_example=input_example,
)
return model
safe_patch(FLAVOR_NAME, xgboost, "train", train, manage_run=True)
safe_patch(FLAVOR_NAME, xgboost.DMatrix, "__init__", __init__)