Cancellation, Timeouts and Errors

import asyncio


async def generate(prompt: str) -> str:
    try:
        await asyncio.sleep(10)                  # waiting on the model
        return "done"
    except asyncio.CancelledError:
        print("cleanup: releasing GPU slot")
        raise
    finally:
        print("finally runs too")


async def main() -> None:
    task = asyncio.create_task(generate("hi"))
    await asyncio.sleep(0.01)
    task.cancel("user pressed stop")
    try:
        await task
    except asyncio.CancelledError:
        print("cancelled:", task.cancelled())


asyncio.run(main())
# Output:
# cleanup: releasing GPU slot
# finally runs too
# cancelled: True

import asyncio


async def bad_call() -> str:
    try:
        await asyncio.sleep(1)
        return "answer"
    except asyncio.CancelledError:          # swallowed
        return "partial"


async def good_call() -> str:
    try:
        await asyncio.sleep(1)
        return "answer"
    except asyncio.CancelledError:
        raise


async def main() -> None:
    for fn in (bad_call, good_call):
        try:
            async with asyncio.timeout(0.05):
                result = await fn()
            print(fn.__name__, "returned", result, "- timeout did not fire")
        except TimeoutError:
            print(fn.__name__, "timed out as expected")


asyncio.run(main())
# Output:
# bad_call returned partial - timeout did not fire
# good_call timed out as expected

monitor.cancel()
try:
    await monitor
except asyncio.CancelledError:
    pass

monitor.cancel()
try:
    await monitor
except asyncio.CancelledError:
    current = asyncio.current_task()
    if current is not None and current.cancelling():
        raise                                   # we are being cancelled as well: propagate

import asyncio


async def retrieve() -> list[str]:
    await asyncio.sleep(0.01)
    return ["chunk-1", "chunk-2"]


async def generate(chunks: list[str]) -> str:
    await asyncio.sleep(0.5)
    return "answer"


async def main() -> None:
    try:
        async with asyncio.timeout(0.1) as cm:
            chunks = await retrieve()
            print("retrieved", len(chunks), "deadline set:", cm.when() is not None)
            await generate(chunks)
    except TimeoutError:
        print("whole RAG request exceeded 100 ms; expired =", cm.expired())


asyncio.run(main())
# Output:
# retrieved 2 deadline set: True
# whole RAG request exceeded 100 ms; expired = True

async def with_idle_timeout(stream, idle: float):
    it = stream.__aiter__()
    while True:
        try:
            async with asyncio.timeout(idle):
                chunk = await it.__anext__()
        except StopAsyncIteration:
            return
        yield chunk

import asyncio


async def write_batch(batch: list[str], log: list[str]) -> None:
    await asyncio.sleep(0.05)               # vector-store upsert
    log.append(f"wrote {len(batch)}")


async def handler(log: list[str]) -> None:
    write = asyncio.ensure_future(write_batch(["a", "b"], log))
    try:
        await asyncio.shield(write)
    except asyncio.CancelledError:
        await write                         # let the write finish, then propagate
        raise


async def main() -> None:
    log: list[str] = []
    task = asyncio.create_task(handler(log))
    await asyncio.sleep(0.01)
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        pass
    print(log)


asyncio.run(main())
# Output: ['wrote 2']

async def stream_answer(client, prompt: str):
    stream = await client.open_stream(prompt)
    try:
        async for event in stream:
            yield event
    finally:
        # Runs on normal end, error, consumer break (with aclosing), and cancellation.
        # A second cancel during this await would interrupt it; keep it short or shield it.
        await asyncio.shield(stream.aclose())

import asyncio


async def embed(doc: str) -> list[float]:
    await asyncio.sleep(0)
    if not doc:
        raise ValueError("empty document")
    return [float(len(doc))]


async def main() -> None:
    docs = ["alpha", "", "gamma"]
    results = await asyncio.gather(*(embed(d) for d in docs), return_exceptions=True)
    vectors = {i: r for i, r in enumerate(results) if not isinstance(r, BaseException)}
    failures = {i: repr(r) for i, r in enumerate(results) if isinstance(r, BaseException)}
    print(vectors)
    print(failures)


