def test_if_getting_input_example_fails(logger):
error_msg = "NoneType has no whatever"
def throws():
raise Exception(error_msg)
input_example, signature = resolve_input_example_and_signature(
throws, infer_model_signature, True, True, logger)
assert input_example is None
assert signature is None
calls = [
call("Failed to gather input example: " + error_msg),
call("Failed to infer model signature: " +
"could not sample data to infer model signature: " + error_msg),
]
assert logger.warning.has_calls(calls)
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 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)
mlflow.log_artifact(filepath)
finally:
plt.close(fig)
shutil.rmtree(tmpdir)
autologging_client = MlflowAutologgingQueueingClient()
# logging booster params separately via mlflow.log_params to extract key/value pairs
# and make it easier to compare them across runs.
booster_params = args[0] if len(args) > 0 else kwargs["params"]
autologging_client.log_params(run_id=mlflow.active_run().info.run_id, params=booster_params)
unlogged_params = [
"params",
"train_set",
"valid_sets",
"valid_names",
"fobj",
"feval",
"init_model",
"evals_result",
"learning_rates",
"callbacks",
]
params_to_log_for_fn = get_mlflow_run_params_for_fn_args(
original, args, kwargs, unlogged_params
)
autologging_client.log_params(
run_id=mlflow.active_run().info.run_id, params=params_to_log_for_fn
)
param_logging_operations = autologging_client.flush(synchronous=False)
all_arg_names = _get_arg_names(original)
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)
autologging_client.log_metrics(
run_id=mlflow.active_run().info.run_id,
metrics={
"stopped_iteration": extra_step,
# best_iteration is set even if training does not stop early.
"best_iteration": model.best_iteration,
},
)
# iteration starts from 1 in LightGBM.
last_iter_results = eval_results[model.best_iteration - 1]
autologging_client.log_metrics(
run_id=mlflow.active_run().info.run_id,
metrics=last_iter_results,
step=extra_step,
)
early_stopping_logging_operations = autologging_client.flush(synchronous=False)
# 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)
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,
)
log_model(
model, artifact_path="model", signature=signature, input_example=input_example,
)
param_logging_operations.await_completion()
if early_stopping:
early_stopping_logging_operations.await_completion()
return model
def train(*args, **kwargs):
def record_eval_results(eval_results):
"""
Create a callback function that records evaluation results.
"""
def callback(env):
res = {}
for data_name, eval_name, value, _ in env.evaluation_result_list:
key = data_name + "-" + eval_name
res[key] = value
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)
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
auto_end_run = True
else:
auto_end_run = False
original = gorilla.get_original_attribute(lightgbm, "train")
# 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")
callback = record_eval_results(eval_results)
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)
# logging metrics on each iteration.
for idx, metrics in enumerate(eval_results):
try_mlflow_log(mlflow.log_metrics, metrics, step=idx)
# 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)
try_mlflow_log(mlflow.log_metric, "stopped_iteration", len(eval_results))
# best_iteration is set even if training does not stop early.
try_mlflow_log(mlflow.log_metric, "best_iteration", model.best_iteration)
# iteration starts from 1 in LightGBM.
try_mlflow_log(
mlflow.log_metrics, eval_results[model.best_iteration - 1], 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: # pylint: disable=broad-except
_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
input_example, signature = resolve_input_example_and_signature(
get_input_example,
infer_model_signature,
log_input_example,
log_model_signature,
_logger,
)
try_mlflow_log(
log_model,
model,
artifact_path="model",
signature=signature,
input_example=input_example,
)
if auto_end_run:
try_mlflow_log(mlflow.end_run)
return model
def train(*args, **kwargs):
def record_eval_results(eval_results, metrics_logger):
"""
Create a callback function that records evaluation results.
"""
def callback(env):
metrics_logger.record_metrics(dict(env.evaluation_result_list),
env.iteration)
eval_results.append(dict(env.evaluation_result_list))
return callback
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
auto_end_run = True
else:
auto_end_run = False
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)
original = gorilla.get_original_attribute(xgboost, "train")
# 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)
try_mlflow_log(mlflow.log_metric, "stopped_iteration",
len(eval_results) - 1)
try_mlflow_log(mlflow.log_metric, "best_iteration",
model.best_iteration)
try_mlflow_log(mlflow.log_metrics,
eval_results[model.best_iteration],
step=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: # pylint: disable=broad-except
_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,
)
if auto_end_run:
try_mlflow_log(mlflow.end_run)
return model
def train(original, *args, **kwargs):
def record_eval_results(eval_results, metrics_logger):
"""
Create a callback function that records evaluation results.
"""
# TODO: Remove `replace("SNAPSHOT", "dev")` once the following issue is addressed:
# https://github.com/dmlc/xgboost/issues/6984
if Version(xgboost.__version__.replace("SNAPSHOT", "dev")) >= Version("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
from cycler import cycler
features = np.array(features)
# Structure the supplied `importance` values as a `num_features`-by-`num_classes` matrix
importances_per_class_by_feature = np.array(importance)
if importances_per_class_by_feature.ndim <= 1:
# In this case, the supplied `importance` values are not given per class. Rather,
# one importance value is given per feature. For consistency with the assumed
# `num_features`-by-`num_classes` matrix structure, we coerce the importance
# values to a `num_features`-by-1 matrix
indices = np.argsort(importance)
# Sort features and importance values by magnitude during transformation to a
# `num_features`-by-`num_classes` matrix
features = features[indices]
importances_per_class_by_feature = np.array(
[[importance] for importance in importances_per_class_by_feature[indices]]
)
# In this case, do not include class labels on the feature importance plot because
# only one importance value has been provided per feature, rather than an
# one importance value for each class per feature
label_classes_on_plot = False
else:
importance_value_magnitudes = np.abs(importances_per_class_by_feature).sum(axis=1)
indices = np.argsort(importance_value_magnitudes)
features = features[indices]
importances_per_class_by_feature = importances_per_class_by_feature[indices]
label_classes_on_plot = True
num_classes = importances_per_class_by_feature.shape[1]
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
h = h + 0.1 * num_classes if num_classes > 1 else h
fig, ax = plt.subplots(figsize=(w, h))
# When importance values are provided for each class per feature, we want to ensure
# that the same color is used for all bars in the bar chart that have the same class
colors_to_cycle = plt.rcParams["axes.prop_cycle"].by_key()["color"][:num_classes]
color_cycler = cycler(color=colors_to_cycle)
ax.set_prop_cycle(color_cycler)
# The following logic operates on one feature at a time, adding a bar to the bar chart
# for each class that reflects the importance of the feature to predictions of that
# class
feature_ylocs = np.arange(num_features)
# Define offsets on the y-axis that are used to evenly space the bars for each class
# around the y-axis position of each feature
offsets_per_yloc = np.linspace(-0.5, 0.5, num_classes) / 2 if num_classes > 1 else [0]
for feature_idx, (feature_yloc, importances_per_class) in enumerate(
zip(feature_ylocs, importances_per_class_by_feature)
):
for class_idx, (offset, class_importance) in enumerate(
zip(offsets_per_yloc, importances_per_class)
):
(bar,) = ax.barh(
feature_yloc + offset,
class_importance,
align="center",
# Set the bar height such that importance value bars for a particular
# feature are spaced properly relative to each other (no overlap or gaps)
# and relative to importance value bars for other features
height=(0.5 / max(num_classes - 1, 1)),
)
if label_classes_on_plot and feature_idx == 0:
# Only set a label the first time a bar for a particular class is plotted to
# avoid duplicate legend entries. If we were to set a label for every bar,
# the legend would contain `num_features` labels for each class.
bar.set_label("Class {}".format(class_idx))
ax.set_yticks(feature_ylocs)
ax.set_yticklabels(features)
ax.set_xlabel("Importance")
ax.set_title("Feature Importance ({})".format(importance_type))
if label_classes_on_plot:
ax.legend()
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)
mlflow.log_artifact(filepath)
finally:
plt.close(fig)
shutil.rmtree(tmpdir)
autologging_client = MlflowAutologgingQueueingClient()
# logging booster params separately to extract key/value pairs and make it easier to
# compare them across runs.
booster_params = args[0] if len(args) > 0 else kwargs["params"]
autologging_client.log_params(run_id=mlflow.active_run().info.run_id, params=booster_params)
unlogged_params = [
"params",
"dtrain",
"evals",
"obj",
"feval",
"evals_result",
"xgb_model",
"callbacks",
"learning_rates",
]
params_to_log_for_fn = get_mlflow_run_params_for_fn_args(
original, args, kwargs, unlogged_params
)
autologging_client.log_params(
run_id=mlflow.active_run().info.run_id, params=params_to_log_for_fn
)
param_logging_operations = autologging_client.flush(synchronous=False)
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)
autologging_client.log_metrics(
run_id=mlflow.active_run().info.run_id,
metrics={
"stopped_iteration": extra_step - 1,
"best_iteration": model.best_iteration,
},
)
autologging_client.log_metrics(
run_id=mlflow.active_run().info.run_id,
metrics=eval_results[model.best_iteration],
step=extra_step,
)
early_stopping_logging_operations = autologging_client.flush(synchronous=False)
# 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)
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,
)
log_model(
model, artifact_path="model", signature=signature, input_example=input_example,
)
param_logging_operations.await_completion()
if early_stopping:
early_stopping_logging_operations.await_completion()
return model
def train(original, *args, **kwargs):
def record_eval_results(eval_results, metrics_logger):
"""
Create a callback function that records evaluation results.
"""
# TODO: Remove `replace("SNAPSHOT", "dev")` once the following issue is addressed:
# https://github.com/dmlc/xgboost/issues/6984
if Version(xgboost.__version__.replace("SNAPSHOT", "dev")) >= Version("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
def _log_posttraining_metadata(autologging_client, 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 autologging_client: An instance of `MlflowAutologgingQueueingClient` used for
efficiently logging run data to MLflow Tracking.
: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.
"""
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))
(X, y_true,
sample_weight) = _get_args_for_metrics(estimator.fit, args, kwargs)
# log common metrics and artifacts for estimators (classifier, regressor)
logged_metrics = _log_estimator_content(
autologging_client=autologging_client,
estimator=estimator,
prefix=_TRAINING_PREFIX,
run_id=mlflow.active_run().info.run_id,
X=X,
y_true=y_true,
sample_weight=sample_weight,
)
if y_true is None and not logged_metrics:
_logger.warning(
"Training metrics will not be recorded because training labels were not specified."
" To automatically record training metrics, provide training labels as inputs to"
" the model training function.")
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()`)
input_example = X[:INPUT_EXAMPLE_SAMPLE_ROWS]
return 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,
)
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:
log_model(
estimator.best_estimator_,
artifact_path="best_estimator",
signature=signature,
input_example=input_example,
)
if hasattr(estimator, "best_score_"):
autologging_client.log_metrics(
run_id=mlflow.active_run().info.run_id,
metrics={"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()
}
autologging_client.log_params(
run_id=mlflow.active_run().info.run_id,
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(
autologging_client=autologging_client,
cv_estimator=estimator,
parent_run=mlflow.active_run(),
max_tuning_runs=max_tuning_runs,
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 _log_posttraining_metadata(estimator, spark_model, params, input_df):
if _is_parameter_search_estimator(estimator):
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: AUTOLOGGING_INTEGRATION_NAME})
_create_child_runs_for_parameter_search(
parent_estimator=estimator,
parent_model=spark_model,
parent_run=mlflow.active_run(),
child_tags=child_tags,
)
except Exception:
import traceback
msg = (
"Encountered exception during creation of child runs for parameter search."
" Child runs may be missing. Exception: {}".format(
traceback.format_exc()))
_logger.warning(msg)
estimator_param_maps = _get_tuning_param_maps(
estimator,
estimator._autologging_metadata.uid_to_indexed_name_map)
metrics_dict, best_index = _get_param_search_metrics_and_best_index(
estimator, spark_model)
_log_parameter_search_results_as_artifact(
estimator_param_maps, metrics_dict,
mlflow.active_run().info.run_id)
# Log best_param_map as JSON artifact
best_param_map = estimator_param_maps[best_index]
mlflow.log_dict(best_param_map,
artifact_file="best_parameters.json")
# Log best_param_map as autologging parameters as well
_log_estimator_params({
f"best_{param_name}": param_value
for param_name, param_value in best_param_map.items()
})
if log_models:
if _should_log_model(spark_model):
from mlflow.models import infer_signature
from mlflow.pyspark.ml._autolog import (
cast_spark_df_with_vector_to_array,
get_feature_cols,
)
from mlflow.spark import _find_and_set_features_col_as_vector_if_needed
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
def _get_input_example_as_pd_df():
feature_cols = list(
get_feature_cols(input_df.schema, spark_model))
limited_input_df = input_df.select(feature_cols).limit(
INPUT_EXAMPLE_SAMPLE_ROWS)
return cast_spark_df_with_vector_to_array(
limited_input_df).toPandas()
def _infer_model_signature(input_example_slice):
input_slice_df = _find_and_set_features_col_as_vector_if_needed(
spark.createDataFrame(input_example_slice),
spark_model)
model_output = spark_model.transform(input_slice_df).drop(
*input_slice_df.columns)
return infer_signature(input_example_slice,
model_output.toPandas())
input_example, signature = resolve_input_example_and_signature(
_get_input_example_as_pd_df,
_infer_model_signature,
log_input_examples,
log_model_signatures,
_logger,
)
mlflow.spark.log_model(
spark_model,
artifact_path="model",
registered_model_name=registered_model_name,
input_example=input_example,
signature=signature,
)
if _is_parameter_search_model(spark_model):
mlflow.spark.log_model(
spark_model.bestModel,
artifact_path="best_model",
)
else:
_logger.warning(
_get_warning_msg_for_skip_log_model(spark_model))
