From 1d6ca50c5bc66ae4e80241c5c1d9aa222fcf4557 Mon Sep 17 00:00:00 2001
From: "Colin L. Rice" <clr@fb.com>
Date: Wed, 20 Nov 2024 13:59:34 -0700
Subject: [PATCH] config: Throw if justknobs value is not a boolean (#139488)

This helps avoid an issue, where someone uses a mutable type that
justknobs does not support within the code. And then it gets overriden
to a different type
Pull Request resolved: https://github.com/pytorch/pytorch/pull/139488
Approved by: https://github.com/ezyang
---
 test/test_utils_config_module.py | 9 ++++++++-
 torch/utils/_config_module.py    | 4 ++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/test/test_utils_config_module.py b/test/test_utils_config_module.py
index 50e391225b8..d17ef155aef 100644
--- a/test/test_utils_config_module.py
+++ b/test/test_utils_config_module.py
@@ -11,7 +11,7 @@ from typing import Optional
 
 from torch.testing._internal import fake_config_module as config
 from torch.testing._internal.common_utils import run_tests, TestCase
-from torch.utils._config_module import _UNSET_SENTINEL
+from torch.utils._config_module import _UNSET_SENTINEL, Config
 
 
 class TestConfigModule(TestCase):
@@ -329,6 +329,13 @@ torch.testing._internal.fake_config_module._save_config_ignore = ['e_ignored']""
             self.assertFalse(config.e_bool)
         self.assertTrue(config.e_bool)
 
+    def test_bad_jk_type(self):
+        with self.assertRaises(
+            AssertionError,
+            msg="AssertionError: justknobs only support booleans, thisisnotvalid is not a boolean",
+        ):
+            Config(default="bad", justknob="fake_knob")
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/torch/utils/_config_module.py b/torch/utils/_config_module.py
index 269ca4e59f8..ee7197afda2 100644
--- a/torch/utils/_config_module.py
+++ b/torch/utils/_config_module.py
@@ -71,6 +71,10 @@ class Config:
         self.env_name_default = env_name_default
         self.env_name_force = env_name_force
         self.value_type = value_type
+        if self.justknob is not None:
+            assert isinstance(
+                self.default, bool
+            ), f"justknobs only support booleans, {self.default} is not a boolean"
 
 
 # Types saved/loaded in configs