SageMaker Pipelines Recipe

Run a two-step SageMaker pipeline that trains and scores a Kumo Engine model with MLOps turned on. When you finish this page, each pipeline execution trains the model in a SageMaker training job, stores it under an S3 artifact root, records the fit in SageMaker managed MLflow, and registers it as a new version of users-churn that holds the champion alias. A processing job then scores your users with trainer.predict(), as in Score in Batch, using the model that the same execution trained, and stores the prediction run under the artifact root. By default every fit moves champion, so a new version is used without review; see Review Before Promoting.

The pipeline comes from examples/mlops/sagemaker in the Kumo Engine repository. Its two steps run fit_job.py and score_job.py, the job scripts in examples/mlops, on an example graph of users, items, and orders. They run the same kumo.MLOps workflow as Turn On MLOps. To use your own graph, edit build_graph, QUERY, and TEXT_COLUMNS in jobs.py as described in Adapt the Template, then upload the code again; the job image already includes the text extra, and --query overrides QUERY.

Before You Start

You need:

• A SageMaker managed MLflow tracking server, and its ARN. Track Runs shows where to find the ARN and how to start a stopped server.
• An S3 bucket. The examples keep the job code under s3://<bucket>/kumo/code/, the input data under s3://<bucket>/kumo/data/, and the artifact root at s3://<bucket>/kumo/artifacts/. The pipeline writes SageMaker’s own files to s3://<bucket>/kumo/sagemaker/, next to the artifact root.
• The input data as three Parquet files under the data prefix: users.parquet with primary key USER_ID, items.parquet with primary key ITEM_ID, and orders.parquet with primary key ORDER_ID, time column ORDER_TIME, and the foreign keys USER_ID and ITEM_ID. The default query needs a few months of orders.
• The kumo-engine wheel, from PyPI or the Kumo Engine releases page. Turn On MLOps shows how to check that a build includes MLOps. The job image installs it.
• Docker, and an Amazon ECR repository for the job image.
• The SageMaker Python SDK version 3 on Python 3.10 or later, installed with pip install "sagemaker>=3", and an AWS region configured, for example through AWS_REGION.
• An execution role that the pipeline’s jobs run as, and permissions for your own identity, as listed in Grant Permissions.

Grant Permissions

Your identity, which creates and starts the pipeline, needs sagemaker:CreatePipeline, sagemaker:UpdatePipeline, sagemaker:StartPipelineExecution, and iam:PassRole on the execution role. It also needs write access to the sagemaker/ prefix, and s3:ListBucket on the bucket or s3:ListAllMyBuckets.

The execution role must trust sagemaker.amazonaws.com, and needs:

Build the Job Image

Clone the Kumo Engine repository, then build the image from the repository root. <wheel dir> is a directory that holds the kumo-engine wheel. This example builds a CPU image, which runs on CPU instance types such as ml.m5.xlarge:

git clone https://github.com/NVIDIA/kumo-engine.git
cd kumo-engine
docker build --platform linux/amd64 --build-arg TORCH=cpu \
  -f examples/mlops/sagemaker/Dockerfile -t kumo-engine-sagemaker <wheel dir>

For GPU instance types, such as ml.g5.xlarge, leave out --build-arg TORCH=cpu. The build ends by printing the installed Kumo Engine and PyTorch versions.

Create the repository and push the image. Replace <account>, <region>, and <tag>:

aws ecr create-repository --repository-name kumo-engine-sagemaker
aws ecr get-login-password | docker login --username AWS \
  --password-stdin <account>.dkr.ecr.<region>.amazonaws.com
docker tag kumo-engine-sagemaker <account>.dkr.ecr.<region>.amazonaws.com/kumo-engine-sagemaker:<tag>
docker push <account>.dkr.ecr.<region>.amazonaws.com/kumo-engine-sagemaker:<tag>

docker login prints Login Succeeded, and the push ends with a line that shows the image digest.

Upload the Code and Data

Copy the job scripts to the code prefix, and the three Parquet files to the data prefix:

aws s3 sync examples/mlops s3://<bucket>/kumo/code/
aws s3 cp users.parquet s3://<bucket>/kumo/data/
aws s3 cp items.parquet s3://<bucket>/kumo/data/
aws s3 cp orders.parquet s3://<bucket>/kumo/data/

aws s3 ls s3://<bucket>/kumo/code/ lists the job scripts, including fit_job.py, score_job.py, and jobs.py.

Create and Start the Pipeline

pipeline.py creates the pipeline, or updates it if it exists, and --start starts an execution. Both steps run the image you pass as --image, so choose instance types that match it: CPU types for the CPU image, GPU types for the GPU image.

The Predict step is a processing job, which uses the <instance type> for processing job usage quota in Service Quotas, separate from the training job quota. For GPU instance types that quota is often 0, so request an increase before you score on a GPU instance type.

Replace the role ARN, the image URI, the bucket, and the tracking server ARN:

python examples/mlops/sagemaker/pipeline.py --role <execution-role-arn> \
  --image <account>.dkr.ecr.<region>.amazonaws.com/kumo-engine-sagemaker:<tag> \
  --code s3://<bucket>/kumo/code/ --data-root s3://<bucket>/kumo/data/ \
  --artifact-root s3://<bucket>/kumo/artifacts/ \
  --tracking-arn <tracking-server-arn> \
  --fit-instance-type ml.m5.xlarge --score-instance-type ml.m5.xlarge \
  --start

The command prints the execution ARN. The pipeline, named kumo-engine-users-churn, appears in SageMaker Studio under Pipelines, with the steps Fit and Predict. To check the execution from the command line, run aws sagemaker describe-pipeline-execution --pipeline-execution-arn <execution-arn>; PipelineExecutionStatus changes from Executing to Succeeded.

These options change the defaults:

Check the Result

When the execution succeeds, both steps show Succeeded in Studio, and:

• The MLflow experiment kumo-engine-users-churn holds the fit run, named after the new version, such as fit-users-churn-v1, and the prediction run. Track Runs shows how to open the MLflow UI on SageMaker.
• The registered model users-churn has a new version that holds champion.
• The artifact root holds the snapshot, the training table, the model, and the prediction run, under snapshots/, training_tables/, models/, and predictions/.

Predict scores the model that its own Fit step trained, even if a later execution has moved the alias since. It loads that model by its id rather than by alias, so the prediction run is named like a run of an unregistered model, as listed in Run Names. This hand-off of the trained model from Fit to Predict has not yet been run on SageMaker.

The Predict step writes no predictions table. Load a prediction run from the artifact root with kumo.PredictionResult.load_from_connector, as described in Read Artifacts Back.

Review Before Promoting

With the default --alias champion, every fit becomes the champion. To review each version first, pass --alias candidate. Each fit then moves candidate, the pipeline still scores the model it trained, and you move champion yourself, as described in Promote a Version.

Run on a Schedule

A pipeline runs when you start it, and every execution both trains and scores. To start it on a schedule, create an Amazon EventBridge schedule that starts the pipeline, as described in Schedule Pipeline Runs. Run pipeline.py again without --start to update the pipeline after you change its options.

Next Steps