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, and nothing tells the user why.

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.

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.

Idea

Warn in the forked child, when it was forked from a process in which cuda.core initialized CUDA. Sketch, untested:

# cuda/core/_utils/cuda_utils.pyx
_cuda_initialized_pid = None  # set where cuda.core calls cuInit


def mark_cuda_initialized():
    global _cuda_initialized_pid
    _cuda_initialized_pid = os.getpid()


def _warn_if_forked_after_cuda_init():
    global _cuda_initialized_pid
    if _cuda_initialized_pid is None or _cuda_initialized_pid == os.getpid():
        return
    _cuda_initialized_pid = None  # warn once per child
    warnings.warn("this process was forked from one that had initialized CUDA ...", UserWarning)


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

mark_cuda_initialized() would be called after cuInit in Device_ensure_cuda_initialized (_device.pyx) and in _memory/_pinned_memory_resource.pyx.

Expected behavior:

• fork parent: each child warns at startup, whether or not IPC is involved.
• forkserver with CUDA calls at the top level of the main module: each child warns.
• Normal spawn and forkserver programs: no warning, because no process that forks has initialized CUDA.

Limits

• Benign forks warn too. A process that initialized CUDA and forks a child that never uses CUDA would also warn. Deferring the warning to the child's first CUDA call avoids that, with more plumbing in the entry points.
• Other initializers are invisible. CUDA initialized by PyTorch or by direct cuda.bindings calls never reaches mark_cuda_initialized.
• A global import-time hook. It adds one register_at_fork hook to every process that imports cuda.core. That is a maintainer decision.

Cheaper alternative

Document, for both spawn and forkserver, that the main module must not initialize CUDA outside the __main__ guard. 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__.