cuda.core: warn when a child is forked from a process that initialized CUDA #3127

Problem

CUDA state does not survive fork. A process forked from one that already initialized CUDA has undefined behavior when it uses CUDA: NVIDIA's CUDA Guidelines for Robust and Safety-Critical Programming, section 1.6 [fork], states: "CUDA does not duplicate any of its internal data structures or threads to other processes upon a call to fork. Any use of CUDA after a call to fork and before a subsequent call to exec has undefined behavior." (https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/, category Required, all compute capabilities).

With multiprocessing this shows up in two ways:

• fork start method: every child is forked from the parent. cuda.core already warns here, but only from check_multiprocessing_start_method, which runs while pickling an IPC object. A program that forks and uses CUDA without IPC gets no warning.

• forkserver start method: the forkserver is a fresh interpreter, and children fork from it, not from the user's process. By default the server preloads the user's main module (forkserver.py: the default preload list is ['__main__'], imported with spawn.import_main_path). If that import runs CUDA calls at the top level of the script, outside an if __name__ == "__main__": guard, the server holds CUDA state and every child inherits it. The failure appears in the children, far from the cause. forkserver is the Linux default since Python 3.14.

  The __main__ preload only takes effect on newer Pythons. Python 3.12's forkserver.py looks for a main_path key, but spawn.get_preparation_data emits init_main_from_path, so the server never imports __main__ and each child re-imports it the way spawn does. Checked: 3.12.3 does not preload __main__; 3.13.14 and 3.14.6 do. On every version, an explicit set_forkserver_preload([...]) of a module that initializes CUDA has the same effect.

Today the child gets no hint about the cause. Its first CUDA call fails with CUDA_ERROR_NOT_INITIALIZED, whose message says the driver "has not been initialized with cuInit() or that initialization has failed".

spawn is not affected in the same way: each child re-imports the main module in its own fresh process, and nothing forks from it.

A parent that initialized CUDA before starting spawn or forkserver workers is fine. The warning must not fire for that pattern, so it can't live in the parent.

Proposal

Warn in a forked child when it reaches its first cuda.core CUDA entry point, if it was forked from a process in which cuda.core initialized CUDA. The work is split three ways:

• Parent side: record the PID right after cuda.core calls cuInit.
• Fork hook: an os.register_at_fork(after_in_child=...) hook only compares PIDs and sets a flag. It does not warn, write output, or take locks at the fork boundary.
• Child side: the entry points that call cuInit check the flag and warn once.

Deferring the warning to the child's first CUDA call keeps children that never use CUDA silent. The check must run before the _is_cuInit guard in Device_ensure_cuda_initialized, because a forked child inherits _is_cuInit = True and would otherwise skip the function body entirely.

# cuda/core/_utils/cuda_utils.pyx

# PID of the process in which cuda.core called cuInit, or None.
cdef object _cuda_initialized_pid = None
# Set in a forked child whose parent had initialized CUDA; reported lazily on
# the child's first CUDA entry point, so forks that never use CUDA stay silent.
cdef bint _forked_after_cuda_init = False


cdef int mark_cuda_initialized() except -1:
    global _cuda_initialized_pid
    _cuda_initialized_pid = os.getpid()
    return 0


def _note_fork_after_cuda_init() -> None:
    # Runs in the child right after fork(): only record state here.
    global _forked_after_cuda_init
    if _cuda_initialized_pid is not None and _cuda_initialized_pid != os.getpid():
        _forked_after_cuda_init = True


cdef int warn_if_forked_after_cuda_init() except -1:
    global _forked_after_cuda_init
    if not _forked_after_cuda_init:
        return 0
    _forked_after_cuda_init = False  # warn once per child
    warnings.warn(
        "This process was forked from a process that had already initialized CUDA. ...",
        CUDAWarning,
        stacklevel=1,  # Cython callers add no Python frames: this is the user's call site
    )
    return 0


if hasattr(os, "register_at_fork"):
    os.register_at_fork(after_in_child=_note_fork_after_cuda_init)

