(torchx/components) Fix entrypoint loading to deal with deferred loading of modules to enable component registration to work properly

Author: kiukchung

docs/source/advanced.rst

@@ -156,7 +156,7 @@ resource can then be used in the following manner:
 
 Registering Custom Components
 -------------------------------
-It is possible to author and register a custom set of components with the
+You can author and register a custom set of components with the
 ``torchx`` CLI as builtins to the CLI. This makes it possible to customize
 a set of components most relevant to your team or organization and support
 it as a CLI ``builtin``. This way users will see your custom components
@@ -166,7 +166,63 @@ when they run
 
  $ torchx builtins
 
-Custom components can be registered via the following modification of the ``entry_points``:
+Custom components can be registered via ``[torchx.components]`` entrypoints.
+If ``my_project.bar`` had the following directory structure:
+
+::
+
+ $PROJECT_ROOT/my_project/bar/
+     |- baz.py
+
+And ``baz.py`` had a single component (function) called ``trainer``:
+
+::
+
+ # baz.py
+ import torchx.specs as specs
+
+ def trainer(...) -> specs.AppDef: ...
+
+
+And the entrypoints were added as:
+
+.. testcode::
+
+   # setup.py
+   ...
+   entry_points={
+       "torchx.components": [
+           "foo = my_project.bar",
+       ],
+   }
+
+TorchX will search the module ``my_project.bar`` for all defined components and group the found
+components under the ``foo.*`` prefix. In this case, the component ``my_project.bar.baz.trainer``
+would be registered with the name ``foo.baz.trainer``.
+
+.. note::
+    Only python packages (those directories with an ``__init__.py`` file)
+    are searched for and TorchX makes no attempt to recurse into namespace packages
+    (directories without a ``__init__.py`` file).
+    However you may register a top level namespace package.
+
+``torchx`` CLI will display registered components via:
+
+.. code-block:: shell-session
+
+ $ torchx builtins
+ Found 1 builtin components:
+ 1. foo.baz.trainer
+
+The custom component can then be used as:
+
+.. code-block:: shell-session
+
+ $ torchx run foo.baz.trainer -- --name "test app"
+
+
+When you register your own components, TorchX will not include its own builtins. To add TorchX's
+builtin components you must specify another entry as:
 
 
 .. testcode::
@@ -176,32 +232,60 @@ Custom components can be registered via the following modification of the ``entr
    entry_points={
        "torchx.components": [
            "foo = my_project.bar",
+           "torchx = torchx.components",
        ],
    }
 
-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.
+This will add back the TorchX builtins but with a ``torchx.*`` component name prefix (e.g. ``torchx.dist.ddp``
+versus the default ``dist.ddp``).
+
+If there are two registry entries pointing to the same component, for instance
 
-.. 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.
+.. testcode::
 
+   # setup.py
+   ...
+   entry_points={
+       "torchx.components": [
+           "foo = my_project.bar",
+           "test = my_project",
+       ],
+   }
 
-After registration, torchx cli will display registered components via:
+
+There will be two sets of overlapping components for those components in ``my_project.bar`` with different
+prefix aliases: ``foo.*`` and ``test.bar.*``. Concretely,
 
 .. code-block:: shell-session
 
  $ torchx builtins
+ Found 2 builtin components:
+ 1. foo.baz.trainer
+ 2. test.bar.baz.trainer
 
-If ``my_project.bar`` had the following directory structure:
+To omit groupings and make the component names shorter, use underscore (e.g ``_`` or ``_0``, ``_1``, etc).
+For example:
 
-::
+.. testcode::
 
- $PROJECT_ROOT/my_project/bar/
-     |- baz.py
+   # setup.py
+   ...
+   entry_points={
+       "torchx.components": [
+           "_0 = my_project.bar",
+           "_1 = torchx.components",
+       ],
+   }
 
-And ``baz.py`` defines a component (function) called ``trainer``. Then the component can be run as a job in the following manner:
+This has the effect of exposing the trainer component as ``baz.trainer`` (as opposed to ``foo.baz.trainer``)
+and adds back the builtin components as in the vanilla installation of torchx, without the ``torchx.*`` prefix.
 
 .. code-block:: shell-session
 
