add spark session

joblib · May 30, 2024 · 1e311e8 · 1e311e8
1 parent 668f062
commit 1e311e8
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/joblibspark/backend.py b/joblibspark/backend.py
@@ -40,6 +40,7 @@
 from pyspark import cloudpickle
 from pyspark.util import VersionUtils
 
+from .utils import get_spark_session
 
 def register():
     """
@@ -72,10 +73,7 @@ def __init__(self, **backend_args):
         super(SparkDistributedBackend, self).__init__(**backend_args)
         self._pool = None
         self._n_jobs = None
-        self._spark = SparkSession \
-            .builder \
-            .appName("JoblibSparkBackend") \
-            .getOrCreate()
+        self._spark = get_spark_session()
         self._spark_context = self._spark.sparkContext
         self._job_group = "joblib-spark-job-group-" + str(uuid.uuid4())
         self._spark_pinned_threads_enabled = isinstance(