sglang 0.5.2rc2__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/entrypoints/http_server.py

@@ -27,9 +27,9 @@ import tempfile
 import threading
 import time
 from http import HTTPStatus
-from typing import Any, AsyncIterator, Callable, Dict, List, Optional
+from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union
 
-import setproctitle
+from sglang.srt.tracing.trace import process_tracing_init, trace_set_thread_info
 
 # Fix a bug of Python threading
 setattr(threading, "_register_atexit", lambda *args, **kwargs: None)
@@ -47,11 +47,7 @@ from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import ORJSONResponse, Response, StreamingResponse
 
-from sglang.srt.disaggregation.utils import (
-    FAKE_BOOTSTRAP_HOST,
-    DisaggregationMode,
-    register_disaggregation_server,
-)
+from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode
 from sglang.srt.entrypoints.engine import _launch_subprocesses
 from sglang.srt.entrypoints.openai.protocol import (
     ChatCompletionRequest,
@@ -74,9 +70,11 @@ from sglang.srt.managers.io_struct import (
     AbortReq,
     CloseSessionReqInput,
     ConfigureLoggingReq,
+    DestroyWeightsUpdateGroupReqInput,
     EmbeddingReqInput,
     GenerateReqInput,
     GetWeightsByNameReqInput,
+    InitWeightsSendGroupForRemoteInstanceReqInput,
     InitWeightsUpdateGroupReqInput,
     LoadLoRAAdapterReqInput,
     OpenSessionReqInput,
@@ -84,6 +82,7 @@ from sglang.srt.managers.io_struct import (
     ProfileReqInput,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
+    SendWeightsToRemoteInstanceReqInput,
     SeparateReasoningReqInput,
     SetInternalStateReq,
     SlowDownReqInput,
@@ -95,9 +94,10 @@ from sglang.srt.managers.io_struct import (
     VertexGenerateReqInput,
 )
 from sglang.srt.managers.multi_tokenizer_mixin import (
-    MultiTokenizerManager,
-    deserialize_data,
+    MultiTokenizerRouter,
+    TokenizerWorker,
     get_main_process_id,
+    monkey_patch_uvicorn_multiprocessing,
     read_from_shared_memory,
     write_data_for_multi_tokenizer,
 )
@@ -127,7 +127,7 @@ HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20))
 # Store global states
 @dataclasses.dataclass
 class _GlobalState:
-    tokenizer_manager: TokenizerManager
+    tokenizer_manager: Union[TokenizerManager, MultiTokenizerRouter, TokenizerWorker]
     template_manager: TemplateManager
     scheduler_info: Dict
 
@@ -140,21 +140,6 @@ def set_global_state(global_state: _GlobalState):
     _global_state = global_state
 
 
-# Function to set up all middlewares for multi-tokenizer compatibility
-def setup_middlewares(api_key: Optional[str], enable_metrics: bool):
-    """Setup all middlewares for both single and multi-process modes"""
-    worker_pid = os.getpid()
-
-    if api_key:
-        add_api_key_middleware(app, api_key)
-        logger.info(f"Worker {worker_pid} added API key middleware")
-
-    if enable_metrics:
-        add_prometheus_middleware(app)
-        enable_func_timer()
-        logger.info(f"Worker {worker_pid} added prometheus middleware")
-
-
 async def init_multi_tokenizer() -> ServerArgs:
     """Read args information from shm and init tokenizer manager for current process"""
     pid = os.getpid()
@@ -162,18 +147,22 @@ async def init_multi_tokenizer() -> ServerArgs:
     logger.info(f"current worker_id: {pid}, main processID: {main_pid}")
 
     # Read configuration from shared memory
-    port_args_data = read_from_shared_memory(f"port_args_{main_pid}")
-    server_args_data = read_from_shared_memory(f"server_args_{main_pid}")
-    scheduler_info_data = read_from_shared_memory(f"scheduler_info_{main_pid}")
-    port_args, server_args = deserialize_data(port_args_data, server_args_data)
-    scheduler_info = scheduler_info_data
+    port_args, server_args, scheduler_info = read_from_shared_memory(
+        f"multi_tokenizer_args_{main_pid}"
+    )
+    server_args: ServerArgs
+
+    # API key authentication is not supported in multi-tokenizer mode
+    assert (
+        server_args.api_key is None
+    ), "API key is not supported in multi-tokenizer mode"
 
     port_args.tokenizer_ipc_name = (
         f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}"
     )
 
     # Launch multi-tokenizer manager process