- $ torchx run foo.baz.trainer -- --name "test app"
+ $ torchx builtins
+ Found 11 builtin components:
+ 1. baz.trainer
+ 2. dist.ddp
+ 3. utils.python
+ 4. ... <more builtins from torchx.components.* ...>
+

torchx/specs/finder.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and 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.

torchx/specs/test/components/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and 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 torchx
+from torchx import specs
+
+
+def comp_a() -> specs.AppDef:
+    return specs.AppDef(
+        name="foo.comp_a",
+        roles=[
+            specs.Role(
+                name="foo.comp_a",
+                image=torchx.IMAGE,
+                entrypoint="echo",
+                args=["hello world"],
+            )
+        ],
+    )

torchx/specs/test/components/a/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and 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.

torchx/specs/test/components/a/b/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and 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 torchx
+from torchx import specs
+
+
+def d() -> specs.AppDef:
+    return specs.AppDef(
+        name="foo.b.c.d",
+        roles=[
+            specs.Role(
+                name="foo.b.c.d",
+                image=torchx.IMAGE,
+                entrypoint="echo",
+                args=["hello world"],
+            )
+        ],
+    )

torchx/specs/test/components/a/b/c.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and 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.

torchx/specs/test/components/c/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and 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 torchx
+from torchx import specs
+
+
+def e() -> specs.AppDef:
+    return specs.AppDef(
+        name="bar.e",
+        roles=[
+            specs.Role(
+                name="bar.e",
+                image=torchx.IMAGE,
+                entrypoint="echo",
+                args=["hello world"],
+            )
+        ],
+    )

torchx/specs/test/components/c/d.py

@@ -7,13 +7,14 @@
 
 import os
 import shutil
-import sys
 import tempfile
 import unittest
 from pathlib import Path
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import torchx.specs.finder as finder
+
+from importlib_metadata import EntryPoints
 from torchx.runner import get_runner
 from torchx.runtime.tracking import FsspecResultTracker
 from torchx.specs.api import AppDef, AppState, Role
@@ -26,8 +27,11 @@
     get_components,
     ModuleComponentsFinder,
 )
+from torchx.util.test.entrypoints_test import EntryPoint_from_text
 from torchx.util.types import none_throws
 
