CODEWITHSUNDEEP
pythondjangopython-asyncio
M S
M S

Reputation: 1

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 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?

Upvotes: -1

Views: 36

Answers (1)

Łukasz Kwieciński
Łukasz Kwieciński

Reputation: 15728

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)

Upvotes: 0

Related Questions