Repository: sktime/sktime
Pull request #5537 (2023-11-06 16:03:06)

[ENH] support for custom `joblib` backends in parallelization
Adds support for custom `joblib` backends in the `sktime` unified parallelization interface. Fixes https://github.com/sktime/sktime/issues/5535.

It seems that using custom backends was previously possible via parameter passing, albeit not properly tested as a possibility.

This PR adds back the possibility via a `"joblib"` parallelization mechanism.

I would be thankful if you could suggest how to best test this, @smirngregHM - preferably without adding soft dependencies.



fkiraly: @SmirnGregHM, have you had opportunity to test whether this works for your use case? And/or whether you have suggestions on changing this? I'd rather not merge and find out it does not work for you...

For instance, do you think that `backend="joblib"` should also work?
SmirnGregHM: @fkiraly sorry for the delay

I have just checked your solution and it works for me. Interface looks a little cumbersome to me though. I'd say, any joblib backend should be available with the same interface. It will also be consistent with [`joblib.Parallel`](https://joblib.readthedocs.io/en/latest/generated/joblib.Parallel.html) interface. But I also see the benefits of having more options than just joblib and benefits of keeping the old interface of passing `'loky'` or `'multiprocessing'` to `backend`.

What would you think about renaming `'joblib_custom'` to `'joblib'`? Then one can do any of:

```python
backend='threading', ... 

backend='joblib', backend_params={'backend': 'threading'}, ...

backend='joblib', backend_params={'backend': 'spark'}, ...
```


fkiraly: I see - that seems like an improvement!
This would also allow us, if we decide so, to remove direct specification of `joblib` backends.
fkiraly: done

Patch:
diff --git a/sktime/base/_meta.py b/sktime/base/_meta.py
index 78901df76..6e49c88e6 100644
--- a/sktime/base/_meta.py
+++ b/sktime/base/_meta.py
@@ -97,29 +97,31 @@ class _HeterogenousMetaEstimator:
 
     def _get_params(self, attr, deep=True, fitted=False):
         if fitted:
-            method = "_get_fitted_params"
-            methodd = "get_fitted_params"
+            private_method = "_get_fitted_params"
+            public_method = "get_fitted_params"
             deepkw = {}
         else:
-            method = "get_params"
-            methodd = "get_params"
+            private_method = "get_params"
+            public_method = "get_params"
             deepkw = {"deep": deep}
 
-        out = getattr(super(), method)(**deepkw)
+        out = getattr(super(), private_method)(**deepkw)
         if deep and hasattr(self, attr):
             estimators = getattr(self, attr)
             estimators = [(x[0], x[1]) for x in estimators]
             out.update(estimators)
             for name, estimator in estimators:
                 # checks estimator has the method we want to call
-                cond1 = hasattr(estimator, methodd)
+                cond1 = hasattr(estimator, public_method)
                 # checks estimator is fitted if calling get_fitted_params
                 is_fitted = hasattr(estimator, "is_fitted") and estimator.is_fitted
                 # if we call get_params and not get_fitted_params, this is True
                 cond2 = not fitted or is_fitted
                 # check both conditions together
                 if cond1 and cond2:
-                    for key, value in getattr(estimator, methodd)(**deepkw).items():
+                    for key, value in getattr(estimator, public_method)(
+                        **deepkw
+                    ).items():
                         out[f"{name}__{key}"] = value
         return out
 
diff --git a/sktime/datatypes/_vectorize.py b/sktime/datatypes/_vectorize.py
index 5d6d73e3b..f8755d435 100644
--- a/sktime/datatypes/_vectorize.py
+++ b/sktime/datatypes/_vectorize.py
@@ -511,6 +511,7 @@ class VectorizedDF:
 
             - "None": executes loop sequentally, simple list comprehension
             - "loky", "multiprocessing" and "threading": uses ``joblib.Parallel`` loops
+            - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``
             - "dask": uses ``dask``, requires ``dask`` package in environment
             - "dask_lazy": same as "dask", but returns delayed object instead
 
diff --git a/sktime/forecasting/base/_base.py b/sktime/forecasting/base/_base.py
index 46d226617..136c43a8d 100644
--- a/sktime/forecasting/base/_base.py
+++ b/sktime/forecasting/base/_base.py
@@ -115,7 +115,8 @@ class BaseForecaster(BaseEstimator):
             backend to use for parallelization when broadcasting/vectorizing, one of
 
             - "None": executes loop sequentally, simple list comprehension
