R/parallel_threads_.R

In qcluster: Clustering via Quadratic Scoring

.qcluster_thread_env_vars <- function() {
  c(
    "OMP_NUM_THREADS",
    "OPENBLAS_NUM_THREADS",
    "MKL_NUM_THREADS",
    "BLIS_NUM_THREADS",
    "GOTO_NUM_THREADS",
    "VECLIB_MAXIMUM_THREADS",
    "OMP_THREAD_LIMIT"
  )
}

.qcluster_restore_thread_env <- function(old_env) {
  unset <- is.na(old_env)
  if (any(unset)) {
    Sys.unsetenv(names(old_env)[unset])
  }
  if (any(!unset)) {
    do.call(Sys.setenv, as.list(old_env[!unset]))
  }

  invisible(NULL)
}

.qcluster_set_single_thread_env <- function() {
  vars <- .qcluster_thread_env_vars()
  old_env <- Sys.getenv(vars, unset = NA_character_)

  values <- rep.int("1", length(vars))
  names(values) <- vars
  do.call(Sys.setenv, as.list(values))

  old_env
}

.qcluster_configure_single_thread_workers <- function(clst) {
  vars <- .qcluster_thread_env_vars()
  parallel::clusterCall(clst, function(vars) {
    values <- rep.int("1", length(vars))
    names(values) <- vars
    do.call(Sys.setenv, as.list(values))

    if (requireNamespace("RhpcBLASctl", quietly = TRUE)) {
      RhpcBLASctl::blas_set_num_threads(1L)
      if (exists("omp_set_num_threads", envir = asNamespace("RhpcBLASctl"), mode = "function", inherits = FALSE)) {
        RhpcBLASctl::omp_set_num_threads(1L)
      }
    }

    invisible(Sys.getenv(vars, unset = ""))
  }, vars)

  invisible(NULL)
}

.qcluster_make_psock_cluster <- function(ncores) {
  old_env <- .qcluster_set_single_thread_env()
  on.exit(.qcluster_restore_thread_env(old_env), add = TRUE)
  clst <- parallel::makeCluster(ncores, type = "PSOCK")
  .qcluster_configure_single_thread_workers(clst)
  clst
}