Fix SageMaker MME cross-model invoke/unload via X-Amzn-SageMaker-Target-Model header (TRI-1563) by akhilraj9 · Pull Request #8923 · triton-inference-server/server

akhilraj9 · 2026-08-10T08:50:40Z

Summary

The SageMaker multi-model endpoint (MME) adapter authorized invoke and unload
requests using the model identity in the URL (/models/<hash>), but performed
the actual inference / unload using the X-Amzn-SageMaker-Target-Model header,
without verifying the two referred to the same model. A request scoped to one
loaded model could therefore invoke or unload a different loaded model, and the
unload path erased the URL-hash registry entry while unloading the header-selected
model — leaving runtime and repository state inconsistent.

This is a confused-deputy / incorrect-authorization issue (CWE-441 / CWE-863). In
multi-tenant or shared-model deployments it breaks tenant isolation (cross-model
inference, cross-model unload/DoS) and corrupts model lifecycle bookkeeping.

Root cause

The registry (sagemaker_models_list_) stored only hash → repo_path, discarding
which model each hash was actually loaded as. At invoke/unload time there was
nothing to validate the header against, so the header was trusted as the action
target.

Fix — bind identity at load, validate on use

Extend the registry value from repo_path to { repo_path, target_model }
(new SageMakerModelInfo), recording the model each hash was loaded under.
On invoke and unload, act on the registered target_model, never on the
header. If the request carries an X-Amzn-SageMaker-Target-Model header that
does not match the registered value, reject with 400.
Unload now evicts the registered model and erases that same entry, so runtime
and repository state stay consistent.

Legitimate traffic is unaffected: SageMaker sets the header consistently with the
loaded model (or omits it), so the mismatch path only triggers on a contradictory
request.

Changes

src/sagemaker_server.{h,cc} — the fix
qa/L0_sagemaker/sagemaker_multi_model_test.py — regression test
(test_sm_5b_identity_confusion_poc): mismatched header on invoke and unload
must return 400, and the other model must remain loaded
qa/L0_sagemaker/test.sh — bump expected MME test count 7 → 8

Testing

L0_sagemaker passes (8/8), including the new regression test. Verified end-to-end
on the internal CI pipeline (build + L0_sagemaker green).
Pipeline ID: 61910130

…ader (TRI-1563) Validate X-Amzn-SageMaker-Target-Model against the target_model registered for the URL hash and reject mismatches with 400.

Asserts invoke and unload reject a mismatched X-Amzn-SageMaker-Target-Model header with 400 and leave the other model loaded.

greptile-apps · 2026-08-10T08:53:02Z

Greptile Summary

The PR binds SageMaker MME invoke and unload operations to the model identity recorded during load, rejecting contradictory target-model headers.

Stores each loaded model’s repository path and target identity together.
Uses the registered identity for inference and unload operations.
Adds exact status and response-body assertions for mismatched invoke and unload requests.
Updates the expected SageMaker MME test count.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
src/sagemaker_server.cc	Resolves invoke and unload targets from registry state and rejects mismatched target-model headers before performing model operations.
src/sagemaker_server.h	Extends each MME registry entry to retain both its repository path and load-time target identity.
qa/L0_sagemaker/sagemaker_multi_model_test.py	Adds regression coverage with exact mismatch-specific status and response assertions and verifies the other model remains loaded.
qa/L0_sagemaker/test.sh	Updates the expected multi-model test count to include the new regression test.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[MME request URL hash] --> B[Lookup registered model]
  B --> C{Target-Model header present?}
  C -- No --> E[Use registered target model]
  C -- Yes --> D{Header matches registered identity?}
  D -- No --> F[Return HTTP 400]
  D -- Yes --> E
  E --> G[Invoke or unload registered model]

_{Reviews (5): Last reviewed commit: "review: address MME target-model review ..." | Re-trigger Greptile}

…TRI-1563) Per review: the PoC accepted any status >= 400; assert the specific 400 and the mismatch error payload so the test proves the identity-check branch ran.

… year The sys.path-append-before-imports pattern (pre-existing since 2022) trips flake8 E402 once the PR touches the file; suppress with # noqa: E402 and update the copyright year to 2026.

…psirt-trtllm-multi-model-identity-confusion-sagemaker # Conflicts: # qa/L0_sagemaker/test.sh

yinggeh · 2026-08-25T07:08:34Z

+    EVBufferAddErrorJson(req->buffer_out, err);
+    evhtp_send_reply(req, EVHTP_RES_BADREQ); /* 400 */
+    TRITONSERVER_ErrorDelete(err);


Can use macro HTTP_RESPOND_IF_ERR

Good call — switched both the invoke and unload rejection paths to HTTP_RESPOND_IF_ERR. Same 400 + JSON body, a lot less boilerplate. Done.

yinggeh · 2026-08-25T07:10:16Z

+    /* Use model name hash as the key, as expected in the SageMaker MME
+     * contract, and remember the target_model the repository was actually
+     * loaded under so invoke/unload can validate the request header against it
+     * (TRI-1563). */


