[Frontend] support AWS SageMaker inference id

docs/source/serving/openai_compatible_server.md

@@ -123,7 +123,7 @@ completion = client.chat.completions.create(
 
 ## Extra HTTP Headers
 
-Only `X-Request-Id` HTTP request header is supported for now. It can be enabled
+Only `X-Request-Id` and `X-Amzn-SageMaker-Inference-Id` HTTP request headers are supported for now. It can be enabled
 with `--enable-request-id-headers`.
 
 > Note that enablement of the headers can impact performance significantly at high QPS

vllm/entrypoints/openai/api_server.py

@@ -706,10 +706,17 @@ async def authentication(request: Request, call_next):
 
         @app.middleware("http")
         async def add_request_id(request: Request, call_next):
-            request_id = request.headers.get(
-                "X-Request-Id") or uuid.uuid4().hex
+            request_id = request.headers.get("X-Request-Id")
+            sagemaker_request_id = request.headers.get(
+                "X-Amzn-SageMaker-Inference-Id")
+
             response = await call_next(request)
-            response.headers["X-Request-Id"] = request_id
+
+            response.headers["X-Request-Id"] = request_id or uuid.uuid4().hex
+            if sagemaker_request_id is not None:
+                response.headers[
+                    "X-Amzn-SageMaker-Inference-Id"] = sagemaker_request_id
+
             return response
 
     for middleware in args.middleware:

vllm/entrypoints/openai/serving_engine.py

@@ -504,7 +504,9 @@ def _base_request_id(raw_request: Optional[Request],
         if raw_request is None:
             return default
 
-        return raw_request.headers.get("X-Request-Id", default)
+        return (raw_request.headers.get("X-Request-Id")
+                or raw_request.headers.get("X-Amzn-SageMaker-Inference-Id")
+                or default)
 
     @staticmethod
     def _get_decoded_token(logprob: Logprob,