import mlflow
[docs]class MLflowConfig:
"""
Logs each fit of hyperparameters in a running instance of mlflow: https://mlflow.org/
"""
def __init__(
self,
tracking_uri,
experiment,
run_name,
save_models=False,
registry_uri=None,
tags=None,
):
"""
Parameters
----------
tracking_uri: str
Address of local or remote tracking server.
experiment: str
Case sensitive name of an experiment to be activated.
run_name: str
Name of new run (stored as a mlflow.runName tag).
save_models: bool, default=False
If ``True``, it will log the estimator into mlflow artifacts
registry_uri: str, default=None
Address of local or remote model registry server.
tags: dict, default=None
Dictionary of tag_name: String -> value.
"""
self.client = mlflow.tracking.MlflowClient()
self.tracking_uri = tracking_uri
self.experiment = experiment
self.run_name = run_name
self.save_models = save_models
self.tags = tags
self.registry_uri = registry_uri
mlflow.set_registry_uri(self.registry_uri)
mlflow.set_tracking_uri(self.tracking_uri)
mlflow.set_experiment(self.experiment)
self.experiment_id = mlflow.get_experiment_by_name(
self.experiment
).experiment_id
if self.tags is not None:
mlflow.set_tags(self.tags)
[docs] def create_run(self, parameters, score, estimator):
"""
Parameters
----------
parameters: dict
A dictionary with the keys as the hyperparameter name and the value as the current value setting
score:
The cross-validation score achieved by the current parameters
estimator: estimator object
The current sklearn estimator that is being fitted
"""
with mlflow.start_run(
experiment_id=self.experiment_id, nested=True, run_name=self.run_name
):
for parameter, value in parameters.items():
mlflow.log_param(key=parameter, value=value)
mlflow.log_metric(key="score", value=score)
if self.save_models:
mlflow.sklearn.log_model(estimator, "model")