Amazon SageMaker

Deploy Kumo Relational NIM as an Amazon SageMaker AI endpoint and send a prediction using Python. This recipe builds a derived image that adds SageMaker health and prediction routes to the original NIM. It runs one model on one GPU per endpoint.

Before you start

• A Bash shell, Docker with the Buildx plugin capable of building linux/amd64 images, AWS CLI v2, Python 3, and sufficient Docker storage to pull and build the NIM image. The build machine does not need a GPU or an x86-64 CPU. Docker Desktop supports cross-platform builds on macOS and Windows. On Windows, run these Bash commands in WSL2 with Docker Desktop integration enabled.
• An NGC account and an NGC personal API key with access to the model image.
• An AWS identity with permission to manage SageMaker endpoints and ECR images, invoke the endpoint, and pass a SageMaker execution role using iam:PassRole.
• An existing SageMaker execution role that trusts sagemaker.amazonaws.com and can pull your ECR image and write CloudWatch logs.
• SageMaker endpoint quota and available capacity for one ml.g6e.2xlarge instance in your Region. This instance has one 48 GB L40S GPU.

1. Sign in and choose names

Sign in using your organization’s AWS CLI authentication method. Replace the execution role ARN below with the role supplied by your administrator. Choose a new endpoint name for each deployment. Run the shell commands in the same terminal. For a named AWS profile, set export AWS_PROFILE=YOUR_PROFILE first.

export AWS_DEFAULT_REGION=us-west-2
export ENDPOINT_NAME=kumo-relational-$(date +%Y%m%d%H%M%S)
export EXECUTION_ROLE_ARN=arn:aws:iam::ACCOUNT_ID:role/ROLE_NAME
export INSTANCE_TYPE=ml.g6e.2xlarge
export NIM_IMAGE=nvcr.io/nim/nvidia/kumo-relational:1.1.0
export ECR_REPOSITORY=kumo-relational-sagemaker
aws sts get-caller-identity
docker login nvcr.io --username '$oauthtoken'

Check that the AWS command shows the intended account. Paste your NGC API key at Docker’s password prompt.

2. Build the SageMaker image

Create an empty directory and save the following three files in it. The proxy forwards SageMaker health checks to /v1/health/ready and prediction requests to /v1/predictions. The original NIM and proxy run inside a single container.

ARG NIM_IMAGE=nvcr.io/nim/nvidia/kumo-relational:1.1.0
FROM caddy:2.10.2@sha256:c3d7ee5d2b11f9dc54f947f68a734c84e9c9666c92c88a7f30b9cba5da182adb AS proxy
FROM ${NIM_IMAGE}
USER 0:0
COPY --from=proxy /usr/bin/caddy /usr/local/bin/caddy
COPY Caddyfile start.sh /opt/sagemaker/
ENV NIM_HTTP_API_PORT=8000
EXPOSE 8080
ENTRYPOINT ["/bin/bash", "/opt/sagemaker/start.sh"]
CMD ["serve"]

{
	admin off
	auto_https off
	persist_config off
}
:8080 {
	@health {
		method GET
		path /ping
	}
	handle @health {
		rewrite * /v1/health/ready
		reverse_proxy 127.0.0.1:8000
	}
	@predict {
		method POST
		path /invocations
	}
	handle @predict {
		rewrite * /v1/predictions
		reverse_proxy 127.0.0.1:8000
	}
	handle {
		respond "Not found" 404
	}
}

