From 7c8ae2b81f0f000c28698e3ac2d40a3a51e5ea9b Mon Sep 17 00:00:00 2001 From: aviezerl Date: Mon, 22 Jun 2026 16:23:42 +0300 Subject: [PATCH] fix: pin BLAS to one thread in PWM kernels to stop thread explosion compute_pwm(), compute_local_pwm() and calc_seq_pwm() (and thus extract_pwm()) run a per-sequence BLAS dgemm inside an RcppParallel/TBB parallelFor over sequences. With a multi-threaded BLAS (MKL, OpenBLAS) every TBB worker spawns its own BLAS thread team, so a single call opens n_threads * n_blas_threads OS threads - e.g. 64x16 = 1025 at 16 threads on a 128-core node, and thousands at the set_parallel() default. That oversubscribes the machine and, on a cluster job with a thread/process (cgroup pids) limit, makes the call FAIL with a thread-creation error (reported via iceqream::infer_trajectory_motifs -> extract_pwm). The per-sequence dgemm is tiny and the parallelFor over sequences is the right granularity, so the inner BLAS must be serial. New helper local_serial_blas() pins the BLAS to one thread for the duration of these kernels and restores the previous count on exit (the rest of the session keeps its threaded BLAS). Verified on MKL: a 16-thread extract_pwm drops from 1025 to 17 OS threads, results identical, throughput same or slightly better. Adds RhpcBLASctl to Imports (the only portable cross-BLAS thread control). Bumps to 0.0.10. Claude-Session: https://claude.ai/code/session_01KfEW7AGsSzjhMqiY8htdnq --- DESCRIPTION | 7 ++++--- NEWS.md | 12 ++++++++++++ R/pssm-utils.R | 6 ++++++ R/pwm.R | 4 ++++ R/utils.R | 24 ++++++++++++++++++++++++ 5 files changed, 50 insertions(+), 3 deletions(-) diff --git a/DESCRIPTION b/DESCRIPTION index e51a8cd..b814e02 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: prego Title: PWM Regression Optimizer -Version: 0.0.9 +Version: 0.0.10 Authors@R: c( person("Aviezer", "Lifshitz", , "aviezer.lifshitz@weizmann.ac.il", role = c("aut", "cre")), person("Amos", "Tanay", , "amos.tanay@weizmann.ac.il", role = "aut") @@ -30,8 +30,9 @@ Imports: magrittr, methods, RcppParallel, - withr -Suggests: + withr, + RhpcBLASctl +Suggests: testthat (>= 3.0.0), doMC LinkingTo: diff --git a/NEWS.md b/NEWS.md index c940bd9..4651a6a 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,5 +1,17 @@ # prego 0.0.10 +* Fix: PWM scoring no longer opens thousands of threads / fails on core-limited + machines. `compute_pwm()`, `compute_local_pwm()` and `calc_seq_pwm()` (hence + `extract_pwm()`) run a small per-sequence BLAS `dgemm` inside an + RcppParallel/TBB `parallelFor`. With a multi-threaded BLAS (MKL, OpenBLAS) each + TBB worker spawned its own BLAS thread team, so one call opened + `n_threads * n_blas_threads` OS threads (thousands on a many-core node) - which + oversubscribed the machine and could fail outright on a cluster job with a + thread/process (cgroup pids) limit. The inner BLAS is now pinned to a single + thread for the duration of these calls (the TBB loop over sequences is the + parallelism layer), and the previous BLAS thread count is restored afterwards. + Results are unchanged; throughput is the same or better. Adds a dependency on + the `RhpcBLASctl` package. * Added `return_all` parameter to `regress_pwm` (multi-kmer path). When TRUE, returns every candidate-kmer regression (sorted by validation score) instead of just the best one - useful for getting N independent motifs without the residual-rounds approach used by `motif_num > 1`. When `sample_for_kmers = TRUE`, each candidate is refit on the full data. * Improved docs for `regress_pwm` (clarified the three operating modes, fixed `n_motifs`/`comb_modle` typos in the return-value section). diff --git a/R/pssm-utils.R b/R/pssm-utils.R index 654fd56..a96d3ba 100644 --- a/R/pssm-utils.R +++ b/R/pssm-utils.R @@ -24,6 +24,9 @@ #' #' @export compute_pwm <- function(sequences, pssm, spat = NULL, spat_min = 1, spat_max = NULL, bidirect = TRUE, prior = 0.01, func = "logSumExp") { + # Keep the per-sequence BLAS dgemm serial; the TBB parallelFor over + # sequences is the parallelism layer (see local_serial_blas()). + local_serial_blas() if (is.null(spat)) { spat <- data.frame(bin = 0, spat_factor = 1) binsize <- nchar(sequences[[1]]) @@ -114,6 +117,9 @@ compute_pwm <- function(sequences, pssm, spat = NULL, spat_min = 1, spat_max = N #' @inheritParams compute_pwm #' @export compute_local_pwm <- function(sequences, pssm, spat = NULL, spat_min = 0, spat_max = NULL, bidirect = TRUE, prior = 0.01, return_list = FALSE) { + # Keep the per-sequence BLAS dgemm serial; the TBB parallelFor over + # sequences is the parallelism layer (see local_serial_blas()). + local_serial_blas() if (is.null(spat)) { spat <- data.frame(bin = 0, spat_factor = 1) binsize <- nchar(sequences[[1]]) diff --git a/R/pwm.R b/R/pwm.R index c1b7859..d1ad9a1 100644 --- a/R/pwm.R +++ b/R/pwm.R @@ -39,6 +39,10 @@ seqs_to_onehot <- function(seqs) { #' #' @export calc_seq_pwm <- function(sequences, mdb, bidirect = TRUE) { + # Keep the per-sequence BLAS dgemm serial; the TBB parallelFor over + # sequences is the parallelism layer (see local_serial_blas()). + local_serial_blas() + # Input validation if (!is.character(sequences)) { stop("sequences must be a character vector") diff --git a/R/utils.R b/R/utils.R index 704657b..3d55577 100644 --- a/R/utils.R +++ b/R/utils.R @@ -54,6 +54,30 @@ set_parallel <- function(thread_num = max(1, round(parallel::detectCores() * 0.8 invisible(NULL) } +# Pin the BLAS to a single thread for the remainder of the calling function. +# +# prego's PWM kernels (compute_pwm / compute_local_pwm / calc_seq_pwm) run a +# small BLAS `dgemm` per sequence INSIDE an RcppParallel/TBB `parallelFor` over +# sequences. With a multi-threaded BLAS (MKL, OpenBLAS), each TBB worker spawns +# its own BLAS thread team, so a single call opens `n_threads * n_blas_threads` +# OS threads - thousands on a many-core node. That oversubscribes the machine +# and, on a cluster job with a thread/process (cgroup pids) limit, makes the +# call FAIL with a thread-creation error. The per-sequence dgemm is tiny and the +# TBB loop over sequences is the right level of parallelism, so the inner BLAS +# must be serial. Restores the previous BLAS thread count on exit, so the rest +# of the session keeps its threaded BLAS. No-op if RhpcBLASctl is unavailable. +local_serial_blas <- function(.local_envir = parent.frame()) { + if (!requireNamespace("RhpcBLASctl", quietly = TRUE)) { + return(invisible(NULL)) + } + old <- tryCatch(RhpcBLASctl::blas_get_num_procs(), error = function(e) NULL) + RhpcBLASctl::blas_set_num_threads(1) + if (!is.null(old) && old > 1) { + withr::defer(RhpcBLASctl::blas_set_num_threads(old), envir = .local_envir) + } + invisible(NULL) +} + safe_llply <- function(.data, .fun, ..., .parallel = FALSE) { tryCatch( {