R/sagemaker.R

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

#' Deploy a vetiver model API to Amazon SageMaker
#'
#' @description Use `vetiver_deploy_sagemaker()` to deploy a [vetiver_model()]
#' that has been versioned and stored via [vetiver_pin_write()] as a Plumber API
#' on [Amazon SageMaker](https://aws.amazon.com/sagemaker/).
#'
#' @inheritParams vetiver_sm_build
#' @inheritParams vetiver_sm_model
#' @inheritParams vetiver_sm_endpoint
#' @param ... Not currently used.
#' @param build_args A list of optional arguments passed to
#' [vetiver_sm_build()] such as the model `version` or the `compute_type`.
#' @param endpoint_args A list of optional arguments passed to
#' [vetiver_sm_endpoint()] such as `accelerator_type` or `data_capture_config`.
#' @param repo_name The name for the AWS ECR repository to store the model.
#'
#' @details
#' This function stores your model deployment image in the same bucket used
#' by `board`.
#'
#' The function `vetiver_deploy_sagemaker()` uses:
#' - [vetiver_sm_build()] to build and push a Docker image to ECR,
#' - [vetiver_sm_model()] to create a SageMaker model, and
#' - [vetiver_sm_endpoint()] to deploy a SageMaker model endpoint.
#'
#' These modular functions are available for more advanced use cases.
#' 
#' If you are working locally, you will likely need to explicitly set up
#' your execution role to work correctly. Check out 
#' ["Execution role requirements"](https://dyfanjones.r-universe.dev/smdocker)
#' in the smdocker documentation, and especially note that the bucket containing 
#' your vetiver model needs to be added as a resource in your IAM role policy.
#'
#' @return
#' The deployed [vetiver_endpoint_sagemaker()].
#'
#' @seealso [vetiver_sm_build()], [vetiver_sm_model()], [vetiver_sm_endpoint()]
#' @export
#' @examples
#' if (FALSE) {
#' library(pins)
#' b <- board_s3(bucket = "my-existing-bucket")
#' cars_lm <- lm(mpg ~ ., data = mtcars)
#' v <- vetiver_model(cars_lm, "cars_linear")
#' vetiver_pin_write(b, v)
#'
#' endpoint <- vetiver_deploy_sagemaker(
#'     board = b,
#'     name = "cars_linear",
#'     instance_type = "ml.t2.medium",
#'     predict_args = list(type = "class", debug = TRUE)
#' )
#' }
#'
vetiver_deploy_sagemaker <- function(board,
                                     name,
                                     instance_type,
                                     ...,
                                     predict_args = list(),
                                     docker_args = list(),
                                     build_args = list(),
                                     endpoint_args = list(),
                                     repo_name = glue("vetiver-sagemaker-{name}")) {

    check_dots_empty()
    if (!inherits(board, "pins_board_s3")) {
        stop_input_type(board, "an S3 pins board")
    }

    repo_name <- ifelse(
        grepl(":", repo_name),
        repo_name,
        glue("{repo_name}:{strftime(Sys.time(), '%Y-%m-%d')}")
    )

    # build image and push to aws ecr
    build_args <- compact(c(
        list(board = board,
             name = name,
             predict_args = predict_args,
             docker_args = docker_args,
             repository = repo_name,
             bucket = board$bucket),
        build_args
    ))
    image_uri <- do.call(vetiver_sm_build, build_args)

    tags <- sm_check_tags(endpoint_args$tags)
    tags <- list_modify(
        tags,
        "vetiver:pin_board" = glue("s3://{board$bucket}/{board$prefix %||% ''}"),
        "vetiver:r-ver" = getRversion()
    )
    endpoint_args$tags <- NULL

    # create sagemaker model
    model_name <- vetiver_sm_model(image_uri, tags = tags)

    # create sagemaker endpoint
    endpoint_args <- compact(c(
        list(
            model_name = model_name,
            instance_type = instance_type,
            tags = tags),
        endpoint_args
    ))
    endpoint <- do.call(vetiver_sm_endpoint, endpoint_args)
    return(endpoint)
}