set -euo pipefail
pids=()
cleanup() {
    trap '' TERM INT
    if ((${#pids[@]})); then
        for pid in "${pids[@]}"; do
            kill -TERM -- "-$pid" 2>/dev/null || true
        done
        wait "${pids[@]}" 2>/dev/null || true
    fi
}
trap cleanup EXIT
trap 'exit 0' TERM INT
setsid caddy run --config /opt/sagemaker/Caddyfile --adapter caddyfile &
pids+=("$!")
setsid /opt/nvidia/nvidia_entrypoint.sh /bin/bash -c "$SERVER_START_SCRIPT_PATH" &
pids+=("$!")
wait -n "${pids[@]}" || true
exit 1

Build the image and upload it to Amazon ECR. If the ECR repository already exists, omit create-repository below.

set -e
ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text)
REGISTRY="$ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com"
export IMAGE_URI="$REGISTRY/$ECR_REPOSITORY:$ENDPOINT_NAME"
aws ecr create-repository --repository-name "$ECR_REPOSITORY"
aws ecr get-login-password | docker login --username AWS --password-stdin "$REGISTRY"
docker buildx build --platform linux/amd64 --provenance=false --sbom=false \
  --build-arg NIM_IMAGE="$NIM_IMAGE" \
  --output "type=image,name=$IMAGE_URI,push=true,oci-mediatypes=false" .

The build explicitly publishes a Docker image manifest compatible with SageMaker.

3. Create the endpoint

Register the image and create one instance:

aws sagemaker create-model   --model-name "$ENDPOINT_NAME"   --execution-role-arn "$EXECUTION_ROLE_ARN"   --primary-container "Image=$IMAGE_URI"
aws sagemaker create-endpoint-config   --endpoint-config-name "$ENDPOINT_NAME"   --production-variants "VariantName=AllTraffic,ModelName=$ENDPOINT_NAME,InitialInstanceCount=1,InstanceType=$INSTANCE_TYPE,ContainerStartupHealthCheckTimeoutInSeconds=1800"
aws sagemaker create-endpoint   --endpoint-name "$ENDPOINT_NAME"   --endpoint-config-name "$ENDPOINT_NAME"
aws sagemaker wait endpoint-in-service --endpoint-name "$ENDPOINT_NAME"

The waiter returns successfully when the endpoint is InService. Image download and model initialization can take several minutes. If the waiter times out, inspect the endpoint before retrying or deleting it:

aws sagemaker describe-endpoint --endpoint-name "$ENDPOINT_NAME"   --query '{Status:EndpointStatus,Reason:FailureReason}'

For quota or capacity failures, request quota or retry when the instance is available. For container startup errors, inspect CloudWatch logs under /aws/sagemaker/Endpoints/<endpoint-name>.

4. Make a prediction

Create a Python environment and install the AWS SDK. The crt extra enables AWS console login credentials, including aws login:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install "boto3[crt]>=1.41.0"

Save the following synthetic request as request.json. It supplies four labeled context rows and predicts a label for one new row.

{
  "model": "kumo-relational",
  "task": {
    "kind": "classification",
    "target": {"column_name": "label", "dtype": "string", "classes": ["no", "yes"], "positive_class": "yes"},
    "entity_table_names": ["accounts"],
    "anchor_time_column": "anchor_time"
  },
  "schema": {
    "instance_table": {
      "columns": {
        "instance_id": {"dtype": "int64", "stype": "ID", "nullable": false},
        "amount": {"dtype": "float64", "stype": "numerical"},
        "label": {"dtype": "string", "stype": "categorical"},
        "account_id": {"dtype": "int64", "stype": "ID", "nullable": false},
        "anchor_time": {"dtype": "timestamp[us]", "stype": "timestamp", "nullable": false}
      },
      "primary_key": "instance_id"
    },
    "related_tables": {
      "accounts": {
        "columns": {
          "instance_id": {"dtype": "int64", "stype": "ID", "nullable": false},
          "account_id": {"dtype": "int64", "stype": "ID", "nullable": false},
          "segment": {"dtype": "string", "stype": "categorical"}
        },
        "primary_key": ["instance_id", "account_id"]
      }
    },
    "relationships": [
      {
        "source_columns": ["instance_id", "account_id"],
        "target_table": "accounts",
        "target_columns": ["instance_id", "account_id"]
      }
    ]
  },
  "context": {
    "instance_table": {
      "format": "arrays",
      "columns": ["instance_id", "amount", "label", "account_id", "anchor_time"],
      "rows": [
        [0, 10.0, "no", 100, "2026-01-01T00:00:00Z"],
        [1, 250.0, "yes", 101, "2026-01-01T00:00:00Z"],
        [2, 12.0, "no", 102, "2026-01-01T00:00:00Z"],
        [3, 300.0, "yes", 103, "2026-01-01T00:00:00Z"]
      ]
    },
    "related_tables": {
      "accounts": {
        "format": "arrays",
        "columns": ["instance_id", "account_id", "segment"],
        "rows": [[0, 100, "small"], [1, 101, "enterprise"], [2, 102, "small"], [3, 103, "enterprise"]]
      }
    }
  },
  "predict": {
    "instance_table": {
      "format": "arrays",
      "columns": ["instance_id", "amount", "account_id", "anchor_time"],
      "rows": [[4, 275.0, 104, "2026-01-01T00:00:00Z"]]
    },
    "related_tables": {
      "accounts": {
        "format": "arrays",
        "columns": ["instance_id", "account_id", "segment"],
        "rows": [[4, 104, "enterprise"]]
      }
    }
  },
  "output": {"fields": ["prediction", "probabilities"]}
}

Save this as predict.py and run python predict.py in the same terminal:

import json
import os
from pathlib import Path
import boto3
from botocore.config import Config
runtime = boto3.client(
    "sagemaker-runtime",
    config=Config(read_timeout=70, retries={"total_max_attempts": 1}),
)
response = runtime.invoke_endpoint(
    EndpointName=os.environ["ENDPOINT_NAME"],
    ContentType="application/json",
    Body=Path("request.json").read_bytes(),
)
result = json.loads(response["Body"].read())
print(json.dumps(result, indent=2))

The response contains a predictions array with one prediction and class probabilities. This small example verifies inference, not model accuracy on your data. AWS credentials authenticate the request; an NGC key is not used for inference.

This recipe supports complete prediction requests through boto3. It does not expose model discovery or cached-session APIs. The existing SDM client HTTP examples cannot connect to this endpoint by substituting a SageMaker URL. For data preparation and modeling examples, see the public Kumo Relational notebook, which uses a direct NIM HTTP connection.

Synchronous requests and responses must each fit within 6 MiB, and inference must finish within 60 seconds. See SageMaker invocation limits.

5. Delete the deployment

Endpoints incur charges while running, including when idle. Delete this endpoint and its configuration when finished:

aws sagemaker delete-endpoint --endpoint-name "$ENDPOINT_NAME"
aws sagemaker wait endpoint-deleted --endpoint-name "$ENDPOINT_NAME"
aws sagemaker delete-endpoint-config --endpoint-config-name "$ENDPOINT_NAME"
aws sagemaker delete-model --model-name "$ENDPOINT_NAME"
aws ecr batch-delete-image --repository-name "$ECR_REPOSITORY"   --image-ids "imageTag=$ENDPOINT_NAME"

The execution role, ECR repository, CloudWatch logs, and local image remain available for reuse. Remove them separately when no longer needed.