R/sagemaker-utils.R

In tidymodels/vetiver: Version, Share, Deploy, and Monitor Models

sm_check_vpc_config <- function(vpc_config, call = caller_env()) {
    if (is_empty(vpc_config)) {
        return(NULL)
    }
    subnet <- vpc_config$Subnets
    if (is_empty(subnet) || !inherits(subnet, "list")) {
        stop_input_type(subnet, "a list", call = call)
    }
    security_group_ids <- vpc_config$SecurityGroupIds
    if (is_empty(security_group_ids) || !inherits(security_group_ids, "list")) {
        stop_input_type(security_group_ids, "a list", call = call)
    }
    return(vpc_config)
}

base_name_from_image <- function(image) {
    # NOTE: model name needs to meet regex:
    # ^[a-zA-Z0-9]([\-a-zA-Z0-9]*[a-zA-Z0-9])?
    # Should there be a check or a clean up of name coming from ecr image?
    m <- regexec("^(.+/)?([^:/]+)(:[^:]+)?$", image)
    image_name <- if (!is.null(m)) unlist(regmatches(image, m))[[3]] else image
    now <- Sys.time()
    timestamp <- gsub("\\.", "-", strftime(now, "%Y-%m-%d-%H-%M-%OS3"))
    image_name <- substring(image_name, 1, (63L - nchar(timestamp) - 1))
    return(sprintf("%s-%s", image_name, timestamp))
}

sm_req_endpoint_config <- function(model_name,
                                   endpoint_name,
                                   instance_type,
                                   initial_instance_count,
                                   accelerator_type,
                                   volume_size,
                                   model_data_download_timeout,
                                   tags,
                                   kms_key,
                                   data_capture_config) {

    volume_size <- return_null_or_integer(volume_size)
    model_data_download_timeout <- return_null_or_integer(model_data_download_timeout)

    request <- compact(list(
        EndpointConfigName = endpoint_name,
        ProductionVariants = list(compact(list(
            ModelName = model_name,
            VariantName = "AllTraffic",
            InitialVariantWeight = 1,
            AcceleratorType = accelerator_type,
            InitialInstanceCount = initial_instance_count,
            InstanceType = instance_type,
            VolumeSizeInGB = volume_size,
            ModelDataDownloadTimeoutInSeconds = model_data_download_timeout
        ))),
        Tags = tags,
        KmsKeyId = kms_key,
        DataCaptureConfig = data_capture_config
    ))

    return(request)
}

return_null_or_integer <- function(x) {
    if (is.null(x)) NULL else as.integer(x)
}

# sagemaker helper functions
sm_create_endpoint <- function(client,
                               endpoint_name,
                               config_name,
                               tags = list(),
                               wait = TRUE) {
    cli::cli_inform("Creating endpoint with name {.val {endpoint_name}}")

    client$create_endpoint(
        EndpointName = endpoint_name, EndpointConfigName = config_name, Tags = tags
    )
    if (isTRUE(wait)) {
        sm_wait_for_endpoint(client, endpoint_name)
    }
    return(endpoint_name)
}

# developed from:
# https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/session.py#L3753-L3786
sm_wait_for_endpoint <- function(client, endpoint, poll = 30, call = caller_env()) {
    desc <- sagemaker_deploy_done(client, endpoint)
    while (is_empty(desc)) {
        Sys.sleep(poll)
        desc <- sagemaker_deploy_done(client, endpoint)
    }
    status <- desc$EndpointStatus
    if (status != "InService") {
        reason <- desc$FailureReason
        message <- c(
            "Error hosting endpoint {.val {endpoint}}: {status}",
            "Reason: {reason}"
        )
        if (grepl("CapacityError", as.character(reason))) {
            cli::cli_abort(message, class = "CapacityError", call = call)
        }
        cli::cli_abort(message, class = "UnexpectedStatusException", call = call)
    }
    return(desc)
}

# developed from:
# https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/session.py#L5809-L5828
sagemaker_deploy_done <- function(client, endpoint_name) {
    hosting_status_codes <- list(
        "OutOfService" = "x",
        "Creating" = "-",
        "Updating" = "-",
        "InService" = "!",
        "RollingBack" = "<",
        "Deleting" = "o",
        "Failed" = "*"
    )
    in_progress_statuses <- c("Creating", "Updating")

    desc <- client$describe_endpoint(EndpointName = endpoint_name)
    status <- desc$EndpointStatus

    msg <- hosting_status_codes[[status]]
    if (is.null(msg)) msg <- "?"

    write_to_console(msg, sep = "")

    if (status %in% in_progress_statuses) {
        return(NULL)
    }

    write_to_console("")
    return(desc)
}

# developed from:
# https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/_studio.py#L21-L113
STUDIO_PROJECT_CONFIG <- ".sagemaker-code-config"

sm_check_tags <- function(tags, call = caller_env(), arg = caller_arg(tags)) {
    if (is_empty(tags)) {
        return(list())
    }
    if (!inherits(tags, "list")) {
        stop_input_type(tags, "a list", call = call)
    }
    if (!is_named(tags)) {
        cli::cli_abort('{.arg {arg}} must have valid names, like {.code list("my-tag" = "my-value")}', call = call)
    }
    return(tags)
}

sm_format_tags <- function(tags) {
    tags <- lapply(names(tags), function(n) list(Key = n, Value = tags[[n]]))
    return(tags)
}

.sm_append_project_tags <- function(tags = NULL, working_dir = NULL) {
    path <- .sm_find_config(working_dir)
    if (is.null(path)) {
        return(tags)
    }

    config <- .sm_load_config(path)
    if (is.null(config)) {
        return(tags)
    }

    additional_tags <- .sm_parse_tags(config)
    if (is.null(additional_tags)) {
        return(tags)
    }

    all_tags <- tags %||% list()
    all_tags <- c(all_tags, additional_tags)

    return(all_tags)
}

.sm_find_config <- function(working_dir = NULL) {
    tryCatch(
        {
            wd <- working_dir %||% getwd()
            path <- NULL
            while (is.null(path) && !grepl("/", wd)) {
                candidate <- fs::path(wd, STUDIO_PROJECT_CONFIG)
                if (fs::file_exists(candidate)) {
                    path <- candidate
                }
                wd <- fs::path_dir(candidate)
            }
            return(path)
        },
        error = function(e) {
            return(NULL)
        }
    )
}

.sm_load_config <- function(path) {
    if (!fs::file_exists(path)) {
        return(NULL)
    }
    tryCatch(
        {
            config <- jsonlite::read_json(path)
            return(config)
        },
        error = function(e) {
            return(NULL)
        }
    )
}

.sm_parse_tags <- function(config) {
    if (!is_empty(config$sagemakerProjectId) || !is_empty(config$sagemakerProjectName)) {
        return(list(
            list("Key" = "sagemaker:project-id", "Value" = config$sagemakerProjectId),
            list("Key" = "sagemaker:project-name", "Value" = config$sagemakerProjectName)
        ))
    } else {
        return(NULL)
    }
}

# allow to write to jupyter console while code still executing
# https://github.com/paws-r/paws/pull/602
write_to_console <- function(msg, sep = "\n") {
    on.exit(flush.console())
    writeLines(msg, sep = sep)
}