How to make the function stop immediately, when in websocket stop requested?

I want to stop running the generating function immediately, whenever stop_requested in the websocket.

class ImageGeneration(BaseAIGeneration):
    async def process(self, websocket, prompt):
        if websocket.stop_requested:
            return None
        await super().process(websocket, prompt)
        if websocket.stop_requested:
            return None
        response = await self.generate_image(prompt=prompt, model=self.model, size=size)
        if websocket_instance.stop_requested or not response:
            return None
        if response and isinstance(response, list):
            image_url = response[0].url
        return image_url

And generate_image function currently is synchronous function

async def generate_image(
    prompt: str, 
    model: str, 
    size: str = "1024x1024"
):
    response = await client.images.generate(
        model=model,
        prompt=prompt,
        size=size,
        n=1
    )
    return response.data

Currently it waits until generate_image fully runs, and only then stops.

I was thinking of creating a task, that will gonna check if the response:

while not task.done():
   if websocket_instance.stop_requested: 
      task.cancel()
   await asyncio.sleep(0.1)

But for me this solution looks weird, and I do not think it is the correct way, how to make the function stop immediately, when in websocket stop requested?

You can use synchronization primivites (especially the Event or Condition instead of that while-loop:

stop_ev = asyncio.Event()

await stop_ev.wait()
task.cancel()

# whenever you stop the websocket:
websocket.stop_requested = True
stop_ev.set()

If you want a lower level alternative you can use futures:

loop = asyncio.get_running_loop()
fut = loop.create_future()

await fut
task.cancel()

# whenever you stop the websocket:
websocket.stop_requested = True
fut.set_result(None)
Вернуться на верх