asyncio.run(main())
# Output:
# {0: [5.0], 2: [5.0]}
# {1: "ValueError('empty document')"}

import asyncio


class RateLimited(Exception):
    pass


async def call(i: int) -> int:
    await asyncio.sleep(0.01)
    if i == 1:
        raise RateLimited("429 on call 1")
    if i == 2:
        raise TimeoutError("call 2 timed out")
    return i


async def main() -> None:
    try:
        async with asyncio.TaskGroup() as tg:
            for i in range(4):
                tg.create_task(call(i))
    except* RateLimited as eg:
        print("rate limited:", [str(e) for e in eg.exceptions])
    except* TimeoutError as eg:
        print("timeouts:", [str(e) for e in eg.exceptions])


asyncio.run(main())
# Output:
# rate limited: ['429 on call 1']
# timeouts: ['call 2 timed out']

import asyncio
import email.utils
import random
import time
from typing import Awaitable, Callable, Optional, TypeVar

T = TypeVar("T")


class TransientError(Exception):
    def __init__(self, msg: str, retry_after: Optional[float] = None) -> None:
        super().__init__(msg)
        self.retry_after = retry_after


def parse_retry_after(value: Optional[str]) -> Optional[float]:
    """Retry-After is either delta-seconds or an HTTP date."""
    if not value:
        return None
    try:
        return max(0.0, float(value))
    except ValueError:
        dt = email.utils.parsedate_to_datetime(value)
        return max(0.0, dt.timestamp() - time.time())


async def with_retry(call: Callable[[], Awaitable[T]], *, attempts: int = 5, base: float = 0.5,
                     cap: float = 30.0, budget: float = 120.0) -> T:
    start = time.monotonic()
    for attempt in range(attempts):
        try:
            return await call()
        except TransientError as err:
            if attempt == attempts - 1:
                raise
            delay = err.retry_after if err.retry_after is not None else random.uniform(0, min(cap, base * 2 ** attempt))
            if time.monotonic() - start + delay > budget:
                raise
            await asyncio.sleep(delay)
    raise AssertionError("unreachable")


async def main() -> None:
    calls = {"n": 0}

    async def flaky() -> str:
        calls["n"] += 1
        if calls["n"] < 3:
            raise TransientError("503 overloaded", retry_after=0.01)
        return "ok"

    print(await with_retry(flaky), "after", calls["n"], "calls")
    print(parse_retry_after("2"))


asyncio.run(main())
# Output:
# ok after 3 calls
# 2.0

import asyncio


class CircuitOpen(Exception):
    pass


class CircuitBreaker:
    def __init__(self, threshold: int = 5, reset_after: float = 30.0) -> None:
        self.threshold, self.reset_after = threshold, reset_after
        self.failures, self.opened_at = 0, None

    async def call(self, fn):
        loop = asyncio.get_running_loop()
        if self.opened_at is not None:
            if loop.time() - self.opened_at < self.reset_after:
                raise CircuitOpen("vector store unavailable; using keyword fallback")
            self.opened_at = None                  # half-open: let one call through
        try:
            result = await fn()
        except Exception:
            self.failures += 1
            if self.failures >= self.threshold:
                self.opened_at = loop.time()
            raise
        self.failures = 0
        return result

import asyncio
import signal


async def serve_forever(stop: asyncio.Event) -> None:
    while not stop.is_set():
        await asyncio.sleep(0.1)          # stand-in for consuming a job queue


async def main() -> None:
    stop = asyncio.Event()
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, stop.set)     # Unix only
    worker = asyncio.create_task(serve_forever(stop))
    await stop.wait()
    # Drain: stop taking new jobs, give in-flight work a deadline, then cancel.
    try:
        async with asyncio.timeout(25):            # below the orchestrator's kill grace period
            await worker
    except TimeoutError:
        worker.cancel()

import asyncio
import logging

logger = logging.getLogger("app")


def handle_loop_error(loop: asyncio.AbstractEventLoop, context: dict) -> None:
    exc = context.get("exception")
    logger.error("asyncio: %s", context.get("message"), exc_info=exc,
                 extra={"task": repr(context.get("task") or context.get("future"))})


async def main() -> None:
    asyncio.get_running_loop().set_exception_handler(handle_loop_error)