Build a model improvement flywheel with Amazon SageMaker HyperPod Data Capture
AWS Machine Learning Blog0 views
Foundation models deployed on Amazon SageMaker HyperPod meet inputs that no pre-curated dataset covers: domain-specific jargon, unexpected formats, new products, and changing user behavior. Without a record of production inference data, teams can’t detect drift, find failure patterns, or use real interactions to improve the model.
Amazon SageMaker HyperPod now provides Data Capture, a built-in capability that records live inference inputs and outputs to Amazon Simple Storage Service (Amazon S3) at three independent tiers, with minimal latency impact and zero impact on model availability.
In this post, we walk through one complete example on a single SageMaker HyperPod cluster. A home-appliance support router captures its production traffic, Amazon Athena queries surface a new-product blind spot, a SQL policy labels the captured records, and a fine-tune on that data teaches the router the new product. We measure accuracy before and after, on data that could not have existed before the model was deployed. The inference operator serves the model on one node, and the training operator runs the fine-tuning on another.
We use synthetic data for the walkthrough. The scripts, manifests, SQL, synthetic datasets, and expected outputs are available in the accompanying GitHub repository.
Three-tier Data Capture architecture
HyperPod Data Capture provides three independent tiers, each at a different point in the request path. You can turn on any combination in the same dataCapture block.
Tier 1: SageMaker endpoint
Tier 2: Load balancer
Tier 3: Model pod
What’s captured
Full request/response payloads
Application Load Balancer (ALB) access-log fields: request line, status codes, timings, client IP, user agent
Full request/response payloads plus event metadata
Tiers 1 and 3 both store full payloads, so enabling both stores every record twice. Use Tier 1 for compliance and Model Monitor, Tier 2 when you only need traffic metadata, and Tier 3 when you intend to use the payloads. This post uses Tier 3 only. Because Tier 3 records complete payloads, which can include names, addresses, and order numbers, protect that data before it becomes training data. Use sampling, AWS Key Management Service (AWS KMS) encryption (kmsKeyId), S3 lifecycle rules, and personally identifiable information (PII) redaction (for example, with Amazon Comprehend).
Two available NVIDIA GPUs, one for serving and one for training. We used an ml.g6e.8xlarge for serving and an ml.g6e.xlarge for training (one 48 GB NVIDIA L40S each). The serving pod holds its GPU for the life of the endpoint, so with only one GPU the training pod stays Pending.
kubectl configured for the cluster.
An S3 bucket in the same Region.
Confirm both operators are installed. The command returns two rows:
The following diagram shows the five steps of the walkthrough on one SageMaker HyperPod cluster: capture every request and response to Amazon S3, analyze the captures with Athena, label them with a SQL policy, fine-tune with the training operator on a second GPU, and redeploy side by side with capture still on.
Figure 1: The five-step flywheel on a single SageMaker HyperPod cluster, from capture through Athena analysis, labeling, fine-tuning, and redeploy
The inference operator adds three sidecars to the model pod. Every request that passes through the NGINX proxy is copied to an OpenTelemetry collector, batched, and written to S3 through a Mountpoint-S3 volume. Athena reads those files in place. The training operator runs the fine-tune as a HyperPodPyTorchJob on a second GPU, and the new model is served by a second endpoint that also has capture on.
Step 0: Set your variables
Every later step reads these. Edit the first block, then paste the whole thing:
# --- edit these --- export EKS_CLUSTER_NAME="my-hyperpod-eks" export AWS_REGION="us-east-1" export BUCKET="amzn-s3-demo-bucket" export SERVE_INSTANCE_TYPE="ml.g6e.8xlarge" # node type for the first endpoint export TRAIN_INSTANCE_TYPE="ml.g6e.xlarge" # node type for the fine-tuning job # --- leave as-is --- export AWS_DEFAULT_REGION="$AWS_REGION" export ENDPOINT_NAME="appliance-router" export MODEL_LOCATION="models/qwen2.5-3b-instruct" export MODEL_V2_LOCATION="models/qwen2.5-3b-router-v2" export TRAIN_SA="flywheel" # service account with S3 access via EKS Pod Identity aws eks update-kubeconfig --name "$EKS_CLUSTER_NAME" --region "$AWS_REGION" render() { python3 -c 'import os,sys; sys.stdout.write(os.path.expandvars(sys.stdin.read()))'; }
Step 1: Stage the base model in S3
We use Qwen2.5-3B-Instruct because it is small enough to fine-tune in about a minute and fits on a single L40S.
The Mountpoint-S3 CSI driver role. Tier 3 capture writes through a Mountpoint-S3 volume, and the driver authenticates as its own IAM roles for service accounts (IRSA) role, not as the pod’s service account. It needs s3:PutObject, s3:AbortMultipartUpload, and s3:DeleteObject on your capture bucket, plus kms:Decrypt, kms:GenerateDataKey, and kms:DescribeKey if you set kmsKeyId. If this is missing, the endpoint stays healthy and requests succeed, but no captures appear. See Mountpoint for Amazon S3 CSI driver.
The inference operator execution role. It needs s3:PutObject on your capture bucket, plus kms:Decrypt and kms:GenerateDataKey if you set kmsKeyId. The capture files themselves are written through the CSI driver role, so when no captures appear, check that role first.
The training pod’s service account ($TRAIN_SA). Read access to the base model and training data, and write access for the merged model, through EKS Pod Identity.
Athena. Read access to the capture prefix, plus a query-results location.
Enabling Data Capture
The next three steps deploy the router with Tier 3 capture on, check that the pod runs the capture sidecars, and confirm that records reach Amazon S3.
Step 3: Deploy the model with Tier 3 capture enabled
This is the whole feature: the dataCapture.modelPod.enabled: true block. Everything else is an ordinary vLLM deployment.
The preceding modelPod block sets 100% sampling and captures both input and output, which are the defaults. To tune it, add any of these optional fields under modelPod:
modelPod: enabled: true initialSamplingPercentage: 50 # capture 50% of requests (default 100) bufferConfig: batchSize: 100 # records buffered per upload batch (default 10) flushIntervalSeconds: 60 # max seconds before a flush (default 60) payloadConfig: maxPayloadSizeKB: 1024 # truncate payloads larger than this (default: no limit) kmsKeyId: "arn:aws:kms:...:key/..." # encrypt captures at rest
Step 4: Wait for the pod
A capture-enabled pod runs four containers: the model, sidecar-reverse-proxy (NGINX), otel-collector, and s3-uploader. The model container and the pod’s app label are both named after modelName.
kubectl get pods -l app=$ENDPOINT_NAME -o custom-columns=\ 'NAME:.metadata.name,READY:.status.containerStatuses[*].ready,CONTAINERS:.spec.containers[*].name,NODE:.spec.nodeName' NAME READY CONTAINERS NODE appliance-router-849f466476-4f5v8 true,true,true,true appliance-router,s3-uploader,sidecar-reverse-proxy,otel-collector hyperpod-i-0f3fe2684afd6a583
In our run, the pod was ready about five minutes after kubectl apply, most of it spent downloading the weights.
Step 5: Send a request and confirm captures land in S3
Send an inference request through the capture proxy on port 8081, then, after the flush interval, list the capture prefix. The model itself listens on port 8000. Requests sent there bypass the proxy and are not captured.
POD=$(kubectl get pod -l app=$ENDPOINT_NAME -o jsonpath='{.items[0].metadata.name}') kubectl port-forward "pod/$POD" 18081:8081 >/dev/null 2>&1 & sleep 5 curl -s http://localhost:18081/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{"model":"'"$ENDPOINT_NAME"'","messages":[{"role":"user","content":"my fridge is not cooling"}],"max_tokens":24}' echo "Waiting for the capture buffer to flush..."; sleep 75 aws s3 ls "s3://${BUCKET}/capture/" --recursive | grep '\.jsonl$' | head 2026-09-28 15:47:14 634356 capture/20e4425092f90c7f/pod/AllTraffic/2026/09/28/15/28-15-47-12-999-ac31d31b-5156-46ad-9ccc-e8a040ab9ab2.jsonl
You see one or more .jsonl objects under {s3Uri}/{uniqueID}/pod/{variant}/{YYYY}/{MM}/{DD}/{HH}/. The uniqueID is generated per endpoint. Each line is one inference in the SageMaker data-capture format shown next.
The captured record format
The request and response bodies are stored base64-encoded, as the encoding field shows, so reading them back requires base64-decoding and then parsing JSON:
Decoded, endpointInput.data is the request body the client sent and endpointOutput.data is the response the model returned.
The scenario: A support-ticket router meets a product launch
The deployed model is a support-ticket router for a home-appliance retailer. It uses vLLM guided JSON decoding to classify each ticket into one of six categories: ProductFridge, ProductWasher, ProductOven, ProductDishwasher, ProductMicrowave, or the catch-all Other. Guided decoding means it can only return a valid object with one of those labels:
{"category": "ProductWasher", "confidence": 1}
After the model was deployed, the retailer launched its own line of air conditioners, the Arctis series (Arctis 9K, 12K, Mini, and so on). The router has no category for them. It can’t say “this is new.” It can only pick the closest label it knows. Unless every request and response is recorded, the evidence disappears the moment the response is returned.
The synthetic traffic
We generated two synthetic datasets:
Production traffic: 900 tickets, replayed in order through the capture proxy over 15 minutes, using the router’s original six-label prompt. The first 400 come before the launch and mention no air conditioners. In the remaining 500, the share of Arctis tickets ramps from 15% to 55%.
Evaluation set: 280 labeled tickets, 40 per category, including the new ProductAC. It is held out by content: it uses symptoms, sentence templates, and Arctis model names that never appear in the traffic.
Each production request looks like this (system prompt abbreviated):
{"model": "appliance-router", "temperature": 0, "max_tokens": 40, "messages": [{"role": "system", "content": "You are a support-ticket router for a home-appliance retailer. ... category is one of: ProductFridge, ProductWasher, ProductOven, ProductDishwasher, ProductMicrowave, Other. ..."}, {"role": "user", "content": "Hi, our Arctis 12K blows warm air since last week. Order #66759."}], "guided_json": {"type": "object", "required": ["category", "confidence"], "properties": {"category": {"enum": ["ProductFridge", "ProductWasher", "ProductOven", "ProductDishwasher", "ProductMicrowave", "Other"]}, "confidence": {"type": "number"}}}}
To replay the traffic, clone the accompanying repository, generate the datasets, and send them through the capture proxy. scripts/pod.sh runs the client inside the model container, so no load balancer or DNS is needed. The client prints the three timestamps that the Athena steps use: RUN_START, LAUNCH_TS (the first post-launch ticket), and RUN_END.
git clone https://github.com/awslabs/awsome-distributed-ai.git cd awsome-distributed-ai/examples/use-cases/sagemaker-hyperpod-data-capture-flywheel (cd data && python3 ../scripts/gen_dataset.py) scripts/pod.sh $ENDPOINT_NAME drive /tmp/tickets_traffic.jsonl --seconds 900 RUN_START=2026-09-28 15:52:40 sent 100/900 ... LAUNCH_TS=2026-09-28 15:59:20 ... RUN_END=2026-09-28 16:07:40
Export the three values that your run printed, and wait one flush interval:
Offline benchmarks still look unchanged after the launch, because they are frozen in the past. The shift is visible only in the captured production traffic. This section shows how a data scientist finds it with Athena, without provisioning any infrastructure.
Step 6: Set your variables and configure an Athena helper function
export ATHENA_DB="flywheel_demo" export ATHENA_OUTPUT="s3://${BUCKET}/athena-results/" # Each endpoint writes under its own uniqueID. Pick the prefix whose captured requests name this # endpoint, not whichever prefix sorts first (the Step 15 endpoint adds a second one). for id in $(aws s3 ls "s3://${BUCKET}/capture/" | awk '/PRE/{print $2}'); do key=$(aws s3 ls "s3://${BUCKET}/capture/${id}pod/AllTraffic/" --recursive | awk '/\.jsonl$/{print $4; exit}') [ -n "$key" ] || continue model=$(aws s3 cp "s3://${BUCKET}/${key}" - 2>/dev/null | head -n 1 | python3 -c 'import sys, json, base64; r = json.loads(sys.stdin.readline()); print(json.loads(base64.b64decode(r["captureData"]["endpointInput"]["data"]))["model"])') if [ "$model" = "$ENDPOINT_NAME" ]; then export CAPTURE_LOCATION="s3://${BUCKET}/capture/${id}pod/AllTraffic/"; break; fi done echo "$CAPTURE_LOCATION" # s3://<bucket>/capture/<uniqueID>/pod/AllTraffic/ # run_athena "<SQL>": runs one statement, prints tab-separated rows, and sets LAST_QID run_athena() { local qid st qid=$(aws athena start-query-execution --query-string "$1" \ --query-execution-context "Database=$ATHENA_DB" \ --result-configuration "OutputLocation=$ATHENA_OUTPUT" \ --output text --query 'QueryExecutionId') || return 1 while :; do st=$(aws athena get-query-execution --query-execution-id "$qid" \ --output text --query 'QueryExecution.Status.State') case "$st" in SUCCEEDED) break ;; FAILED|CANCELLED) aws athena get-query-execution --query-execution-id "$qid" \ --output text --query 'QueryExecution.Status.StateChangeReason' >&2; return 1 ;; esac sleep 2 done LAST_QID=$qid aws athena get-query-results --query-execution-id "$qid" \ --output text --query 'ResultSet.Rows[*].Data[*].VarCharValue' } run_athena "CREATE DATABASE IF NOT EXISTS ${ATHENA_DB}"
Step 7: Declare the captures to Athena (two layers)
The first layer is the envelope. An external table describes the outer JSON of each record, and partition projection over year/month/day/hour maps the hourly S3 prefixes to partitions without a crawler or ADD PARTITION. The data fields are declared as plain strings, because they are base64 text.
The second layer is a view that decodes each record once. from_utf8(from_base64(...)) turns the payloads back into JSON. Without it, every extraction returns NULL. The ticket is the user message (messages[1], because messages[0] is the system prompt). The model’s answer is a JSON string inside the response JSON, so it needs a second extraction.
run_athena " CREATE OR REPLACE VIEW preds AS WITH decoded AS ( SELECT CAST(from_iso8601_timestamp(eventMetadata.inferenceTime) AS timestamp) AS ts, from_utf8(from_base64(captureData.endpointInput.data)) AS req, from_utf8(from_base64(captureData.endpointOutput.data)) AS resp FROM inference_captures WHERE year = '${RUN_START:0:4}' AND month = '${RUN_START:5:2}' AND day = '${RUN_START:8:2}' ) SELECT ts, json_extract_scalar(req, '\$.messages[1].content') AS ticket_text, json_extract_scalar(json_extract_scalar(resp, '\$.choices[0].message.content'), '\$.category') AS predicted_category, CAST(json_extract_scalar(json_extract_scalar(resp, '\$.choices[0].message.content'), '\$.confidence') AS double) AS confidence FROM decoded WHERE ts >= timestamp '${RUN_START}' AND ts < timestamp '${RUN_END}' " run_athena "SELECT count(*) AS n, count_if(predicted_category IS NULL) AS unparsed FROM preds" n unparsed 900 0
Two notes on the view:
Always filter on partitions. Athena reads raw JSON, so a query without a partition filter scans every object you have ever captured. The view filters to one day, and the time bound limits it to our traffic run: the same prefix also holds the request from Step 5.
messages[1]assumes one system message followed by one user message. Captured traffic is whatever clients actually sent, so adjust the extraction for multi-turn or few-shot requests.
(Output is tab-separated. We aligned it here for readability. Our run compressed “weeks” into 15 minutes. In production, bucket bydate_trunc('day', ts)or'week'instead.)
The following chart plots the per-minute share of tickets routed to Other as bars, and the cumulative share as a line, with the launch marked.
Figure 2: Per-minute and cumulative share of tickets routed to Other, before and after the Arctis launch
Before the launch, Other stays under 7% per minute. It jumps to 15% in the launch minute and keeps climbing, reaching 47–65% in the last three minutes. The cumulative line in the chart only reaches 22.6%, because it averages the post-launch tickets with all the traffic from before the launch. That is why the per-minute rate, not the cumulative share, is what you would put behind an alarm.
Split at the launch, the whole category mix looks like this:
run_athena " SELECT IF(ts < timestamp '${LAUNCH_TS}', '1-before launch', '2-after launch') AS period, predicted_category, count(*) AS n FROM preds GROUP BY 1, 2 ORDER BY 1, 3 DESC" period predicted_category n 1-before launch ProductFridge 91 1-before launch ProductWasher 83 1-before launch ProductDishwasher 74 1-before launch ProductOven 74 1-before launch ProductMicrowave 59 1-before launch Other 19 2-after launch Other 184 2-after launch ProductFridge 84 2-after launch ProductOven 66 2-after launch ProductDishwasher 64 2-after launch ProductWasher 58 2-after launch ProductMicrowave 44
Other went from 4.8% of tickets before the launch (19 of 400) to 36.8% after it (184 of 500), and became the largest category.
Step 9: What is inside the Other bucket?
run_athena " SELECT word, count(*) AS freq FROM preds CROSS JOIN UNNEST(regexp_extract_all(lower(ticket_text), '[a-z0-9-]{4,}')) AS t(word) WHERE predicted_category = 'Other' AND word NOT IN ('this','that','with','from','have','since','week','help','need','again','time', 'month','second','take','look','someone','order','last','after','when','onto') GROUP BY 1 ORDER BY 2 DESC, 1 LIMIT 10" word freq arctis 141 unit 40 mini 33 cooling 30 link 30 lost 30 loudly 30 rattles 30 wi-fi 30 blows 28
The Other bucket isn’t a mix of unrelated tickets. It’s dominated by one product name, arctis, and the vocabulary of a cooling appliance, which points to a new category the router has no label for.
Step 10: How big is it, and does the model’s confidence flag it?
Search the whole traffic stream for the new product, not just the Other bucket:
run_athena " SELECT IF(ts < timestamp '${LAUNCH_TS}', '1-before launch', '2-after launch') AS period, predicted_category, count(*) AS n FROM preds WHERE regexp_like(lower(ticket_text), 'arctis|air condition') GROUP BY 1, 2 ORDER BY 1, 3 DESC" period predicted_category n 2-after launch Other 164 2-after launch ProductFridge 14
There are 178 air-conditioner tickets, 19.8% of all traffic. None of them appeared before the launch. They fall into two groups:
164 went toOther. Under the six-label schema, that is the correct answer: Other exists for exactly this case. The router isn’t broken. The taxonomy is missing a category, so these tickets reach a generic queue instead of the team that supports the product.
14 went toProductFridge. These are true misroutes: a refrigerator team is now getting air-conditioner tickets.
Now test the model’s confidence on these known misroutes:
run_athena " SELECT predicted_category, confidence, count(*) AS n FROM preds WHERE regexp_like(lower(ticket_text), 'arctis|air condition') GROUP BY 1, 2 ORDER BY 1, 2 DESC" predicted_category confidence n Other 1.0 157 Other 0.9 7 ProductFridge 0.95 1 ProductFridge 0.9 9 ProductFridge 0.8 2 ProductFridge 0.7 1 ProductFridge 0.0 1
Of the 14 misrouted tickets, 10 reported a confidence of 0.9 or higher. A labeling rule such as “keep predictions with confidence >= 0.9” would have kept 71% of the known-wrong labels. That’s because confidence is a field we asked the model to write, not a calibrated probability. If you need a real uncertainty signal, request token logprobs from vLLM and read the probability at the category token instead of the model’s self-report. Budget for the extra payload size, because the logprobs are captured with every response.
From diagnosis to labeled training data
A data scientist doesn’t hand-label hundreds of tickets. Instead, they write the Athena finding down as a labeling policy: a SQL view that assigns a corrected label to every captured record. We relabel the misrouted air-conditioner tickets to the new class.
Step 11: The labeling policy
run_athena " CREATE OR REPLACE VIEW labeled_training_data AS WITH t AS (SELECT ticket_text, predicted_category, lower(ticket_text) AS txt FROM preds) SELECT ticket_text, predicted_category, CASE -- 1. The new class, identified by keyword evidence, not by what the model predicted WHEN regexp_like(txt, '\\barctis\\b|air condition|\\bac unit\\b|mini split') AND NOT regexp_like(txt, 'fridge|refrigerat|freezer|washer|\\boven\\b|\\brange\\b|dishwasher|microwave') THEN 'ProductAC' -- 2. Keep the production decision when nothing hints at the new product WHEN NOT regexp_like(txt, '\\barctis\\b|air condition|\\bac\\b|mini split') THEN predicted_category -- 3. Mixed signals: send to human review (Amazon SageMaker Ground Truth) ELSE NULL END AS corrected_category FROM t"
The doubled backslashes (\\b) are shell escaping inside the double-quoted string. Athena receives \b, a regular-expression word boundary.
Check what the policy produced:
run_athena " SELECT coalesce(corrected_category, '(human review)') AS label, count(*) AS n, count(DISTINCT ticket_text) AS distinct_texts, count_if(corrected_category <> predicted_category) AS relabeled FROM labeled_training_data GROUP BY 1 ORDER BY 2 DESC" label n distinct_texts relabeled ProductAC 178 114 178 ProductFridge 161 110 0 ProductWasher 141 96 0 ProductOven 140 97 0 ProductDishwasher 138 95 0 ProductMicrowave 103 82 0 Other 39 8 0
This heuristic works because the new category is clear-cut (keyword-identifiable). For ambiguous tickets or overlapping categories, the NULL rows would route to Ground Truth for human labeling.
Step 12: Build the training file
Export the labeled view and convert it to chat-format JSONL. The allowed category set now includes ProductAC, the same seven-label prompt clients will send after the redeploy. Save this as build_training_file.py:
cat > build_training_file.py <<'PY' import csv, json, sys CATS = ["ProductFridge", "ProductWasher", "ProductOven", "ProductDishwasher", "ProductMicrowave", "ProductAC", "Other"] PROMPT = ("You are a support-ticket router for a home-appliance retailer. Route each ticket to the product it " 'concerns. Respond with ONLY a JSON object {"category": "...", "confidence": <0-1>} where category is ' "one of: " + ", ".join(CATS) + '. Use "Other" if the ticket does not clearly concern one of these products.') n = 0 with open(sys.argv[1], newline="") as fin, open(sys.argv[2], "w") as fout: for row in csv.DictReader(fin): fout.write(json.dumps({"messages": [ {"role": "system", "content": PROMPT}, {"role": "user", "content": row["ticket_text"]}, {"role": "assistant", "content": json.dumps({"category": row["corrected_category"], "confidence": 1.0})}, ]}) + "\n"); n += 1 print(f"wrote {n} training examples -> {sys.argv[2]}") PY run_athena "SELECT ticket_text, corrected_category FROM labeled_training_data WHERE corrected_category IS NOT NULL" > /dev/null aws s3 cp "${ATHENA_OUTPUT}${LAST_QID}.csv" labeled.csv --only-show-errors # Athena's own CSV: header + proper quoting python3 build_training_file.py labeled.csv train.jsonl aws s3 cp train.jsonl "s3://${BUCKET}/training/train.jsonl" wrote 900 training examples -> train.jsonl
Fine-tuning on captured data with the HyperPod training operator
We fine-tune on the same cluster that serves the model, using the HyperPod training operator on the second GPU. The first endpoint keeps serving throughout.
Step 13: The training script
The script downloads the base model and training data from S3 and trains a low-rank adaptation (LoRA) adapter. The loss is computed only on the assistant turn, so the model learns to produce the label rather than repeat the prompt. The script then merges the adapter back into the base weights, so vLLM can serve the result unchanged, and uploads the merged model to S3. The job mounts the script into the training pod from a ConfigMap.
cat > train.py <<'PY' import os, boto3, torch from datasets import load_dataset from peft import LoraConfig, get_peft_model from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments s3 = boto3.client("s3") def split(uri): b, _, k = uri[5:].partition("/"); return b, k # 1. Stage base model + training data from S3 bucket, prefix = split(os.environ["BASE_MODEL_S3"]) for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix): for o in page.get("Contents", []): dst = os.path.join("/work/base", os.path.relpath(o["Key"], prefix)) os.makedirs(os.path.dirname(dst), exist_ok=True); s3.download_file(bucket, o["Key"], dst) s3.download_file(*split(os.environ["TRAIN_FILE_S3"]), "/work/train.jsonl") print("[stage] base model and training data downloaded", flush=True) # 2. LoRA SFT, loss only on the assistant turn (the label), not the prompt tok = AutoTokenizer.from_pretrained("/work/base") def encode(ex): full = tok.apply_chat_template(ex["messages"], tokenize=False) prompt = tok.apply_chat_template(ex["messages"][:-1], tokenize=False, add_generation_prompt=True) ids = tok(full, add_special_tokens=False)["input_ids"] n = len(tok(prompt, add_special_tokens=False)["input_ids"]) return {"input_ids": ids, "labels": [-100] * n + ids[n:]} ds = load_dataset("json", data_files="/work/train.jsonl", split="train") print(f"[data] {len(ds)} examples", flush=True) ds = ds.map(encode, remove_columns=ds.column_names) def collate(batch): m = max(len(b["input_ids"]) for b in batch) pad = lambda xs, v: [x + [v] * (m - len(x)) for x in xs] return {"input_ids": torch.tensor(pad([b["input_ids"] for b in batch], tok.pad_token_id)), "labels": torch.tensor(pad([b["labels"] for b in batch], -100)), "attention_mask": torch.tensor(pad([[1] * len(b["input_ids"]) for b in batch], 0))} model = AutoModelForCausalLM.from_pretrained("/work/base", torch_dtype=torch.bfloat16, device_map="cuda:0") model = get_peft_model(model, LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])) model.print_trainable_parameters() Trainer(model=model, train_dataset=ds, data_collator=collate, args=TrainingArguments( output_dir="/tmp/ckpt", num_train_epochs=float(os.getenv("EPOCHS", "2")), learning_rate=1e-4, per_device_train_batch_size=8, bf16=True, logging_steps=10, save_strategy="no", report_to=[], warmup_ratio=0.1, lr_scheduler_type="cosine")).train() # 3. Merge the adapter into the base weights (a plain model vLLM can serve) and upload it model.merge_and_unload().save_pretrained("/work/out", safe_serialization=True) tok.save_pretrained("/work/out") bucket, prefix = split(os.environ["OUTPUT_S3"]) for f in os.listdir("/work/out"): s3.upload_file(os.path.join("/work/out", f), bucket, f"{prefix.rstrip('/')}/{f}") print(f"[done] merged model uploaded to {os.environ['OUTPUT_S3']}", flush=True) PY kubectl create configmap router-train --from-file=train.py
logMonitoringConfigurationis the training operator’s distinguishing feature. If no loss line appears within 20 minutes of start, or the log goes quiet for 10 minutes mid-run while the process is still alive, the operator restarts the job instead of letting it burn GPU hours.
The job installs three Python packages when it starts. The training node needs outbound internet access, for example through a NAT gateway, and so does pulling the vllm/vllm-openai image from Docker Hub. On a cluster without internet egress, bake the packages into a custom image in Amazon Elastic Container Registry (Amazon ECR) and drop the pip install.
Training took 60 seconds on one L40S. LoRA at rank 16 trains about 29.9 million parameters, under 1% of the 3.1-billion-parameter model. The average training loss was 0.0123, and the per-step loss reached 0.0 within the first epoch, which shows the model memorizes 900 templated examples readily. On real data, hold out a validation split and stop early. Here we rely on the held-out evaluation set.
Step 15: Deploy the fine-tuned model side by side, with capture still on
Instead of editing the running endpoint, deploy the new model as a second endpoint from the same manifest, onto the GPU the training job just released. The first endpoint keeps serving until you cut clients over, so the redeploy causes no downtime, and rolling back means pointing clients back at the first endpoint.
The new endpoint writes captures under its own uniqueID. Because ENDPOINT_NAME is now appliance-router-v2, rerunning the CAPTURE_LOCATION lookup from Step 6 finds that prefix, so you can point a second Athena table at it for the next round of analysis. When you move clients over, update the endpoint, the label set, and the guided-JSON enum in the same release.
Measuring the lift
We ran the same 280 held-out tickets against the fine-tuned endpoint:
Fine-tuning taught the router the new category and caused no regression on the five established classes. It also fixed five Other errors that existed before the launch which the launched router sent to ProductFridge.
Because our evaluation set comes from the same generator as the training data, we also wrote 24 probe tickets by hand, in styles the generator never produces. The fine-tuned model scored 20/24:
MISS ProductFridge -> ProductAC: The cooling fan in my fridge is squealing. MISS Other -> ProductAC: When will the Arctis 12K be back in stock? MISS Other -> ProductAC: Can I get a quote to install central air in my house? MISS Other -> ProductAC: Do you price match the Arctis at other stores? overall 20/24 = 83.33%
It routes all eight air-conditioner probes correctly, but it learned two shortcuts: “cooling” means air conditioner, and any mention of Arctis means ProductAC, including stock and price-match questions. Both come from the training data: the labeling policy sent every Arctis ticket to ProductAC, and the captured traffic contained no pre-sales questions. With capture on for the new endpoint, those tickets are exactly what the next round of capture collects.
How much latency does capture add?
We measured this directly: the same 400 routing requests, sent from inside the pod either to port 8000 (no capture) or to port 8081 (through the capture proxy). We ran three rounds at 8 concurrent requests, plus one serial round:
Path
Concurrency
p50
p99
Throughput
Direct to vLLM (8000), no capture
8
200.7 ms
249.4 ms
39.1 req/s
Through capture proxy (8081)
8
200.0 ms
249.8 ms
39.1 req/s
Direct to vLLM (8000), no capture
1
164.4 ms
206.9 ms
6.0 req/s
Through capture proxy (8081)
1
164.4 ms
207.0 ms
6.0 req/s
Rows at concurrency 8 are the average of three rounds. Round-to-round variation (±1 ms at p50) was larger than the difference between the two paths.
Other uses for captured data
After inference data lands in S3, the same records feed many workflows:
Improve the model: Use supervised fine-tuning (shown here), distill a large model’s outputs into a smaller model, or build preference pairs for direct preference optimization (DPO).
Understand the model: Run Athena analytics (shown here), build golden regression sets from real prompts, set drift alarms, discover emerging intents, and detect malformed outputs and refusals.
Operate and govern: Route ambiguous samples to Ground Truth for labeling, keep audit trails (Tier 1 is recommended), replay the exact input behind an incident, and analyze traffic and payloads for cost attribution.
