Add automatic full-card PTA OOM injection (no PTA recovery)

torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp

@@ -23,6 +23,7 @@
 #include "torch_npu/csrc/core/npu/CachingHostAllocator.h"
 #include "torch_npu/csrc/aten/NPUOpApiNativeFunctions.h"
 #include "torch_npu/csrc/aten/NPUNativeFunctions.h"
+#include "torch_npu/csrc/core/npu/interface/AclInterface.h"
 #include "third_party/op-plugin/op_plugin/utils/op_api_common.h"
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -49,7 +50,7 @@ void copy_between_host_and_device_opapi(at::Tensor& dst, const at::Tensor& src,
         void* currentPtr = torch_npu::utils::is_npu(dst) ? src.data_ptr() : dst.data_ptr();
         process_non_blocking_copy(storage, currentPtr, stream, kind);
     } else {
-        aclError error = aclrtSynchronizeStream(stream);
+        aclError error = c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream);
         auto ret = CalcuOpUtil::AclrtMemcpyWithModeSwitch(
             std::make_pair(dst.storage().unsafeGetStorageImpl(), dst.storage_offset() * dst.itemsize()), nbytes,
             std::make_pair(src.storage().unsafeGetStorageImpl(), src.storage_offset() * src.itemsize()), nbytes, kind);

torch_npu/csrc/core/npu/NPUCachingAllocator.cpp

@@ -107,7 +107,6 @@ 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;
 static char SHAREABLE_HANDLE_VERSION = 1;
 enum ShareableHandleType : char {
     SHAREABLE_NPU_MALLOC = 'c',
@@ -1152,9 +1151,8 @@ 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.");
+    {
+        maybeThrowPtaOom("malloc", device);
 
         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
@@ -2937,6 +2935,9 @@ class DeviceCachingAllocator {
                 } else {
                     TORCH_NPU_MEMORY_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event.get());
                 }
+                if (check_error) {
+                    maybeThrowPtaOom("NPUCachingAllocator::synchronize_and_free_events");
+                }
 #ifndef BUILD_LIBTORCH
                 const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace();
                 if (C10_UNLIKELY(trigger)) {

torch_npu/csrc/core/npu/NPUEvent.cpp

@@ -2,6 +2,7 @@
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
 #include "torch_npu/csrc/core/npu/NPUGuard.h"
 #include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/core/npu/NPUEventManager.h"
 #include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
 #include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"
@@ -202,6 +203,7 @@ void NPUEvent::synchronize() const
         }
         NPU_CHECK_ERROR(aclrtSynchronizeEvent(event_));
         ASCEND_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event_);
+        maybeThrowPtaOom("NPUEvent::synchronize");
 #ifndef BUILD_LIBTORCH
         const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
         if (C10_UNLIKELY(trigger)) {

torch_npu/csrc/core/npu/NPUException.cpp

@@ -1,6 +1,12 @@
+#include <atomic>
+#include <chrono>
+#include <cstdlib>
+#include <thread>
+
 #include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
-#include "torch_npu/csrc/core/npu/NPUStream.h"
+#include "torch_npu/csrc/core/npu/NPURecovery.h"
 #include "torch_npu/csrc/core/npu/NpuVariables.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 
@@ -433,4 +439,133 @@ bool isCannOOM(const std::string &errMsg)
     return false;
 }
 
