update commands loading following native namespace packages

Author: IlyasMoutawwakil

optimum/commands/optimum_cli.py

@@ -17,11 +17,10 @@
 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
-from .export import ExportCommand
+from optimum.commands.base import BaseOptimumCLICommand, CommandInfo, RootOptimumCLICommand
+from optimum.commands.env import EnvironmentCommand
+from optimum.commands.export.base import ExportCommand
+from optimum.utils import logging
 
 
 logger = logging.get_logger()
@@ -110,29 +109,37 @@ def dynamic_load_commands_in_register() -> (
 ):
     """
     Loads a list of command classes to register to the CLI from the `optimum/commands/register/` directory.
-    It will look in any python file if there is a `REGISTER_COMMANDS` list, and load the commands to register
-    accordingly.
-    At this point, nothing is actually registered in the root CLI.
+
 
     Returns:
         `List[Tuple[Union[Type[BaseOptimumCLICommand], CommandInfo], Optional[Type[BaseOptimumCLICommand]]]]`: A list
         of tuples with two elements. The first element corresponds to the command to register, it can either be a
         subclass of `BaseOptimumCLICommand` or a `CommandInfo`. The second element corresponds to the parent command,
         where `None` means that the parent command is the root CLI command.
     """
-    commands_to_register = []
-    register_dir_path = Path(__file__).parent / "register"
-    for filename in register_dir_path.iterdir():
-        if filename.is_dir() or filename.suffix != ".py":
-            if filename.name not in ["__pycache__", "README.md"]:
-                logger.warning(
-                    f"Skipping {filename} because only python files are allowed when registering commands dynamically."
-                )
+    registration_files = []
+    for dist in importlib.metadata.distributions():
+        if dist.metadata["Name"] is None:
+            continue
+        if dist.metadata["Name"] == "optimum-benchmark":
+            continue
+        if not dist.metadata["Name"].startswith("optimum-"):
             continue
-        module_name = f".register.{filename.stem}"
-        module = importlib.import_module(module_name, package="optimum.commands")
-        commands_to_register_in_file = getattr(module, "REGISTER_COMMANDS", [])
-        for command_idx, command in enumerate(commands_to_register_in_file):
+
+        dist_name = dist.metadata["Name"]
+        dist_module_name = dist_name.replace("-", ".")
+        dist_module = importlib.import_module(dist_module_name)
+        dist_module_path = Path(dist_module.__file__).parent
+        commands_register_path = dist_module_path.parent / "commands" / "register"
+        for file in commands_register_path.iterdir():
+            if file.suffix == ".py" and file.stem != "__init__":
+                registration_files.append(file.stem)
+
+    commands_to_register = []
+    for register_file in registration_files:
+        submodule = importlib.import_module(f"optimum.commands.register.{register_file}")
+        commands_to_register_in_module = getattr(submodule, "REGISTER_COMMANDS", [])
+        for command_idx, command in enumerate(commands_to_register_in_module):
             if isinstance(command, tuple):
                 command_or_command_info, parent_command_cls = command
             else:
@@ -143,9 +150,10 @@ def dynamic_load_commands_in_register() -> (
                 command_or_command_info, BaseOptimumCLICommand
             ):
                 raise ValueError(
-                    f"The command at index {command_idx} in {filename} is not of the right type: {type(command_or_command_info)}."
+                    f"The command at index {command_idx} in the `commands.register` module of {dist_name} is not of the right type: {type(command_or_command_info)}."
                 )
             commands_to_register.append((command_or_command_info, parent_command_cls))
+
     return commands_to_register
 
 
@@ -181,9 +189,6 @@ def main():
     for subcommand_cls in OPTIMUM_CLI_ROOT_SUBCOMMANDS:
         register_optimum_cli_subcommand(subcommand_cls, parent_command=root)
 
-    # 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(