Feature Idea: Encapsulate local env and run script on Databricks

[zacdav-db]

1. renv to capture env + deps
2. specify an init if needed
3. upload all contents
4. define job spec (init + libs). first step is to install sys deps + reconstitute renv.
5. submit run job (optional)

[zacdav-db]

Initial test of this idea:

# Run R code on Databricks via one-off Jobs
# Zac Davies (zac@databricks.com)

# ---- brainstorm --------------------------------------------------------------
# 1. capture the existing dependencies and force into an renv lockfile as a
#    a way to reconstitute environment (packages, files)

# 2. workspace file system works for relative imports, this means we can leave
#    code as-is and just force inject an R notebook that reconstitutes env and
#    runs the code via source

# 3. move assets to databricks workspace file system (wsfs)

# 4. if job, trigger a one-off job - objects that need to be kept are captured
#    into a `rds` file and pushed to a volume and downloaded on finish
# ------------------------------------------------------------------------------

get_job_objects <- function(path) {
  
  # download and decode
  base64_obj <- brickster::db_workspace_export(path)
  decoded_raw <- base64enc::base64decode(base64_obj$content)
  
  # write to temp file and read in
  tmp_file <- fs::file_temp()
  writeBin(decoded_raw, tmp_file)
  objects <- readRDS(tmp_file)
  unlink(tmp_file)
  
  objects
}


resolve_lockfile <- function(dir) { 
  lockfile <- fs::path(dir, "renv", ext = "lock")
  
  # use existing lock file or try to create one
  if (fs::file_exists(lockfile)) {
    cli::cli_alert_info("Using existing {.path {lockfile}}")
  } else {
    cli::cli_progress_step(
      msg = "{.header Creating {.pkg renv} snapshot}",
      msg_done = "{.header {.pkg renv} snapshot completed}",
      msg_failed = "{.header Failed to create {.pkg renv} snapshot}",
    )
    deps <- renv::dependencies(path = dir, root = dir, quiet = TRUE)
    env <- renv::snapshot(
      project = dir,
      lockfile = lockfile,
      packages = deps$Package,
      prompt = FALSE
    )
  }
  lockfile
}


brickster_browse_workspace_url <- function(dir) {
  
  host <- brickster::db_host()
  wsid <- brickster::db_current_workspace_id()
  folder_id <- brickster::db_workspace_get_status(dir)$resource_id
  
  glue::glue("https://{host}/browse/folders/{folder_id}?o={wsid}")
  
}

brickster_job_dir <- function(user, uuid) {
  
  dir <- fs::path("/Workspace/Users/", user, "__brickster_jobs", uuid)
  
  # note: checking dir status to create is slower and requires error handling
  # easier to create `__brickster` directory and it'll just no-op if exists
  brickster::db_workspace_mkdirs(dir)
  
  dir
  
}

brickster_runner_notebook <- function(user) {
  fs::path("/Workspace/Users/", user, "__brickster_jobs", "runner")
}

copy_files_to_dbrx <- function(local_files, remote_dir) {
  
  # TODO: check no files are greater than 10MB before uploading
  
  adj_files <- local_files |>
    dplyr::mutate(
      remote_path = stringr::str_replace(
        string = path,
        pattern = fs::path_common(path),
        replacement = remote_dir
      )
    )
  
  # create all required directories
  dirs_needed <- adj_files |> 
    dplyr::filter(type == "directory") |>
    purrr::pluck("remote_path")
  
  purrr::walk(dirs_needed, ~brickster::db_workspace_mkdirs(.x))
  
  # upload files
  adj_files |>  
    dplyr::filter(type == "file") |>
    dplyr::select(path, remote_path) |>
    purrr::pwalk(function(path, remote_path, ...) {
      cli::cli_progress_step(
        msg = "{.header Uploading: }{.path {path}} to {.path {remote_path}}",
        msg_done = "{.header Uploaded: }{.path {path}}",
        msg_failed = "{.header Upload Failed: }{.path {path}}"
      )
      brickster::db_workspace_import(path = remote_path, file = path)
      cli::cli_progress_done()
    })
  
  cli::cli_end()
  NULL
  
}


# specify script to run
# for now, assumes {brickster} is setup correctly for auth in env
# TODO: support specific files being uploaded instead of whole directory
run_on_dbrx <- function(cluster_id,
                        entry_point,
                        dir = fs::path_wd(),
                        # files = NULL,
                        # envvars = list(),
                        return_objects = c(),
                        object_path = NULL,
                        wait_for_completion = TRUE) {
  
  
  # TODO: when specifying `files` param `dir` is ignored
  # TODO: print message when `files` are used
  
  # TODO: build renv snapshot
  # TODO: check that all files referenced are scanned too
  resolve_lockfile(dir = dir)
  
  
  # generate a UUID for the job run
  uuid <- uuid::UUIDgenerate(n = 1, use.time = TRUE)
  
  # who is current user based on credentials?
  user <- brickster::db_current_user()$userName

  job_dir <- brickster_job_dir(user, uuid)
  
  # TODO: wsfs is restrictive compared to a volume for file size (10mb)
  # move to databricks - use path of user by default under `__brickster_jobs/{{uuid}}`
  
  # TODO: handles when files are specified, need table like `dir_info()`
  files <- fs::dir_info(path = dir, recurse = TRUE)[, 1:3]
  
  copy_files_to_dbrx(local_files = files, remote_dir = job_dir)
  
  # print out link to folder where uploads are
  wsfs_url <- brickster_browse_workspace_url(job_dir)
  cli::cli_alert_success(text = "Files uploaded to: {.url {wsfs_url}}")
  
  # trigger one off job of the notebook with parameters
  # define a job task
  task <- brickster::job_task(
    task_key = "brickster_job_execution",
    description = glue::glue("Running {job_dir}"),
    existing_cluster_id = cluster_id,
    task = brickster::notebook_task(
      notebook_path = brickster_runner_notebook(user = user),
      base_parameters = list(
        return_objects = paste0(return_objects, collapse = ";"),
        job_dir = job_dir,
        job_entry_point = entry_point
      )
    )
  )
  
  job_run <- brickster::db_jobs_runs_submit(
    tasks = brickster::job_tasks(task),
    run_name = uuid
  )
  
  # TODO: optionally poll for completion and return objects if requsted
  
  # return defined objects
  job_objs <- fs::path(job_dir, "__returned_objects.rds")
  
  list(job_id = job_run$job_id, job_objs_path = job_objs)
  
}

results <- run_on_dbrx(
  entry_point = "analysis.R",
  dir = "/Users/zachary.davies/Documents/projects/test-project-job-dbx-ft",
  cluster_id = "1017-210529-ij03ow0p",
  return_objects = c("files", "plot")
)

get_job_objects(results$job_objs_path)