FIX Passes global configuration when spawning joblib jobs (#17634)

* WIP

* BUG Passes context to joblib jobs

* BUG Fix

* BUG Fix

* CLN Moves delayed into fixes

* CLN Adds linter

* REV Revert diff

* DOC Adds comment for removal

* Use a simple implementation

benchmarks/bench_saga.py

@@ -7,7 +7,8 @@
 import time
 import os
 
-from joblib import delayed, Parallel
+from joblib import Parallel
+from sklearn.utils.fixes import delayed
 import matplotlib.pyplot as plt
 import numpy as np
 

build_tools/circle/linting.sh

@@ -172,3 +172,11 @@ then
     echo "$doctest_directive"
     exit 1
 fi
+
+joblib_import="$(git grep -l -A 10 -E "joblib import.+delayed" -- "*.py" ":!sklearn/utils/_joblib.py" ":!sklearn/utils/fixes.py")"
+
+if [ ! -z "$joblib_import" ]; then
+    echo "Use from sklearn.utils.fixes import delayed instead of joblib delayed. The following files contains imports to joblib.delayed:"
+    echo "$joblib_import"
+    exit 1
+fi

sklearn/calibration.py

@@ -13,7 +13,7 @@
 
 from math import log
 import numpy as np
-from joblib import delayed, Parallel
+from joblib import Parallel
 
 from scipy.special import expit
 from scipy.special import xlogy
@@ -24,6 +24,7 @@
                    MetaEstimatorMixin)
 from .preprocessing import label_binarize, LabelBinarizer
 from .utils import check_array, indexable, column_or_1d
+from .utils.fixes import delayed
 from .utils.validation import check_is_fitted, check_consistent_length
 from .utils.validation import _check_sample_weight
 from .pipeline import Pipeline

sklearn/cluster/_mean_shift.py

@@ -16,10 +16,11 @@
 
 import numpy as np
 import warnings
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from collections import defaultdict
 from ..utils.validation import check_is_fitted, _deprecate_positional_args
+from ..utils.fixes import delayed
 from ..utils import check_random_state, gen_batches, check_array
 from ..base import BaseEstimator, ClusterMixin
 from ..neighbors import NearestNeighbors

sklearn/compose/_column_transformer.py

@@ -12,7 +12,7 @@
 import numbers
 import numpy as np
 from scipy import sparse
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from ..base import clone, TransformerMixin
 from ..utils._estimator_html_repr import _VisualBlock
