Refactor components discovery logic (#88)

Summary:
Pull Request resolved: #88

Move component search registration into a separate module and share it across runner and torch_cli
Introduce new format for ``[torchx.components]`` :

    foo = my_project.some_module.bar

Where ``my_project.some_module.bar`` is a python module. The components will be recursively found in the dir for ``my_project.some_module.bar.__file__`` file.

Reviewed By: kiukchung

Differential Revision: D29202878

fbshipit-source-id: af07a9bf4d3cf6038fc6d0f1c704c025176a2da9

Author: aivanou

docs/source/configure.rst

@@ -107,16 +107,36 @@ when they run
 
  $ torchx builtins
 
-Custom components can be registered via the following endpoint:
+Custom components can be registered via the following modification of the ``entry_points.txt``:
 
 ::
 
  [torchx.components]
- custom_component = my_module.components:my_component
+ foo = my_project.bar
 
+The line above registers a group ``foo`` that is associated with the module ``my_project.bar``.
+Torchx will recursively traverse lowest level dir associated with the ``my_project.bar`` and will find
+all defined components.
 
-Custom components can be executed in the following manner:
+.. note:: If there are two registry entries, e.g. ``foo = my_project.bar`` and ``test = my_project``
+          there will be two sets of overlapping components with different aliases.
+
+
+After registration, torchx cli will display registered components via:
+
+.. code-block:: shell-session
+
+ $ torchx builtins
+
+If ``my_project.bar`` had the following directory structure:
+
+::
+
+ $PROJECT_ROOT/my_project/bar/
+     |- baz.py
+
+And `baz.py` defines a component (function) called `trainer`. Then the component can be run as a job in the following manner:
 
 .. code-block:: shell-session
 
- $ torchx run --scheduler local --scheduler_args image_fetcher=...,root_dir=/tmp custom_component -- --name "test app"
+ $ torchx run foo.baz.trainer -- --name "test app"

torchx/cli/cmd_run.py

@@ -5,20 +5,15 @@
 # LICENSE file in the root directory of this source tree.
 
 import argparse
-import glob
-import importlib
-import os
-from dataclasses import dataclass
-from inspect import getmembers, isfunction
-from typing import Dict, List, Optional, Union
+from dataclasses import asdict
+from pprint import pformat
+from typing import Dict, List, Union, cast
 
 import torchx.specs as specs
 from pyre_extensions import none_throws
 from torchx.cli.cmd_base import SubCommand
 from torchx.runner import get_runner
-from torchx.specs.file_linter import get_fn_docstring, validate
-from torchx.util import entrypoints
-from torchx.util.io import COMPONENTS_DIR, get_abspath, read_conf_file
+from torchx.specs.finder import get_components, _Component
 from torchx.util.types import to_dict
 
 
@@ -38,96 +33,19 @@ def _parse_run_config(arg: str) -> specs.RunConfig:
     return conf
 
 
-def _to_module(filepath: str) -> str:
-    path, _ = os.path.splitext(filepath)
-    return path.replace(os.path.sep, ".")
-
-
-def _get_builtin_description(filepath: str, function_name: str) -> Optional[str]:
-    source = read_conf_file(filepath)
-    if len(validate(source, torchx_function=function_name)) != 0:
-        return None
-
-    func_definition, _ = none_throws(get_fn_docstring(source, function_name))
-    return func_definition
-
-
-@dataclass
-class BuiltinComponent:
-    definition: str
-    description: str
-
-
-def _get_component_definition(module: str, function_name: str) -> str:
-    if module.startswith("torchx.components"):
-        module = module.split("torchx.components.")[1]
-    return f"{module}.{function_name}"
-
-
-def _to_relative(filepath: str) -> str:
-    if os.path.isabs(filepath):
-        # make path torchx/components/$suffix out of the abs
-        rel_path = filepath.split(str(COMPONENTS_DIR))[1]
-        return f"{str(COMPONENTS_DIR)}{rel_path}"
-    else:
-        return os.path.join(str(COMPONENTS_DIR), filepath)
-
-
-def _get_components_from_file(filepath: str) -> List[BuiltinComponent]:
-    components_path = _to_relative(filepath)
-    components_module_path = _to_module(components_path)
-    module = importlib.import_module(components_module_path)
-    functions = getmembers(module, isfunction)
-    buitin_functions = []
-    for function_name, _ in functions:
-        # Ignore private functions.
-        if function_name.startswith("_"):
-            continue
-        component_desc = _get_builtin_description(filepath, function_name)
-        if component_desc:
-            definition = _get_component_definition(
-                components_module_path, function_name
-            )
-            builtin_component = BuiltinComponent(
-                definition=definition,
-                description=component_desc,
-            )
-            buitin_functions.append(builtin_component)
-    return buitin_functions
-
-
-def _allowed_path(path: str) -> bool:
-    filename = os.path.basename(path)
-    if filename.startswith("_"):
-        return False
-    return True
-
-
-def _builtins() -> List[BuiltinComponent]:
-    components_dir = entrypoints.load(
-        "torchx.file", "get_dir_path", default=get_abspath
-    )(COMPONENTS_DIR)
-
-    builtins: List[BuiltinComponent] = []
-    search_pattern = os.path.join(components_dir, "**", "*.py")
-    for filepath in glob.glob(search_pattern, recursive=True):
-        if not _allowed_path(filepath):
-            continue
-        components = _get_components_from_file(filepath)
-        builtins += components
-    return builtins
-
-
 class CmdBuiltins(SubCommand):
     def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
-        pass  # no arguments
+        pass
+
+    def _builtins(self) -> Dict[str, _Component]:
+        return get_components()
 
     def run(self, args: argparse.Namespace) -> None:
-        builtin_configs = _builtins()
-        num_builtins = len(builtin_configs)
+        builtin_components = self._builtins()
+        num_builtins = len(builtin_components)
         print(f"Found {num_builtins} builtin configs:")
-        for i, component in enumerate(builtin_configs):
-            print(f" {i + 1:2d}. {component.definition} - {component.description}")
+        for i, component in enumerate(builtin_components.values()):
+            print(f" {i + 1:2d}. {component.name}")
 
 
 class CmdRun(SubCommand):
@@ -172,15 +90,23 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
     def run(self, args: argparse.Namespace) -> None:
         # TODO: T91790598 - remove the if condition when all apps are migrated to pure python
         runner = get_runner()
-        app_handle = runner.run_from_path(
+        result = runner.run_component(
             args.conf_file,
             args.conf_args,
             args.scheduler,
             args.scheduler_args,
             dryrun=args.dryrun,
         )
 
-        if not args.dryrun:
+        if args.dryrun:
+            app_dryrun_info = cast(specs.AppDryRunInfo, result)
+            print("=== APPLICATION ===")
+            print(pformat(asdict(app_dryrun_info._app), indent=2, width=80))
+
+            print("=== SCHEDULER REQUEST ===")
+            print(app_dryrun_info)
+        else:
+            app_handle = cast(specs.AppHandle, result)
             if args.scheduler == "local":
                 runner.wait(app_handle)
             else:

torchx/cli/test/cmd_run_test.py

@@ -15,7 +15,7 @@
 from typing import Generator
 from unittest.mock import MagicMock, patch
 
-from torchx.cli.cmd_run import CmdBuiltins, CmdRun, _builtins, _parse_run_config
+from torchx.cli.cmd_run import CmdBuiltins, CmdRun, _parse_run_config
 
 
 @contextmanager
@@ -109,7 +109,8 @@ def test_run(self) -> None:
         cmd_builtins.run(args)
 
     def test_builtins(self) -> None:
-        builtins = _builtins()
+        cmd_builtins = CmdBuiltins()
+        builtins = cmd_builtins._builtins()
         # make sure there's at least one
         # there will always be one (example.torchx)
         self.assertTrue(len(builtins) > 0)

torchx/components/__init__.py

@@ -20,7 +20,7 @@
 
   # using via sdk
   from torchx.runner import get_runner
-  get_runner().run_from_path("distributed.ddp", app_args=[], scheduler="local", ...)
+  get_runner().run_component("distributed.ddp", app_args=[], scheduler="local", ...)
 
   # using via torchx-cli
 

torchx/runner/api.py

@@ -6,13 +6,11 @@
 # LICENSE file in the root directory of this source tree.
 
 import getpass
-import importlib
 import json
+import logging
 import time
-from dataclasses import asdict
 from datetime import datetime
-from pprint import pformat
-from typing import Any, Dict, Iterable, List, Optional, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from pyre_extensions import none_throws
 from torchx.runner.events import log_event
@@ -28,16 +26,22 @@
     UnknownAppException,
     from_file,
     from_function,
-    from_module,
     make_app_handle,
     parse_app_handle,
     runopts,
 )
-from torchx.util import entrypoints
+from torchx.specs.finder import get_component
+
+logger: logging.Logger = logging.getLogger(__name__)
 
 
 NONE: str = "<NONE>"
 
+# Unique identifier of the component, can be either
+# component name, e.g.: utils.echo or path to component function:
+# /tmp/foobar.py:component
+ComponentId = str
+
 
 class Runner:
     """
@@ -65,81 +69,75 @@ def __init__(
         self._wait_interval = wait_interval
         self._apps: Dict[AppHandle, AppDef] = {}
 
-    def run_from_path(
+    def run_component(
         self,
-        component_path: str,
+        component_name: ComponentId,
         app_args: List[str],
         scheduler: SchedulerBackend = "default",
         cfg: Optional[RunConfig] = None,
         dryrun: bool = False,
-    ) -> AppHandle:
+        # pyre-ignore[24]: Allow generic AppDryRunInfo
+    ) -> Union[AppHandle, AppDryRunInfo]:
         """Resolves and runs the application in the specified mode.
 
-        Retrieves application based on ``component_path`` and runs it in the specified mode.
+        Retrieves application based on ``component_name`` and runs it in the specified mode.
 
-        The ``component_path`` has the following resolution order(from high-pri to low-pri):
+        The ``component_name`` has the following resolution order(from high-pri to low-pri):
             * User-registered components. Users can register components via
                 https://packaging.python.org/specifications/entry-points/. Method looks for
                 entrypoints in the group ``torchx.components``.
-            * File-based components in format: ``$FILE_PATH:FUNCTION_NAME``. Both relative and
-                absolute paths supported.
             * Builtin components relative to `torchx.components`. The path to the component should
                 be module name relative to `torchx.components` and function name in a format:
                 ``$module.$function``.
+            * File-based components in format: ``$FILE_PATH:FUNCTION_NAME``. Both relative and
+                absolute paths supported.
 
         Usage:
 
         ::
 
-         runner.run_from_path("distributed.ddp", ...) - will be resolved to
+         runner.run_component("distributed.ddp", ...) - will be resolved to
             ``torchx.components.distributed`` module and ``ddp`` function.
 
 
-         runner.run_from_path("~/home/components.py:my_component", ...) - will be resolved to
+         runner.run_component("~/home/components.py:my_component", ...) - will be resolved to
             ``~/home/components.py`` file and ``my_component`` function.
 
+        Args:
+            component_name: The name of the component to lookup
+            app_args: Arguments of the component function
+            scheduler: Scheduler to execute component
+            cfg: Scheduler run configuration
+            dryrun: If True, the will return ``torchx.specs.AppDryRunInfo``
+
 
         Returns:
             An application handle that is used to call other action APIs on the app, or ``<NONE>``
             if it dryrun specified.
 
         Raises:
             AppNotReRunnableException: if the session/scheduler does not support re-running attached apps
-            ValueError: if the ``component_path`` is failed to resolve.
+            ValueError: if the ``component_name`` is failed to resolve.
         """
-        app_fn = entrypoints.load("torchx.components", component_path, default=NONE)
-        if app_fn != NONE:
-            app = from_function(app_fn, app_args)
-        elif ":" in component_path:
-            file_path, function_name = component_path.split(":")
+        component_def = get_component(component_name)
+        if component_def:
+            app = from_function(component_def.fn, app_args)
+        elif ":" in component_name:
+            file_path, function_name = component_name.split(":")
             if not function_name:
                 raise ValueError(
-                    f"Incorrect path: {component_path}. Missing function name after `:`"
+                    f"Incorrect path: {component_name}. Missing function name after `:`"
                 )
             app = from_file(file_path, function_name, app_args)
         else:
-            function_name = component_path.split(".")[-1]
-            component_module = component_path[
-                0 : len(component_path) - len(function_name) - 1
-            ]
-            full_module = f"torchx.components.{component_module}"
-            try:
-                app_module = importlib.import_module(full_module)
-            except ModuleNotFoundError:
-                raise ValueError(
-                    f"Incorrect component path: {component_path}. Valid values: \n"
-                    "* $file_path:function_name, Example: ~/home/file.py:my_component \n"
-                    "* Module and function name to the builtin components, e.g. : distributed.ddp"
-                )
-            app = from_module(app_module, function_name, app_args)
+            raise ValueError(
+                f"Component `{component_name}` not found. Please make sure it is one of the "
+                "builtins: `torchx builtins`. Or registered via `[torchx.components]` "
+                "entry point (see: https://pytorch.org/torchx/latest/configure.html)"
+            )
 
         if dryrun:
-            print("=== APPLICATION ===")
-            print(pformat(asdict(app), indent=2, width=80))
-
-            print("=== SCHEDULER REQUEST ===")
-            print(self.dryrun(app, scheduler, cfg))
-            return "<NONE>"
+            return self.dryrun(app, scheduler, cfg)
         else:
             return self.run(app, scheduler, cfg)