#' Deploy a vetiver model API to Amazon SageMaker with modular functions
#'
#' @description
#' Use the function [vetiver_deploy_sagemaker()] for basic deployment on
#' [SageMaker](https://aws.amazon.com/sagemaker/), or these three functions 
#' together for more advanced use cases:
#' - `vetiver_sm_build()` generates and builds a Docker image on SageMaker for
#' a vetiver model
#' - `vetiver_sm_model()` creates an Amazon SageMaker model
#' - `vetiver_sm_endpoint()` deploys an Amazon SageMaker model endpoint
#'
#' @inheritParams vetiver_prepare_docker
#' @param board An AWS S3 board created with [pins::board_s3()]. This board
#' must be in the correct region for your SageMaker instance.
#' @param repository The ECR repository and tag for the image as a character.
#' Defaults to `sagemaker-studio-${domain_id}:latest`.
#' @param compute_type The [CodeBuild](https://aws.amazon.com/codebuild/)
#' compute type as a character. Defaults to `BUILD_GENERAL1_SMALL`.
#' @param role The ARN IAM role name (as a character) to be used with:
#' - CodeBuild for `vetiver_sm_build()`
#' - the SageMaker model for `vetiver_sm_model()`
#'
#' Defaults to the SageMaker Studio execution role.
#' @param bucket The S3 bucket to use for sending data to CodeBuild as a
#' character. Defaults to the SageMaker SDK default bucket.
#' @param vpc_id ID of the VPC that will host the CodeBuild project such as
#' `"vpc-05c09f91d48831c8c"`.
#' @param subnet_ids List of subnet IDs for the CodeBuild project, such as
#' `list("subnet-0b31f1863e9d31a67")`.
#' @param security_group_ids List of security group IDs for the CodeBuild
#' project, such as `list("sg-0ce4ec0d0414d2ddc")`.
#' @param log A logical to show the logs of the running CodeBuild build.
#' Defaults to `TRUE`.
#' @param ... [Docker build parameters](https://docs.docker.com/engine/reference/commandline/build/#options>)
#' (Use "_" instead of "-"; for example, Docker optional parameter
#' `build-arg` becomes `build_arg`).
#'
#' @details The function `vetiver_sm_build()` generates the files necessary to
#' build a Docker container to deploy a vetiver model in SageMaker and then
#' builds the image on [AWS CodeBuild](https://aws.amazon.com/codebuild/). The
#' resulting image is stored in [AWS ECR](https://aws.amazon.com/ecr/).
#' This function creates a Plumber file and Dockerfile appropriate for
#' SageMaker, for example, with `path = "/invocations"` and `port = 8080`.
#'
#' If you run into problems with Docker rate limits, then either
#' - authenticate to Docker from SageMaker, or
#' - use a [public ECR base image](https://gallery.ecr.aws/docker/library/r-base),
#' passed through `docker_args`
#'
#' @seealso [vetiver_prepare_docker()], [vetiver_deploy_sagemaker()], [vetiver_endpoint_sagemaker()]
#' @examples
#' if (FALSE) {
#' library(pins)
#' b <- board_s3(bucket = "my-existing-bucket")
#' cars_lm <- lm(mpg ~ ., data = mtcars)
#' v <- vetiver_model(cars_lm, "cars_linear")
#' vetiver_pin_write(b, v)
#'
#' new_image_uri <- vetiver_sm_build(
#'     board = b,
#'     name = "cars_linear",
#'     predict_args = list(type = "class", debug = TRUE),
#'     docker_args = list(
#'         base_image = "FROM public.ecr.aws/docker/library/r-base:4.2.2"
#'     )
#' )
#'
#' model_name <- vetiver_sm_model(
#'     new_image_uri,
#'     tags = list("my_custom_tag" = "fuel_efficiency")
#' )
#'
#' vetiver_sm_endpoint(model_name, "ml.t2.medium")
#' }
#' @return `vetiver_sm_build()` returns the AWS ECR image URI and
#' `vetiver_sm_model()` returns the model name (both as characters).
#' `vetiver_sm_endpoint()` returns a new [vetiver_endpoint_sagemaker()] object.
#' @export
vetiver_sm_build <- function(board,
                             name,
                             version = NULL,
                             path = fs::dir_create(tempdir(), "vetiver"),
                             predict_args = list(),
                             docker_args = list(),
                             repository = NULL,
                             compute_type = c(
                                 "BUILD_GENERAL1_SMALL",
                                 "BUILD_GENERAL1_MEDIUM",
                                 "BUILD_GENERAL1_LARGE",
                                 "BUILD_GENERAL1_2XLARGE"
                             ),
                             role = NULL,
                             bucket = NULL,
                             vpc_id = NULL,
                             subnet_ids = list(),
                             security_group_ids = list(),
                             log = TRUE,
                             ...) {
    check_installed("smdocker")
    check_dots_used()
    compute_type <- arg_match(compute_type)
    docker_args <- list_modify(docker_args, port = 8080)
    predict_args <- list_modify(predict_args, path = "/invocations")

    vetiver_prepare_docker(
        board = board,
        name = name,
        version = version,
        path = path,
        predict_args = predict_args,
        docker_args = docker_args
    )

    image_uri <- smdocker::sm_build(
        repository = repository,
        compute_type = compute_type,
        role = role,
        dir = path,
        bucket = bucket,
        vpc_id = vpc_id,
        subnet_ids = subnet_ids,
        security_group_ids = security_group_ids,
        log = log,
        ...
    )

    return(image_uri)
}

