Add JustKnobs

test/cpp/api/CMakeLists.txt

@@ -41,6 +41,7 @@ set(TORCH_API_TEST_SOURCES
   ${TORCH_API_TEST_DIR}/torch_include.cpp
   ${TORCH_API_TEST_DIR}/inference_mode.cpp
   ${TORCH_API_TEST_DIR}/grad_mode.cpp
+  ${TORCH_API_TEST_DIR}/JustKnobs.cpp
   ${TORCH_API_TEST_DIR}/operations.cpp
   ${TORCH_API_TEST_DIR}/nested_int.cpp
 )

test/cpp/api/JustKnobs.cpp

@@ -0,0 +1,25 @@
+#include <gtest/gtest.h>
+
+#include <torch/csrc/JustKnobs.h>
+
+TEST(JustKnobsTest, BooleanKnobReturnsDefault) {
+  // These tests apply to both internal and OSS builds because
+  // 1. In internal builds, a non-existent knob should throw an exception and
+  // hence use the default.
+  // 2. In OSS builds, we always return the default.
+  static torch::jk::BooleanKnob knob("path/to:knob_that_does_not_exist");
+  ASSERT_EQ(knob(true), true);
+  ASSERT_EQ(
+      knob(
+          false,
+          // Optionally, provide a hashval and/or a switchval
+          "hashval_for_consistent_randomization",
+          "switchval_for_overrides"),
+      false);
+}
+
+TEST(JustKnobsTest, IntegerKnobReturnsDefault) {
+  static torch::jk::IntegerKnob knob("path/to:knob_that_does_not_exist");
+  ASSERT_EQ(knob(42), 42);
+  ASSERT_EQ(knob(-100, "switchval_for_overrides"), -100);
+}

torch/csrc/JustKnobs.h

@@ -0,0 +1,86 @@
+#pragma once
+
+#include <cstdint>
+#include <string_view>
+
+/*
+    JustKnobs is a lightweight configuration system. It supports
+    re-configuring values in real-time (i.e., without a job needing
+    to restart).
+
+    Below, we have a
+    1. Meta-only implementation. Delegates to our underlying library.
+    2. Default OSS implementation. Always returns the default.
+
+    Please see the corresponding test file for basic usage and concepts.
+*/
+
+namespace torch::jk {
+
+#ifdef FBCODE_CAFFE2
+
+#include <justknobs/JustKnobs.h>
+
+class BooleanKnob {
+ public:
+  explicit BooleanKnob(std::string_view name) : real_knob_(name) {}
+
+  template <typename... Args>
+  bool operator()(bool default_value, Args&&... args) const {
+    try {
+      return real_knob_(std::forward<Args>(args)...);
+    } catch (...) {
+      return default_value;
+    }
+  }
+
+ private:
+  ::facebook::jk::BooleanKnob real_knob_;
+};
+
+class IntegerKnob {
+ public:
+  explicit IntegerKnob(std::string_view name) : real_knob_(name) {}
+
+  template <typename... Args>
+  int64_t operator()(int64_t default_value, Args&&... args) const {
+    try {
+      return real_knob_(std::forward<Args>(args)...);
+    } catch (...) {
+      return default_value;
+    }
+  }
+
+ private:
+  ::facebook::jk::IntegerKnob real_knob_;
+};
+
+#else
+
+class BooleanKnob {
+ public:
+  explicit BooleanKnob(std::string_view name) {
+    (void)name;
+  }
+
+  template <typename... Args>
+  bool operator()(bool default_value, Args&&...) const {
+    return default_value;
+  }
+};
+
+class IntegerKnob {
+ public:
+  explicit IntegerKnob(std::string_view name) {
+    (void)name;
+  }
+
+  template <typename... Args>
+  int64_t operator()(int64_t default_value, Args&&...) const {
+    return default_value;
+  }
+};
+
+#endif
+
+} // namespace torch::jk