[MEMORY] retry GPU memory allocation if fragmented (apache#16194)

larroy · Sep 28, 2019 · c068f3f · c068f3f
1 parent 8147279
commit c068f3f
Showing 1 changed file with 20 additions and 4 deletions.
diff --git a/src/storage/pooled_storage_manager.h b/src/storage/pooled_storage_manager.h
@@ -153,8 +153,16 @@ void GPUPooledStorageManager::Alloc(Storage::Handle* handle) {
 
     void* ret = nullptr;
     cudaError_t e = cudaMalloc(&ret, size);
-    if (e != cudaSuccess && e != cudaErrorCudartUnloading) {
-      LOG(FATAL) << "cudaMalloc failed: " << cudaGetErrorString(e);
+    if (e != cudaSuccess) {
+      if (e == cudaErrorMemoryAllocation) {
+        ReleaseAll();
+        e = cudaMalloc(&ret, size);
+        if (e != cudaSuccess && e != cudaErrorCudartUnloading) {
+          LOG(FATAL) << "cudaMalloc retry failed: " << cudaGetErrorString(e);
+        }
+      } else if (e != cudaErrorCudartUnloading) {
+        LOG(FATAL) << "cudaMalloc failed: " << cudaGetErrorString(e);
+      }
     }
     used_memory_ += size;
     handle->dptr = ret;
@@ -328,8 +336,16 @@ void GPUPooledRoundedStorageManager::Alloc(Storage::Handle* handle) {
 
     void* ret = nullptr;
     cudaError_t e = cudaMalloc(&ret, size);
-    if (e != cudaSuccess && e != cudaErrorCudartUnloading) {
-      LOG(FATAL) << "cudaMalloc failed: " << cudaGetErrorString(e);
+    if (e != cudaSuccess) {
+      if (e == cudaErrorMemoryAllocation) {
+        ReleaseAll();
+        e = cudaMalloc(&ret, size);
+        if (e != cudaSuccess && e != cudaErrorCudartUnloading) {
+          LOG(FATAL) << "cudaMalloc retry failed: " << cudaGetErrorString(e);
+        }
+      } else if (e != cudaErrorCudartUnloading) {
+        LOG(FATAL) << "cudaMalloc failed: " << cudaGetErrorString(e);
+      }
     }
     used_memory_ += size;
     handle->dptr = ret;