From 689ec97bac9e5a327fb12215a364d7078408a70c Mon Sep 17 00:00:00 2001 From: Yitao Li Date: Tue, 23 Mar 2021 20:15:14 +0000 Subject: [PATCH] tag nested runs with parent run ID Signed-off-by: Yitao Li --- mlflow/R/mlflow/R/tracking-runs.R | 8 +++++++- mlflow/R/mlflow/tests/testthat/test-tracking-runs.R | 3 +++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/mlflow/R/mlflow/R/tracking-runs.R b/mlflow/R/mlflow/R/tracking-runs.R index 509bd23c10087..732da4dbd917f 100644 --- 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 index bd744bb90b113..5ffd86b983c09 100644 --- 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()) })