Lazy load commands from plugins #3901

ankatiyar · 2024-05-30T11:12:09Z

Description

Partly resolve #1476

As discussed in #1476, the initialisation of KedroCLI takes up a significant chunk of time. This is especially evident if you have a lot of kedro plugins installed in the environment as KedroCLI which is a CommandCollection, loads up all the commands from the plugins before the command is processed. This PR is to add a lazy way to load the commands from the plugins.

Development notes

Add a property called plugin_groups which returns a dict of entry_point_name to the EntryPoint object.
This group is passed to the super class CommandCollection which is a custom version of click.CommandCollection used by Kedro and defined in kedro/framework/cli/utils.py

Update to custom CommandCollection:

Add a new main function which receives the command to be run as args, eg [ "catalog", "list"] or ["airflow", "create"]
if the first element of the args is not in the command, then load commands from plugins
run the command by passing to super().main()

Loading of plugins:

I tried to implement a sort of "smart" loading where if the command arg, i.e. airflow or mlflow or viz partially matches the entry point names in the lazy_group dict keys, load the plugin, check if the command now exists and exit.
Note: The entry point names are decided by the plugins, eg. kedro-airflow's ep name is airflow which fully matches the command name too but for kedro-viz and kedro-mlflow the entry point names don't match the commands. eg. Command: viz Entry point name: kedro-viz

Otherwise, load all plugins one by one, exit if the command exists after loading plugins.

TODO

Fix code coverage on unit tests, which is tricky since this code is only used when a plugin is installed but we don't have tests like these in our unit tests.

Developer Certificate of Origin

We need all contributions to comply with the Developer Certificate of Origin (DCO). All commits must be signed off by including a Signed-off-by line in the commit message. See our wiki for guidance.

If your PR is blocked due to unsigned commits, then you must follow the instructions under "Rebase the branch" on the GitHub Checks page for your PR. This will retroactively add the sign-off to all unsigned commits and allow the DCO check to pass.

Checklist

Read the contributing guidelines
Signed off each commit with a Developer Certificate of Origin (DCO)
Opened this PR as a 'Draft Pull Request' if it is work-in-progress
Updated the documentation to reflect the code changes
Added a description of this change in the RELEASE.md file
Added tests to cover my changes
Checked if this change will affect Kedro-Viz, and if so, communicated that with the Viz team

Signed-off-by: Ankita Katiyar <ankitakatiyar2401@gmail.com>

astrojuanlu · 2024-06-03T12:08:48Z

To note, as @ankatiyar noted internally, "bad" lookups will take a lot of time: for example, kedro infoo (notice typo) took 17 seconds in my computer and tried to import kedro-viz, kedro-mlflow, and more.

Going forward, is there a smarter way we can declare new CLI commands, so that we can look them up in entry points without actually importing anything? (Feel free to open a new issue about this)

Signed-off-by: Ankita Katiyar <ankitakatiyar2401@gmail.com>

ankatiyar · 2024-06-10T13:53:49Z

The code coverage is not complete but opening this up for review anyway to gather feedback on the implementation.

merelcht

I added some suggestions/questions to the code, but I actually don't think it's a good idea that now the plugin commands aren't loaded when you do kedro --help. I've built a small plugin myself following https://docs.kedro.org/en/stable/extend_kedro/plugins.html#example-of-a-simple-plugin and the problem you get here is that the commands belong to the kedro group, but don't show at all. In this case there's no other top-level group you can call to see what sub-commands exist. On top of that the overwriting behaviour is changed, so if I now create my own micropkg command inside this kedrojson plugin it doesn't overwrite the original kedro micropkg where in the non lazy-loading version it does do the overwriting.

While this solution does lead to performance gain, it also degrades the plugin command discoverability and therefore IMO usability of the plugins as a whole. Features like https://docs.kedro.org/en/stable/extend_kedro/plugins.html#global-and-project-commands will not exist anymore after this change.

All in all, I don't think the performance gains are worth losing the plugin features.

merelcht · 2024-06-11T09:41:50Z

kedro/framework/cli/cli.py

