How Making Heavy PyTorch CUDA Processing Async Caused a Memory Leak—and How We Fixed It
Hello! This is the Qualiteg Product Development team!
This is the story of how simply converting a synchronous method into an asynchronous one (async) hit us with an unexpected memory leak*.
It happened while we were developing a video processing system that uses deep learning models.
We had a requirement to "report processing progress in real time over WebSocket," and I casually figured, "I just need to use async/await, right?"—and fell straight into an unexpected trap.
Despite using a professional-grade GPU, the system crashed from running out of memory.
In this article, I want to share the cause, the solution, and the lessons learned in detail. I hope it helps anyone facing a similar problem.
*Strictly speaking, this is not a "memory leak" but "delayed memory release." The practical impact is the same, however, so for convenience this article refers to it as a memory leak.
Background: Why Progress Notifications Need to Be Asynchronous
What Modern Web Applications Demand
In modern web application development, showing real-time progress for long-running operations has become standard practice for improving user experience.
With the spread of generative AI in particular, operations that take several minutes to tens of minutes are no longer unusual. Our own MotionVox also runs for a long time when generating longer videos. Progress notifications are needed so users know how far along the process is and how much longer it will take.
Understanding Why an Asynchronous progress_listener Is Necessary
At first, I had not thought deeply about why progress notifications must be asynchronous, but as the implementation progressed, the necessity became clear.
First, WebSocket communication is inherently asynchronous. For example, handling WebSocket in FastAPI looks like this:
from fastapi import FastAPI, WebSocket
from datetime import datetime
import asyncio
app = FastAPI()
@app.websocket("/ws/{task_id}")
async def websocket_endpoint(websocket: WebSocket, task_id: str):
"""WebSocket endpoint - establishes bidirectional communication with the client"""
await websocket.accept() # Accept the connection (asynchronous)
# Define an async function for progress notifications
async def send_progress(percent: float, message: str, details: dict = None):
"""Async function that sends progress info over the WebSocket"""
payload = {
"type": "progress",
"task_id": task_id,
"percent": percent,
"message": message,
"timestamp": datetime.now().isoformat(),
"details": details or {}
}
# Send JSON over the WebSocket (this is an async operation)
await websocket.send_json(payload)
try:
# Run the heavy processing (this is the heart of the problem)
result = await process_heavy_video_task(
task_id=task_id,
progress_listener=send_progress # pass the async function
)
# Notify completion
await websocket.send_json({
"type": "complete",
"task_id": task_id,
"result": result,
"timestamp": datetime.now().isoformat()
})
except Exception as e:
# Notify the error
await websocket.send_json({
"type": "error",
"task_id": task_id,
"error": str(e),
"timestamp": datetime.now().isoformat()
})
finally:
await websocket.close()
Moreover, in real production environments, progress notification is not just about sending over WebSocket—multiple asynchronous operations need to run at the same time.
The First Implementation: A Naive Approach and the Thinking Behind It
Step 1: The Original Synchronous Processing
The video processing system I first implemented was standard synchronous code of the kind you often see in PyTorch tutorials. It worked without any problems.
import torch
import torchvision.transforms as transforms
from typing import List, Any
import numpy as np
def process_video_frames(
model: torch.nn.Module,
frames: List[np.ndarray],
device: str = "cuda"
) -> List[torch.Tensor]:
"""
Synchronous function that processes video frames
Args:
model: PyTorch deep learning model
frames: list of video frames (numpy arrays)
device: execution device ("cuda" or "cpu")
Returns:
List of processed frames
"""
# Define preprocessing (image normalization, etc.)
transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((512, 512)), # Match the model's input size
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
processed_frames = []
total_frames = len(frames)
print(f"Starting processing of {total_frames} frames...")
# Process each frame sequentially
for i, frame in enumerate(frames):
# Preprocessing: convert the numpy array to a tensor
input_tensor = transform(frame)
input_tensor = input_tensor.unsqueeze(0) # Add a batch dimension
# Transfer to the GPU
input_tensor = input_tensor.to(device)
# Run inference (no gradient computation needed)
with torch.no_grad():
output = model(input_tensor)
# Postprocessing: move the result back to the CPU
result = output.cpu()
processed_frames.append(result)
# Simple progress display (every 100 frames)
if (i + 1) % 100 == 0:
memory_allocated = torch.cuda.memory_allocated() / 1024**3
print(f"Processed {i+1}/{total_frames} frames, "
f"GPU Memory: {memory_allocated:.2f}GB")
print(f"Processing complete! Total frames: {len(processed_frames)}")
return processed_frames
This implementation was simple, easy to understand, and worked fine for small videos.
The Moment I Thought "This Will Be Easy"
Once the requirement became "show processing progress on a web page in real time," asynchronous notification like the code above became necessary.
The video generation had originally been implemented as synchronous code, so my thinking went: "if the progress_listener that sends progress notifications is passed in as an async function, then I just need to make the processing function that calls it async as well." At first glance, this seems like a perfectly logical and simple solution.
Python has the async/await syntax, so calling an async function from another async function is easy—just use the await keyword.
In other words, just add async to the video generation code that had been synchronous, and call the asynchronous progress_listener from inside it. Problem solved... or so it should have been.
So, full of confidence, I wrote code like the following...
async def process_video_frames_async(
model: torch.nn.Module,
frames: List[np.ndarray],
progress_listener: Any, # asynchronous callback function
device: str = "cuda"
) -> List[torch.Tensor]:
"""
Async function that processes video frames (the problematic version)
This implementation has a serious flaw!
"""
# Define preprocessing (same as the sync version)
transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
processed_frames = []
total_frames = len(frames)
print(f"Starting async processing of {total_frames} frames...")
start_time = asyncio.get_event_loop().time()
for i, frame in enumerate(frames):
loop_start = asyncio.get_event_loop().time()
# === Synchronous processing starts here ===
# Preprocessing: convert the numpy array to a tensor
input_tensor = transform(frame)
input_tensor = input_tensor.unsqueeze(0)
# Transfer to the GPU
input_tensor = input_tensor.to(device)
# Run inference
with torch.no_grad():
output = model(input_tensor)
# Postprocessing: move the result back to the CPU
result = output.cpu()
processed_frames.append(result)
# === Synchronous processing ends here ===
# Notify progress asynchronously! Problem solved! (or so I thought)
current_progress = (i + 1) / total_frames * 100
elapsed_time = asyncio.get_event_loop().time() - start_time
estimated_total = elapsed_time / (i + 1) * total_frames
remaining_time = estimated_total - elapsed_time
await progress_listener(
percent=current_progress,
message=f"Processing frame {i+1}/{total_frames}",
details={
"current_frame": i + 1,
"total_frames": total_frames,
"elapsed_seconds": round(elapsed_time, 2),
"remaining_seconds": round(remaining_time, 2),
"fps": round((i + 1) / elapsed_time, 2) if elapsed_time > 0 else 0,
"memory_gb": torch.cuda.memory_allocated() / 1024**3
}
)
# Detailed logging every 10 frames
if (i + 1) % 10 == 0:
loop_time = asyncio.get_event_loop().time() - loop_start
memory_used = torch.cuda.memory_allocated() / 1024**3
print(f"Frame {i+1}: Loop time: {loop_time:.3f}s, "
f"GPU Memory: {memory_used:.2f}GB")
return processed_frames
The code was simple, the type checker raised no warnings, and there were no syntax errors.
"That's a wrap."
Or so it should have been...
The Unexpected Result: Memory Exploded
First Run in the Test Environment
Processing a test video (30 seconds, 900 frames) in the development environment worked fine. "See, it really was easy," I thought with relief.
But when I ran it on a production-scale video (10 minutes, 18,000 frames), something frightening happened.
Starting async processing of 18000 frames...
Frame 10: Loop time: 0.125s, GPU Memory: 1.82GB
Frame 100: Loop time: 0.156s, GPU Memory: 3.45GB
Frame 500: Loop time: 0.203s, GPU Memory: 8.72GB
Frame 1000: Loop time: 0.298s, GPU Memory: 16.34GB
Frame 1500: Loop time: 0.412s, GPU Memory: 24.81GB
Frame 2000: Loop time: 0.589s, GPU Memory: 33.15GB
Frame 2500: Loop time: 0.834s, GPU Memory: 41.62GB
Frame 2800: Loop time: 1.205s, GPU Memory: 47.23GB
RuntimeError: CUDA out of memory. Tried to allocate 512.00 MiB
(GPU 0; 47.54 GiB total capacity; 46.89 GiB already allocated;
324.00 MiB free; 47.12 GiB reserved in total by PyTorch)
The Abnormal Behavior We Observed
Beyond the growth in memory usage, other strange phenomena appeared as well:
Dramatic slowdown in processing speed
- First frame: 0.125 seconds
- Frame 2,800: 1.205 seconds (about 10x slower)
CPU memory also grew
- The system monitor showed RAM usage climbing steadily
The process stopped responding
- After a while, it reached a state where even Ctrl+C could not stop it...
Root Cause Analysis: Why Did the Memory Leak (Strictly, Delayed Release) Occur?
Cause 1: A Design That Accumulates All Results in Memory
The first and most obvious problem was the sloppy implementation of accumulating every result in a list. Reading recent papers and their reference code, this pattern is genuinely everywhere—but it is something you should never do in production, and I did it anyway.
(In my defense, the plan was to start with this rough implementation just to verify behavior, and to add proper batching and chunking later—but since it more or less worked, I forgot to fix it. I had even filed an issue for it and still missed it.)
However, this alone did not explain everything.
That is because the synchronous version had the exact same implementation, and the synchronous version worked correctly.
processed_frames = [] # All frame results are held here
for frame in frames: # with 18,000 frames...
output = model(frame) # if each output is about 50MB...
processed_frames.append(output) # 900GB needed in total!?
Let's do the actual math:
- Input for one frame: 512×512×3×4 bytes = about 3MB
- Intermediate layer outputs of the model: about 20MB
- Final output: about 30MB
- Total: about 50MB per frame
- 18,000 frames × 50MB = 900GB
This is clearly a problem, but it should have been just as much of a problem in the synchronous version.
"But it worked in the synchronous version!"
So why did only the asynchronous version crash?
Cause 2: The Peculiarities of Memory Management in async Functions, Plus PyTorch's CUDA Cache
This was the biggest pitfall of all.
Python's async functions manage memory differently from ordinary functions.
When an await is reached, the coroutine's "current local variable state" is saved in its frame at that point, and any variables that might be referenced by subsequent code are not released.
As a result, large tensors on the GPU and inference results tend to remain in memory even after the next iteration begins.
On top of that, PyTorch is designed so that even after it finishes using GPU memory (VRAM), it does not return it to the OS right away but keeps it as an internal cache. In the synchronous version, variables reach the end of their lifetime quickly, so the cache gets overwritten and rarely grows much. In the asynchronous version, however, variable lifetimes are extended, the cache region balloons, and the result is a rapid increase in VRAM usage.
# Memory management in a synchronous function
def sync_process():
for i in range(1000):
# Create a large tensor (4MB)
big_tensor = torch.randn(1000, 1000).cuda()
# Process with the model (another 4MB)
result = model(big_tensor)
# The key point: at the end of the loop, big_tensor is overwritten on
# the next iteration, and the old one immediately becomes eligible for garbage collection
# Python's reference count drops to 0 and the memory is freed
However, things work differently in an async function.
# Memory management in an asynchronous function
async def async_process():
for i in range(1000):
# Create a large tensor (4MB)
big_tensor = torch.randn(1000, 1000).cuda()
# Process with the model (another 4MB)
result = model(big_tensor)
# This is the heart of the problem!
await something() # await suspends the coroutine's execution
# At this point, Python must save the coroutine's state, so references to big_tensor and result are retained.
In the synchronous version, references are dropped as soon as the loop iteration ends, so PyTorch's cache region is reused immediately. In the asynchronous version, variables with extended lifetimes keep occupying the cache region.
As a result, the combination of extended variable lifetimes and PyTorch's VRAM cache makes GPU memory snowball.
To add a bit more detail on how coroutine state is saved, it looks something like this:
async def detailed_async_process():
# All of these variables are saved in the coroutine's frame
local_var_1 = create_large_object() # 100MB
local_var_2 = create_another_object() # 50MB
for i in range(100):
loop_var = create_loop_object() # 10MB
# At the await point, the following are saved:
# - local_var_1, local_var_2
# - loop_var
# - the loop counter i
# - all other local variables
await async_operation()
# Problem: even after loop_var is overwritten,
# the previous loop_var may still be referenced
Cause 3: The Compatibility Problem Between PyTorch's CUDA Operations and asyncio
PyTorch's CUDA operations are inherently synchronous, blocking operations.
These operations are in fact asynchronous internally, using CUDA streams, but at the Python level they appear synchronous.
# Internal behavior of PyTorch CUDA operations
tensor_gpu = tensor.cuda() # Enqueued on the CUDA stream at this point
output = model(tensor_gpu) # Runs on the GPU (Python waits)
result = output.cpu() # GPU-to-CPU transfer (synchronization point)
async functions that perform these blocking operations run into the following problems:
Delayed timing of memory release
async def memory_leak():
for i in range(1000):
# Allocate GPU memory
gpu_tensor = create_gpu_tensor()
# Blocking operation
process_gpu_tensor(gpu_tensor)
# Memory should be freed here, but...
await notify_progress()
# in reality it may not have been freed yet
Blocking the event loop
async def problematic():
# This heavy operation blocks the event loop
result = heavy_gpu_operation() # other coroutines cannot run
await something() # control finally passes to other coroutines
Cause 4: Garbage Collection Timing
Python's garbage collection (GC) normally kicks in as soon as a reference count reaches zero, but when circular references exist, it has to wait for the periodic GC pass.
In asynchronous functions, the coroutine object itself holds references to variables, which makes GC timing harder to predict.
import gc
import sys
async def gc_timing_issue():
for i in range(1000):
obj = LargeObject()
# Check the reference count
print(f"Reference count: {sys.getrefcount(obj)}")
# Normally 2 (the obj variable + a temporary reference inside getrefcount)
await something()
# After the await, the reference count may have increased
print(f"Reference count after await: {sys.getrefcount(obj)}")
# It can be 3 or more (a reference from the coroutine frame)
# Memory is not freed unless GC is run manually
gc.collect()
So, going async brings you face to face with coroutine-specific memory-management complexity you normally never think about. It is worth remembering that merely adding the "async" keyword changes quite a lot under the hood.
Implementing the Solutions: Several Approaches in Detail
Now that the causes have come into view, let's think through the solutions.
Solution 1: A Fundamental Fix with Streaming Processing
First, before we even talk about sync versus async, the outrageous implementation that holds all data in memory has to be fixed. I switched to streaming processing that writes processed data to disk immediately. Somewhat unexpectedly, this turned out to be the most effective solution. Looping around while stacking up huge amounts of memory is a bad idea to begin with—especially in coroutines.
import os
import json
import gc
import tempfile
from pathlib import Path
from datetime import datetime
from typing import List, Optional, Callable
import numpy as np
import torch
def get_transform():
import torchvision.transforms as T
return T.Compose([
T.ToPILImage(),
T.Resize((512, 512)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
async def process_video_frames_streaming(
model: torch.nn.Module,
frames: List[np.ndarray],
output_dir: Optional[str] = None,
progress_listener: Optional[Callable] = None,
batch_size: int = 32,
device: str = "cuda",
use_amp: bool = False,
) -> str:
# Prepare the output directory
if output_dir is None:
output_dir = tempfile.mkdtemp(prefix="video_frames_")
else:
os.makedirs(output_dir, exist_ok=True)
output_path = Path(output_dir)
# Save metadata
metadata = {
"total_frames": len(frames),
"batch_size": batch_size,
"model_name": model.__class__.__name__,
"device": device,
"processed_at": datetime.now().isoformat(),
}
with open(output_path / "metadata.json", "w", encoding="utf-8") as f:
json.dump(metadata, f, ensure_ascii=False, indent=2)
total_frames = len(frames)
# Important: create the transform once, outside the loop
transform = get_transform()
# Put the model in inference mode
model.eval()
if device == "cuda":
model.to(device)
print(f"Starting streaming processing of {total_frames} frames")
print(f"Output directory: {output_dir}")
print(f"Batch size: {batch_size}")
# Choose the autocast dtype for your environment (float16/TF32 etc. on Ampere or later)
autocast_dtype = torch.float16 if use_amp else None
# Process batch by batch
for batch_start in range(0, total_frames, batch_size):
batch_end = min(batch_start + batch_size, total_frames)
batch = frames[batch_start:batch_end]
# ---- Start of batch processing ----
# inference_mode is the most lightweight for inference (even more optimized than no_grad)
with torch.inference_mode():
for local_idx, frame in enumerate(batch):
global_idx = batch_start + local_idx
# Preprocessing (CPU)
input_tensor = transform(frame) # CxHxW, float32 on CPU
# Pinning memory makes non_blocking transfer effective (optional)
input_tensor = input_tensor.pin_memory() if device == "cuda" else input_tensor
input_tensor = input_tensor.unsqueeze(0) # NxCxHxW
# Transfer to the GPU (non_blocking=True reduces waiting)
if device == "cuda":
input_tensor = input_tensor.to(device, non_blocking=True)
# Inference (mixed precision if needed)
if use_amp and device == "cuda":
with torch.autocast(device_type="cuda", dtype=autocast_dtype):
output = model(input_tensor)
else:
output = model(input_tensor)
# Move back to the CPU and save (freed from memory immediately)
result_cpu = output.detach().to("cpu", copy=True)
frame_filename = output_path / f"frame_{global_idx:06d}.pt"
# Optionally shrink further (e.g. half precision): result_cpu = result_cpu.half()
torch.save(
{
"frame_index": global_idx,
"output": result_cpu,
"input_shape": tuple(frame.shape),
"processing_time": datetime.now().isoformat(),
},
frame_filename,
)
# Free explicitly
del input_tensor, output, result_cpu
# After each batch, tidy up GPU memory / CPU garbage collection (frequency depends on your use case)
if device == "cuda":
torch.cuda.empty_cache()
# torch.cuda.ipc_collect() is usually unnecessary on recent PyTorch
gc.collect()
# Log memory status
memory_stats = None
if device == "cuda":
allocated = torch.cuda.memory_allocated() / 1024**3
reserved = torch.cuda.memory_reserved() / 1024**3
free = torch.cuda.mem_get_info()[0] / 1024**3
memory_stats = {"allocated": allocated, "reserved": reserved, "free": free}
print(
f"Batch {batch_start}-{batch_end}: "
f"GPU Memory - Allocated: {allocated:.2f}GB, "
f"Reserved: {reserved:.2f}GB, Free: {free:.2f}GB"
)
# Notify progress asynchronously (per batch)
if progress_listener:
progress = batch_end / total_frames * 100.0
await progress_listener(
percent=progress,
message=f"Processed {batch_end}/{total_frames} frames",
details={
"batch_start": batch_start,
"batch_end": batch_end,
"batch_size": len(batch),
"output_dir": str(output_dir),
"memory_stats": memory_stats,
},
)
# Yield control back to the event loop (important)
await asyncio.sleep(0)
print(f"Processing complete! Results saved to: {output_dir}")
# Completion notification
if progress_listener:
await progress_listener(
percent=100.0,
message="Processing complete",
details={"total_frames": total_frames, "output_dir": str(output_dir), "success": True},
)
return str(output_dir)
Solution 2: Proper Asynchronization Using a Thread Pool
The other approach keeps the PyTorch processing synchronous, runs it on a separate thread with run_in_executor, and still delivers progress notifications through asynchronous callbacks.
Problem: You Cannot Call an Async Function from Another Thread
Offloading PyTorch's synchronous work to another thread is fine as far as it goes, but this approach stumbles on one point: a synchronous function running on another thread cannot call the asynchronous progress_listener.
# This does not work
def pytorch_process_sync(frames, progress_listener):
"""Synchronous function that runs on a separate thread"""
for i, frame in enumerate(frames):
result = model(frame)
# Error! You cannot call an async function from a sync function
await progress_listener(i) # SyntaxError
Solution: Process in Batches and Notify Asynchronously
So I tried a slightly tricky design: at each batch boundary, control returns to the main thread, and the asynchronous callback is invoked there.
import asyncio
from concurrent.futures import ThreadPoolExecutor
class VideoProcessor:
"""
Uses a thread pool while sending progress via asynchronous callbacks
"""
def __init__(self, model: torch.nn.Module, device: str = "cuda"):
self.model = model
self.device = device
self.executor = ThreadPoolExecutor(max_workers=1)
def _process_batch_sync(
self,
batch: List[np.ndarray],
start_idx: int
) -> dict:
"""
Process a batch synchronously (runs on a separate thread)
Sends no progress notifications; just returns results and metadata
"""
results = []
processing_times = []
for local_idx, frame in enumerate(batch):
start_time = time.time()
with torch.no_grad():
# PyTorch processing
tensor = self.transform(frame).unsqueeze(0).to(self.device)
output = self.model(tensor)
result = output.cpu()
# Save
global_idx = start_idx + local_idx
output_path = f"output/frame_{global_idx:06d}.pt"
torch.save(result, output_path)
results.append(output_path)
# Clear memory
del tensor, output, result
processing_times.append(time.time() - start_time)
# Clear GPU memory
if self.device == "cuda":
torch.cuda.empty_cache()
# Return the batch results
return {
"saved_paths": results,
"processing_times": processing_times,
"memory_used": torch.cuda.memory_allocated() / 1024**3 if self.device == "cuda" else 0,
"thread_id": threading.current_thread().ident
}
async def process_frames_async(
self,
frames: List[np.ndarray],
progress_listener: Callable, # asynchronous callback
batch_size: int = 32
) -> List[str]:
"""
The main asynchronous interface
Processes batch by batch and sends progress asynchronously after each batch
"""
loop = asyncio.get_event_loop()
total_frames = len(frames)
all_saved_paths = []
total_processing_time = 0
for batch_start in range(0, total_frames, batch_size):
batch_end = min(batch_start + batch_size, total_frames)
batch = frames[batch_start:batch_end]
# Process the batch on a separate thread
batch_result = await loop.run_in_executor(
self.executor,
self._process_batch_sync,
batch,
batch_start
)
# Accumulate results
all_saved_paths.extend(batch_result["saved_paths"])
total_processing_time += sum(batch_result["processing_times"])
# We are back on the main thread here,
# so we can call the asynchronous callback!
progress = batch_end / total_frames * 100
await progress_listener(
percent=progress,
message=f"Processed batch {batch_start}-{batch_end}/{total_frames}",
details={
"current_batch": batch_start // batch_size + 1,
"total_batches": (total_frames + batch_size - 1) // batch_size,
"batch_processing_time": sum(batch_result["processing_times"]),
"average_frame_time": sum(batch_result["processing_times"]) / len(batch),
"memory_gb": batch_result["memory_used"],
"fps": len(all_saved_paths) / total_processing_time if total_processing_time > 0 else 0,
"thread_id": batch_result["thread_id"]
}
)
# Completion notification
await progress_listener(
percent=100,
message="Processing complete",
details={
"total_frames": total_frames,
"total_time": total_processing_time,
"average_fps": total_frames / total_processing_time
}
)
return all_saved_paths
The Benefits of This Approach
This approach has one important advantage: the event loop is never blocked.
Because PyTorch's heavy processing runs on a separate thread, the main thread can keep handling other WebSocket requests and API calls. Even while one user runs a heavy video-processing job, responses to other users stay fast, and the responsiveness of the whole system is preserved.
In other words, you avoid an implementation where heavy synchronous work gets mixed into your asynchronous code.
Incidentally,
when you offload GPU work to another thread in PyTorch, the CUDA context is shared between threads, but
- CUDA calls are fundamentally not thread-safe in multithreaded environments (the official documentation warns about this)
- —hitting the GPU from multiple threads at once can cause performance degradation or unexpected errors
, so keep this in mind.
Usage Example
# Usage example with FastAPI
@app.post("/process-video")
async def process_video_endpoint(video_id: str):
processor = VideoProcessor(model)
# Async callback that sends progress over the WebSocket
async def send_progress(percent, message, details):
await websocket.send_json({
"type": "progress",
"video_id": video_id,
"percent": percent,
"message": message,
"details": details,
"timestamp": datetime.now().isoformat()
})
# Also record to the database (asynchronous)
await db.execute(
"UPDATE video_tasks SET progress = ? WHERE id = ?",
(percent, video_id)
)
# Run the processing (with the async callback)
results = await processor.process_frames_async(
frames=load_video_frames(video_id),
progress_listener=send_progress,
batch_size=32
)
return {"status": "complete", "results": results}
Solution 3: When You Need Finer-Grained Progress Notifications
If you really do need per-frame progress notifications rather than per-batch, another option is to use asyncio.run_coroutine_threadsafe.
class DetailedProgressProcessor:
"""
Provides detailed per-frame progress notifications
"""
def _process_with_detailed_progress(
self,
frames: List[np.ndarray],
loop: asyncio.AbstractEventLoop,
async_progress_listener: Callable
) -> List[str]:
"""
Runs on a separate thread and sends progress for every frame
"""
results = []
total = len(frames)
for i, frame in enumerate(frames):
# PyTorch processing
with torch.no_grad():
tensor = self.transform(frame).to(self.device)
output = self.model(tensor)
result = output.cpu()
# Save
path = f"output/frame_{i:06d}.pt"
torch.save(result, path)
results.append(path)
# The magic that calls an async function from another thread!
future = asyncio.run_coroutine_threadsafe(
async_progress_listener(
percent=(i + 1) / total * 100,
message=f"Processing frame {i+1}/{total}",
details={"frame_index": i}
),
loop # the main thread's event loop
)
# Wait for the result if needed (optional)
try:
future.result(timeout=0.1) # Expect completion within 100ms
except TimeoutError:
# Processing continues even if the progress notification is delayed
pass
# Clear memory periodically
if i % 10 == 0 and self.device == "cuda":
torch.cuda.empty_cache()
return results
async def process_with_detailed_progress(
self,
frames: List[np.ndarray],
progress_listener: Callable
) -> List[str]:
"""
Processing with detailed progress notifications
"""
loop = asyncio.get_event_loop()
# Run the processing on a separate thread
results = await loop.run_in_executor(
None,
self._process_with_detailed_progress,
frames,
loop, # pass the current event loop
progress_listener
)
return results
Caveats When Calling "torch.cuda.empty_cache()"
One thing to keep in mind about torch.cuda.empty_cache: the code comments say "clear memory periodically," but this does not physically return memory to the OS; it merely empties PyTorch's cache region.
So while it is useful when you need to forcibly increase free space, calling it too often can actually increase the cost of re-allocating memory.
Performance Comparison
Actual test results (using asynchronous callbacks):
| Implementation | GPU Memory | Notification Granularity | Implementation Complexity |
|---|---|---|---|
| Original async implementation (broken) | Grows → OOM | Per frame | ★☆☆☆ |
| Streaming (Solution 1) | Steady at 3.2GB | Per batch | ★★☆☆ |
| Thread pool (Solution 2) | Steady at 3.5GB | Per batch | ★★★☆ |
Summary
The Pitfalls Where Asynchronous Programming Meets GPU Processing
The most important lesson from this experience is that casually converting synchronous code to async can land you in coroutine-specific memory-management traps. In an ordinary synchronous function, variables are overwritten on each loop iteration and memory is released naturally. In an async function, however, once an await is involved, any local variables at that point that may still be referenced by subsequent code are retained by the coroutine object, causing unexpected memory retention. This is a fatal problem that is easy to overlook if your thinking stops at "just make it async and you can send progress notifications."
Another key realization was that continuously allocating GPU memory can produce a memory-leak-like state due to the interplay between garbage collection timing and CUDA memory management. PyTorch's CUDA operations are inherently synchronous; forcing them to run inside an async function threw off the timing of memory release, and even a large amount of VRAM ended up exhausted.
A Practical Approach to Solving It
If you run into this problem, start by breaking up memory usage with streaming processing—that felt like the right first move. Rather than processing all frames at once, processing in batches and saving to disk immediately keeps delayed memory release from reaching dangerous levels. For our requirements, this fix alone got everything through unit tests and load tests, and the problem was solved. It also kept changes to the existing code to a minimum.
If that still cannot meet your requirements, consider separating the processing with a thread pool and run_in_executor into its own thread. Moving the PyTorch work to a separate thread prevents the event loop from blocking while still letting the asynchronous callbacks work correctly. The implementation gets somewhat more complex, but it gives you more flexible control.
Finally, it is important to implement explicit memory management and never skip it.
Even if it looks redundant, treat it as necessary "insurance" for stable operation: delete variables with del statements and run torch.cuda.empty_cache() periodically.
By applying these measures step by step, you can safely make "modern asynchronous processing" and "heavy GPU processing" coexist. This is not a problem you can solve with the surface-level understanding of "just use async/await," but with the right approach it can always be solved, as this case shows.
Thank you for reading all the way to the end once again! See you next time!