-            - "loky", "multiprocessing" and "threading": uses ``joblib`` ``Parallel``
+            - "loky", "multiprocessing" and "threading": uses ``joblib.Parallel``
+            - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``
             - "dask": uses ``dask``, requires ``dask`` package in environment
         """,
         "backend:parallel:params": """
@@ -124,12 +125,20 @@ class BaseForecaster(BaseEstimator):
             Valid keys depend on the value of ``backend:parallel``:
 
             - "None": no additional parameters, ``backend_params`` is ignored
-            - "loky", "multiprocessing" and "threading":
-              any valid keys for ``joblib.Parallel`` can be passed here,
-              e.g., ``n_jobs``, with the exception of ``backend`` which is directly
-              controlled by ``backend:parallel``
-            - "dask": any valid keys for ``dask.compute``
-              can be passed, e.g., ``scheduler``
+            - "loky", "multiprocessing" and "threading": default ``joblib`` backends
+              any valid keys for ``joblib.Parallel`` can be passed here, e.g.,
+              ``n_jobs``, with the exception of ``backend`` which is directly
+              controlled by ``backend``.
+              If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+              will default to ``joblib`` defaults.
+            - "joblib": custom and 3rd party ``joblib`` backends,
+              e.g., ``spark``. Any valid keys for ``joblib.Parallel``
+              can be passed here, e.g., ``n_jobs``,
+            ``backend`` must be passed as a key of ``backend_params`` in this case.
+              If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+              will default to ``joblib`` defaults.
+            - "dask": any valid keys for ``dask.compute`` can be passed,
+              e.g., ``scheduler``
         """,
     }
 
diff --git a/sktime/forecasting/compose/_reduce.py b/sktime/forecasting/compose/_reduce.py
index 1fa634987..9f179fc26 100644
--- a/sktime/forecasting/compose/_reduce.py
+++ b/sktime/forecasting/compose/_reduce.py
@@ -1584,68 +1584,31 @@ def _create_fcst_df(target_date, origin_df, fill=None):
     -------
     A pandas dataframe or series
     """
-    oi = origin_df.index
-    if not isinstance(oi, pd.MultiIndex):
-        if isinstance(origin_df, pd.Series):
-            if fill is None:
-                template = pd.Series(np.zeros(len(target_date)), index=target_date)
-            else:
-                template = pd.Series(fill, index=target_date)
-            template.name = origin_df.name
-            return template
-        else:
-            if fill is None:
-                template = pd.DataFrame(
-                    np.zeros((len(target_date), len(origin_df.columns))),
-                    index=target_date,
-                    columns=origin_df.columns.to_list(),
-                )
-            else:
-                template = pd.DataFrame(
-                    fill, index=target_date, columns=origin_df.columns.to_list()
-                )
-            return template
+    if not isinstance(target_date, ForecastingHorizon):
+        ix = pd.Index(target_date)
+        fh = ForecastingHorizon(ix, is_relative=False)
     else:
-        idx = origin_df.index.to_frame(index=False)
-        instance_names = idx.columns[0:-1].to_list()
-        time_names = idx.columns[-1]
-        idx = idx[instance_names].drop_duplicates()
-
-        timeframe = pd.DataFrame(target_date, columns=[time_names])
-        target_frame = idx.merge(timeframe, how="cross")
-        if hasattr(target_date, "freq"):
-            freq_inferred = target_date.freq
-            mi = (
-                target_frame.groupby(instance_names, as_index=True)
-                .apply(
-                    lambda df: df.drop(instance_names, axis=1)
-                    .set_index(time_names)
-                    .asfreq(freq_inferred)
-                )
-                .index
-            )
-        else:
-            mi = (
-                target_frame.groupby(instance_names, as_index=True)
-                .apply(lambda df: df.drop(instance_names, axis=1).set_index(time_names))
-                .index
-            )
+        fh = target_date.to_absolute()
 
