Repository: sktime/sktime

Issue:
[BUG] Modify the `ModuleNotFoundError` message for `mlflow` to contain a more accurate installation method
I can understand the rationale for not including it in `all_extras`. However, I propose including `mlflow` in `all_extras` since there are use cases like #5324 where people might use it to simply store their models locally, among other things (thereby making `mlflow` a utility package).

I'm open to being convinced otherwise, but it seems `mlflow` should also be bundled with other soft dependencies.


Discussion:
The key argument imo is that:

* `mlflow` is only one of multiple deployment frameworks
* users can select to install it with the `mlflow` dep set already

I'm sure @yarnabrina has some (more informed) opinions though as he has invested most work in re-organizing the dependency sets.
I'm not sure if my opinion is more informed or more biased, but I do have one based on my quite recent personal experience.

`mlflow` has a bunch of dependencies itself, which caused a lot of conflicts for me with the other packages I was using. `all_extras` have plenty of packages itself, so I'm concerned that those along with `mlflow` dependencies may lead to some weird conflicts or solve to combination of older versions. Given these issues, I've switched to `mlflow-skinny` to reduce the number of issues, but still had to introduce few pinnings here and there to address the remaining ones.

So, my personal preference will be to leave it out of the extras, and if someone wants to use `mlflow`, then can install themselves separately.
Ah yes, I can actually echo that. When @benjaminbluhm added his early version of `mlflow` integration to `sktime`, he first added `mlflow` and some other stuff to the `all_extras` if I am not mistaken.

This led to strange interactions between version bounds, that's why we split it off as a separate dependency set.
In this case, I recommend adding a conditional logic inside `_check_soft_dependencies()`. Currently, it gives the following message on the absence of `mlflow`:
```
ModuleNotFoundError: No module named 'mlflow'. 'mlflow' is a soft dependency and not included in the base sktime installation. Please run: `pip install mlflow` to install the mlflow package. To install all soft dependencies, run: `pip install sktime[all_extras]`
```

Whereas a more accurate description would be:
```
ModuleNotFoundError: No module named `mlflow`. `mlflow` is an extra dependency and is not included in the base sktime installation. Please run `pip install mlflow` or `pip install sktime[mlflow]` to install the package.
```

Resolving patch:
diff --git a/sktime/utils/mlflow_sktime.py b/sktime/utils/mlflow_sktime.py
--- a/sktime/utils/mlflow_sktime.py
+++ b/sktime/utils/mlflow_sktime.py
@@ -52,9 +52,12 @@
 import sktime
 from sktime import utils
 from sktime.utils.multiindex import flatten_multiindex
-from sktime.utils.validation._dependencies import _check_soft_dependencies
+from sktime.utils.validation._dependencies import (
+    _check_mlflow_dependencies,
+    _check_soft_dependencies,
+)
 
-if _check_soft_dependencies("mlflow", severity="warning"):
+if _check_mlflow_dependencies(severity="warning"):
     from mlflow import pyfunc
 
 FLAVOR_NAME = "mlflow_sktime"
@@ -93,7 +96,7 @@ def get_default_pip_requirements(include_cloudpickle=False):
     Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment
     that, at a minimum, contains these requirements.
     """
-    _check_soft_dependencies("mlflow", severity="error")
+    _check_mlflow_dependencies(severity="error")
     from mlflow.utils.requirements_utils import _get_pinned_requirement
 
     pip_deps = [_get_pinned_requirement("sktime")]
@@ -111,7 +114,7 @@ def get_default_conda_env(include_cloudpickle=False):
     The default Conda environment for MLflow Models produced by calls to
     :func:`save_model()` and :func:`log_model()`
     """
