[BUG] Subprocess GPU mode loses the pool's GPU assignment
Describe the bug
With COCOINDEX_RUN_GPU_IN_SUBPROCESS=1, coco.current_gpu() returns None inside the function, so code that does f"cuda:{coco.current_gpu() or 0}" always ends up on GPU 0. The pool assigns a GPU, but execute_in_subprocess() only pickles (fn, args, kwargs) and _current_gpus is a ContextVar, so nothing reaches the child. The in-process path passes it through _run_with_gpu_context().
Related: _get_pool() is a single ProcessPoolExecutor(max_workers=1) shared by all GPUs, so the calls also run one at a time.
I know this is a known limitation (listed in #2224, and there's a UserWarning for it). Filing it so there's something to track. The warning says calls run on the same GPU, but the child actually gets no assignment at all.
To Reproduce
Confirmed on a real device first. One RTX 4070, torch 2.14.0+cu126, three calls doing a real matmul in subprocess mode:
call assigned ran on mem MB pid
0 None None - 17544
1 None None - 17544
2 None None - 17544assigned is None, so the function can't pick a device at all; it only works because user code falls back to cuda:0.
Cost of the serialization, 100 tasks of fixed work across 8 GPUs:
wall throughput workers GPUs the calls saw
shipped 8.02s 12.5 t/s 1 [None]All 100 queued behind one worker. Seven GPUs idle.
A script that reproduces the None without needing a GPU is below. It exits 1 while the bug is present.
import asyncio
import os
import sys
import time
import cocoindex as coco
from cocoindex._internal.runner import GPURunner
N = 4
HOLD = 1.0
def probe(tag: str) -> dict:
"""Runs under the GPU runner; reports the GPU it can actually see."""
info = {
"tag": tag,
"pid": os.getpid(),
"gpu": coco.current_gpu(),
"cvd": os.environ.get("CUDA_VISIBLE_DEVICES"),
}
time.sleep(HOLD) # hold the GPU so all N calls overlap
return info
class RecordingGPURunner(GPURunner):
"""GPURunner that also records what the pool handed out."""
def __init__(self, fraction: float = 1.0) -> None:
super().__init__(fraction)
self.assigned: list[int] = []
async def _acquire_gpu(self) -> int:
gpu_id = await super()._acquire_gpu()
self.assigned.append(gpu_id)
return gpu_id
async def run_mode(subprocess_mode: bool):
os.environ["COCOINDEX_RUN_GPU_IN_SUBPROCESS"] = "1" if subprocess_mode else "0"
runner = RecordingGPURunner(1.0) # fresh instance re-reads the env var
t0 = time.perf_counter()
rows = await asyncio.gather(
*(runner.run_sync_fn(probe, f"job{i}") for i in range(N))
)
return rows, runner.assigned, time.perf_counter() - t0
def report(title, rows, assigned, secs):
print(f"\n--- {title} ---")
print(f"{'job':<7}{'pool assigned':<16}{'job saw':<10}{'pid'}")
for r, g in zip(rows, assigned):
print(f"{r['tag']:<7}GPU {g:<12}{str(r['gpu']):<10}{r['pid']}")
seen = sorted({r["gpu"] for r in rows}, key=lambda x: (x is None, x))
print(f"pool handed out : {sorted(assigned)}")
print(f"jobs actually saw: {seen}")
print(f"worker processes : {len({r['pid'] for r in rows})}")
print(f"wall clock : {secs:.2f}s (ideal {HOLD:.1f}s with {N} GPUs)")
return seen, secs
async def main() -> int:
coco.configure_gpu_pool(N)
print(f"pretending {N} GPUs | {N} calls holding a GPU for {HOLD:.1f}s each")
seen_in, t_in = report("IN-PROCESS (default)", *await run_mode(False))
seen_sub, t_sub = report(
"SUBPROCESS (COCOINDEX_RUN_GPU_IN_SUBPROCESS=1)", *await run_mode(True)
)
print(f"\nin-process : saw {seen_in} in {t_in:.2f}s")
print(f"subprocess : saw {seen_sub} in {t_sub:.2f}s ({t_sub / t_in:.1f}x slower)")
failed = seen_sub == [None] and seen_in != [None]
print("FAIL - subprocess mode lost the pool's GPU assignment" if failed
else "PASS - assignment propagated")
return 1 if failed else 0
if __name__ == "__main__":
sys.exit(asyncio.run(main()))Expected behavior
coco.current_gpu() reports the assigned id in subprocess mode too, and N GPUs run N calls in parallel.
CocoIndex Version
1.0.21, and runner.py is unchanged on main (3b4e54c3).
Additional context
I have a fix working locally: the assignment carried in the payload and restored in the child, and the pool sized so GPUPool, not the executor, decides concurrency. Same 100-task run with it applied:
wall throughput workers GPUs the calls saw
shipped 8.02s 12.5 t/s 1 [None]
fixed 1.30s 77.2 t/s 9 [0,1,2,3,4,5,6,7]On a single GPU with fraction=1.0 it is a wash (6.48s vs 6.36s, one worker either way), so the change costs nothing where there is no parallelism to unlock. The real-CUDA check passes: the tensor lands on the device current_gpu() reported. Multi-GPU numbers above are 8 simulated devices , I only have one physical card, so multi-device routing itself is not hardware-verified.
Before I open a PR, is this in scope for #2277, or a follow-up? And #2336 mentions the same cancellation and crash-retry problems on this path, though it keeps generic coco.GPU on the existing executor, so I wasn't sure if that supervisor is meant to cover this later.
Source: cocoindex-io/cocoindex