(torchx/config) rename TORCHX_CONFIG to TORCHXCONFIG, improve unittest and documentation for config.find_configs (#507)

Summary:
Pull Request resolved: #507

Renames env var `TORCHX_CONFIG` to `TORCHXCONFIG` to keep it consistent with the naming pattern of `.torchxconfig`. Improves documentation for `config.find_config` and unittests.

Fixes the above that was introduced in D36692686 (5a88b4d) that landed before I got a chance to review it.

Reviewed By: d4l3k

Differential Revision: D36849299

fbshipit-source-id: edd6db88ff20535b50741d509a1d2dfe57b974dd

Author: Kiuk Chung

torchx/runner/config.py

@@ -152,8 +152,8 @@ def my_component(a: int) -> specs.AppDef:
 
 
 CONFIG_FILE = ".torchxconfig"
-ENV_TORCHX_CONFIG = "TORCHX_CONFIG"
 CONFIG_PREFIX_DELIM = ":"
+ENV_TORCHXCONFIG = "TORCHXCONFIG"
 
 _NONE = "None"
 
@@ -431,20 +431,29 @@ def get_config(
 
 def find_configs(dirs: Optional[Iterable[str]] = None) -> List[str]:
     """
-    find_configs returns all the .torchxconfig files it finds:
-        - if environment variable TORCHX_CONFIG is specified, then
-          it is expected to exists and will be loaded as the only config
-        - otherwise, directories are checked to look for .torchxconfig and loaded. If
-          directories is empty it checks the local directory.
+    Finds and returns the filepath to ``.torchxconfig`` files based
+    on the following logic:
+
+    1. If the environment variable ``TORCHXCONFIG`` exists, then its value
+       is returned in a single-element list and the directories specified through
+       the ``dirs`` parameter is NOT searched.
+    2. Otherwise, a ``.torchxconfig`` file is looked for in ``dirs`` and
+       the filepaths to existing config files are returned. If ``dirs`` is
+       not specified or is empty then, this method looks for a ``.torchxconfig`` file
+       in CWD (current working dir) and returns the filepath to it if one exists.
+
     """
 
-    config_files = []
-    config = os.getenv(ENV_TORCHX_CONFIG, "")
-    if len(config) > 0:
+    config = os.getenv(ENV_TORCHXCONFIG)
+    if config is not None:
         configfile = Path(config)
-        assert configfile.exists(), f"{str(configfile)} expected but not found"
-        config_files.append(str(configfile))
+        if not configfile.is_file():
+            raise FileNotFoundError(
+                f"`{ENV_TORCHXCONFIG}={config}` does not exist or is not a file."
+            )
+        return [str(configfile)]
     else:
+        config_files = []
         if not dirs:
             dirs = [str(Path.cwd())]
         for d in dirs:

torchx/runner/test/config_test.py

@@ -18,6 +18,8 @@
 from torchx.runner.config import (
     apply,
     dump,
+    ENV_TORCHXCONFIG,
+    find_configs,
     get_config,
     get_configs,
     load,
@@ -259,6 +261,44 @@ def test_get_configs(self) -> None:
             ),
         )
 
+    def test_find_configs(self) -> None:
+        config_dir = Path(self.test_dir)
+        cwd_dir = config_dir / "cwd"
+        custom_dir = config_dir / "custom"
+
+        cwd_dir.mkdir()
+        custom_dir.mkdir()
+
+        base_config = config_dir / ".torchxconfig"
+        cwd_config = cwd_dir / ".torchxconfig"
+        override_config = custom_dir / ".torchxconfig"
+
+        base_config.touch()
+        cwd_config.touch()
+        override_config.touch()
+
+        # should find configs in the specified dirs
+        configs = find_configs(dirs=[str(config_dir)])
+        self.assertEqual([str(base_config)], configs)
+
+        # should find config in cwd if no dirs is specified
+        with patch(PATH_CWD, return_value=str(cwd_dir)):
+            configs = find_configs()
+            self.assertEqual([str(cwd_config)], configs)
+
+        # if TORCHXCONFIG env var exists then should just return the config specified
+        with patch.dict(os.environ, {ENV_TORCHXCONFIG: str(override_config)}):
+            configs = find_configs(dirs=[str(config_dir)])
+            self.assertEqual([str(override_config)], configs)
+
+        # if TORCHXCONFIG points to a non-existing file, then assert exception
+        with patch.dict(
+            os.environ,
+            {ENV_TORCHXCONFIG: str(config_dir / ".torchxconfig_nonexistent")},
+        ):
+            with self.assertRaises(FileNotFoundError):
+                find_configs(dirs=[str(config_dir)])
+
     def test_get_config(self) -> None:
         configdir0 = Path(self.test_dir) / "test_load_component_defaults" / "0"
         configdir1 = Path(self.test_dir) / "test_load_component_defaults" / "1"
@@ -286,6 +326,20 @@ def test_get_config(self) -> None:
             get_config(prefix="badprefix", name="dist.ddp", key="j", dirs=dirs),
         )
 
+        # check that if TORCHXCONFIG is set then only that config is loaded
+        override_config = Path(self.test_dir) / ".torchxconfig_custom"
+        override_config_contents = """
+[component:dist.ddp]
+image = foobar_custom
+        """
+        self._write(str(override_config), override_config_contents)
+
+        with patch.dict(os.environ, {ENV_TORCHXCONFIG: str(override_config)}):
+            self.assertDictEqual(
+                {"image": "foobar_custom"},
+                get_configs(prefix="component", name="dist.ddp", dirs=dirs),
+            )
+
     def test_load(self) -> None:
         cfg = {}
         load(scheduler="local_cwd", f=StringIO(_CONFIG), cfg=cfg)
@@ -328,33 +382,6 @@ def test_apply_dirs(self, _) -> None:
         self.assertEqual(100, cfg.get("i"))
         self.assertEqual(1.2, cfg.get("f"))
 
-    @patch(
-        TORCHX_GET_SCHEDULERS,
-        return_value={"test": TestScheduler()},
-    )
-    def test_apply_from_annotated_config(self, _) -> None:
-        def mock_getenv(x: str, y: Optional[str] = None) -> Optional[str]:
-            if x != "TORCHX_CONFIG":
-                return os.environ.get(x, y)
-            else:
-                return f"{self.test_dir}/another_torchx_config"
-
-        with patch(PATH_CWD, return_value=Path(self.test_dir)):
-            with patch(
-                "torchx.runner.config.os.getenv",
-                mock_getenv,
-            ):
-                cfg: Dict[str, CfgVal] = {"s": "runtime_value"}
-                apply(
-                    scheduler="test",
-                    cfg=cfg,
-                    # these dirs will be ignored
-                    dirs=[str(Path(self.test_dir) / "home" / "bob"), self.test_dir],
-                )
-                self.assertEqual("runtime_value", cfg.get("s"))
-                self.assertEqual(200, cfg.get("i"))
-                self.assertEqual(None, cfg.get("f"))
-
     def test_dump_invalid_scheduler(self) -> None:
         with self.assertRaises(ValueError):
             dump(f=StringIO(), schedulers=["does-not-exist"])