[engine] support async and streaming (#1614)
This commit is contained in:
Byron Hsu
@@ -716,6 +716,58 @@ class Engine:
|
||||
logprob_start_len: Optional[Union[List[int], int]] = None,
|
||||
top_logprobs_num: Optional[Union[List[int], int]] = None,
|
||||
lora_path: Optional[List[Optional[str]]] = None,
|
||||
stream: bool = False,
|
||||
):
|
||||
# TODO (ByronHsu): refactor to reduce the duplicated code
|
||||
|
||||
obj = GenerateReqInput(
|
||||
text=prompt,
|
||||
sampling_params=sampling_params,
|
||||
return_logprob=return_logprob,
|
||||
logprob_start_len=logprob_start_len,
|
||||
top_logprobs_num=top_logprobs_num,
|
||||
lora_path=lora_path,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
# get the current event loop
|
||||
loop = asyncio.get_event_loop()
|
||||
ret = loop.run_until_complete(generate_request(obj, None))
|
||||
|
||||
if stream is True:
|
||||
STREAM_END_SYMBOL = "data: [DONE]"
|
||||
STREAM_CHUNK_START_SYMBOL = "data:"
|
||||
|
||||
def generator_wrapper():
|
||||
offset = 0
|
||||
loop = asyncio.get_event_loop()
|
||||
generator = ret.body_iterator
|
||||
while True:
|
||||
chunk = loop.run_until_complete(generator.__anext__())
|
||||
|
||||
if chunk.startswith(STREAM_END_SYMBOL):
|
||||
break
|
||||
else:
|
||||
data = json.loads(chunk[len(STREAM_CHUNK_START_SYMBOL) :])
|
||||
data["text"] = data["text"][offset:]
|
||||
offset += len(data["text"])
|
||||
yield data
|
||||
|
||||
# we cannot yield in the scope of generate() because python does not allow yield + return in the same function
|
||||
# however, it allows to wrap the generator as a subfunction and return
|
||||
return generator_wrapper()
|
||||
else:
|
||||
return ret
|
||||
|
||||
async def async_generate(
|
||||
self,
|
||||
prompt: Union[str, List[str]],
|
||||
sampling_params: Optional[Dict] = None,
|
||||
return_logprob: Optional[Union[List[bool], bool]] = False,
|
||||
logprob_start_len: Optional[Union[List[int], int]] = None,
|
||||
top_logprobs_num: Optional[Union[List[int], int]] = None,
|
||||
lora_path: Optional[List[Optional[str]]] = None,
|
||||
stream: bool = False,
|
||||
):
|
||||
obj = GenerateReqInput(
|
||||
text=prompt,
|
||||
@@ -724,13 +776,37 @@ class Engine:
|
||||
logprob_start_len=logprob_start_len,
|
||||
top_logprobs_num=top_logprobs_num,
|
||||
lora_path=lora_path,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
# get the current event loop
|
||||
loop = asyncio.get_event_loop()
|
||||
return loop.run_until_complete(generate_request(obj, None))
|
||||
ret = await generate_request(obj, None)
|
||||
|
||||
if stream is True:
|
||||
STREAM_END_SYMBOL = "data: [DONE]"
|
||||
STREAM_CHUNK_START_SYMBOL = "data:"
|
||||
|
||||
generator = ret.body_iterator
|
||||
|
||||
async def generator_wrapper():
|
||||
|
||||
offset = 0
|
||||
|
||||
while True:
|
||||
chunk = await generator.__anext__()
|
||||
|
||||
if chunk.startswith(STREAM_END_SYMBOL):
|
||||
break
|
||||
else:
|
||||
data = json.loads(chunk[len(STREAM_CHUNK_START_SYMBOL) :])
|
||||
data["text"] = data["text"][offset:]
|
||||
offset += len(data["text"])
|
||||
yield data
|
||||
|
||||
return generator_wrapper()
|
||||
else:
|
||||
return ret
|
||||
|
||||
def shutdown(self):
|
||||
kill_child_process(os.getpid(), including_parent=False)
|
||||
|
||||
# TODO (ByronHsu): encode and async generate
|
||||
# TODO (ByronHsu): encode
|
||||
|
||||
144
python/sglang/test/few_shot_gsm8k_engine.py
Normal file
144
python/sglang/test/few_shot_gsm8k_engine.py
Normal file
@@ -0,0 +1,144 @@
|
||||
import argparse
|
||||
import ast
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.api import set_default_backend
|
||||
from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
|
||||
from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl
|
||||
|
||||
INVALID = -9999999
|
||||
|
||||
|
||||
def get_one_example(lines, i, include_answer):
|
||||
ret = "Question: " + lines[i]["question"] + "\nAnswer:"
|
||||
if include_answer:
|
||||
ret += " " + lines[i]["answer"]
|
||||
return ret
|
||||
|
||||
|
||||
def get_few_shot_examples(lines, k):
|
||||
ret = ""
|
||||
for i in range(k):
|
||||
ret += get_one_example(lines, i, True) + "\n\n"
|
||||
return ret
|
||||
|
||||
|
||||
def get_answer_value(answer_str):
|
||||
answer_str = answer_str.replace(",", "")
|
||||
numbers = re.findall(r"\d+", answer_str)
|
||||
if len(numbers) < 1:
|
||||
return INVALID
|
||||
try:
|
||||
return ast.literal_eval(numbers[-1])
|
||||
except SyntaxError:
|
||||
return INVALID
|
||||
|
||||
|
||||
async def concurrent_generate(engine, prompts, sampling_param):
|
||||
tasks = []
|
||||
for prompt in prompts:
|
||||
tasks.append(asyncio.create_task(engine.async_generate(prompt, sampling_param)))
|
||||
|
||||
outputs = await asyncio.gather(*tasks)
|
||||
return outputs
|
||||
|
||||
|
||||
def run_eval(args):
|
||||
# Select backend
|
||||
engine = sgl.Engine(model_path=args.model_path, log_level="error")
|
||||
|
||||
if args.local_data_path is None:
|
||||
# Read data
|
||||
url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl"
|
||||
filename = download_and_cache_file(url)
|
||||
else:
|
||||
filename = args.local_data_path
|
||||
|
||||
lines = list(read_jsonl(filename))
|
||||
|
||||
# Construct prompts
|
||||
num_questions = args.num_questions
|
||||
num_shots = args.num_shots
|
||||
few_shot_examples = get_few_shot_examples(lines, num_shots)
|
||||
|
||||
questions = []
|
||||
labels = []
|
||||
for i in range(len(lines[:num_questions])):
|
||||
questions.append(get_one_example(lines, i, False))
|
||||
labels.append(get_answer_value(lines[i]["answer"]))
|
||||
assert all(l != INVALID for l in labels)
|
||||
arguments = [{"question": q} for q in questions]
|
||||
|
||||
# construct the prompts
|
||||
prompts = []
|
||||
for i, arg in enumerate(arguments):
|
||||
q = arg["question"]
|
||||
prompt = few_shot_examples + q
|
||||
prompts.append(prompt)
|
||||
|
||||
sampling_param = {
|
||||
"stop": ["Question", "Assistant:", "<|separator|>"],
|
||||
"max_new_tokens": 512,
|
||||
"temperature": 0,
|
||||
}
|
||||
|
||||
# Run requests
|
||||
tic = time.time()
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
outputs = loop.run_until_complete(
|
||||
concurrent_generate(engine, prompts, sampling_param)
|
||||
)
|
||||
|
||||
# End requests
|
||||
latency = time.time() - tic
|
||||
|
||||
# Shutdown the engine
|
||||
engine.shutdown()
|
||||
|
||||
# Parse output
|
||||
preds = []
|
||||
|
||||
for output in outputs:
|
||||
preds.append(get_answer_value(output["text"]))
|
||||
|
||||
# Compute accuracy
|
||||
acc = np.mean(np.array(preds) == np.array(labels))
|
||||
invalid = np.mean(np.array(preds) == INVALID)
|
||||
|
||||
# Compute speed
|
||||
num_output_tokens = sum(
|
||||
output["meta_info"]["completion_tokens"] for output in outputs
|
||||
)
|
||||
output_throughput = num_output_tokens / latency
|
||||
|
||||
# Print results
|
||||
print(f"Accuracy: {acc:.3f}")
|
||||
print(f"Invalid: {invalid:.3f}")
|
||||
print(f"Latency: {latency:.3f} s")
|
||||
print(f"Output throughput: {output_throughput:.3f} token/s")
|
||||
|
||||
return {
|
||||
"accuracy": acc,
|
||||
"latency": latency,
|
||||
"output_throughput": output_throughput,
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--model-path", type=str, default="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
)
|
||||
parser.add_argument("--local-data-path", type=Optional[str], default=None)
|
||||
parser.add_argument("--num-shots", type=int, default=5)
|
||||
parser.add_argument("--num-questions", type=int, default=200)
|
||||
args = parser.parse_args()
|
||||
metrics = run_eval(args)
|
||||
Reference in New Issue
Block a user