(torchx/components) Fix entrypoint loading to deal with deferred load…

[kiukchung]

Fixes a bug + enhances custom component registration.

1. [bugfix] torchx not allowing components to be registered via [torchx.components] entrypoint because of a change in utils.entrypoints.load_entrypoints() that defers the entrypoint load by wrapping the entrypoints with a _defer_load_fn() but assumes that all entrypoints are functions (e.g. foo.bar:baz) and does not handle modules (e.g. foo.bar). Adds graceful handling of module entrypoints.

2. [enhancement] Allows users to specify an "empty" group (e.g. no prefix) for component names by using the _* "special" prefix in the alias key (similar to "hidden" variables in python). Example

   [torchx.components]
    _0 = torchx.components.dist
    _1 = torchx.components.utils
   

   names dist.ddp as just ddp and utils.python as just python

3. [change] If custom components are registered via entrypoints, skips loading builtins. Users can still load the builtins by adding an extra line in their entrypoint as:

   [torchx.components]
   foo = bar.baz
   _ = torchx.components
   

   NOTE: this is a BC incompatible change, HOWEVER given that component registration via entrypoints was broken, I'm assuming that no one was registering their own components.

4. [docs] Reflected the changes above in the docs page

5. [todo fix] resolves [file-linter] Refactor file-linter and components-finder #261 since we no longer have components_base.py file
   Test plan:

---

Added unittests + revised unittests for utils.entrypoints and specs.finder

[d4l3k]

@kiukchung lots of CI failures (docs, pyre, unit) are the relevant ones I think

[d4l3k]

LGTM

[d4l3k]

may want to mention you can do this in the config as well?

[kiukchung]

the one supported by .torchxconfig is super weird (I want to deprecate it in favor of supporting "alias" -> "module_name" style custom components from .torchxconfig). Currently what you can do is to have a section like

# .torchxconfig
[torchx.file]
get_file_contents = my.custom.file_reader:read

And when you run

torchx run my/custom.py:component_fn

We call my.custom.file_reader.read("my/custom.py") which we expect to load the component_fn into the global namespace. Took me a while to wrap my head around it, so I don't expect anyone to actually use it.

Instead I'd like to support something like this from .torchxconfig:

# .torchxconfig
[torchx.components]
foo = my.custom.component
_ = torchx.components

which would be the same as doing

# setup.py

entrypoints = {
  "torchx.components": [
        "foo = my.custom.component",
        "_ = torchx.components",
   ],
}

[d4l3k]

this now works recursively? got it, interesting OK

[kiukchung]

it was always loading recursively (by file not by module though)

[d4l3k]

do we want to drop the dist. and utils. prefixes for TorchX components now? not a blocker for this PR

[kiukchung]

yeah so we could drop the "module" name from the builtin components (as discussed offline we'd have to keep the alias dist.*,util.* etc while hiding them from torchx builtins). Should be easy to do now since we can make the builtins load:

ModuleComponentsFinder("torchx.components.dist", group="")
ModuleComponentsFinder("torchx.components.utils", group="")
# ... and so on...

[codecov]

Codecov Report

Merging #695 (9e30ed6) into main (c74bb9c) will increase coverage by 0.06%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #695      +/-   ##
==========================================
+ Coverage   92.39%   92.46%   +0.06%     
==========================================
  Files          83       86       +3     
  Lines        5670     5666       -4     
==========================================
  Hits         5239     5239              
+ Misses        431      427       -4     

Impacted Files	Coverage Δ
torchx/specs/finder.py	99.33% <100.00%> (+2.34%)	⬆️
torchx/specs/test/components/a/__init__.py	100.00% <100.00%> (ø)
torchx/specs/test/components/a/b/c.py	100.00% <100.00%> (ø)
torchx/specs/test/components/c/d.py	100.00% <100.00%> (ø)
torchx/util/entrypoints.py	100.00% <100.00%> (ø)
torchx/components/dist.py	96.36% <0.00%> (-0.07%)	⬇️
torchx/schedulers/kubernetes_mcad_scheduler.py	94.02% <0.00%> (-0.02%)	⬇️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[kiukchung]

Also resolves: #694 by passing os.getenv["CONTAINER_REPO"] as the container_repo argument to push_image() in scripts/kube_dist_trainer.py