From cde063848d0a45309b8f18f10e004e5e9af3d97b Mon Sep 17 00:00:00 2001 From: Ciaran Ryan-Anderson Date: Thu, 6 Aug 2026 15:15:50 -0600 Subject: [PATCH 1/6] Load cuQuantum libraries with RTLD_LOCAL and explain fork-after-CUDA-init failures --- crates/pecos-cuquantum-sys/src/loader.rs | 35 +++--- crates/pecos-cuquantum/src/error.rs | 54 ++++++++- docs/user-guide/cuda-setup.md | 15 +++ .../src/pecos/simulators/custatevec/state.py | 29 +++++ .../tests/pecos/unit/test_custatevec_state.py | 107 ++++++++++++++++++ 5 files changed, 221 insertions(+), 19 deletions(-) create mode 100644 python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py diff --git a/crates/pecos-cuquantum-sys/src/loader.rs b/crates/pecos-cuquantum-sys/src/loader.rs index 3d0574d51..3d181ef54 100644 --- a/crates/pecos-cuquantum-sys/src/loader.rs +++ b/crates/pecos-cuquantum-sys/src/loader.rs @@ -50,8 +50,9 @@ macro_rules! load_sym { #[allow(non_snake_case)] pub struct CuQuantumBackend { // Keep libraries alive. - // Rust drops fields in declaration order, so dependents must come before - // their dependencies. cuda_rt is last because everything depends on it. + // Rust drops fields in declaration order, so cuDensityMat and cuTensorNet must + // come before their cuTensor dependency. cuda_rt remains last so its direct + // CUDA runtime symbols stay loaded until the other library handles are released. _cudensitymat: Library, _cutensornet: Library, _cutensor: Option, @@ -269,20 +270,23 @@ fn cutensor_search_paths() -> Vec { paths } -/// Load a shared library with RTLD_GLOBAL so its symbols are visible to subsequent loads. -/// This is necessary because cuQuantum libs have transitive dependencies (e.g. -/// libcutensornet depends on libcutensor) that need to resolve via the global symbol table. +/// Load a shared library without adding its symbols to the process-global lookup scope. +/// +/// cuTensor is preloaded by full path before libraries with a `NEEDED` entry for its SONAME. ELF +/// loaders reuse that already-loaded object when resolving the entry, so global symbol visibility +/// is unnecessary. CUDA libraries export unmangled C symbols; keeping them local prevents this +/// loader from polluting symbol resolution for other CUDA-using libraries in the same process. #[cfg(unix)] -fn load_global>(path: P) -> Result { +fn load_local>(path: P) -> Result { // RTLD_NOW: resolve all symbols immediately - // RTLD_GLOBAL: make symbols available for subsequent dlopen calls - let flags = libc::RTLD_NOW | libc::RTLD_GLOBAL; + // RTLD_LOCAL: keep exported symbols out of the process-global lookup scope + let flags = libc::RTLD_NOW | libc::RTLD_LOCAL; let lib = unsafe { UnixLibrary::open(Some(path.as_ref()), flags) }?; Ok(lib.into()) } #[cfg(not(unix))] -fn load_global>(path: P) -> Result { +fn load_local>(path: P) -> Result { unsafe { Library::new(path.as_ref()) } } @@ -296,7 +300,7 @@ fn try_load_lib(names: &[&str], search_dirs: &[PathBuf]) -> LoadResult for dir in search_dirs { let path = dir.join(name); log::debug!("Trying to load {name} from: {}", path.display()); - match load_global(&path) { + match load_local(&path) { Ok(lib) => { log::info!("Loaded {name} from: {}", path.display()); return Ok(lib); @@ -306,7 +310,7 @@ fn try_load_lib(names: &[&str], search_dirs: &[PathBuf]) -> LoadResult } // Fall back to bare name (system linker search) log::debug!("Trying system path for {name}"); - if let Ok(lib) = load_global(*name) { + if let Ok(lib) = load_local(*name) { log::info!("Loaded {name} from system path"); return Ok(lib); } @@ -328,16 +332,17 @@ fn load_all() -> Result { let cq_paths = cuquantum_search_paths(); let ct_paths = cutensor_search_paths(); - // Load CUDA runtime first (transitive dependency for everything else). + // Load CUDA runtime for the API symbols used directly by this backend. As dependency + // preloading, this is only a best-effort compatibility measure for variants that dynamically + // depend on libcudart; the shipped cuQuantum 25.11 libraries link the runtime statically. // Try versioned soname first -- unversioned symlink may not exist on runtime-only installs. let cuda_rt = try_load_lib( &["libcudart.so.13", "libcudart.so.12", "libcudart.so"], &cuda_paths, )?; - // Load cuTensor before cuTensorNet (transitive dependency). - // The handle must stay alive in CuQuantumBackend so dlclose doesn't - // unload the library while cuTensorNet still needs its symbols. + // Load cuTensor before cuTensorNet and cuDensityMat, whose dynamic `NEEDED` entries resolve + // its SONAME to this object. The handle must stay alive so it is not unloaded first. let cutensor = try_load_lib(&["libcutensor.so.2", "libcutensor.so"], &ct_paths).ok(); let custatevec = try_load_lib(&["libcustatevec.so.1", "libcustatevec.so"], &cq_paths)?; diff --git a/crates/pecos-cuquantum/src/error.rs b/crates/pecos-cuquantum/src/error.rs index 0153316ef..111972c8b 100644 --- a/crates/pecos-cuquantum/src/error.rs +++ b/crates/pecos-cuquantum/src/error.rs @@ -52,7 +52,11 @@ pub enum CuQuantumError { /// cuStateVec-specific error #[derive(Error, Debug, Clone, Copy, PartialEq, Eq)] pub enum StateVecError { - #[error("Not initialized")] + #[error( + "Not initialized. CUDA could not initialize in this process; a common cause is running \ + in a forked child of a process that already initialized CUDA (CUDA contexts do not \ + survive fork) — use the multiprocessing \"spawn\" start method." + )] NotInitialized, #[error("Allocation failed")] @@ -142,7 +146,11 @@ impl From for CuQuantumError { /// cuStabilizer-specific error #[derive(Error, Debug, Clone, Copy, PartialEq, Eq)] pub enum StabilizerError { - #[error("Not initialized")] + #[error( + "Not initialized. CUDA could not initialize in this process; a common cause is running \ + in a forked child of a process that already initialized CUDA (CUDA contexts do not \ + survive fork) — use the multiprocessing \"spawn\" start method." + )] NotInitialized, #[error("Allocation failed")] @@ -201,7 +209,11 @@ impl From for CuQuantumError { /// cuTensorNet-specific error #[derive(Error, Debug, Clone, Copy, PartialEq, Eq)] pub enum TensorNetError { - #[error("Not initialized")] + #[error( + "Not initialized. CUDA could not initialize in this process; a common cause is running \ + in a forked child of a process that already initialized CUDA (CUDA contexts do not \ + survive fork) — use the multiprocessing \"spawn\" start method." + )] NotInitialized, #[error("Allocation failed")] @@ -289,7 +301,11 @@ impl From for CuQuantumError { /// cuDensityMat-specific error #[derive(Error, Debug, Clone, Copy, PartialEq, Eq)] pub enum DensityMatError { - #[error("Not initialized")] + #[error( + "Not initialized. CUDA could not initialize in this process; a common cause is running \ + in a forked child of a process that already initialized CUDA (CUDA contexts do not \ + survive fork) — use the multiprocessing \"spawn\" start method." + )] NotInitialized, #[error("Allocation failed")] @@ -461,6 +477,36 @@ mod tests { assert!(msg.contains("Invalid value")); } + #[test] + fn test_cuquantum_not_initialized_display_has_fork_hint_only_for_that_status() { + let not_initialized_messages = [ + CuQuantumError::from(custatevecStatus_t::CUSTATEVEC_STATUS_NOT_INITIALIZED).to_string(), + CuQuantumError::from(custabilizerStatus_t::CUSTABILIZER_STATUS_NOT_INITIALIZED) + .to_string(), + CuQuantumError::from(cutensornetStatus_t::CUTENSORNET_STATUS_NOT_INITIALIZED) + .to_string(), + CuQuantumError::from(cudensitymatStatus_t::CUDENSITYMAT_STATUS_NOT_INITIALIZED) + .to_string(), + ]; + for message in not_initialized_messages { + assert!(message.contains("forked child")); + assert!(message.contains("\"spawn\" start method")); + } + + let other_messages = [ + CuQuantumError::from(custatevecStatus_t::CUSTATEVEC_STATUS_INVALID_VALUE).to_string(), + CuQuantumError::from(custabilizerStatus_t::CUSTABILIZER_STATUS_INVALID_VALUE) + .to_string(), + CuQuantumError::from(cutensornetStatus_t::CUTENSORNET_STATUS_INVALID_VALUE).to_string(), + CuQuantumError::from(cudensitymatStatus_t::CUDENSITYMAT_STATUS_INVALID_VALUE) + .to_string(), + ]; + for message in other_messages { + assert!(!message.contains("forked child")); + assert!(!message.contains("\"spawn\" start method")); + } + } + #[test] fn test_stabilizer_status_conversion() { let err = StabilizerError::from(custabilizerStatus_t::CUSTABILIZER_STATUS_INVALID_VALUE); diff --git a/docs/user-guide/cuda-setup.md b/docs/user-guide/cuda-setup.md index b326aa626..de937573f 100644 --- a/docs/user-guide/cuda-setup.md +++ b/docs/user-guide/cuda-setup.md @@ -27,6 +27,21 @@ Both approaches require: - CUDA Toolkit (system-level installation) - cuQuantum SDK (for Rust bindings) or Python packages (for Python bindings) +## CUDA Simulators and Multiprocessing + +When a parent process has initialized CUDA, its children must use the multiprocessing +`spawn` start method rather than `fork`. A forked child cannot use the parent's CUDA +context because CUDA contexts do not survive fork. Depending on the backend, this +fails with `cudaErrorInitializationError`, `CUDA_ERROR_NOT_INITIALIZED`, or a +cuQuantum `Not initialized` error from a cuStateVec, cuStabilizer, cuTensorNet, or +cuDensityMat handle. + +PECOS's built-in multiprocessing engine already uses `spawn`. Host programs that +embed PECOS in their own process pools must also select `spawn` if the parent may +initialize CUDA before creating workers. The Rust cuQuantum bindings load their CUDA +libraries with local symbol visibility so they do not disturb other CUDA-using +libraries in the same process. + ## System Requirements ### Hardware Requirements diff --git a/python/quantum-pecos/src/pecos/simulators/custatevec/state.py b/python/quantum-pecos/src/pecos/simulators/custatevec/state.py index 47306ae3b..dae50aaa0 100644 --- a/python/quantum-pecos/src/pecos/simulators/custatevec/state.py +++ b/python/quantum-pecos/src/pecos/simulators/custatevec/state.py @@ -29,6 +29,14 @@ ) from pecos.simulators.sim_class_types import StateVector +# cudaErrorInitializationError from the CUDA runtime API. CuPy exposes only the +# numeric CUDA status through CUDARuntimeError.status, not the named enum constant. +_CUDA_ERROR_INITIALIZATION = 3 + +# CUDA_ERROR_NOT_INITIALIZED from the CUDA driver API. CuPy exposes the numeric +# CUDA status through CUDADriverError.status. +_CUDA_DRIVER_ERROR_NOT_INITIALIZED = 3 + if TYPE_CHECKING: import sys @@ -66,6 +74,26 @@ def __init__(self, num_qubits: int, _seed: int | None = None) -> None: self.bindings = bindings.gate_dict self.num_qubits = num_qubits + try: + self._initialize_cuda() + except (cp.cuda.runtime.CUDARuntimeError, cp.cuda.driver.CUDADriverError) as exc: + runtime_initialization_error = ( + isinstance(exc, cp.cuda.runtime.CUDARuntimeError) and exc.status == _CUDA_ERROR_INITIALIZATION + ) + driver_not_initialized_error = ( + isinstance(exc, cp.cuda.driver.CUDADriverError) and exc.status == _CUDA_DRIVER_ERROR_NOT_INITIALIZED + ) + if runtime_initialization_error or driver_not_initialized_error: + msg = ( + "CUDA could not initialize in this process; a common cause is running in a forked child of a " + "process that already initialized CUDA (CUDA contexts do not survive fork) — use the " + 'multiprocessing "spawn" start method.' + ) + raise RuntimeError(msg) from exc + raise + + def _initialize_cuda(self) -> None: + """Allocate the state vector and initialize its CUDA resources.""" # Set data type as double precision complex numbers self.cp_type = cp.complex128 self.cuda_type = cudaDataType.CUDA_C_64F # == cp.complex128 @@ -103,6 +131,7 @@ def __init__(self, num_qubits: int, _seed: int | None = None) -> None: ) # CuStateVec handle initialization + # Errors raised directly by cusv.create() stay uncaught so cuQuantum preserves their native type. self.libhandle = cusv.create() self.stream = cp.cuda.Stream() cusv.set_stream(self.libhandle, self.stream.ptr) diff --git a/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py b/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py new file mode 100644 index 000000000..646292643 --- /dev/null +++ b/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py @@ -0,0 +1,107 @@ +# Copyright 2026 The PECOS Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. + +"""Tests for actionable CUDA initialization errors from the legacy CuPy simulator. + +The tests inject stub optional dependencies so the error handling remains covered +whether or not CuPy and cuQuantum are installed in the test environment. +""" + +from __future__ import annotations + +import sys +import types +from unittest.mock import Mock + +import pytest +from pecos.simulators.custatevec import state as custatevec_state + + +class _StubCudaStatusError(RuntimeError): + """CuPy-compatible CUDA error carrying a numeric status.""" + + def __init__(self, status: int) -> None: + super().__init__(status) + self.status = status + + +class _StubRuntimeError(_StubCudaStatusError): + """Stand-in for ``cupy.cuda.runtime.CUDARuntimeError``.""" + + +class _StubDriverError(_StubCudaStatusError): + """Stand-in for ``cupy.cuda.driver.CUDADriverError``.""" + + +def _inject_optional_dependency_stubs(monkeypatch: pytest.MonkeyPatch) -> types.ModuleType: + """Inject the subset of CuPy and cuQuantum metadata used before allocation.""" + cp = types.ModuleType("cupy") + cuda = types.ModuleType("cupy.cuda") + runtime = types.ModuleType("cupy.cuda.runtime") + driver = types.ModuleType("cupy.cuda.driver") + runtime.CUDARuntimeError = _StubRuntimeError + runtime.runtimeGetVersion = Mock(return_value=12000) + driver.CUDADriverError = _StubDriverError + cuda.runtime = runtime + cuda.driver = driver + cuda.Device = Mock() + cp.cuda = cuda + cp.complex128 = object() + cp.zeros = Mock(return_value=[0, 0]) + + monkeypatch.setitem(sys.modules, "cupy", cp) + monkeypatch.setitem(sys.modules, "cupy.cuda", cuda) + monkeypatch.setitem(sys.modules, "cupy.cuda.runtime", runtime) + monkeypatch.setitem(sys.modules, "cupy.cuda.driver", driver) + monkeypatch.setattr(custatevec_state, "cp", cp) + monkeypatch.setattr(custatevec_state, "require_custatevec", Mock()) + monkeypatch.setattr(custatevec_state, "cudaDataType", types.SimpleNamespace(CUDA_C_64F=object())) + monkeypatch.setattr(custatevec_state, "ComputeType", types.SimpleNamespace(COMPUTE_64F=object())) + return cp + + +def test_cuda_initialization_errors_explain_fork_and_preserve_cause(monkeypatch) -> None: + """Translate the runtime and driver initialization statuses from CuPy calls.""" + cp = _inject_optional_dependency_stubs(monkeypatch) + + runtime_error = cp.cuda.runtime.CUDARuntimeError(3) + cp.zeros.side_effect = runtime_error + + with pytest.raises(RuntimeError, match='multiprocessing "spawn" start method') as exc_info: + custatevec_state.CuStateVec(1) + + assert type(exc_info.value) is RuntimeError + assert "forked child" in str(exc_info.value) + assert exc_info.value.__cause__ is runtime_error + + driver_error = cp.cuda.driver.CUDADriverError(3) + cp.zeros.side_effect = None + cp.cuda.Device.side_effect = driver_error + + with pytest.raises(RuntimeError, match='multiprocessing "spawn" start method') as exc_info: + custatevec_state.CuStateVec(1) + + assert type(exc_info.value) is RuntimeError + assert "forked child" in str(exc_info.value) + assert exc_info.value.__cause__ is driver_error + + +def test_other_cuda_runtime_status_propagates_unchanged(monkeypatch) -> None: + """Do not translate CUDARuntimeError statuses unrelated to initialization.""" + cp = _inject_optional_dependency_stubs(monkeypatch) + cuda_error = cp.cuda.runtime.CUDARuntimeError(1) # cudaErrorInvalidValue + cp.zeros.side_effect = cuda_error + + with pytest.raises(cp.cuda.runtime.CUDARuntimeError) as exc_info: + custatevec_state.CuStateVec(1) + + assert exc_info.value is cuda_error From bde3a0b4eb0b64860ec88f7f5714a5ff85efae0a Mon Sep 17 00:00:00 2001 From: Ciaran Ryan-Anderson Date: Thu, 6 Aug 2026 15:36:20 -0600 Subject: [PATCH 2/6] Add fork-poison guard, spawn-discipline test, and probe invariant docs --- crates/pecos-cuquantum-sys/src/loader.rs | 6 ++ docs/user-guide/cuda-setup.md | 9 +- .../src/pecos/simulators/_cuda_fork_guard.py | 75 ++++++++++++++++ .../pecos/simulators/cuda_stabilizer/state.py | 3 + .../pecos/simulators/cuda_statevec/state.py | 3 + .../src/pecos/simulators/custatevec/state.py | 14 +-- .../tests/pecos/unit/test_cuda_fork_guard.py | 85 +++++++++++++++++++ .../tests/pecos/unit/test_custatevec_state.py | 9 ++ .../test_multiprocessing_spawn_discipline.py | 43 ++++++++++ 9 files changed, 240 insertions(+), 7 deletions(-) create mode 100644 python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py create mode 100644 python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py create mode 100644 python/quantum-pecos/tests/pecos/unit/test_multiprocessing_spawn_discipline.py diff --git a/crates/pecos-cuquantum-sys/src/loader.rs b/crates/pecos-cuquantum-sys/src/loader.rs index 3d181ef54..cc0a8e63a 100644 --- a/crates/pecos-cuquantum-sys/src/loader.rs +++ b/crates/pecos-cuquantum-sys/src/loader.rs @@ -179,11 +179,17 @@ unsafe impl Sync for CuQuantumBackend {} static BACKEND: OnceLock> = OnceLock::new(); /// Load cuQuantum libraries. Thread-safe, loads only once. +/// +/// Availability probing must only open libraries; it must never create CUDA handles or contexts, +/// because initializing the driver in a parent process poisons every later forked child. pub fn try_load() -> Result<&'static CuQuantumBackend, &'static CuQuantumLoadError> { BACKEND.get_or_init(load_all).as_ref() } /// Check if cuQuantum is available at runtime. +/// +/// This probe only opens libraries and must not create CUDA handles, contexts, or otherwise +/// initialize the driver before a caller may fork worker processes. pub fn is_available() -> bool { try_load().is_ok() } diff --git a/docs/user-guide/cuda-setup.md b/docs/user-guide/cuda-setup.md index de937573f..d02215915 100644 --- a/docs/user-guide/cuda-setup.md +++ b/docs/user-guide/cuda-setup.md @@ -40,7 +40,14 @@ PECOS's built-in multiprocessing engine already uses `spawn`. Host programs that embed PECOS in their own process pools must also select `spawn` if the parent may initialize CUDA before creating workers. The Rust cuQuantum bindings load their CUDA libraries with local symbol visibility so they do not disturb other CUDA-using -libraries in the same process. +libraries in the same process. Availability checks only load those libraries; they +do not create CUDA handles or contexts or initialize the CUDA driver. + +PECOS records when its CUDA simulator wrappers begin CUDA initialization. If one of +those simulators is later constructed in a child forked from that process, PECOS +fails fast with the spawn guidance above. PECOS cannot detect CUDA initialization by +another library in the parent; those cases still surface the guided error from the +CUDA or cuQuantum layer. ## System Requirements diff --git a/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py b/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py new file mode 100644 index 000000000..aac5222ca --- /dev/null +++ b/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py @@ -0,0 +1,75 @@ +# Copyright 2026 The PECOS Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. + +"""Detect processes forked after PECOS initialized CUDA. + +The Python CUDA simulator wrappers mark their first CUDA call and use this guard +to fail before touching an inherited, unusable CUDA context. Direct users of the +raw ``pecos_rslib_cuda`` extension deliberately bypass this guard; its guided +``Not initialized`` errors remain their backstop. + +Registering fork hooks does not load CUDA libraries, create handles, or initialize +the CUDA driver. +""" + +from __future__ import annotations + +import os +import warnings + +__all__ = ["CUDA_FORK_ERROR_MESSAGE", "check_fork_poison", "mark_cuda_initialized"] + +CUDA_FORK_ERROR_MESSAGE = ( + "CUDA could not initialize in this process; a common cause is running in a forked child of a process that " + 'already initialized CUDA (CUDA contexts do not survive fork) — use the multiprocessing "spawn" start method.' +) + +_cuda_initialized = False +_forked_after_cuda_init = False +_fork_warning_emitted = False + + +def _warn_before_fork() -> None: + """Warn once when this process forks after a PECOS CUDA call.""" + global _fork_warning_emitted + if _cuda_initialized and not _fork_warning_emitted: + _fork_warning_emitted = True + warnings.warn( + "This process is forking after CUDA initialization, so CUDA will be unusable in the forked child. " + + CUDA_FORK_ERROR_MESSAGE, + RuntimeWarning, + stacklevel=2, + ) + + +def _mark_forked_child() -> None: + """Record that the child inherited evidence of parent CUDA initialization.""" + global _forked_after_cuda_init + if _cuda_initialized: + _forked_after_cuda_init = True + + +def mark_cuda_initialized() -> None: + """Record that this process is about to make a real CUDA call.""" + global _cuda_initialized + _cuda_initialized = True + + +def check_fork_poison() -> None: + """Fail before CUDA use in a child forked after parent CUDA initialization.""" + if _forked_after_cuda_init: + raise RuntimeError(CUDA_FORK_ERROR_MESSAGE) + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(before=_warn_before_fork) + os.register_at_fork(after_in_child=_mark_forked_child) diff --git a/python/quantum-pecos/src/pecos/simulators/cuda_stabilizer/state.py b/python/quantum-pecos/src/pecos/simulators/cuda_stabilizer/state.py index 0b68fd7fe..29031370f 100644 --- a/python/quantum-pecos/src/pecos/simulators/cuda_stabilizer/state.py +++ b/python/quantum-pecos/src/pecos/simulators/cuda_stabilizer/state.py @@ -21,6 +21,7 @@ from pecos_rslib_cuda import CuStabilizer as CuStabilizerRs +from pecos.simulators._cuda_fork_guard import check_fork_poison, mark_cuda_initialized from pecos.simulators.cuda_stabilizer import bindings from pecos.simulators.sim_class_types import Stabilizer @@ -68,6 +69,7 @@ def __init__(self, num_qubits: int, seed: int | None = None) -> None: num_qubits: Number of qubits to simulate. seed: Optional random seed for reproducibility. """ + check_fork_poison() if not isinstance(num_qubits, int): msg = "``num_qubits`` should be of type ``int``." raise TypeError(msg) @@ -78,6 +80,7 @@ def __init__(self, num_qubits: int, seed: int | None = None) -> None: self.num_qubits = num_qubits # Create the Rust backend + mark_cuda_initialized() if seed is not None: self.backend = CuStabilizerRs.with_seed(num_qubits, seed) else: diff --git a/python/quantum-pecos/src/pecos/simulators/cuda_statevec/state.py b/python/quantum-pecos/src/pecos/simulators/cuda_statevec/state.py index a856a80d1..1409b0517 100644 --- a/python/quantum-pecos/src/pecos/simulators/cuda_statevec/state.py +++ b/python/quantum-pecos/src/pecos/simulators/cuda_statevec/state.py @@ -21,6 +21,7 @@ from pecos_rslib_cuda import CuStateVec as CuStateVecRs +from pecos.simulators._cuda_fork_guard import check_fork_poison, mark_cuda_initialized from pecos.simulators.cuda_statevec import bindings from pecos.simulators.sim_class_types import StateVector @@ -63,6 +64,7 @@ def __init__(self, num_qubits: int, seed: int | None = None) -> None: num_qubits: Number of qubits to simulate. seed: Optional random seed for reproducibility. """ + check_fork_poison() if not isinstance(num_qubits, int): msg = "``num_qubits`` should be of type ``int``." raise TypeError(msg) @@ -73,6 +75,7 @@ def __init__(self, num_qubits: int, seed: int | None = None) -> None: self.num_qubits = num_qubits # Create the Rust backend + mark_cuda_initialized() if seed is not None: self.backend = CuStateVecRs.with_seed(num_qubits, seed) else: diff --git a/python/quantum-pecos/src/pecos/simulators/custatevec/state.py b/python/quantum-pecos/src/pecos/simulators/custatevec/state.py index dae50aaa0..c80315e5d 100644 --- a/python/quantum-pecos/src/pecos/simulators/custatevec/state.py +++ b/python/quantum-pecos/src/pecos/simulators/custatevec/state.py @@ -19,6 +19,11 @@ from typing import TYPE_CHECKING +from pecos.simulators._cuda_fork_guard import ( + CUDA_FORK_ERROR_MESSAGE, + check_fork_poison, + mark_cuda_initialized, +) from pecos.simulators.custatevec import bindings from pecos.simulators.custatevec._cuquantum_compat import ( ComputeType, @@ -61,6 +66,7 @@ def __init__(self, num_qubits: int, _seed: int | None = None) -> None: num_qubits (int): Number of qubits being represented. _seed (int): Seed for randomness (kept for API compatibility, not used in GPU-based simulator). """ + check_fork_poison() # Fail loudly (and only here, at construction) if CuPy / bindings-era cuQuantum # is unavailable -- importing pecos must stay non-fatal without CUDA installed. require_custatevec() @@ -84,12 +90,7 @@ def __init__(self, num_qubits: int, _seed: int | None = None) -> None: isinstance(exc, cp.cuda.driver.CUDADriverError) and exc.status == _CUDA_DRIVER_ERROR_NOT_INITIALIZED ) if runtime_initialization_error or driver_not_initialized_error: - msg = ( - "CUDA could not initialize in this process; a common cause is running in a forked child of a " - "process that already initialized CUDA (CUDA contexts do not survive fork) — use the " - 'multiprocessing "spawn" start method.' - ) - raise RuntimeError(msg) from exc + raise RuntimeError(CUDA_FORK_ERROR_MESSAGE) from exc raise def _initialize_cuda(self) -> None: @@ -101,6 +102,7 @@ def _initialize_cuda(self) -> None: # Allocate the statevector in GPU and initialize it to |0> self.cupy_vector = None + mark_cuda_initialized() self.reset() #################################################### diff --git a/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py b/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py new file mode 100644 index 000000000..0e24a3ce5 --- /dev/null +++ b/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py @@ -0,0 +1,85 @@ +# Copyright 2026 The PECOS Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. + +"""GPU-free tests for CUDA initialization state inherited across ``fork``.""" + +from __future__ import annotations + +import os +import warnings + +import pytest +from pecos.simulators import _cuda_fork_guard as guard + +pytestmark = pytest.mark.skipif(not hasattr(os, "fork"), reason="os.fork is unavailable on this platform") + + +@pytest.fixture(autouse=True) +def _reset_guard_state(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep the guard's process-local state isolated between tests.""" + monkeypatch.setattr(guard, "_cuda_initialized", False) + monkeypatch.setattr(guard, "_forked_after_cuda_init", False) + monkeypatch.setattr(guard, "_fork_warning_emitted", False) + + +def _wait_for_child(pid: int) -> int: + """Wait for a forked child and return its conventional exit code.""" + _, status = os.waitpid(pid, 0) + return os.waitstatus_to_exitcode(status) + + +def test_marked_parent_poisons_forked_child() -> None: + """A child inherits the poison marker when its parent marked CUDA use.""" + guard.mark_cuda_initialized() + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + pid = os.fork() + + if pid == 0: + try: + guard.check_fork_poison() + except RuntimeError as exc: + os._exit(0 if str(exc) == guard.CUDA_FORK_ERROR_MESSAGE else 2) + os._exit(1) + + assert _wait_for_child(pid) == 0 + + +def test_unmarked_parent_does_not_poison_forked_child() -> None: + """Forking before any marked CUDA call leaves the child usable.""" + pid = os.fork() + if pid == 0: + try: + guard.check_fork_poison() + except RuntimeError: + os._exit(1) + os._exit(0) + + assert _wait_for_child(pid) == 0 + + +def test_marked_parent_warns_only_once_before_fork() -> None: + """Repeated forks after marked CUDA use emit one parent-side warning.""" + guard.mark_cuda_initialized() + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", RuntimeWarning) + for _ in range(2): + pid = os.fork() + if pid == 0: + os._exit(0) + assert _wait_for_child(pid) == 0 + + fork_warnings = [warning for warning in caught if warning.category is RuntimeWarning] + assert len(fork_warnings) == 1 + assert "forking after CUDA initialization" in str(fork_warnings[0].message) + assert guard.CUDA_FORK_ERROR_MESSAGE in str(fork_warnings[0].message) diff --git a/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py b/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py index 646292643..a564b89aa 100644 --- a/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py +++ b/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py @@ -23,6 +23,7 @@ from unittest.mock import Mock import pytest +from pecos.simulators import _cuda_fork_guard as guard from pecos.simulators.custatevec import state as custatevec_state @@ -42,6 +43,14 @@ class _StubDriverError(_StubCudaStatusError): """Stand-in for ``cupy.cuda.driver.CUDADriverError``.""" +@pytest.fixture(autouse=True) +def _reset_guard_state(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep synthetic failed initialization from marking later tests.""" + monkeypatch.setattr(guard, "_cuda_initialized", False) + monkeypatch.setattr(guard, "_forked_after_cuda_init", False) + monkeypatch.setattr(guard, "_fork_warning_emitted", False) + + def _inject_optional_dependency_stubs(monkeypatch: pytest.MonkeyPatch) -> types.ModuleType: """Inject the subset of CuPy and cuQuantum metadata used before allocation.""" cp = types.ModuleType("cupy") diff --git a/python/quantum-pecos/tests/pecos/unit/test_multiprocessing_spawn_discipline.py b/python/quantum-pecos/tests/pecos/unit/test_multiprocessing_spawn_discipline.py new file mode 100644 index 000000000..561449101 --- /dev/null +++ b/python/quantum-pecos/tests/pecos/unit/test_multiprocessing_spawn_discipline.py @@ -0,0 +1,43 @@ +# Copyright 2026 The PECOS Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. + +"""Prevent platform-default or explicit-fork multiprocessing in PECOS source.""" + +from __future__ import annotations + +import re +from pathlib import Path + +_SOURCE_ROOT = Path(__file__).resolve().parents[3] / "src" / "pecos" +_ALLOWLIST: frozenset[Path] = frozenset() +_FORK_HAZARD_PATTERNS = { + "get_context(fork)": re.compile(r"\bget_context\s*\(\s*(['\"])fork\1\s*\)"), + "set_start_method(fork)": re.compile(r"\bset_start_method\s*\(\s*(['\"])fork\1\s*\)"), + "multiprocessing.Pool": re.compile(r"\bmultiprocessing\.Pool\s*\("), + "mp.Pool": re.compile(r"\bmp\.Pool\s*\("), + "bare Pool": re.compile(r"(? None: + """Require explicit non-fork multiprocessing throughout PECOS Python source.""" + violations = [] + for source_path in sorted(_SOURCE_ROOT.rglob("*.py")): + relative_path = source_path.relative_to(_SOURCE_ROOT) + if relative_path in _ALLOWLIST: + continue + for line_number, line in enumerate(source_path.read_text(encoding="utf-8").splitlines(), start=1): + for pattern_name, pattern in _FORK_HAZARD_PATTERNS.items(): + if pattern.search(line): + violations.append(f"{relative_path}:{line_number}: {pattern_name}: {line.strip()}") + + assert not violations, "Fork-unsafe multiprocessing usage found:\n" + "\n".join(violations) From a00cf30cb1026d105f309dab015e8a767fd6b8fe Mon Sep 17 00:00:00 2001 From: Ciaran Ryan-Anderson Date: Thu, 6 Aug 2026 18:51:08 -0600 Subject: [PATCH 3/6] Trigger CI From f34817ac8641a53abc09427fabff5461015248ab Mon Sep 17 00:00:00 2001 From: Ciaran Ryan-Anderson Date: Thu, 6 Aug 2026 19:14:52 -0600 Subject: [PATCH 4/6] Use a state dict for the fork-guard latch to satisfy the lint gate --- .../src/pecos/simulators/_cuda_fork_guard.py | 23 +++++++++---------- .../tests/pecos/unit/test_cuda_fork_guard.py | 8 ++++--- .../tests/pecos/unit/test_custatevec_state.py | 8 ++++--- 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py b/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py index aac5222ca..9ea8d46c2 100644 --- a/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py +++ b/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py @@ -33,16 +33,17 @@ 'already initialized CUDA (CUDA contexts do not survive fork) — use the multiprocessing "spawn" start method.' ) -_cuda_initialized = False -_forked_after_cuda_init = False -_fork_warning_emitted = False +_state = { + "cuda_initialized": False, + "forked_after_cuda_init": False, + "fork_warning_emitted": False, +} def _warn_before_fork() -> None: """Warn once when this process forks after a PECOS CUDA call.""" - global _fork_warning_emitted - if _cuda_initialized and not _fork_warning_emitted: - _fork_warning_emitted = True + if _state["cuda_initialized"] and not _state["fork_warning_emitted"]: + _state["fork_warning_emitted"] = True warnings.warn( "This process is forking after CUDA initialization, so CUDA will be unusable in the forked child. " + CUDA_FORK_ERROR_MESSAGE, @@ -53,20 +54,18 @@ def _warn_before_fork() -> None: def _mark_forked_child() -> None: """Record that the child inherited evidence of parent CUDA initialization.""" - global _forked_after_cuda_init - if _cuda_initialized: - _forked_after_cuda_init = True + if _state["cuda_initialized"]: + _state["forked_after_cuda_init"] = True def mark_cuda_initialized() -> None: """Record that this process is about to make a real CUDA call.""" - global _cuda_initialized - _cuda_initialized = True + _state["cuda_initialized"] = True def check_fork_poison() -> None: """Fail before CUDA use in a child forked after parent CUDA initialization.""" - if _forked_after_cuda_init: + if _state["forked_after_cuda_init"]: raise RuntimeError(CUDA_FORK_ERROR_MESSAGE) diff --git a/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py b/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py index 0e24a3ce5..4f8ec1d00 100644 --- a/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py +++ b/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py @@ -26,9 +26,11 @@ @pytest.fixture(autouse=True) def _reset_guard_state(monkeypatch: pytest.MonkeyPatch) -> None: """Keep the guard's process-local state isolated between tests.""" - monkeypatch.setattr(guard, "_cuda_initialized", False) - monkeypatch.setattr(guard, "_forked_after_cuda_init", False) - monkeypatch.setattr(guard, "_fork_warning_emitted", False) + monkeypatch.setattr( + guard, + "_state", + {"cuda_initialized": False, "forked_after_cuda_init": False, "fork_warning_emitted": False}, + ) def _wait_for_child(pid: int) -> int: diff --git a/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py b/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py index a564b89aa..015103195 100644 --- a/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py +++ b/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py @@ -46,9 +46,11 @@ class _StubDriverError(_StubCudaStatusError): @pytest.fixture(autouse=True) def _reset_guard_state(monkeypatch: pytest.MonkeyPatch) -> None: """Keep synthetic failed initialization from marking later tests.""" - monkeypatch.setattr(guard, "_cuda_initialized", False) - monkeypatch.setattr(guard, "_forked_after_cuda_init", False) - monkeypatch.setattr(guard, "_fork_warning_emitted", False) + monkeypatch.setattr( + guard, + "_state", + {"cuda_initialized": False, "forked_after_cuda_init": False, "fork_warning_emitted": False}, + ) def _inject_optional_dependency_stubs(monkeypatch: pytest.MonkeyPatch) -> types.ModuleType: From 580c9c31b550ab5377b77d8d49fc098ceb68c524 Mon Sep 17 00:00:00 2001 From: Ciaran Ryan-Anderson Date: Thu, 6 Aug 2026 21:15:12 -0600 Subject: [PATCH 5/6] Fold delta-review findings: guard-wiring tests, parent invariant, MPS guard, AST spawn scan, reset checks --- docs/user-guide/cuda-setup.md | 3 +- .../src/pecos/simulators/_cuda_fork_guard.py | 6 +- .../pecos/simulators/cuda_stabilizer/state.py | 1 + .../pecos/simulators/cuda_statevec/state.py | 1 + .../src/pecos/simulators/mps_pytket/state.py | 3 + .../tests/pecos/unit/test_cuda_fork_guard.py | 20 ++- .../tests/pecos/unit/test_custatevec_state.py | 57 ++++++- .../pecos/unit/test_mps_cuda_fork_guard.py | 114 ++++++++++++++ .../test_multiprocessing_spawn_discipline.py | 149 ++++++++++++++++-- 9 files changed, 334 insertions(+), 20 deletions(-) create mode 100644 python/quantum-pecos/tests/pecos/unit/test_mps_cuda_fork_guard.py diff --git a/docs/user-guide/cuda-setup.md b/docs/user-guide/cuda-setup.md index d02215915..6fb6f36d3 100644 --- a/docs/user-guide/cuda-setup.md +++ b/docs/user-guide/cuda-setup.md @@ -47,7 +47,8 @@ PECOS records when its CUDA simulator wrappers begin CUDA initialization. If one those simulators is later constructed in a child forked from that process, PECOS fails fast with the spawn guidance above. PECOS cannot detect CUDA initialization by another library in the parent; those cases still surface the guided error from the -CUDA or cuQuantum layer. +CUDA or cuQuantum layer. When a marked parent forks, PECOS also emits a one-time +`RuntimeWarning` with the same guidance before the fork. ## System Requirements diff --git a/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py b/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py index 9ea8d46c2..b6155feef 100644 --- a/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py +++ b/python/quantum-pecos/src/pecos/simulators/_cuda_fork_guard.py @@ -59,7 +59,11 @@ def _mark_forked_child() -> None: def mark_cuda_initialized() -> None: - """Record that this process is about to make a real CUDA call.""" + """Record that this process is about to make a real CUDA call. + + The mark deliberately latches if that call fails because partial driver + initialization still makes later forked children unsafe. + """ _state["cuda_initialized"] = True diff --git a/python/quantum-pecos/src/pecos/simulators/cuda_stabilizer/state.py b/python/quantum-pecos/src/pecos/simulators/cuda_stabilizer/state.py index 29031370f..98ed8f6ee 100644 --- a/python/quantum-pecos/src/pecos/simulators/cuda_stabilizer/state.py +++ b/python/quantum-pecos/src/pecos/simulators/cuda_stabilizer/state.py @@ -88,6 +88,7 @@ def __init__(self, num_qubits: int, seed: int | None = None) -> None: def reset(self) -> Self: """Reset the quantum state to |0...0>.""" + check_fork_poison() self.backend.reset() return self diff --git a/python/quantum-pecos/src/pecos/simulators/cuda_statevec/state.py b/python/quantum-pecos/src/pecos/simulators/cuda_statevec/state.py index 1409b0517..47890b466 100644 --- a/python/quantum-pecos/src/pecos/simulators/cuda_statevec/state.py +++ b/python/quantum-pecos/src/pecos/simulators/cuda_statevec/state.py @@ -83,6 +83,7 @@ def __init__(self, num_qubits: int, seed: int | None = None) -> None: def reset(self) -> Self: """Reset the quantum state to |0...0>.""" + check_fork_poison() self.backend.reset() return self diff --git a/python/quantum-pecos/src/pecos/simulators/mps_pytket/state.py b/python/quantum-pecos/src/pecos/simulators/mps_pytket/state.py index 1a91a8fa2..1df5f6d9b 100644 --- a/python/quantum-pecos/src/pecos/simulators/mps_pytket/state.py +++ b/python/quantum-pecos/src/pecos/simulators/mps_pytket/state.py @@ -26,6 +26,7 @@ MPSxGate, ) +from pecos.simulators._cuda_fork_guard import check_fork_poison, mark_cuda_initialized from pecos.simulators.mps_pytket import bindings from pecos.simulators.sim_class_types import StateTN @@ -66,6 +67,7 @@ def __init__(self, num_qubits: int, **mps_params: SimulatorInitParams) -> None: For detailed documentation, see pytket-cutensornet Config class: https://docs.quantinuum.com/tket/extensions/pytket-cutensornet/ """ + check_fork_poison() if not isinstance(num_qubits, int): msg = "``num_qubits`` should be of type ``int``." raise TypeError(msg) @@ -80,6 +82,7 @@ def __init__(self, num_qubits: int, **mps_params: SimulatorInitParams) -> None: self.dtype = self.config._complex_t # cuTensorNet handle initialization + mark_cuda_initialized() self.libhandle = CuTensorNetHandle() # Initialise the MPS on state |0> diff --git a/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py b/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py index 4f8ec1d00..bbc00429d 100644 --- a/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py +++ b/python/quantum-pecos/tests/pecos/unit/test_cuda_fork_guard.py @@ -51,9 +51,13 @@ def test_marked_parent_poisons_forked_child() -> None: guard.check_fork_poison() except RuntimeError as exc: os._exit(0 if str(exc) == guard.CUDA_FORK_ERROR_MESSAGE else 2) - os._exit(1) + except BaseException: + os._exit(3) + else: + os._exit(1) assert _wait_for_child(pid) == 0 + guard.check_fork_poison() def test_unmarked_parent_does_not_poison_forked_child() -> None: @@ -62,9 +66,10 @@ def test_unmarked_parent_does_not_poison_forked_child() -> None: if pid == 0: try: guard.check_fork_poison() - except RuntimeError: + except BaseException: os._exit(1) - os._exit(0) + else: + os._exit(0) assert _wait_for_child(pid) == 0 @@ -78,7 +83,14 @@ def test_marked_parent_warns_only_once_before_fork() -> None: for _ in range(2): pid = os.fork() if pid == 0: - os._exit(0) + try: + guard.check_fork_poison() + except RuntimeError: + os._exit(0) + except BaseException: + os._exit(2) + else: + os._exit(1) assert _wait_for_child(pid) == 0 fork_warnings = [warning for warning in caught if warning.category is RuntimeWarning] diff --git a/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py b/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py index 015103195..4e3c59196 100644 --- a/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py +++ b/python/quantum-pecos/tests/pecos/unit/test_custatevec_state.py @@ -61,13 +61,25 @@ def _inject_optional_dependency_stubs(monkeypatch: pytest.MonkeyPatch) -> types. driver = types.ModuleType("cupy.cuda.driver") runtime.CUDARuntimeError = _StubRuntimeError runtime.runtimeGetVersion = Mock(return_value=12000) + runtime.deviceGetDefaultMemPool = Mock(return_value=1) + runtime.memPoolSetAttribute = Mock() + runtime.cudaMemPoolAttrReleaseThreshold = object() driver.CUDADriverError = _StubDriverError cuda.runtime = runtime cuda.driver = driver - cuda.Device = Mock() + cuda.Device = Mock( + return_value=types.SimpleNamespace(attributes={"MemoryPoolsSupported": True}, id=0), + ) + cuda.Stream = Mock(return_value=types.SimpleNamespace(ptr=1)) cp.cuda = cuda cp.complex128 = object() cp.zeros = Mock(return_value=[0, 0]) + cusv = types.SimpleNamespace( + create=Mock(return_value=object()), + set_stream=Mock(), + set_device_mem_handler=Mock(), + destroy=Mock(), + ) monkeypatch.setitem(sys.modules, "cupy", cp) monkeypatch.setitem(sys.modules, "cupy.cuda", cuda) @@ -77,9 +89,52 @@ def _inject_optional_dependency_stubs(monkeypatch: pytest.MonkeyPatch) -> types. monkeypatch.setattr(custatevec_state, "require_custatevec", Mock()) monkeypatch.setattr(custatevec_state, "cudaDataType", types.SimpleNamespace(CUDA_C_64F=object())) monkeypatch.setattr(custatevec_state, "ComputeType", types.SimpleNamespace(COMPUTE_64F=object())) + monkeypatch.setattr(custatevec_state, "cusv", cusv) return cp +def test_custatevec_constructor_marks_cuda_initialized(monkeypatch) -> None: + """Successful construction records the wrapper's first CUDA call.""" + cp = _inject_optional_dependency_stubs(monkeypatch) + + sim = custatevec_state.CuStateVec(1) + + assert vars(guard)["_state"]["cuda_initialized"] + cp.zeros.assert_called_once() + custatevec_state.cusv.create.assert_called_once() + assert sim.libhandle is not None + sim.libhandle = None + + +def test_custatevec_constructor_checks_fork_poison_before_cuda(monkeypatch) -> None: + """A poisoned child fails before dependency checks or CUDA calls.""" + cp = _inject_optional_dependency_stubs(monkeypatch) + vars(guard)["_state"]["forked_after_cuda_init"] = True + + with pytest.raises(RuntimeError) as exc_info: + custatevec_state.CuStateVec(1) + + assert str(exc_info.value) == guard.CUDA_FORK_ERROR_MESSAGE + custatevec_state.require_custatevec.assert_not_called() + cp.zeros.assert_not_called() + cp.cuda.runtime.runtimeGetVersion.assert_not_called() + cp.cuda.Device.assert_not_called() + custatevec_state.cusv.create.assert_not_called() + + +def test_custatevec_reset_checks_fork_poison(monkeypatch) -> None: + """An inherited simulator fails before reset touches its CUDA state.""" + _inject_optional_dependency_stubs(monkeypatch) + sim = custatevec_state.CuStateVec(1) + vars(guard)["_state"]["forked_after_cuda_init"] = True + + with pytest.raises(RuntimeError) as exc_info: + sim.reset() + + assert str(exc_info.value) == guard.CUDA_FORK_ERROR_MESSAGE + sim.libhandle = None + + def test_cuda_initialization_errors_explain_fork_and_preserve_cause(monkeypatch) -> None: """Translate the runtime and driver initialization statuses from CuPy calls.""" cp = _inject_optional_dependency_stubs(monkeypatch) diff --git a/python/quantum-pecos/tests/pecos/unit/test_mps_cuda_fork_guard.py b/python/quantum-pecos/tests/pecos/unit/test_mps_cuda_fork_guard.py new file mode 100644 index 000000000..002c4a780 --- /dev/null +++ b/python/quantum-pecos/tests/pecos/unit/test_mps_cuda_fork_guard.py @@ -0,0 +1,114 @@ +# Copyright 2026 The PECOS Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. + +"""GPU-free constructor wiring tests for the pytket MPS CUDA simulator.""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path +from unittest.mock import Mock + +import pytest +from pecos.simulators import _cuda_fork_guard as guard + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "src/pecos/simulators/mps_pytket/state.py" + + +@pytest.fixture(autouse=True) +def _reset_guard_state(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep synthetic CUDA initialization state isolated between tests.""" + monkeypatch.setattr( + guard, + "_state", + {"cuda_initialized": False, "forked_after_cuda_init": False, "fork_warning_emitted": False}, + ) + + +def _load_mps_module(monkeypatch: pytest.MonkeyPatch) -> types.SimpleNamespace: + """Load MPS state with stub pytket and cuTensorNet modules.""" + pytket = types.ModuleType("pytket") + pytket.__path__ = [] + pytket.Qubit = Mock(side_effect=lambda index: index) + extensions = types.ModuleType("pytket.extensions") + extensions.__path__ = [] + cutensornet = types.ModuleType("pytket.extensions.cutensornet") + cutensornet.__path__ = [] + structured_state = types.ModuleType("pytket.extensions.cutensornet.structured_state") + + config = Mock(return_value=types.SimpleNamespace(_complex_t=complex)) + handle = types.SimpleNamespace(destroy=Mock()) + handle_factory = Mock(return_value=handle) + logger = types.SimpleNamespace(info=Mock()) + mps_factory = Mock(return_value=types.SimpleNamespace(_logger=logger)) + structured_state.Config = config + structured_state.CuTensorNetHandle = handle_factory + structured_state.MPSxGate = mps_factory + + mps_package = types.ModuleType("pecos.simulators.mps_pytket") + mps_package.__path__ = [] + bindings = types.ModuleType("pecos.simulators.mps_pytket.bindings") + bindings.gate_dict = {} + mps_package.bindings = bindings + + for name, module in { + "pytket": pytket, + "pytket.extensions": extensions, + "pytket.extensions.cutensornet": cutensornet, + "pytket.extensions.cutensornet.structured_state": structured_state, + "pecos.simulators.mps_pytket": mps_package, + "pecos.simulators.mps_pytket.bindings": bindings, + }.items(): + monkeypatch.setitem(sys.modules, name, module) + + spec = importlib.util.spec_from_file_location("_test_mps_cuda_fork_guard", _MODULE_PATH) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + module.MPS.__del__ = Mock() + return types.SimpleNamespace( + module=module, + config=config, + handle_factory=handle_factory, + mps_factory=mps_factory, + qubit=pytket.Qubit, + ) + + +def test_mps_constructor_marks_cuda_initialized(monkeypatch) -> None: + """Successful handle construction records the MPS wrapper's CUDA call.""" + stubs = _load_mps_module(monkeypatch) + + sim = stubs.module.MPS(2) + + assert vars(guard)["_state"]["cuda_initialized"] + stubs.handle_factory.assert_called_once_with() + stubs.mps_factory.assert_called_once() + assert sim.libhandle is stubs.handle_factory.return_value + + +def test_mps_constructor_checks_fork_poison_before_cuda(monkeypatch) -> None: + """A poisoned child fails before MPS configuration or CUDA handle creation.""" + stubs = _load_mps_module(monkeypatch) + vars(guard)["_state"]["forked_after_cuda_init"] = True + + with pytest.raises(RuntimeError) as exc_info: + stubs.module.MPS(2) + + assert str(exc_info.value) == guard.CUDA_FORK_ERROR_MESSAGE + stubs.config.assert_not_called() + stubs.handle_factory.assert_not_called() + stubs.mps_factory.assert_not_called() + stubs.qubit.assert_not_called() diff --git a/python/quantum-pecos/tests/pecos/unit/test_multiprocessing_spawn_discipline.py b/python/quantum-pecos/tests/pecos/unit/test_multiprocessing_spawn_discipline.py index 561449101..ccb676446 100644 --- a/python/quantum-pecos/tests/pecos/unit/test_multiprocessing_spawn_discipline.py +++ b/python/quantum-pecos/tests/pecos/unit/test_multiprocessing_spawn_discipline.py @@ -14,30 +14,153 @@ from __future__ import annotations -import re +import ast from pathlib import Path _SOURCE_ROOT = Path(__file__).resolve().parents[3] / "src" / "pecos" _ALLOWLIST: frozenset[Path] = frozenset() -_FORK_HAZARD_PATTERNS = { - "get_context(fork)": re.compile(r"\bget_context\s*\(\s*(['\"])fork\1\s*\)"), - "set_start_method(fork)": re.compile(r"\bset_start_method\s*\(\s*(['\"])fork\1\s*\)"), - "multiprocessing.Pool": re.compile(r"\bmultiprocessing\.Pool\s*\("), - "mp.Pool": re.compile(r"\bmp\.Pool\s*\("), - "bare Pool": re.compile(r"(? str | None: + """Return a dotted name for a simple name or attribute expression.""" + if isinstance(expression, ast.Name): + return expression.id + if isinstance(expression, ast.Attribute): + prefix = _qualified_name(expression.value) + if prefix is not None: + return f"{prefix}.{expression.attr}" + return None + + +def _string_argument(call: ast.Call, keyword_name: str) -> str | None: + """Return a call's first positional or selected keyword string argument.""" + argument = ( + call.args[0] + if call.args + else next( + (keyword.value for keyword in call.keywords if keyword.arg == keyword_name), + None, + ) + ) + if isinstance(argument, ast.Constant) and isinstance(argument.value, str): + return argument.value + return None + + +def _find_fork_hazards(source: str) -> list[tuple[int, str]]: + """Find fork-prone process creation in parsed Python source.""" + tree = ast.parse(source) + source_lines = source.splitlines() + multiprocessing_aliases = {"multiprocessing", "mp"} + os_aliases = {"os"} + multiprocessing_imports: dict[str, str] = {} + process_pool_executor_names = {"ProcessPoolExecutor"} + hazards: set[tuple[int, str]] = set() + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "multiprocessing": + multiprocessing_aliases.add(alias.asname or alias.name) + elif alias.name == "os": + os_aliases.add(alias.asname or alias.name) + elif isinstance(node, ast.ImportFrom): + if node.module == "multiprocessing": + for alias in node.names: + local_name = alias.asname or alias.name + multiprocessing_imports[local_name] = alias.name + if alias.name in {"Pool", "Process"}: + hazards.add((node.lineno, f"from multiprocessing import {alias.name}")) + elif node.module == "concurrent.futures": + for alias in node.names: + if alias.name == "ProcessPoolExecutor": + process_pool_executor_names.add(alias.asname or alias.name) + + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + function_name = _qualified_name(node.func) + if function_name is None: + continue + name_parts = function_name.split(".") + root_name = name_parts[0] + leaf_name = name_parts[-1] + source_line = source_lines[node.lineno - 1].lstrip() if source_lines else "" + definition_line = source_line.startswith(("def ", "async def ", "class ")) + + imported_name = multiprocessing_imports.get(function_name) + is_multiprocessing_member = root_name in multiprocessing_aliases and len(name_parts) == 2 + if leaf_name == "get_context" or imported_name == "get_context": + method = _string_argument(node, "method") + if not node.args and method is None: + hazards.add((node.lineno, "get_context() uses the platform default")) + elif method == "fork": + hazards.add((node.lineno, 'get_context("fork")')) + elif leaf_name == "set_start_method" or imported_name == "set_start_method": + if _string_argument(node, "method") == "fork": + hazards.add((node.lineno, 'set_start_method("fork")')) + elif root_name in os_aliases and len(name_parts) == 2 and leaf_name == "fork": + hazards.add((node.lineno, "os.fork()")) + elif is_multiprocessing_member and leaf_name in {"Pool", "Process"}: + hazards.add((node.lineno, f"{root_name}.{leaf_name}()")) + elif isinstance(node.func, ast.Name) and not definition_line: + if function_name == "Pool" or imported_name == "Pool": + hazards.add((node.lineno, f"bare or imported {function_name}()")) + elif imported_name == "Process": + hazards.add((node.lineno, f"imported {function_name}()")) + + is_process_pool_executor = leaf_name == "ProcessPoolExecutor" or function_name in process_pool_executor_names + if is_process_pool_executor and not any(keyword.arg == "mp_context" for keyword in node.keywords): + hazards.add((node.lineno, "ProcessPoolExecutor() without mp_context")) + + return sorted(hazards) + + +def test_hazard_detector_catches_variants_without_safe_false_positives() -> None: + """Pin supported syntax variants and the intentional precision exclusions.""" + hazardous_sources = [ + 'multiprocessing.get_context("fork", force=True)', + "mp.get_context( 'fork' , extra=True)", + 'multiprocessing.set_start_method("fork", force=True)', + "multiprocessing.get_context()", + "os.fork()", + "multiprocessing.Process(target=work)", + "mp.Pool(processes=2)", + "Pool(processes=2)", + "from multiprocessing import Pool as WorkerPool", + "from multiprocessing import Process as WorkerProcess", + "ProcessPoolExecutor(max_workers=2)", + "ProcessPoolExecutor(\n max_workers=2,\n)", + ] + for source in hazardous_sources: + assert _find_fork_hazards(source), source + + safe_sources = [ + "# multiprocessing.Pool(processes=2)", + "def Pool(processes):\n return processes", + "class ProcessPoolExecutor:\n pass", + 'multiprocessing.get_context("spawn").Pool(processes=2)', + "ProcessPoolExecutor(max_workers=2, mp_context=spawn_context)", + "ProcessPoolExecutor(\n max_workers=2,\n mp_context=spawn_context,\n)", + ] + for source in safe_sources: + assert not _find_fork_hazards(source), source def test_pecos_source_avoids_fork_based_multiprocessing() -> None: """Require explicit non-fork multiprocessing throughout PECOS Python source.""" + assert _SOURCE_ROOT.is_dir(), f"PECOS source root does not exist: {_SOURCE_ROOT}" + source_paths = sorted(_SOURCE_ROOT.rglob("*.py")) + assert source_paths, f"No Python source files found under {_SOURCE_ROOT}" + violations = [] - for source_path in sorted(_SOURCE_ROOT.rglob("*.py")): + for source_path in source_paths: relative_path = source_path.relative_to(_SOURCE_ROOT) if relative_path in _ALLOWLIST: continue - for line_number, line in enumerate(source_path.read_text(encoding="utf-8").splitlines(), start=1): - for pattern_name, pattern in _FORK_HAZARD_PATTERNS.items(): - if pattern.search(line): - violations.append(f"{relative_path}:{line_number}: {pattern_name}: {line.strip()}") + source = source_path.read_text(encoding="utf-8") + for line_number, hazard in _find_fork_hazards(source): + source_line = source.splitlines()[line_number - 1].strip() + violations.append(f"{relative_path}:{line_number}: {hazard}: {source_line}") assert not violations, "Fork-unsafe multiprocessing usage found:\n" + "\n".join(violations) From fc38a626cdb0452e1e5cdc5d4c444de63e1b0375 Mon Sep 17 00:00:00 2001 From: Ciaran Ryan-Anderson Date: Thu, 6 Aug 2026 22:21:12 -0600 Subject: [PATCH 6/6] Restore the custatevec reset fork guard lost to a pre-commit mutation-check revert --- python/quantum-pecos/src/pecos/simulators/custatevec/state.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/quantum-pecos/src/pecos/simulators/custatevec/state.py b/python/quantum-pecos/src/pecos/simulators/custatevec/state.py index c80315e5d..0a52f5bba 100644 --- a/python/quantum-pecos/src/pecos/simulators/custatevec/state.py +++ b/python/quantum-pecos/src/pecos/simulators/custatevec/state.py @@ -157,6 +157,7 @@ def free( def reset(self) -> Self: """Reset the quantum state for another run without reinitializing.""" + check_fork_poison() # Initialize all qubits in the zero state if self.cupy_vector is not None: self.cupy_vector[:] = 0