@@ -172,13 +174,19 @@ def main(
                click.echo(hint)
            sys.exit(exc.code)

+    @property
+    def plugin_groups(self) -> dict[str, EntryPoint]:


Can you add a docstring here to explain what the property is for?

merelcht · 2024-06-11T09:42:22Z

kedro/framework/cli/cli.py

        """Property which loads all global command groups from plugins and
        combines them with the built-in ones (eventually overriding the
        built-in ones if they are redefined by plugins).


I'm guessing this needs to be updated, because it's not loading the plugins anymore?

I'm also wondering whether the "overriding" behaviour mentioned here is changed with the lazy loading? Update: I tested this and the behaviour is indeed changed. With the lazy loading it seems like the overwritten version isn't called.

merelcht · 2024-06-11T09:43:47Z

kedro/framework/cli/utils.py

@@ -90,6 +90,13 @@ def wrapit(func: Any) -> Any:
    return wrapit


+def _partial_match(plugin_names: list[str], command_name: str) -> str | None:


Can you add a docstring here as well?

ankatiyar · 2024-06-11T12:03:11Z

Fair point @merelcht! 🤔
I was also trying to implement a way where --help does trigger the loading of the plugins.
About overwriting of core Kedro commands eg. micropkg, that might also be possible with lazy loading. Let me explore these functionalities and update the PR (if it works!)!

noklam · 2024-06-17T10:00:36Z

kedro/framework/cli/utils.py

+        **extra: Any,
+    ) -> Any:
+        # Load plugins if the command is not found in the current sources
+        if args and args[0] not in self.list_commands(None):  # type: ignore[arg-type]


What does self.list_commands(None) do? From the docs it seems like it expects a ctx.

Can this be replaced with get_command instead? https://click.palletsprojects.com/en/8.1.x/api/#click.MultiCommand.get_command

ankatiyar · 2024-06-18T12:13:19Z

Unfortunately, I don't think the overriding behaviour of CLI is possible with lazy loading of plugins - i.e. if a plugin has overriding CLI commands for Kedro core commands such as kedro catalog list etc since the premise of the proposal is to only load plugins if a command does not exist in KedroCLI.
I think it's possible to trigger the loading of plugins and display plugin commands with kedro --help but I would like to get the team's opinions on if we should be going ahead with this at all
cc @astrojuanlu @noklam @merelcht

astrojuanlu · 2024-06-18T13:07:51Z

The python CLI has --help and --help-all. Maybe for the next breaking release we could make --help display only the "core" commands, and --help-all display everything.

For retaining the current behavior while attaining much faster load times for all the rest of commands, I'd be 👍🏼 on force-triggering loading of all plugins in kedro --help

ankatiyar · 2024-06-19T14:35:32Z

Moving this back to the drafts, will close it if no other comments come in. I'll focus on getting kedro-org/kedro-viz#1920 and #3883 ready for review first!

ankatiyar · 2024-06-21T08:10:02Z

Closing this in favour of the other solution!

Lazy load commands from plugins

81ee250

Signed-off-by: Ankita Katiyar <ankitakatiyar2401@gmail.com>

ankatiyar mentioned this pull request May 30, 2024

[spike] Improve Kedro CLI startup time #1476

Closed

ankatiyar added 2 commits May 30, 2024 15:39

Combine into plugin_groups

71de119

Signed-off-by: Ankita Katiyar <ankitakatiyar2401@gmail.com>

Fix tests

7c09aa7

Signed-off-by: Ankita Katiyar <ankitakatiyar2401@gmail.com>

ankatiyar added 3 commits June 6, 2024 11:54

Try smart loading first

157bffb

Signed-off-by: Ankita Katiyar <ankitakatiyar2401@gmail.com>

lint + tests

46648d8

Signed-off-by: Ankita Katiyar <ankitakatiyar2401@gmail.com>

Lint

2576d1f

Signed-off-by: Ankita Katiyar <ankitakatiyar2401@gmail.com>

ankatiyar marked this pull request as ready for review June 10, 2024 13:52

ankatiyar requested a review from merelcht as a code owner June 10, 2024 13:52

ankatiyar requested review from astrojuanlu and noklam June 10, 2024 13:52

ankatiyar requested a review from DimedS June 10, 2024 13:55

merelcht requested changes Jun 11, 2024

View reviewed changes

noklam reviewed Jun 17, 2024

View reviewed changes

ankatiyar closed this Jun 21, 2024

ankatiyar deleted the lazy-load-plugins branch August 21, 2024 11:08

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Lazy load commands from plugins #3901

Lazy load commands from plugins #3901

ankatiyar commented May 30, 2024 •

edited

Loading

astrojuanlu commented Jun 3, 2024

ankatiyar commented Jun 10, 2024 •

edited

Loading

merelcht left a comment •

edited

Loading

merelcht Jun 11, 2024

merelcht Jun 11, 2024

merelcht Jun 11, 2024

merelcht Jun 11, 2024

ankatiyar commented Jun 11, 2024 •

edited

Loading

noklam Jun 17, 2024

noklam Jun 17, 2024

ankatiyar commented Jun 18, 2024

astrojuanlu commented Jun 18, 2024

ankatiyar commented Jun 19, 2024

ankatiyar commented Jun 21, 2024

		@@ -90,6 +90,13 @@ def wrapit(func: Any) -> Any:
		return wrapit


		def _partial_match(plugin_names: list[str], command_name: str) -> str \| None:

Lazy load commands from plugins #3901

Lazy load commands from plugins #3901

Conversation

ankatiyar commented May 30, 2024 • edited Loading

Description

Development notes

TODO

Developer Certificate of Origin

Checklist

astrojuanlu commented Jun 3, 2024

ankatiyar commented Jun 10, 2024 • edited Loading

merelcht left a comment • edited Loading

Choose a reason for hiding this comment

merelcht Jun 11, 2024

Choose a reason for hiding this comment

merelcht Jun 11, 2024

Choose a reason for hiding this comment

merelcht Jun 11, 2024

Choose a reason for hiding this comment

merelcht Jun 11, 2024

Choose a reason for hiding this comment

ankatiyar commented Jun 11, 2024 • edited Loading

noklam Jun 17, 2024

Choose a reason for hiding this comment

noklam Jun 17, 2024

Choose a reason for hiding this comment

ankatiyar commented Jun 18, 2024

astrojuanlu commented Jun 18, 2024

ankatiyar commented Jun 19, 2024

ankatiyar commented Jun 21, 2024

ankatiyar commented May 30, 2024 •

edited

Loading

ankatiyar commented Jun 10, 2024 •

edited

Loading

merelcht left a comment •

edited

Loading

ankatiyar commented Jun 11, 2024 •

edited

Loading