+namespace {
+static std::atomic<bool> g_pta_oom_injected{false};
+static std::atomic<bool> g_pta_oom_timer_started{false};
+static std::atomic<bool> g_pta_oom_timer_expired{false};
+static std::chrono::steady_clock::time_point g_pta_oom_start_time;
+static constexpr int64_t kDefaultTriggerAfterSeconds = 360;
+
+bool isPtaOomDebugEnabled()
+{
+    const static bool enabled = []() -> bool {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_DEBUG");
+        return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0);
+    }();
+    return enabled;
+}
+
+int64_t getTriggerAfterSeconds()
+{
+    const static int64_t trigger_seconds = []() -> int64_t {
+        char *seconds_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_AFTER_SECONDS");
+        if (seconds_val != nullptr) {
+            return strtol(seconds_val, nullptr, 10);
+        }
+        char *minutes_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_AFTER_MINUTES");
+        if (minutes_val != nullptr) {
+            return strtol(minutes_val, nullptr, 10) * 60;
+        }
+        char *timer_mode = c10_npu::option::get_and_log_env("PTA_OOM_TIMER");
+        if (timer_mode != nullptr && strtol(timer_mode, nullptr, 10) != 0) {
+            return kDefaultTriggerAfterSeconds;
+        }
+        return 0;
+    }();
+    return trigger_seconds;
+}
+
+void throwFullCardPtaOom(const char *context, int device);
+
+void ensurePtaOomTimerStarted()
+{
+    if (g_pta_oom_timer_started.load()) {
+        return;
+    }
+    const int64_t trigger_seconds = getTriggerAfterSeconds();
+    if (trigger_seconds <= 0) {
+        return;
+    }
+    bool expected = false;
+    if (!g_pta_oom_timer_started.compare_exchange_strong(expected, true)) {
+        return;
+    }
+    g_pta_oom_start_time = std::chrono::steady_clock::now();
+    ASCEND_LOGI("PTA OOM timer started, will trigger after %lld seconds",
+        static_cast<long long>(trigger_seconds));
+    std::thread([trigger_seconds]() {
+        std::this_thread::sleep_for(std::chrono::seconds(trigger_seconds));
+        if (!g_pta_oom_injected.load()) {
+            g_pta_oom_timer_expired.store(true);
+            ASCEND_LOGI("PTA OOM timer expired after %lld seconds, pending throw on next hook",
+                static_cast<long long>(trigger_seconds));
+        }
+    }).detach();
+}
+
+int64_t getElapsedSecondsSincePtaOomStart()
+{
+    const auto now = std::chrono::steady_clock::now();
+    return std::chrono::duration_cast<std::chrono::seconds>(now - g_pta_oom_start_time).count();
+}
+
+const bool kPtaOomTimerArmed = []() {
+    if (getTriggerAfterSeconds() > 0) {
+        ensurePtaOomTimerStarted();
+    }
+    return true;
+}();
+
+void throwFullCardPtaOom(const char *context, int device)
+{
+    if (g_pta_oom_injected.exchange(true)) {
+        return;
+    }
+
+    if (device < 0) {
+        NPU_CHECK_ERROR(c10_npu::GetDevice(&device));
+    }
+
+    NPUCachingAllocator::markAllBlockUnsafe(device);
+    c10_npu::set_npu_data_unsafe_flag(true);
+
+    const int64_t trigger_seconds = getTriggerAfterSeconds();
+    auto retmsg = std::string("NPU out of memory. Injected full-card PTA OOM on NPU ") +
+        std::to_string(device) +
+        ". All existing tensors on this device are marked unsafe. "
+        "Triggered after " + std::to_string(trigger_seconds) + " seconds";
+    if (context != nullptr && context[0] != '\0') {
+        retmsg += ", context is ";
+        retmsg += context;
+    }
+    retmsg += ". ";
+    retmsg += PTA_ERROR(ErrCode::MEMORY);
+    TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
+}
+} // namespace
+
+void maybeThrowPtaOom(const char *context, int device)
+{
+    const int64_t trigger_seconds = getTriggerAfterSeconds();
+    if (trigger_seconds <= 0 || g_pta_oom_injected.load()) {
+        return;
+    }
+    ensurePtaOomTimerStarted();
+    const bool expired = g_pta_oom_timer_expired.load() ||
+        (g_pta_oom_timer_started.load() && getElapsedSecondsSincePtaOomStart() >= trigger_seconds);
+    if (!expired) {
+        if (isPtaOomDebugEnabled()) {
+            const int64_t elapsed = getElapsedSecondsSincePtaOomStart();
+            if (elapsed > 0 && elapsed % 60 == 0) {
+                ASCEND_LOGI("PTA OOM timer: elapsed=%lld/%lld seconds context=%s",
+                    static_cast<long long>(elapsed),
+                    static_cast<long long>(trigger_seconds),
+                    context != nullptr ? context : "");
+            }
+        }
+        return;
+    }
+    throwFullCardPtaOom(context, device);
+}
+
 } // namespace c10_npu

torch_npu/csrc/core/npu/NPUException.h

