batchtools_slurm(): Use 'sacct --jobs=batch_id' to find queued jobs that not yet show up with 'squeue'

Author: HenrikBengtsson

DESCRIPTION

@@ -1,5 +1,5 @@
 Package: future.batchtools
-Version: 0.20.0-9004
+Version: 0.20.0-9005
 Depends:
     R (>= 3.2.0),
     parallelly,

NAMESPACE

@@ -60,6 +60,7 @@ export(loggedError)
 export(loggedOutput)
 export(makeClusterFunctionsBash)
 importFrom(batchtools,Worker)
+importFrom(batchtools,assertRegistry)
 importFrom(batchtools,batchExport)
 importFrom(batchtools,batchMap)
 importFrom(batchtools,cfBrewTemplate)
@@ -84,6 +85,7 @@ importFrom(batchtools,makeClusterFunctionsTORQUE)
 importFrom(batchtools,makeRegistry)
 importFrom(batchtools,makeSubmitJobResult)
 importFrom(batchtools,removeRegistry)
+importFrom(batchtools,runOSCommand)
 importFrom(batchtools,saveRegistry)
 importFrom(batchtools,setJobNames)
 importFrom(batchtools,submitJobs)

R/BatchtoolsFutureBackend-class.R

@@ -535,6 +535,14 @@ status <- function(future, ...) {
 
   jobid <- config$jobid
   if (is.na(jobid)) return("not submitted")
+
+  ## Optionally filter by the scheduler's job ID, if it exists
+  batch_id <- reg[["status"]][["batch.id"]]
+  ## Pass this to cluster functions listJobsQueued() and listJobsRunning()
+  ## via an R option, because we cannot pass as an argument.
+  options(future.batchtools.batch_id = batch_id)
+  on.exit(options(future.batchtools.batch_id = NULL), add = TRUE)
+
   status <- get_status(reg = reg, ids = jobid)
   status <- (unlist(status) == 1L)
   status <- status[status]

R/BatchtoolsTemplateFutureBackend-class.R

@@ -71,7 +71,7 @@ BatchtoolsTemplateFutureBackend <- function(type, scheduler.latency = 1.0, fs.la
     lsf      = makeClusterFunctionsLSF,
     openlava = makeClusterFunctionsOpenLava,
     sge      = makeClusterFunctionsSGE,
-    slurm    = makeClusterFunctionsSlurm,
+    slurm    = makeClusterFunctionsSlurm2,
     torque   = makeClusterFunctionsTORQUE,
                makeClusterFunctions 
   )

R/batchtools_slurm.R

@@ -151,3 +151,120 @@ slurm_version <- local({
     version
   }
 })
+
+
+# Patch Slurm cluster functions listJobsQueued() and listJobsRunning()
+# to use `sacct` instead of `squeue`
+#' @importFrom batchtools assertRegistry runOSCommand
+#' @importFrom utils tail
+patchClusterFunctionsSlurm <- function(cf) {
+  OSError <- import_from("OSError", package = "batchtools")
+  stopifnot(inherits(cf, "ClusterFunctions"))
+
+  env <- environment(cf[["listJobsQueued"]])
+  array.jobs <- env[["array.jobs"]]
+  getClusters <- env[["getClusters"]]
+  nodename <- env[["nodename"]]
+  
+  listJobs <- function(reg, args) {      
+    assertRegistry(reg, writeable = FALSE)
+    args <- c(args, "--user=$USER", "--noheader", "--parsable2", "--allocations", "--format=JobID")
+    clusters <- getClusters(reg)
+    if (length(clusters) > 0) {
+       args <- c(args, sprintf("--clusters=%s", clusters))
+    }
+    res <- runOSCommand("sacct", args, nodename = nodename)
+    if (res$exit.code > 0L) {
+      OSError("Listing of jobs failed", res)
+    }
+    if (length(clusters) > 0) {
+      res <- tail(res$output, -1L)
+    } else {
+      res <- res$output
+    }
+    res
+  } ## listJobs()
+  
+  cf$listJobsQueued <- function(reg) {
+    ## List PENDING (PD) and REQUEUED (RQ) jobs
+    listJobs(reg, "--state=PD,RQ")
+  }                                                          
+  
+  cf$listJobsRunning <- function(reg) {
+    ## List RUNNING (R), SUSPENDED (S), RESIZING (RS) jobs
+    listJobs(reg, "--state=R,S,RS")
+  }
+
+  cf
+} ## patchClusterFunctionsSlurm()
+
+
+
+# Patch Slurm cluster functions listJobsQueued() and listJobsRunning()
+# to use `sacct` instead of `squeue`
+#' @importFrom batchtools assertRegistry runOSCommand
+#' @importFrom utils tail
+patchClusterFunctionsSlurm2 <- function(cf) {
+  OSError <- import_from("OSError", package = "batchtools")
+  stopifnot(inherits(cf, "ClusterFunctions"))
+
+  env <- environment(cf[["listJobsQueued"]])
+  array.jobs <- env[["array.jobs"]]
+  getClusters <- env[["getClusters"]]
+  nodename <- env[["nodename"]]
+  org_listJobsQueued <- env[["listJobsQueued"]]
+  
+  isJobQueued <- function(reg, batch_id) {
+    stopifnot(length(batch_id) == 1L, !is.na(batch_id), nzchar(batch_id))
+    
+    ## FIXME: Add also --starttime=<start time>, because 'sacct' only returns jobs ran today
+    args <- c("--user=$USER", "--noheader", "--parsable2", "--allocations", "--format=State", sprintf("--jobs=%s", batch_id))
+    clusters <- getClusters(reg)
+    if (length(clusters) > 0) {
+       args <- c(args, sprintf("--clusters=%s", clusters))
+    }
+    res <- runOSCommand("sacct", args, nodename = nodename)
+    if (res$exit.code > 0L) {
+      OSError("Failed to check if job is pending", res)
+    }
+    if (length(clusters) > 0) {
+      res <- tail(res$output, -1L)
+    } else {
+      res <- res$output
+    }
+
+    if (length(res) == 0) return(FALSE)
+    
+    res %in% c("PENDING", "REQUEUED")
+  } ## isJobQueued()
+  
+  cf$listJobsQueued <- function(reg) {
+    batch_id <- getOption("future.batchtools.batch_id", NULL)
+    
+    ## Queued jobs according to 'squeue'
+    jobs <- org_listJobsQueued(reg)
+    if (is.null(batch_id)) return(jobs)
+
+    ## Is the job queued?
+    if (length(jobs) > 0) {
+      jobs <- intersect(jobs, as.character(batch_id))
+      if (length(jobs) > 0) return(jobs)
+    }
+
+    ## Ask 'sacct' it if is PENDING or REQUEUED
+    if (isJobQueued(reg, batch_id)) jobs <- as.character(batch_id)
+
+    jobs
+  }
+
+  cf
+} ## patchClusterFunctionsSlurm2()
+
+
+#' @importFrom batchtools makeClusterFunctionsSlurm
+
+makeClusterFunctionsSlurm2 <- function(template = "slurm", array.jobs = TRUE, nodename = "localhost", scheduler.latency = 1, fs.latency = 65, ...) {
+  cf <- makeClusterFunctionsSlurm(template = template, array.jobs = array.jobs, nodename = nodename, scheduler.latency = scheduler.latency, fs.latency =fs.latency, ...)
+  cf <- patchClusterFunctionsSlurm2(cf)
+  cf
+}