API reference

mlflow_autogluon.autolog

Autologging support for AutoGluon.

Patches TabularPredictor.fit so that a single call to :func:mlflow_autogluon.autolog records parameters, leaderboard metrics, artifacts, and (optionally) the fitted predictor itself to MLflow, mirroring the behavior of MLflow's built-in autologging integrations.

Because this is a community flavor, mlflow.autolog() does not enable it; call mlflow_autogluon.autolog() explicitly before fitting.

autolog(log_models=True, log_model_signatures=True, log_input_examples=False, log_datasets=True, log_leaderboard=True, log_fit_summary=False, registered_model_name=None, extra_tags=None, disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False)

Enable automatic logging for AutoGluon TabularPredictor.fit calls.

After calling this function, every fit call logs to the active MLflow run (a run is created automatically when none is active):

• Parameters: predictor configuration (label, problem type, eval metric) and fit arguments (presets, time_limit, hyperparameters, bagging and stacking settings, and so on).
• Metrics: validation score and fit time per trained model from the leaderboard, plus the best model's validation score and total fit time.
• Tags: AutoGluon version and best model name.
• Artifacts: the leaderboard as CSV, optionally the fit summary as JSON, and the fitted predictor logged with the autogluon flavor.

Parameters:

Name	Type	Description	Default
log_models	bool	If True, log the fitted predictor as an MLflow model.	True
log_model_signatures	bool	If True, infer a model signature from a sample of the training data and the predictor's output and attach it to logged models.	True
log_input_examples	bool	If True, save a small sample of the training features as the logged model's input example.	False
log_datasets	bool	If True, attach the training data to the run as an MLflow dataset input.	True
log_leaderboard	bool	If True, log the leaderboard as a CSV artifact and per-model metrics.	True
log_fit_summary	bool	If True, log predictor.fit_summary() as a JSON artifact.	False
registered_model_name	str | None	If given, logged models are also registered under this name in the model registry.	None
extra_tags	dict[str, Any] | None	Dict of extra tags to set on autologged runs.	None
disable	bool	If True, disable the integration.	False
exclusive	bool	If True, autologged content is not logged to user-created fluent runs.	False
disable_for_unsupported_versions	bool	If True, disable autologging for untested AutoGluon versions.	False
silent	bool	If True, suppress all MLflow event logs and warnings from autologging.	False

Source code in mlflow_autogluon/autolog.py

@autologging_integration(FLAVOR_NAME)
def autolog(
    log_models: bool = True,
    log_model_signatures: bool = True,
    log_input_examples: bool = False,
    log_datasets: bool = True,
    log_leaderboard: bool = True,
    log_fit_summary: bool = False,
    registered_model_name: str | None = None,
    extra_tags: dict[str, Any] | None = None,
    disable: bool = False,
    exclusive: bool = False,
    disable_for_unsupported_versions: bool = False,
    silent: bool = False,
) -> None:  # pylint: disable=unused-argument
    """Enable automatic logging for AutoGluon ``TabularPredictor.fit`` calls.

    After calling this function, every ``fit`` call logs to the active MLflow
    run (a run is created automatically when none is active):

    - Parameters: predictor configuration (label, problem type, eval metric)
      and fit arguments (presets, time_limit, hyperparameters, bagging and
      stacking settings, and so on).
    - Metrics: validation score and fit time per trained model from the
      leaderboard, plus the best model's validation score and total fit time.
    - Tags: AutoGluon version and best model name.
    - Artifacts: the leaderboard as CSV, optionally the fit summary as JSON,
      and the fitted predictor logged with the ``autogluon`` flavor.

    Args:
        log_models: If ``True``, log the fitted predictor as an MLflow model.
        log_model_signatures: If ``True``, infer a model signature from a
            sample of the training data and the predictor's output and attach
            it to logged models.
        log_input_examples: If ``True``, save a small sample of the training
            features as the logged model's input example.
        log_datasets: If ``True``, attach the training data to the run as an
            MLflow dataset input.
        log_leaderboard: If ``True``, log the leaderboard as a CSV artifact
            and per-model metrics.
        log_fit_summary: If ``True``, log ``predictor.fit_summary()`` as a
            JSON artifact.
        registered_model_name: If given, logged models are also registered
            under this name in the model registry.
        extra_tags: Dict of extra tags to set on autologged runs.
        disable: If ``True``, disable the integration.
        exclusive: If ``True``, autologged content is not logged to
            user-created fluent runs.
        disable_for_unsupported_versions: If ``True``, disable autologging for
            untested AutoGluon versions.
        silent: If ``True``, suppress all MLflow event logs and warnings from
            autologging.
    """
    patched_any = False
    for module_name, class_name in _PREDICTOR_REGISTRY.values():
        try:
            module = importlib.import_module(module_name)
        except ImportError:
            continue
        predictor_class = getattr(module, class_name)
        safe_patch(FLAVOR_NAME, predictor_class, "fit", _patched_fit, manage_run=True)
        patched_any = True
    if not patched_any:
        raise ImportError(
            "No AutoGluon predictor packages found. Install at least one of: "
            + ", ".join(m for m, _ in _PREDICTOR_REGISTRY.values())
        )

