From 069602e028d01b5287f719ccd8fae02026842a05 Mon Sep 17 00:00:00 2001
From: Bram Wasti <bwasti@fb.com>
Date: Tue, 26 Jan 2021 13:03:23 -0800
Subject: [PATCH] [torch vitals] Initial implementation (#51047)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51047

If the environment variable `TORCH_VITAL` is set to a non-zero length string, the vitals a dumped at program end.

The API is very similar to google's logging

Test Plan: buck test //caffe2/aten:vitals

Reviewed By: bitfort

Differential Revision: D25791248

fbshipit-source-id: 0b40da7d22c31d2c4b2094f0dcb1229a35338ac2
---
 aten/src/ATen/core/Vitals.cpp | 28 ++++++++++++++
 aten/src/ATen/core/Vitals.h   | 44 +++++++++++++++++++++
 aten/src/ATen/test/vitals.cpp | 73 +++++++++++++++++++++++++++++++++++
 3 files changed, 145 insertions(+)
 create mode 100644 aten/src/ATen/core/Vitals.cpp
 create mode 100644 aten/src/ATen/core/Vitals.h
 create mode 100644 aten/src/ATen/test/vitals.cpp

diff --git a/aten/src/ATen/core/Vitals.cpp b/aten/src/ATen/core/Vitals.cpp
new file mode 100644
index 000000000000..da1a208ab501
--- /dev/null
+++ b/aten/src/ATen/core/Vitals.cpp
@@ -0,0 +1,28 @@
+#include <ATen/core/Vitals.h>
+#include <cstdlib>
+
+TorchVitalAttr& TorchVital::create(const std::string& attr) {
+  if (!torchVitalEnabled()) {
+    static TorchVitalAttr disabled;
+    return disabled;
+  }
+  auto iter = attrs.find(attr);
+  if (iter == attrs.end()) {
+    auto r = attrs.emplace(std::make_pair(attr, TorchVitalAttr()));
+    return r.first->second;
+  }
+  return iter->second;
+}
+
+bool torchVitalEnabled() {
+  // If this is a performance hit, make `enabled` variable static
+  // and return `const bool&` instead
+  bool enabled = []() {
+    auto e = getenv("TORCH_VITAL");
+    if (e != nullptr) {
+      return strlen(e) > 0;
+    }
+    return false;
+  }();
+  return enabled;
+}
diff --git a/aten/src/ATen/core/Vitals.h b/aten/src/ATen/core/Vitals.h
new file mode 100644
index 000000000000..19d3436ebfef
--- /dev/null
+++ b/aten/src/ATen/core/Vitals.h
@@ -0,0 +1,44 @@
+#pragma once
+#include <cstring>
+#include <iostream>
+#include <sstream>
+#include <unordered_map>
+
+bool torchVitalEnabled();
+
+struct TorchVitalAttr {
+  // always initialized to empty
+  std::string value = "";
+  template <typename T>
+  TorchVitalAttr& operator<<(const T& t) {
+    if (torchVitalEnabled()) {
+      std::stringstream ss;
+      ss << t;
+      value += ss.str();
+    }
+    return *this;
+  }
+};
+
+struct TorchVital {
+  std::string name;
+  std::unordered_map<std::string, TorchVitalAttr> attrs;
+
+  explicit TorchVital(std::string n) : name(std::move(n)) {}
+  TorchVital() = delete;
+
+  TorchVitalAttr& create(const std::string& attr);
+
+  ~TorchVital() {
+    for (const auto& m : attrs) {
+      std::cout << "[TORCH_VITAL] " << name << "." << m.first << "\t\t "
+                << m.second.value << "\n";
+    }
+  }
+};
+
+#define TORCH_VITAL_DECLARE(name) extern TorchVital TorchVital_##name;
+
+#define TORCH_VITAL_DEFINE(name) TorchVital TorchVital_##name(#name);
+
+#define TORCH_VITAL(name, attr) TorchVital_##name.create(#attr)
diff --git a/aten/src/ATen/test/vitals.cpp b/aten/src/ATen/test/vitals.cpp
new file mode 100644
index 000000000000..d39e64ebf23b
--- /dev/null
+++ b/aten/src/ATen/test/vitals.cpp
@@ -0,0 +1,73 @@
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+#include <ATen/core/Vitals.h>
+#include <cstdlib>
+
+TEST(Vitals, Basic) {
+  std::stringstream buffer;
+
+  std::streambuf* sbuf = std::cout.rdbuf();
+  std::cout.rdbuf(buffer.rdbuf());
+  {
+    setenv("TORCH_VITAL", "1", 1);
+    TORCH_VITAL_DEFINE(Testing);
+    TORCH_VITAL(Testing, Attribute0) << 1;
+    TORCH_VITAL(Testing, Attribute1) << "1";
+    TORCH_VITAL(Testing, Attribute2) << 1.0f;
+    TORCH_VITAL(Testing, Attribute3) << 1.0;
+    auto t = at::ones({1, 1});
+    TORCH_VITAL(Testing, Attribute4) << t;
+  }
+  std::cout.rdbuf(sbuf);
+
+  auto s = buffer.str();
+  ASSERT_TRUE(s.find("Testing.Attribute0\t\t 1") != std::string::npos);
+  ASSERT_TRUE(s.find("Testing.Attribute1\t\t 1") != std::string::npos);
+  ASSERT_TRUE(s.find("Testing.Attribute2\t\t 1") != std::string::npos);
+  ASSERT_TRUE(s.find("Testing.Attribute3\t\t 1") != std::string::npos);
+  ASSERT_TRUE(s.find("Testing.Attribute4\t\t  1") != std::string::npos);
+}
+
+TEST(Vitals, MultiString) {
+  std::stringstream buffer;
+
+  std::streambuf* sbuf = std::cout.rdbuf();
+  std::cout.rdbuf(buffer.rdbuf());
+  {
+    setenv("TORCH_VITAL", "1", 1);
+    TORCH_VITAL_DEFINE(Testing);
+    TORCH_VITAL(Testing, Attribute0) << 1 << " of " << 2;
+    TORCH_VITAL(Testing, Attribute1) << 1;
+    TORCH_VITAL(Testing, Attribute1) << " of ";
+    TORCH_VITAL(Testing, Attribute1) << 2;
+  }
+  std::cout.rdbuf(sbuf);
+
+  auto s = buffer.str();
+  ASSERT_TRUE(s.find("Testing.Attribute0\t\t 1 of 2") != std::string::npos);
+  ASSERT_TRUE(s.find("Testing.Attribute1\t\t 1 of 2") != std::string::npos);
+}
+
+TEST(Vitals, OnAndOff) {
+  for (auto i = 0; i < 2; ++i) {
+    std::stringstream buffer;
+
+    std::streambuf* sbuf = std::cout.rdbuf();
+    std::cout.rdbuf(buffer.rdbuf());
+    {
+      setenv("TORCH_VITAL", i ? "1" : "", 1);
+      TORCH_VITAL_DEFINE(Testing);
+      TORCH_VITAL(Testing, Attribute0) << 1;
+    }
+    std::cout.rdbuf(sbuf);
+
+    auto s = buffer.str();
+    auto f = s.find("Testing.Attribute0\t\t 1");
+    if (i) {
+      ASSERT_TRUE(f != std::string::npos);
+    } else {
+      ASSERT_TRUE(f == std::string::npos);
+    }
+  }
+}