tag nested runs with parent run ID

Signed-off-by: Yitao Li <yitao@rstudio.com>
mlflow · Mar 23, 2021 · 689ec97 · 689ec97
1 parent b2c2f69
commit 689ec97
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/mlflow/R/mlflow/R/tracking-runs.R b/mlflow/R/mlflow/R/tracking-runs.R
@@ -540,7 +540,6 @@ mlflow_record_logged_model <- function(model_spec, run_id = NULL, client = NULL)
 #'
 #' @export
 mlflow_start_run <- function(run_id = NULL, experiment_id = NULL, start_time = NULL, tags = NULL, client = NULL, nested = FALSE) {
-
   # When `client` is provided, this function acts as a wrapper for `runs/create` and does not register
   #  an active run.
   if (!is.null(client)) {
@@ -563,6 +562,13 @@ mlflow_start_run <- function(run_id = NULL, experiment_id = NULL, start_time = N
          call. = FALSE
     )
   }
+  if (nested && !is.null(active_run_id)) {
+    tags <- as.list(tags)
+    # create a tag containing the parent run ID so that MLflow UI can display
+    # nested runs properly
+    tags[["mlflow.parentRunId"]] <- active_run_id
+  }
+
 
   existing_run_id <- run_id %||% {
     env_run_id <- Sys.getenv("MLFLOW_RUN_ID")

diff --git a/mlflow/R/mlflow/tests/testthat/test-tracking-runs.R b/mlflow/R/mlflow/tests/testthat/test-tracking-runs.R
@@ -77,6 +77,9 @@ test_that("mlflow_start_run()/mlflow_end_run() works properly with nested runs",
     expect_equal(mlflow:::mlflow_get_active_run_id(), runs[[i]]$run_uuid)
     run <- mlflow_end_run(client = client, run_id = runs[[i]]$run_uuid)
     expect_identical(run$run_uuid, runs[[i]]$run_uuid)
+    if (i > 1) {
+      expect_equal(run$tags$`mlflow.parentRunId`, runs[[i - 1]]$run_uuid)
+    }
   }
   expect_null(mlflow:::mlflow_get_active_run_id())
 })