From 900e66c5e3717ae4e2011f743f6718880d6969d2 Mon Sep 17 00:00:00 2001
From: David Corvoysier <david@huggingface.co>
Date: Wed, 5 Jun 2024 11:09:02 +0200
Subject: [PATCH] feat(commands): allow subpackages to declare their
 subcommands

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.
---
 optimum/commands/__init__.py    |  2 +-
 optimum/commands/optimum_cli.py | 56 ++++++++++++++++++++++++++++++---
 optimum/subpackages.py          |  2 +-
 3 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/optimum/commands/__init__.py b/optimum/commands/__init__.py
index 540ea4dd86..37bacd87fc 100644
--- a/optimum/commands/__init__.py
+++ b/optimum/commands/__init__.py
@@ -16,4 +16,4 @@
 from .env import EnvironmentCommand
 from .export import ExportCommand, ONNXExportCommand, TFLiteExportCommand
 from .onnxruntime import ONNXRuntimeCommand, ONNXRuntimeOptimizeCommand, ONNXRuntimeQuantizeCommand
-from .optimum_cli import register_optimum_cli_subcommand
+from .optimum_cli import optimum_cli_subcommand
diff --git a/optimum/commands/optimum_cli.py b/optimum/commands/optimum_cli.py
index 4bae9bb5f8..87ada7c895 100644
--- a/optimum/commands/optimum_cli.py
+++ b/optimum/commands/optimum_cli.py
@@ -17,6 +17,7 @@
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Type, Union
 
+from ..subpackages import load_subpackages
 from ..utils import logging
 from .base import BaseOptimumCLICommand, CommandInfo, RootOptimumCLICommand
 from .env import EnvironmentCommand
@@ -26,7 +27,48 @@
 
 logger = logging.get_logger()
 
-OPTIMUM_CLI_SUBCOMMANDS = [ExportCommand, EnvironmentCommand, ONNXRuntimeCommand]
+# The table below contains the optimum-cli root subcommands provided by the optimum package
+OPTIMUM_CLI_ROOT_SUBCOMMANDS = [ExportCommand, EnvironmentCommand, ONNXRuntimeCommand]
+
+# The table below is dynamically populated when loading subpackages
+_OPTIMUM_CLI_SUBCOMMANDS = []
+
+
+def optimum_cli_subcommand(parent_command: Optional[Type[BaseOptimumCLICommand]] = None):
+    """
+    A decorator to declare optimum-cli subcommands.
+
+    The declaration of an optimum-cli subcommand looks like this:
+
+    ```
+    @optimum_cli_subcommand()
+    class MySubcommand(BaseOptimumCLICommand):
+        <implementation>
+    ```
+
+    or
+
+    ```
+    @optimum_cli_subcommand(ExportCommand)
+    class MySubcommand(BaseOptimumCLICommand):
+        <implementation>
+    ```
+
+    Args:
+        parent_command: (`Optional[Type[BaseOptimumCLICommand]]`):
+            The class of the parent command or None if this is a top-level command. Defaults to None.
+
+    """
+
+    if parent_command is not None and not issubclass(parent_command, BaseOptimumCLICommand):
+        raise ValueError(f"The parent command {parent_command} must be a subclass of BaseOptimumCLICommand")
+
+    def wrapper(subcommand):
+        if not issubclass(subcommand, BaseOptimumCLICommand):
+            raise ValueError(f"The subcommand {subcommand} must be a subclass of BaseOptimumCLICommand")
+        _OPTIMUM_CLI_SUBCOMMANDS.append((subcommand, parent_command))
+
+    return wrapper
 
 
 def resolve_command_to_command_instance(
@@ -137,15 +179,19 @@ def main():
     root = RootOptimumCLICommand("Optimum CLI tool", usage="optimum-cli")
     parser = root.parser
 
-    for subcommand_cls in OPTIMUM_CLI_SUBCOMMANDS:
+    for subcommand_cls in OPTIMUM_CLI_ROOT_SUBCOMMANDS:
         register_optimum_cli_subcommand(subcommand_cls, parent_command=root)
 
-    commands_in_register = dynamic_load_commands_in_register()
+    # Load subpackages to give them a chance to declare their own subcommands
+    load_subpackages()
+
+    # Register subcommands declared by the subpackages or found in the register files under commands/register
+    commands_to_register = _OPTIMUM_CLI_SUBCOMMANDS + dynamic_load_commands_in_register()
     command2command_instance = resolve_command_to_command_instance(
-        root, [parent_command_cls for _, parent_command_cls in commands_in_register if parent_command_cls is not None]
+        root, [parent_command_cls for _, parent_command_cls in commands_to_register if parent_command_cls is not None]
     )
 
-    for command_or_command_info, parent_command in commands_in_register:
+    for command_or_command_info, parent_command in commands_to_register:
         if parent_command is None:
             parent_command_instance = root
         else:
diff --git a/optimum/subpackages.py b/optimum/subpackages.py
index 11e8707953..9f905cba44 100644
--- a/optimum/subpackages.py
+++ b/optimum/subpackages.py
@@ -56,4 +56,4 @@ def load_subpackages():
     """
     OPTIMUM_NAMESPACE = "optimum"
     OPTIMUM_SUBPACKAGE_LOADER = "subpackage"
-    return load_namespace_modules(OPTIMUM_NAMESPACE, OPTIMUM_SUBPACKAGE_LOADER)
+    load_namespace_modules(OPTIMUM_NAMESPACE, OPTIMUM_SUBPACKAGE_LOADER)