Deploy GPU-based AI Guardrails to offload LLM security detection from the AI Gateway and improve throughput in production.

Choose the deployment path that fits your environment for installing GPU-based guardrails:

• Option A – SageMaker path: Push Docker image to ECR, deploy SageMaker endpoint, protect it with Cognito (M2M) and expose via API Gateway.
• Option B – EC2/VM path: Run the GPU container directly on a GPU host, optionally front with HAProxy, and optionally keep API Gateway and Cognito for consistent authentication.

Option A — Deploy on AWS SageMaker (Managed)

Prerequisites

• AWS CLI and permissions for ECR, SageMaker, API Gateway, IAM, and Cognito
• GPU-capable SageMaker instance family quota (for example ml.g5.xlarge)
• AI Guardrails LLM Docker image tarball or access to a Docker registry

Step 1 — Push the LLM image to Amazon ECR

1. Create an ECR repository with AES-256 encryption.
   For example: aisecurity/aisecurityllm-docker-image

# Authenticate Docker to ECR
aws ecr get-login-password --region ${AWS_REGION} | docker login --username AWS --password-stdin \ ${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_REGION}.amazonaws.com
# If you have a local tarball
IMAGE_TAG=${IMAGE_TAG}
docker load -i aisecurityllm-${IMAGE_TAG}.tgz
docker tag <loaded-image>:${IMAGE_TAG} \ ${ECR_URI}:${IMAGE_TAG}
docker push ${ECR_URI}:${IMAGE_TAG}

Step 2 — Create the SageMaker execution role

Create the role with a trust policy and attach the AmazonS3FullAccess and AmazonSageMakerFullAccess managed policies.

# Trust SageMaker and attach policies (AmazonS3FullAccess, AmazonSageMakerFullAccess), add PassRole if API GW will assume it

Step 3 — Create the Model, Endpoint Config, and Endpoint

