diff --git a/qa/L0_sagemaker/test.sh b/qa/L0_sagemaker/test.sh index c86b278e38..74822f2abc 100755 --- a/qa/L0_sagemaker/test.sh +++ b/qa/L0_sagemaker/test.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -612,6 +612,196 @@ set -e kill $SERVER_PID wait $SERVE_PID +### Restricted API regression for SageMaker endpoint ### +# Verify that --http-restricted-api applies to SageMaker model management +# endpoints (load, unload, list, get) while leaving health and inference +# unrestricted. + +SERVER_LOG="./sagemaker_restricted_api_server.log" +SERVER_ARGS="--allow-sagemaker=true --allow-http=true \ + --allow-grpc=false --allow-metrics=false \ + --model-repository=`pwd`/models \ + --model-control-mode=explicit \ + --load-model=sm_model \ + --http-restricted-api=model-repository:X-SM-Auth=secret" +run_server_nowait +sagemaker_wait_for_server_ready $SERVER_PID 10 +if [ "$WAIT_RET" != "0" ]; then + echo -e "\n***\n*** Failed to start $SERVER\n***" + kill $SERVER_PID || true + cat $SERVER_LOG + exit 1 +fi + +set +e + +# Health should succeed without restricted header +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out localhost:8080/ping` +if [ "$code" != "200" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected /ping to succeed without restricted header\n***" + RET=1 +fi + +# Inference should succeed without restricted header +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8080/invocations \ + -H "Content-Type: application/json" \ + -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]}]}'` +if [ "$code" != "200" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected /invocations inference to succeed without restricted header\n***" + RET=1 +fi + +# List models without auth header should be blocked +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out localhost:8080/models` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected GET /models to return 403 without restricted header (got $code)\n***" + RET=1 +else + grep "This API is restricted" ./curl.out + if [ $? -ne 0 ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected restriction error message in response body\n***" + RET=1 + fi +fi + +# Get model without auth header should be blocked +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out localhost:8080/models/sm_model` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected GET /models/ to return 403 without restricted header (got $code)\n***" + RET=1 +fi + +# Load model without auth header should be blocked +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8080/models \ + -H "Content-Type: application/json" \ + -d '{"model_name":"test","url":"/opt/ml/models/123/model"}'` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected POST /models (load) to return 403 without restricted header (got $code)\n***" + RET=1 +fi + +# Unload model without auth header should be blocked +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X DELETE localhost:8080/models/sm_model` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected DELETE /models (unload) to return 403 without restricted header (got $code)\n***" + RET=1 +fi + +# List models WITH correct auth header should succeed +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -H "X-SM-Auth: secret" localhost:8080/models` +if [ "$code" == "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected GET /models with auth header to pass restriction check\n***" + RET=1 +fi + +# Get model WITH correct auth header should pass restriction check +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -H "X-SM-Auth: secret" localhost:8080/models/sm_model` +if [ "$code" == "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected GET /models/ with auth header to pass restriction check\n***" + RET=1 +fi + +# Wrong auth header value should be rejected +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -H "X-SM-Auth: wrong" localhost:8080/models` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected wrong auth header value to return 403 (got $code)\n***" + RET=1 +fi + +# Verify core HTTP endpoint is also restricted by the same flag +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8000/v2/repository/index` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected core HTTP repository index to be restricted (got $code)\n***" + RET=1 +fi + +set -e + +kill $SERVER_PID +wait $SERVER_PID + +### HTTP max input size enforcement on SageMaker endpoint ### +# Verify that --http-max-input-size is enforced on the SageMaker /invocations +# path, not just the core HTTP endpoint. + +SERVER_LOG="./sagemaker_max_input_size_server.log" +SERVER_ARGS="--allow-sagemaker=true --allow-http=true \ + --allow-grpc=false --allow-metrics=false \ + --model-repository=`pwd`/models \ + --model-control-mode=explicit \ + --load-model=sm_model \ + --http-max-input-size=256" +run_server_nowait +sagemaker_wait_for_server_ready $SERVER_PID 10 +if [ "$WAIT_RET" != "0" ]; then + echo -e "\n***\n*** Failed to start $SERVER\n***" + kill $SERVER_PID || true + cat $SERVER_LOG + exit 1 +fi + +set +e + +# Small payload under 256 bytes should succeed +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8080/invocations \ + -H "Content-Type: application/json" \ + -d '{"inputs":[{"name":"INPUT0","datatype":"INT32","shape":[1,1],"data":[1]},{"name":"INPUT1","datatype":"INT32","shape":[1,1],"data":[1]}]}'` +if [ "$code" != "200" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected small payload to succeed on SageMaker endpoint (got $code)\n***" + RET=1 +fi + +# Large payload over 256 bytes should be rejected +rm -f ./curl.out +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]}]}' +code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8080/invocations \ + -H "Content-Type: application/json" \ + -d "$LARGE_PAYLOAD"` +if [ "$code" == "200" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected oversized payload to be rejected on SageMaker endpoint\n***" + RET=1 +fi + +# Same limit should apply to core HTTP endpoint +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST localhost:8000/v2/models/sm_model/infer \ + -H "Content-Type: application/json" \ + -d "$LARGE_PAYLOAD"` +if [ "$code" == "200" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected oversized payload to be rejected on core HTTP endpoint\n***" + RET=1 +fi + +set -e + +kill $SERVER_PID +wait $SERVER_PID + unlink /opt/ml/model rm -rf /opt/ml/model diff --git a/src/main.cc b/src/main.cc index 90264b1a8b..597a97f17c 100644 --- a/src/main.cc +++ b/src/main.cc @@ -180,10 +180,22 @@ StartSagemakerService( triton::server::TraceManager* trace_manager, const std::shared_ptr& shm_manager) { + size_t max_input_size = triton::server::HTTP_DEFAULT_MAX_INPUT_SIZE; + triton::server::RestrictedFeatures restricted_apis{}; +#ifdef TRITON_ENABLE_HTTP + // Reuse HTTP server settings for SageMaker endpoint behavior. In + // particular, --http-restricted-api and --http-max-input-size also apply + // to SageMaker requests. Without TRITON_ENABLE_HTTP, SageMaker falls back to + // default input size and unrestricted APIs (no command-line configuration for + // restricted APIs). + max_input_size = g_triton_params.http_max_input_size_; + restricted_apis = g_triton_params.http_restricted_apis_; +#endif // TRITON_ENABLE_HTTP + TRITONSERVER_Error* err = triton::server::SagemakerAPIServer::Create( server, trace_manager, shm_manager, g_triton_params.sagemaker_port_, g_triton_params.sagemaker_address_, g_triton_params.sagemaker_thread_cnt_, - service); + max_input_size, restricted_apis, service); if (err == nullptr) { err = (*service)->Start(); } diff --git a/src/sagemaker_server.cc b/src/sagemaker_server.cc index 637be6fc9c..8315ea594e 100644 --- a/src/sagemaker_server.cc +++ b/src/sagemaker_server.cc @@ -41,6 +41,16 @@ namespace triton { namespace server { } \ } while (false) +#define RETURN_AND_RESPOND_IF_RESTRICTED(REQ, RESTRICTED_CATEGORY) \ + do { \ + auto const& is_restricted_api = \ + restricted_apis_.IsRestricted(RESTRICTED_CATEGORY); \ + auto const& restriction = restricted_apis_.Get(RESTRICTED_CATEGORY); \ + if (is_restricted_api && RespondIfRestricted((REQ), restriction)) { \ + return; \ + } \ + } while (false) + namespace { void @@ -173,11 +183,15 @@ SagemakerAPIServer::Handle(evhtp_request_t* req) if (multi_model_name.empty()) { LOG_VERBOSE(1) << "SageMaker request: LIST ALL MODELS"; + RETURN_AND_RESPOND_IF_RESTRICTED( + req, RestrictedCategory::MODEL_REPOSITORY); SageMakerMMEListModel(req); return; } else { LOG_VERBOSE(1) << "SageMaker request: GET MODEL"; + RETURN_AND_RESPOND_IF_RESTRICTED( + req, RestrictedCategory::MODEL_REPOSITORY); SageMakerMMEGetModel(req, multi_model_name.c_str()); return; } @@ -213,6 +227,8 @@ SagemakerAPIServer::Handle(evhtp_request_t* req) if (action.empty()) { LOG_VERBOSE(1) << "SageMaker request: LOAD MODEL"; + RETURN_AND_RESPOND_IF_RESTRICTED( + req, RestrictedCategory::MODEL_REPOSITORY); std::unordered_map parse_load_map; ParseSageMakerRequest(req, &parse_load_map, "load"); if (!parse_load_map.empty()) { @@ -224,6 +240,9 @@ SagemakerAPIServer::Handle(evhtp_request_t* req) case htp_method_DELETE: { // UNLOAD MODEL LOG_VERBOSE(1) << "SageMaker request: UNLOAD MODEL"; + + RETURN_AND_RESPOND_IF_RESTRICTED( + req, RestrictedCategory::MODEL_REPOSITORY); req->method = htp_method_POST; SageMakerMMEUnloadModel(req, multi_model_name.c_str()); @@ -253,10 +272,12 @@ SagemakerAPIServer::Create( triton::server::TraceManager* trace_manager, const std::shared_ptr& shm_manager, const int32_t port, const std::string address, const int thread_cnt, + const size_t max_input_size, const RestrictedFeatures& restricted_apis, std::unique_ptr* http_server) { http_server->reset(new SagemakerAPIServer( - server, trace_manager, shm_manager, port, address, thread_cnt)); + server, trace_manager, shm_manager, port, address, thread_cnt, + max_input_size, restricted_apis)); const std::string addr = address + ":" + std::to_string(port); LOG_INFO << "Started Sagemaker HTTPService at " << addr; diff --git a/src/sagemaker_server.h b/src/sagemaker_server.h index 1aef149605..f4392e67b6 100644 --- a/src/sagemaker_server.h +++ b/src/sagemaker_server.h @@ -46,6 +46,7 @@ class SagemakerAPIServer : public HTTPAPIServer { triton::server::TraceManager* trace_manager, const std::shared_ptr& smb_manager, const int32_t port, const std::string address, const int thread_cnt, + const size_t max_input_size, const RestrictedFeatures& restricted_apis, std::unique_ptr* sagemaker_server); class SagemakeInferRequestClass : public InferRequestClass { @@ -74,10 +75,12 @@ class SagemakerAPIServer : public HTTPAPIServer { const std::shared_ptr& server, triton::server::TraceManager* trace_manager, const std::shared_ptr& shm_manager, - const int32_t port, const std::string address, const int thread_cnt) + const int32_t port, const std::string address, const int thread_cnt, + const size_t max_input_size, const RestrictedFeatures& restricted_apis) : HTTPAPIServer( server, trace_manager, shm_manager, port, false /* reuse_port */, - address, "" /* header_forward_pattern */, thread_cnt), + address, "" /* header_forward_pattern */, thread_cnt, + max_input_size, restricted_apis), ping_regex_(R"(/ping)"), invocations_regex_(R"(/invocations)"), models_regex_(R"(/models(?:/)?([^/]+)?(/invoke)?)"), model_path_regex_(