#' @param image_uri The AWS ECR image URI for the Amazon SageMaker Model to be
#' created (for example, as returned by [vetiver_sm_build()]).
#' @param model_name The Amazon SageMaker model name to be deployed.
#' @param vpc_config A list containing the VPC configuration for the Amazon
#' SageMaker model [API VpcConfig](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_VpcConfig.html)
#' (optional).
#' * `Subnets`: List of subnet ids
#' * `SecurityGroupIds`: List of security group ids
#' @param enable_network_isolation A logical to specify whether the container
#' will run in network isolation mode. Defaults to `FALSE`.
#' @param tags A named list of tags for labeling the Amazon SageMaker model or
#' model endpint to be created.
#' @rdname vetiver_sm_build
#' @export
vetiver_sm_model <- function(image_uri,
                             model_name,
                             role = NULL,
                             vpc_config = list(),
                             enable_network_isolation = FALSE,
                             tags = list()) {
    check_installed(c("smdocker", "paws.machine.learning"))
    config <- smdocker::smdocker_config()
    sagemaker_client <- paws.machine.learning::sagemaker(config)

    if (is_missing(model_name)) {
        model_name <- base_name_from_image(image_uri)
    }

    if (is.null(role)) {
        role <- smdocker::sagemaker_get_execution_role()
    }
    tags <- sm_check_tags(tags)
    tags <- sm_format_tags(tags)

    request <- list(
        "ModelName" = model_name,
        "ExecutionRoleArn" = role,
        "PrimaryContainer" = list("Image" = image_uri)
    )
    request$Tags <- .sm_append_project_tags(tags)
    request$VpcConfig <- sm_check_vpc_config(vpc_config)
    if (isTRUE(enable_network_isolation)) {
        request$EnableNetworkIsolation <- TRUE
    }
    # create model
    do.call(sagemaker_client$create_model, request)

    return(model_name)
}