AWS_REGION=${AWS_REGION}
IMAGE_TAG=${IMAGE_TAG}
MODEL_NAME=aisecurity-llm-model-${IMAGE_TAG//./-}
ECR_IMAGE_URI=${ECR_URI}:${IMAGE_TAG}
SAGEMAKER_EXECUTION_ROLE_ARN= ${SAGEMAKER_EXECUTION_ROLE_ARN}
# Create the model
aws sagemaker create-model \
  --model-name ${MODEL_NAME} \
  --primary-container Image=${ECR_IMAGE_URI} \
  --execution-role-arn ${SAGEMAKER_EXECUTION_ROLE_ARN} \
  --region ${AWS_REGION}
# Create the endpoint configuration
ENDPOINT_CONFIG_NAME=aisecuriy-llm-endpoint-config-${IMAGE_TAG//./-}
aws sagemaker create-endpoint-config \
  --endpoint-config-name ${ENDPOINT_CONFIG_NAME} \
  --production-variants VariantName=AllTraffic,ModelName=${MODEL_NAME},InitialInstanceCount=1,InstanceType=ml.g5.xlarge \
  --region ${AWS_REGION}
# Create the endpoint
ENDPOINT_NAME=aisecurity-llm-endpoint
aws sagemaker create-endpoint \
  --endpoint-name ${ENDPOINT_NAME} \
  --endpoint-config-name ${ENDPOINT_CONFIG_NAME} \
  --region ${AWS_REGION}
# Wait until InService
aws sagemaker wait endpoint-in-service --endpoint-name ${ENDPOINT_NAME} --region ${AWS_REGION}

SageMaker Model and Endpoint Configuration objects are immutable. For image updates, create a new Model and Endpoint Configuration. Then update the existing Endpoint to the new configuration.

Step 4 — Configure JWT with Amazon Cognito (M2M)

1. Create a User Pool and a Resource Server with a custom scope.
   For example: predict/Sagemaker.invoke
2. Create a Machine-to-Machine App Client with the custom scope and client_credentials grant.

TOKEN_URL="https://<your-domain>.auth.${AWS_REGION}.amazoncognito.com/oauth2/token"
curl -s -X POST "${TOKEN_URL}" \
  -H "Content-Type: application/x-www-form-urlencoded" \
  -d "grant_type=client_credentials&client_id=${CLIENT_ID}&client_secret=${CLIENT_SECRET}&scope=predict/Sagemaker.invoke"

Step 5 — Expose the endpoint via Amazon API Gateway

1. Create an IAM role for API Gateway (assume-role trust to apigateway.amazonaws.com) and allow sagemaker:InvokeEndpoint.
2. Create a REST API, add a resource (for example /predict), attach a Cognito authorizer with scope predict/Sagemaker.invoke.
   # Path override example: endpoints/aisecurity-llm-endpoint/invocations # Integration type: AWS, service: SageMaker Runtime, method: POST # Execution role: the role from step 1 (or reuse the SageMaker execution role with updated trust)
3. Create method and integration responses for 200/400/500, then deploy the API to a stage and note the Invoke URL.

Step 6 – Validate the end-to-end flow

INVOKE_URL="https://<api-id>.execute-api.${AWS_REGION}.amazonaws.com/<stage>/predict"
ACCESS_TOKEN="<from Cognito>"
curl -s -X POST ${INVOKE_URL} \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer ${ACCESS_TOKEN}" \
  -d '{
    "input_data": "the capital of France",
    "signature": "<base64_signature_for_input_data>",
    "verbose": true
  }' | jq .

Option B — Deploy on GPU EC2 or VM (Self-managed)

When to choose this

• Much lower cost than SageMaker (e.g., g5.xlarge EC2 at about $1.00/h versus over $900/month for a continuously running SageMaker endpoint).
• Option to stop/start instances and scale horizontally using HAProxy or NLB.

Prerequisites

• GPU instance (for example g5.xlarge, g6.xlarge) or on-prem VM with NVIDIA GPU
• NVIDIA driver and nvidia-container-toolkit installed on the host
• Docker and access to the aisecurityllm image

Step 1 – Install NVIDIA driver and container toolkit (Ubuntu)

sudo apt update
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
echo "deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://nvidia.github.io/libnvidia-container/stable/deb/$(dpkg --print-architecture) /" | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit ubuntu-drivers-common
sudo ubuntu-drivers autoinstall
# Reboot if required, then verify:
nvidia-smi

Step 2 – Run the GPU container

# Load the image if using a tarball
IMAGE_TAG=${IMAGE_TAG}
sudo docker load -i aisecurityllm-${IMAGE_TAG}.tgz
# Start the service on port 8080 with GPU
sudo docker run -d --gpus all \
  --shm-size=1g \
  -e CUDA_VISIBLE_DEVICES=0 \
  -e TORCH_COMPILE_DISABLE=1 \
  -e PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 \
  -e TOKENIZERS_PARALLELISM=false \
  -p 8080:8080 \
  <loaded-image>:${IMAGE_TAG}

Step 3 – Health check and test

# Health
curl -s http://<host>:8080/ping
# Inference (signature must match input_data)
curl -s -X POST http://<host>:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{
     "input_data":"the capital of France",
     "signature":"<base64_signature>",
     "verbose":true
     }' | jq .

Optional: Scale-out with HAProxy

For higher throughput, deploy multiple GPU hosts and balance traffic with HAProxy. You can optionally front HAProxy with an AWS Network Load Balancer and API Gateway (reusing the Cognito M2M flow from Option A) to provide a single, secured external endpoint.

HAProxy example

# /etc/haproxy/haproxy.cfg (excerpt)
frontend fe_http
  bind :80
  default_backend be_invocations
backend be_invocations
  balance roundrobin
  option httpchk GET /ping
  http-check expect status 200
  server gpu-node-1 <host1-ip>:8080 check
  server gpu-node-2 <host2-ip>:8080 check
  server gpu-node-3 <host3-ip>:8080 check
  # ...add/remove servers to scale

API Gateway + Cognito with HAProxy backend

• Create a VPC Link to an internal NLB in front of HAProxy, add a new API resource (for example /haproxy).
• Keep your Cognito authorizer; define scopes such as haproxy.invoke.
• Invoke with the Cognito bearer token; API Gateway forwards to HAProxy, which load-balances requests to GPU nodes.

Performance Reference

The following throughput numbers were observed in testing. Actual performance depends on content length, profile configuration, and AI Gateway host sizing.

• Latency: typically 0.5–1.3s in tests; varies by profile and content.
• Cost: SageMaker hosting is convenient but costly and cannot be easily paused; EC2/HAProxy clusters can scale down or shut off when idle.

Configuration in Netskope (AIG CLI quick check)

Verify LLM service connectivity via the AIG Configuration Wizard (aig-cli).

Common checks:

• Configure with Host URL only (no auth) to a direct EC2/HAProxy URL ending with /invocations.
• Configure with Host URL, JWT URL, client_id, client_secret, and scope for API Gateway and Cognito setup.

Troubleshooting