| 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') |
|
|
