Skip to content

Async/Await: Asynchronous I/O for High Concurrency

Overview

Async/await enables handling 1000s of concurrent connections with a single thread: - Event loop: Manages multiple concurrent operations - Coroutines: Functions that can pause and resume - Non-blocking I/O: Requests don't block the event loop - Scalability: Handle 10,000+ concurrent connections with minimal overhead

Perfect for building high-concurrency inference APIs and data loading pipelines.


Async Basics

Coroutines and Awaitherating

import asyncio

# Coroutine (async function)
async def fetch_data(url):
    """Async function that returns a coroutine."""
    print(f"Fetching {url}...")
    await asyncio.sleep(1)  # Simulate I/O
    print(f"Done fetching {url}")
    return f"Data from {url}"

# Run coroutine
async def main():
    # Await executes coroutine and waits
    result = await fetch_data("example.com")
    print(result)

# Execute async main
asyncio.run(main())

Concurrent Execution with gather

import asyncio

async def task(name, delay):
    print(f"{name} starting")
    await asyncio.sleep(delay)
    print(f"{name} done")
    return f"Result from {name}"

async def main():
    # Run multiple tasks concurrently
    results = await asyncio.gather(
        task("A", 1),
        task("B", 2),
        task("C", 1)
    )
    # Total time: max(1, 2, 1) = 2 seconds
    # Not 1+2+1 = 4 seconds!
    return results

results = asyncio.run(main())
print(results)

Cancellation and Timeout

import asyncio

async def long_task():
    try:
        await asyncio.sleep(10)
    except asyncio.CancelledError:
        print("Task cancelled!")
        raise

async def main():
    # Cancel after timeout
    task = asyncio.create_task(long_task())

    try:
        await asyncio.wait_for(task, timeout=1)
    except asyncio.TimeoutError:
        print("Task timed out")
        task.cancel()

Async I/O Patterns

Async HTTP Requests

import asyncio
import aiohttp

async def fetch_url(session, url):
    """Async HTTP request."""
    async with session.get(url, timeout=5) as response:
        return await response.text()

async def fetch_multiple(urls):
    """Fetch multiple URLs concurrently."""
    async with aiohttp.ClientSession() as session:
        tasks = [fetch_url(session, url) for url in urls]
        results = await asyncio.gather(*tasks)
    return results

# Usage
urls = ['https://example.com', 'https://example.org', ...]
results = asyncio.run(fetch_multiple(urls))

Async Database Queries

import asyncio
import aiomysql

async def query_db(sql):
    """Async database query."""
    conn = await aiomysql.connect(host='localhost', db='mydb')
    async with conn.cursor() as cursor:
        await cursor.execute(sql)
        result = await cursor.fetchall()
    conn.close()
    return result

async def concurrent_queries():
    """Multiple database queries concurrently."""
    results = await asyncio.gather(
        query_db("SELECT * FROM users"),
        query_db("SELECT * FROM posts"),
        query_db("SELECT * FROM comments")
    )
    return results

results = asyncio.run(concurrent_queries())

Async in ML: Inference Server

Simple Async Inference Server

import asyncio
import torch
from aiohttp import web

class AsyncInferenceServer:
    def __init__(self, model_path):
        self.model = torch.load(model_path)
        self.model.eval()

    async def infer(self, data):
        """Non-blocking inference."""
        # Move to thread pool to avoid blocking event loop
        loop = asyncio.get_event_loop()
        x = torch.tensor(data)

        result = await loop.run_in_executor(
            None,
            lambda: self._infer_impl(x)
        )
        return result

    def _infer_impl(self, x):
        """Actual inference (runs in thread)."""
        with torch.no_grad():
            return self.model(x).numpy()

    async def handle_request(self, request):
        """HTTP request handler."""
        data = await request.json()
        result = await self.infer(data['input'])
        return web.json_response({'output': result.tolist()})

if __name__ == '__main__':
    server = AsyncInferenceServer('model.pth')
    app = web.Application()
    app.router.add_post('/infer', server.handle_request)
    web.run_app(app, port=8080)

# Can handle 1000s of concurrent requests!

Async Data Pipeline