+_METADATA_EPS: str = "torchx.util.entrypoints.metadata.entry_points"
+
 
 def _test_component(name: str, role_name: str = "worker") -> AppDef:
     """
@@ -58,92 +62,126 @@ def invalid_component(name, role_name: str = "worker") -> AppDef:
     )
 
 
-class DirComponentsFinderTest(unittest.TestCase):
-    def test_get_components(self) -> None:
-        components = _load_components()
-        self.assertTrue(len(components) > 1)
-        component = components["utils.echo"]
-        self.assertEqual("utils.echo", component.name)
-        self.assertEqual(
-            "Echos a message to stdout (calls echo)", component.description
+class FinderTest(unittest.TestCase):
+    _ENTRY_POINTS: EntryPoints = EntryPoints(
+        EntryPoint_from_text(
+            """
+[torchx.components]
+_ = torchx.specs.test.finder_test
+        """
         )
-        self.assertEqual("echo", component.fn_name)
-        self.assertIsNotNone(component.fn)
+    )
+
+    def tearDown(self) -> None:
+        # clear the globals since find_component() has side-effects
+        # and we load a bunch of mocks for components in the tests below
+        finder._components = None
+
+    def test_module_relname(self) -> None:
+        import torchx.specs.test.components as c
+        import torchx.specs.test.components.a as ca
+
+        self.assertEqual("", finder.module_relname(c, relative_to=c))
+        self.assertEqual("a", finder.module_relname(ca, relative_to=c))
+        with self.assertRaises(ValueError):
+            finder.module_relname(c, relative_to=ca)
 
     def test_get_component_by_name(self) -> None:
         component = none_throws(get_component("utils.echo"))
         self.assertEqual("utils.echo", component.name)
         self.assertEqual("echo", component.fn_name)
         self.assertIsNotNone(component.fn)
 
-    def test_get_invalid_component_by_name(self) -> None:
-        test_torchx_group = {"foobar": sys.modules[__name__]}
-        finder._components = None
-        with patch("torchx.specs.finder.entrypoints") as entrypoints_mock:
-            entrypoints_mock.load_group.return_value = test_torchx_group
-            with self.assertRaises(ComponentValidationException):
-                get_component("foobar.finder_test.invalid_component")
+    @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
+    def test_get_invalid_component_by_name(self, _: MagicMock) -> None:
+        with self.assertRaises(ComponentValidationException):
+            get_component("invalid_component")
 
-    def test_get_unknown_component_by_name(self) -> None:
-        test_torchx_group = {"foobar": sys.modules[__name__]}
-        finder._components = None
-        with patch("torchx.specs.finder.entrypoints") as entrypoints_mock:
-            entrypoints_mock.load_group.return_value = test_torchx_group
-            with self.assertRaises(ComponentNotFoundException):
-                get_component("foobar.finder_test.unknown_component")
-
-    def test_get_invalid_component(self) -> None:
-        test_torchx_group = {"foobar": sys.modules[__name__]}
-        with patch("torchx.specs.finder.entrypoints") as entrypoints_mock:
-            entrypoints_mock.load_group.return_value = test_torchx_group
-            components = _load_components()
-        foobar_component = components["foobar.finder_test.invalid_component"]
+    @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
+    def test_get_unknown_component_by_name(self, _: MagicMock) -> None:
+        with self.assertRaises(ComponentNotFoundException):
+            get_component("unknown_component")
+
+    @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
+    def test_get_invalid_component(self, _: MagicMock) -> None:
+        components = _load_components()
+        foobar_component = components["invalid_component"]
         self.assertEqual(1, len(foobar_component.validation_errors))
 
-    def test_get_entrypoints_components(self) -> None:
-        test_torchx_group = {"foobar": sys.modules[__name__]}
-        with patch("torchx.specs.finder.entrypoints") as entrypoints_mock:
-            entrypoints_mock.load_group.return_value = test_torchx_group
-            components = _load_components()
-        foobar_component = components["foobar.finder_test._test_component"]
+    @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
+    def test_get_entrypoints_components(self, _: MagicMock) -> None:
+        components = _load_components()
+        foobar_component = components["_test_component"]
         self.assertEqual(_test_component, foobar_component.fn)
         self.assertEqual("_test_component", foobar_component.fn_name)
-        self.assertEqual("foobar.finder_test._test_component", foobar_component.name)
+        self.assertEqual("_test_component", foobar_component.name)
         self.assertEqual("Test component", foobar_component.description)
 
-    def test_get_base_module_name(self) -> None:
-        finder = ModuleComponentsFinder(sys.modules[__name__], "")
-        expected_name = "torchx.specs.test"
-        actual_name = finder._get_base_module_name(sys.modules[__name__])
-        self.assertEqual(expected_name, actual_name)
-
-    def test_get_base_module_name_for_init_module(self) -> None:
-        finder = ModuleComponentsFinder("", "")
-        expected_name = "torchx.specs"
-        actual_name = finder._get_base_module_name(sys.modules["torchx.specs"])
-        self.assertEqual(expected_name, actual_name)
-
-    def test_get_component_name(self) -> None:
-        finder = ModuleComponentsFinder("", group="foobar")
-        actual_name = finder._get_component_name(
-            "test.main_module", "test.main_module.sub_module.bar", "get_component"
-        )
-        expected_name = "foobar.sub_module.bar.get_component"
-        self.assertEqual(expected_name, actual_name)
-
-    def test_strip_init(self) -> None:
-        finder = ModuleComponentsFinder("", "")
-        self.assertEqual("foobar", finder._strip_init("foobar.__init__"))
-        self.assertEqual("", finder._strip_init("__init__"))
-        self.assertEqual("foobar", finder._strip_init("foobar"))
-
-    def test_get_module_name(self) -> None:
-        finder = ModuleComponentsFinder("", "")
-        actual_name = finder._get_module_name(
-            "/test/path/main_module/foobar.py", "/test/path", "main"
+    @patch(
+        _METADATA_EPS,
+        return_value=EntryPoints(
+            EntryPoint_from_text(
+                """
+[torchx.components]
+foo = torchx.specs.test.components.a
+bar = torchx.specs.test.components.c.d
+"""
+            )
+        ),
+    )
+    def test_load_custom_components(self, _: MagicMock) -> None:
+        components = _load_components()
+
+        # the name of the appdefs returned by each component
+        # is the expected component name
+        for actual_name, comp in components.items():
+            expected_name = comp.fn().name
+            self.assertEqual(expected_name, actual_name)
+
+        self.assertEqual(3, len(components))
+
+    @patch(
+        _METADATA_EPS,
+        return_value=EntryPoints(
+            EntryPoint_from_text(
+                """
+[torchx.components]
+_0 = torchx.specs.test.components.a
+_1 = torchx.specs.test.components.c.d
+"""
+            )
+        ),
+    )
+    def test_load_custom_components_nogroup(self, _: MagicMock) -> None:
+        components = _load_components()
+
+        # test component names are hardcoded expecting
+        # test.components.* to be grouped under foo.*
+        # and components.a_namepace.* to be grouped under bar.*
+        # since we are testing _* (no group prefix) remove the first prefix
+        for actual_name, comp in components.items():
+            expected_name = comp.fn().name.split(".", maxsplit=1)[1]
+            self.assertEqual(expected_name, actual_name)
+
+    def test_load_builtins(self) -> None:
+        components = _load_components()
+
+        # if nothing registered in entrypoints, then builtins should be loaded
+        expected = {
+            c.name for c in ModuleComponentsFinder("torchx.components", group="").find()
+        }
+        self.assertEqual(components.keys(), expected)
+
+    def test_load_builtin_echo(self) -> None:
+        components = _load_components()
+        self.assertTrue(len(components) > 1)
+        component = components["utils.echo"]
+        self.assertEqual("utils.echo", component.name)
+        self.assertEqual(
+            "Echos a message to stdout (calls echo)", component.description
         )
-        expected_name = "main.main_module.foobar"
-        self.assertEqual(expected_name, actual_name)
+        self.assertEqual("echo", component.fn_name)
+        self.assertIsNotNone(component.fn)
 
 
 def current_file_path() -> str:

torchx/specs/test/finder_test.py

@@ -39,7 +39,10 @@ def load(group: str, name: str, default=None):
 
 def _defer_load_ep(ep: EntryPoint) -> object:
     def run(*args: object, **kwargs: object) -> object:
-        return ep.load()(*args, **kwargs)
+        if ep.attr is None:  # this is a module
+            return ep.load()
+        else:
+            return ep.load()(*args, **kwargs)
 
     return run
 
@@ -51,7 +54,10 @@ def load_group(
 ):
     """
     Loads all the entry points specified by ``group`` and returns
