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.

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

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

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.

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

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

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)