fix FileNotFound errors when reading builtin configs from a xar, use entrypoint to register custom fb plugins for torchx

Summary: need to use parutils to load bundled builtin config files...

Reviewed By: tierex

Differential Revision: D28614050

fbshipit-source-id: a1869da90ffe07ad18fb3ac360f7215c8f568b6c

Author: Kiuk Chung

torchx/cli/cmd_run.py

@@ -10,14 +10,16 @@
 import warnings
 from dataclasses import asdict
 from os import path
+from pathlib import Path
 from pprint import pformat
-from typing import Callable, Iterable, List, Type
+from typing import Callable, Iterable, List, Optional, Type
 
 import torchx.specs as specs
 import yaml
 from torchx.cli.cmd_base import SubCommand
 from torchx.cli.conf_helpers import parse_args_children
 from torchx.runner import get_runner
+from torchx.util import entrypoints
 
 
 class UnsupportFeatureError(Exception):
@@ -94,35 +96,53 @@ def _parse_run_config(arg: str) -> specs.RunConfig:
 
 # TODO kiuk@ move read_conf_file + _builtins to the Runner once the Runner is API stable
 
-_CONFIG_DIR: str = path.join(path.dirname(__file__), "config")
+_CONFIG_DIR: Path = Path("torchx/cli/config")
 _CONFIG_EXT = ".torchx"
 
 
+def get_file_contents(conf_file: str) -> Optional[str]:
+    """
+    Reads the ``conf_file`` relative to the root of the project.
+    Returns ``None`` if ``$root/$conf_file`` does not exist.
+    Example: ``get_file("torchx/cli/config/foo.txt")``
+    """
+
+    root = path.dirname(__file__).replace(__name__.replace(".", path.sep), "")
+    abspath = path.join(root, conf_file)
+    if path.exists(abspath):
+        with open(abspath, "r") as f:
+            return f.read()
+    else:
+        return None
+
+
 def read_conf_file(conf_file: str) -> str:
-    builtin_conf = path.join(_CONFIG_DIR, conf_file)
+    builtin_conf = entrypoints.load(
+        "torchx.file",
+        "get_file_contents",
+        default=get_file_contents,
+    )(str(_CONFIG_DIR / conf_file))
 
     # user provided conf file precedes the builtin config
     # just print a warning but use the user provided one
     if path.exists(conf_file):
-        if path.exists(builtin_conf):
+        if builtin_conf:
             warnings.warn(
                 f"The provided config file: {conf_file} overlaps"
                 f" with a built-in. It is recommended that you either"
                 f" rename the config file or use abs path."
                 f" Will use: {path.abspath(conf_file)} for this run."
             )
-    else:  # conf_file does not exist fallback to builtin
-        conf_file = builtin_conf
-
-    if not path.exists(conf_file):
+        with open(conf_file, "r") as f:
+            return f.read()
+    elif builtin_conf:  # conf_file does not exist fallback to builtin
+        return builtin_conf
+    else:  # neither conf_file nor builtin exists, raise error
         raise FileNotFoundError(
-            f"{conf_file} does not exist and/or is not a builtin."
+            f"{conf_file} does not exist and is not a builtin."
             " For a list of available builtins run `torchx builtins`"
         )
 
-    with open(conf_file, "r") as f:
-        return f.read()
-
 
 def _builtins() -> List[str]:
     builtins: List[str] = []
@@ -143,7 +163,7 @@ def run(self, args: argparse.Namespace) -> None:
         num_builtins = len(builtin_configs)
         print(f"Found {num_builtins} builtin configs:")
         for i, name in enumerate(builtin_configs):
-            print(f" {i+1:2d}. {name}")
+            print(f" {i + 1:2d}. {name}")
 
 
 class CmdRun(SubCommand):

torchx/util/__init__.py

@@ -0,0 +1,41 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from importlib import metadata
+from importlib.metadata import EntryPoint
+from typing import Dict
+
+
+# pyre-ignore-all-errors[3, 2]
+def load(group: str, name: str, default=None):
+    """
+    Loads the entry point specified by
+
+    ::
+
+     [group]
+     name1 = this.is:a_function
+     -- or --
+     name2 = this.is.a.module
+
+    In case such an entry point is not found, an optional
+    default is returned. If the default is not specified
+    and the entry point is not found, then this method
+    raises an error.
+    """
+
+    entrypoints = metadata.entry_points()
+
+    if group not in entrypoints and default:
+        return default
+
+    eps: Dict[str, EntryPoint] = {ep.name: ep for ep in entrypoints[group]}
+
+    if name not in eps and default:
+        return default
+    else:
+        ep = eps[name]
+        return ep.load()

torchx/util/entrypoints.py

@@ -0,0 +1,46 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from importlib.metadata import EntryPoint
+from typing import Dict
+from unittest.mock import MagicMock, patch
+
+from torchx.util.entrypoints import load
+
+
+def foobar() -> str:
+    return "foobar"
+
+
+def barbaz() -> str:
+    return "barbaz"
+
+
+_ENTRY_POINT_TXT: str = """
+[entrypoints.test]
+foo = torchx.util.test.entrypoints_test:foobar
+"""
+
+_ENTRY_POINTS: Dict[str, EntryPoint] = {
+    # pyre-ignore[16]
+    "entrypoints.test": EntryPoint._from_text(_ENTRY_POINT_TXT)
+}
+
+_METADATA_EPS: str = "torchx.util.entrypoints.metadata.entry_points"
+
+
+class EntryPointsTest(unittest.TestCase):
+    @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
+    def test_load(self, mock_md_eps: MagicMock) -> None:
+        print(type(load("entrypoints.test", "foo")))
+        self.assertEqual("foobar", load("entrypoints.test", "foo")())
+
+    @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
+    def test_load_with_default(self, mock_md_eps: MagicMock) -> None:
+        self.assertEqual("barbaz", load("entrypoints.test", "missing", barbaz)())
+        self.assertEqual("barbaz", load("entrypoints.missing", "foo", barbaz)())
+        self.assertEqual("barbaz", load("entrypoints.missing", "missing", barbaz)())