@@ -336,6 +336,8 @@ std::string handleSuspectRemoteError(int errorCode);
 
 bool isCannOOM(const std::string &errMsg);
 
+void maybeThrowPtaOom(const char *context = nullptr, int device = -1);
+
 bool ShouldAppendDeviceErrorVerbose();
 
 void clear_device_error_info();

torch_npu/csrc/core/npu/NPUFunctions.cpp

@@ -6,6 +6,7 @@
 #include "torch_npu/csrc/core/npu/NPUAffinityController.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "torch_npu/csrc/core/npu/GetCANNInfo.h"
+#include "torch_npu/csrc/core/npu/interface/AclInterface.h"
 #include "third_party/acl/inc/acl/acl_rt.h"
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -363,7 +364,7 @@ void stream_synchronize(aclrtStream stream)
         trigger->traceNpuStreamSynchronization(reinterpret_cast<uintptr_t>(stream));
     }
 #endif
-    NPU_CHECK_ERROR(aclrtSynchronizeStream(stream));
+    NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream));
 }
 
 aclError SetDeviceResLimit(int32_t device, int32_t type, uint32_t value)

torch_npu/csrc/core/npu/NPUQueue.cpp

@@ -8,6 +8,7 @@
 #include "torch_npu/csrc/framework/OpCommand.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "torch_npu/csrc/core/npu/NPUEventManager.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/logging/LogContext.h"
 
 #ifndef BUILD_LIBTORCH
@@ -348,6 +349,9 @@ NPUStatus Repository::MakeSureQueueEmpty(bool check_error)
             throw std::runtime_error(runtime_error);
         }
     }
+    if (check_error) {
+        maybeThrowPtaOom("MakeSureQueueEmpty", device_idx);
+    }
     logger->debug("MakeSureQueueEmpty: clearing successful, device = %d, write_idx = %u, read_idx = %u, status = %d",
         device_idx, write_idx.idx, read_idx.idx, GetStatus());
 

torch_npu/csrc/core/npu/NPUStream.cpp

@@ -14,6 +14,7 @@
 #include "torch_npu/csrc/core/npu/NPUGuard.h"
 #include "torch_npu/csrc/core/npu/NPUQueue.h"
 #include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
 #include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"
@@ -508,6 +509,10 @@ NPUStatus emptyAllNPUStream(bool check_error)
         }
     }
 
+    if (check_error) {
+        maybeThrowPtaOom("emptyAllNPUStream");
+    }
+
     return NPU_STATUS_SUCCESS;
 }
 
@@ -559,6 +564,9 @@ bool npuSynchronizeDevice(bool check_error)
             ASCEND_LOGE("MakeSureQueueEmpty fail, ret: %s", ret.c_str());
         }
     }
