> This page is for Kumo Engine. > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/sdgm/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/sdgm/_mcp/server. # SageMaker Pipelines Recipe > Run a SageMaker pipeline that trains a Kumo Engine model, registers it in SageMaker managed MLflow, and scores with it 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](/sdgm/kumo-engine/user-code-guide/operate-models/score), 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](#review-before-promoting). The pipeline comes from [`examples/mlops/sagemaker`](https://github.com/NVIDIA/kumo-engine/tree/main/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](/sdgm/kumo-engine/user-code-guide/operate-models/turn-on). To use your own graph, edit `build_graph`, `QUERY`, and `TEXT_COLUMNS` in `jobs.py` as described in [Adapt the Template](/sdgm/kumo-engine/user-code-guide/operate-models/databricks#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](/sdgm/kumo-engine/user-code-guide/operate-models/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:///kumo/code/`, the input data under `s3:///kumo/data/`, and the artifact root at `s3:///kumo/artifacts/`. The pipeline writes SageMaker's own files to `s3:///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](https://github.com/NVIDIA/kumo-engine/releases) page. [Turn On MLOps](/sdgm/kumo-engine/user-code-guide/operate-models/turn-on#install-the-mlops-extras) 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). ## 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: | For | Permissions | | -------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | Jobs | `sagemaker:CreateTrainingJob`, `DescribeTrainingJob`, `StopTrainingJob`, `CreateProcessingJob`, `DescribeProcessingJob`, `StopProcessingJob`, and `AddTags`, and `iam:PassRole` on itself | | S3 | Read access to the code and data prefixes, and read and write access to the artifact root and the `sagemaker/` prefix | | Image and logs | `ecr:GetAuthorizationToken` and pull access to the ECR repository, and CloudWatch Logs | | MLflow | The SageMaker managed MLflow permissions listed in [Track Runs](/sdgm/kumo-engine/user-code-guide/operate-models/track-runs#choose-a-tracker) | ## Build the Job Image Clone the Kumo Engine repository, then build the image from the repository root. `` 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`: ```bash 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 ``` 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 ``, ``, and ``: ```bash aws ecr create-repository --repository-name kumo-engine-sagemaker aws ecr get-login-password | docker login --username AWS \ --password-stdin .dkr.ecr..amazonaws.com docker tag kumo-engine-sagemaker .dkr.ecr..amazonaws.com/kumo-engine-sagemaker: docker push .dkr.ecr..amazonaws.com/kumo-engine-sagemaker: ``` `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: ```bash aws s3 sync examples/mlops s3:///kumo/code/ aws s3 cp users.parquet s3:///kumo/data/ aws s3 cp items.parquet s3:///kumo/data/ aws s3 cp orders.parquet s3:///kumo/data/ ``` `aws s3 ls s3:///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 ` 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: ```bash python examples/mlops/sagemaker/pipeline.py --role \ --image .dkr.ecr..amazonaws.com/kumo-engine-sagemaker: \ --code s3:///kumo/code/ --data-root s3:///kumo/data/ \ --artifact-root s3:///kumo/artifacts/ \ --tracking-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 `; `PipelineExecutionStatus` changes from `Executing` to `Succeeded`. These options change the defaults: | Option | Default | | ---------------- | ----------------------------------------------------------------------------------- | | `--name` | `kumo-engine-users-churn`. Use your own name on a shared account | | `--experiment` | `kumo-engine-users-churn`, the MLflow experiment | | `--model-name` | `users-churn` | | `--alias` | `champion`, the alias each fit moves to its new version | | `--query` | `PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.USER_ID` | | `--table-prefix` | Empty. A prefix of the Parquet file names, such as `demo_` for `demo_users.parquet` | | `--on-error` | `raise`, so a failed store, tracking call, or registration fails the step | ## 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](/sdgm/kumo-engine/user-code-guide/operate-models/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](/sdgm/kumo-engine/user-code-guide/operate-models/track-runs#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](/sdgm/kumo-engine/user-code-guide/operate-models/turn-on#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](/sdgm/kumo-engine/user-code-guide/operate-models/registry#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](https://docs.aws.amazon.com/sagemaker/latest/dg/pipeline-eventbridge.html). Run `pipeline.py` again without `--start` to update the pipeline after you change its options. ## Next Steps #### [Register and Promote Models](/sdgm/kumo-engine/user-code-guide/operate-models/registry) Compare a new version with the champion before you move the alias. #### [MLOps Reference](/sdgm/kumo-engine/user-code-guide/operate-models/reference) Look up what each stored artifact records and where it is stored.