MLFlowLogger#
- class composer.loggers.MLFlowLogger(experiment_name=None, run_name=None, tags=None, tracking_uri=None, rank_zero_only=True, flush_interval=10, model_registry_prefix='', model_registry_uri=None)[source]#
Log to MLflow.
- Parameters
experiment_name โ (str, optional): MLflow experiment name. If not set it will be use the MLflow environment variable or a default value
run_name โ (str, optional): MLflow run name. If not set it will be the same as the Trainer run name
tags โ (dict, optional): MLflow tags to log with the run
tracking_uri (str | Path, optional) โ MLflow tracking uri, the URI to the remote or local endpoint where logs are stored (If none it is set to MLflow default)
rank_zero_only (bool, optional) โ Whether to log only on the rank-zero process (default:
True
).flush_interval (int) โ The amount of time, in seconds, that MLflow must wait between logging batches of metrics. Any metrics that are recorded by Composer during this interval are enqueued, and the queue is flushed when the interval elapses (default:
10
).model_registry_prefix (str, optional) โ The prefix to use when registering models. If registering to Unity Catalog, must be in the format
{catalog_name}.{schema_name}
. (default: empty string)model_registry_uri (str, optional) โ The URI of the model registry to use. To register models to Unity Catalog, set to
databricks-uc
. (default: None)
- log_model(flavor, **kwargs)[source]#
Log a model to MLflow.
- Parameters
flavor (str) โ The MLflow model flavor to use. Currently only
'transformers'
is supported.**kwargs โ Keyword arguments to pass to the MLflow model logging function.
- Raises
NotImplementedError โ If
flavor
is not'transformers'
.
- register_model(model_uri, name, await_registration_for=300, tags=None)[source]#
Register a model to model registry.
- Parameters
model_uri (str) โ The URI of the model to register.
name (str) โ The name of the model to register. Will be appended to
model_registry_prefix
.await_registration_for (Optional[int], optional) โ The number of seconds to wait for the model to be registered. Defaults to 300.
tags (Dict[str, Any], optional) โ A dictionary of tags to add to the model. Defaults to None.
registry_uri (str, optional) โ The URI of the model registry. Defaults to โdatabricks-ucโ which will register to the Databricks Unity Catalog.
- Returns
ModelVersion โ The registered model.
- save_model(flavor, **kwargs)[source]#
Save a model to MLflow.
- Parameters
flavor (str) โ The MLflow model flavor to use. Currently only
'transformers'
is supported.**kwargs โ Keyword arguments to pass to the MLflow model saving function.
- Raises
NotImplementedError โ If
flavor
is not'transformers'
.