R/sbatch_options.R

In EpiModel/slurmworkflow: Build workflows for the Slurm workload manager

# Helper functions related to the SBATCH options.

#' Generate the "#SBATCH --argument=value" statements from an `sbatch_opts` list
#'
#' @keywords internal
#' @noRd
make_sbatch_lines <- function(sbatch_opts) {
  sbatch_lines <- mapply(
    function(opt, val) paste0("#SBATCH --", opt, "=", val),
    names(sbatch_opts), sbatch_opts
  )
  c("# Generated by the `slurmworkflow` R package", sbatch_lines)
}

#' Generate a string of "--argument=value" statements from an `sbatch_opts` list
#'
#' Necessary for the "next step options"
#'
#' @keywords internal
#' @noRd
make_sbatch_statement <- function(sbatch_opts) {
  sbatch_lines <- mapply(
    function(opt, val) paste0("--", opt, "=", val),
    names(sbatch_opts), sbatch_opts
  )
  sbatch_statement <- paste0(sbatch_lines, collapse = " ")
  sbatch_statement
}

#' @keywords internal
#' @noRd
update_sbatch_opts <- function(cur_opts, update_opts) {
  sbatch_opts <- remove_conflicting_sbatch_opts(cur_opts, update_opts)
  sbatch_opts[names(update_opts)] <- update_opts
  validate_sbatch_opts(sbatch_opts)
}

#' Remove from `cur_opts` the options that woud conflict with the ones in
#' `update_opts`. (e.g. "mem" in `cur_opts` and `mem-per-cpu` in `update_opts`)
#'
#' @keywords internal
#' @noRd
remove_conflicting_sbatch_opts <- function(cur_opts, update_opts) {
  update_opts <- clean_sbatch_opts(update_opts)
  conflicting_opts <- Filter(
    function(x) any(x %in% names(update_opts)),
    get_exclusive_sbatch_opts()
  )
  cur_opts[unlist(conflicting_opts)] <- NULL
  cur_opts
}

#' Check and cleanup the `sbatch_opts` list
#' @keywords internal
#' @noRd
validate_sbatch_opts <- function(sbatch_opts) {
  sbatch_opts <- clean_sbatch_opts(sbatch_opts)
  check_mandatory_sbatch_opts(sbatch_opts)
  check_valid_sbatch_opts(sbatch_opts)
  check_exclusive_sbatch_opts(sbatch_opts)
  sbatch_opts
}

#' @keywords internal
#' @noRd
clean_sbatch_opts <- function(sbatch_opts) {
  sbatch_opts <- Filter(function(x) !is.null(x), sbatch_opts)
  sbatch_opts
}

#' @keywords internal
#' @noRd
check_valid_sbatch_opts <- function(sbatch_opts) {
  incorrect_opts <- ! names(sbatch_opts) %in% get_sbatch_opts_list()
  if (any(incorrect_opts)) {
    stop(
      "The following SBATCH options are not valid: \n    `",
      paste0(names(sbatch_opts)[incorrect_opts], collapse = "`, `"), "`\n",
      "Use only the long form (e.g. `account` and not `A`)."
    )
  }
}

#' Check if all mandatory SBATCH options are set
#'
#' For now, no option is mandatory. It might be the case later.
#' (previously `partition` and account were but it caused problems  with HPC
#' without accounting plugins for slurm)
#'
#' @keywords internal
#' @noRd
check_mandatory_sbatch_opts <- function(sbatch_opts) {
  mandatory_opts <- c()
  missing_opts <- ! mandatory_opts %in% names(sbatch_opts)
  if (any(missing_opts)) {
    stop(
      "The following mandatory SBATCH options are missing:\n    `",
      paste0(mandatory_opts[missing_opts], collapse = "`, `"), "` \n"
    )
  }
}

#' @keywords internal
#' @noRd
check_exclusive_sbatch_opts <- function(sbatch_opts) {
  for (exclusive_opts in get_exclusive_sbatch_opts()) {
    exclusive_pos <- which(exclusive_opts %in% names(sbatch_opts))
    if (length(exclusive_pos) > 1) {
      stop(
        "The following SBATCH options: \n    `",
        paste0(exclusive_opts[exclusive_pos], collapse = "`, `"), "`\n",
        "are mutually exclusive. \n`"
      )
    }
  }
}


#' Return a list of mutually exclusive SBATCH options
#'
#' @keywords internal
#' @noRd
get_exclusive_sbatch_opts <- function() {
  list(
    mem_opts = c("mem", "mem-per-cpu", "mem-per-gpu"),
    thread_opts = c(
      "extra-node-info", "hint", "threads-per-core", "ntasks-per-core")
  )
}

#' List all valid SBATCH options
#'
#' @keywords internal
#' @noRd
get_sbatch_opts_list <- function() {
  c(
    "account",
    "acctg-freq",
    "array",
    "batch",
    "bb",
    "bbf",
    "begin",
    "chdir",
    "cluster-constraint",
    "clusters",
    "comment",
    "constraint",
    "container",
    "contiguous",
    "core-spec",
    "cores-per-socket",
    "cpu-freq",
    "cpus-per-gpu",
    "cpus-per-task",
    "deadline",
    "delay-boot",
    "dependency",
    "distribution",
    "error",
    "exclude",
    "exclusive",
    "export",
    "export-file",
    "extra-node-info",
    "get-user-env",
    "gid",
    "gpu-bind",
    "gpu-freq",
    "gpus",
    "gpus-per-node",
    "gpus-per-socket",
    "gpus-per-task",
    "gres",
    "gres-flags",
    "help",
    "hint",
    "hold",
    "ignore-pbs",
    "input",
    "job-name",
    "kill-on-invalid-dep",
    "licenses",
    "mail-type",
    "mail-user",
    "mcs-label",
    "mem",
    "mem-bind",
    "mem-per-cpu",
    "mem-per-gpu",
    "mincpus",
    "network",
    "nice",
    "no-requeue",
    "nodefile",
    "nodelist",
    "nodes",
    "ntasks",
    "ntasks-per-core",
    "ntasks-per-gpu",
    "ntasks-per-node",
    "ntasks-per-socket",
    "open-mode",
    "output",
    "overcommit",
    "oversubscribe",
    "parsable",
    "partition",
    "power",
    "priority",
    "profile",
    "propagate",
    "qos",
    "quiet",
    "reboot",
    "requeue",
    "reservation",
    "signal",
    "sockets-per-node",
    "spread-job",
    "switches",
    "test-only",
    "thread-spec",
    "threads-per-core",
    "time",
    "time-min",
    "tmp",
    "uid",
    "usage",
    "use-min-nodes",
    "verbose",
    "wait",
    "wait-all-nodes",
    "wckey",
    "wrap"
  )
}