Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions examples/tensorrt_llm/common/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import weakref
from enum import Enum
from queue import Queue
from typing import Callable, Optional, TypedDict, Union
from typing import Any, Callable, Coroutine, Optional, TypedDict, Union

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -54,7 +54,9 @@ class ConversationMessage(TypedDict):
class ManagedThread(threading.Thread):
def __init__(
self,
task: Optional[Union[Callable[..., bool], weakref.WeakMethod]],
task: Optional[
Union[Callable[..., Coroutine[Any, Any, bool]], weakref.WeakMethod]
],
error_queue: Optional[Queue] = None,
name: Optional[str] = None,
loop: Optional[asyncio.AbstractEventLoop] = None,
Expand All @@ -74,7 +76,9 @@ def set_loop(self, loop: asyncio.AbstractEventLoop):

def run(self):
while not self.stop_event.is_set():
task: Optional[Union[Callable[..., bool], weakref.WeakMethod]] = self.task
task: Optional[
Union[Callable[..., Coroutine[Any, Any, bool]], weakref.WeakMethod]
] = self.task
if isinstance(task, weakref.WeakMethod):
task = task()
if task is None:
Expand Down
Loading