Add --stream-response-default-include-usage server flag (#16711)
This commit is contained in:
@@ -341,7 +341,6 @@ async def lifespan(fast_api_app: FastAPI):
|
||||
_global_state.tokenizer_manager,
|
||||
_global_state.template_manager,
|
||||
enable_prompt_tokens_details=True,
|
||||
enable_force_include_usage=True,
|
||||
tool_server=tool_server,
|
||||
)
|
||||
except Exception:
|
||||
|
||||
@@ -41,6 +41,7 @@ from sglang.srt.entrypoints.openai.utils import (
|
||||
process_cached_tokens_details_from_ret,
|
||||
process_hidden_states_from_ret,
|
||||
process_routed_experts_from_ret,
|
||||
should_include_usage,
|
||||
to_openai_style_logprobs,
|
||||
)
|
||||
from sglang.srt.function_call.core_types import ToolCallItem
|
||||
@@ -655,6 +656,11 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
|
||||
stream_started = False
|
||||
try:
|
||||
include_usage, continuous_usage_stats = should_include_usage(
|
||||
request.stream_options,
|
||||
self.tokenizer_manager.server_args.stream_response_default_include_usage,
|
||||
)
|
||||
|
||||
async for content in self.tokenizer_manager.generate_request(
|
||||
adapted_request, raw_request
|
||||
):
|
||||
@@ -743,10 +749,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
)
|
||||
|
||||
# Add usage stats if continuous_usage_stats is enabled
|
||||
if (
|
||||
request.stream_options
|
||||
and request.stream_options.continuous_usage_stats
|
||||
):
|
||||
if continuous_usage_stats:
|
||||
chunk.usage = UsageProcessor.calculate_token_usage(
|
||||
prompt_tokens=prompt_tokens.get(index, 0),
|
||||
completion_tokens=completion_tokens.get(index, 0),
|
||||
@@ -767,6 +770,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
content,
|
||||
request,
|
||||
has_tool_calls,
|
||||
continuous_usage_stats,
|
||||
):
|
||||
if chunk:
|
||||
yield chunk
|
||||
@@ -798,10 +802,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
)
|
||||
|
||||
# Add usage stats if continuous_usage_stats is enabled
|
||||
if (
|
||||
request.stream_options
|
||||
and request.stream_options.continuous_usage_stats
|
||||
):
|
||||
if continuous_usage_stats:
|
||||
chunk.usage = UsageProcessor.calculate_token_usage(
|
||||
prompt_tokens=prompt_tokens.get(index, 0),
|
||||
completion_tokens=completion_tokens.get(index, 0),
|
||||
@@ -881,7 +882,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
yield f"data: {routed_experts_chunk.model_dump_json()}\n\n"
|
||||
|
||||
# Additional usage chunk
|
||||
if request.stream_options and request.stream_options.include_usage:
|
||||
if include_usage:
|
||||
usage = UsageProcessor.calculate_streaming_usage(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
@@ -1313,6 +1314,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
content: Dict[str, Any],
|
||||
request: ChatCompletionRequest,
|
||||
has_tool_calls: Dict[int, bool],
|
||||
continuous_usage_stats: bool = False,
|
||||
):
|
||||
"""Process tool calls in streaming response"""
|
||||
if index not in parser_dict:
|
||||
@@ -1351,7 +1353,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
)
|
||||
|
||||
# Add usage stats if continuous_usage_stats is enabled
|
||||
if request.stream_options and request.stream_options.continuous_usage_stats:
|
||||
if continuous_usage_stats:
|
||||
prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
|
||||
completion_tokens = content["meta_info"].get("completion_tokens", 0)
|
||||
chunk.usage = UsageProcessor.calculate_token_usage(
|
||||
@@ -1401,7 +1403,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
)
|
||||
|
||||
# Add usage stats if continuous_usage_stats is enabled
|
||||
if request.stream_options and request.stream_options.continuous_usage_stats:
|
||||
if continuous_usage_stats:
|
||||
prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
|
||||
completion_tokens = content["meta_info"].get("completion_tokens", 0)
|
||||
chunk.usage = UsageProcessor.calculate_token_usage(
|
||||
|
||||
@@ -23,6 +23,7 @@ from sglang.srt.entrypoints.openai.utils import (
|
||||
process_cached_tokens_details_from_ret,
|
||||
process_hidden_states_from_ret,
|
||||
process_routed_experts_from_ret,
|
||||
should_include_usage,
|
||||
to_openai_style_logprobs,
|
||||
)
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
@@ -225,6 +226,11 @@ class OpenAIServingCompletion(OpenAIServingBase):
|
||||
|
||||
stream_started = False
|
||||
try:
|
||||
include_usage, continuous_usage_stats = should_include_usage(
|
||||
request.stream_options,
|
||||
self.tokenizer_manager.server_args.stream_response_default_include_usage,
|
||||
)
|
||||
|
||||
async for content in self.tokenizer_manager.generate_request(
|
||||
adapted_request, raw_request
|
||||
):
|
||||
@@ -318,10 +324,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
|
||||
)
|
||||
|
||||
# Add usage stats if continuous_usage_stats is enabled
|
||||
if (
|
||||
request.stream_options
|
||||
and request.stream_options.continuous_usage_stats
|
||||
):
|
||||
if continuous_usage_stats:
|
||||
chunk.usage = UsageProcessor.calculate_token_usage(
|
||||
prompt_tokens=prompt_tokens.get(index, 0),
|
||||
completion_tokens=completion_tokens.get(index, 0),
|
||||
@@ -371,7 +374,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
|
||||
yield f"data: {routed_experts_chunk.model_dump_json()}\n\n"
|
||||
|
||||
# Handle final usage chunk
|
||||
if request.stream_options and request.stream_options.include_usage:
|
||||
if include_usage:
|
||||
usage = UsageProcessor.calculate_streaming_usage(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
|
||||
@@ -76,7 +76,6 @@ class OpenAIServingResponses(OpenAIServingChat):
|
||||
template_manager: TemplateManager,
|
||||
*,
|
||||
enable_prompt_tokens_details: bool = False,
|
||||
enable_force_include_usage: bool = False,
|
||||
tool_server: Optional[ToolServer] = None,
|
||||
) -> None:
|
||||
super().__init__(tokenizer_manager, template_manager)
|
||||
@@ -84,7 +83,6 @@ class OpenAIServingResponses(OpenAIServingChat):
|
||||
# template_manager is already set by parent class
|
||||
self.reasoning_parser = self.tokenizer_manager.server_args.reasoning_parser
|
||||
self.enable_prompt_tokens_details = enable_prompt_tokens_details
|
||||
self.enable_force_include_usage = enable_force_include_usage
|
||||
|
||||
# Get default sampling params from model config if available
|
||||
self.default_sampling_params = {}
|
||||
|
||||
@@ -6,6 +6,7 @@ from sglang.srt.entrypoints.openai.protocol import (
|
||||
ChatCompletionRequest,
|
||||
CompletionRequest,
|
||||
LogProbs,
|
||||
StreamOptions,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -73,6 +74,23 @@ def process_hidden_states_from_ret(
|
||||
return hidden_states
|
||||
|
||||
|
||||
def should_include_usage(
|
||||
stream_options: StreamOptions | None, stream_response_default_include_usage: bool
|
||||
) -> tuple[bool, bool]:
|
||||
# When stream_options are specified in the request
|
||||
if stream_options:
|
||||
include_usage = (
|
||||
stream_options.include_usage or stream_response_default_include_usage
|
||||
)
|
||||
continuous_usage_stats = bool(stream_options.continuous_usage_stats)
|
||||
else:
|
||||
include_usage, continuous_usage_stats = (
|
||||
stream_response_default_include_usage,
|
||||
False,
|
||||
)
|
||||
return include_usage, continuous_usage_stats
|
||||
|
||||
|
||||
def process_routed_experts_from_ret(
|
||||
ret_item: Dict[str, Any],
|
||||
request: Union[
|
||||
|
||||
@@ -374,6 +374,7 @@ class ServerArgs:
|
||||
pp_max_micro_batch_size: Optional[int] = None
|
||||
pp_async_batch_depth: int = 0
|
||||
stream_interval: int = 1
|
||||
stream_response_default_include_usage: bool = False
|
||||
incremental_streaming_output: bool = False
|
||||
enable_streaming_session: bool = False
|
||||
random_seed: Optional[int] = None
|
||||
@@ -4155,6 +4156,12 @@ class ServerArgs:
|
||||
action="store_true",
|
||||
help="Whether to output as a sequence of disjoint segments.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stream-response-default-include-usage",
|
||||
action="store_true",
|
||||
help="Include usage in every streaming response "
|
||||
"(even when stream_options is not specified).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stream-output",
|
||||
action=DeprecatedStoreTrueAction,
|
||||
|
||||
Reference in New Issue
Block a user