-    the entry points as a map of ``name (str) -> entrypoint.load()``.
+    the entry points as a map of ``name (str) -> deferred_load_fn``.
+    where the ``deferred_load_fn`` (as the name implies) defers the
+    loading of the entrypoint (e.g. ``entrypoint.load()``) until the
+    caller explicitly executes the funtion.
 
     For the following ``entry_point.txt``:
 
@@ -61,9 +67,21 @@ def load_group(
      bar = this.is:a_fn
      baz = this.is:b_fn
 
-    1. ``load_group("foo")`` -> ``{"bar", this.is.a_fn, "baz": this.is.b_fn}``
+    1. ``load_group("foo")["bar"]("baz")`` -> equivalent to calling ``this.is.a_fn("baz")``
     1. ``load_group("food")`` -> ``None``
-    1. ``load_group("food", default={"hello": this.is.c_fn})`` -> ``{"hello": this.is.c_fn}``
+    1. ``load_group("food", default={"hello": this.is.c_fn})["hello"]("world")`` -> equivalent to calling ``this.is.c_fn("world")``
+
+
+    If the entrypoint is a module (versus a function as shown above), then calling the ``deferred_load_fn``
+    simply loads the module and ignores any ``*args`` or ``**kwargs`` passed. For example:
+
+    ::
+
+     [foo]
+     bar = this.is.a.module
+
+    1. ``load_group("foo")["bar"]()`` -> loads ``this.is.a.module`` and returns a ``module`` type
+    1. ``load_group("foo")["bar"]("baz", hello="world")`` -> same as above (ignores ``*args`` and ``**kwargs``)
 
     """
 

torchx/util/entrypoints.py

@@ -6,6 +6,7 @@
 
 import unittest
 from configparser import ConfigParser
+from types import ModuleType
 from typing import List
 from unittest.mock import MagicMock, patch
 
@@ -54,6 +55,11 @@ def barbaz() -> str:
 baz = torchx.util.test.entrypoints_test:missing_attr
 """
 
+_EP_GRP_MOD_TXT: str = """
+[ep.grp.mod.test]
+baz = torchx.util.test.entrypoints_test
+"""
+
 _EP_GRP_IGN_MOD_TXT: str = """
 [ep.grp.missing.mod.test]
 baz = torchx.util.test.entrypoints_test.missing_module
@@ -62,6 +68,7 @@ def barbaz() -> str:
     EntryPoint_from_text(_EP_TXT)
     + EntryPoint_from_text(_EP_GRP_TXT)
     + EntryPoint_from_text(_EP_GRP_IGN_ATTR_TXT)
+    + EntryPoint_from_text(_EP_GRP_MOD_TXT)
     + EntryPoint_from_text(_EP_GRP_IGN_MOD_TXT)
 )
 
@@ -70,21 +77,21 @@ def barbaz() -> str:
 
 class EntryPointsTest(unittest.TestCase):
     @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
-    def test_load(self, mock_md_eps: MagicMock) -> None:
+    def test_load(self, _: MagicMock) -> None:
         print(type(load("entrypoints.test", "foo")))
         self.assertEqual("foobar", load("entrypoints.test", "foo")())
 
         with self.assertRaisesRegex(KeyError, "baz"):
             load("entrypoints.test", "baz")()
 
     @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
-    def test_load_with_default(self, mock_md_eps: MagicMock) -> None:
+    def test_load_with_default(self, _: 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)())
 
     @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
-    def test_load_group(self, mock_md_eps: MagicMock) -> None:
+    def test_load_group(self, _: MagicMock) -> None:
         eps = load_group("ep.grp.test")
         self.assertEqual(2, len(eps), eps)
         self.assertEqual("foobar", eps["foo"]())
@@ -93,8 +100,16 @@ def test_load_group(self, mock_md_eps: MagicMock) -> None:
         eps = load_group("ep.grp.test.missing")
         self.assertIsNone(eps)
 
+        eps = load_group("ep.grp.mod.test")
+        module = eps["baz"]()
+        self.assertEqual(ModuleType, type(module))
+        self.assertEqual("torchx.util.test.entrypoints_test", module.__name__)
+
+        # module's deferred load function should ignore *args and **kwargs
+        self.assertEqual(module, eps["baz"]("ignored", should="ignore"))
+
     @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
-    def test_load_group_with_default(self, mock_md_eps: MagicMock) -> None:
+    def test_load_group_with_default(self, _: MagicMock) -> None:
         eps = load_group("ep.grp.test", {"foo": barbaz, "bar": foobar})
         self.assertEqual(2, len(eps))
         self.assertEqual("foobar", eps["foo"]())
@@ -106,7 +121,7 @@ def test_load_group_with_default(self, mock_md_eps: MagicMock) -> None:
         self.assertEqual("foobar", eps["bar"]())
 
     @patch(_METADATA_EPS, return_value=_ENTRY_POINTS)
-    def test_load_group_not_missing(self, mock_md_eps: MagicMock) -> None:
+    def test_load_group_missing(self, _: MagicMock) -> None:
         with self.assertRaises(AttributeError):
             load_group("ep.grp.missing.attr.test")["baz"]()