#' @param instance_type Type of EC2 instance to use; see
#' [Amazon SageMaker pricing](https://aws.amazon.com/sagemaker/pricing/).
#' @param endpoint_name The name to use for the Amazon SageMaker model endpoint
#' to be created, if to be different from `model_name`.
#' @param initial_instance_count The initial number of instances to run
#' in the endpoint.
#' @param accelerator_type Type of Elastic Inference accelerator to
#' attach to an endpoint for model loading and inference, for
#' example, `"ml.eia1.medium"`.
#' @param kms_key The ARN of the KMS key used to encrypt the data on the
#' storage volume attached to the instance hosting the endpoint.
#' @param data_capture_config A list for configuration to control how Amazon
#' SageMaker captures inference data.
#' @param volume_size The size, in GB, of the ML storage volume attached to
#' the individual inference instance associated with the production variant.
#' Currently only Amazon EBS gp2 storage volumes are supported.
#' @param model_data_download_timeout The timeout value, in seconds, to download
#' and extract model data from Amazon S3.
#' @param wait A logical for whether to wait for the endpoint to be deployed.
#' Defaults to `TRUE`.
#' @rdname vetiver_sm_build
#' @export
vetiver_sm_endpoint <- function(model_name,
                                instance_type,
                                endpoint_name = NULL,
                                initial_instance_count = 1,
                                accelerator_type = NULL,
                                tags = list(),
                                kms_key = NULL,
                                data_capture_config = list(),
                                volume_size = NULL,
                                model_data_download_timeout = NULL,
                                wait = TRUE) {
    check_installed(c("smdocker", "paws.machine.learning"))

    config <- smdocker::smdocker_config()
    sagemaker_client <- paws.machine.learning::sagemaker(config)

    endpoint_name <- endpoint_name %||% model_name

    tags <- sm_check_tags(tags)
    tags <- sm_format_tags(tags)

    request <- sm_req_endpoint_config(
        model_name,
        endpoint_name,
        instance_type,
        initial_instance_count,
        accelerator_type,
        volume_size,
        model_data_download_timeout,
        .sm_append_project_tags(tags),
        kms_key,
        data_capture_config
    )

    # create endpoint config
    resp <- do.call(
        sagemaker_client$create_endpoint_config,
        request
    )

    # create endpoint
    endpoint_name <- sm_create_endpoint(
        sagemaker_client, model_name, endpoint_name, tags, wait
    )

    return(vetiver_endpoint_sagemaker(model_name))
}

#' Delete Amazon SageMaker model, endpoint, and endpoint configuration
#'
#' Use this function to delete the Amazon SageMaker components used in a
#' [vetiver_endpoint_sagemaker()] object. This function does _not_ delete
#' any pinned model object in S3.
#'
#' @param object The model API endpoint object to be deleted, created with
#' [vetiver_endpoint_sagemaker()].
#' @param delete_model Delete the SageMaker model? Defaults to `TRUE`.
#' @param delete_endpoint Delete both the endpoint and endpoint configuration?
#' Defaults to `TRUE`.
#' @return `TRUE`, invisibly
#' @seealso [vetiver_deploy_sagemaker()], [vetiver_sm_build()], [vetiver_endpoint_sagemaker()]
#' @export
vetiver_sm_delete <- function(object, delete_model = TRUE, delete_endpoint = TRUE) {
    check_installed(c("smdocker", "paws.machine.learning"))

    config <- smdocker::smdocker_config()
    sagemaker_client <- paws.machine.learning::sagemaker(config)

    endpoint_name <- object$model_endpoint

    if (is_true(delete_endpoint)) {
        tryCatch(
            {
                endpoint_describe <- sagemaker_client$describe_endpoint(
                    EndpointName = endpoint_name
                )
                endpoint_config_name <- endpoint_describe$EndpointConfigName
                sagemaker_client$delete_endpoint_config(endpoint_config_name)
            },
            error = function(err) {
                cli::cli_warn("Unable to delete {.val {endpoint_name}} endpoint configuration.")
            }
        )
        sagemaker_client$delete_endpoint(endpoint_name)
    }
    if (is_true(delete_model)) {
        sagemaker_client$delete_model(endpoint_name)
    }
    return(invisible(TRUE))
}

