wehub-resource-sync
347 lines
12 KiB
Python
347 lines
12 KiB
Python
import logging
|
|
import os
|
|
from copy import deepcopy
|
|
from typing import TYPE_CHECKING, Dict, Optional
|
|
|
|
from packaging import version
|
|
|
|
from ray._private.dict import flatten_dict
|
|
|
|
if TYPE_CHECKING:
|
|
from mlflow.entities import Run
|
|
from mlflow.tracking import MlflowClient
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class _MLflowLoggerUtil:
|
|
"""Util class for setting up and logging to MLflow.
|
|
|
|
Use this util for any library that needs MLflow logging/tracking logic
|
|
such as Ray Tune or Ray Train.
|
|
"""
|
|
|
|
def __init__(self):
|
|
import mlflow
|
|
|
|
self._mlflow = mlflow
|
|
self.experiment_id = None
|
|
|
|
def __deepcopy__(self, memo=None):
|
|
# mlflow is a module, and thus cannot be copied
|
|
_mlflow = self._mlflow
|
|
self.__dict__.pop("_mlflow")
|
|
dict_copy = deepcopy(self.__dict__, memo)
|
|
copied_object = _MLflowLoggerUtil()
|
|
copied_object.__dict__.update(dict_copy)
|
|
self._mlflow = _mlflow
|
|
copied_object._mlflow = _mlflow
|
|
return copied_object
|
|
|
|
def setup_mlflow(
|
|
self,
|
|
tracking_uri: Optional[str] = None,
|
|
registry_uri: Optional[str] = None,
|
|
experiment_id: Optional[str] = None,
|
|
experiment_name: Optional[str] = None,
|
|
tracking_token: Optional[str] = None,
|
|
artifact_location: Optional[str] = None,
|
|
create_experiment_if_not_exists: bool = True,
|
|
) -> None:
|
|
"""
|
|
Sets up MLflow.
|
|
|
|
Sets the Mlflow tracking uri & token, and registry URI. Also sets
|
|
the MLflow experiment that the logger should use, and possibly
|
|
creates new experiment if it does not exist.
|
|
|
|
Args:
|
|
tracking_uri: The tracking URI for the MLflow tracking
|
|
server.
|
|
registry_uri: The registry URI for the MLflow model registry.
|
|
experiment_id: The id of an already existing MLflow
|
|
experiment to use for logging. If None is passed in
|
|
here and the MFLOW_EXPERIMENT_ID is not set, or the
|
|
experiment with this id does not exist,
|
|
``experiment_name`` will be used instead. This argument takes
|
|
precedence over ``experiment_name`` if both are passed in.
|
|
experiment_name: The experiment name to use for logging.
|
|
If None is passed in here, the MLFLOW_EXPERIMENT_NAME environment
|
|
variable is used to determine the experiment name.
|
|
If the experiment with the name already exists with MLflow,
|
|
it will be reused. If not, a new experiment will be created
|
|
with the provided name if
|
|
``create_experiment_if_not_exists`` is set to True.
|
|
tracking_token: Tracking token used to authenticate with MLflow.
|
|
artifact_location: The location to store run artifacts.
|
|
If not provided, MLFlow picks an appropriate default.
|
|
Ignored if experiment already exists.
|
|
create_experiment_if_not_exists: Whether to create an
|
|
experiment with the provided name if it does not already
|
|
exist. Defaults to True.
|
|
|
|
Raises:
|
|
ValueError: ``experiment_id`` and ``experiment_name`` are both ``None``.
|
|
"""
|
|
if tracking_token:
|
|
os.environ["MLFLOW_TRACKING_TOKEN"] = tracking_token
|
|
|
|
self._mlflow.set_tracking_uri(tracking_uri)
|
|
self._mlflow.set_registry_uri(registry_uri)
|
|
|
|
# First check experiment_id.
|
|
experiment_id = (
|
|
experiment_id
|
|
if experiment_id is not None
|
|
else os.environ.get("MLFLOW_EXPERIMENT_ID")
|
|
)
|
|
if experiment_id is not None:
|
|
from mlflow.exceptions import MlflowException
|
|
|
|
try:
|
|
self._mlflow.get_experiment(experiment_id=experiment_id)
|
|
logger.debug(
|
|
f"Experiment with provided id {experiment_id} "
|
|
"exists. Setting that as the experiment."
|
|
)
|
|
self.experiment_id = experiment_id
|
|
return
|
|
except MlflowException:
|
|
pass
|
|
|
|
# Then check experiment_name.
|
|
experiment_name = (
|
|
experiment_name
|
|
if experiment_name is not None
|
|
else os.environ.get("MLFLOW_EXPERIMENT_NAME")
|
|
)
|
|
if experiment_name is not None and self._mlflow.get_experiment_by_name(
|
|
name=experiment_name
|
|
):
|
|
logger.debug(
|
|
f"Experiment with provided name {experiment_name} "
|
|
"exists. Setting that as the experiment."
|
|
)
|
|
self.experiment_id = self._mlflow.get_experiment_by_name(
|
|
experiment_name
|
|
).experiment_id
|
|
return
|
|
|
|
# An experiment with the provided id or name does not exist.
|
|
# Create a new experiment if applicable.
|
|
if experiment_name and create_experiment_if_not_exists:
|
|
logger.debug(
|
|
"Existing experiment not found. Creating new "
|
|
f"experiment with name: {experiment_name}"
|
|
)
|
|
self.experiment_id = self._mlflow.create_experiment(
|
|
name=experiment_name, artifact_location=artifact_location
|
|
)
|
|
return
|
|
|
|
if create_experiment_if_not_exists:
|
|
raise ValueError(
|
|
f"Experiment with the provided experiment_id: "
|
|
f"{experiment_id} does not exist and no "
|
|
f"experiment_name provided. At least one of "
|
|
f"these has to be provided."
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Experiment with the provided experiment_id: "
|
|
f"{experiment_id} or experiment_name: "
|
|
f"{experiment_name} does not exist. Please "
|
|
f"create an MLflow experiment and provide "
|
|
f"either its id or name."
|
|
)
|
|
|
|
def _parse_dict(self, dict_to_log: Dict) -> Dict:
|
|
"""Parses provided dict to convert all values to float.
|
|
|
|
MLflow can only log metrics that are floats. This does not apply to
|
|
logging parameters or artifacts.
|
|
|
|
Args:
|
|
dict_to_log: The dictionary containing the metrics to log.
|
|
|
|
Returns:
|
|
A dictionary containing the metrics to log with all values being
|
|
converted to floats, or skipped if not able to be converted.
|
|
"""
|
|
new_dict = {}
|
|
for key, value in dict_to_log.items():
|
|
try:
|
|
value = float(value)
|
|
new_dict[key] = value
|
|
except (ValueError, TypeError):
|
|
logger.debug(
|
|
"Cannot log key {} with value {} since the "
|
|
"value cannot be converted to float.".format(key, value)
|
|
)
|
|
continue
|
|
|
|
return new_dict
|
|
|
|
def start_run(
|
|
self,
|
|
run_name: Optional[str] = None,
|
|
tags: Optional[Dict] = None,
|
|
set_active: bool = False,
|
|
) -> "Run":
|
|
"""Starts a new run and possibly sets it as the active run.
|
|
|
|
Args:
|
|
run_name: Name of the new MLflow run to create.
|
|
tags: Tags to set for the new run.
|
|
set_active: Whether to set the new run as the active run.
|
|
If an active run already exists, then that run is returned.
|
|
|
|
Returns:
|
|
The newly created MLflow run.
|
|
"""
|
|
import mlflow
|
|
from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME
|
|
|
|
if tags is None:
|
|
tags = {}
|
|
|
|
if set_active:
|
|
return self._start_active_run(run_name=run_name, tags=tags)
|
|
|
|
client = self._get_client()
|
|
# If `mlflow==1.30.0` and we don't use `run_name`, then MLflow might error. For
|
|
# more information, see #29749.
|
|
if version.parse(mlflow.__version__) >= version.parse("1.30.0"):
|
|
run = client.create_run(
|
|
run_name=run_name, experiment_id=self.experiment_id, tags=tags
|
|
)
|
|
else:
|
|
tags[MLFLOW_RUN_NAME] = run_name
|
|
run = client.create_run(experiment_id=self.experiment_id, tags=tags)
|
|
|
|
return run
|
|
|
|
def _start_active_run(
|
|
self, run_name: Optional[str] = None, tags: Optional[Dict] = None
|
|
) -> "Run":
|
|
"""Starts a run and sets it as the active run if one does not exist.
|
|
|
|
If an active run already exists, then returns it.
|
|
"""
|
|
active_run = self._mlflow.active_run()
|
|
if active_run:
|
|
return active_run
|
|
|
|
return self._mlflow.start_run(
|
|
run_name=run_name, experiment_id=self.experiment_id, tags=tags
|
|
)
|
|
|
|
def _run_exists(self, run_id: str) -> bool:
|
|
"""Check if run with the provided id exists."""
|
|
from mlflow.exceptions import MlflowException
|
|
|
|
try:
|
|
self._mlflow.get_run(run_id=run_id)
|
|
return True
|
|
except MlflowException:
|
|
return False
|
|
|
|
def _get_client(self) -> "MlflowClient":
|
|
"""Returns an ml.tracking.MlflowClient instance to use for logging."""
|
|
tracking_uri = self._mlflow.get_tracking_uri()
|
|
registry_uri = self._mlflow.get_registry_uri()
|
|
|
|
from mlflow.tracking import MlflowClient
|
|
|
|
return MlflowClient(tracking_uri=tracking_uri, registry_uri=registry_uri)
|
|
|
|
def log_params(self, params_to_log: Dict, run_id: Optional[str] = None):
|
|
"""Logs the provided parameters to the run specified by run_id.
|
|
|
|
If no ``run_id`` is passed in, then logs to the current active run.
|
|
If there is not active run, then creates a new run and sets it as
|
|
the active run.
|
|
|
|
Args:
|
|
params_to_log: Dictionary of parameters to log.
|
|
run_id: The ID of the run to log to.
|
|
"""
|
|
params_to_log = flatten_dict(params_to_log)
|
|
|
|
if run_id and self._run_exists(run_id):
|
|
client = self._get_client()
|
|
for key, value in params_to_log.items():
|
|
client.log_param(run_id=run_id, key=key, value=value)
|
|
|
|
else:
|
|
for key, value in params_to_log.items():
|
|
self._mlflow.log_param(key=key, value=value)
|
|
|
|
def log_metrics(
|
|
self, step: int, metrics_to_log: Dict, run_id: Optional[str] = None
|
|
):
|
|
"""Logs the provided metrics to the run specified by run_id.
|
|
|
|
|
|
If no ``run_id`` is passed in, then logs to the current active run.
|
|
If there is not active run, then creates a new run and sets it as
|
|
the active run.
|
|
|
|
Args:
|
|
step: Step at which the metrics are logged.
|
|
metrics_to_log: Dictionary of metrics to log.
|
|
run_id: The ID of the run to log to.
|
|
"""
|
|
metrics_to_log = flatten_dict(metrics_to_log)
|
|
metrics_to_log = self._parse_dict(metrics_to_log)
|
|
|
|
if run_id and self._run_exists(run_id):
|
|
client = self._get_client()
|
|
for key, value in metrics_to_log.items():
|
|
client.log_metric(run_id=run_id, key=key, value=value, step=step)
|
|
|
|
else:
|
|
for key, value in metrics_to_log.items():
|
|
self._mlflow.log_metric(key=key, value=value, step=step)
|
|
|
|
def save_artifacts(self, dir: str, run_id: Optional[str] = None):
|
|
"""Saves directory as artifact to the run specified by run_id.
|
|
|
|
If no ``run_id`` is passed in, then saves to the current active run.
|
|
If there is not active run, then creates a new run and sets it as
|
|
the active run.
|
|
|
|
Args:
|
|
dir: Path to directory containing the files to save.
|
|
run_id: The ID of the run to log to.
|
|
"""
|
|
if run_id and self._run_exists(run_id):
|
|
client = self._get_client()
|
|
client.log_artifacts(run_id=run_id, local_dir=dir)
|
|
else:
|
|
self._mlflow.log_artifacts(local_dir=dir)
|
|
|
|
def end_run(self, status: Optional[str] = None, run_id: Optional[str] = None):
|
|
"""Terminates the run specified by run_id.
|
|
|
|
If no ``run_id`` is passed in, then terminates the
|
|
active run if one exists.
|
|
|
|
Args:
|
|
status: The status to set when terminating the run.
|
|
run_id: The ID of the run to terminate.
|
|
|
|
"""
|
|
if (
|
|
run_id
|
|
and self._run_exists(run_id)
|
|
and not (
|
|
self._mlflow.active_run()
|
|
and self._mlflow.active_run().info.run_id == run_id
|
|
)
|
|
):
|
|
client = self._get_client()
|
|
client.set_terminated(run_id=run_id, status=status)
|
|
else:
|
|
self._mlflow.end_run(status=status)
|