Skip to content

Commit a5b45da

Browse files
mattwittwermattwittwer
authored
fix: apply HTTP restrictions to the SageMaker endpoint (#8686)
1 parent b14c139 commit a5b45da

4 files changed

Lines changed: 231 additions & 5 deletions

File tree

‎qa/L0_sagemaker/test.sh‎

Lines changed: 191 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#!/bin/bash
2-
# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved.
2+
# Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved.
33
#
44
# Redistribution and use in source and binary forms, with or without
55
# modification, are permitted provided that the following conditions
@@ -612,6 +612,196 @@ set -e
612612
kill $SERVER_PID
613613
wait $SERVE_PID
614614

615+
### Restricted API regression for SageMaker endpoint ###
616+
# Verify that --http-restricted-api applies to SageMaker model management
617+
# endpoints (load, unload, list, get) while leaving health and inference
618+
# unrestricted.
619+
620+
SERVER_LOG="./sagemaker_restricted_api_server.log"
621+
SERVER_ARGS="--allow-sagemaker=true --allow-http=true \
622+
--allow-grpc=false --allow-metrics=false \
623+
--model-repository=`pwd`/models \
624+
--model-control-mode=explicit \
625+
--load-model=sm_model \
626+
--http-restricted-api=model-repository:X-SM-Auth=secret"
627+
run_server_nowait
628+
sagemaker_wait_for_server_ready $SERVER_PID 10
629+
if [ "$WAIT_RET" != "0" ]; then
630+
echo -e "\n***\n*** Failed to start $SERVER\n***"
631+
kill $SERVER_PID || true
632+
cat $SERVER_LOG
633+
exit 1
634+
fi
635+
636+
set +e
637+
638+
# Health should succeed without restricted header
639+
rm -f ./curl.out
640+
code=`curl -s -w %{http_code} -o ./curl.out localhost:8080/ping`
641+
if [ "$code" != "200" ]; then
642+
cat ./curl.out
643+
echo -e "\n***\n*** Failed. Expected /ping to succeed without restricted header\n***"
644+
RET=1
645+
fi
646+
647+
# Inference should succeed without restricted header
648+
rm -f ./curl.out
649+
code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8080/invocations \
650+
-H "Content-Type: application/json" \
651+
-d '{"inputs":[{"name":"INPUT0","datatype":"INT32","shape":[1,16],"data":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]},{"name":"INPUT1","datatype":"INT32","shape":[1,16],"data":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]}]}'`
652+
if [ "$code" != "200" ]; then
653+
cat ./curl.out
654+
echo -e "\n***\n*** Failed. Expected /invocations inference to succeed without restricted header\n***"
655+
RET=1
656+
fi
657+
658+
# List models without auth header should be blocked
659+
rm -f ./curl.out
660+
code=`curl -s -w %{http_code} -o ./curl.out localhost:8080/models`
661+
if [ "$code" != "403" ]; then
662+
cat ./curl.out
663+
echo -e "\n***\n*** Failed. Expected GET /models to return 403 without restricted header (got $code)\n***"
664+
RET=1
665+
else
666+
grep "This API is restricted" ./curl.out
667+
if [ $? -ne 0 ]; then
668+
cat ./curl.out
669+
echo -e "\n***\n*** Failed. Expected restriction error message in response body\n***"
670+
RET=1
671+
fi
672+
fi
673+
674+
# Get model without auth header should be blocked
675+
rm -f ./curl.out
676+
code=`curl -s -w %{http_code} -o ./curl.out localhost:8080/models/sm_model`
677+
if [ "$code" != "403" ]; then
678+
cat ./curl.out
679+
echo -e "\n***\n*** Failed. Expected GET /models/<name> to return 403 without restricted header (got $code)\n***"
680+
RET=1
681+
fi
682+
683+
# Load model without auth header should be blocked
684+
rm -f ./curl.out
685+
code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8080/models \
686+
-H "Content-Type: application/json" \
687+
-d '{"model_name":"test","url":"/opt/ml/models/123/model"}'`
688+
if [ "$code" != "403" ]; then
689+
cat ./curl.out
690+
echo -e "\n***\n*** Failed. Expected POST /models (load) to return 403 without restricted header (got $code)\n***"
691+
RET=1
692+
fi
693+
694+
# Unload model without auth header should be blocked
695+
rm -f ./curl.out
696+
code=`curl -s -w %{http_code} -o ./curl.out -X DELETE localhost:8080/models/sm_model`
697+
if [ "$code" != "403" ]; then
698+
cat ./curl.out
699+
echo -e "\n***\n*** Failed. Expected DELETE /models (unload) to return 403 without restricted header (got $code)\n***"
700+
RET=1
701+
fi
702+
703+
# List models WITH correct auth header should succeed
704+
rm -f ./curl.out
705+
code=`curl -s -w %{http_code} -o ./curl.out -H "X-SM-Auth: secret" localhost:8080/models`
706+
if [ "$code" == "403" ]; then
707+
cat ./curl.out
708+
echo -e "\n***\n*** Failed. Expected GET /models with auth header to pass restriction check\n***"
709+
RET=1
710+
fi
711+
712+
# Get model WITH correct auth header should pass restriction check
713+
rm -f ./curl.out
714+
code=`curl -s -w %{http_code} -o ./curl.out -H "X-SM-Auth: secret" localhost:8080/models/sm_model`
715+
if [ "$code" == "403" ]; then
716+
cat ./curl.out
717+
echo -e "\n***\n*** Failed. Expected GET /models/<name> with auth header to pass restriction check\n***"
718+
RET=1
719+
fi
720+
721+
# Wrong auth header value should be rejected
722+
rm -f ./curl.out
723+
code=`curl -s -w %{http_code} -o ./curl.out -H "X-SM-Auth: wrong" localhost:8080/models`
724+
if [ "$code" != "403" ]; then
725+
cat ./curl.out
726+
echo -e "\n***\n*** Failed. Expected wrong auth header value to return 403 (got $code)\n***"
727+
RET=1
728+
fi
729+
730+
# Verify core HTTP endpoint is also restricted by the same flag
731+
rm -f ./curl.out
732+
code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8000/v2/repository/index`
733+
if [ "$code" != "403" ]; then
734+
cat ./curl.out
735+
echo -e "\n***\n*** Failed. Expected core HTTP repository index to be restricted (got $code)\n***"
736+
RET=1
737+
fi
738+
739+
set -e
740+
741+
kill $SERVER_PID
742+
wait $SERVER_PID
743+
744+
### HTTP max input size enforcement on SageMaker endpoint ###
745+
# Verify that --http-max-input-size is enforced on the SageMaker /invocations
746+
# path, not just the core HTTP endpoint.
747+
748+
SERVER_LOG="./sagemaker_max_input_size_server.log"
749+
SERVER_ARGS="--allow-sagemaker=true --allow-http=true \
750+
--allow-grpc=false --allow-metrics=false \
751+
--model-repository=`pwd`/models \
752+
--model-control-mode=explicit \
753+
--load-model=sm_model \
754+
--http-max-input-size=256"
755+
run_server_nowait
756+
sagemaker_wait_for_server_ready $SERVER_PID 10
757+
if [ "$WAIT_RET" != "0" ]; then
758+
echo -e "\n***\n*** Failed to start $SERVER\n***"
759+
kill $SERVER_PID || true
760+
cat $SERVER_LOG
761+
exit 1
762+
fi
763+
764+
set +e
765+
766+
# Small payload under 256 bytes should succeed
767+
rm -f ./curl.out
768+
code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8080/invocations \
769+
-H "Content-Type: application/json" \
770+
-d '{"inputs":[{"name":"INPUT0","datatype":"INT32","shape":[1,1],"data":[1]},{"name":"INPUT1","datatype":"INT32","shape":[1,1],"data":[1]}]}'`
771+
if [ "$code" != "200" ]; then
772+
cat ./curl.out
773+
echo -e "\n***\n*** Failed. Expected small payload to succeed on SageMaker endpoint (got $code)\n***"
774+
RET=1
775+
fi
776+
777+
# Large payload over 256 bytes should be rejected
778+
rm -f ./curl.out
779+
LARGE_PAYLOAD='{"inputs":[{"name":"INPUT0","datatype":"INT32","shape":[1,16],"data":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]},{"name":"INPUT1","datatype":"INT32","shape":[1,16],"data":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]}]}'
780+
code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8080/invocations \
781+
-H "Content-Type: application/json" \
782+
-d "$LARGE_PAYLOAD"`
783+
if [ "$code" == "200" ]; then
784+
cat ./curl.out
785+
echo -e "\n***\n*** Failed. Expected oversized payload to be rejected on SageMaker endpoint\n***"
786+
RET=1
787+
fi
788+
789+
# Same limit should apply to core HTTP endpoint
790+
rm -f ./curl.out
791+
code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8000/v2/models/sm_model/infer \
792+
-H "Content-Type: application/json" \
793+
-d "$LARGE_PAYLOAD"`
794+
if [ "$code" == "200" ]; then
795+
cat ./curl.out
796+
echo -e "\n***\n*** Failed. Expected oversized payload to be rejected on core HTTP endpoint\n***"
797+
RET=1
798+
fi
799+
800+
set -e
801+
802+
kill $SERVER_PID
803+
wait $SERVER_PID
804+
615805
unlink /opt/ml/model
616806
rm -rf /opt/ml/model
617807

‎src/main.cc‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,10 +180,22 @@ StartSagemakerService(
180180
triton::server::TraceManager* trace_manager,
181181
const std::shared_ptr<triton::server::SharedMemoryManager>& shm_manager)
182182
{
183+
size_t max_input_size = triton::server::HTTP_DEFAULT_MAX_INPUT_SIZE;
184+
triton::server::RestrictedFeatures restricted_apis{};
185+
#ifdef TRITON_ENABLE_HTTP
186+
// Reuse HTTP server settings for SageMaker endpoint behavior. In
187+
// particular, --http-restricted-api and --http-max-input-size also apply
188+
// to SageMaker requests. Without TRITON_ENABLE_HTTP, SageMaker falls back to
189+
// default input size and unrestricted APIs (no command-line configuration for
190+
// restricted APIs).
191+
max_input_size = g_triton_params.http_max_input_size_;
192+
restricted_apis = g_triton_params.http_restricted_apis_;
193+
#endif // TRITON_ENABLE_HTTP
194+
183195
TRITONSERVER_Error* err = triton::server::SagemakerAPIServer::Create(
184196
server, trace_manager, shm_manager, g_triton_params.sagemaker_port_,
185197
g_triton_params.sagemaker_address_, g_triton_params.sagemaker_thread_cnt_,
186-
service);
198+
max_input_size, restricted_apis, service);
187199
if (err == nullptr) {
188200
err = (*service)->Start();
189201
}

‎src/sagemaker_server.cc‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,16 @@ namespace triton { namespace server {
4141
} \
4242
} while (false)
4343

44+
#define RETURN_AND_RESPOND_IF_RESTRICTED(REQ, RESTRICTED_CATEGORY) \
45+
do { \
46+
auto const& is_restricted_api = \
47+
restricted_apis_.IsRestricted(RESTRICTED_CATEGORY); \
48+
auto const& restriction = restricted_apis_.Get(RESTRICTED_CATEGORY); \
49+
if (is_restricted_api && RespondIfRestricted((REQ), restriction)) { \
50+
return; \
51+
} \
52+
} while (false)
53+
4454
namespace {
4555

4656
void
@@ -173,11 +183,15 @@ SagemakerAPIServer::Handle(evhtp_request_t* req)
173183
if (multi_model_name.empty()) {
174184
LOG_VERBOSE(1) << "SageMaker request: LIST ALL MODELS";
175185

186+
RETURN_AND_RESPOND_IF_RESTRICTED(
187+
req, RestrictedCategory::MODEL_REPOSITORY);
176188
SageMakerMMEListModel(req);
177189
return;
178190
} else {
179191
LOG_VERBOSE(1) << "SageMaker request: GET MODEL";
180192

193+
RETURN_AND_RESPOND_IF_RESTRICTED(
194+
req, RestrictedCategory::MODEL_REPOSITORY);
181195
SageMakerMMEGetModel(req, multi_model_name.c_str());
182196
return;
183197
}
@@ -213,6 +227,8 @@ SagemakerAPIServer::Handle(evhtp_request_t* req)
213227
if (action.empty()) {
214228
LOG_VERBOSE(1) << "SageMaker request: LOAD MODEL";
215229

230+
RETURN_AND_RESPOND_IF_RESTRICTED(
231+
req, RestrictedCategory::MODEL_REPOSITORY);
216232
std::unordered_map<std::string, std::string> parse_load_map;
217233
ParseSageMakerRequest(req, &parse_load_map, "load");
218234
if (!parse_load_map.empty()) {
@@ -224,6 +240,9 @@ SagemakerAPIServer::Handle(evhtp_request_t* req)
224240
case htp_method_DELETE: {
225241
// UNLOAD MODEL
226242
LOG_VERBOSE(1) << "SageMaker request: UNLOAD MODEL";
243+
244+
RETURN_AND_RESPOND_IF_RESTRICTED(
245+
req, RestrictedCategory::MODEL_REPOSITORY);
227246
req->method = htp_method_POST;
228247

229248
SageMakerMMEUnloadModel(req, multi_model_name.c_str());
@@ -253,10 +272,12 @@ SagemakerAPIServer::Create(
253272
triton::server::TraceManager* trace_manager,
254273
const std::shared_ptr<SharedMemoryManager>& shm_manager, const int32_t port,
255274
const std::string address, const int thread_cnt,
275+
const size_t max_input_size, const RestrictedFeatures& restricted_apis,
256276
std::unique_ptr<HTTPServer>* http_server)
257277
{
258278
http_server->reset(new SagemakerAPIServer(
259-
server, trace_manager, shm_manager, port, address, thread_cnt));
279+
server, trace_manager, shm_manager, port, address, thread_cnt,
280+
max_input_size, restricted_apis));
260281

261282
const std::string addr = address + ":" + std::to_string(port);
262283
LOG_INFO << "Started Sagemaker HTTPService at " << addr;

‎src/sagemaker_server.h‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ class SagemakerAPIServer : public HTTPAPIServer {
4646
triton::server::TraceManager* trace_manager,
4747
const std::shared_ptr<SharedMemoryManager>& smb_manager,
4848
const int32_t port, const std::string address, const int thread_cnt,
49+
const size_t max_input_size, const RestrictedFeatures& restricted_apis,
4950
std::unique_ptr<HTTPServer>* sagemaker_server);
5051

5152
class SagemakeInferRequestClass : public InferRequestClass {
@@ -74,10 +75,12 @@ class SagemakerAPIServer : public HTTPAPIServer {
7475
const std::shared_ptr<TRITONSERVER_Server>& server,
7576
triton::server::TraceManager* trace_manager,
7677
const std::shared_ptr<SharedMemoryManager>& shm_manager,
77-
const int32_t port, const std::string address, const int thread_cnt)
78+
const int32_t port, const std::string address, const int thread_cnt,
79+
const size_t max_input_size, const RestrictedFeatures& restricted_apis)
7880
: HTTPAPIServer(
7981
server, trace_manager, shm_manager, port, false /* reuse_port */,
80-
address, "" /* header_forward_pattern */, thread_cnt),
82+
address, "" /* header_forward_pattern */, thread_cnt,
83+
max_input_size, restricted_apis),
8184
ping_regex_(R"(/ping)"), invocations_regex_(R"(/invocations)"),
8285
models_regex_(R"(/models(?:/)?([^/]+)?(/invoke)?)"),
8386
model_path_regex_(

0 commit comments

Comments
 (0)