Как сделать так, чтобы функция немедленно останавливалась, когда в websocket запрашивается остановка?

Я хочу немедленно прекратить выполнение генерирующей функции всякий раз, когда stop_requestедет в 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

И функция generate_image в настоящее время является синхронной функцией

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

В настоящее время он ожидает полного запуска generate_image и только после этого останавливается.

Я подумывал о создании задачи, которая будет проверять, будет ли получен ответ:

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

Но для меня это решение выглядит странным, и я не думаю, что это правильный способ, как заставить функцию немедленно остановиться, когда в websocket запрашивается остановка?

Вы можете использовать примитивы синхронизации (особенно Event или Condition вместо этого цикла while:

stop_ev = asyncio.Event()

await stop_ev.wait()
task.cancel()

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

Если вам нужна альтернатива более низкого уровня, вы можете использовать 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)
Вернуться на верх