feat: add resample stages to tuning callbacks

DESCRIPTION

@@ -50,6 +50,9 @@ Suggests:
     rpart,
     testthat (>= 3.0.0),
     xgboost
+Remotes:
+    mlr-org/mlr3,
+    mlr-org/bbotk@mirai
 VignetteBuilder:
     knitr
 Config/testthat/edition: 3

R/ArchiveAsyncTuning.R

@@ -170,8 +170,7 @@ ArchiveAsyncTuning = R6Class("ArchiveAsyncTuning",
       # cache benchmark result
       if (self$rush$n_finished_tasks > private$.benchmark_result$n_resample_results) {
         bmrs = map(self$finished_data$resample_result, as_benchmark_result)
-        init = BenchmarkResult$new()
-        private$.benchmark_result = Reduce(function(lhs, rhs) lhs$combine(rhs), bmrs, init = init)
+        private$.benchmark_result = Reduce(function(lhs, rhs) lhs$combine(rhs), bmrs)
       }
       private$.benchmark_result
     }

R/CallbackAsyncTuning.R

@@ -1,4 +1,4 @@
-#' @title Create Asynchronous Tuning Callback
+#' @title Asynchronous Tuning Callback
 #'
 #' @description
 #' Specialized [bbotk::CallbackAsync] for asynchronous tuning.
