Fix crash (#141)

* feat: fix crash for first token * feat: return threads errors, return stats if requested only

Author: itayariel1

workflows/chatbot/inference/generate.py

@@ -2,6 +2,7 @@
 import copy, time
 from datetime import datetime
 import torch
+from queue import Queue
 import re, os, logging
 from threading import Thread
 import contextlib
@@ -539,9 +540,11 @@ def predict_stream(**params):
  force_words_ids = params["force_words_ids"] if "force_words_ids" in params else None
  use_hpu_graphs = params["use_hpu_graphs"] if "use_hpu_graphs" in params else False
  use_cache = params["use_cache"] if "use_cache" in params else True
+ return_stats = params["return_stats"] if "return_stats" in params else False
  prompt = params["prompt"]
  model = MODELS[model_name]["model"]
  tokenizer = MODELS[model_name]["tokenizer"]
+ errors_queue = Queue()
  task = params.get("task", "")
 
  if task != "":
@@ -586,25 +589,28 @@ def predict_stream(**params):
  )
 
  def generate_output():
- with torch.no_grad():
- with torch.cpu.amp.autocast(
- enabled=True, dtype=torch.bfloat16, cache_enabled=True
- ):
- generation_kwargs = dict(
- streamer=streamer,
- generation_config=generation_config,
- return_dict_in_generate=True,
- )
- generation_kwargs["stopping_criteria"] = StoppingCriteriaList(
- [
- StopOnTokens(
- min_length=max(max_new_tokens - 20, 0),
- start_length=input_token_len,
- stop_token_id=stop_token_ids,
- )
- ]
- )
- return model.generate(**input_tokens, **generation_kwargs)
+ try:
+ with torch.no_grad():
+ with torch.cpu.amp.autocast(
+ enabled=True, dtype=torch.bfloat16, cache_enabled=True
+ ):
+ generation_kwargs = dict(
+ streamer=streamer,
+ generation_config=generation_config,
+ return_dict_in_generate=True,
+ )
+ generation_kwargs["stopping_criteria"] = StoppingCriteriaList(
+ [
+ StopOnTokens(
+ min_length=max(max_new_tokens - 20, 0),
+ start_length=input_token_len,
+ stop_token_id=stop_token_ids,
+ )
+ ]
+ )
+ return model.generate(**input_tokens, **generation_kwargs)
+ except Exception as e:
+ errors_queue.put(e)
 
  generation_thread = Thread(target=generate_output)
  generation_thread.start()
@@ -655,21 +661,23 @@ def generate_output():
  # generation_config.top_p = top_p
  generation_config.temperature = temperature
  generation_config.repetition_penalty = repetition_penalty
-
  def generate_output():
- with torch.no_grad():
- return model.generate(
- **input_tokens,
- **generate_kwargs,
- streamer=streamer,
- generation_config=generation_config,
- return_dict_in_generate=True,
- output_scores=True,
- max_new_tokens=max_new_tokens,
- lazy_mode=True,
- hpu_graphs=use_hpu_graphs,
- ignore_eos=False,
- )
+ try:
+ with torch.no_grad():
+ return model.generate(
+ **input_tokens,
+ **generate_kwargs,
+ streamer=streamer,
+ generation_config=generation_config,
+ return_dict_in_generate=True,
+ output_scores=True,
+ max_new_tokens=max_new_tokens,
+ lazy_mode=True,
+ hpu_graphs=use_hpu_graphs,
+ ignore_eos=False,
+ )
+ except Exception as e:
+ errors_queue.put(e)
 
  generation_thread = Thread(target=generate_output)
  generation_thread.start()
@@ -679,6 +687,14 @@ def generate_output():
  )
  output_word_len = 0
 
+ generation_thread.join(0.1)
+ if generation_thread.is_alive():
+ pass
+ else:
+ thread_exception = errors_queue.get()
+ raise thread_exception
+ # prevent crash if no words are coming out
+ first_token_output_time = datetime.now()
  for new_text in streamer:
  if len(new_text) == 0:
  continue
@@ -697,14 +713,15 @@ def generate_output():
  if output_word_len != 1
  else 0
  )
- stats = {
- "input_token_len": input_token_len,
- "output_word_len": output_word_len,
- "duration": duration,
- "first_word_latency": first_word_latency,
- "msecond_per_word": msecond_per_word,
- }
- yield "END_OF_STREAM_STATS={}".format(stats)
+ if return_stats:
+ stats = {
+ "input_token_len": input_token_len,
+ "output_word_len": output_word_len,
+ "duration": duration,
+ "first_word_latency": first_word_latency,
+ "msecond_per_word": msecond_per_word,
+ }
+ yield "END_OF_STREAM_STATS={}".format(stats)
 
 
 def predict(**params):