#' Post new data to a deployed SageMaker model endpoint and return predictions
#'
#' @param object A SageMaker model endpoint object created with [vetiver_endpoint_sagemaker()].
#' @param new_data New data for making predictions, such as a data frame.
#' @param ... Extra arguments passed to [paws.machine.learning::sagemakerruntime_invoke_endpoint()]
#'
#' @return A tibble of model predictions with as many rows as in `new_data`.
#' @seealso [augment.vetiver_endpoint_sagemaker()]
#' @export
#' @examples
#' if (FALSE) {
#'   endpoint <- vetiver_endpoint_sagemaker("sagemaker-demo-model")
#'   predict(endpoint, mtcars[4:7, -1])
#' }
#' @export
predict.vetiver_endpoint_sagemaker <- function(object, new_data, ...) {
    check_installed(c("jsonlite", "smdocker", "paws.machine.learning"))
    data_json <- jsonlite::toJSON(new_data, na = "string")
    config <- smdocker::smdocker_config()
    sm_runtime <- paws.machine.learning::sagemakerruntime(config)
    tryCatch(
        {
            resp <- sm_runtime$invoke_endpoint(object$model_endpoint, data_json, ...)
            resp <- resp$Body
        },
        error = function(error) {
            error_code <- error$error_response$ErrorCode
            if (!is.null(error_code) && error_code == "NO_SUCH_ENDPOINT") {
                cli::cli_abort("Model endpoint {.val {object$model_endpoint}} not found.")
            }
            stop(error)
        }
    )
    con <- rawConnection(resp)
    on.exit(close(con))
    resp <- jsonlite::fromJSON(con)
    return(tibble::as_tibble(resp))
}

#' Post new data to a deployed SageMaker model endpoint and augment with predictions
#'
#' @param x A SageMaker model endpoint object created with [vetiver_endpoint_sagemaker()].
#' @inheritParams predict.vetiver_endpoint_sagemaker
#'
#' @return The `new_data` with added prediction column(s).
#' @seealso [predict.vetiver_endpoint_sagemaker()]
#' @export
#' @examples
#'
#' if (FALSE) {
#'   endpoint <- vetiver_endpoint_sagemaker("sagemaker-demo-model")
#'   augment(endpoint, mtcars[4:7, -1])
#' }
#'
augment.vetiver_endpoint_sagemaker <- function(x, new_data, ...) {
    preds <- predict(x, new_data = new_data, ...)
    vctrs::vec_cbind(tibble::as_tibble(new_data), preds)
}


#' Create a SageMaker model API endpoint object for prediction
#'
#' This function creates a model API endpoint for prediction from a Sagemaker Model.
#' No HTTP calls are made until you actually
#' [`predict()`][predict.vetiver_endpoint_sagemaker()] with your endpoint.
#'
#' @param model_endpoint The name of the Amazon SageMaker model endpoint.
#' @return A new `vetiver_endpoint_sagemaker` object
#'
#' @examplesIf rlang::is_installed("smdocker")
#' vetiver_endpoint_sagemaker("vetiver-sagemaker-demo-model")
#' @export
vetiver_endpoint_sagemaker <- function(model_endpoint) {
    check_installed("smdocker")
    config <- smdocker::smdocker_config()
    check_character(model_endpoint)
    check_character(config$region)
    new_vetiver_endpoint_sagemaker(model_endpoint, config$region)
}

new_vetiver_endpoint_sagemaker <- function(model_endpoint = character(),
                                           region = character()) {
    structure(
        list(model_endpoint = model_endpoint, region = region),
        class = "vetiver_endpoint_sagemaker"
    )
}

#' @export
format.vetiver_endpoint_sagemaker <- function(x, ...) {
    cli::cli_format_method({
        cli::cli_h3("A SageMaker model endpoint for prediction:")
        cli::cli_text("Model endpoint: {x$model_endpoint}")
        cli::cli_text("Region: {x$region}")
    })
}

#' @export
print.vetiver_endpoint_sagemaker <- function(x, ...) {
    cat(format(x), sep = "\n")
    invisible(x)
}