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, synchronous=False, log_system_metrics=True, ignore_metrics=None, ignore_hyperparameters=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: โ€˜โ€™)

  • 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)

  • synchronous (bool, optional) โ€“ Whether to log synchronously. If True, Mlflow will log synchronously to the MLflow backend. If False, Mlflow will log asynchronously. (default: False)

  • log_system_metrics (bool, optional) โ€“ Whether to log system metrics. If True, Mlflow will log system metrics (CPU/GPU/memory/network usage) during training. (default: True)

  • ignore_metrics (List[str], optional) โ€“ A list of glob patterns for metrics to ignore when logging. (default: None)

  • ignore_hyperparameters (List[str], optional) โ€“ A list of glob patterns for hyperparameters to ignore when logging. (default: None)

log_model(flavor, **kwargs)[source]#

Log a model to MLflow.

Parameters

  • flavor (Literal['transformers']) โ€“ 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 (int, optional) โ€“ The number of seconds to wait for the model to be registered. Defaults to 300.

  • tags (Optional[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 None which will register to the Databricks Unity Catalog.

Returns

ModelVersion โ€“ The registered model.

register_model_with_run_id(model_uri, name, await_creation_for=300, tags=None)[source]#

Similar to register_model, but uses a different MLflow API to allow passing in the run id.

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_creation_for (int, optional) โ€“ The number of seconds to wait for the model to be registered. Defaults to 300.

  • tags (Optional[Dict[str, Any]], optional) โ€“ A dictionary of tags to add to the model. Defaults to None.

save_model(flavor, **kwargs)[source]#

Save a model to MLflow.

Note: The 'peft' flavor is experimental and the API is subject to change without warning.

Parameters

  • flavor (Literal['transformers', 'peft']) โ€“ The MLflow model flavor to use. Currently only 'transformers' and 'peft' are supported.

  • **kwargs โ€“ Keyword arguments to pass to the MLflow model saving function.

Raises

NotImplementedError โ€“ If flavor is not 'transformers' or 'peft'.