From d15c42663ae75a6adc660831e1457b80c8fa5e69 Mon Sep 17 00:00:00 2001 From: Young Joon Lee Date: Tue, 4 Jul 2023 15:30:33 +0900 Subject: [PATCH] fix(joblib): optimize test_batcher.py --- tests/hyfi/joblib/test_batcher.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/hyfi/joblib/test_batcher.py b/tests/hyfi/joblib/test_batcher.py index a74d7af2..782a669a 100644 --- a/tests/hyfi/joblib/test_batcher.py +++ b/tests/hyfi/joblib/test_batcher.py @@ -1,22 +1,24 @@ from hyfi.joblib.batch.apply import Apply +from hyfi.joblib.batch.apply_batch import ApplyBatch from hyfi.joblib import JobLibConfig def batcher_test(backend): print(f"Testing backend: {backend}") - joblib = JobLibConfig( - **{"distributed_framework": {"backend": backend}, "num_workers": 2} - ) + joblib = JobLibConfig(backend=backend, num_workers=2, minibatch_size=2) # b = Batcher(minibatch_size=2, backend=backend, procs=2) + joblib.initialize() b = joblib._batcher_instance_ import numpy as np a = Apply(np.power, b, [2], {}) print(a.transform([1, 2, 3, 4])) + a = ApplyBatch(np.power, b, [2], {}) + print(a.transform([1, 2, 3, 4])) def test_bacher_backends(): - backends = ["serial", "threading", "multiprocessing", "ray", "loky", "joblilb"] + backends = ["serial", "threading", "multiprocessing"] # sourcery skip: no-loop-in-tests for backend in backends: batcher_test(backend)