Skip to content

Commit

Permalink
Add support for joblib 1.4.2 (#55)
Browse files Browse the repository at this point in the history
* fix with joblib 1.4.2

* fix lint

* fix workflow
  • Loading branch information
lu-wang-dl authored Jun 5, 2024
1 parent 60d109e commit e716ac4
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 7 deletions.
10 changes: 8 additions & 2 deletions .github/workflows/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,18 @@ jobs:
PIN_MODE: [false, true]
PYSPARK_VERSION: ["3.0.3", "3.1.3", "3.2.3", "3.3.2", "3.4.0"]
include:
- PYSPARK_VERSION: "3.4.0"
- PYSPARK_VERSION: "3.5.1"
PYTHON_VERSION: "3.11"
JOBLIB_VERSION: "1.3.0"
- PYSPARK_VERSION: "3.4.0"
- PYSPARK_VERSION: "3.5.1"
PYTHON_VERSION: "3.11"
JOBLIB_VERSION: "1.4.2"
- PYSPARK_VERSION: "3.5.1"
PYTHON_VERSION: "3.12"
JOBLIB_VERSION: "1.3.0"
- PYSPARK_VERSION: "3.5.1"
PYTHON_VERSION: "3.12"
JOBLIB_VERSION: "1.4.2"
exclude:
- PYSPARK_VERSION: "3.0.3"
PIN_MODE: true
Expand Down
15 changes: 10 additions & 5 deletions joblibspark/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,15 @@
import AutoBatchingMixin, ParallelBackendBase, register_parallel_backend, SequentialBackend

try:
from joblib._parallel_backends import SafeFunction
# joblib >=1.4.0
from joblib._utils import _TracebackCapturingWrapper as SafeFunction
except ImportError:
# joblib >= 1.3.0
from joblib._parallel_backends import PoolManagerMixin
SafeFunction = None
try:
from joblib._parallel_backends import SafeFunction
except ImportError:
# joblib >= 1.3.0
from joblib._parallel_backends import PoolManagerMixin
SafeFunction = None

from py4j.clientserver import ClientServer

Expand Down Expand Up @@ -216,7 +220,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 e716ac4

Please sign in to comment.