[torch vitals] Initial implementation (#51047)

Summary:
Pull Request resolved: #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

@@ -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;
+}

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)

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);
+    }
+  }
+}