fix: pin BLAS to one thread in PWM kernels to stop thread explosion

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: 

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.
 * Fix: `compute_pwm()` gave a sequence a different score depending on the other
   sequences in the batch. The motif-scan window was capped to the length of the
   *first* sequence for the whole batch, so a longer sequence sitting behind a

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)
         # One flat bin spanning the longest sequence, so every sequence is
@@ -118,6 +121,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]])

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")

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(
         {