import asyncio
import torch
from torch.utils.data import Dataset, DataLoader

class AsyncDataPipeline:
    def __init__(self, data_loader, model, batch_size=32):
        self.data_loader = data_loader
        self.model = model
        self.batch_size = batch_size
        self.queue = asyncio.Queue(maxsize=10)

    async def load_data(self):
        """Async data loading."""
        for batch_x, batch_y in self.data_loader:
            await self.queue.put((batch_x, batch_y))
        await self.queue.put(None)  # End signal

    async def process_data(self):
        """Async inference."""
        results = []
        while True:
            item = await self.queue.get()
            if item is None:
                break

            batch_x, batch_y = item

            # Non-blocking inference
            loop = asyncio.get_event_loop()
            output = await loop.run_in_executor(
                None,
                lambda: self.model(batch_x)
            )
            results.append(output)

        return results

    async def run(self):
        """Run data pipeline."""
        await asyncio.gather(
            self.load_data(),
            self.process_data()
        )

# Usage
pipeline = AsyncDataPipeline(train_loader, model)
asyncio.run(pipeline.run())

Async Batch Collection

Collect Requests and Batch

import asyncio
import torch

class AsyncBatcher:
    def __init__(self, model, batch_size=32, timeout=0.1):
        self.model = model
        self.batch_size = batch_size
        self.timeout = timeout
        self.request_queue = asyncio.Queue()
        self.response_futures = {}

    async def collect_and_process(self):
        """Collect requests and batch process."""
        while True:
            batch = []
            request_ids = []

            # Collect up to batch_size requests
            start = asyncio.get_event_loop().time()
            while len(batch) < self.batch_size:
                try:
                    elapsed = asyncio.get_event_loop().time() - start
                    remaining = max(0, self.timeout - elapsed)

                    if remaining <= 0:
                        break

                    request_id, data = await asyncio.wait_for(
                        self.request_queue.get(),
                        timeout=remaining
                    )
                    batch.append(data)
                    request_ids.append(request_id)
                except asyncio.TimeoutError:
                    break

            # Process batch
            if batch:
                batch_tensor = torch.stack(batch)

                with torch.no_grad():
                    outputs = self.model(batch_tensor)

                # Return results
                for req_id, output in zip(request_ids, outputs):
                    self.response_futures[req_id].set_result(output)

    async def infer(self, request_id, data):
        """Submit inference request."""
        future = asyncio.Future()
        self.response_futures[request_id] = future

        await self.request_queue.put((request_id, data))
        return await future

if __name__ == '__main__':
    model = torch.nn.Linear(10, 5)
    batcher = AsyncBatcher(model, batch_size=32)

    # Start batch collection in background
    asyncio.create_task(batcher.collect_and_process())

    # Submit requests
    results = asyncio.run(
        asyncio.gather(*[
            batcher.infer(i, torch.randn(10))
            for i in range(100)
        ])
    )

Async vs Threading vs Multiprocessing

Aspect Async Threading Multiprocessing
Concurrency 1000s easy 100s max Many heavy
CPU parallelism No No (GIL) Yes
I/O efficiency Excellent Good Overhead
Memory Low Low High
Complexity Medium Medium High
Use case Web servers I/O + some compute CPU-bound

Best Practices

Rule 1: Never Block the Event Loop

# BAD: Blocking operation in async
async def bad_request():
    time.sleep(1)  # Blocks entire event loop!
    return "result"

# GOOD: Use run_in_executor
async def good_request():
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, time.sleep, 1)
    return "result"

Rule 2: Use Proper Exception Handling

async def safe_gather():
    """Handle exceptions in concurrent tasks."""
    tasks = [...]

    # Catch individual exceptions
    results = await asyncio.gather(*tasks, return_exceptions=True)

    # Check for exceptions
    for result in results:
        if isinstance(result, Exception):
            print(f"Task failed: {result}")

Summary: When to Use Async

✓ Web servers (1000+ concurrent clients) ✓ High-concurrency APIs ✓ I/O-heavy workloads ✓ Data pipelines

✗ CPU-bound computation ✗ Simple scripts ✗ Low concurrency needs