@@ -25,6 +25,7 @@
 from ..utils.metaestimators import _BaseComposition
 from ..utils.validation import check_array, check_is_fitted
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 
 
 __all__ = [

sklearn/covariance/_graph_lasso.py

@@ -13,13 +13,14 @@
 
 import numpy as np
 from scipy import linalg
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from . import empirical_covariance, EmpiricalCovariance, log_likelihood
 
 from ..exceptions import ConvergenceWarning
 from ..utils.validation import check_random_state
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 # mypy error: Module 'sklearn.linear_model' has no attribute '_cd_fast'
 from ..linear_model import _cd_fast as cd_fast  # type: ignore
 from ..linear_model import lars_path_gram

sklearn/decomposition/_dict_learning.py

@@ -11,14 +11,15 @@
 
 import numpy as np
 from scipy import linalg
-from joblib import Parallel, delayed, effective_n_jobs
+from joblib import Parallel, effective_n_jobs
 
 from ..base import BaseEstimator, TransformerMixin
 from ..utils import deprecated
 from ..utils import (check_array, check_random_state, gen_even_slices,
                      gen_batches)
 from ..utils.extmath import randomized_svd, row_norms
 from ..utils.validation import check_is_fitted, _deprecate_positional_args
+from ..utils.fixes import delayed
 from ..linear_model import Lasso, orthogonal_mp_gram, LassoLars, Lars
 
 

sklearn/decomposition/_lda.py

@@ -14,13 +14,14 @@
 import numpy as np
 import scipy.sparse as sp
 from scipy.special import gammaln, logsumexp
-from joblib import Parallel, delayed, effective_n_jobs
+from joblib import Parallel, effective_n_jobs
 
 from ..base import BaseEstimator, TransformerMixin
 from ..utils import check_random_state, gen_batches, gen_even_slices
 from ..utils.validation import check_non_negative
 from ..utils.validation import check_is_fitted
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 
 from ._online_lda_fast import (mean_change, _dirichlet_expectation_1d,
                                _dirichlet_expectation_2d)

sklearn/ensemble/_bagging.py

@@ -10,7 +10,7 @@
 from abc import ABCMeta, abstractmethod
 from warnings import warn
 
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from ._base import BaseEnsemble, _partition_estimators
 from ..base import ClassifierMixin, RegressorMixin
@@ -23,6 +23,7 @@
 from ..utils.random import sample_without_replacement
 from ..utils.validation import has_fit_parameter, check_is_fitted, \
     _check_sample_weight, _deprecate_positional_args
+from ..utils.fixes import delayed
 
 
 __all__ = ["BaggingClassifier",

sklearn/ensemble/_forest.py

@@ -48,7 +48,7 @@ class calls the ``fit`` method of each sub-estimator on random samples
 import numpy as np
 from scipy.sparse import issparse
 from scipy.sparse import hstack as sparse_hstack
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from ..base import ClassifierMixin, RegressorMixin, MultiOutputMixin
 from ..metrics import r2_score
@@ -59,6 +59,7 @@ class calls the ``fit`` method of each sub-estimator on random samples
 from ..utils import check_random_state, check_array, compute_sample_weight
 from ..exceptions import DataConversionWarning
 from ._base import BaseEnsemble, _partition_estimators
+from ..utils.fixes import delayed
 from ..utils.fixes import _joblib_parallel_args
 from ..utils.multiclass import check_classification_targets
 from ..utils.validation import check_is_fitted, _check_sample_weight

sklearn/ensemble/_stacking.py

@@ -7,7 +7,7 @@
 from copy import deepcopy
 
 import numpy as np
-from joblib import Parallel, delayed
+from joblib import Parallel
 import scipy.sparse as sparse
 
 from ..base import clone
@@ -33,6 +33,7 @@
 from ..utils.validation import check_is_fitted
 from ..utils.validation import column_or_1d
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 
 
 class _BaseStacking(TransformerMixin, _BaseHeterogeneousEnsemble,

sklearn/ensemble/_voting.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from ..base import ClassifierMixin
 from ..base import RegressorMixin
@@ -33,6 +33,7 @@
 from ..utils.validation import _deprecate_positional_args
 from ..exceptions import NotFittedError
 from ..utils._estimator_html_repr import _VisualBlock
+from ..utils.fixes import delayed
 
 
 class _BaseVoting(TransformerMixin, _BaseHeterogeneousEnsemble):

sklearn/feature_selection/_rfe.py

@@ -8,12 +8,13 @@
 
 import numpy as np
 import numbers
-from joblib import Parallel, delayed, effective_n_jobs
+from joblib import Parallel, effective_n_jobs
 
 from ..utils.metaestimators import if_delegate_has_method
 from ..utils.metaestimators import _safe_split
 from ..utils.validation import check_is_fitted
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 from ..base import BaseEstimator
 from ..base import MetaEstimatorMixin
 from ..base import clone

sklearn/inspection/_permutation_importance.py

@@ -1,13 +1,13 @@
 """Permutation importance for estimators."""
 import numpy as np
 from joblib import Parallel
-from joblib import delayed
 
 from ..metrics import check_scoring
 from ..utils import Bunch
 from ..utils import check_random_state
 from ..utils import check_array
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 
 
 def _weights_scorer(scorer, estimator, X, y, sample_weight):

sklearn/inspection/_plot/partial_dependence.py

@@ -6,14 +6,15 @@
 import numpy as np
 from scipy import sparse
 from scipy.stats.mstats import mquantiles
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from .. import partial_dependence
 from ...base import is_regressor
 from ...utils import check_array
 from ...utils import check_matplotlib_support  # noqa
 from ...utils import _safe_indexing
 from ...utils.validation import _deprecate_positional_args
+from ...utils.fixes import delayed
 
 
 @_deprecate_positional_args

sklearn/linear_model/_base.py

@@ -23,7 +23,7 @@
 from scipy import optimize
 from scipy import sparse
 from scipy.special import expit
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from ..base import (BaseEstimator, ClassifierMixin, RegressorMixin,
                     MultiOutputMixin)
@@ -37,6 +37,7 @@
 from ..utils._seq_dataset import ArrayDataset32, CSRDataset32
 from ..utils._seq_dataset import ArrayDataset64, CSRDataset64
 from ..utils.validation import check_is_fitted, _check_sample_weight
+from ..utils.fixes import delayed
 from ..preprocessing import normalize as f_normalize
 
 # TODO: bayesian_ridge_regression and bayesian_regression_ard

sklearn/linear_model/_coordinate_descent.py

@@ -12,7 +12,7 @@
 
 import numpy as np
 from scipy import sparse
-from joblib import Parallel, delayed, effective_n_jobs
+from joblib import Parallel, effective_n_jobs
 
 from ._base import LinearModel, _pre_fit
 from ..base import RegressorMixin, MultiOutputMixin
@@ -25,6 +25,7 @@
 from ..utils.validation import check_is_fitted, _check_sample_weight
 from ..utils.validation import column_or_1d
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 
 # mypy error: Module 'sklearn.linear_model' has no attribute '_cd_fast'
 from . import _cd_fast as cd_fast  # type: ignore

sklearn/linear_model/_least_angle.py

@@ -15,7 +15,7 @@
 import numpy as np
 from scipy import linalg, interpolate
 from scipy.linalg.lapack import get_lapack_funcs
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from ._base import LinearModel
 from ..base import RegressorMixin, MultiOutputMixin
@@ -25,6 +25,7 @@
 from ..model_selection import check_cv
 from ..exceptions import ConvergenceWarning
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 
 SOLVE_TRIANGULAR_ARGS = {'check_finite': False}
 

sklearn/linear_model/_logistic.py

@@ -16,7 +16,7 @@
 import numpy as np
 from scipy import optimize, sparse
 from scipy.special import expit, logsumexp
-from joblib import Parallel, delayed, effective_n_jobs
+from joblib import Parallel, effective_n_jobs
 
 from ._base import LinearClassifierMixin, SparseCoefMixin, BaseEstimator
 from ._sag import sag_solver
@@ -32,6 +32,7 @@
 from ..utils.validation import _deprecate_positional_args
 from ..utils.multiclass import check_classification_targets
 from ..utils.fixes import _joblib_parallel_args
+from ..utils.fixes import delayed
 from ..model_selection import check_cv
 from ..metrics import get_scorer
 

sklearn/linear_model/_omp.py

@@ -11,12 +11,13 @@
 import numpy as np
 from scipy import linalg
 from scipy.linalg.lapack import get_lapack_funcs
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from ._base import LinearModel, _pre_fit
 from ..base import RegressorMixin, MultiOutputMixin
 from ..utils import as_float_array, check_array
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 from ..model_selection import check_cv
 
 premature = """ Orthogonal matching pursuit ended prematurely due to linear

sklearn/linear_model/_stochastic_gradient.py

@@ -9,7 +9,7 @@
 
 from abc import ABCMeta, abstractmethod
 
-from joblib import Parallel, delayed
+from joblib import Parallel
 
 from ..base import clone, is_classifier
 from ._base import LinearClassifierMixin, SparseCoefMixin
@@ -20,6 +20,7 @@
 from ..utils.multiclass import _check_partial_fit_first_call
 from ..utils.validation import check_is_fitted, _check_sample_weight
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 from ..exceptions import ConvergenceWarning
 from ..model_selection import StratifiedShuffleSplit, ShuffleSplit
 

sklearn/linear_model/_theil_sen.py

@@ -15,12 +15,13 @@
 from scipy import linalg
 from scipy.special import binom
 from scipy.linalg.lapack import get_lapack_funcs
-from joblib import Parallel, delayed, effective_n_jobs
+from joblib import Parallel, effective_n_jobs
 
 from ._base import LinearModel
 from ..base import RegressorMixin
 from ..utils import check_random_state
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 from ..exceptions import ConvergenceWarning
 
 _EPSILON = np.finfo(np.double).eps

sklearn/manifold/_mds.py

@@ -6,7 +6,7 @@
 # License: BSD
 
 import numpy as np
-from joblib import Parallel, delayed, effective_n_jobs
+from joblib import Parallel, effective_n_jobs
 
 import warnings
 
@@ -15,6 +15,7 @@
 from ..utils import check_random_state, check_array, check_symmetric
 from ..isotonic import IsotonicRegression
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 
 
 def _smacof_single(dissimilarities, metric=True, n_components=2, init=None,

sklearn/metrics/pairwise.py

@@ -17,7 +17,7 @@
 from scipy.spatial import distance
 from scipy.sparse import csr_matrix
 from scipy.sparse import issparse
-from joblib import Parallel, delayed, effective_n_jobs
+from joblib import Parallel, effective_n_jobs
 
 from ..utils.validation import _num_samples
 from ..utils.validation import check_non_negative
@@ -29,6 +29,7 @@
 from ..preprocessing import normalize
 from ..utils._mask import _get_mask
 from ..utils.validation import _deprecate_positional_args
+from ..utils.fixes import delayed
 from ..utils.fixes import sp_version, parse_version
 
 from ._pairwise_fast import _chi2_kernel_fast, _sparse_manhattan

sklearn/model_selection/_search.py

@@ -32,12 +32,13 @@
 from ._validation import _insert_error_scores
 from ._validation import _normalize_score_results
 from ..exceptions import NotFittedError
-from joblib import Parallel, delayed
+from joblib import Parallel
 from ..utils import check_random_state
 from ..utils.random import sample_without_replacement
 from ..utils.validation import indexable, check_is_fitted, _check_fit_params
 from ..utils.validation import _deprecate_positional_args
 from ..utils.metaestimators import if_delegate_has_method
+from ..utils.fixes import delayed
 from ..metrics._scorer import _check_multimetric_scoring
 from ..metrics import check_scoring
 from ..utils import deprecated