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'.