Both cdef functions are declared in cuda_utils.pxd and cimported at the two places where cuda.core calls cuInit:

# cuda/core/_device.pyx
cdef inline int Device_ensure_cuda_initialized() except? -1:
    # A forked child inherits _is_cuInit, so this must run before the guard.
    warn_if_forked_after_cuda_init()
    global _is_cuInit
    if _is_cuInit is False:
        with _lock, nogil:
            HANDLE_RETURN(cydriver.cuInit(0))
            _is_cuInit = True
        mark_cuda_initialized()
        ...

# cuda/core/_memory/_pinned_memory_resource.pyx, PinnedMemoryResource.from_allocation_handle
warn_if_forked_after_cuda_init()
HANDLE_RETURN(cydriver.cuInit(0))
mark_cuda_initialized()

Because the check is a cdef call plus a bool test, the fast path costs nothing measurable.

Validation

A prototype of this sketch was built and run on an RTX 6000 Ada (driver 595.99.02, CUDA 13.4, cuda-bindings 13.4.4), under Python 3.12.3 and 3.14.6. Each scenario runs as a separate process with PYTHONWARNINGS=always, and the harness counts the warnings:

Also checked:

• The warning names the user's line (for example dev = Device(0)) and appears before the CUDA error.
• With warnings.filterwarnings("error", category=CUDAWarning), the child's Device(0) raises. A second call does not warn again, and the parent is unaffected.
• Device() stayed at about 128 ns and Device(0) at about 103-106 ns per call, which is within noise.
• No regressions in tests/test_multiprocessing_warning.py, tests/test_device.py, tests/memory_ipc or tests/test_memory.py (615 passed), with RuntimeWarning escalated to an error.

Warning category

Emit cuda.core.CUDAWarning, not UserWarning, for this diagnostic. CUDAWarning is cuda.core's documented channel for CUDA-related conditions it reports rather than raises, and it lets users escalate every such condition uniformly with warnings.filterwarnings("error", category=cuda.core.CUDAWarning). A hard, unconditional error is the wrong default here: fork after init is "undefined behavior," not "guaranteed failure" (some driver versions tolerate it), and the existing check_multiprocessing_start_method only reads the global start method at pickle time, which does not prove a dangerous fork will follow. Keeping it a warning that users can promote to an error preserves that choice.

The existing check_multiprocessing_start_method currently emits UserWarning; it should emit CUDAWarning too, so the fork diagnostics share one escalation path. This is a small, separable change from the child-side hook above.

Limits

• Inherited objects bypass the check. The warning fires only at the entry points that call cuInit. A child that only uses objects inherited from the parent, such as a Stream, never reaches them. In the validation run such a child got neither a warning nor an error. Covering this would mean checking the flag on more entry points.
• Other initializers are invisible. CUDA initialized by PyTorch or by direct cuda.bindings calls never reaches mark_cuda_initialized.
• Free-threaded builds may warn twice. Two threads reaching their first CUDA call at the same time can both see the flag set. That is harmless, and an atomic flag (as in Use atomic flag for buffer memory-attribute initialization #2216) would rule it out.

Cheaper alternative

Document that the main module must not initialize CUDA outside the __main__ guard. For forkserver (and fork) this is a correctness requirement: top-level CUDA init contaminates the forkserver, and every child inherits broken state. For spawn it is not undefined behavior -- each fresh child merely re-initializes CUDA on import -- but the same guard is still the right habit, because forkserver is the Linux default since Python 3.14, so an unguarded top-level init that is only wasteful under spawn becomes unsafe the moment the default start method applies. Python's multiprocessing docs describe the guard, but cuda.core's docs never mention it or forkserver.

Related

• cuda.core: test VMM IPC under the forkserver start method #3111 and test(cuda.core): cover the forkserver start method for VMM IPC buffers #3124 add the forkserver tests for VMM IPC buffers. They came out of the review in cuda.core: share VirtualMemoryResource buffers between processes #3028 (comment), where this question first came up.
• cuda.core: transfer VMM IPC allocations one at a time #3110 would remove the Process special case in VirtualMemoryBuffer.__reduce__.