-    _check_soft_dependencies("mlflow", severity="error")
+    _check_mlflow_dependencies(severity="error")
     from mlflow.utils.environment import _mlflow_conda_env
 
     return _mlflow_conda_env(
@@ -211,7 +214,7 @@ def save_model(
     >>> loaded_model = mlflow_sktime.load_model(model_uri=model_path)  # doctest: +SKIP
     >>> loaded_model.predict(fh=[1, 2, 3])  # doctest: +SKIP
     """  # noqa: E501
-    _check_soft_dependencies("mlflow", severity="error")
+    _check_mlflow_dependencies(severity="error")
     from mlflow.exceptions import MlflowException
     from mlflow.models import Model
     from mlflow.models.model import MLMODEL_FILE_NAME
@@ -420,7 +423,7 @@ def log_model(
     ...     sktime_model=forecaster,
     ...     artifact_path=artifact_path)  # doctest: +SKIP
     """  # noqa: E501
-    _check_soft_dependencies("mlflow", severity="error")
+    _check_mlflow_dependencies(severity="error")
     from mlflow.models import Model
 
     if await_registration_for is None:
@@ -493,7 +496,7 @@ def load_model(model_uri, dst_path=None):
     ...     path=model_path)
     >>> loaded_model = mlflow_sktime.load_model(model_uri=model_path)  # doctest: +SKIP
     """  # noqa: E501
-    _check_soft_dependencies("mlflow", severity="error")
+    _check_mlflow_dependencies(severity="error")
     from mlflow.tracking.artifact_utils import _download_artifact_from_uri
     from mlflow.utils.model_utils import (
         _add_code_from_conf_to_system_path,
@@ -519,7 +522,7 @@ def load_model(model_uri, dst_path=None):
 
 
 def _save_model(model, path, serialization_format):
-    _check_soft_dependencies("mlflow", severity="error")
+    _check_mlflow_dependencies(severity="error")
     from mlflow.exceptions import MlflowException
     from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
 
@@ -542,7 +545,7 @@ def _save_model(model, path, serialization_format):
 
 
 def _load_model(path, serialization_format):
-    _check_soft_dependencies("mlflow", severity="error")
+    _check_mlflow_dependencies(severity="error")
     from mlflow.exceptions import MlflowException
     from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
 
@@ -585,7 +588,7 @@ def _load_pyfunc(path):
     ----------
     .. [1] https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#mlflow.pyfunc.load_model
     """  # noqa: E501
-    _check_soft_dependencies("mlflow", severity="error")
+    _check_mlflow_dependencies(severity="error")
     from mlflow.exceptions import MlflowException
     from mlflow.utils.model_utils import _get_flavor_configuration
 
@@ -621,7 +624,7 @@ def _load_pyfunc(path):
 
 class _SktimeModelWrapper:
     def __init__(self, sktime_model):
-        _check_soft_dependencies("mlflow", severity="error")
+        _check_mlflow_dependencies(severity="error")
         self.sktime_model = sktime_model
 
     def predict(self, X):
diff --git a/sktime/utils/validation/_dependencies.py b/sktime/utils/validation/_dependencies.py
--- a/sktime/utils/validation/_dependencies.py
+++ b/sktime/utils/validation/_dependencies.py
@@ -40,7 +40,7 @@ def _check_soft_dependencies(
         should be provided if import name differs from package name
     severity : str, "error" (default), "warning", "none"
         behaviour for raising errors or warnings
-        "error" - raises a `ModuleNotFoundException` if one of packages is not installed
+        "error" - raises a `ModuleNotFoundError` if one of packages is not installed
         "warning" - raises a warning if one of packages is not installed
             function returns False if one of packages is not installed, otherwise True
         "none" - does not raise exception or warning
@@ -216,7 +216,7 @@ def _check_dl_dependencies(msg=None, severity="error"):
         error message to be returned in the `ModuleNotFoundError`, overrides default
     severity : str, "error" (default), "warning", "none"
         behaviour for raising errors or warnings
-        "error" - raises a ModuleNotFoundException if one of packages is not installed
+        "error" - raises a ModuleNotFoundError if one of packages is not installed
         "warning" - raises a warning if one of packages is not installed
             function returns False if one of packages is not installed, otherwise True
         "none" - does not raise exception or warning
@@ -254,6 +254,49 @@ def _check_dl_dependencies(msg=None, severity="error"):
             )
 
 
+def _check_mlflow_dependencies(
+    msg=None, severity="error", suppress_import_stdout=False
+):
+    """Check if `mlflow` and its dependencies are installed.
+
+    Parameters
+    ----------
+    msg: str, optional, default= default message (msg below)
+        error message to be returned when `ModuleNotFoundError` is raised.
+    severity: str, either of "error", "warning" or "none"
+        behaviour for raising errors or warnings
+        "error" - raises a `ModuleNotFound` if mlflow-related packages are not found.
+        "warning" - raises a warning message if any mlflow-related package is not
+            installed also returns False. In case all packages are present,
+            returns True.
+        "none" - does not raise any exception or warning and simply returns True
+            if all packages are installed otherwise return False.
+
+    Raise
+    -----
+    ModuleNotFoundError
+        User Friendly error with a suggested action to install mlflow dependencies
+
+    Returns
+    -------
+    boolean - whether all mlflow-related packages are installed.
+    """
+    if not isinstance(msg, str):
+        msg = (
+            "`mlflow` is an extra dependency and is not included "
+            "in the base sktime installation. "
+            "Please run `pip install mlflow` "
+            "or `pip install sktime[mlflow]` to install the package."
+        )
+
+    return _check_soft_dependencies(
+        "mlflow",
+        msg=msg,
+        severity=severity,
+        suppress_import_stdout=suppress_import_stdout,
+    )
+
+
 def _check_python_version(obj, package=None, msg=None, severity="error"):
     """Check if system python version is compatible with requirements of obj.
 
@@ -343,7 +386,7 @@ def _check_estimator_deps(obj, msg=None, severity="error"):
         error message to be returned in the `ModuleNotFoundError`, overrides default
     severity : str, "error" (default), "warning", or "none"
         behaviour for raising errors or warnings
-        "error" - raises a ModuleNotFoundException if environment is incompatible
+        "error" - raises a `ModuleNotFoundError` if environment is incompatible
         "warning" - raises a warning if environment is incompatible
             function returns False if environment is incompatible, otherwise True
         "none" - does not raise exception or warning


Test patch: