Fix HFRunner hang when subprocess dies during init (#21582)
This commit is contained in:
@@ -15,6 +15,7 @@
|
||||
import json
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import queue as queue_mod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Optional, Tuple, Union
|
||||
|
||||
@@ -411,7 +412,16 @@ class HFRunner:
|
||||
self.in_queue.put(
|
||||
(prompts, image_data, max_new_tokens, lora_paths, token_ids_logprob)
|
||||
)
|
||||
return self.out_queue.get()
|
||||
while True:
|
||||
try:
|
||||
return self.out_queue.get(timeout=5)
|
||||
except queue_mod.Empty:
|
||||
if not self.model_proc.is_alive() and self.out_queue.empty():
|
||||
exitcode = self.model_proc.exitcode
|
||||
raise RuntimeError(
|
||||
f"HFRunner subprocess died with exit code {exitcode} "
|
||||
f"before producing output"
|
||||
)
|
||||
|
||||
def terminate(self):
|
||||
self.model_proc.terminate()
|
||||
|
||||
Reference in New Issue
Block a user