feat: make PTA OOM injection controllable

torch_npu/csrc/core/npu/NPUCachingAllocator.cpp

@@ -1,5 +1,7 @@
 #include <algorithm>
+#include <atomic>
 #include <bitset>
+#include <cstdlib>
 #include <deque>
 #include <map>
 #include <memory>
@@ -23,6 +25,8 @@
 #include "torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h"
 #include "torch_npu/csrc/core/npu/NPURecovery.h"
 #include "torch_npu/csrc/core/npu/NPUGuard.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
+#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "NPUBlockHandle.h"
 #include "torch_npu/csrc/core/npu/NpuVariables.h"
 #include "torch_npu/csrc/core/npu/GetCANNInfo.h"
@@ -107,7 +111,9 @@ const std::string kMinDriverVersion = "25.0.RC1";     // minimum driver version
 const std::string kCannModule = "CANN";               // cann module name
 constexpr int kPrecision = 4;                         // precision of the memory usage information
 constexpr size_t kLazyQuerySize = 512;                // lazy query event size
-static int64_t g_malloc_call_count = 0;
+constexpr int64_t kDefaultPtaOomTriggerCount = 30004;
+static std::atomic<int64_t> g_pta_oom_candidate_count{0};
+static std::atomic<bool> g_pta_oom_triggered{false};
 static char SHAREABLE_HANDLE_VERSION = 1;
 enum ShareableHandleType : char {
     SHAREABLE_NPU_MALLOC = 'c',
@@ -116,6 +122,40 @@ enum ShareableHandleType : char {
 
 using StatTypes = std::array<bool, static_cast<size_t>(StatType::NUM_TYPES)>;
 
+int64_t getPtaOomTriggerCount()
+{
+    const static int64_t trigger_count = []() -> int64_t {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_COUNT");
+        return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : kDefaultPtaOomTriggerCount;
+    }();
+    return trigger_count;
+}
+
+void maybeThrowPtaOom(int device, size_t size)
+{
+    if (g_pta_oom_triggered.load() || size == 0) {
+        return;
+    }
+
+    const int64_t trigger_count = getPtaOomTriggerCount();
+    if (trigger_count <= 0) {
+        return;
+    }
+
+    if (c10_npu::currentStreamCaptureStatus() != c10_npu::CaptureStatus::None) {
+        return;
+    }
+
+    const int64_t current_count = ++g_pta_oom_candidate_count;
+    if (current_count > trigger_count && current_count < trigger_count + 2) {
+        g_pta_oom_triggered.store(true);
+        auto retmsg = std::string("NPU out of memory. Injected PTA OOM after ") +
+            std::to_string(current_count) + " normal tensor allocations. Tried to allocate " +
+            format_size(size) + " on NPU " + std::to_string(device) + ".";
+        TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
+    }
+}
+
 void update_stat(Stat &stat, int64_t amount)
 {
     stat.current += amount;
@@ -1152,12 +1192,7 @@ class DeviceCachingAllocator {
     // Thus, do not call a public method from another public method.
 
     Block *malloc(int device, size_t orig_size, aclrtStream stream, uint8_t allocator_type = 0)
-    {           
-        g_malloc_call_count++;
-        auto retmsg = std::string("NPU out of memory. Tried to allocate more than 1EB memory.");
-        if (g_malloc_call_count > 30004 && g_malloc_call_count < 30006) {
-            TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
-        }
+    {
         TORCH_NPU_MEMORY_LOGD("Allocating memory: size=%zu, device=%d", orig_size, device);
         // done outside the lock because we don't know what locks the recorder needs
         // to have...
@@ -3518,6 +3553,7 @@ class NpuCachingAllocator : public NPUAllocator {
 
         int device = 0;
         NPU_CHECK_ERROR(c10_npu::GetDevice(&device));
+        maybeThrowPtaOom(device, size);
         LazySetDevice(device);
         void *devPtr = nullptr;
         void (*deleteFunc)(void *) = &local_raw_delete;