I wouldn't put ticket number unless there is future work.

Removed the ticket number here and in the two related comments. Done.

yinggeh · 2026-08-25T07:11:37Z

+  struct SageMakerModelInfo {
+    // Parent path of the registered model repository
+    // (e.g. /opt/ml/models/<hash>). Needed to unregister the repo on unload.
+    std::string repo_path;


Made both repo_path and target_model const — they're set once at load and only read after. Done.

yinggeh · 2026-08-25T07:12:52Z

            "Expected status code 404, received {}".format(r.status_code),
        )

+    def test_sm_5b_identity_confusion_poc(self):


What's sm_5b?

sm = sagemaker; 5b just orders it right after test_sm_5_model_unload (unittest runs tests lexicographically). Renamed to test_sm_5b_target_model_header_mismatch with a docstring so it's self-explanatory, and dropped "poc" — it's a permanent regression test now, not a one-off.

yinggeh · 2026-08-25T07:13:33Z

-import tritonclient.http as httpclient
+import numpy as np  # noqa: E402
+import requests  # noqa: E402
+import test_util as tu  # noqa: E402


Move sys.path.append("../common") before this line for fewer # noqa: E402.

Done — moved the stdlib/third-party imports above sys.path.append("../common") so only test_util needs # noqa: E402 now (8 → 1). Matches the existing pattern in e.g. L0_cuda_graph/trt_cuda_graph_test.py.

Use HTTP_RESPOND_IF_ERR macro, const struct members, drop ticket refs, rename test_sm_5b, reduce noqa. No behavior change.

akhilraj9 added 2 commits August 9, 2026 08:51

fix: bind SageMaker MME invoke/unload to the loaded model, not the he…

e0cb40a

…ader (TRI-1563) Validate X-Amzn-SageMaker-Target-Model against the target_model registered for the URL hash and reject mismatches with 400.

test: add TRI-1563 SageMaker MME identity-confusion regression test

2281637

Asserts invoke and unload reject a mismatched X-Amzn-SageMaker-Target-Model header with 400 and leave the other model loaded.

greptile-apps Bot reviewed Aug 10, 2026

View reviewed changes

Comment thread qa/L0_sagemaker/sagemaker_multi_model_test.py Outdated

akhilraj9 assigned pskiran1 and unassigned pskiran1 Aug 10, 2026

akhilraj9 requested review from pskiran1, whoisj and yinggeh August 10, 2026 10:39

akhilraj9 added 3 commits August 10, 2026 14:06

test: assert exact 400 + error payload on MME target-model mismatch (…

0d544cc

…TRI-1563) Per review: the PoC accepted any status >= 400; assert the specific 400 and the mismatch error payload so the test proves the identity-check branch ran.

test: noqa flake8 E402 on sagemaker MME test imports + bump copyright…

da09e04

… year The sys.path-append-before-imports pattern (pre-existing since 2022) trips flake8 E402 once the PR touches the file; suppress with # noqa: E402 and update the copyright year to 2026.

Merge remote-tracking branch 'origin/main' into asaraswathi/tri-1563-…

fd626c9

…psirt-trtllm-multi-model-identity-confusion-sagemaker # Conflicts: # qa/L0_sagemaker/test.sh

akhilraj9 self-assigned this Aug 24, 2026

yinggeh reviewed Aug 25, 2026

View reviewed changes

review: address MME target-model review comments

efcb811

Use HTTP_RESPOND_IF_ERR macro, const struct members, drop ticket refs, rename test_sm_5b, reduce noqa. No behavior change.

Navigation Menu

Uh oh!

Fix SageMaker MME cross-model invoke/unload via X-Amzn-SageMaker-Target-Model header (TRI-1563) - #8923

Fix SageMaker MME cross-model invoke/unload via X-Amzn-SageMaker-Target-Model header (TRI-1563)#8923
akhilraj9 wants to merge 6 commits into
mainfrom
asaraswathi/tri-1563-psirt-trtllm-multi-model-identity-confusion-sagemaker

akhilraj9 commented Aug 10, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Aug 10, 2026 •

edited

Loading

Uh oh!

Uh oh!

yinggeh Aug 25, 2026

Uh oh!

akhilraj9 Aug 26, 2026

Uh oh!

yinggeh Aug 25, 2026

Uh oh!

akhilraj9 Aug 26, 2026

Uh oh!

yinggeh Aug 25, 2026

Uh oh!

akhilraj9 Aug 26, 2026

Uh oh!

yinggeh Aug 25, 2026

Uh oh!

akhilraj9 Aug 26, 2026

Uh oh!

yinggeh Aug 25, 2026

Uh oh!

akhilraj9 Aug 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

akhilraj9 commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Root cause

Fix — bind identity at load, validate on use

Changes

Testing

Uh oh!

greptile-apps Bot commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Milestone

Development

Uh oh!

3 participants

akhilraj9 commented Aug 10, 2026 •

edited

Loading

greptile-apps Bot commented Aug 10, 2026 •

edited

Loading