config: Support str env variables (pytorch#145980)

Summary:
This allows us to use environment variables to set string values. We've added
tests for the specific functionality implemented here. Note that we already
accidentally started setting up configs to use this, so we're just adding the
feature.

Additionally, we're not fully validating the underlying type when we set the
value (and in general, it's more difficult than we would like to do this). Let
me know if people feel strongly, and we can add a PR to do this.
Pull Request resolved: pytorch#145980
Approved by: https://github.com/yushangdi, https://github.com/oulgen

Author: c00w

test/test_utils_config_module.py

@@ -6,6 +6,8 @@
 
 os.environ["ENV_TRUE"] = "1"
 os.environ["ENV_FALSE"] = "0"
+os.environ["ENV_STR"] = "1234"
+os.environ["ENV_STR_EMPTY"] = ""
 
 from typing import Optional
 
@@ -100,6 +102,12 @@ def test_env_name_semantics(self):
         config.e_env_force = False
         self.assertTrue(config.e_env_force)
 
+    def test_env_name_string_semantics(self):
+        self.assertEqual(config.e_env_default_str, "1234")
+        self.assertEqual(config.e_env_default_str_empty, "")
+        config.e_env_default_str = "override"
+        self.assertEqual(config.e_env_default_str, "override")
+
     def test_multi_env(self):
         self.assertTrue(config2.e_env_default_multi)
         self.assertTrue(config2.e_env_force_multi)
@@ -129,6 +137,8 @@ def test_save_config(self):
                 "e_jk_false": False,
                 "e_env_default": True,
                 "e_env_default_FALSE": False,
+                "e_env_default_str": "1234",
+                "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
             },
@@ -161,6 +171,8 @@ def test_save_config_portable(self):
                 "e_jk_false": False,
                 "e_env_default": True,
                 "e_env_default_FALSE": False,
+                "e_env_default_str": "1234",
+                "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
             },
@@ -180,6 +192,8 @@ def test_codegen_config(self):
             """torch.testing._internal.fake_config_module.e_bool = False
 torch.testing._internal.fake_config_module.e_env_default = True
 torch.testing._internal.fake_config_module.e_env_default_FALSE = False
+torch.testing._internal.fake_config_module.e_env_default_str = '1234'
+torch.testing._internal.fake_config_module.e_env_default_str_empty = ''
 torch.testing._internal.fake_config_module.e_env_force = True
 torch.testing._internal.fake_config_module._save_config_ignore = ['e_ignored']""",
         )
@@ -202,7 +216,7 @@ def test_codegen_config_function(self):
         )
 
     def test_get_hash(self):
-        hash_value = b"\xf2C\xdbo\x99qq\x12\x11\xf7\xb4\xeewVpZ"
+        hash_value = b"\x87\xf7\xc6\x1di\x7f\x96-\x85\xdc\x04\xd5\xd0\xf6\x1c\x87"
         self.assertEqual(
             config.get_hash(),
             hash_value,
@@ -259,6 +273,8 @@ def test_dict_copy_semantics(self):
                 "e_jk_false": False,
                 "e_env_default": True,
                 "e_env_default_FALSE": False,
+                "e_env_default_str": "1234",
+                "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
             },
@@ -288,6 +304,8 @@ def test_dict_copy_semantics(self):
                 "e_jk_false": False,
                 "e_env_default": True,
                 "e_env_default_FALSE": False,
+                "e_env_default_str": "1234",
+                "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
             },
@@ -317,6 +335,8 @@ def test_dict_copy_semantics(self):
                 "e_jk_false": False,
                 "e_env_default": True,
                 "e_env_default_FALSE": False,
+                "e_env_default_str": "1234",
+                "e_env_default_str_empty": "",
                 "e_env_force": True,
                 "e_optional": True,
             },

torch/testing/_internal/fake_config_module.py

@@ -24,6 +24,10 @@
 e_jk_false: bool = Config(justknob="does_not_exist", default=False)
 e_env_default: bool = Config(env_name_default="ENV_TRUE", default=False)
 e_env_default_FALSE: bool = Config(env_name_default="ENV_FALSE", default=True)
+e_env_default_str: bool = Config(env_name_default="ENV_STR", default="default")
+e_env_default_str_empty: bool = Config(
+    env_name_default="ENV_STR_EMPTY", default="default"
+)
 e_env_force: bool = Config(env_name_force="ENV_TRUE", default=False)
 e_aliased_bool: bool = Config(
     alias="torch.testing._internal.fake_config_module2.e_aliasing_bool"

torch/utils/_config_module.py

@@ -159,13 +159,13 @@ def Config(
         )
 
 
-def _read_env_variable(name: str) -> Optional[bool]:
+def _read_env_variable(name: str) -> Optional[Union[bool, str]]:
     value = os.environ.get(name)
     if value == "1":
         return True
     if value == "0":
         return False
-    return None
+    return value
 
 
 def install_config_module(module: ModuleType) -> None: