From d47be235d3ff6f928798acd7963103cb9045cabc Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Wed, 5 Jun 2024 15:30:40 +0200 Subject: [PATCH] refactor(onnxruntime): add subpackage and move commands 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. --- optimum/commands/__init__.py | 1 - optimum/commands/optimum_cli.py | 3 +-- optimum/onnxruntime/runs/__init__.py | 6 +++--- optimum/onnxruntime/subpackage/__init__.py | 1 + .../subpackage/commands}/__init__.py | 2 -- .../subpackage/commands}/base.py | 4 +++- .../subpackage/commands}/optimize.py | 4 ++-- .../subpackage/commands}/quantize.py | 6 +++--- optimum/subpackages.py | 7 +++++++ 9 files changed, 20 insertions(+), 14 deletions(-) create mode 100644 optimum/onnxruntime/subpackage/__init__.py rename optimum/{commands/onnxruntime => onnxruntime/subpackage/commands}/__init__.py (87%) rename optimum/{commands/onnxruntime => onnxruntime/subpackage/commands}/base.py (91%) rename optimum/{commands/onnxruntime => onnxruntime/subpackage/commands}/optimize.py (96%) rename optimum/{commands/onnxruntime => onnxruntime/subpackage/commands}/quantize.py (95%) diff --git a/optimum/commands/__init__.py b/optimum/commands/__init__.py index 37bacd87fc..8a2a276d1c 100644 --- a/optimum/commands/__init__.py +++ b/optimum/commands/__init__.py @@ -15,5 +15,4 @@ from .base import BaseOptimumCLICommand, CommandInfo, RootOptimumCLICommand from .env import EnvironmentCommand from .export import ExportCommand, ONNXExportCommand, TFLiteExportCommand -from .onnxruntime import ONNXRuntimeCommand, ONNXRuntimeOptimizeCommand, ONNXRuntimeQuantizeCommand from .optimum_cli import optimum_cli_subcommand diff --git a/optimum/commands/optimum_cli.py b/optimum/commands/optimum_cli.py index 87ada7c895..64a7075c6c 100644 --- a/optimum/commands/optimum_cli.py +++ b/optimum/commands/optimum_cli.py @@ -22,13 +22,12 @@ from .base import BaseOptimumCLICommand, CommandInfo, RootOptimumCLICommand from .env import EnvironmentCommand from .export import ExportCommand -from .onnxruntime import ONNXRuntimeCommand logger = logging.get_logger() # The table below contains the optimum-cli root subcommands provided by the optimum package -OPTIMUM_CLI_ROOT_SUBCOMMANDS = [ExportCommand, EnvironmentCommand, ONNXRuntimeCommand] +OPTIMUM_CLI_ROOT_SUBCOMMANDS = [ExportCommand, EnvironmentCommand] # The table below is dynamically populated when loading subpackages _OPTIMUM_CLI_SUBCOMMANDS = [] diff --git a/optimum/onnxruntime/runs/__init__.py b/optimum/onnxruntime/runs/__init__.py index 1d98294934..d21db2a4ac 100644 --- a/optimum/onnxruntime/runs/__init__.py +++ b/optimum/onnxruntime/runs/__init__.py @@ -110,9 +110,9 @@ def __init__(self, run_config): model_class = FeaturesManager.get_model_class_for_feature(get_autoclass_name(self.task)) self.torch_model = model_class.from_pretrained(run_config["model_name_or_path"]) - self.return_body[ - "model_type" - ] = self.torch_model.config.model_type # return_body is initialized in parent class + self.return_body["model_type"] = ( + self.torch_model.config.model_type + ) # return_body is initialized in parent class def _launch_time(self, trial): batch_size = trial.suggest_categorical("batch_size", self.batch_sizes) diff --git a/optimum/onnxruntime/subpackage/__init__.py b/optimum/onnxruntime/subpackage/__init__.py new file mode 100644 index 0000000000..7029af7132 --- /dev/null +++ b/optimum/onnxruntime/subpackage/__init__.py @@ -0,0 +1 @@ +from .commands import ONNXRuntimeCommand diff --git a/optimum/commands/onnxruntime/__init__.py b/optimum/onnxruntime/subpackage/commands/__init__.py similarity index 87% rename from optimum/commands/onnxruntime/__init__.py rename to optimum/onnxruntime/subpackage/commands/__init__.py index 1b9c24c3b2..44facf5ea5 100644 --- a/optimum/commands/onnxruntime/__init__.py +++ b/optimum/onnxruntime/subpackage/commands/__init__.py @@ -14,5 +14,3 @@ # limitations under the License. from .base import ONNXRuntimeCommand -from .optimize import ONNXRuntimeOptimizeCommand -from .quantize import ONNXRuntimeQuantizeCommand diff --git a/optimum/commands/onnxruntime/base.py b/optimum/onnxruntime/subpackage/commands/base.py similarity index 91% rename from optimum/commands/onnxruntime/base.py rename to optimum/onnxruntime/subpackage/commands/base.py index 53e3245ea4..df4414c19d 100644 --- a/optimum/commands/onnxruntime/base.py +++ b/optimum/onnxruntime/subpackage/commands/base.py @@ -14,11 +14,13 @@ # limitations under the License. """optimum.onnxruntime command-line interface base classes.""" -from .. import BaseOptimumCLICommand, CommandInfo +from optimum.commands import BaseOptimumCLICommand, CommandInfo, optimum_cli_subcommand + from .optimize import ONNXRuntimeOptimizeCommand from .quantize import ONNXRuntimeQuantizeCommand +@optimum_cli_subcommand() class ONNXRuntimeCommand(BaseOptimumCLICommand): COMMAND = CommandInfo( name="onnxruntime", diff --git a/optimum/commands/onnxruntime/optimize.py b/optimum/onnxruntime/subpackage/commands/optimize.py similarity index 96% rename from optimum/commands/onnxruntime/optimize.py rename to optimum/onnxruntime/subpackage/commands/optimize.py index 5890e0a07c..1dd82f0ee2 100644 --- a/optimum/commands/onnxruntime/optimize.py +++ b/optimum/onnxruntime/subpackage/commands/optimize.py @@ -75,8 +75,8 @@ def parse_args(parser: "ArgumentParser"): return parse_args_onnxruntime_optimize(parser) def run(self): - from ...onnxruntime.configuration import AutoOptimizationConfig, ORTConfig - from ...onnxruntime.optimization import ORTOptimizer + from ...configuration import AutoOptimizationConfig, ORTConfig + from ...optimization import ORTOptimizer if self.args.output == self.args.onnx_model: raise ValueError("The output directory must be different than the directory hosting the ONNX model.") diff --git a/optimum/commands/onnxruntime/quantize.py b/optimum/onnxruntime/subpackage/commands/quantize.py similarity index 95% rename from optimum/commands/onnxruntime/quantize.py rename to optimum/onnxruntime/subpackage/commands/quantize.py index 2613cb33ba..6f6d843cc7 100644 --- a/optimum/commands/onnxruntime/quantize.py +++ b/optimum/onnxruntime/subpackage/commands/quantize.py @@ -17,7 +17,7 @@ from pathlib import Path from typing import TYPE_CHECKING -from .. import BaseOptimumCLICommand +from optimum.commands import BaseOptimumCLICommand if TYPE_CHECKING: @@ -69,8 +69,8 @@ def parse_args(parser: "ArgumentParser"): return parse_args_onnxruntime_quantize(parser) def run(self): - from ...onnxruntime.configuration import AutoQuantizationConfig, ORTConfig - from ...onnxruntime.quantization import ORTQuantizer + from ...configuration import AutoQuantizationConfig, ORTConfig + from ...quantization import ORTQuantizer if self.args.output == self.args.onnx_model: raise ValueError("The output directory must be different than the directory hosting the ONNX model.") diff --git a/optimum/subpackages.py b/optimum/subpackages.py index 9f905cba44..4ad02bc827 100644 --- a/optimum/subpackages.py +++ b/optimum/subpackages.py @@ -8,6 +8,8 @@ import importlib_metadata from importlib.util import find_spec, module_from_spec +from .utils import is_onnxruntime_available + logger = logging.getLogger(__name__) @@ -54,6 +56,11 @@ def load_subpackages(): This module is then in charge of registering the subpackage commands. """ + OPTIMUM_NAMESPACE = "optimum" OPTIMUM_SUBPACKAGE_LOADER = "subpackage" load_namespace_modules(OPTIMUM_NAMESPACE, OPTIMUM_SUBPACKAGE_LOADER) + + # Load subpackages from internal modules not explicitly defined as namespace packages + if is_onnxruntime_available(): + from .onnxruntime import subpackage