Async/Await¶
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
async def bad_request():
time.sleep(1) # Blocks entire event loop!
return "result"
# GOOD
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
-
Related Topics¶
- 03 Global Interpreter Lock (Gil) - Why async avoids GIL
- 04 Multithreading Vs Multiprocessing - Comparison with threading
- 06 Inference Optimization Patterns - Inference servers