Allow optimum to discover and load subpackages#1894
Conversation
d47be23 to
35e5e79
Compare
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
33cf651 to
70660ad
Compare
dacorvo
changed the title
As an alternative to directly adding their commands in a register.py file under the root optimum directory, this adds a decorator to declare a subcommand that can be used by subpackages when they are loaded. This will fix the issue of subcommands 'disappearing' when optimum is upgraded without reinstalling the subpackage.
The onnxruntime commands are moved into a subpackage loader directory. This subpackage directory is only loaded (and its commands added) when the onnxruntime is available. This avoids wrongly indicating that the onnxruntime commands are available when the package is actually not installed.
70660ad to
40e4f8d
Compare
fxmarty
left a comment
fxmarty
left a comment
There was a problem hiding this comment.
LGTM, apart that optimum-cli has a 1-3s delay to start running (already on main) on my laptop, which should be ideally improved in a future PR
| from typing import TYPE_CHECKING | ||
|
|
||
| from .. import BaseOptimumCLICommand | ||
| from optimum.commands import BaseOptimumCLICommand |
There was a problem hiding this comment.
That's mostly because the relative path would have been complicated, but also to serve as a model for actual subpackages that can live outside the optimum tree.
| SUBPACKAGE_LOADER = "subpackage" | ||
| load_namespace_modules("optimum", SUBPACKAGE_LOADER) | ||
|
|
||
| # Load subpackages from internal modules not explicitly defined as namespace packages | ||
| loader_name = "." + SUBPACKAGE_LOADER |
There was a problem hiding this comment.
That's the name of the module that is imported in each subpackage to "load" them. I used a constant because I use it twice and was expecting maybe some discussion about the naming.
| from .env import EnvironmentCommand | ||
| from .export import ExportCommand, ONNXExportCommand, TFLiteExportCommand | ||
| from .onnxruntime import ONNXRuntimeCommand, ONNXRuntimeOptimizeCommand, ONNXRuntimeQuantizeCommand | ||
| from .optimum_cli import register_optimum_cli_subcommand |
There was a problem hiding this comment.
Would need to make sure register_optimum_cli_subcommand is not used in subpackages
There was a problem hiding this comment.
There are no calls in optimum and I quickly checked the optimum-XX to verify that none of them was calling it either, but I was not 100 % sure. If that is an issue I can make it public again.
There was a problem hiding this comment.
Existing subpackages are using the "register" method at the moment.
michaelbenayoun
left a comment
michaelbenayoun
left a comment
There was a problem hiding this comment.
That is super cool!
Thanks for this improvement @dacorvo
| if sys.version_info >= (3, 8): | ||
| from importlib import metadata as importlib_metadata | ||
| else: | ||
| import importlib_metadata |
There was a problem hiding this comment.
That raises the question for later: what is the minimal python version we want to support. I have seen different things accross differents huggingface libraries.
There was a problem hiding this comment.
Python 3.8 is the oldest supported, EOL is october 2024: https://devguide.python.org/versions/
I guess we could only support 3.8 and newer.
5 participants
What does this PR do?
This introduces explicitly the concept of
optimumsubpackages that are loaded "dynamically" by theoptimum-cli(and possibly in the future by other high-level abstraction modules that would require backend-specific implementations).An
optimumsubpackage is primarily a package under theoptimumnamespace, but it can be also a module inside theoptimumpackage itself, even if that module does not have its ownpyproject.tomlorsetup.py(as a real namespace package would).By convention, loading a subpackage means importing a
subpackagemodule available at the root of the package.This pull-request first adds a
load_subpackageshelper that will go through the installed namespace packages and try to "load" them.It then adds a decorator to declare custom subcommands, using the same convention as the one already used in the
registerfiles.This is by using this decorator that each subpackage can add its own subcommands, as it will be called by the
optimum-cliwhen evaluating its commands.In order to validate the new helpers, the
onnxruntimecommands are moved inside anonnxruntime/subpackagedirectory that is explicitly loaded byload_subpackages, but only it theonnxruntimepackage is available.This improves the UX when using the
optimum-cli:optimuminstallation,onnxruntimecommands will only be available if the package is actually installed.