Skip to content

Commit f1f6b31

Browse files
committed
Use env var to determine number of workers
1 parent efe6ba9 commit f1f6b31

File tree

1 file changed

+29
-2
lines changed

1 file changed

+29
-2
lines changed

python/lsst/resources/_resourcePath.py

+29-2
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,14 @@
1919
import io
2020
import locale
2121
import logging
22+
import multiprocessing
2223
import os
2324
import posixpath
2425
import re
2526
import shutil
2627
import tempfile
2728
import urllib.parse
29+
from functools import cache
2830
from pathlib import Path, PurePath, PurePosixPath
2931
from random import Random
3032

@@ -59,6 +61,31 @@
5961
MAX_WORKERS = 10
6062

6163

64+
@cache
65+
def _get_num_workers() -> int:
66+
f"""Calculate the number of workers to use.
67+
68+
Returns
69+
-------
70+
num : `int`
71+
The number of workers to use. Will use the value of the
72+
``LSST_RESOURCES_NUM_WORKERS`` environment variable if set. Will fall
73+
back to using the CPU count (plus 2) but capped at {MAX_WORKERS}.
74+
"""
75+
num_workers = -1
76+
env_value = os.getenv("LSST_RESOURCES_NUM_WORKERS")
77+
if env_value is not None:
78+
try:
79+
num_workers = int(env_value)
80+
except TypeError:
81+
pass
82+
if num_workers == -1:
83+
# Look at the processor count and add 2.
84+
num_workers = multiprocessing.cpu_count() + 2
85+
# But don't ever return more than the maximum allowed.
86+
return min([num_workers, MAX_WORKERS])
87+
88+
6289
class ResourcePath: # numpydoc ignore=PR02
6390
"""Convenience wrapper around URI parsers.
6491
@@ -883,7 +910,7 @@ def _mexists(cls, uris: Iterable[ResourcePath]) -> dict[ResourcePath, bool]:
883910
existence : `dict` of [`ResourcePath`, `bool`]
884911
Mapping of original URI to boolean indicating existence.
885912
"""
886-
exists_executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
913+
exists_executor = concurrent.futures.ThreadPoolExecutor(max_workers=_get_num_workers())
887914
future_exists = {exists_executor.submit(uri.exists): uri for uri in uris}
888915

889916
results: dict[ResourcePath, bool] = {}
@@ -926,7 +953,7 @@ def mtransfer(
926953
A dict of all the transfer attempts with a boolean indicating
927954
whether the transfer succeeded for the target URI.
928955
"""
929-
exists_executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
956+
exists_executor = concurrent.futures.ThreadPoolExecutor(max_workers=_get_num_workers())
930957
future_transfers = {
931958
exists_executor.submit(
932959
to_uri.transfer_from,

0 commit comments

Comments
 (0)