diff --git a/cuda_bindings/tests/test_cuda.py b/cuda_bindings/tests/test_cuda.py index d85724fe49a..8282f21673b 100644 --- a/cuda_bindings/tests/test_cuda.py +++ b/cuda_bindings/tests/test_cuda.py @@ -17,6 +17,11 @@ from cuda.bindings._internal.utils import FunctionNotFoundError from cuda_python_test_helpers import driver_version_less_than +# Cap for memory pools created by tests. A pool created with maxSize 0 reserves +# virtual address space of about twice the device memory, and that reservation +# can fail on systems with a limited virtual address space. +POOL_SIZE = 2 * 1024 * 1024 # 2 MiB + def supportsMemoryPool(): err, isSupported = cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrMemoryPoolsSupported, 0) @@ -265,6 +270,7 @@ def test_cuda_memPool_attr(): poolProps.allocType = cuda.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED poolProps.location.id = 0 poolProps.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + poolProps.maxSize = POOL_SIZE attr_list = [None] * 8 err, pool = cuda.cuMemPoolCreate(poolProps) @@ -460,11 +466,29 @@ def test_cuda_mem_range_attr(device): assert err == cuda.CUresult.CUDA_SUCCESS +@pytest.fixture +def xfail_on_graph_mempool_oom(): + """Return a check that xfails the calling test if the graph memory pool cannot be created. + + The graph memory pool reserves virtual address space of about twice the + device memory, and unlike pools from cuMemPoolCreate it cannot be capped. + On systems with a limited virtual address space this reservation can fail + with CUDA_ERROR_OUT_OF_MEMORY, especially when another uncapped pool + already exists in the process. + """ + + def check(err): + if err == cuda.CUresult.CUDA_ERROR_OUT_OF_MEMORY: + pytest.xfail("graph memory pool creation failed with CUDA_ERROR_OUT_OF_MEMORY") + + return check + + @pytest.mark.skipif( driver_version_less_than(11040) or not supportsMemoryPool(), reason="Mempool for graphs not supported" ) @pytest.mark.thread_unsafe(reason="used high memory can be higher if threaded.") -def test_cuda_graphMem_attr(device): +def test_cuda_graphMem_attr(device, xfail_on_graph_mempool_oom): err, stream = cuda.cuStreamCreate(0) assert err == cuda.CUresult.CUDA_SUCCESS @@ -485,6 +509,7 @@ def test_cuda_graphMem_attr(device): assert destroy_err == cuda.CUresult.CUDA_SUCCESS (destroy_err,) = cuda.cuStreamDestroy(stream) assert destroy_err == cuda.CUresult.CUDA_SUCCESS + xfail_on_graph_mempool_oom(err) assert err == cuda.CUresult.CUDA_SUCCESS err, freeNode = cuda.cuGraphAddMemFreeNode(graph, [allocNode], 1, params.dptr) assert err == cuda.CUresult.CUDA_SUCCESS diff --git a/cuda_bindings/tests/test_cudart.py b/cuda_bindings/tests/test_cudart.py index b16391210ed..f1c263f7c13 100644 --- a/cuda_bindings/tests/test_cudart.py +++ b/cuda_bindings/tests/test_cudart.py @@ -13,6 +13,11 @@ from cuda.bindings import runtime from cuda_python_test_helpers import driver_version_less_than +# Cap for memory pools created by tests. A pool created with maxSize 0 reserves +# virtual address space of about twice the device memory, and that reservation +# can fail on systems with a limited virtual address space. +POOL_SIZE = 2 * 1024 * 1024 # 2 MiB + def isSuccess(err): return err == cudart.cudaError_t.cudaSuccess @@ -520,6 +525,7 @@ def test_cudart_MemPool_attr(): poolProps.allocType = cudart.cudaMemAllocationType.cudaMemAllocationTypePinned poolProps.location.id = 0 poolProps.location.type = cudart.cudaMemLocationType.cudaMemLocationTypeDevice + poolProps.maxSize = POOL_SIZE attr_list = [None] * 8 err, pool = cudart.cudaMemPoolCreate(poolProps)