Skip to content

Commit

Permalink
fix with joblib 1.4.2
Browse files Browse the repository at this point in the history
  • Loading branch information
lu-wang-dl committed May 30, 2024
1 parent 60d109e commit a98340c
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions joblibspark/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,13 +23,15 @@
import uuid
from packaging.version import Version, parse

import joblib
from joblib.parallel \
import AutoBatchingMixin, ParallelBackendBase, register_parallel_backend, SequentialBackend

try:
if parse(joblib.__version__) >= Version('1.4.0'):
from joblib._utils import _TracebackCapturingWrapper as SafeFunction
elif parse(joblib.__version__) < Version('1.3.0'):
from joblib._parallel_backends import SafeFunction
except ImportError:
# joblib >= 1.3.0
else:
from joblib._parallel_backends import PoolManagerMixin
SafeFunction = None

Expand Down Expand Up @@ -216,7 +218,8 @@ def mapper_fn(_):
)

return self._get_pool().apply_async(
SafeFunction(run_on_worker_and_fetch_result), callback=callback
SafeFunction(run_on_worker_and_fetch_result),
callback=callback, error_callback=callback
)


Expand Down

0 comments on commit a98340c

Please sign in to comment.