fix: Fix race condition in CUDA Simulator

[ccam80]

Description

There is a race condition in the CUDA simulator, specifically in the swapped_cuda_module context manager.

I use the simulator for quick-running CI to avoid using up precious free GPU minutes. Occasionally, I get this error:

AttributeError: tid=[0, 13, 0] ctaid=[0, 0, 0]: module 'numba.cuda' has no attribute 'local'

It is raised from a different thread each time. The error arose more commonly after I began allocating arrays in a small helper function in its own module. The error is similar to the one raised in numba/numba#1844.

Each thread in the simulator is a threading.Thread object, so they share memory. Every time a device function is called, it is wrapped in this context manager:

numba-cuda/numba_cuda/numba/cuda/simulator/kernelapi.py

Lines 494 to 509 in aff41e9

	@contextmanager
	def swapped_cuda_module(fn, fake_cuda_module):
	from numba import cuda
	fn_globs = fn.__globals__
	# get all globals that is the "cuda" module
	orig = dict((k, v) for k, v in fn_globs.items() if v is cuda)
	# build replacement dict
	repl = dict((k, fake_cuda_module) for k, v in orig.items())
	# replace
	fn_globs.update(repl)
	try:
	yield
	finally:
	# revert
	fn_globs.update(orig)

Race:

Thread A and Thread B are executing device functions in the same python module. They don't need to be the same function. They must be in a separate file from the kernel definition, as the kernel replaces references on entry, run all threads, and restores only after all threads have exited.

1. Thread A launches and swaps numba.cuda for fake_cuda, yields.
2. Thread B launches and gets orig = {} and repl = {}, as no references to cuda exist in it's __globals__ dict. Thread B yields.
3. Thread A exits, replacing fake_cuda with numba.cuda
4. Thread B calls e.g. cuda.local.array, and sees replaced reference to numba.cuda. local is not imported as part of numba.cuda when NUMBA_ENABLE_CUDASIM==1, so the error is thrown.

MWE

The Gist below contains a script that reliably causes the error on my machine. It takes ~200s to hit the race on my machine, typically, so I have not added it to the test suite. It does seem to fail faster on xdist, but it has a very long runtime when it doesn't fail.

Reproducer
Place all three files in the same directory, and run cudasim_race_mwe.

Fix

This PR implements a per-module lock and reference count, so that the first entrance to the context for a module replaces cuda -> fake_cuda, and the last thread to exit restores fake_cuda -> cuda. There may be a performance hit associated for simulated kernels with many device function calls from many modules, but this should be small, as all threads except for the first entrant and last exit perform a single integer comparison and increment/decrement an integer counter under the lock. The short "benchmark" run in the MWE did not change duration between the patched and unpatched versions on my machine.

@contextmanager
def swapped_cuda_module(fn, fake_cuda_module):
    from numba import cuda
    fn_globs = fn.__globals__
    gid = id(fn_globs)

    # Use a lock per-modules to avoid cross-locking other modules
    lock = _globals_locks[gid]

    with lock:
        # Scan and replace globals with fake module on first entrance only
        if _swap_refcount[gid] == 0:
            orig = {k: v for k, v in fn_globs.items() if v is cuda}
            _swap_orig[gid] = orig
            for k in orig:
                fn_globs[k] = fake_cuda_module

        # Increment the reference counter on every entrance
        _swap_refcount[gid] += 1
    try:
        yield
    finally:
        with lock:
            # Decrement "number of modules using fake CUDA" counter on exit
            _swap_refcount[gid] -= 1

            # Last thread to leave the context restores real cuda
            if _swap_refcount[gid] == 0:
                fn_globs.update(_swap_orig.pop(gid))
                del _swap_refcount[gid]
                del _globals_locks[gid]

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

Implemented per-module locking with reference counting to fix race condition in swapped_cuda_module context manager. The fix prevents threads from interfering with each other's module swaps by tracking how many threads are using each module's globals and only performing the swap on first entry and restoration on last exit.

• Added _locks_register_lock to safely create per-module locks
• Added _globals_locks defaultdict to store per-module locks
• Added _swap_refcount to track number of active threads per module
• Added _swap_orig to store original module references
• Modified swapped_cuda_module to use reference counting pattern

Confidence Score: 4/5

• Safe to merge after fixing the lock cleanup - the core race condition fix is sound
• The threading logic correctly addresses the race condition with proper locking and reference counting, but there's a memory leak where _globals_locks[gid] is never deleted when refcount reaches zero
• numba_cuda/numba/cuda/simulator/kernelapi.py needs the lock cleanup fix at line 534

Important Files Changed

Filename	Overview
numba_cuda/numba/cuda/simulator/kernelapi.py	Added per-module locking and reference counting to fix race condition in context manager, but missing cleanup of lock registry causing memory leak

[Copilot]

Pull request overview

This PR fixes a race condition in the CUDA simulator's swapped_cuda_module context manager that occurs when multiple simulated threads simultaneously call device functions from the same Python module. The fix implements per-module locks and reference counting to ensure thread-safe swapping and restoration of module globals.

Key Changes:

• Added per-module locking mechanism using defaultdict(threading.Lock) to synchronize access to module globals
• Implemented reference counting to track the number of active threads using the fake CUDA module in each Python module
• Modified the swap logic so only the first entering thread performs the swap, and only the last exiting thread performs the restoration

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[greptile-apps]

Additional Comments (1)

1. numba_cuda/numba/cuda/simulator/kernelapi.py, line 493-504 (link)

   logic: race condition: defaultdict access is not thread-safe when key doesn't exist

   when multiple threads simultaneously call device functions from the same module for the first time, they race at line 504. threads can create different lock objects for the same gid, defeating the per-module locking

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[gmarkall]

/ok to test 1d3a392

[greptile-apps]

Additional Comments (1)

1. numba_cuda/numba/cuda/simulator/kernelapi.py, line 531 (link)

   logic: missing cleanup for _globals_locks[gid] - the lock itself is never deleted, causing memory leak for each unique module that uses the simulator

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[gmarkall]

/ok to test 8fb2b52

[gmarkall]

Thanks for the PR! I can't seem to find a way to hit the race with your reproducer, but the theory you outlined and the fix seem to make sense, so I'm comfortable with it.

I pushed a small fix so that the format check passes, and I'll set this to merge once the CI run is done.

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: memory leak: _globals_locks[gid] never cleaned up

Suggested change

	del _swap_refcount[gid]
	del _swap_refcount[gid]
	del _globals_locks[gid]