1361 lines
58 KiB
Python
1361 lines
58 KiB
Python
import datetime
|
|
import importlib
|
|
import logging
|
|
import os
|
|
import re
|
|
import shutil
|
|
import sys
|
|
import warnings
|
|
|
|
import cloudpickle
|
|
import yaml
|
|
|
|
import mlflow
|
|
from mlflow.entities import SourceType, ViewType
|
|
from mlflow.environment_variables import MLFLOW_RECIPES_EXECUTION_TARGET_STEP_NAME
|
|
from mlflow.exceptions import BAD_REQUEST, INVALID_PARAMETER_VALUE, MlflowException
|
|
from mlflow.models import Model
|
|
from mlflow.recipes.artifacts import (
|
|
DataframeArtifact,
|
|
HyperParametersArtifact,
|
|
ModelArtifact,
|
|
RunArtifact,
|
|
)
|
|
from mlflow.recipes.cards import BaseCard
|
|
from mlflow.recipes.step import BaseStep, StepClass
|
|
from mlflow.recipes.utils.execution import get_step_output_path
|
|
from mlflow.recipes.utils.metrics import (
|
|
_get_builtin_metrics,
|
|
_get_custom_metrics,
|
|
_get_error_fn,
|
|
_get_extended_task,
|
|
_get_model_type_from_template,
|
|
_get_primary_metric,
|
|
_load_custom_metrics,
|
|
transform_multiclass_metrics_dict,
|
|
)
|
|
from mlflow.recipes.utils.step import (
|
|
get_merged_eval_metrics,
|
|
get_pandas_data_profiles,
|
|
validate_classification_config,
|
|
)
|
|
from mlflow.recipes.utils.tracking import (
|
|
TrackingConfig,
|
|
apply_recipe_tracking_config,
|
|
get_recipe_tracking_config,
|
|
get_run_tags_env_vars,
|
|
log_code_snapshot,
|
|
)
|
|
from mlflow.recipes.utils.wrapped_recipe_model import WrappedRecipeModel
|
|
from mlflow.tracking import MlflowClient
|
|
from mlflow.tracking.fluent import _get_experiment_id
|
|
from mlflow.utils.databricks_utils import get_databricks_env_vars, get_databricks_run_url
|
|
from mlflow.utils.file_utils import TempDir
|
|
from mlflow.utils.mlflow_tags import (
|
|
MLFLOW_RECIPE_PROFILE_NAME,
|
|
MLFLOW_RECIPE_STEP_NAME,
|
|
MLFLOW_RECIPE_TEMPLATE_NAME,
|
|
MLFLOW_SOURCE_TYPE,
|
|
)
|
|
from mlflow.utils.string_utils import strip_prefix
|
|
|
|
_REBALANCING_CUTOFF = 5000
|
|
_REBALANCING_DEFAULT_RATIO = 0.3
|
|
_USER_DEFINED_TRAIN_STEP_MODULE = "steps.train"
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TrainStep(BaseStep):
|
|
MODEL_ARTIFACT_RELATIVE_PATH = "model"
|
|
SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH = "sk_model"
|
|
PREDICTED_TRAINING_DATA_RELATIVE_PATH = "predicted_training_data.parquet"
|
|
|
|
def __init__(self, step_config, recipe_root, recipe_config=None):
|
|
super().__init__(step_config, recipe_root)
|
|
self.tracking_config = TrackingConfig.from_dict(self.step_config)
|
|
self.recipe_config = recipe_config
|
|
|
|
def _validate_and_apply_step_config(self):
|
|
if "using" in self.step_config:
|
|
if self.step_config["using"] not in ["custom", "automl/flaml"]:
|
|
raise MlflowException(
|
|
f"Invalid train step configuration value {self.step_config['using']} for "
|
|
f"key 'using'. Supported values are: ['custom', 'automl/flaml']",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
self.step_config["using"] = "custom"
|
|
|
|
if "tuning" in self.step_config:
|
|
if "enabled" in self.step_config["tuning"] and isinstance(
|
|
self.step_config["tuning"]["enabled"], bool
|
|
):
|
|
self.step_config["tuning_enabled"] = self.step_config["tuning"]["enabled"]
|
|
else:
|
|
raise MlflowException(
|
|
"The 'tuning' configuration in the train step must include an "
|
|
"'enabled' key whose value is either true or false.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
if self.step_config["tuning_enabled"]:
|
|
if "sample_fraction" in self.step_config["tuning"]:
|
|
sample_fraction = float(self.step_config["tuning"]["sample_fraction"])
|
|
if 0 < sample_fraction <= 1.0:
|
|
self.step_config["sample_fraction"] = sample_fraction
|
|
else:
|
|
raise MlflowException(
|
|
"The tuning 'sample_fraction' configuration in the train step "
|
|
"must be between 0 and 1.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
self.step_config["sample_fraction"] = 1.0
|
|
|
|
if "algorithm" not in self.step_config["tuning"]:
|
|
self.step_config["tuning"]["algorithm"] = "hyperopt.rand.suggest"
|
|
|
|
if "parallelism" not in self.step_config["tuning"]:
|
|
self.step_config["tuning"]["parallelism"] = 1
|
|
|
|
if "max_trials" not in self.step_config["tuning"]:
|
|
raise MlflowException(
|
|
"The 'max_trials' configuration in the train step must be provided.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
if "parameters" not in self.step_config["tuning"]:
|
|
raise MlflowException(
|
|
"The 'parameters' configuration in the train step must be provided "
|
|
" when tuning is enabled.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
else:
|
|
self.step_config["tuning_enabled"] = False
|
|
|
|
if "estimator_params" not in self.step_config:
|
|
self.step_config["estimator_params"] = {}
|
|
|
|
self.target_col = self.step_config.get("target_col")
|
|
if self.target_col is None:
|
|
raise MlflowException(
|
|
"Missing target_col config in recipe config.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
self.recipe = self.step_config.get("recipe")
|
|
if self.recipe is None:
|
|
raise MlflowException(
|
|
"Missing recipe config in recipe config.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
self.positive_class = self.step_config.get("positive_class")
|
|
self.extended_task = _get_extended_task(self.recipe, self.positive_class)
|
|
self.rebalance_training_data = self.step_config.get("rebalance_training_data", True)
|
|
self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)
|
|
if "estimator_method" not in self.step_config and self.step_config["using"] == "custom":
|
|
raise MlflowException(
|
|
"Missing 'estimator_method' configuration in the train step, "
|
|
"which is using 'custom'.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
self.predict_prefix = self.step_config.get("predict_prefix", "predicted_")
|
|
self.predict_scores_for_all_classes = self.step_config.get(
|
|
"predict_scores_for_all_classes", True
|
|
)
|
|
self.primary_metric = _get_primary_metric(
|
|
self.step_config.get("primary_metric"), self.extended_task
|
|
)
|
|
builtin_metrics = _get_builtin_metrics(self.extended_task)
|
|
custom_metrics = _get_custom_metrics(self.step_config, self.extended_task)
|
|
self.user_defined_custom_metrics = {metric.name: metric for metric in custom_metrics}
|
|
self.evaluation_metrics = {metric.name: metric for metric in builtin_metrics}
|
|
self.evaluation_metrics.update(self.user_defined_custom_metrics)
|
|
self.evaluation_metrics_greater_is_better = {
|
|
metric.name: metric.greater_is_better for metric in builtin_metrics
|
|
}
|
|
self.evaluation_metrics_greater_is_better.update(
|
|
{metric.name: metric.greater_is_better for metric in custom_metrics}
|
|
)
|
|
if self.primary_metric is not None and self.primary_metric not in self.evaluation_metrics:
|
|
raise MlflowException(
|
|
f"The primary metric '{self.primary_metric}' is a custom metric, but its"
|
|
" corresponding custom metric configuration is missing from `recipe.yaml`.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
self.code_paths = [os.path.join(self.recipe_root, "steps")]
|
|
|
|
@classmethod
|
|
def construct_search_space_from_yaml(cls, params):
|
|
from hyperopt import hp
|
|
|
|
search_space = {}
|
|
for param_name, param_details in params.items():
|
|
if "values" in param_details:
|
|
param_details_to_pass = param_details["values"]
|
|
search_space[param_name] = hp.choice(param_name, param_details_to_pass)
|
|
elif "distribution" in param_details:
|
|
hp_tuning_fn = getattr(hp, param_details["distribution"])
|
|
param_details_to_pass = param_details.copy()
|
|
param_details_to_pass.pop("distribution")
|
|
search_space[param_name] = hp_tuning_fn(param_name, **param_details_to_pass)
|
|
else:
|
|
raise MlflowException(
|
|
f"Parameter {param_name} must contain either a list of 'values' or a "
|
|
f"'distribution' following hyperopt parameter expressions",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return search_space
|
|
|
|
@classmethod
|
|
def is_tuning_param_equal(cls, tuning_param, logged_param):
|
|
if isinstance(tuning_param, bool):
|
|
return tuning_param == bool(logged_param)
|
|
elif isinstance(tuning_param, int):
|
|
return tuning_param == int(logged_param)
|
|
elif isinstance(tuning_param, float):
|
|
return tuning_param == float(logged_param)
|
|
elif isinstance(tuning_param, str):
|
|
return tuning_param.strip() == logged_param.strip()
|
|
else:
|
|
return tuning_param == logged_param
|
|
|
|
def _fitted_estimator(self, *args):
|
|
if "calibrate_proba" in self.step_config:
|
|
return self._calibrated_classifier_fitted_estimator(*args)
|
|
else:
|
|
return self._label_encoded_fitted_estimator(*args)
|
|
|
|
def _calibrated_classifier_fitted_estimator(self, estimator, X_train, y_train):
|
|
original_estimator = estimator.fit(X_train, y_train)
|
|
if "classification" in self.recipe:
|
|
from sklearn.calibration import CalibratedClassifierCV
|
|
|
|
estimator = CalibratedClassifierCV(
|
|
estimator, method=self.step_config["calibrate_proba"]
|
|
)
|
|
estimator.fit(X_train, y_train)
|
|
|
|
return estimator, {"original_estimator": original_estimator}
|
|
|
|
def _label_encoded_fitted_estimator(self, estimator, X_train, y_train):
|
|
from sklearn.preprocessing import LabelEncoder
|
|
|
|
label_encoder = None
|
|
target_column_class_labels = None
|
|
encoded_y_train = y_train
|
|
if "classification" in self.recipe:
|
|
label_encoder = LabelEncoder()
|
|
label_encoder.fit(y_train)
|
|
target_column_class_labels = label_encoder.classes_
|
|
import pandas as pd
|
|
|
|
encoded_y_train = pd.Series(label_encoder.transform(y_train))
|
|
|
|
def inverse_label_encoder(predicted_output):
|
|
if not label_encoder:
|
|
return predicted_output
|
|
|
|
return label_encoder.inverse_transform(predicted_output)
|
|
|
|
estimator.fit(X_train, encoded_y_train)
|
|
original_predict = estimator.predict
|
|
|
|
def wrapped_predict(*args, **kwargs):
|
|
return inverse_label_encoder(original_predict(*args, **kwargs))
|
|
|
|
estimator.predict = wrapped_predict
|
|
|
|
return estimator, {"target_column_class_labels": target_column_class_labels}
|
|
|
|
def _run(self, output_directory):
|
|
def my_warn(*args, **kwargs):
|
|
timestamp = datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S")
|
|
stacklevel = 1 if "stacklevel" not in kwargs else kwargs["stacklevel"]
|
|
frame = sys._getframe(stacklevel)
|
|
filename = frame.f_code.co_filename
|
|
lineno = frame.f_lineno
|
|
original_message = args[0] if len(args) > 0 else kwargs.get("message", "")
|
|
message = f"{timestamp} {filename}:{lineno}: {original_message}\n"
|
|
with open(os.path.join(output_directory, "warning_logs.txt"), "a") as f:
|
|
f.write(message)
|
|
|
|
original_warn = warnings.warn
|
|
warnings.warn = my_warn
|
|
try:
|
|
import numpy as np
|
|
import pandas as pd
|
|
import sklearn
|
|
from sklearn.pipeline import make_pipeline
|
|
from sklearn.utils.class_weight import compute_class_weight
|
|
|
|
from mlflow.models import infer_signature
|
|
|
|
with open(os.path.join(output_directory, "warning_logs.txt"), "w"):
|
|
pass
|
|
|
|
apply_recipe_tracking_config(self.tracking_config)
|
|
|
|
transformed_training_data_path = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name="transform",
|
|
relative_path="transformed_training_data.parquet",
|
|
)
|
|
train_df = pd.read_parquet(transformed_training_data_path)
|
|
validate_classification_config(
|
|
self.task, self.positive_class, train_df, self.target_col
|
|
)
|
|
self.using_rebalancing = False
|
|
if self.extended_task == "classification/binary":
|
|
classes = np.unique(train_df[self.target_col])
|
|
class_weights = compute_class_weight(
|
|
class_weight="balanced",
|
|
classes=classes,
|
|
y=train_df[self.target_col],
|
|
)
|
|
self.original_class_weights = dict(zip(classes, class_weights))
|
|
if self.rebalance_training_data and len(classes) == 2:
|
|
if len(train_df) > _REBALANCING_CUTOFF:
|
|
self.using_rebalancing = True
|
|
train_df = self._rebalance_classes(train_df)
|
|
else:
|
|
_logger.info(
|
|
f"Training data has less than {_REBALANCING_CUTOFF} rows, "
|
|
f"skipping rebalancing."
|
|
)
|
|
|
|
X_train, y_train = train_df.drop(columns=[self.target_col]), train_df[self.target_col]
|
|
|
|
transformed_validation_data_path = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name="transform",
|
|
relative_path="transformed_validation_data.parquet",
|
|
)
|
|
validation_df = pd.read_parquet(transformed_validation_data_path)
|
|
|
|
raw_training_data_path = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name="split",
|
|
relative_path="train.parquet",
|
|
)
|
|
raw_train_df = pd.read_parquet(raw_training_data_path)
|
|
raw_X_train = raw_train_df.drop(columns=[self.target_col])
|
|
|
|
raw_validation_data_path = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name="split",
|
|
relative_path="validation.parquet",
|
|
)
|
|
raw_validation_df = pd.read_parquet(raw_validation_data_path)
|
|
|
|
transformer_path = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name="transform",
|
|
relative_path="transformer.pkl",
|
|
)
|
|
|
|
tags = {
|
|
MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.RECIPE),
|
|
MLFLOW_RECIPE_TEMPLATE_NAME: self.step_config["recipe"],
|
|
MLFLOW_RECIPE_PROFILE_NAME: self.step_config["profile"],
|
|
MLFLOW_RECIPE_STEP_NAME: MLFLOW_RECIPES_EXECUTION_TARGET_STEP_NAME.get(),
|
|
}
|
|
|
|
run_name = self.tracking_config.run_name
|
|
best_estimator_params = None
|
|
mlflow.autolog(log_models=False, silent=True)
|
|
with mlflow.start_run(run_name=run_name, tags=tags) as run:
|
|
estimator = self._resolve_estimator(
|
|
X_train, y_train, validation_df, run, output_directory
|
|
)
|
|
best_estimator_params = estimator.get_params()
|
|
fitted_estimator, additional_fitted_args = self._fitted_estimator(
|
|
estimator, X_train, y_train
|
|
)
|
|
logged_estimator = self._log_estimator_to_mlflow(fitted_estimator, X_train)
|
|
|
|
# Create a recipe consisting of the transformer+model for test data evaluation
|
|
with open(transformer_path, "rb") as f:
|
|
transformer = cloudpickle.load(f)
|
|
mlflow.sklearn.log_model(
|
|
transformer, "transform/transformer", code_paths=self.code_paths
|
|
)
|
|
|
|
trained_pipeline = make_pipeline(transformer, fitted_estimator)
|
|
# Creating a wrapped recipe model which exposes a single predict function
|
|
# so it can output both predict and predict_proba(for a classification problem)
|
|
# at the same time.
|
|
wrapped_model = WrappedRecipeModel(
|
|
self.predict_scores_for_all_classes,
|
|
self.predict_prefix,
|
|
target_column_class_labels=additional_fitted_args.get(
|
|
"target_column_class_labels"
|
|
),
|
|
)
|
|
|
|
model_uri = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name=self.name,
|
|
relative_path=TrainStep.MODEL_ARTIFACT_RELATIVE_PATH,
|
|
)
|
|
sklearn_model_uri = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name=self.name,
|
|
relative_path=TrainStep.SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH,
|
|
)
|
|
if os.path.exists(model_uri):
|
|
shutil.rmtree(model_uri)
|
|
if os.path.exists(sklearn_model_uri):
|
|
shutil.rmtree(sklearn_model_uri)
|
|
|
|
# Saving the sklearn model as a separate output since `mlflow.evaluate()`, which is
|
|
# used in evaluate step of the recipe, needs this model's sklearn representation
|
|
# to computes metrics (the pyfunc representation of the user-facing model logged to
|
|
# MLflow Tracking is not currently compatible with `mlflow.evaluate()`)
|
|
mlflow.sklearn.save_model(trained_pipeline, sklearn_model_uri)
|
|
artifacts = {"model_path": sklearn_model_uri}
|
|
with TempDir() as tmp:
|
|
# Saving a temp model so that the output schema (signature) of the model's
|
|
# pyfunc representation can be inferred and included when logging the model
|
|
# to MLflow Tracking. Unfortunately, there is currently no easy way to infer
|
|
# the model's signature without first saving a copy of it, and there is no easy
|
|
# way to add an inferred signature to an existing model
|
|
pyfunc_model_tmp_path = os.path.join(tmp.path(), "pyfunc_model")
|
|
mlflow.pyfunc.save_model(
|
|
path=pyfunc_model_tmp_path,
|
|
python_model=wrapped_model,
|
|
artifacts=artifacts,
|
|
)
|
|
tempModel = mlflow.pyfunc.load_model(pyfunc_model_tmp_path)
|
|
model_schema = infer_signature(
|
|
raw_X_train, tempModel.predict(raw_X_train.copy())
|
|
)
|
|
mlflow.pyfunc.save_model(
|
|
path=model_uri,
|
|
python_model=wrapped_model,
|
|
artifacts=artifacts,
|
|
signature=model_schema,
|
|
code_path=self.code_paths,
|
|
)
|
|
model = mlflow.pyfunc.load_model(model_uri)
|
|
# Adding a sklearn flavor to the pyfunc model so models could be loaded easily
|
|
# using mlflow.sklearn.load_model
|
|
tmp_model_info = Model.load(model_uri)
|
|
model_data_subpath = os.path.join(
|
|
"artifacts", TrainStep.SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH, "model.pkl"
|
|
)
|
|
model_info = Model(
|
|
artifact_path="train/model",
|
|
run_id=run.info.run_id,
|
|
utc_time_created=tmp_model_info.utc_time_created,
|
|
flavors=tmp_model_info.flavors,
|
|
signature=tmp_model_info.signature, # ModelSignature
|
|
saved_input_example_info=tmp_model_info.saved_input_example_info,
|
|
model_uuid=tmp_model_info.model_uuid,
|
|
mlflow_version=tmp_model_info.mlflow_version,
|
|
metadata=tmp_model_info.metadata,
|
|
)
|
|
model_info.add_flavor(
|
|
mlflow.sklearn.FLAVOR_NAME,
|
|
pickled_model=model_data_subpath,
|
|
sklearn_version=sklearn.__version__,
|
|
serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE,
|
|
code="code",
|
|
)
|
|
model_info.save(f"{model_uri}/MLmodel")
|
|
mlflow.log_artifacts(model_uri, "train/model")
|
|
|
|
with open(os.path.join(output_directory, "run_id"), "w") as f:
|
|
f.write(run.info.run_id)
|
|
log_code_snapshot(
|
|
self.recipe_root, run.info.run_id, recipe_config=self.recipe_config
|
|
)
|
|
|
|
eval_metrics = {}
|
|
for dataset_name, (dataset, metric_prefix) in {
|
|
"training": (train_df, "training_"),
|
|
"validation": (validation_df, "val_"),
|
|
}.items():
|
|
eval_config = {
|
|
"log_model_explainability": False,
|
|
"metric_prefix": metric_prefix,
|
|
}
|
|
if self.positive_class is not None:
|
|
eval_config["pos_label"] = self.positive_class
|
|
eval_result = mlflow.evaluate(
|
|
model=logged_estimator.model_uri,
|
|
data=dataset,
|
|
targets=self.target_col,
|
|
model_type=_get_model_type_from_template(self.recipe),
|
|
evaluators="default",
|
|
extra_metrics=_load_custom_metrics(
|
|
self.recipe_root,
|
|
self.evaluation_metrics.values(),
|
|
),
|
|
evaluator_config=eval_config,
|
|
)
|
|
eval_result.save(os.path.join(output_directory, f"eval_{dataset_name}"))
|
|
eval_metrics[dataset_name] = {
|
|
strip_prefix(k, metric_prefix): v for k, v in eval_result.metrics.items()
|
|
}
|
|
|
|
target_data = raw_validation_df[self.target_col]
|
|
prediction_result = model.predict(raw_validation_df.drop(self.target_col, axis=1))
|
|
|
|
use_probability_for_error_rate = False
|
|
if isinstance(prediction_result, pd.DataFrame) and {
|
|
f"{self.predict_prefix}label",
|
|
f"{self.predict_prefix}score",
|
|
}.issubset(prediction_result.columns):
|
|
if self.positive_class:
|
|
prediction_result_for_error = prediction_result[
|
|
f"{self.predict_prefix}score_{self.positive_class}"
|
|
].values
|
|
# use_probability_for_error_rate to true to compute error function
|
|
# based on positive class
|
|
use_probability_for_error_rate = True
|
|
else:
|
|
prediction_result_for_error = prediction_result[
|
|
f"{self.predict_prefix}score"
|
|
].values
|
|
|
|
prediction_result = prediction_result[f"{self.predict_prefix}label"].values
|
|
else:
|
|
prediction_result_for_error = prediction_result
|
|
error_fn = _get_error_fn(
|
|
self.recipe,
|
|
use_probability=use_probability_for_error_rate,
|
|
positive_class=self.positive_class,
|
|
)
|
|
pred_and_error_df = pd.DataFrame(
|
|
{
|
|
"target": target_data,
|
|
"prediction": prediction_result,
|
|
"error": error_fn(prediction_result_for_error, target_data.to_numpy()),
|
|
}
|
|
)
|
|
calibrated_plot = None
|
|
train_predictions = model.predict(raw_train_df.drop(self.target_col, axis=1))
|
|
if isinstance(train_predictions, pd.DataFrame) and {
|
|
f"{self.predict_prefix}label",
|
|
f"{self.predict_prefix}score",
|
|
}.issubset(train_predictions.columns):
|
|
predicted_training_data = raw_train_df.assign(
|
|
predicted_data=train_predictions[f"{self.predict_prefix}label"],
|
|
predicted_score=train_predictions[f"{self.predict_prefix}score"].values,
|
|
)
|
|
if self.positive_class:
|
|
worst_examples_df = BaseStep._generate_worst_examples_dataframe(
|
|
raw_train_df,
|
|
train_predictions[f"{self.predict_prefix}label"].values,
|
|
error_fn(
|
|
train_predictions[
|
|
f"{self.predict_prefix}score_{self.positive_class}"
|
|
].values,
|
|
raw_train_df[self.target_col].to_numpy(),
|
|
),
|
|
self.target_col,
|
|
)
|
|
|
|
if "calibrate_proba" in self.step_config and hasattr(
|
|
additional_fitted_args.get("original_estimator"), "predict_proba"
|
|
):
|
|
from sklearn.calibration import CalibrationDisplay
|
|
|
|
calibrated_plot = CalibrationDisplay.from_estimator(
|
|
additional_fitted_args.get("original_estimator"),
|
|
raw_train_df.drop(self.target_col, axis=1),
|
|
raw_train_df[self.target_col],
|
|
pos_label=self.positive_class,
|
|
)
|
|
else:
|
|
# compute worst examples data_frame only if positive class exists
|
|
worst_examples_df = pd.DataFrame()
|
|
|
|
else:
|
|
predicted_training_data = raw_train_df.assign(predicted_data=train_predictions)
|
|
worst_examples_df = BaseStep._generate_worst_examples_dataframe(
|
|
raw_train_df,
|
|
train_predictions,
|
|
error_fn(train_predictions, raw_train_df[self.target_col].to_numpy()),
|
|
self.target_col,
|
|
)
|
|
predicted_training_data.to_parquet(
|
|
os.path.join(output_directory, TrainStep.PREDICTED_TRAINING_DATA_RELATIVE_PATH)
|
|
)
|
|
|
|
leaderboard_df = None
|
|
try:
|
|
leaderboard_df = self._get_leaderboard_df(run, eval_metrics)
|
|
except Exception as e:
|
|
_logger.warning(
|
|
"Failed to build model leaderboard due to unexpected failure: %s", e
|
|
)
|
|
tuning_df = None
|
|
if self.step_config["tuning_enabled"]:
|
|
try:
|
|
tuning_df = self._get_tuning_df(run, params=best_estimator_params.keys())
|
|
except Exception as e:
|
|
_logger.warning(
|
|
"Failed to build tuning results table due to unexpected failure: %s", e
|
|
)
|
|
|
|
card = self._build_step_card(
|
|
eval_metrics=eval_metrics,
|
|
pred_and_error_df=pred_and_error_df,
|
|
model_schema=model_schema,
|
|
run_id=run.info.run_id,
|
|
model_uri=model_uri,
|
|
worst_examples_df=worst_examples_df,
|
|
train_df=raw_train_df,
|
|
output_directory=output_directory,
|
|
leaderboard_df=leaderboard_df,
|
|
tuning_df=tuning_df,
|
|
calibrated_plot=calibrated_plot,
|
|
)
|
|
card.save_as_html(output_directory)
|
|
for step_name in ("ingest", "split", "transform", "train"):
|
|
self._log_step_card(run.info.run_id, step_name)
|
|
return card
|
|
finally:
|
|
warnings.warn = original_warn
|
|
|
|
def _get_user_defined_estimator(self, X_train, y_train, validation_df, run, output_directory):
|
|
sys.path.append(self.recipe_root)
|
|
estimator_fn = getattr(
|
|
importlib.import_module(_USER_DEFINED_TRAIN_STEP_MODULE),
|
|
self.step_config["estimator_method"],
|
|
)
|
|
estimator_hardcoded_params = self.step_config["estimator_params"]
|
|
|
|
# if using rebalancing pass in original class weights to preserve original distribution
|
|
if self.using_rebalancing:
|
|
estimator_hardcoded_params["class_weight"] = self.original_class_weights
|
|
|
|
if self.step_config["tuning_enabled"]:
|
|
estimator_hardcoded_params, best_hp_params = self._tune_and_get_best_estimator_params(
|
|
run.info.run_id,
|
|
estimator_hardcoded_params,
|
|
estimator_fn,
|
|
X_train,
|
|
y_train,
|
|
validation_df,
|
|
)
|
|
best_combined_params = dict(estimator_hardcoded_params, **best_hp_params)
|
|
estimator = estimator_fn(best_combined_params)
|
|
all_estimator_params = estimator.get_params()
|
|
default_params_keys = all_estimator_params.keys() - best_combined_params.keys()
|
|
default_params = {k: all_estimator_params[k] for k in default_params_keys}
|
|
self._write_best_parameters_outputs(
|
|
output_directory,
|
|
best_hp_params=best_hp_params,
|
|
best_hardcoded_params=estimator_hardcoded_params,
|
|
default_params=default_params,
|
|
)
|
|
elif len(estimator_hardcoded_params) > 0:
|
|
estimator = estimator_fn(estimator_hardcoded_params)
|
|
all_estimator_params = estimator.get_params()
|
|
default_params_keys = all_estimator_params.keys() - estimator_hardcoded_params.keys()
|
|
default_params = {k: all_estimator_params[k] for k in default_params_keys}
|
|
self._write_best_parameters_outputs(
|
|
output_directory,
|
|
best_hardcoded_params=estimator_hardcoded_params,
|
|
default_params=default_params,
|
|
)
|
|
else:
|
|
estimator = estimator_fn()
|
|
default_params = estimator.get_params()
|
|
self._write_best_parameters_outputs(output_directory, default_params=default_params)
|
|
return estimator
|
|
|
|
def _resolve_estimator_plugin(self, plugin_str, X_train, y_train, output_directory):
|
|
plugin_str = plugin_str.replace("/", ".")
|
|
estimator_fn = importlib.import_module(
|
|
f"mlflow.recipes.steps.{plugin_str}"
|
|
).get_estimator_and_best_params
|
|
estimator, best_parameters = estimator_fn(
|
|
X_train,
|
|
y_train,
|
|
self.task,
|
|
self.extended_task,
|
|
self.step_config,
|
|
self.recipe_root,
|
|
self.evaluation_metrics,
|
|
self.primary_metric,
|
|
)
|
|
self.best_estimator_name = estimator.__class__.__name__
|
|
self.best_estimator_class = (
|
|
f"{estimator.__class__.__module__}.{estimator.__class__.__name__}"
|
|
)
|
|
self.best_parameters = best_parameters
|
|
self._write_best_parameters_outputs(output_directory, automl_params=best_parameters)
|
|
return estimator
|
|
|
|
def _resolve_estimator(self, X_train, y_train, validation_df, run, output_directory):
|
|
using_plugin = self.step_config.get("using", "custom")
|
|
if using_plugin == "custom":
|
|
return self._get_user_defined_estimator(
|
|
X_train, y_train, validation_df, run, output_directory
|
|
)
|
|
else:
|
|
return self._resolve_estimator_plugin(using_plugin, X_train, y_train, output_directory)
|
|
|
|
def _get_leaderboard_df(self, run, eval_metrics):
|
|
import pandas as pd
|
|
|
|
mlflow_client = MlflowClient()
|
|
exp_id = _get_experiment_id()
|
|
|
|
primary_metric_greater_is_better = self.evaluation_metrics[
|
|
self.primary_metric
|
|
].greater_is_better
|
|
primary_metric_order = "DESC" if primary_metric_greater_is_better else "ASC"
|
|
|
|
search_max_results = 100
|
|
search_result = mlflow_client.search_runs(
|
|
experiment_ids=exp_id,
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
max_results=search_max_results,
|
|
order_by=[f"metrics.{self.primary_metric} {primary_metric_order}"],
|
|
)
|
|
|
|
metric_names = self.evaluation_metrics.keys()
|
|
|
|
leaderboard_items = []
|
|
for old_run in search_result:
|
|
if all(metric_key in old_run.data.metrics for metric_key in metric_names):
|
|
leaderboard_items.append(
|
|
{
|
|
"Run ID": old_run.info.run_id,
|
|
"Run Time": datetime.datetime.fromtimestamp(
|
|
old_run.info.start_time // 1000
|
|
),
|
|
**{
|
|
metric_name: old_run.data.metrics[metric_key]
|
|
for metric_name, metric_key in zip(metric_names, metric_names)
|
|
},
|
|
}
|
|
)
|
|
|
|
top_leaderboard_items = [
|
|
{"Model Rank": i + 1, **t} for i, t in enumerate(leaderboard_items[:2])
|
|
]
|
|
|
|
if (
|
|
len(top_leaderboard_items) == 2
|
|
and top_leaderboard_items[0][self.primary_metric]
|
|
== top_leaderboard_items[1][self.primary_metric]
|
|
):
|
|
# If top1 and top2 model primary metrics are equal,
|
|
# then their rank are both 1.
|
|
top_leaderboard_items[1]["Model Rank"] = "1"
|
|
|
|
top_leaderboard_item_index_values = ["Best", "2nd Best"][: len(top_leaderboard_items)]
|
|
|
|
latest_model_item = {
|
|
"Run ID": run.info.run_id,
|
|
"Run Time": datetime.datetime.fromtimestamp(run.info.start_time // 1000),
|
|
**eval_metrics["validation"],
|
|
}
|
|
|
|
for i, leaderboard_item in enumerate(leaderboard_items):
|
|
latest_value = latest_model_item[self.primary_metric]
|
|
historical_value = leaderboard_item[self.primary_metric]
|
|
if (primary_metric_greater_is_better and latest_value >= historical_value) or (
|
|
not primary_metric_greater_is_better and latest_value <= historical_value
|
|
):
|
|
latest_model_item["Model Rank"] = str(i + 1)
|
|
break
|
|
else:
|
|
latest_model_item["Model Rank"] = f"> {len(leaderboard_items)}"
|
|
|
|
# metric columns order: primary metric, custom metrics, builtin metrics.
|
|
def sorter(m):
|
|
if m == self.primary_metric:
|
|
return 0, m
|
|
elif self.evaluation_metrics[m].custom_function is not None:
|
|
return 1, m
|
|
else:
|
|
return 2, m
|
|
|
|
metric_columns = sorted(metric_names, key=sorter)
|
|
|
|
return (
|
|
pd.DataFrame.from_records(
|
|
[latest_model_item, *top_leaderboard_items],
|
|
columns=["Model Rank", *metric_columns, "Run Time", "Run ID"],
|
|
)
|
|
.apply(
|
|
lambda s: s.map(lambda x: f"{x:.6g}") if s.name in metric_names else s,
|
|
axis=0,
|
|
)
|
|
.set_axis(["Latest"] + top_leaderboard_item_index_values, axis="index")
|
|
.transpose()
|
|
)
|
|
|
|
def _get_tuning_df(self, run, params=None):
|
|
exp_id = _get_experiment_id()
|
|
primary_metric_tag = f"metrics.{self.primary_metric}"
|
|
order_str = (
|
|
"DESC" if self.evaluation_metrics_greater_is_better[self.primary_metric] else "ASC"
|
|
)
|
|
tuning_runs = mlflow.search_runs(
|
|
[exp_id],
|
|
filter_string=f"tags.mlflow.parentRunId like '{run.info.run_id}'",
|
|
order_by=[f"{primary_metric_tag} {order_str}", "attribute.start_time ASC"],
|
|
)
|
|
if params:
|
|
params = [f"params.{param}" for param in params]
|
|
tuning_runs = tuning_runs.filter([f"metrics.{self.primary_metric}", *params])
|
|
else:
|
|
tuning_runs = tuning_runs.filter([f"metrics.{self.primary_metric}"])
|
|
tuning_runs = tuning_runs.reset_index().rename(
|
|
columns={"index": "Model Rank", primary_metric_tag: self.primary_metric}
|
|
)
|
|
tuning_runs["Model Rank"] += 1
|
|
return tuning_runs.head(10)
|
|
|
|
def _build_step_card(
|
|
self,
|
|
eval_metrics,
|
|
pred_and_error_df,
|
|
model_schema,
|
|
run_id,
|
|
model_uri,
|
|
worst_examples_df,
|
|
train_df,
|
|
output_directory,
|
|
leaderboard_df=None,
|
|
tuning_df=None,
|
|
calibrated_plot=None,
|
|
):
|
|
import pandas as pd
|
|
from sklearn import set_config
|
|
from sklearn.utils import estimator_html_repr
|
|
|
|
card = BaseCard(self.recipe_name, self.name)
|
|
# Tab 0: model performance summary.
|
|
metric_df = (
|
|
get_merged_eval_metrics(
|
|
eval_metrics,
|
|
ordered_metric_names=[
|
|
self.primary_metric,
|
|
*self.user_defined_custom_metrics.keys(),
|
|
],
|
|
)
|
|
.reset_index()
|
|
.rename(columns={"index": "Metric"})
|
|
)
|
|
|
|
def row_style(row):
|
|
if row.Metric == self.primary_metric or row.Metric in self.user_defined_custom_metrics:
|
|
return pd.Series("font-weight: bold", row.index)
|
|
else:
|
|
return pd.Series("", row.index)
|
|
|
|
metric_table_html = BaseCard.render_table(
|
|
metric_df.style.format({"training": "{:.6g}", "validation": "{:.6g}"}).apply(
|
|
row_style, axis=1
|
|
)
|
|
)
|
|
|
|
# Tab 1: Model performance.
|
|
card.add_tab(
|
|
"Model Performance",
|
|
"<h3 class='section-title'>Summary Metrics (Validation)</h3>{{ METRICS }} ",
|
|
).add_html("METRICS", metric_table_html)
|
|
|
|
if not self.skip_data_profiling:
|
|
# Tab 2: Prediction and error data profile.
|
|
pred_and_error_df_profile = get_pandas_data_profiles(
|
|
[
|
|
[
|
|
"Predictions and Errors (Validation Dataset)",
|
|
pred_and_error_df.reset_index(drop=True),
|
|
]
|
|
]
|
|
)
|
|
card.add_tab("Data Profile (Predictions)", "{{PROFILE}}").add_pandas_profile(
|
|
"PROFILE", pred_and_error_df_profile
|
|
)
|
|
# Tab 3: Model architecture.
|
|
set_config(display="diagram")
|
|
model_repr = estimator_html_repr(
|
|
mlflow.sklearn.load_model(
|
|
os.path.join(model_uri, "artifacts", TrainStep.SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH)
|
|
)
|
|
)
|
|
card.add_tab("Model Architecture", "{{MODEL_ARCH}}").add_html("MODEL_ARCH", model_repr)
|
|
|
|
# Tab 4: Inferred model (transformer + estimator) schema.
|
|
def render_schema(inputs, title):
|
|
from mlflow.types import ColSpec
|
|
|
|
table = BaseCard.render_table(
|
|
{
|
|
"Name": " " + (spec.name or "-"),
|
|
"Type": repr(spec.type) if isinstance(spec, ColSpec) else repr(spec),
|
|
}
|
|
for spec in inputs
|
|
)
|
|
return f'<div style="margin: 5px"><h2>{title}</h2>{table}</div>'
|
|
|
|
schema_tables = [render_schema(model_schema.inputs.inputs, "Inputs")]
|
|
if model_schema.outputs:
|
|
schema_tables += [render_schema(model_schema.outputs.inputs, "Outputs")]
|
|
|
|
card.add_tab("Model Schema", "{{MODEL_SCHEMA}}").add_html(
|
|
"MODEL_SCHEMA",
|
|
'<div style="display: flex">{tables}</div>'.format(tables="\n".join(schema_tables)),
|
|
)
|
|
|
|
# Tab 5: Examples with Largest Prediction Error
|
|
if not worst_examples_df.empty:
|
|
(
|
|
card.add_tab("Worst Predictions", "{{ WORST_EXAMPLES_TABLE }}").add_html(
|
|
"WORST_EXAMPLES_TABLE", BaseCard.render_table(worst_examples_df)
|
|
)
|
|
)
|
|
|
|
if calibrated_plot:
|
|
calibrated_plot_location = os.path.join(output_directory, "calibrated_plot_location")
|
|
calibrated_plot.figure_.savefig(calibrated_plot_location, format="png")
|
|
(
|
|
card.add_tab("Prob. Calibration", "{{ CALIBRATED_PLOT }}").add_image(
|
|
"CALIBRATED_PLOT", calibrated_plot_location
|
|
)
|
|
)
|
|
|
|
# Tab 6: Worst predictions profile vs train profile.
|
|
if not self.skip_data_profiling:
|
|
worst_prediction_profile = get_pandas_data_profiles(
|
|
[
|
|
["Worst Predictions", worst_examples_df.reset_index(drop=True)],
|
|
["Train", train_df.reset_index(drop=True)],
|
|
]
|
|
)
|
|
card.add_tab(
|
|
"Data Profile (Worst vs Train)", "{{ WORST_EXAMPLES_COMP }}"
|
|
).add_pandas_profile("WORST_EXAMPLES_COMP", worst_prediction_profile)
|
|
|
|
# Tab 7: Leaderboard
|
|
if leaderboard_df is not None:
|
|
(
|
|
card.add_tab("Leaderboard", "{{ LEADERBOARD_TABLE }}").add_html(
|
|
"LEADERBOARD_TABLE", BaseCard.render_table(leaderboard_df, hide_index=False)
|
|
)
|
|
)
|
|
|
|
# Tab 8: Best Parameters (AutoML and Tuning)
|
|
is_automl_run = self.step_config["using"].startswith("automl")
|
|
best_parameters_yaml = os.path.join(output_directory, "best_parameters.yaml")
|
|
|
|
if os.path.exists(best_parameters_yaml):
|
|
best_parameters_card_tab = card.add_tab(
|
|
f"Best Parameters {' (AutoML)' if is_automl_run else ''}",
|
|
"{{ BEST_PARAMETERS }} ",
|
|
)
|
|
|
|
if is_automl_run:
|
|
automl_estimator_str = (
|
|
f"<b>Best estimator:</b><br>"
|
|
f"<pre>{self.best_estimator_name}</pre><br>"
|
|
f"<b>Best estimator class:</b><br>"
|
|
f"<pre>{self.best_estimator_class}</pre><br><br>"
|
|
)
|
|
else:
|
|
automl_estimator_str = ""
|
|
|
|
with open(best_parameters_yaml) as f:
|
|
best_parameters = f.read()
|
|
best_parameters_card_tab.add_html(
|
|
"BEST_PARAMETERS",
|
|
f"{automl_estimator_str}<b>Best parameters:</b><br>"
|
|
f"<pre>{best_parameters}</pre><br><br>",
|
|
)
|
|
|
|
# Tab 9: HP trials
|
|
if tuning_df is not None:
|
|
tuning_trials_card_tab = card.add_tab(
|
|
"Tuning Trials",
|
|
"{{ SEARCH_SPACE }}" + "{{ TUNING_TABLE_TITLE }}" + "{{ TUNING_TABLE }}",
|
|
)
|
|
tuning_params = yaml.dump(self.step_config["tuning"]["parameters"])
|
|
tuning_trials_card_tab.add_html(
|
|
"SEARCH_SPACE",
|
|
f"<b>Tuning search space:</b> <br><pre>{tuning_params}</pre><br><br>",
|
|
)
|
|
tuning_trials_card_tab.add_html("TUNING_TABLE_TITLE", "<b>Tuning results:</b><br>")
|
|
tuning_trials_card_tab.add_html(
|
|
"TUNING_TABLE",
|
|
BaseCard.render_table(
|
|
tuning_df.style.apply(
|
|
lambda row: pd.Series("font-weight: bold", row.index),
|
|
axis=1,
|
|
subset=tuning_df.index[0],
|
|
),
|
|
hide_index=True,
|
|
),
|
|
)
|
|
|
|
# Tab 10: Warning log outputs.
|
|
warning_output_path = os.path.join(output_directory, "warning_logs.txt")
|
|
if os.path.exists(warning_output_path):
|
|
warnings_output_tab = card.add_tab("Warning Logs", "{{ STEP_WARNINGS }}")
|
|
with open(warning_output_path) as f:
|
|
warnings_output_tab.add_html("STEP_WARNINGS", f"<pre>{f.read()}</pre>")
|
|
|
|
# Tab 11: Run summary.
|
|
run_card_tab = card.add_tab(
|
|
"Run Summary",
|
|
"{{ RUN_ID }} " + "{{ MODEL_URI }}" + "{{ EXE_DURATION }}" + "{{ LAST_UPDATE_TIME }}",
|
|
)
|
|
model_uri_path = f"runs:/{run_id}/train/model"
|
|
run_url = get_databricks_run_url(
|
|
tracking_uri=mlflow.get_tracking_uri(),
|
|
run_id=run_id,
|
|
)
|
|
model_url = get_databricks_run_url(
|
|
tracking_uri=mlflow.get_tracking_uri(),
|
|
run_id=run_id,
|
|
artifact_path=re.sub(rf"^.*?{run_id}", "", model_uri_path),
|
|
)
|
|
|
|
if run_url is not None:
|
|
run_card_tab.add_html(
|
|
"RUN_ID", f"<b>MLflow Run ID:</b> <a href={run_url}>{run_id}</a><br><br>"
|
|
)
|
|
else:
|
|
run_card_tab.add_markdown("RUN_ID", f"**MLflow Run ID:** `{run_id}`")
|
|
|
|
if model_url is not None:
|
|
run_card_tab.add_html(
|
|
"MODEL_URI", f"<b>MLflow Model URI:</b> <a href={model_url}>{model_uri_path}</a>"
|
|
)
|
|
else:
|
|
run_card_tab.add_markdown("MODEL_URI", f"**MLflow Model URI:** `{model_uri_path}`")
|
|
|
|
return card
|
|
|
|
@classmethod
|
|
def from_recipe_config(cls, recipe_config, recipe_root):
|
|
step_config = {}
|
|
if recipe_config.get("steps", {}).get("train", {}) is not None:
|
|
step_config.update(recipe_config.get("steps", {}).get("train", {}))
|
|
if recipe_config.get("custom_metrics") is not None:
|
|
step_config["custom_metrics"] = recipe_config["custom_metrics"]
|
|
if recipe_config.get("primary_metric") is not None:
|
|
step_config["primary_metric"] = recipe_config["primary_metric"]
|
|
step_config["recipe"] = recipe_config.get("recipe")
|
|
step_config["profile"] = recipe_config.get("profile")
|
|
step_config["target_col"] = recipe_config.get("target_col")
|
|
if "positive_class" in recipe_config:
|
|
step_config["positive_class"] = recipe_config.get("positive_class")
|
|
step_config.update(
|
|
get_recipe_tracking_config(
|
|
recipe_root_path=recipe_root,
|
|
recipe_config=recipe_config,
|
|
).to_dict()
|
|
)
|
|
return cls(step_config, recipe_root, recipe_config=recipe_config)
|
|
|
|
@property
|
|
def name(self):
|
|
return "train"
|
|
|
|
@property
|
|
def environment(self):
|
|
environ = get_databricks_env_vars(tracking_uri=self.tracking_config.tracking_uri)
|
|
environ.update(get_run_tags_env_vars(recipe_root_path=self.recipe_root))
|
|
return environ
|
|
|
|
def get_artifacts(self):
|
|
return [
|
|
ModelArtifact("model", self.recipe_root, self.name, self.tracking_config.tracking_uri),
|
|
RunArtifact("run", self.recipe_root, self.name, self.tracking_config.tracking_uri),
|
|
HyperParametersArtifact("best_parameters", self.recipe_root, self.name),
|
|
DataframeArtifact(
|
|
"predicted_training_data",
|
|
self.recipe_root,
|
|
self.name,
|
|
TrainStep.PREDICTED_TRAINING_DATA_RELATIVE_PATH,
|
|
),
|
|
]
|
|
|
|
def step_class(self):
|
|
return StepClass.TRAINING
|
|
|
|
def _tune_and_get_best_estimator_params(
|
|
self,
|
|
parent_run_id,
|
|
estimator_hardcoded_params,
|
|
estimator_fn,
|
|
X_train,
|
|
y_train,
|
|
validation_df,
|
|
):
|
|
tuning_params = self.step_config["tuning"]
|
|
try:
|
|
from hyperopt import SparkTrials, Trials, fmin, space_eval
|
|
except ModuleNotFoundError:
|
|
raise MlflowException(
|
|
"Hyperopt not installed and is required if tuning is enabled",
|
|
error_code=BAD_REQUEST,
|
|
)
|
|
|
|
# wrap training in objective fn
|
|
def objective(X_train, y_train, validation_df, hyperparameter_args, on_worker=False):
|
|
run_name = self.tracking_config.run_name
|
|
if on_worker:
|
|
client = MlflowClient()
|
|
parent_tags = client.get_run(parent_run_id).data.tags
|
|
child_run = client.create_run(
|
|
_get_experiment_id(),
|
|
tags={**parent_tags, "mlflow.parentRunId": parent_run_id},
|
|
run_name=run_name,
|
|
)
|
|
run_args = {"run_id": child_run.info.run_id}
|
|
else:
|
|
run_args = {"run_name": run_name, "nested": True}
|
|
with mlflow.start_run(**run_args) as tuning_run:
|
|
estimator_args = dict(estimator_hardcoded_params, **hyperparameter_args)
|
|
estimator = estimator_fn(estimator_args)
|
|
|
|
sample_fraction = self.step_config["sample_fraction"]
|
|
|
|
# if sparktrials, then read from broadcast
|
|
if tuning_params["parallelism"] > 1:
|
|
X_train = X_train.value
|
|
y_train = y_train.value
|
|
validation_df = validation_df.value
|
|
|
|
X_train_sampled = X_train.sample(frac=sample_fraction, random_state=42)
|
|
y_train_sampled = y_train.sample(frac=sample_fraction, random_state=42)
|
|
|
|
fitted_estimator, _ = self._fitted_estimator(
|
|
estimator, X_train_sampled, y_train_sampled
|
|
)
|
|
|
|
logged_estimator = self._log_estimator_to_mlflow(
|
|
fitted_estimator, X_train_sampled, on_worker=on_worker
|
|
)
|
|
|
|
eval_result = mlflow.evaluate(
|
|
model=logged_estimator.model_uri,
|
|
data=validation_df,
|
|
targets=self.target_col,
|
|
model_type=_get_model_type_from_template(self.recipe),
|
|
evaluators="default",
|
|
extra_metrics=_load_custom_metrics(
|
|
self.recipe_root,
|
|
self.evaluation_metrics.values(),
|
|
),
|
|
evaluator_config={
|
|
"log_model_explainability": False,
|
|
"pos_label": self.positive_class,
|
|
},
|
|
)
|
|
autologged_params = mlflow.get_run(run_id=tuning_run.info.run_id).data.params
|
|
manual_log_params = {}
|
|
for param_name, param_value in estimator_args.items():
|
|
if param_name in autologged_params:
|
|
if not TrainStep.is_tuning_param_equal(
|
|
param_value, autologged_params[param_name]
|
|
):
|
|
_logger.warning(
|
|
f"Failed to log search space parameter due to conflict. "
|
|
f"Attempted to log search space parameter {param_name} as "
|
|
f"{param_value}, but {param_name} is already logged as "
|
|
f"{autologged_params[param_name]} during training. "
|
|
f"Are you passing `estimator_params` properly to the "
|
|
f" estimator?"
|
|
)
|
|
else:
|
|
manual_log_params[param_name] = param_value
|
|
|
|
if len(manual_log_params) > 0:
|
|
mlflow.log_params(manual_log_params)
|
|
|
|
# return +/- metric
|
|
transformed_metrics = transform_multiclass_metrics_dict(
|
|
eval_result.metrics, self.extended_task
|
|
)
|
|
sign = -1 if self.evaluation_metrics_greater_is_better[self.primary_metric] else 1
|
|
return sign * transformed_metrics[self.primary_metric]
|
|
|
|
search_space = TrainStep.construct_search_space_from_yaml(tuning_params["parameters"])
|
|
algo_type, algo_name = tuning_params["algorithm"].rsplit(".", 1)
|
|
tuning_algo = getattr(importlib.import_module(algo_type, "hyperopt"), algo_name)
|
|
max_trials = tuning_params["max_trials"]
|
|
parallelism = tuning_params["parallelism"]
|
|
|
|
if parallelism > 1:
|
|
from pyspark.sql import SparkSession
|
|
|
|
spark_session = SparkSession.builder.config(
|
|
"spark.databricks.mlflow.trackHyperopt.enabled", "false"
|
|
).getOrCreate()
|
|
sc = spark_session.sparkContext
|
|
|
|
X_train = sc.broadcast(X_train)
|
|
y_train = sc.broadcast(y_train)
|
|
validation_df = sc.broadcast(validation_df)
|
|
|
|
hp_trials = SparkTrials(parallelism, spark_session=spark_session)
|
|
on_worker = True
|
|
else:
|
|
hp_trials = Trials()
|
|
on_worker = False
|
|
|
|
fmin_kwargs = {
|
|
"fn": lambda params: objective(
|
|
X_train, y_train, validation_df, params, on_worker=on_worker
|
|
),
|
|
"space": search_space,
|
|
"algo": tuning_algo,
|
|
"max_evals": max_trials,
|
|
"trials": hp_trials,
|
|
}
|
|
if "early_stop_fn" in tuning_params:
|
|
try:
|
|
train_module_name, early_stop_fn_name = tuning_params["early_stop_fn"].rsplit(
|
|
".", 1
|
|
)
|
|
except ValueError:
|
|
early_stop_fn_name = tuning_params["early_stop_fn"]
|
|
train_module_name = _USER_DEFINED_TRAIN_STEP_MODULE
|
|
early_stop_fn = getattr(importlib.import_module(train_module_name), early_stop_fn_name)
|
|
fmin_kwargs["early_stop_fn"] = early_stop_fn
|
|
best_hp_params = fmin(**fmin_kwargs)
|
|
best_hp_params = space_eval(search_space, best_hp_params)
|
|
best_hp_estimator_loss = hp_trials.best_trial["result"]["loss"]
|
|
if len(estimator_hardcoded_params) > 1:
|
|
hardcoded_estimator_loss = objective(
|
|
X_train, y_train, validation_df, estimator_hardcoded_params
|
|
)
|
|
|
|
if best_hp_estimator_loss < hardcoded_estimator_loss:
|
|
best_hardcoded_params = {
|
|
param_name: param_value
|
|
for param_name, param_value in estimator_hardcoded_params.items()
|
|
if param_name not in best_hp_params
|
|
}
|
|
else:
|
|
best_hp_params = {}
|
|
best_hardcoded_params = estimator_hardcoded_params
|
|
else:
|
|
best_hardcoded_params = {}
|
|
return (best_hardcoded_params, best_hp_params)
|
|
|
|
def _log_estimator_to_mlflow(self, estimator, X_train_sampled, on_worker=False):
|
|
from mlflow.models import infer_signature
|
|
|
|
if hasattr(estimator, "best_score_") and (type(estimator.best_score_) in [int, float]):
|
|
mlflow.log_metric("best_cv_score", estimator.best_score_)
|
|
if hasattr(estimator, "best_params_"):
|
|
mlflow.log_params(estimator.best_params_)
|
|
|
|
if on_worker:
|
|
mlflow.log_params(estimator.get_params())
|
|
estimator_tags = {
|
|
"estimator_name": estimator.__class__.__name__,
|
|
"estimator_class": (
|
|
estimator.__class__.__module__ + "." + estimator.__class__.__name__
|
|
),
|
|
}
|
|
mlflow.set_tags(estimator_tags)
|
|
estimator_schema = infer_signature(
|
|
X_train_sampled, estimator.predict(X_train_sampled.copy())
|
|
)
|
|
return mlflow.sklearn.log_model(
|
|
estimator,
|
|
f"{self.name}/estimator",
|
|
signature=estimator_schema,
|
|
code_paths=self.code_paths,
|
|
)
|
|
|
|
def _write_best_parameters_outputs(
|
|
self,
|
|
output_directory,
|
|
best_hp_params=None,
|
|
best_hardcoded_params=None,
|
|
automl_params=None,
|
|
default_params=None,
|
|
):
|
|
if best_hp_params or best_hardcoded_params or automl_params or default_params:
|
|
best_parameters_path = os.path.join(output_directory, "best_parameters.yaml")
|
|
if os.path.exists(best_parameters_path):
|
|
os.remove(best_parameters_path)
|
|
with open(best_parameters_path, "a") as file:
|
|
self._write_one_param_output(automl_params or {}, file, "automl parameters")
|
|
self._write_one_param_output(best_hp_params or {}, file, "tuned hyperparameters")
|
|
self._write_one_param_output(
|
|
best_hardcoded_params or {}, file, "hardcoded parameters"
|
|
)
|
|
self._write_one_param_output(default_params or {}, file, "default parameters")
|
|
mlflow.log_artifact(best_parameters_path, artifact_path="train")
|
|
|
|
def _write_one_param_output(self, params, file, caption):
|
|
if params:
|
|
file.write(f"# {caption} \n")
|
|
self._safe_dump_with_numeric_values(params, file, default_flow_style=False)
|
|
file.write("\n")
|
|
|
|
def _safe_dump_with_numeric_values(self, data, file, **kwargs):
|
|
import numpy as np
|
|
|
|
processed_data = {}
|
|
for key, value in data.items():
|
|
if isinstance(value, np.floating):
|
|
processed_data[key] = float(value)
|
|
elif isinstance(value, np.integer):
|
|
processed_data[key] = int(value)
|
|
elif isinstance(value, dict):
|
|
processed_data[key] = str(value)
|
|
else:
|
|
processed_data[key] = str(value)
|
|
|
|
if len(processed_data) > 0:
|
|
yaml.safe_dump(processed_data, file, **kwargs)
|
|
|
|
def _rebalance_classes(self, train_df):
|
|
import pandas as pd
|
|
|
|
resampling_minority_percentage = self.step_config.get(
|
|
"resampling_minority_percentage", _REBALANCING_DEFAULT_RATIO
|
|
)
|
|
|
|
df_positive_class = train_df[train_df[self.target_col] == self.positive_class]
|
|
df_negative_class = train_df[train_df[self.target_col] != self.positive_class]
|
|
|
|
if len(df_positive_class) > len(df_negative_class):
|
|
df_minority_class, df_majority_class = df_negative_class, df_positive_class
|
|
else:
|
|
df_minority_class, df_majority_class = df_positive_class, df_negative_class
|
|
|
|
original_minority_percentage = len(df_minority_class) / len(train_df)
|
|
if original_minority_percentage >= resampling_minority_percentage:
|
|
_logger.info(
|
|
f"Class imbalance of {original_minority_percentage:.2f} "
|
|
f"is better than {resampling_minority_percentage}, no need to rebalance"
|
|
)
|
|
return train_df
|
|
|
|
_logger.info(
|
|
f"Detected class imbalance: minority class percentage is "
|
|
f"{original_minority_percentage:.2f}"
|
|
)
|
|
|
|
majority_class_target = int(
|
|
len(df_minority_class)
|
|
* (1 - resampling_minority_percentage)
|
|
/ resampling_minority_percentage
|
|
)
|
|
df_majority_downsampled = df_majority_class.sample(majority_class_target)
|
|
_logger.info(
|
|
f"After downsampling: minority class percentage is {resampling_minority_percentage:.2f}"
|
|
)
|
|
return pd.concat([df_minority_class, df_majority_downsampled], axis=0).sample(frac=1)
|