mlflow_autogluon.flavor.save_model(ag_model, path, conda_env=None, code_paths=None, mlflow_model=None, signature=None, input_example=None, pip_requirements=None, extra_pip_requirements=None, metadata=None)

Save a fitted AutoGluon predictor to a local path in MLflow model format.

Parameters:

Name	Type	Description	Default
ag_model	Any	Fitted TabularPredictor, TimeSeriesPredictor, or MultiModalPredictor instance.	required
path	str	Local filesystem destination for the MLflow model.	required
conda_env	dict[str, Any] | str | None	Conda environment dict or path to a conda YAML file.	None
code_paths	list[str] | None	Local code paths to package with the model.	None
mlflow_model	Model | None	Existing :class:mlflow.models.Model to add flavors to.	None
signature	ModelSignature | None	:class:mlflow.models.ModelSignature describing input/output.	None
input_example	Any | None	Example model input, saved alongside the model.	None
pip_requirements	list[str] | str | None	Override for the default pip requirements.	None
extra_pip_requirements	list[str] | str | None	Additional pip requirements.	None
metadata	dict[str, Any] | None	Custom metadata dict stored in the MLmodel file.	None

Source code in mlflow_autogluon/flavor.py

def save_model(
    ag_model: Any,
    path: str,
    conda_env: dict[str, Any] | str | None = None,
    code_paths: list[str] | None = None,
    mlflow_model: Model | None = None,
    signature: ModelSignature | None = None,
    input_example: Any | None = None,
    pip_requirements: list[str] | str | None = None,
    extra_pip_requirements: list[str] | str | None = None,
    metadata: dict[str, Any] | None = None,
) -> None:
    """Save a fitted AutoGluon predictor to a local path in MLflow model format.

    Args:
        ag_model: Fitted ``TabularPredictor``, ``TimeSeriesPredictor``, or
            ``MultiModalPredictor`` instance.
        path: Local filesystem destination for the MLflow model.
        conda_env: Conda environment dict or path to a conda YAML file.
        code_paths: Local code paths to package with the model.
        mlflow_model: Existing :class:`mlflow.models.Model` to add flavors to.
        signature: :class:`mlflow.models.ModelSignature` describing input/output.
        input_example: Example model input, saved alongside the model.
        pip_requirements: Override for the default pip requirements.
        extra_pip_requirements: Additional pip requirements.
        metadata: Custom metadata dict stored in the MLmodel file.
    """
    model_type = _validate_ag_model(ag_model)
    _validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)

    path = os.path.abspath(path)
    _validate_and_prepare_target_save_path(path)
    code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)

    if mlflow_model is None:
        mlflow_model = Model()
    mlflow_model.signature = _with_default_params_schema(signature)
    if input_example is not None:
        _save_example(mlflow_model, input_example, path)
    if metadata is not None:
        mlflow_model.metadata = metadata

    model_data_path = os.path.join(path, _MODEL_DATA_SUBPATH)
    _persist_predictor(ag_model, model_type, model_data_path)

    pyfunc.add_to_model(
        mlflow_model,
        loader_module="mlflow_autogluon.flavor",
        conda_env=_CONDA_ENV_FILE_NAME,
        python_env=_PYTHON_ENV_FILE_NAME,
        code=code_dir_subpath,
    )
    mlflow_model.add_flavor(
        FLAVOR_NAME,
        autogluon_version=_get_autogluon_version(model_type),
        model_type=model_type,
        data=_MODEL_DATA_SUBPATH,
        code=code_dir_subpath,
    )
    mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))

    if conda_env is None:
        default_reqs = (
            get_default_pip_requirements(model_type) if pip_requirements is None else None
        )
        conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
            default_reqs, pip_requirements, extra_pip_requirements
        )
    else:
        conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)

    with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
        yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
    if pip_constraints:
        write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
    write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
    _PythonEnv.current().to_yaml(os.path.join(path, _PYTHON_ENV_FILE_NAME))

mlflow_autogluon.flavor.log_model(ag_model, artifact_path=None, conda_env=None, code_paths=None, registered_model_name=None, signature=None, input_example=None, await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS, pip_requirements=None, extra_pip_requirements=None, metadata=None, name=None, **kwargs)

Log a fitted AutoGluon predictor as an MLflow artifact for the current run.

Parameters:

Name	Type	Description	Default
ag_model	Any	Fitted TabularPredictor, TimeSeriesPredictor, or MultiModalPredictor instance.	required
artifact_path	str | None	Run-relative artifact path (MLflow 2.x convention).	None
name	str | None	Model name (MLflow 3.x convention). Falls back to artifact_path on MLflow versions that do not support it.	None
registered_model_name	str | None	If given, register the model under this name.	None
kwargs	Any	Remaining arguments are forwarded to :func:save_model.	{}

Returns:

Name	Type	Description
A	ModelInfo	class:mlflow.models.model.ModelInfo describing the logged model.

