ai-content-maker/.venv/Lib/site-packages/numba/cuda/tests/cudapy/test_multiprocessing.py

47 lines
1.2 KiB
Python

import os
import multiprocessing as mp
import numpy as np
from numba import cuda
from numba.cuda.testing import skip_on_cudasim, CUDATestCase
import unittest
has_mp_get_context = hasattr(mp, 'get_context')
is_unix = os.name == 'posix'
def fork_test(q):
from numba.cuda.cudadrv.error import CudaDriverError
try:
cuda.to_device(np.arange(1))
except CudaDriverError as e:
q.put(e)
else:
q.put(None)
@skip_on_cudasim('disabled for cudasim')
class TestMultiprocessing(CUDATestCase):
@unittest.skipUnless(has_mp_get_context, 'requires mp.get_context')
@unittest.skipUnless(is_unix, 'requires Unix')
def test_fork(self):
"""
Test fork detection.
"""
cuda.current_context() # force cuda initialize
# fork in process that also uses CUDA
ctx = mp.get_context('fork')
q = ctx.Queue()
proc = ctx.Process(target=fork_test, args=[q])
proc.start()
exc = q.get()
proc.join()
# there should be an exception raised in the child process
self.assertIsNotNone(exc)
self.assertIn('CUDA initialized before forking', str(exc))
if __name__ == '__main__':
unittest.main()