-        if fill is None:
-            template = pd.DataFrame(
-                np.zeros((len(target_date) * idx.shape[0], len(origin_df.columns))),
-                index=mi,
-                columns=origin_df.columns.to_list(),
-            )
-        else:
-            template = pd.DataFrame(
-                fill,
-                index=mi,
-                columns=origin_df.columns.to_list(),
-            )
+    index = fh.get_expected_pred_idx(origin_df)
+
+    if isinstance(origin_df, pd.Series):
+        columns = [origin_df.name]
+    else:
+        columns = origin_df.columns.to_list()
+
+    if fill is None:
+        values = 0
+    else:
+        values = fill
+
+    res = pd.DataFrame(values, index=index, columns=columns)
+
+    if isinstance(origin_df, pd.Series) and not isinstance(index, pd.MultiIndex):
+        res = res.iloc[:, 0]
+        res.name = origin_df.name
 
-        template = template.astype(origin_df.dtypes.to_dict())
-        return template
+    return res
 
 
 def _coerce_col_str(X):
diff --git a/sktime/forecasting/model_evaluation/_functions.py b/sktime/forecasting/model_evaluation/_functions.py
index 155caabd8..613360937 100644
--- a/sktime/forecasting/model_evaluation/_functions.py
+++ b/sktime/forecasting/model_evaluation/_functions.py
@@ -413,6 +413,7 @@ def evaluate(
 
         - "None": executes loop sequentally, simple list comprehension
         - "loky", "multiprocessing" and "threading": uses ``joblib.Parallel`` loops
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``
         - "dask": uses ``dask``, requires ``dask`` package in environment
         - "dask_lazy": same as "dask",
           but changes the return to (lazy) ``dask.dataframe.DataFrame``.
@@ -435,12 +436,18 @@ def evaluate(
         Valid keys depend on the value of ``backend``:
 
         - "None": no additional parameters, ``backend_params`` is ignored
-        - "loky", "multiprocessing" and "threading":
-            any valid keys for ``joblib.Parallel`` can be passed here,
-            e.g., ``n_jobs``, with the exception of ``backend``
-            which is directly controlled by ``backend``
+        - "loky", "multiprocessing" and "threading": default ``joblib`` backends
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          with the exception of ``backend`` which is directly controlled by ``backend``.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``.
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          ``backend`` must be passed as a key of ``backend_params`` in this case.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
         - "dask": any valid keys for ``dask.compute`` can be passed,
-            e.g., ``scheduler``
+          e.g., ``scheduler``
 
     Returns
     -------
diff --git a/sktime/forecasting/model_selection/_tune.py b/sktime/forecasting/model_selection/_tune.py
index 7b70b6d1e..fd0fcd585 100644
--- a/sktime/forecasting/model_selection/_tune.py
+++ b/sktime/forecasting/model_selection/_tune.py
@@ -464,6 +464,7 @@ class ForecastingGridSearchCV(BaseGridSearch):
 
         - "None": executes loop sequentally, simple list comprehension
         - "loky", "multiprocessing" and "threading": uses ``joblib.Parallel`` loops
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``
         - "dask": uses ``dask``, requires ``dask`` package in environment
 
         Recommendation: Use "dask" or "loky" for parallel evaluate.
@@ -498,10 +499,16 @@ class ForecastingGridSearchCV(BaseGridSearch):
         Valid keys depend on the value of ``backend``:
 
         - "None": no additional parameters, ``backend_params`` is ignored
-        - "loky", "multiprocessing" and "threading":
-            any valid keys for ``joblib.Parallel`` can be passed here,
-            e.g., ``n_jobs``, with the exception of ``backend``
-            which is directly controlled by ``backend``
+        - "loky", "multiprocessing" and "threading": default ``joblib`` backends
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          with the exception of ``backend`` which is directly controlled by ``backend``.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``.
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          ``backend`` must be passed as a key of ``backend_params`` in this case.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
         - "dask": any valid keys for ``dask.compute`` can be passed, e.g., ``scheduler``
 
     Attributes
@@ -796,6 +803,7 @@ class ForecastingRandomizedSearchCV(BaseGridSearch):
 
         - "None": executes loop sequentally, simple list comprehension
         - "loky", "multiprocessing" and "threading": uses ``joblib.Parallel`` loops
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``
         - "dask": uses ``dask``, requires ``dask`` package in environment
 
         Recommendation: Use "dask" or "loky" for parallel evaluate.
@@ -830,10 +838,16 @@ class ForecastingRandomizedSearchCV(BaseGridSearch):
         Valid keys depend on the value of ``backend``:
 
         - "None": no additional parameters, ``backend_params`` is ignored
-        - "loky", "multiprocessing" and "threading":
-            any valid keys for ``joblib.Parallel`` can be passed here,
-            e.g., ``n_jobs``, with the exception of ``backend``
-            which is directly controlled by ``backend``
+        - "loky", "multiprocessing" and "threading": default ``joblib`` backends
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          with the exception of ``backend`` which is directly controlled by ``backend``.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``.
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          ``backend`` must be passed as a key of ``backend_params`` in this case.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
         - "dask": any valid keys for ``dask.compute`` can be passed, e.g., ``scheduler``
 
     Attributes
@@ -1055,6 +1069,7 @@ class ForecastingSkoptSearchCV(BaseGridSearch):
 
         - "None": executes loop sequentally, simple list comprehension
         - "loky", "multiprocessing" and "threading": uses ``joblib.Parallel`` loops
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``
         - "dask": uses ``dask``, requires ``dask`` package in environment
 
         Recommendation: Use "dask" or "loky" for parallel evaluate.
@@ -1085,10 +1100,16 @@ class ForecastingSkoptSearchCV(BaseGridSearch):
         Valid keys depend on the value of ``backend``:
 
         - "None": no additional parameters, ``backend_params`` is ignored
-        - "loky", "multiprocessing" and "threading":
-            any valid keys for ``joblib.Parallel`` can be passed here,
-            e.g., ``n_jobs``, with the exception of ``backend``
-            which is directly controlled by ``backend``
+        - "loky", "multiprocessing" and "threading": default ``joblib`` backends
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          with the exception of ``backend`` which is directly controlled by ``backend``.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``.
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          ``backend`` must be passed as a key of ``backend_params`` in this case.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
         - "dask": any valid keys for ``dask.compute`` can be passed, e.g., ``scheduler``
 
     Attributes
diff --git a/sktime/transformations/base.py b/sktime/transformations/base.py
index 94ed57546..126ce81f4 100644
--- a/sktime/transformations/base.py
+++ b/sktime/transformations/base.py
@@ -166,7 +166,8 @@ class BaseTransformer(BaseEstimator):
             backend to use for parallelization when broadcasting/vectorizing, one of
 
             - "None": executes loop sequentally, simple list comprehension
-            - "loky", "multiprocessing" and "threading": uses ``joblib`` ``Parallel``
+            - "loky", "multiprocessing" and "threading": uses ``joblib.Parallel``
+            - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``
             - "dask": uses ``dask``, requires ``dask`` package in environment
         """,
         "backend:parallel:params": """
@@ -175,10 +176,18 @@ class BaseTransformer(BaseEstimator):
             Valid keys depend on the value of ``backend:parallel``:
 
             - "None": no additional parameters, ``backend_params`` is ignored
-            - "loky", "multiprocessing" and "threading":
-              any valid keys for ``joblib.Parallel`` can be passed here,
-              e.g., ``n_jobs``, with the exception of ``backend`` which is directly
-              controlled by ``backend:parallel``
+            - "loky", "multiprocessing" and "threading": default ``joblib`` backends
+              any valid keys for ``joblib.Parallel`` can be passed here, e.g.,
+              ``n_jobs``, with the exception of ``backend`` which is directly
+              controlled by ``backend``.
+              If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+              will default to ``joblib`` defaults.
+            - "joblib": custom and 3rd party ``joblib`` backends,
+              e.g., ``spark``. Any valid keys for ``joblib.Parallel``
+              can be passed here, e.g., ``n_jobs``,
+            ``backend`` must be passed as a key of ``backend_params`` in this case.
+              If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+              will default to ``joblib`` defaults.
             - "dask": any valid keys for ``dask.compute``
               can be passed, e.g., ``scheduler``
         """,
diff --git a/sktime/utils/parallel.py b/sktime/utils/parallel.py
index ab839145c..6a8fb4090 100644
--- a/sktime/utils/parallel.py
+++ b/sktime/utils/parallel.py
@@ -38,6 +38,7 @@ def parallelize(fun, iter, meta=None, backend=None, backend_params=None):
 
         - "None": executes loop sequentally, simple list comprehension
         - "loky", "multiprocessing" and "threading": uses ``joblib`` ``Parallel`` loops
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``
         - "dask": uses ``dask``, requires ``dask`` package in environment
         - "dask_lazy": same as ``"dask"``, but returns delayed object instead of list
 
@@ -46,9 +47,16 @@ def parallelize(fun, iter, meta=None, backend=None, backend_params=None):
         Valid keys depend on the value of ``backend``:
 
         - "None": no additional parameters, ``backend_params`` is ignored
-        - "loky", "multiprocessing" and "threading":
+        - "loky", "multiprocessing" and "threading": default ``joblib`` backends
           any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
-          with the exception of ``backend`` which is directly controlled by ``backend``
+          with the exception of ``backend`` which is directly controlled by ``backend``.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
+        - "joblib": custom and 3rd party ``joblib`` backends, e.g., ``spark``.
+          any valid keys for ``joblib.Parallel`` can be passed here, e.g., ``n_jobs``,
+          ``backend`` must be passed as a key of ``backend_params`` in this case.
+          If ``n_jobs`` is not passed, it will default to ``-1``, other parameters
+          will default to ``joblib`` defaults.
         - "dask": any valid keys for ``dask.compute`` can be passed, e.g., ``scheduler``
     """
     if meta is None:
@@ -72,6 +80,7 @@ backend_dict = {
     "loky": "joblib",
     "multiprocessing": "joblib",
     "threading": "joblib",
+    "joblib": "joblib",
     "dask": "dask",
     "dask_lazy": "dask",
 }
@@ -92,7 +101,29 @@ def _parallelize_joblib(fun, iter, meta, backend, backend_params):
     from joblib import Parallel, delayed
 
     par_params = backend_params.copy()
-    par_params["backend"] = backend
+    if "backend" not in par_params:
+        # if user selects custom joblib backend but does not specify backend explicitly,
+        # raise a ValueError
+        if backend == "joblib":
+            raise ValueError(
+                '"joblib" was selected as first layer parallelization backend, '
+                "but no backend string was "
+                '"passed in the backend parameters dict, e.g., "spark". '
+                "Please specify a backend to joblib as a key-value pair "
+                "in the backend_params arg or the backend:parallel:params config "
+                'when using "joblib". '
+                'For clarity, "joblib" should only be used for two-layer '
+                "backend dispatch, where the first layer is joblib, "
+                "and the second layer is a custom backend of joblib, e.g., spark."
+                "For first-party joblib backends, please use the backend string "
+                'of sktime directly, e.g., by specifying "multiprocessing" or "loky".'
+            )
+        # in all other cases, we ensure the backend parameter is one of
+        # "loky", "multiprocessing" or "threading", as passed via backend
+        else:
+            par_params["backend"] = backend
+    elif backend != "joblib":
+        par_params["backend"] = backend
 
     if "n_jobs" not in par_params:
         par_params["n_jobs"] = -1


Test patch:
diff --git a/sktime/forecasting/compose/tests/test_reduce_global.py b/sktime/forecasting/compose/tests/test_reduce_global.py
index c332178dd..e2a78d278 100644
--- a/sktime/forecasting/compose/tests/test_reduce_global.py
+++ b/sktime/forecasting/compose/tests/test_reduce_global.py
@@ -106,6 +106,7 @@ def y_dict():
     # Create integer index data
     y_numeric = y_train.copy()
     y_numeric.index = pd.to_numeric(y_numeric.index)
+    y_numeric.index.names = [None]  # setting None to cover "no index name" case
     y_dict["y_numeric"] = y_numeric
 
     return y_dict
@@ -139,7 +140,7 @@ def check_eval(test_input, expected):
         ),
         (
             "y_train",
-            [None],
+            ["Period"],
         ),
         (
             "y_numeric",
@@ -186,7 +187,7 @@ def test_recursive_reduction(y, index_names, y_dict):
         ),
         (
             "y_train",
-            [None],
+            ["Period"],
         ),
         (
             "y_numeric",
@@ -232,7 +233,7 @@ def test_direct_reduction(y, index_names, y_dict):
         ),
         (
             "y_train",
-            [None],
+            ["Period"],
         ),
         (
             "y_numeric",