Source code in mlflow_autogluon/flavor.py

def log_model(
    ag_model: Any,
    artifact_path: str | None = None,
    conda_env: dict[str, Any] | str | None = None,
    code_paths: list[str] | None = None,
    registered_model_name: str | None = None,
    signature: ModelSignature | None = None,
    input_example: Any | None = None,
    await_registration_for: int = DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
    pip_requirements: list[str] | str | None = None,
    extra_pip_requirements: list[str] | str | None = None,
    metadata: dict[str, Any] | None = None,
    name: str | None = None,
    **kwargs: Any,
) -> ModelInfo:
    """Log a fitted AutoGluon predictor as an MLflow artifact for the current run.

    Args:
        ag_model: Fitted ``TabularPredictor``, ``TimeSeriesPredictor``, or
            ``MultiModalPredictor`` instance.
        artifact_path: Run-relative artifact path (MLflow 2.x convention).
        name: Model name (MLflow 3.x convention). Falls back to ``artifact_path``
            on MLflow versions that do not support it.
        registered_model_name: If given, register the model under this name.
        kwargs: Remaining arguments are forwarded to :func:`save_model`.

    Returns:
        A :class:`mlflow.models.model.ModelInfo` describing the logged model.
    """
    log_params = inspect.signature(Model.log).parameters
    if name is not None and "name" in log_params:
        # MLflow 3.x: artifact_path is still a required positional, pass None
        # alongside name to opt into the new naming convention.
        log_kwargs = {"artifact_path": None, "name": name}
    else:
        log_kwargs = {"artifact_path": artifact_path or name or "model"}

    return Model.log(
        flavor=inspect.getmodule(save_model),
        ag_model=ag_model,
        conda_env=conda_env,
        code_paths=code_paths,
        registered_model_name=registered_model_name,
        signature=signature,
        input_example=input_example,
        await_registration_for=await_registration_for,
        pip_requirements=pip_requirements,
        extra_pip_requirements=extra_pip_requirements,
        metadata=metadata,
        **log_kwargs,
        **kwargs,
    )

mlflow_autogluon.flavor.load_model(model_uri, dst_path=None)

Load a native AutoGluon predictor from an MLflow model URI.

Parameters:

Name	Type	Description	Default
model_uri	str	URI of the MLflow model, e.g. runs:/<run_id>/model or models:/<name>/<version>.	required
dst_path	str | None	Optional local destination for downloaded artifacts.	None

Returns:

Type	Description
Any	The restored AutoGluon predictor (TabularPredictor,
Any	TimeSeriesPredictor, or MultiModalPredictor).

Source code in mlflow_autogluon/flavor.py

def load_model(model_uri: str, dst_path: str | None = None) -> Any:
    """Load a native AutoGluon predictor from an MLflow model URI.

    Args:
        model_uri: URI of the MLflow model, e.g. ``runs:/<run_id>/model`` or
            ``models:/<name>/<version>``.
        dst_path: Optional local destination for downloaded artifacts.

    Returns:
        The restored AutoGluon predictor (``TabularPredictor``,
        ``TimeSeriesPredictor``, or ``MultiModalPredictor``).
    """
    local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
    flavor_conf = _get_flavor_configuration(
        model_path=local_model_path, flavor_name=FLAVOR_NAME
    )
    _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
    ag_model_path = os.path.join(local_model_path, flavor_conf.get("data", _MODEL_DATA_SUBPATH))
    return _load_model_from_data_path(
        ag_model_path, flavor_conf.get("model_type", _MODEL_TYPE_TABULAR)
    )

mlflow_autogluon.flavor.get_default_pip_requirements(model_type=_MODEL_TYPE_TABULAR)

Return the default pip requirements for models produced by this flavor.

Parameters:

Name	Type	Description	Default
model_type	str	One of "tabular", "timeseries", or "multimodal".	_MODEL_TYPE_TABULAR

Source code in mlflow_autogluon/flavor.py

def get_default_pip_requirements(model_type: str = _MODEL_TYPE_TABULAR) -> list[str]:
    """Return the default pip requirements for models produced by this flavor.

    Args:
        model_type: One of ``"tabular"``, ``"timeseries"``, or ``"multimodal"``.
    """
    module_name, _ = _registry_entry(model_type)
    return [f"{module_name}=={_get_autogluon_version(model_type)}"]

mlflow_autogluon.flavor.get_default_conda_env(model_type=_MODEL_TYPE_TABULAR)

Return the default conda environment for models produced by this flavor.

Parameters:

Name	Type	Description	Default
model_type	str	One of "tabular", "timeseries", or "multimodal".	_MODEL_TYPE_TABULAR

Source code in mlflow_autogluon/flavor.py

def get_default_conda_env(model_type: str = _MODEL_TYPE_TABULAR) -> dict[str, Any]:
    """Return the default conda environment for models produced by this flavor.

    Args:
        model_type: One of ``"tabular"``, ``"timeseries"``, or ``"multimodal"``.
    """
    return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements(model_type))