diff --git a/mllam_data_prep/cli.py b/mllam_data_prep/cli.py index 7c356c1..43dd26d 100644 --- a/mllam_data_prep/cli.py +++ b/mllam_data_prep/cli.py @@ -1,4 +1,5 @@ import argparse +import math import os from pathlib import Path @@ -60,8 +61,9 @@ def call(args=None): ) # get the number of system cores n_system_cores = os.cpu_count() - # compute the number of cores to use - n_local_cores = int(args.dask_distributed_local_core_fraction * n_system_cores) + # compute the number of cores to use, using ceil to ensure at least 1 + # worker is used even on single-core machines (e.g. standard GitHub CI) + n_local_cores = math.ceil(args.dask_distributed_local_core_fraction * n_system_cores) # get the total system memory total_memory = psutil.virtual_memory().total # compute the memory per worker diff --git a/tests/test_distributed.py b/tests/test_distributed.py index 727f871..bdecdfa 100644 --- a/tests/test_distributed.py +++ b/tests/test_distributed.py @@ -28,6 +28,7 @@ def distributed(): "args", [ ["example.danra.yaml", "--dask-distributed-local-core-fraction", "1.0"], + ["example.danra.yaml", "--dask-distributed-local-core-fraction", "0.5"], ["example.danra.yaml", "--dask-distributed-local-core-fraction", "0.0"], ["example.danra.yaml"], ],