MLflow Utility Functions

dagshub.mlflow.patch_mlflow(include: List[str] | None = None, exclude: List[str] | None = None, raise_on_log_model=True, patch_top_level=True, patch_mlflow_client=True)

Patch MLflow functions, making some of them stop raising exceptions, instead logging them to console. This is useful for long runs that have MLflow failing occasionally, so you don’t have to restart the run if logging failed.

By default, all top-level mlflow.log_... and MlflowClient.log_... functions are patched.

Parameters:

  • include – List of full names of functions you want to patch additionally to the default ones.

  • exclude –

    List of functions you DON’T want to patch. Use this if you need to make sure, for example, all log_artifact functions raise an exception on failure:

    patch_mlflow(exclude=["mlflow.log_artifact", "mlflow.log_artifacts"])

  • raise_on_log_model –

    If True, patched log calls still raise an exception if they are called from a log_model function.

    This ensures that even if you want to ignore failed calls to log_artifact or log_figure etc., failed calls to log_model in particular will still raise errors. This should work for all supported ML frameworks, as long as the function is named log_model.

  • patch_top_level – Whether to patch the top level mlflow functions.

  • patch_mlflow_client – Whether to patch the MlflowClient class.

dagshub.mlflow.unpatch_mlflow()

Removes the failsafe MLflow patching, returning all MLflow functions to their original state