Skip to content

Commit b6fecdd

Browse files
committed
fix cpu count
1 parent c1d0d05 commit b6fecdd

1 file changed

Lines changed: 17 additions & 1 deletion

File tree

src/cell_eval/_evaluator.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,19 @@
1818
logger = logging.getLogger(__name__)
1919

2020

21+
def _available_cpus() -> int:
22+
"""Return CPUs the current process is allowed to use.
23+
24+
Uses ``os.sched_getaffinity`` on Linux so SLURM/cgroup/taskset limits are
25+
respected; falls back to ``mp.cpu_count`` on macOS/Windows where that API
26+
is unavailable (those platforms typically run locally without cgroup caps).
27+
"""
28+
try:
29+
return len(os.sched_getaffinity(0))
30+
except AttributeError:
31+
return mp.cpu_count()
32+
33+
2134
class MetricsEvaluator:
2235
"""
2336
Evaluates benchmarking metrics of a predicted and real anndata object.
@@ -70,6 +83,9 @@ def __init__(
7083
# Enable a global string cache for categorical columns
7184
pl.enable_string_cache()
7285

86+
if num_threads == -1:
87+
num_threads = _available_cpus()
88+
7389
if os.path.exists(outdir):
7490
logger.warning(
7591
f"Output directory {outdir} already exists, potential overwrite occurring"
@@ -91,7 +107,7 @@ def __init__(
91107
anndata_pair=self.anndata_pair,
92108
de_pred=de_pred,
93109
de_real=de_real,
94-
num_threads=num_threads if num_threads != -1 else mp.cpu_count(),
110+
num_threads=num_threads,
95111
allow_discrete=allow_discrete,
96112
outdir=outdir,
97113
prefix=prefix,

0 commit comments

Comments
 (0)