File size: 2,355 Bytes
fc32096 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | import json
import time
def completed_outputs(engine, prompts, samplings, stream=False):
if not stream:
for index, output in enumerate(engine.generate(prompts,samplings,use_tqdm=False)):
yield index, output
return
identifiers = engine.enqueue(prompts,samplings,use_tqdm=False)
processor = getattr(engine.llm_engine,'output_processor',None)
if processor is not None:
indices = {processor.request_states[request_id].external_req_id:index for index,request_id in enumerate(identifiers)}
else:
indices = {request_id:index for index,request_id in enumerate(identifiers)}
finished = set()
finished_tokens = 0
start = last = time.monotonic()
while engine.llm_engine.has_unfinished_requests():
for output in engine.llm_engine.step():
if not output.finished:
continue
if output.request_id not in indices or output.request_id in finished:
raise ValueError('Unknown or duplicate completed request')
finished.add(output.request_id)
finished_tokens += sum(len(completion.token_ids) for completion in getattr(output,'outputs',[]))
yield indices[output.request_id],output
now = time.monotonic()
if now-last >= 30:
progress = {'type':'queue_progress','completed':len(finished),'queued':len(indices),
'unfinished':engine.llm_engine.get_num_unfinished_requests(),'seconds':round(now-start,1),
'completed_response_tokens':finished_tokens}
if processor is not None:
states = [processor.request_states[request_id] for request_id in identifiers
if request_id in processor.request_states]
detokenizers = [getattr(state,'detokenizer',None) for state in states]
if all(detokenizer is not None for detokenizer in detokenizers):
progress['unfinished_response_tokens'] = sum(d.num_output_tokens() for d in detokenizers)
progress['generated_response_tokens'] = finished_tokens + progress['unfinished_response_tokens']
print(json.dumps(progress),flush=True)
last = now
if len(finished) != len(indices):
raise RuntimeError('Engine stopped with incomplete requests')
|