+    if (check_error) {
+        maybeThrowPtaOom("npuSynchronizeDevice");
+    }
     auto acl_ret = c10_npu::acl::AclrtSynchronizeDeviceWithTimeout();
     if (acl_ret != ACL_ERROR_NONE) {
         CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(acl_ret);

torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp

@@ -12,6 +12,8 @@
 #include "torch_npu/csrc/core/NPUSerialization.h"
 #include "torch_npu/csrc/core/npu/NPUHooksInterface.h"
 #include "torch_npu/csrc/core/npu/NPUEventManager.h"
+#include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -221,6 +223,7 @@ void NPUGuardImpl::synchronizeEvent(void* event) const
 
     NPU_CHECK_ERROR_WITHOUT_UCE(aclrtSynchronizeEvent(npu_event));
     ASCEND_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", npu_event);
+    maybeThrowPtaOom("NPUGuardImpl::synchronizeEvent");
 #ifndef BUILD_LIBTORCH
     const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
     if (C10_UNLIKELY(trigger)) {

torch_npu/csrc/core/npu/interface/AclInterface.cpp

@@ -8,6 +8,7 @@
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "torch_npu/csrc/core/npu/NPUException.h"
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/core/npu/GetCANNInfo.h"
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -594,8 +595,9 @@ aclError AclrtSynchronizeStreamWithTimeout(aclrtStream stream) {
     typedef aclError (*AclrtSynchronizeStreamWithTimeout)(aclrtStream, int32_t);
     static AclrtSynchronizeStreamWithTimeout func = (AclrtSynchronizeStreamWithTimeout)GET_FUNC(aclrtSynchronizeStreamWithTimeout);
     int32_t timeout = c10_npu::option::OptionsManager::GetACLExecTimeout();
+    aclError ret = ACL_ERROR_NONE;
     if (func != nullptr) {
-        return func(stream, timeout);
+        ret = func(stream, timeout);
     } else {
         TORCH_NPU_WARN_ONCE(func, "Failed to find function", "aclrtSynchronizeStreamWithTimeout");
         typedef aclError (*AclrtSynchronizeStream)(aclrtStream);
@@ -604,8 +606,10 @@ aclError AclrtSynchronizeStreamWithTimeout(aclrtStream stream) {
             func_backup = (AclrtSynchronizeStream)GET_FUNC(aclrtSynchronizeStream);
         }
         TORCH_CHECK(func_backup, "Failed to find function", "aclrtSynchronizeStreamWithTimeout and aclrtSynchronizeStream", PROF_ERROR(ErrCode::NOT_FOUND));
-        return func_backup(stream);
+        ret = func_backup(stream);
     }
+    c10_npu::maybeThrowPtaOom("AclrtSynchronizeStreamWithTimeout");
+    return ret;
 }
 
 aclError AclrtDestroyStreamForce(aclrtStream stream) {
@@ -992,8 +996,9 @@ aclError AclrtSynchronizeDeviceWithTimeout(void)
     typedef aclError (*AclrtSynchronizeDeviceWithTimeout)(int32_t);
     static AclrtSynchronizeDeviceWithTimeout func = (AclrtSynchronizeDeviceWithTimeout)GET_FUNC(aclrtSynchronizeDeviceWithTimeout);
     int32_t timeout = c10_npu::option::OptionsManager::GetACLDeviceSyncTimeout();
+    aclError ret = ACL_ERROR_NONE;
     if (func != nullptr) {
-        return func(timeout);
+        ret = func(timeout);
     } else {
         if (timeout > 0) {
             TORCH_NPU_WARN_ONCE("The ACL_DEVICE_SYNC_TIMEOUT does not take effect. If you want to enable this env, please upgrade CANN to the matching version.");
@@ -1004,8 +1009,10 @@ aclError AclrtSynchronizeDeviceWithTimeout(void)
             func_backup = (AclrtSynchronizeDevice)GET_FUNC(aclrtSynchronizeDevice);
         }
         TORCH_CHECK(func_backup, "Failed to find function ", "aclrtSynchronizeDeviceWithTimeout and aclrtSynchronizeDevice", PTA_ERROR(ErrCode::NOT_FOUND));
-        return func_backup();
+        ret = func_backup();
     }
+    c10_npu::maybeThrowPtaOom("AclrtSynchronizeDeviceWithTimeout");
+    return ret;
 }
 
 aclError AclrtEventGetTimestamp(aclrtEvent event, uint64_t *timestamp)
@@ -1076,6 +1083,7 @@ aclError AclmdlRIDebugPrint(aclmdlRI modelRI)
 aclError AclmdlRIExecuteAsync(aclmdlRI modelRI, aclrtStream stream)
 {
     ACL_CALL_LOG("aclmdlRIExecuteAsync", "modelRI=" << modelRI << ", stream=" << stream);
+    c10_npu::maybeThrowPtaOom("AclmdlRIExecuteAsync");
     typedef aclError (*AclmdlRIExecuteAsync)(aclmdlRI, aclrtStream);
     static AclmdlRIExecuteAsync func = nullptr;
     if (func == nullptr) {
@@ -1574,6 +1582,9 @@ aclError AclrtMemcpyAsyncWithCondition(void *dst, size_t destMax, const void *sr
         func = (AclrtMemcpyAsyncWithConditionFunc)GET_FUNC(aclrtMemcpyAsyncWithCondition);
     }
     TORCH_CHECK(func, "Failed to find function ", "aclrtMemcpyAsyncWithCondition", PROF_ERROR(ErrCode::NOT_FOUND));
+    if (kind == aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST) {
+        c10_npu::maybeThrowPtaOom("AclrtMemcpyAsyncWithCondition");
+    }
     return func(dst, destMax, src, count, kind, stream);
 }