dynamic_async: Add breakpoint for when a job is cancelled (#551)

Previously, when the job cancel function was called, the job itself
wasn't notified that it's cancelled. This caused a deadlock in
the calling function due to waiting forever in a queue.

Therefore, add a signal to break out of the generation loop on a
call to cancel.

Signed-off-by: kingbri <[email protected]>

Author: kingbri1

exllamav2/generator/dynamic_async.py

@@ -75,6 +75,7 @@ class ExLlamaV2DynamicJobAsync:
     job: ExLlamaV2DynamicJob
     queue: asyncio.Queue
     generator: ExLlamaV2DynamicGeneratorAsync
+    cancelled: bool = False
 
     def __init__(self, generator: ExLlamaV2DynamicGeneratorAsync, *args: object, **kwargs: object):
         self.generator = generator
@@ -87,6 +88,10 @@ async def put_result(self, result):
 
     async def __aiter__(self):
         while True:
+            # Get out if the job is cancelled
+            if self.cancelled:
+                break
+
             result = await self.queue.get()
             if isinstance(result, Exception):
                 raise result
@@ -96,3 +101,4 @@ async def __aiter__(self):
 
     async def cancel(self):
         await self.generator.cancel(self)
+        self.cancelled = True