Limit the number of single allocation memory warnings.

PiperOrigin-RevId: 194415953

1 files changed, 16 insertions, 5 deletions

diff --git a/tensorflow/core/framework/allocator.cc b/tensorflow/core/framework/allocator.cc
index 1a7e5219cd..1c62d37955 100644
--- a/tensorflow/core/framework/allocator.cc
+++ b/tensorflow/core/framework/allocator.cc
@@ -68,6 +68,9 @@ static const double kLargeAllocationWarningThreshold = 0.1;
 // exceeds this threshold.
 static const double kTotalAllocationWarningThreshold = 0.5;
 
+static const int kMaxSingleAllocationWarnings = 5;
+static const int kMaxTotalAllocationWarnings = 1;
+
 // Cache first invocation to port::AvailableRam, as it can be expensive.
 static int64_t LargeAllocationWarningBytes() {
   static int64_t value = static_cast<int64>(port::AvailableRam() *
@@ -90,14 +93,18 @@ void EnableCPUAllocatorFullStats(bool enable) {
 
 class CPUAllocator : public Allocator {
  public:
-  CPUAllocator() : total_allocation_warning_triggered_(false) {}
+  CPUAllocator()
+      : single_allocation_warning_count_(0),
+        total_allocation_warning_count_(0) {}
 
   ~CPUAllocator() override {}
 
   string Name() override { return "cpu"; }
 
   void* AllocateRaw(size_t alignment, size_t num_bytes) override {
-    if (num_bytes > LargeAllocationWarningBytes()) {
+    if (num_bytes > LargeAllocationWarningBytes() &&
+        single_allocation_warning_count_ < kMaxSingleAllocationWarnings) {
+      ++single_allocation_warning_count_;
       LOG(WARNING) << "Allocation of " << num_bytes << " exceeds "
                    << 100 * kLargeAllocationWarningThreshold
                    << "% of system memory.";
@@ -115,11 +122,11 @@ class CPUAllocator : public Allocator {
           std::max<int64>(stats_.max_alloc_size, alloc_size);
 
       if (stats_.bytes_in_use > TotalAllocationWarningBytes() &&
-          !total_allocation_warning_triggered_) {
+          total_allocation_warning_count_ < kMaxTotalAllocationWarnings) {
+        ++total_allocation_warning_count_;
         LOG(WARNING) << "Total allocated memory " << stats_.bytes_in_use
                      << "exceeds " << 100 * kTotalAllocationWarningThreshold
                      << "% of system memory";
-        total_allocation_warning_triggered_ = true;
       }
     }
     return p;
@@ -154,7 +161,11 @@ class CPUAllocator : public Allocator {
  private:
   mutex mu_;
   AllocatorStats stats_ GUARDED_BY(mu_);
-  bool total_allocation_warning_triggered_ GUARDED_BY(mu_);
+
+  // Use <atomic> for single allocations to avoid mutex contention when
+  // statistics are disabled.
+  std::atomic<int> single_allocation_warning_count_;
+  int total_allocation_warning_count_ GUARDED_BY(mu_);
 
   TF_DISALLOW_COPY_AND_ASSIGN(CPUAllocator);
 };