Building and Serving a Custom Model with Amazon SageMaker, Then Wiring It Into a Bedrock Agent
Train a custom model on Amazon SageMaker, deploy it to a real-time endpoint, then wire it into a Bedrock Agent via Return of Control so the agent can call your model mid-conversation.
Bedrock's model catalog covers a lot of ground, but it doesn't cover everything. If you need a model trained on your own proprietary, tabular data, a churn predictor built on your actual customer history, a fraud score trained on your actual transaction patterns, that's not a foundation model problem. That's a real machine learning problem, and on AWS it belongs to Amazon SageMaker, a genuinely separate platform from Bedrock with its own SDK, its own training and hosting model, and its own way of thinking about a model's lifecycle.
One note before starting, AWS has been sunsetting several older SageMaker AI features, Ground Truth Plus has already ended, and Ground Truth, Clarify, Debugger, and Model Monitor have moved to a no-new-customers state as of July 30, 2026. None of that touches what this article actually uses. Training jobs, real-time endpoints, and Studio are explicitly unaffected, and that's the whole path built below.
This is a hands-on build of the whole path: train a model with the SageMaker Python SDK, deploy it behind a real-time endpoint, and wire that endpoint into a Bedrock agent using return of control, so a conversational agent can call your custom model mid-conversation the same way it would call any other tool.
The mental model first
SageMaker and Bedrock don't share a runtime. Bedrock is built for foundation models, prompting, and retrieval. SageMaker is built for the full custom ML lifecycle, and the two meet at exactly one point in this pipeline, worth being precise about where:
- An Estimator wraps a training script and a compute spec. Calling
.fit()on it runs your training code on managed infrastructure and tracks the job. .deploy()turns a trained model into a live, real-time inference endpoint, aSagemakerRuntimeclient target with its own name, its own instance type, and its own cost, running inside your SageMaker environment, not inside Bedrock.- Return of control is Bedrock's mechanism for client-side tool execution. Instead of a Lambda function running server-side when the agent decides to use a tool, Bedrock hands the tool call back to your application, in the
returnControlfield of anInvokeAgentresponse, and your code executes it and sends the result back in a follow-up call. This is the direct Bedrock equivalent of a client-defined function tool, and it's the mechanism this whole integration is built on.
Prerequisites
- Access to Amazon Bedrock Agents and a SageMaker execution role with permission to train and host.
- Python 3.9+ with both SDKs installed.
1
pip install sagemaker boto3
1
2
3
4
5
import sagemaker
import boto3
session = sagemaker.Session()
role = "arn:aws:iam::123456789012:role/SageMakerExecutionRole"
Step 1: train the model
The SKLearn Estimator wraps a plain scikit-learn training script and runs it on managed infrastructure.
1
2
3
4
5
6
7
8
9
10
11
12
from sagemaker.sklearn.estimator import SKLearn
sk_estimator = SKLearn(
entry_point="train.py",
role=role,
instance_type="ml.m5.large",
instance_count=1,
framework_version="1.2-1",
hyperparameters={"n_estimators": 200, "max_depth": 6},
)
sk_estimator.fit({"train": "s3://acme-ml-data/churn-training/"})
train.py is ordinary scikit-learn code, load the data from the path SageMaker mounts it at, fit a model, and serialize it to the output directory SageMaker expects. The Estimator handles provisioning the training instance, running the script, and uploading the resulting model artifact to S3, none of that is custom infrastructure you build yourself.Step 2: deploy it behind a real-time endpoint
1
2
3
4
5
predictor = sk_estimator.deploy(
initial_instance_count=1,
instance_type="ml.m5.large",
endpoint_name="churn-endpoint",
)
One call, and
.deploy() handles creating the SageMaker model, the endpoint configuration, and the endpoint itself. Worth knowing this is three separate resources under the hood, since production tooling often manages them individually, swapping a new model into an existing endpoint config without recreating the endpoint, but the SDK collapses all of it into one method for the common case.Step 3: confirm it works before anything else touches it
1
2
3
4
5
6
7
8
9
10
11
12
import json
import boto3
runtime = boto3.client("sagemaker-runtime")
test_payload = {"tenure_months": 14, "monthly_spend": 89.50, "support_tickets": 3}
response = runtime.invoke_endpoint(
EndpointName="churn-endpoint",
ContentType="application/json",
Body=json.dumps(test_payload),
)
print(json.loads(response["Body"].read()))
Get a real prediction back directly through
sagemaker-runtime before wiring anything else into it. Debugging a broken endpoint through the extra layer of an agent's return-control loop is meaningfully harder than debugging it directly.Step 4: define the Bedrock agent with a return-of-control action group
This is the actual bridge, and the key detail is what's absent: no Lambda function, no API schema file. The action group just declares a function's name, description, and parameters, and tells Bedrock to hand control back to the caller instead of executing anything itself.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
bedrock_agent = boto3.client("bedrock-agent")
agent = bedrock_agent.create_agent(
agentName="retention-agent",
foundationModel="anthropic.claude-sonnet-4-20250514-v1:0",
agentResourceRoleArn="arn:aws:iam::123456789012:role/BedrockAgentExecutionRole",
instruction="Help the team assess churn risk. Call get_churn_score whenever specific customer numbers are provided.",
)
agent_id = agent["agent"]["agentId"]
bedrock_agent.create_agent_action_group(
agentId=agent_id,
agentVersion="DRAFT",
actionGroupName="churn-scoring",
actionGroupExecutor={"customControl": "RETURN_CONTROL"},
functionSchema={
"functions": [
{
"name": "get_churn_score",
"description": "Predict churn probability for a customer given tenure, spend, and support ticket history.",
"parameters": {
"tenure_months": {"type": "integer", "description": "Months active.", "required": True},
"monthly_spend": {"type": "number", "description": "Average monthly spend in dollars.", "required": True},
"support_tickets": {"type": "integer", "description": "Support tickets in the last 90 days.", "required": True},
},
}
]
},
)
bedrock_agent.prepare_agent(agentId=agent_id)
Step 5: run it, and handle the return-control loop yourself
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import uuid
bedrock_runtime = boto3.client("bedrock-agent-runtime")
def get_churn_score(tenure_months, monthly_spend, support_tickets):
payload = {"tenure_months": tenure_months, "monthly_spend": monthly_spend, "support_tickets": support_tickets}
resp = runtime.invoke_endpoint(EndpointName="churn-endpoint", ContentType="application/json", Body=json.dumps(payload))
return json.loads(resp["Body"].read())
session_id = str(uuid.uuid4())
response = bedrock_runtime.invoke_agent(
agentId=agent_id,
agentAliasId="TSTALIASID",
sessionId=session_id,
inputText="A customer's been with us 14 months, spends about $90/month, and filed 3 tickets recently. Churn risk?",
)
for event in response["completion"]:
if "returnControl" in event:
rc = event["returnControl"]
invocation_id = rc["invocationId"]
func_input = rc["invocationInputs"][0]["functionInvocationInput"]
params = {p["name"]: p["value"] for p in func_input["parameters"]}
result = get_churn_score(
int(params["tenure_months"]), float(params["monthly_spend"]), int(params["support_tickets"])
)
follow_up = bedrock_runtime.invoke_agent(
agentId=agent_id,
agentAliasId="TSTALIASID",
sessionId=session_id,
sessionState={
"invocationId": invocation_id,
"returnControlInvocationResults": [
{
"functionResult": {
"actionGroup": func_input["actionGroup"],
"function": func_input["function"],
"responseBody": {"TEXT": {"body": json.dumps(result)}},
}
}
],
},
)
for chunk_event in follow_up["completion"]:
if "chunk" in chunk_event:
print(chunk_event["chunk"]["bytes"].decode("utf-8"), end="")
The model decides whether the question warrants calling
get_churn_score at all. Your code executes the actual call to the SageMaker endpoint when it does, and the result flows back into the same session, keyed by invocationId, for the model to finish answering with. Nothing about the agent definition changes based on what the function actually does behind the scenes, a database call, a REST API, or, as here, a real-time SageMaker endpoint.Where this fits in the bigger picture
Worth being explicit about something the SDKs don't make obvious on their own: the Bedrock agent and the SageMaker endpoint are two different services with two different boto3 clients,
bedrock-agent-runtime and sagemaker-runtime, and nothing about return of control merges them. Bedrock's returnControl payload tells your application what to call and with what arguments. Your application is entirely responsible for actually calling SageMaker, authenticating to it, and shaping the result back into the format Bedrock expects. There's no shared trust boundary here at all, and that's deliberate, not a gap.That separation means the SageMaker side, training, versioning, redeploying the model, can be owned and iterated on entirely independently of the Bedrock side, by a different team if that's how your organization is structured, without either side needing access to the other's IAM roles.
Production considerations before you commit
- Return-control sessions have a time budget. The agent's turn is waiting on your application to execute the function and send the result back. If the SageMaker endpoint is slow, cold-started, or under-provisioned, that latency is now inside the user's perceived response time for the whole conversation turn, not hidden behind an async job.
- Scope IAM roles separately, and narrowly. The role your application uses to call
sagemaker-runtime:InvokeEndpointshould be scoped to that one endpoint, not broad SageMaker access, and it's a different role entirely from the Bedrock agent's execution role. Treat them as two independent security boundaries because that's what they are. - Blue/green the endpoint, don't overwrite it. Updating
EndpointConfigto point at a new model version, with traffic shifted gradually rather than swapped all at once, is what catches a regression in a newly trained model before it's serving every request. - Real-time endpoint instances are billed whether or not they're scoring. An always-on
ml.m5.largeinstance costs money at idle the same way a SageMaker endpoint always does. If call volume is low and bursty, look at Serverless Inference or asynchronous inference before defaulting to a real-time endpoint sized for peak load. - Parse
invocationInputsdefensively. The example above assumes exactly one function invocation perreturnControlevent, which is the common case, but Bedrock's schema allows for more. Don't hard-code an index assumption into production code without checking what's actually in the list. - Model drift doesn't announce itself. Nothing in this pipeline retrains automatically when real-world data shifts away from what the model was trained on. Log the inputs and outputs of every
get_churn_scorecall somewhere reviewable, and revisit training data on a real cadence, not only once someone notices predictions have gotten worse.
Where this leaves you
Bedrock's model catalog is the right tool for the overwhelming majority of generative AI work, and for teams that never need a model trained on their own proprietary data, it's genuinely the whole story. But the moment the problem is "predict something specific about our own customers, from our own historical data," that's SageMaker's job, and pretending otherwise means either forcing a foundation model to approximate a task it was never built for, or building custom training and serving infrastructure by hand. The actual integration effort here is small: a training job, a deployed endpoint, and one return-of-control action group. The two services don't need to merge for that to work. They just both need to keep doing the one thing each is actually good at, and let a well-defined handoff, your own application code parsing
returnControl and calling invoke_endpoint, do the rest.References
- Amazon Web Services. "Train a Model with a SageMaker Estimator." Amazon SageMaker Developer Guide. docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html
- Amazon Web Services. "Deploy models for real-time inference." Amazon SageMaker Developer Guide. docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints.html
- Amazon Web Services. "Return control to the agent developer by sending elicited information in an InvokeAgent response." Amazon Bedrock User Guide. docs.aws.amazon.com/bedrock/latest/userguide/agents-returncontrol.html
- Amazon Web Services. "InvokeAgent." Amazon Bedrock Agent Runtime API Reference. docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_InvokeAgent.html
- Amazon Web Services. "Chat Completions API on the bedrock-mantle endpoint." Amazon Bedrock User Guide. docs.aws.amazon.com/bedrock/latest/userguide/inference-chat-completions-mantle.html
Enjoyed reading this content? Let the author know!
Your likes, comments, shares, and saves help creators reach more builders.
Loading recommendations
Loading article