Fix cluster load-balancing hang when n_tasks > n_nodes (#591)

Author: shikokuchuo

NEWS.md

@@ -5,6 +5,7 @@
 * `http_config()` gains a `headers` argument, now the primary way to supply HTTP headers (including authentication such as session cookie, bearer token, or API key). `cookie` and `token` are retained as convenience arguments that append `Cookie:` and `Authorization: Bearer` entries to `headers` (thanks @ddl-dkelkhoff, #612).
 * Reduces overhead for synchronous daemons by using an in-process transport.
 * Fixes `.handleSimpleError()` appearing in `$stack.trace` on a `miraiError` (regression in mirai 2.7.0).
+* Fixes load-balanced parallel functions (e.g. `parLapplyLB()`, `foreach::%dopar%`) hanging on a mirai cluster when there are more tasks than nodes, a regression in mirai 2.5.1 (thanks @manforkr, #591).
 
 # mirai 2.7.0
 

R/parallel.R

@@ -89,6 +89,8 @@ make_cluster <- function(n, url = NULL, remote = NULL, ...) {
     daemons(n, dispatcher = FALSE, ..., cleanup = FALSE, .compute = id)
   }
 
+  `[[<-`(..[[id]], "cvs", cv())
+
   cl <- lapply(seq_len(n), create_node, id = id)
   `attributes<-`(cl, list(class = c("miraiCluster", "cluster"), id = id))
 }
@@ -121,7 +123,8 @@ sendData.miraiNode <- function(node, data) {
   value <- data[["data"]]
   tagged <- !is.null(value[["tag"]])
   if (tagged) {
-    cv_reset(envir[["cv"]])
+    orig_cv <- envir[["cv"]]
+    `[[<-`(envir, "cv", envir[["cvs"]])
   }
 
   m <- mirai(
@@ -131,6 +134,7 @@ sendData.miraiNode <- function(node, data) {
     .compute = id
   )
   if (tagged) {
+    `[[<-`(envir, "cv", orig_cv)
     `[[<-`(m, "tag", value[["tag"]])
   }
   `[[<-`(node, "mirai", m)
@@ -143,7 +147,7 @@ recvData.miraiNode <- function(node) call_aio(.subset2(node, "mirai"))
 #' @exportS3Method parallel::recvOneData
 #'
 recvOneData.miraiCluster <- function(cl) {
-  wait(..[[attr(cl, "id")]][["cv"]])
+  wait(..[[attr(cl, "id")]][["cvs"]])
   node <- which.min(lapply(cl, node_unresolved))
   m <- .subset2(.subset2(cl, node), "mirai")
   list(node = node, value = `class<-`(m, NULL))

tests/tests.R

@@ -285,6 +285,14 @@ connection && {
   test_print(cl <- make_cluster(n = 1, url = local_url(), remote = remote_config()))
   test_null(stopCluster(cl))
 }
+# load-balancing across multiple nodes with more tasks than nodes
+connection && NOT_CRAN && {
+  cl <- make_cluster(2)
+  res <- parLapplyLB(cl, 1:6, function(i) i * 2L)
+  test_identical(res, as.list(seq(2L, 12L, by = 2L)))
+  test_identical(parSapplyLB(cl, 1:6, function(i) i + 1L), 2:7)
+  test_null(stopCluster(cl))
+}
 # advanced daemons and dispatcher tests
 connection && NOT_CRAN && {
   test_true(daemons(url = "ws://:0", correctype = 0L, token = TRUE))