update joblib_temp_folder setting

Author: chenyangkang

stemflow/model/AdaSTEM.py

@@ -114,7 +114,8 @@ def __init__(
         lazy_loading_dir: Union[str, None] = None,
         min_class_sample: int = 1,
         ensemble_bootstrap: bool = False,
-        joblib_backend: str = 'loky'
+        joblib_backend: str = 'loky',
+        joblib_temp_folder: Union[None, str] = None
     ):
         """Make an AdaSTEM object
 
@@ -192,6 +193,8 @@ def __init__(
                 Whether to bootstrap the data at each ensemble level to account for uncertainty. Defaults to False.
             joblib_backend:
                 The backend of joblib. Defaults to 'loky'. Other options include 'multiprocessing', 'threading'.
+            joblib_temp_folder:
+                The temporary folder for joblib. If None, falling back to joblib's default directory. If 'lazy_loading_dir', set as the same directory as lazy_loading_dir. If it's string, create a directory and store data into it. Defaults to None.
         Raises:
             AttributeError: Base model do not have method 'fit' or 'predict'
             AttributeError: task not in one of ['regression', 'classification', 'hurdle']
@@ -267,6 +270,7 @@ def __init__(
         n_jobs = check_transform_n_jobs(self, n_jobs)
         self.n_jobs = n_jobs
         self.joblib_backend = joblib_backend
+        self.joblib_temp_folder = joblib_temp_folder
 
         # 7. Plotting params
         self.plot_xlims = plot_xlims
@@ -374,7 +378,7 @@ def split(self, X_train: pd.core.frame.DataFrame, verbosity: Union[None, int] =
         )
 
         if n_jobs > 1 and isinstance(n_jobs, int):
-            parallel = joblib.Parallel(n_jobs=n_jobs, return_as="generator", backend=self.joblib_backend, temp_folder=self.lazy_loading_dir)
+            parallel = joblib.Parallel(n_jobs=n_jobs, return_as="generator", backend=self.joblib_backend, temp_folder=self.joblib_temp_folder)
             output_generator = parallel(
                 joblib.delayed(partial_get_one_ensemble_quadtree)(
                     ensemble_count=ensemble_count, rng=np.random.default_rng(self.rng.integers(1e9) + ensemble_count)
@@ -573,7 +577,7 @@ def mp_train(ensemble, self=self, data=data):
                 res = self.SAC_ensemble_training(index_df=ensemble[1], data=data)
                 return res
 
-            parallel = joblib.Parallel(n_jobs=n_jobs, return_as="generator", backend=self.joblib_backend, temp_folder=self.lazy_loading_dir)
+            parallel = joblib.Parallel(n_jobs=n_jobs, return_as="generator", backend=self.joblib_backend, temp_folder=self.joblib_temp_folder)
             output_generator = parallel(joblib.delayed(mp_train)(i) for i in groups)
 
         # tqdm wrapper
@@ -642,6 +646,15 @@ def fit(
                 shutil.rmtree(self.lazy_loading_dir)
         self.lazy_loading_dir = str(Path(self.lazy_loading_dir.rstrip('/\\')))
 
+        # Setup joblib_temp_folder
+        if self.joblib_temp_folder is None:
+            pass
+        elif self.joblib_temp_folder=='lazy_loading_dir':
+            self.joblib_temp_folder = self.lazy_loading_dir
+        else:
+            if not os.path.exists(self.joblib_temp_folder):
+                os.makedirs(self.joblib_temp_folder)
+        
         verbosity = check_verbosity(self, verbosity)
         check_X_train(X_train)
         check_y_train(y_train)
@@ -804,7 +817,7 @@ def mp_predict(ensemble, self=self, data=data):
                 res = self.SAC_ensemble_predict(index_df=ensemble[1], data=data)
                 return res
 
-            parallel = joblib.Parallel(n_jobs=n_jobs, return_as="generator", backend=self.joblib_backend, temp_folder=self.lazy_loading_dir)
+            parallel = joblib.Parallel(n_jobs=n_jobs, return_as="generator", backend=self.joblib_backend, temp_folder=self.joblib_temp_folder)
             output_generator = parallel(joblib.delayed(mp_predict)(i) for i in groups)
 
         # tqdm wrapper
@@ -1224,7 +1237,7 @@ def assign_feature_importances_by_points(
 
         # assign input spatio-temporal points to stixels
         if n_jobs > 1:
-            parallel = joblib.Parallel(n_jobs=n_jobs, return_as="generator", backend=self.joblib_backend, temp_folder=self.lazy_loading_dir)
+            parallel = joblib.Parallel(n_jobs=n_jobs, return_as="generator", backend=self.joblib_backend, temp_folder=self.joblib_temp_folder)
             output_generator = parallel(joblib.delayed(partial_assign_func)(i) for i in list(range(self.ensemble_fold)))
             if verbosity > 0:
                 output_generator = tqdm(output_generator, total=self.ensemble_fold, desc="Querying ensembles: ")
@@ -1381,7 +1394,8 @@ def __init__(
         lazy_loading_dir = None,
         min_class_sample = 1,
         ensemble_bootstrap = False,
-        joblib_backend = 'loky'
+        joblib_backend = 'loky',
+        joblib_temp_folder = None
     ):
         super().__init__(
             base_model=base_model,
@@ -1416,7 +1430,8 @@ def __init__(
             lazy_loading_dir=lazy_loading_dir,
             min_class_sample=min_class_sample,
             ensemble_bootstrap=ensemble_bootstrap,
-            joblib_backend=joblib_backend
+            joblib_backend=joblib_backend,
+            joblib_temp_folder = joblib_temp_folder
         )
 
         self._estimator_type = 'classifier'
@@ -1569,7 +1584,8 @@ def __init__(
         lazy_loading_dir=None,
         min_class_sample=1,
         ensemble_bootstrap=False,
-        joblib_backend='loky'
+        joblib_backend='loky',
+        joblib_temp_folder=None
     ):
         super().__init__(
             base_model=base_model,
@@ -1604,7 +1620,8 @@ def __init__(
             lazy_loading_dir=lazy_loading_dir,
             min_class_sample=min_class_sample,
             ensemble_bootstrap=ensemble_bootstrap,
-            joblib_backend=joblib_backend
+            joblib_backend=joblib_backend,
+            joblib_temp_folder=joblib_temp_folder
         )
 
         self._estimator_type = 'regressor'

stemflow/model/STEM.py

@@ -51,7 +51,8 @@ def __init__(
         lazy_loading_dir: Union[str, None] = None,
         min_class_sample: int = 1,
         ensemble_bootstrap: bool = False,
-        joblib_backend: str = 'loky'
+        joblib_backend: str = 'loky',
+        joblib_temp_folder: Union[None, str] = None
     ):
         """Make a STEM object
 
