budgetpool

Memory-budget-aware parallel execution pool for Python.

A drop-in replacement for ProcessPoolExecutor that calculates safe worker counts based on available memory, prevents OOM crashes, and provides backpressure on task submission.

Why?

ProcessPoolExecutor spawns workers based on CPU count, ignoring memory. When each worker loads a large dataset or model, this easily causes OOM kills:

# Dangerous — 8 workers × 3GB each = 24GB on a 16GB machine
with ProcessPoolExecutor(max_workers=os.cpu_count()) as pool:
    results = list(pool.map(run_backtest, param_grid))

BudgetPool fixes this by computing safe worker counts from your memory budget:

from budgetpool import BudgetPool

# Safe — fits within 12GB, 2GB per worker → 6 workers max
with BudgetPool(memory_budget_gb=12.0, memory_per_worker_gb=2.0) as pool:
    results = list(pool.map(run_backtest, param_grid))

Install

pip install budgetpool

Quick Start

Basic Usage

from budgetpool import BudgetPool

with BudgetPool(memory_per_worker_gb=2.0) as pool:
    results = list(pool.map(heavy_func, items))
    print(f"Used {pool.num_workers} workers")

With Explicit Budget

with BudgetPool(
    memory_budget_gb=12.0,      # Total budget for all workers
    memory_per_worker_gb=2.0,   # Estimated peak per worker
    max_workers=8,              # CPU cap (optional)
) as pool:
    futures = [pool.submit(process, item) for item in items]
    results = [f.result() for f in futures]

Memory Monitoring

from budgetpool import get_memory_info, safe_worker_count

# Check system memory
info = get_memory_info()
print(f"Total: {info.total_gb:.1f}GB, Available: {info.available_gb:.1f}GB")
print(f"Safe for workers: {info.free_for_workers_gb:.1f}GB")

# Calculate worker count without creating a pool
n = safe_worker_count(memory_per_worker_gb=3.0)
print(f"Safe worker count: {n}")

How It Works

1. Startup: Reads system memory (via psutil), respects cgroup limits in containers
2. Worker calculation: min(budget ÷ per_worker, cpu_count, max_workers)
3. Backpressure: Blocks submit() when pending tasks exceed max_pending (default: 2× workers)
4. Runtime checks: Warns at 85% memory usage, raises MemoryBudgetExceeded at 95%

API Reference

BudgetPool

Methods: submit(), map(), shutdown() — same signatures as ProcessPoolExecutor.

Properties: num_workers, memory_budget_gb, memory_info, stats.

PoolStats

Cumulative statistics available via pool.stats:

safe_worker_count(memory_per_worker_gb, max_workers=None) → int

Standalone function to calculate safe worker count without creating a pool.

get_memory_info() → MemoryInfo

Returns a MemoryInfo dataclass with total_gb, available_gb, used_gb, percent, and free_for_workers_gb.

CLI

Check system memory and safe worker counts:

$ python -m budgetpool status
System Memory
  Total:          16.0 GB
  Available:      6.5 GB
  Used:           8.1 GB (60%)
  Free for workers: 4.5 GB

Safe Worker Counts
  0.5 GB/worker → 8 workers
  1.0 GB/worker → 4 workers
  2.0 GB/worker → 2 workers
  4.0 GB/worker → 1 workers
  8.0 GB/worker → 1 workers

Container Support

budgetpool automatically detects cgroup v1/v2 memory limits, so it works correctly inside Docker containers where psutil.virtual_memory().total would report host memory:

docker run --memory=4g python -c "
from budgetpool import get_memory_info
print(get_memory_info().total_gb)  # → 4.0, not host memory
"

Requirements

• Python 3.10+
• psutil ≥ 5.9.0

License

Apache-2.0