ARTIFACTORY: How to Integrate Artifactory with AWS SageMaker - End-to-End Setup, Training, and Deployment Guide

It is possible to utilize the enterprise management of software artifacts with JFrog Artifactory within the environment of Amazon SageMaker, which helps companies build, train, and deploy machine learning (ML) models for any use case with fully managed infrastructure, tools, and workflows.

This knowledge base article includes an end-to-end guide on integrating JFrog Artifactory with AWS SageMaker to build, train, and deploy machine learning models.

It covers:

• Prerequisites & Setup: Repository setup (PyPI, Docker, Hugging Face), scoped tokens, and AWS networking, Secrets Manager, and Lambda authentication.

• Model Training: Writing training scripts, building custom Docker images via Artifactory, running SageMaker jobs, and saving models.

• Inference Deployment: Creating inference containers, deploying Artifactory-hosted models to SageMaker endpoints, and running predictions.

Part 1: Prerequisites and Environment Setup for Artifactory & AWS SageMaker

For most integration scenarios of SageMaker and Artifactory, there is an administrative setup required pertaining to user permissions and policies to access AWS and Artifactory resources as well as network isolation. The following setup and configuration are recommended and/or required prior to beginning any of the subsequent demos or using the example code.

This tutorial assumes you have the necessary administrative access to an Artifactory instance, and an AWS account/region.

Artifactory Setup

Python packages and model hubs like Hugging Face are commonly used in ML application and model development. To resolve dependencies from a public external repository (via proxy) or to store your own binaries produced during development, set up the appropriate repositories in Artifactory.

Create the following in Artifactory:

1. Remote PyPi Repository: to proxy all PyPi packages that will be needed from pypi.org

   Set Up Remote PyPI Repositories

2. Virtual PyPi Repository: to aggregate requests for remote PyPi packages and any locally hosted or proprietary PyPi packages in a local PyPi repository in Artifactory

   Set Up Virtual PyPI Repositories

NOTE:

Although this tutorial will not require the creation of a Local PyPi Repository, it is best practice to reference the Virtual PyPi Repo when configuring pip.

3. Local Hugging Face Repository: to store the custom ML Model version created by the SageMaker training job

   Set Up Local Hugging Face Repositories

4. Remote Docker Repository: to proxy requests for Docker images needed from Docker Hub

   Remote Docker Repositories

5. Local Docker Repository: to store a custom Docker image to be used for training

   Local Docker Repositories

6. Virtual Docker Repository: to aggregate requests for both the Remote and Local Docker Repositories

   Virtual Docker Repositories

7. Scoped Identity Token: to provide limited and scoped permission to access Artifactory from the SageMaker environment

   Generating Scoped Tokens

AWS Environment and SageMaker Setup

For this tutorial and a working demo, your AWS cloud environment is expected to include the following:

1. VPC, Private Subnet, Security Group, and VPC endpoints: to restrict internet access, monitor and inspect traffic, and control data flow from SageMaker resources

   • Securing Amazon SageMaker Studio connectivity using a private VPC | AWS Machine Learning Blog

   • Create a VPC - Amazon Virtual Private Cloud

   • Connect your VPC to services using AWS PrivateLink - Amazon Virtual Private Cloud

2. SageMaker Domain/Users

   Get started - Amazon SageMaker

   Note the role name and ARN of the SageMaker execution role used during the configuration of the SageMaker Domain as they will be needed in various places during this tutorial.
   SageMaker Roles
   

