add job monitor for SLURM

DESCRIPTION

@@ -49,7 +49,8 @@ Imports:
   rlang,
   utils,
   vctrs,
-  xml2
+  xml2,
+  yaml
 Suggests:
   knitr (>= 1.30),
   markdown (>= 1.1),

NAMESPACE

@@ -7,6 +7,7 @@ export(crew_class_launcher_sge)
 export(crew_class_launcher_slurm)
 export(crew_class_monitor_cluster)
 export(crew_class_monitor_sge)
+export(crew_class_monitor_slurm)
 export(crew_controller_lsf)
 export(crew_controller_pbs)
 export(crew_controller_sge)
@@ -18,6 +19,7 @@ export(crew_launcher_sge)
 export(crew_launcher_slurm)
 export(crew_monitor_cluster)
 export(crew_monitor_sge)
+export(crew_monitor_slurm)
 importFrom(R6,R6Class)
 importFrom(crew,crew_assert)
 importFrom(crew,crew_class_launcher)

NEWS.md

@@ -1,6 +1,6 @@
 # crew.cluster 0.3.0.9000
 
-
+* Add a "monitor" class for SLURM clusters.
 
 # crew.cluster 0.3.0
 

R/crew_monitor_slurm.R

@@ -0,0 +1,119 @@
+#' @title `r lifecycle::badge("experimental")` Create a SLURM monitor object.
+#' @export
+#' @family slurm
+#' @description Create an `R6` object to monitor SLURM cluster jobs.
+#' @inheritParams crew_monitor_cluster
+crew_monitor_slurm <- function(
+  verbose = TRUE,
+  command_list = as.character(Sys.which("squeue")),
+  command_terminate = as.character(Sys.which("scancel"))
+) {
+  out <- crew_class_monitor_slurm$new(
+    verbose = verbose,
+    command_list = command_list,
+    command_terminate = command_terminate
+  )
+  out$validate()
+  out
+}
+
+#' @title `r lifecycle::badge("experimental")` SLURM monitor class
+#' @export
+#' @family slurm
+#' @description SLURM monitor `R6` class
+#' @details See [crew_monitor_slurm()].
+crew_class_monitor_slurm <- R6::R6Class(
+  classname = "crew_class_monitor_slurm",
+  inherit = crew_class_monitor_cluster,
+  cloneable = FALSE,
+  public = list(
+    #' @description List SLURM jobs.
+    #' 
+    #' This function loads the entire SLURM queue for all users, so it may take
+    #' several seconds to execute. It is intended for interactive use, and
+    #' should especially be avoided in scripts where it is called frequently.
+    #' @return A `tibble` with one row per SLURM job and columns with
+    #'   specific details.
+    #' @param user Character of length 1, user name of the jobs to list.
+    jobs = function(user = ps::ps_username()) {
+      # Cannot be tested with automated tests.
+      # Tested in tests/slurm/monitor.R.
+      # nocov start
+      crew::crew_assert(
+        user,
+        is.character(.),
+        length(.) == 1L,
+        !anyNA(.),
+        nzchar(.),
+        message = "'user' must be `NULL` or a character vector of length 1"
+      )
+      text <- system2(
+        private$.command_list,
+        args = shQuote(c("--yaml")),
+        stdout = TRUE,
+        stderr = if_any(private$.verbose, "", FALSE),
+        wait = TRUE
+      )
+      monitor_cols <- c("job_id", "partition", "name", "user_name", "job_state",
+                        "start_time", "node_count", "state_reason")
+      yaml <- yaml::read_yaml(text = text)
+      out <- map(
+        yaml$jobs,
+        ~ tibble::new_tibble(
+          c(
+            map(.x[monitor_cols], ~ unlist(.x) %||% NA),
+            list(
+              nodes = paste(unlist(.x$job_resources$nodes), collapse = ",") %||% NA
+            )
+          )
+        )
+      )
+      out <- do.call(vctrs::vec_rbind, out)
+      out <- out[out$user_name == user,]
+      out <- out[which(out$job_state != "CANCELLED"),]
+      out$job_id <- as.character(out$job_id)
+      out$start_time <- as.POSIXct(out$start_time, origin = "1970-01-01")
+      out
+      # nocov end
+    },
+    #' @description Terminate one or more SLURM jobs.
+    #' @return `NULL` (invisibly).
+    #' @param jobs Character vector of job names or job IDs to terminate.
+    #'   Ignored if `all` is set to `TRUE`.
+    #' @param all Logical of length 1, whether to terminate all the jobs
+    #'   under your user name. This terminates ALL your SLURM jobs,
+    #'   regardless of whether `crew.cluster` launched them,
+    #'   so use with caution!
+    terminate = function(jobs = NULL, all = FALSE) {
+      # Cannot be tested with automated tests.
+      # Tested in tests/slurm/monitor.R.
+      # nocov start
+      crew::crew_assert(
+        jobs %||% "x",
+        is.character(.),
+        !anyNA(.),
+        nzchar(.),
+        message = paste(
+          "'jobs' must be `NULL` or a character vector of",
+          "valid job names or IDs."
+        )
+      )
+      crew::crew_assert(
+        all,
+        isTRUE(.) || isFALSE(.),
+        message = "'all' must be TRUE or FALSE."
+      )
+      args <- shQuote(if_any(all, c("-u", ps::ps_username()), jobs))
+      stream <- if_any(private$.verbose, "", FALSE)
+      system2(
+        command = private$.command_terminate,
+        args = args,
+        stdout = stream,
+        stderr = stream,
+        wait = TRUE
+      )
+      invisible()
+      # nocov end
+    }
+  )
+)

