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