-    tokenizer_manager = MultiTokenizerManager(server_args, port_args)
+    tokenizer_manager = TokenizerWorker(server_args, port_args)
     template_manager = TemplateManager()
     template_manager.initialize_templates(
         tokenizer_manager=tokenizer_manager,
@@ -192,18 +181,29 @@ async def init_multi_tokenizer() -> ServerArgs:
             scheduler_info=scheduler_info,
         )
     )
+
+    if server_args.enable_trace:
+        process_tracing_init(server_args.oltp_traces_endpoint, "sglang")
+        if server_args.disaggregation_mode == "null":
+            thread_label = f"MultiTokenizer-{tokenizer_manager.worker_id}"
+            trace_set_thread_info(thread_label)
+
     return server_args
 
 
 @asynccontextmanager
 async def lifespan(fast_api_app: FastAPI):
-    server_args = getattr(fast_api_app, "server_args", None)
-    if server_args is None:
+    if not getattr(fast_api_app, "is_single_tokenizer_mode", False):
         # Initialize multi-tokenizer support for worker processes
-        fast_api_app.server_args = await init_multi_tokenizer()
-        setup_middlewares(
-            fast_api_app.server_args.api_key, fast_api_app.server_args.enable_metrics
-        )
+        fast_api_app.server_args: ServerArgs = await init_multi_tokenizer()
+
+        # only metrics middleware is supported in multi-tokenizer mode
+        worker_pid = os.getpid()
+        if fast_api_app.server_args.enable_metrics:
+            add_prometheus_middleware(app)
+            enable_func_timer()
+
+        logger.info(f"Worker {worker_pid} added prometheus middleware")
         fast_api_app.warmup_thread = threading.Thread(
             target=_wait_and_warmup,
             args=(
@@ -299,7 +299,23 @@ app.add_middleware(
 
 @app.exception_handler(HTTPException)
 async def validation_exception_handler(request: Request, exc: HTTPException):
-    """Enrich HTTP exception with status code and other details"""
+    """Enrich HTTP exception with status code and other details.
+
+    For /v1/responses, emit OpenAI-style nested error envelope:
+    {"error": {"message": "...", "type": "...", "param": null, "code": <status>}}
+    """
+    # adjust fmt for responses api
+    if request.url.path.startswith("/v1/responses"):
+        nested_error = {
+            "message": exc.detail,
+            "type": HTTPStatus(exc.status_code).phrase,
+            "param": None,
+            "code": exc.status_code,
+        }
+        return ORJSONResponse(
+            content={"error": nested_error}, status_code=exc.status_code
+        )
+
     error = ErrorResponse(
         object="error",
         message=exc.detail,
@@ -312,7 +328,10 @@ async def validation_exception_handler(request: Request, exc: HTTPException):
 # Custom exception handlers to change validation error status codes
 @app.exception_handler(RequestValidationError)
 async def validation_exception_handler(request: Request, exc: RequestValidationError):
-    """Override FastAPI's default 422 validation error with 400"""
+    """Override FastAPI's default 422 validation error with 400.
+
+    For /v1/responses, emit OpenAI-style nested error envelope; for other endpoints keep legacy format.
+    """
     exc_str = str(exc)
     errors_str = str(exc.errors())
 
@@ -321,6 +340,16 @@ async def validation_exception_handler(request: Request, exc: RequestValidationE
     else:
         message = exc_str
 
+    if request.url.path.startswith("/v1/responses"):
+        # adapt specially, for v1/responses API only (notice the error key is different)
+        nested_error = {
+            "message": message,
+            "type": HTTPStatus.BAD_REQUEST.phrase,
+            "param": None,
+            "code": HTTPStatus.BAD_REQUEST.value,
+        }
+        return ORJSONResponse(status_code=400, content={"error": nested_error})
+
     err = ErrorResponse(
         message=message,
         type=HTTPStatus.BAD_REQUEST.phrase,
@@ -681,6 +710,38 @@ async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: R
         )
 
 
+@app.post("/init_weights_send_group_for_remote_instance")
+async def init_weights_send_group_for_remote_instance(
+    obj: InitWeightsSendGroupForRemoteInstanceReqInput, request: Request
+):
+    success, message = (
+        await _global_state.tokenizer_manager.init_weights_send_group_for_remote_instance(
+            obj, request
+        )
+    )
+    content = {"success": success, "message": message}
+    if success:
+        return ORJSONResponse(content, status_code=200)
+    else:
+        return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
+
+
+@app.post("/send_weights_to_remote_instance")
+async def send_weights_to_remote_instance(
+    obj: SendWeightsToRemoteInstanceReqInput, request: Request
+):
+    success, message = (
+        await _global_state.tokenizer_manager.send_weights_to_remote_instance(
+            obj, request
+        )
+    )
+    content = {"success": success, "message": message}
+    if success:
+        return ORJSONResponse(content, status_code=200)
+    else:
+        return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
+
+
 @app.post("/init_weights_update_group")
 async def init_weights_update_group(
     obj: InitWeightsUpdateGroupReqInput, request: Request
@@ -696,6 +757,20 @@ async def init_weights_update_group(
         return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
 
 
+@app.post("/destroy_weights_update_group")
+async def destroy_weights_update_group(
+    obj: DestroyWeightsUpdateGroupReqInput, request: Request
+):
+    """Destroy the parameter update group."""
+    success, message = (
+        await _global_state.tokenizer_manager.destroy_weights_update_group(obj, request)
+    )
+    content = {"success": success, "message": message}
+    return ORJSONResponse(
+        content, status_code=200 if success else HTTPStatus.BAD_REQUEST
+    )
+
+
 @app.post("/update_weights_from_tensor")
 async def update_weights_from_tensor(
     obj: UpdateWeightsFromTensorReqInput, request: Request
@@ -1168,7 +1243,6 @@ def launch_server(
     2. Inter-process communication is done through IPC (each process uses a different port) via the ZMQ library.
     """
     if server_args.tokenizer_worker_num > 1:
-        setproctitle.setproctitle(f"sglang::http_server/multi_tokenizer_router")
         port_args = PortArgs.init_new(server_args)
         port_args.tokenizer_worker_ipc_name = (
             f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}"
@@ -1177,11 +1251,16 @@ def launch_server(
             server_args=server_args, port_args=port_args
         )
     else:
-        setproctitle.setproctitle(f"sglang::http_server/tokenizer_manager")
         tokenizer_manager, template_manager, scheduler_info = _launch_subprocesses(
             server_args=server_args,
         )
 
+        if server_args.enable_trace:
+            process_tracing_init(server_args.oltp_traces_endpoint, "sglang")
+            if server_args.disaggregation_mode == "null":
+                thread_label = "Tokenizer"
+                trace_set_thread_info(thread_label)
+
     set_global_state(
         _GlobalState(
             tokenizer_manager=tokenizer_manager,
@@ -1191,12 +1270,10 @@ def launch_server(
     )
 
     if server_args.tokenizer_worker_num > 1:
-        port_args_shm, server_args_shm, scheduler_info_shm = (
-            write_data_for_multi_tokenizer(
-                port_args,
-                server_args,
-                scheduler_info,
-            )
+        multi_tokenizer_args_shm = write_data_for_multi_tokenizer(
+            port_args,
+            server_args,
+            scheduler_info,
         )
     else:
         # Add api key authorization
@@ -1233,6 +1310,9 @@ def launch_server(
                 "level": "INFO",
                 "propagate": False,
             }
+
+            monkey_patch_uvicorn_multiprocessing()
+
             uvicorn.run(
                 "sglang.srt.entrypoints.http_server:app",
                 host=server_args.host,
@@ -1243,6 +1323,7 @@ def launch_server(
                 workers=server_args.tokenizer_worker_num,
             )
         else:
+            app.is_single_tokenizer_mode = True
             uvicorn.run(
                 app,
                 host=server_args.host,
@@ -1253,10 +1334,8 @@ def launch_server(
             )
     finally:
         if server_args.tokenizer_worker_num > 1:
-            port_args_shm.unlink()
-            server_args_shm.unlink()
-            scheduler_info_shm.unlink()
-            _global_state.tokenizer_manager.clear_tokenizer_mapping()
+            multi_tokenizer_args_shm.unlink()
+            _global_state.tokenizer_manager.socket_mapping.clear_all_sockets()
         else:
             warmup_thread.join()
 
@@ -1405,13 +1484,5 @@ def _wait_and_warmup(
     if server_args.debug_tensor_dump_input_file:
         kill_process_tree(os.getpid())
 
-    if server_args.pdlb_url is not None:
-        register_disaggregation_server(
-            server_args.disaggregation_mode,
-            server_args.port,
-            server_args.disaggregation_bootstrap_port,
-            server_args.pdlb_url,
-        )
-
     if launch_callback is not None:
         launch_callback()

sglang/srt/entrypoints/openai/protocol.py

@@ -16,12 +16,14 @@
 import time
 import uuid
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, TypeAlias, Union
+from typing import Any, Dict, List, NamedTuple, Optional, TypeAlias, Union
 
 from openai.types.responses import (
     ResponseFunctionToolCall,
     ResponseInputItemParam,
     ResponseOutputItem,
+    ResponseOutputMessage,
+    ResponseOutputText,
     ResponseReasoningItem,
 )
 from openai.types.responses.response import ToolChoice
@@ -228,6 +230,15 @@ class CompletionRequest(BaseModel):
 
     # For request id
     rid: Optional[Union[List[str], str]] = None
+    # Extra key for classifying the request (e.g. cache_salt)
+    extra_key: Optional[Union[List[str], str]] = None
+    # Cache salt for request caching
+    cache_salt: Optional[Union[List[str], str]] = None
+    # Priority for the request
+    priority: Optional[int] = None
+
+    # For custom metric labels
+    custom_labels: Optional[Dict[str, str]] = None
 
     @field_validator("max_tokens")
     @classmethod
@@ -334,7 +345,7 @@ class FunctionResponse(BaseModel):
     """Function response."""
 
     name: Optional[str] = None
-    arguments: Optional[str] = None
+    arguments: Optional[str | Dict[str, Any]] = None
 
 
 class ToolCall(BaseModel):
@@ -383,7 +394,7 @@ class Function(BaseModel):
     """Function descriptions."""
 
     description: Optional[str] = Field(default=None, examples=[None])
-    name: Optional[str] = None
+    name: str
     parameters: Optional[object] = None
     strict: bool = False
 
@@ -447,7 +458,7 @@ class ChatCompletionRequest(BaseModel):
         description="Constrains effort on reasoning for reasoning models. "
         "'low' is the least effort, 'high' is the most effort. Reducing reasoning effort can "
         "result in faster responses and fewer tokens used on reasoning in a response. "
-        "Currently only supported for OpenAI models.",
+        "Currently only supported for OpenAI models in the harmony path, i.e GPT-OSS models.",
     )
 
     @model_validator(mode="before")
@@ -540,6 +551,12 @@ class ChatCompletionRequest(BaseModel):
 
     # For request id
     rid: Optional[Union[List[str], str]] = None
+    # Extra key for classifying the request (e.g. cache_salt)
+    extra_key: Optional[Union[List[str], str]] = None
+    # Cache salt for request caching
+    cache_salt: Optional[Union[List[str], str]] = None
+    # Priority for the request
+    priority: Optional[int] = None
 
     # For PD disaggregation
     bootstrap_host: Optional[Union[List[str], str]] = None
@@ -641,6 +658,8 @@ class EmbeddingRequest(BaseModel):
 
     # The request id.
     rid: Optional[Union[List[str], str]] = None
+    # Priority for the request
+    priority: Optional[int] = None
 
 
 class EmbeddingObject(BaseModel):
@@ -769,6 +788,13 @@ class ResponsesRequest(BaseModel):
         description="The request_id related to this request. If the caller does not set it, a random uuid will be generated.",
     )
     priority: int = Field(default=0, description="Request priority")
+    extra_key: Optional[str] = Field(
+        default=None,
+        description="Extra key for classifying the request (e.g. cache_salt)",
+    )
+    cache_salt: Optional[str] = Field(
+        default=None, description="Cache salt for request caching"
+    )
 
     # SGLang-specific sampling parameters
     frequency_penalty: float = 0.0
@@ -857,6 +883,26 @@ class ResponsesResponse(BaseModel):
     tool_choice: str = "auto"
     tools: List[ResponseTool] = Field(default_factory=list)
 
+    # OpenAI compatibility fields. not all are used at the moment.
+    # Recommend checking https://platform.openai.com/docs/api-reference/responses
+    error: Optional[dict] = None
+    incomplete_details: Optional[dict] = None  # TODO(v) support this input
+    instructions: Optional[str] = None
+    max_output_tokens: Optional[int] = None
+    previous_response_id: Optional[str] = None
+    reasoning: Optional[dict] = (
+        # Unused. No model supports this. For GPT-oss, system prompt sets
+        # the field, not server args.
+        None  # {"effort": Optional[str], "summary": Optional[str]}
+    )
+    store: Optional[bool] = None
+    temperature: Optional[float] = None
+    text: Optional[dict] = None  # e.g. {"format": {"type": "text"}}
+    top_p: Optional[float] = None
+    truncation: Optional[str] = None
+    user: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+
     @classmethod
     def from_request(
         cls,
@@ -871,6 +917,41 @@ class ResponsesResponse(BaseModel):
         usage: Optional[UsageInfo],
     ) -> "ResponsesResponse":
         """Create a response from a request."""
+
+        # Determine if the output is plain text only to set text.format
+        def _is_text_only(
+            items: List[
+                Union[
+                    ResponseOutputItem, ResponseReasoningItem, ResponseFunctionToolCall
+                ]
+            ]
+        ) -> bool:
+            if not items:
+                return False
+            for it in items:
+                # tool call -> not pure text.
+                if isinstance(it, ResponseReasoningItem) or isinstance(
+                    it, ResponseFunctionToolCall
+                ):
+                    return False
+                try:
+                    if isinstance(it, ResponseOutputText):
+                        continue
+                    elif isinstance(it, ResponseOutputMessage):
+                        if not it.content:
+                            continue
+                        for c in it.content:
+                            if not isinstance(c, ResponseOutputText):
+                                return False
+                    else:
+                        # Unknown type, not considered text-only
+                        return False
+                except AttributeError:
+                    return False
+            return True
+
+        text_format = {"format": {"type": "text"}} if _is_text_only(output) else None
+
         return cls(
             id=request.request_id,
             created_at=created_time,
@@ -881,6 +962,23 @@ class ResponsesResponse(BaseModel):
             parallel_tool_calls=request.parallel_tool_calls or True,
             tool_choice=request.tool_choice,
             tools=request.tools,
+            # fields for parity with v1/responses
+            error=None,
+            incomplete_details=None,
+            instructions=request.instructions,
+            max_output_tokens=request.max_output_tokens,
+            previous_response_id=request.previous_response_id,  # TODO(v): ensure this is propagated if retrieved from store
+            reasoning={
+                "effort": request.reasoning.effort if request.reasoning else None,
+                "summary": None,  # unused
+            },
+            store=request.store,
+            temperature=request.temperature,
+            text=text_format,  # TODO(v): Expand coverage per https://platform.openai.com/docs/api-reference/responses/list
+            top_p=request.top_p,
+            truncation=request.truncation,
+            user=request.user,
+            metadata=request.metadata or {},
         )
 
 
@@ -919,6 +1017,16 @@ class MessageProcessingResult:
     tool_call_constraint: Optional[Any] = None
 
 
+class ToolCallProcessingResult(NamedTuple):
+    """Result of processing tool calls in a response."""
+
+    tool_calls: Optional[
+        List[Any]
+    ]  # List of ToolCall objects or None if parsing failed
+    remaining_text: str  # Text remaining after parsing tool calls
+    finish_reason: Dict[str, Any]  # Updated finish reason dictionary
+
+
 class ResponseReasoningTextContent(BaseModel):
     text: str
     type: Literal["reasoning_text"] = "reasoning_text"

sglang/srt/entrypoints/openai/serving_base.py

@@ -1,15 +1,20 @@
+from __future__ import annotations
+
 import json
 import logging
 import uuid
 from abc import ABC, abstractmethod
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from fastapi import HTTPException, Request
 from fastapi.responses import ORJSONResponse, StreamingResponse
 
 from sglang.srt.entrypoints.openai.protocol import ErrorResponse, OpenAIServingRequest
 from sglang.srt.managers.io_struct import GenerateReqInput
-from sglang.srt.managers.tokenizer_manager import TokenizerManager
+from sglang.srt.server_args import ServerArgs
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.tokenizer_manager import TokenizerManager
 
 logger = logging.getLogger(__name__)
 
@@ -20,6 +25,14 @@ class OpenAIServingBase(ABC):
 
     def __init__(self, tokenizer_manager: TokenizerManager):
         self.tokenizer_manager = tokenizer_manager
+        self.allowed_custom_labels = (
+            set(
+                self.tokenizer_manager.server_args.tokenizer_metrics_allowed_custom_labels
+            )
+            if isinstance(self.tokenizer_manager.server_args, ServerArgs)
+            and self.tokenizer_manager.server_args.tokenizer_metrics_allowed_custom_labels
+            else None
+        )
 
     async def handle_request(
         self, request: OpenAIServingRequest, raw_request: Request
@@ -33,7 +46,7 @@ class OpenAIServingBase(ABC):
 
             # Convert to internal format
             adapted_request, processed_request = self._convert_to_internal_request(
-                request
+                request, raw_request
             )
 
             # Note(Xinyuan): raw_request below is only used for detecting the connection of the client
@@ -49,6 +62,12 @@ class OpenAIServingBase(ABC):
             return self.create_error_response(
                 message=e.detail, err_type=str(e.status_code), status_code=e.status_code
             )
+        except ValueError as e:
+            return self.create_error_response(
+                message=str(e),
+                err_type="BadRequest",
+                status_code=400,
+            )
         except Exception as e:
             logger.exception(f"Error in request: {e}")
             return self.create_error_response(
@@ -73,10 +92,24 @@ class OpenAIServingBase(ABC):
 
         return f"{self._request_id_prefix()}{uuid.uuid4().hex}"
 
+    def _compute_extra_key(self, request: OpenAIServingRequest) -> Optional[str]:
+        """Compute the final extra_key by concatenating cache_salt and extra_key if both are provided."""
+        parts = []
+        for key in ["cache_salt", "extra_key"]:
+            value = getattr(request, key, None)
+            if value:
+                if not isinstance(value, str):
+                    raise TypeError(
+                        f"Value of {key} must be a string, but got {type(value).__name__}"
+                    )
+                parts.append(value)
+        return "".join(parts) if parts else None
+
     @abstractmethod
     def _convert_to_internal_request(
         self,
         request: OpenAIServingRequest,
+        raw_request: Request = None,
     ) -> tuple[GenerateReqInput, OpenAIServingRequest]:
         """Convert OpenAI request to internal format"""
         pass
@@ -150,3 +183,32 @@ class OpenAIServingBase(ABC):
             code=status_code,
         )
         return json.dumps({"error": error.model_dump()})
+
+    def extract_custom_labels(self, raw_request):
+        if (
+            not self.allowed_custom_labels
+            or not self.tokenizer_manager.server_args.tokenizer_metrics_custom_labels_header
+        ):
+            return None
+
+        custom_labels = None
+        header = (
+            self.tokenizer_manager.server_args.tokenizer_metrics_custom_labels_header
+        )
+        try:
+            raw_labels = (
+                json.loads(raw_request.headers.get(header))
+                if raw_request and raw_request.headers.get(header)
+                else None
+            )
+        except json.JSONDecodeError as e:
+            logger.exception(f"Error in request: {e}")
+            raw_labels = None
+
+        if isinstance(raw_labels, dict):
+            custom_labels = {
+                label: value
+                for label, value in raw_labels.items()
+                if label in self.allowed_custom_labels
+            }
+        return custom_labels