A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[benchmarks] Add option to use unique jsonschema for each request (#14457)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
Mar 8, 2025 at 09:36 UTC
9085aabd627ec828745adb0b1afdca61bc4198aa
1 file changed
+41
-17
benchmarks/benchmark_serving_structured_output.py
+41
-17
@@ -24,11 +24,13 @@ On the client side, run:
24
"""
25
import argparse
26
import asyncio
27
+import copy
28
import dataclasses
29
import json
30
import os
31
import random
32
import time
33
+import uuid
34
import warnings
35
from collections.abc import AsyncGenerator
36
from dataclasses import dataclass
@@ -109,24 +111,43 @@ class SampleRequest:
111
112
def sample_requests(tokenizer: PreTrainedTokenizerBase,
113
args: argparse.Namespace) -> list[SampleRequest]:
112
- if args.dataset == 'json':
114
+ if args.dataset == 'json' or args.dataset == 'json-unique':
115
if args.json_schema_path is None:
116
dir_path = os.path.dirname(os.path.realpath(__file__))
117
args.json_schema_path = os.path.join(dir_path,
118
"structured_schemas",
119
"structured_schema_1.json")
120
+ json_schemas = []
121
with open(args.json_schema_path) as f:
122
schema = json.load(f)
120
- prompt = f"Generate an example of a user profile given the following schema: {json.dumps(schema)}" # noqa: E501
121
- input_len = len(tokenizer(prompt).input_ids)
122
- print(f"Input length of the prompt: {input_len} tokens")
123
+
124
+ if args.dataset == 'json-unique':
125
+ json_schemas = [
126
+ copy.deepcopy(schema) for _ in range(args.num_prompts)
127
+ ]
128
+ for i in range(len(json_schemas)):
129
+ json_schemas[i]["properties"][
130
+ f"__optional_field_{uuid.uuid4()}"] = {
131
+ "type":
132
+ "string",
133
+ "description":
134
+ "An unique optional field to avoid cached schemas"
135
+ }
136
+
137
+ def gen_prompt(index: int):
138
+ schema = json_schemas[index % len(json_schemas)]
139
+ return f"Generate an example of a user profile given the following schema: {json.dumps(schema)}" # noqa: E501
140
+
141
+ def get_schema(index: int):
142
+ return json_schemas[index % len(json_schemas)]
143
+
144
requests = [
124
- SampleRequest(prompt=prompt,
125
- prompt_len=input_len,
145
+ SampleRequest(prompt=gen_prompt(i),
146
+ prompt_len=len(tokenizer(gen_prompt(i)).input_ids),
147
expected_output_len=args.output_len,
127
- schema=schema,
148
+ schema=get_schema(i),
149
structure_type=args.structure_type)
129
- for _ in range(args.num_prompts)
150
+ for i in range(args.num_prompts)
151
]
152
153
elif args.dataset == "grammar":
@@ -821,10 +842,12 @@ if __name__ == "__main__":
842
default="/v1/completions",
843
help="API endpoint.",
844
)
824
- parser.add_argument(
825
- "--dataset",
826
- default='json',
827
- choices=['json', 'grammar', 'regex', 'choice', 'xgrammar_bench'])
845
+ parser.add_argument("--dataset",
846
+ default='json',
847
+ choices=[
848
+ 'json', 'json-unique', 'grammar', 'regex',
849
+ 'choice', 'xgrammar_bench'
850
+ ])
851
parser.add_argument("--json_schema_path",
852
type=str,
853
default=None,
@@ -966,11 +989,12 @@ if __name__ == "__main__":
989
type=float,
990
default=1.0,
991
help="Ratio of Structured Outputs requests")
969
- parser.add_argument("--structured-output-backend",
970
- type=str,
971
- choices=["outlines", "lm-format-enforcer", "xgrammar"],
972
- default="xgrammar",
973
- help="Backend to use for structured outputs")
992
+ parser.add_argument(
993
+ "--structured-output-backend",
994
+ type=str,
995
+ choices=["outlines", "lm-format-enforcer", "xgrammar", "json-unique"],
996
+ default="xgrammar",
997
+ help="Backend to use for structured outputs")
998
999
args = parser.parse_args()
1000
main(args)