Create a custom Docker container image for SageMaker and use it for model training in AWS Step Functions

Setup Amazon ECR and create a new private registry.

If you haven’t already, set up Amazon ECR by following the instructions in Setting up with Amazon ECR in the Amazon ECR User Guide. Each AWS account is provided with a default private Amazon ECR registry.

Create an Amazon ECR private repository.

Follow the instructions in Creating a private repository in the Amazon ECR User Guide.

Note: The repository that you create is where you’ll store your custom Docker container images.

Create a Dockerfile that includes the specifications needed to run your SageMaker processing job.

Create a Dockerfile that includes the specifications needed to run your SageMaker processing job by configuring a Dockerfile. For instructions, see Adapting your own training container in the Amazon SageMaker Developer Guide.

For more information about Dockerfiles, see the Dockerfile Reference in the Docker documentation.

Example Jupyter notebook code cells to create a Dockerfile

Cell 1

# Make docker folder
!mkdir -p docker

Cell 2

%%writefile docker/Dockerfile

FROM python:3.7-slim-buster

RUN pip3 install pandas==0.25.3 scikit-learn==0.21.3
ENV PYTHONUNBUFFERED=TRUE

ENTRYPOINT ["python3"]

Build your Docker container image and push it to Amazon ECR.

For more information, see Building and registering the container in Building your own algorithm container on GitHub.

Example Jupyter notebook code cells to build and register a Docker image

Important: Before running the following cells, make sure that you’ve created a Dockerfile and stored it in the directory called docker. Also, make sure that you’ve created an Amazon ECR repository, and that you replace the ecr_repository value in the first cell with your repository’s name.

Cell 1

import boto3
tag = ':latest'
account_id = boto3.client('sts').get_caller_identity().get('Account')
region = boto3.Session().region_name
ecr_repository = 'byoc'

image_uri = '{}.dkr.ecr.{}.amazonaws.com/{}'.format(account_id, region, ecr_repository + tag)

Cell 2

# Build docker image
!docker build -t $image_uri docker

Cell 3

# Authenticate to ECR
!aws ecr get-login-password --region {region} | docker login --username AWS --password-stdin {account_id}.dkr.ecr.{region}.amazonaws.com

Cell 4

# Push docker image
!docker push $image_uri

Note: You must authenticate your Docker client to your private registry so that you can use the docker push and docker pull commands. These commands push and pull images to and from the repositories in your registry.

Create a Python script that includes your custom processing and model training logic.

Write custom processing logic to run in your data processing script. Then, save it as a Python script named training.py.

For more information, see Bring your own model with SageMaker Script Mode on GitHub.

Example Python script that includes custom processing and model training logic

%%writefile training.py
from numpy import empty
import pandas as pd
import os
from sklearn import datasets, svm
from joblib import dump, load


if __name__ == '__main__':
    digits = datasets.load_digits()
    #create classifier object
    clf = svm.SVC(gamma=0.001, C=100.)
    
    #fit the model
    clf.fit(digits.data[:-1], digits.target[:-1])
    
    #model output in binary format
    output_path = os.path.join('/opt/ml/processing/model', "model.joblib")
    dump(clf, output_path)

Create a Step Functions workflow that includes your SageMaker Processing job as one of the steps.

Install and import the AWS Step Functions Data Science SDK and upload the training.py file to Amazon S3. Then, use the Amazon SageMaker Python SDK to define a processing step in Step Functions.

Important: Make sure that you’ve created an IAM execution role for Step Functions in your AWS account.

Example environment set up and custom training script to upload to Amazon S3

!pip install stepfunctions

import boto3
import stepfunctions
import sagemaker
import datetime

from stepfunctions import steps
from stepfunctions.inputs import ExecutionInput
from stepfunctions.steps import (
    Chain
)
from stepfunctions.workflow import Workflow
from sagemaker.processing import ScriptProcessor, ProcessingInput, ProcessingOutput

sagemaker_session = sagemaker.Session()
bucket = sagemaker_session.default_bucket() 
role = sagemaker.get_execution_role()
prefix = 'byoc-training-model'

# See prerequisites section to create this role
workflow_execution_role = f"arn:aws:iam::{account_id}:role/AmazonSageMaker-StepFunctionsWorkflowExecutionRole"

execution_input = ExecutionInput(
    schema={
        "PreprocessingJobName": str})


input_code = sagemaker_session.upload_data(
    "training.py",
    bucket=bucket,
    key_prefix="preprocessing.py",
)

Example SageMaker processing step definition that uses a custom Amazon ECR image and Python script

Note: Make sure that you use the execution_input parameter to specify the job name. The parameter’s value must be unique each time the job runs. Also, the training.py file’s code is passed as an input parameter to the ProcessingStep, which means that it will be copied inside the container. The destination for the ProcessingInput code is the same as the second argument inside the container_entrypoint.

script_processor = ScriptProcessor(command=['python3'],
                image_uri=image_uri,
                role=role,
                instance_count=1,
                instance_type='ml.m5.xlarge')


processing_step = steps.ProcessingStep(
    "training-step",
    processor=script_processor,
    job_name=execution_input["PreprocessingJobName"],
    inputs=[
        ProcessingInput(
            source=input_code,
            destination="/opt/ml/processing/input/code",
            input_name="code",
        ),
    ],
    outputs=[
        ProcessingOutput(
            source='/opt/ml/processing/model', 
            destination="s3://{}/{}".format(bucket, prefix), 
            output_name='byoc-example')
    ],
    container_entrypoint=["python3", "/opt/ml/processing/input/code/training.py"],
)

Example Step Functions workflow that runs a SageMaker processing job

Note: This example workflow includes the SageMaker processing job step only, not a complete Step Functions workflow. For a full example workflow, see Example notebooks in SageMaker in the AWS Step Functions Data Science SDK documentation.

workflow_graph = Chain([processing_step])

workflow = Workflow(
    name="ProcessingWorkflow",
    definition=workflow_graph,
    role=workflow_execution_role
)

workflow.create()
# Execute workflow
execution = workflow.execute(
    inputs={
        "PreprocessingJobName": str(datetime.datetime.now().strftime("%Y%m%d%H%M-%SS")),  # Each pre processing job (SageMaker processing job) requires a unique name,
    }
)
execution_output = execution.get_output(wait=True)