tests/slurm/monitor.R

@@ -0,0 +1,77 @@
+test_that("SLURM monitor terminate one job at a time", {
+  controller <- crew_controller_slurm(
+    name = "my_workflow",
+    workers = 2L,
+    seconds_idle = 600,
+    script_lines = paste0("module load R/", getRversion()),
+    verbose = TRUE
+  )
+  on.exit(controller$terminate())
+  controller$start()
+  controller$launch(n = 2L)
+  names <- vapply(
+    controller$launcher$workers$handle,
+    function(handle) handle$name,
+    FUN.VALUE = character(1L)
+  )
+  monitor <- crew_monitor_slurm(verbose = TRUE)
+  crew::crew_retry(
+    ~ all(names %in% monitor$jobs()$name),
+    seconds_interval = 1,
+    seconds_timeout = 60,
+    message = "could not submit jobs"
+  )
+  jobs <- monitor$jobs()
+  jobs <- jobs[jobs$name %in% names, ]
+  expect_equal(nrow(jobs), 2L)
+  expect_true(is.character(jobs$job_id))
+  expect_false(anyNA(jobs$job_id))
+  expect_true(all(nzchar(jobs$job_id)))
+  monitor$terminate(jobs = jobs$job_id)
+  crew::crew_retry(
+    ~ !any(jobs$job_id %in% monitor$jobs()$job_id),
+    seconds_interval = 1,
+    seconds_timeout = 60,
+    message = "could not terminate jobs"
+  )
+  expect_false(any(jobs$job_id %in% monitor$jobs()$job_id))
+})
+
+test_that("THIS TEST DELETES ALL USER JOBS! USE WITH CAUTION!", {
+  controller <- crew_controller_slurm(
+    name = "my_workflow",
+    workers = 2L,
+    seconds_idle = 600,
+    script_lines = paste0("module load R/", getRversion()),
+    verbose = TRUE
+  )
+  on.exit(controller$terminate())
+  controller$start()
+  controller$launch(n = 2L)
+  names <- vapply(
+    controller$launcher$workers$handle,
+    function(handle) handle$name,
+    FUN.VALUE = character(1L)
+  )
+  monitor <- crew_monitor_slurm(verbose = TRUE)
+  crew::crew_retry(
+    ~ all(names %in% monitor$jobs()$name),
+    seconds_interval = 1,
+    seconds_timeout = 60,
+    message = "could not submit jobs"
+  )
+  jobs <- monitor$jobs()
+  jobs <- jobs[jobs$name %in% names, ]
+  expect_equal(nrow(jobs), 2L)
+  expect_true(is.character(jobs$job_id))
+  expect_false(anyNA(jobs$job_id))
+  expect_true(all(nzchar(jobs$job_id)))
+  monitor$terminate(all = TRUE)
+  crew::crew_retry(
+    ~ nrow(monitor$jobs()) < 1L,
+    seconds_interval = 1,
+    seconds_timeout = 60,
+    message = "could not terminate jobs"
+  )
+  expect_lt(nrow(monitor$jobs()), 1L)
+})