@@ -127,6 +128,8 @@ def __init__(
                 Whether to bootstrap the data at each ensemble level to account for uncertainty. Defaults to False.
             joblib_backend:
                 The backend of joblib. Defaults to 'loky'. Other options include 'multiprocessing', 'threading'.
+            joblib_temp_folder:
+                The temporary folder for joblib. If None, falling back to joblib's default directory. If 'lazy_loading_dir', set as the same directory as lazy_loading_dir. If it's string, create a directory and store data into it. Defaults to None.
         Raises:
             AttributeError: Base model do not have method 'fit' or 'predict'
             AttributeError: task not in one of ['regression', 'classification', 'hurdle']
@@ -186,7 +189,8 @@ def __init__(
             lazy_loading_dir=lazy_loading_dir,
             min_class_sample=min_class_sample,
             ensemble_bootstrap=ensemble_bootstrap,
-            joblib_backend=joblib_backend
+            joblib_backend=joblib_backend,
+            joblib_temp_folder=joblib_temp_folder
         )
 
         self.grid_len = grid_len
@@ -254,7 +258,8 @@ def __init__(
         lazy_loading_dir: Union[str, None] = None,
         min_class_sample: int = 1,
         ensemble_bootstrap: bool = False,
-        joblib_backend: str = 'loky'
+        joblib_backend: str = 'loky',
+        joblib_temp_folder: Union[None, str] = None
     ):
         super().__init__(
             base_model=base_model,
@@ -289,7 +294,8 @@ def __init__(
             lazy_loading_dir=lazy_loading_dir,
             min_class_sample=min_class_sample,
             ensemble_bootstrap=ensemble_bootstrap,
-            joblib_backend=joblib_backend
+            joblib_backend=joblib_backend,
+            joblib_temp_folder=joblib_temp_folder
         )
 
         self.grid_len = grid_len
@@ -357,7 +363,8 @@ def __init__(
         lazy_loading_dir: Union[str, None] = None,
         min_class_sample: int = 1,
         ensemble_bootstrap: bool = False,
-        joblib_backend: str = 'loky'
+        joblib_backend: str = 'loky',
+        joblib_temp_folder: Union[None, str]= None
     ):
         super().__init__(
             base_model=base_model,
@@ -392,7 +399,8 @@ def __init__(
             lazy_loading_dir=lazy_loading_dir,
             min_class_sample=min_class_sample,
             ensemble_bootstrap=ensemble_bootstrap,
-            joblib_backend=joblib_backend
+            joblib_backend=joblib_backend,
+            joblib_temp_folder=joblib_temp_folder
         )
 
         self.grid_len = grid_len

tests/test_joblib_temp_folder.py

@@ -0,0 +1,37 @@
+import numpy as np
+import pandas as pd
+import os
+
+from stemflow.model.AdaSTEM import AdaSTEM
+from stemflow.model_selection import ST_train_test_split
+
+from .make_models import (
+    make_AdaSTEMClassifier,
+    make_AdaSTEMRegressor,
+    make_parallel_SphereAdaClassifier,
+    make_parallel_STEMClassifier,
+    make_SphereAdaClassifier,
+    make_SphereAdaSTEMRegressor,
+    make_STEMClassifier,
+    make_STEMRegressor,
+)
+from .set_up_data import set_up_data
+
+x_names, (X, y) = set_up_data()
+X_train, X_test, y_train, y_test = ST_train_test_split(
+    X, y, Spatio_blocks_count=100, Temporal_blocks_count=100, random_state=42, test_size=0.3
+)
+
+
+def test_AdaSTEMRegressor_custom_temp_folder1():
+    model = make_AdaSTEMRegressor(lazy_loading=True, joblib_temp_folder='lazy_loading_dir')
+    model = model.fit(X_train, np.where(y_train > 0, 1, 0))
+
+    pred_mean, pred_std = model.predict(X_test.reset_index(drop=True), return_std=True, verbosity=1, n_jobs=1)
+
+def test_AdaSTEMRegressor_custom_temp_folder2():
+    model = make_AdaSTEMRegressor(lazy_loading=True, joblib_temp_folder='./test_tmp_folder')
+    model = model.fit(X_train, np.where(y_train > 0, 1, 0))
+    assert os.path.exists('./test_tmp_folder')
+
+    pred_mean, pred_std = model.predict(X_test.reset_index(drop=True), return_std=True, verbosity=1, n_jobs=1)