3. S3 Bucket (with training data): to store data that will be accessed for ML model training

   Creating a bucket - Amazon Simple Storage Service

   Ensure the SageMaker execution role has permission to access the S3 Bucket. SageMaker Roles

   Download the dataset for training from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz and upload it to the S3 bucket to a directory named “training-data-folder”

   (e.g. s3://<S3 BUCKET NAME>/training-data-folder/)

4. AWS Secret: to centralize and store all of the sensitive data needed for authorization and access to Artifactory and the necessary repositories

   Create an AWS Secrets Manager secret

   A convenient script, deploy-secret.sh that creates an AWS Secret that can be used for all of the demos in this article is available here.

   The following is an explanation for each of the keys that will be needed for this portion of the tutorial:

   username : the username used for authentication with Artifactory

   (Example value: myusername)

   token : the value of the scoped token created for this tutorial (refer to Artifactory setup)

   (Example value: zYzMbggy3zB34s2OeHQlxzuWOuGvo4LFRL7FVpyEtlxykheHQZywhZxZh2L8RIk3)

   host : the url of the Artifactory instance

   (Example value: myartifactory.jfrog.io)

   pypirepo : the name of the virtual PyPi repo in Artifactory

   (Example value: sagemaker-pypi)

   huggingfacerepo : the name of the remote Hugging Face repo in Artifactory

   (Example value: huggingface-remote)

   huggingface-remote-repo : the full url of the Hugging Face remote repo in Artifactory

   (Example value: https://myartifactory.jfrog.io/artifactory/api/huggingfaceml/huggingface-remote)

   huggingface-remote-token : an identity access token for the remote Hugging Face repo in Artifactory

   (Example value: ibFodVFSyWTnAPch2XJNSPHHZatd9rlwHnsONsNY5qk/VyLoU2cxvQVGdyZ58XQO)

   huggingface-local-repo : the full url of the Hugging Face local repo in Artifactory

   (Example value: https://myartifactory.jfrog.io/artifactory/api/huggingfaceml/huggingface-remote

   huggingface-local-token : an identity access token for the local Hugging Face repo in Artifactory

   (Example value: m2JLw7XFg4kA1yMIz1U4n7/gPDKh7znzeXFjXuM1nJX8bnWLPgV5m7tmumuYvhq6)

5. AWS Lambda: to provide access credentials to SageMaker to authenticate to the Docker registry in Artifactory

   Use a Docker registry that requires authentication for training - Amazon SageMaker

   Creating an AWS Lambda for the purpose of authenticating to the Artifactory Docker Registry will require creating related IAM roles and policies. An example script, deploy-lambda.sh that will create all of the necessary resources can be found here in the infrastructure directory along with the lambda function itself in lambda_function.py.

   Once the Lambda function is created, remember to add to the SageMaker execution role the InvokeFunction permission.

   {
               "Action": [
                   "lambda:InvokeFunction"
               ],
               "Effect": "Allow",
               "Resource": [
                   "arn:aws:lambda:${region}:${account_id}:function:${lambda_name}"
               ]
   }

Development Environment

1. Docker: to build the custom training Docker image and push to Artifactory

2. IDE or text editor: to edit the training code and the Dockerfile

3. Git/Github user: to access the example code

4. Artifactory access: to interact with the necessary repositories

5. AWS account access: to interact with AWS resources

6. Python 3, Pip: to run Python scripts and resolve dependencies from Artifactory - Resolve from Artifactory Using Pip

Part 2: How to Build and Train ML Models Using AWS SageMaker and Artifactory

The following instructions will walk through the steps necessary to configure Artifactory and SageMaker in order to build and train an ML model.

A full working example of building and training an ML model using Tensorflow and the SageMaker Python SDK can be found in the train directory here.

The goal of this tutorial is to create, train, and store a model using SageMaker by performing the following steps:

1. Prepare a Python training script

2. Build a Docker container image for training

3. Create and store an ML model in Artifactory with a SageMaker training job

Preparing a Training Script

The code samples for this tutorial show how to train a model using the SageMaker Python SDK. Refer to the training script train.py in the train directory here to see an example and a full explanation. This script will be wrapped in a custom image that will later be used when launching the SageMaker training job.

Among other details, this script is responsible for the following:

• Parsing job input arguments

• Loading the training data

• Configuring and training the ML model

• Retrieving the necessary credentials and other info from the AWS Secret

• Saving the model and related artifacts to Artifactory

Building a Training Container Image

SageMaker training jobs require launching a SageMaker compatible Docker image. For this example, the prepared training script will be wrapped within a Docker image that is set up with the expected directory structure and environment variables to run in the SageMaker environment. Refer to the Dockerfile in the train directory available here.

Note the first several lines of the Dockerfile — to retrieve the base image required, the FROM statement is set up to proxy a request to Docker Hub via the virtual Docker repository set up in Artifactory.

# Artifactory server/registry and repo for resolving Docker images
ARG ARTIFACTORY_DOCKER_REGISTRY
ARG ARTIFACTORY_DOCKER_REPO
# Download an open source TensorFlow Docker image to use as the Base Image
FROM
${ARTIFACTORY_DOCKER_REGISTRY}/${ARTIFACTORY_DOCKER_REPO}/tensorflow/tensorflow:latest-gpu-jupyter

Another item to note in this file is the use of a Docker build secret to supply the pip configuration necessary to resolve all of the required Python package dependencies needed to fulfill the requirements for the training script (and later for inference).

#     Ensure the source file 'pip.conf' in this project with the appropriate
#     index-url and credentials.
RUN --mount=type=secret,id=pipconfig,target=/etc/pip.conf \
    pip install -r /opt/ml/code/requirements.txt

Both the requirements.txt file and the pip.conf file are available in the git repository.

Update the pip.conf file as required with the credentials needed for the virtual PyPi repository in Artifactory, then build, tag and upload the custom Docker image to Artifactory. The following is an example command to accomplish this using the Dockerfile in the github repo. Remember to modify the Artifactory host name, the image name/tag, and the name of the virtual PyPi repository as needed for your environment.

docker build -t \
myartifactory.jfrog.io/sagemaker-docker-virtual/sagemaker/train:1.0_huggingface \
    --secret id=pipconfig,src=pip.conf \
    --build-arg "ARTIFACTORY_DOCKER_REGISTRY=myartifactory.jfrog.io" \
    --build-arg "ARTIFACTORY_DOCKER_REPO=sagemaker-docker-virtual" .
docker push myartifactory.jfrog.io/sagemaker-docker-virtual/sagemaker/train:1.0_huggingface

Note the format of the Docker tag. This format is required to identify the Artifactory registry and repository where the image will be pushed.

Running a SageMaker Training Job

Now that a custom SageMaker-compatible Docker image that includes the training script is available in Artifactory, it is now possible to launch a SageMaker Training Job to train and store an ML Model.

The Python script run-train-job.py in the train directory here demonstrates how to configure a SageMaker Python SDK Estimator with the custom Docker training image and then call its fit method to train the model.

Running this script will assign a unique name to a SageMaker Training Job and launch the specified image to begin training. This example runs about six minutes to complete a single epoch (for demo purposes), and then the model version is uploaded to the local Hugging Face repo in Artifactory.

Part 3: How to Deploy Models to AWS SageMaker Inference Endpoints via Artifactory

The goal of this example is to demonstrate deploying and interacting with the custom ML model created and stored in the previous section by performing the following steps:

1. Prepare an inference script

2. Build an inference container image

3. Deploy ML model with SageMaker Inference

4. Utilize deployed model to make predictions

All of the sample code for inference is available in the inference directory here.

Preparing an Inference Script

Refer to the script inference.py to see an example and full explanation of how to retrieve a custom, trained model from Artifactory and handle queries.

This script is responsible for the following:

• Initializing the handler service that is executed by the model server

• Retrieving the necessary credentials and other info from the AWS Secret

• Retrieving the custom Docker image for inference from Artifactory

• Retrieving the custom model for inference from Artifactory

• Processes the input and output for predictions

Building an Inference Container Image

Similar to a SageMaker training job, a SageMaker-compatible Docker image is required when deploying a model for inference. The inference script is expected to be contained within the Docker image. In this example, the file entrypoint.py is also included that will start the model server. An example Dockerfile is available in the inference directory here.

Note that the inference Dockerfile uses the training image created in the previous section as a base image. This is certainly not required, but convenient for this example because it already contains the required dependencies needed for the inference script.

FROM ${ARTIFACTORY_DOCKER_REGISTRY}/${ARTIFACTORY_DOCKER_REPO}/sagemaker/train:1.0_huggingface

The following is an example command to build, tag, and upload the custom Docker image to Artifactory. Remember to modify the Artifactory host name, the image name/tag, and the name of the virtual PyPi repository as needed for your environment.

docker build -t \
myartifactory.jfrog.io/sagemaker-docker-virtual/sagemaker/inference-service:1.0_huggingface \
    --build-arg "ARTIFACTORY_DOCKER_REGISTRY=myartifactory.jfrog.io" \
    --build-arg "ARTIFACTORY_DOCKER_REPO=sagemaker-docker-virtual" .
docker push
myartifactory.jfrog.io/sagemaker-docker-virtual/sagemaker/inference-service:1.0_huggingface

As with the training image, note the format of the Docker tag for the inference. This format is required to identify the Artifactory registry and repository where the image will be pushed.

Deploying ML Model with SageMaker Inference

With the custom Docker image, scripts, and custom ML model in place, it’s now possible to deploy the model to a SageMaker endpoint. Run the deploy-model.py script which does the work of configuring, initializing, and deploying the model to a specified endpoint.

Using Deployed Model to Make Predictions

Finally, via the SageMaker Python SDK Predictor, it’s possible to run predictions against the SageMaker endpoint. Note that the script test-inference.py in the sample repo uses the same inference endpoint configured in the deploy script to make predictions against the custom ML model that’s already deployed.