Revert "Spike: use vendored loky from joblib."

This reverts commit 618ed20e0f05405f1928ca6727d1dde31fc6ffd8. GitOrigin-RevId: 62c0fcef1a3ea6019ce2aa5185e0187ec26cfa6e
gretelai · Jul 6, 2023 · e71b743 · e71b743
1 parent dbfe117
commit e71b743
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/gretel_synthetics/generate_parallel.py b/src/gretel_synthetics/generate_parallel.py
@@ -10,7 +10,7 @@
 from concurrent import futures
 from typing import List, Optional, Set, Tuple, TYPE_CHECKING, Union
 
-import joblib.externals.loky as loky
+import loky
 
 if TYPE_CHECKING:
     from gretel_synthetics.generate import BaseGenerator, GenText, Settings

diff --git a/tests/test_generate_parallel.py b/tests/test_generate_parallel.py
@@ -24,7 +24,7 @@
         (8, 100, 5, -0.5, 4),
     ],
 )
-@patch("joblib.externals.loky.cpu_count")
+@patch("loky.cpu_count")
 def test_split_work(
     cpu_count, num_cpus, total_lines, chunk_size, parallelism, expected_workers
 ):
@@ -66,7 +66,7 @@ def mock_submit(unused_fn, chunk_size: int, hard_limit: Optional[int] = None):
         for failure_rate in (0.0, 0.25, 0.5)
     ],
 )
-@patch("joblib.externals.loky.ProcessPoolExecutor")
+@patch("loky.ProcessPoolExecutor")
 def test_generate_parallel(pool_mock, num_lines, num_workers, chunk_size, failure_rate):
     pool_instance = pool_mock.return_value
     pool_instance.submit.side_effect = _mock_submitter(failure_rate)
@@ -82,7 +82,7 @@ def test_generate_parallel(pool_mock, num_lines, num_workers, chunk_size, failur
     assert invalid_lines <= ceil(failure_rate * num_lines)
 
 
-@patch("joblib.externals.loky.ProcessPoolExecutor")
+@patch("loky.ProcessPoolExecutor")
 def test_generate_parallel_too_many_invalid(pool_mock):
     pool_instance = pool_mock.return_value
     pool_instance.submit.side_effect = _mock_submitter(0.5)  # 0.5 failure rate