@@ -17,6 +17,26 @@ CallbackAsyncTuning = R6Class("CallbackAsyncTuning",
     #' Called in `ObjectiveTuningAsync$eval()`.
     on_eval_after_xs = NULL,
 
+    #' @field on_resample_begin (`function()`)\cr
+    #' Stage called at the beginning of an evaluation.
+    #' Called in `workhorse()` (internal).
+    on_resample_begin = NULL,
+
+    #' @field on_resample_before_train (`function()`)\cr
+    #' Stage called before training the learner.
+    #' Called in `workhorse()` (internal).
+    on_resample_before_train = NULL,
+
+    #' @field on_resample_before_predict (`function()`)\cr
+    #' Stage called before predicting.
+    #' Called in `workhorse()` (internal).
+    on_resample_before_predict = NULL,
+
+    #' @field on_resample_end (`function()`)\cr
+    #' Stage called at the end of an evaluation.
+    #' Called in `workhorse()` (internal).
+    on_resample_end = NULL,
+
     #' @field on_eval_after_resample (`function()`)\cr
     #' Stage called after hyperparameter configurations are evaluated.
     #' Called in `ObjectiveTuningAsync$eval()`.
@@ -101,6 +121,22 @@ CallbackAsyncTuning = R6Class("CallbackAsyncTuning",
 #'  Called in `ObjectiveTuningAsync$eval()`.
 #'  The functions must have two arguments named `callback` and `context`.
 #'  The argument of `$.eval(xs)` is available in the `context`.
+#' @param on_resample_begin (`function()`)\cr
+#'  Stage called at the beginning of an evaluation.
+#'  Called in `workhorse()` (internal).
+#'  The functions must have two arguments named `callback` and `context`.
+#' @param on_resample_before_train (`function()`)\cr
+#'  Stage called before training the learner.
+#'  Called in `workhorse()` (internal).
+#'  The functions must have two arguments named `callback` and `context`.
+#' @param on_resample_before_predict (`function()`)\cr
+#'  Stage called before predicting.
+#'  Called in `workhorse()` (internal).
+#'  The functions must have two arguments named `callback` and `context`.
+#' @param on_resample_end (`function()`)\cr
+#'  Stage called at the end of an evaluation.
+#'  Called in `workhorse()` (internal).
+#'  The functions must have two arguments named `callback` and `context`.
 #' @param on_eval_after_resample (`function()`)\cr
 #'  Stage called after a hyperparameter configuration is evaluated.
 #'  Called in `ObjectiveTuningAsync$eval()`.
@@ -152,6 +188,10 @@ callback_async_tuning = function(
   on_worker_begin = NULL,
   on_optimizer_before_eval = NULL,
   on_eval_after_xs = NULL,
+  on_resample_begin = NULL,
+  on_resample_before_train = NULL,
+  on_resample_before_predict = NULL,
+  on_resample_end = NULL,
   on_eval_after_resample = NULL,
   on_eval_before_archive = NULL,
   on_optimizer_after_eval = NULL,
@@ -167,6 +207,10 @@ callback_async_tuning = function(
     on_worker_begin,
     on_optimizer_before_eval,
     on_eval_after_xs,
+    on_resample_begin,
+    on_resample_before_train,
+    on_resample_before_predict,
+    on_resample_end,
     on_eval_after_resample,
     on_eval_before_archive,
     on_optimizer_after_eval,
@@ -181,6 +225,10 @@ callback_async_tuning = function(
       "on_worker_begin",
       "on_optimizer_before_eval",
       "on_eval_after_xs",
+      "on_resample_begin",
+      "on_resample_before_train",
+      "on_resample_before_predict",
+      "on_resample_end",
       "on_eval_after_resample",
       "on_eval_before_archive",
       "on_optimizer_after_eval",

R/ObjectiveTuningAsync.R

@@ -31,7 +31,7 @@ ObjectiveTuningAsync = R6Class("ObjectiveTuningAsync",
       lg$debug("Resampling hyperparameter configuration")
 
       # resample hyperparameter configuration
-      private$.resample_result = resample(self$task, self$learner, self$resampling, store_models = self$store_models, allow_hotstart = TRUE, clone = character(0))
+      private$.resample_result = resample(self$task, self$learner, self$resampling, store_models = self$store_models, allow_hotstart = TRUE, clone = character(0), callbacks = self$callbacks)
       call_back("on_eval_after_resample", self$callbacks, self$context)
 
       lg$debug("Aggregating performance")

tests/testthat/test_ArchiveAsyncTuning.R

@@ -12,6 +12,9 @@ test_that("ArchiveAsyncTuning access methods work", {
     terminator = trm("evals", n_evals = 20),
     store_benchmark_result = TRUE
   )
+
+  expect_benchmark_result(instance$archive$benchmark_result)
+
   tuner = tnr("async_random_search")
   tuner$optimize(instance)
 

tests/testthat/test_CallbackAsyncTuning.R

@@ -346,3 +346,142 @@ test_that("on_result in TuningInstanceBatchMultiCrit works", {
   expect_equal(unique(instance$result$classif.ce), 0.7)
 })
 
+# stages in mlr3 workhorse -----------------------------------------------------
+
+test_that("on_resample_begin works", {
+  skip_on_cran()
+  skip_if_not_installed("rush")
+  flush_redis()
+
+  callback = callback_async_tuning("test",
+    on_resample_begin = function(callback, context) {
+      # expect_* does not work
+      assert_task(context$task)
+      assert_learner(context$learner)
+      assert_resampling(context$resampling)
+      checkmate::assert_number(context$iteration)
+      checkmate::assert_null(context$pdatas)
+      context$data_extra = list(success = TRUE)
+    }
+  )
+
+  rush::rush_plan(n_workers = 2)
+  instance = tune(
+    tuner = tnr("async_random_search"),
+    task = tsk("pima"),
+    learner = lrn("classif.rpart", minsplit = to_tune(1, 10)),
+    resampling = rsmp ("holdout"),
+    measures = msr("classif.ce"),
+    term_evals = 2,
+    callbacks = callback)
+
+  expect_class(instance$objective$context, "ContextAsyncTuning")
+
+  walk(as.data.table(instance$archive$benchmark_result)$data_extra, function(data_extra) {
+    expect_true(data_extra$success)
+  })
+})
+
+test_that("on_resample_before_train works", {
+  skip_on_cran()
+  skip_if_not_installed("rush")
+  flush_redis()
+
+  callback = callback_async_tuning("test",
+    on_resample_before_train = function(callback, context) {
+      assert_task(context$task)
+      assert_learner(context$learner)
+      assert_resampling(context$resampling)
+      checkmate::assert_number(context$iteration)
+      checkmate::assert_null(context$pdatas)
+      context$data_extra = list(success = TRUE)
+    }
+  )
+
+  rush::rush_plan(n_workers = 2)
+  instance = tune(
+    tuner = tnr("async_random_search"),
+    task = tsk("pima"),
+    learner = lrn("classif.rpart", minsplit = to_tune(1, 10)),
+    resampling = rsmp ("holdout"),
+    measures = msr("classif.ce"),
+    term_evals = 2,
+    callbacks = callback)
+
+  expect_class(instance$objective$context, "ContextAsyncTuning")
+
+  walk(as.data.table(instance$archive$benchmark_result)$data_extra, function(data_extra) {
+    expect_true(data_extra$success)
+  })
+})
+
+test_that("on_resample_before_predict works", {
+  skip_on_cran()
+  skip_if_not_installed("rush")
+  flush_redis()
+
+  callback = callback_async_tuning("test",
+    on_resample_before_predict = function(callback, context) {
+      assert_task(context$task)
+      assert_learner(context$learner)
+      assert_resampling(context$resampling)
+      checkmate::assert_null(context$pdatas)
+      context$data_extra = list(success = TRUE)
+    }
+  )
+
+  rush::rush_plan(n_workers = 2)
+  instance = tune(
+    tuner = tnr("async_random_search"),
+    task = tsk("pima"),
+    learner = lrn("classif.rpart", minsplit = to_tune(1, 10)),
+    resampling = rsmp ("holdout"),
+    measures = msr("classif.ce"),
+    term_evals = 2,
+    callbacks = callback)
+
+  expect_class(instance$objective$context, "ContextAsyncTuning")
+
+  walk(as.data.table(instance$archive$benchmark_result)$data_extra, function(data_extra) {
+    expect_true(data_extra$success)
+  })
+})
+
+test_that("on_resample_end works", {
+  skip_on_cran()
+  skip_if_not_installed("rush")
+  flush_redis()
+
+  callback = callback_async_tuning("test",
+    on_resample_end = function(callback, context) {
+      # expect_* does not work
+      assert_task(context$task)
+      assert_learner(context$learner)
+      assert_resampling(context$resampling)
+      checkmate::assert_number(context$iteration)
+      checkmate::assert_class(context$pdatas$test, "PredictionData")
+      context$learner$state = mlr3misc::insert_named(context$learner$state, list(state_success = TRUE))
+      context$data_extra = list(success = TRUE)
+    }
+  )
+
+  rush::rush_plan(n_workers = 2)
+  instance = tune(
+    tuner = tnr("async_random_search"),
+    task = tsk("pima"),
+    learner = lrn("classif.rpart", minsplit = to_tune(1, 10)),
+    resampling = rsmp ("holdout"),
+    measures = msr("classif.ce"),
+    term_evals = 2,
+    callbacks = callback)
+
+  expect_class(instance$objective$context, "ContextAsyncTuning")
+
+  walk(as.data.table(instance$archive$benchmark_result)$data_extra, function(data_extra) {
+    expect_true(data_extra$success)
+  })
+
+  walk(instance$archive$benchmark_result$score()$learner, function(learner, ...) {
+    expect_true(learner$state$state_success)
+  })
+})