{Core} Revert #32518 using threads to build command table for performance

src/azure-cli-core/azure/cli/core/__init__.py

@@ -9,8 +9,6 @@
 import os
 import sys
 import timeit
-import concurrent.futures
-from concurrent.futures import ThreadPoolExecutor
 
 from knack.cli import CLI
 from knack.commands import CLICommandsLoader
@@ -36,10 +34,6 @@
 ALWAYS_LOADED_MODULES = []
 # Extensions that will always be loaded if installed. They don't expose commands but hook into CLI core.
 ALWAYS_LOADED_EXTENSIONS = ['azext_ai_examples', 'azext_next']
-# Timeout (in seconds) for loading a single module. Acts as a safety valve to prevent indefinite hangs
-MODULE_LOAD_TIMEOUT_SECONDS = 30
-# Maximum number of worker threads for parallel module loading.
-MAX_WORKER_THREAD_COUNT = 4
 
 
 def _configure_knack():
@@ -203,17 +197,6 @@ def _configure_style(self):
         format_styled_text.theme = theme
 
 
-class ModuleLoadResult:  # pylint: disable=too-few-public-methods
-    def __init__(self, module_name, command_table, group_table, elapsed_time, error=None, traceback_str=None, command_loader=None):
-        self.module_name = module_name
-        self.command_table = command_table
-        self.group_table = group_table
-        self.elapsed_time = elapsed_time
-        self.error = error
-        self.traceback_str = traceback_str
-        self.command_loader = command_loader
-
-
 class MainCommandsLoader(CLICommandsLoader):
 
     # Format string for pretty-print the command module table
@@ -258,11 +241,11 @@ def load_command_table(self, args):
         import pkgutil
         import traceback
         from azure.cli.core.commands import (
-            _load_extension_command_loader, ExtensionCommandSource)
+            _load_module_command_loader, _load_extension_command_loader, BLOCKED_MODS, ExtensionCommandSource)
         from azure.cli.core.extension import (
             get_extensions, get_extension_path, get_extension_modname)
         from azure.cli.core.breaking_change import (
-            import_core_breaking_changes, import_extension_breaking_changes)
+            import_core_breaking_changes, import_module_breaking_changes, import_extension_breaking_changes)
 
         def _update_command_table_from_modules(args, command_modules=None):
             """Loads command tables from modules and merge into the main command table.
@@ -290,10 +273,38 @@ def _update_command_table_from_modules(args, command_modules=None):
                 except ImportError as e:
                     logger.warning(e)
 
-            results = self._load_modules(args, command_modules)
+            count = 0
+            cumulative_elapsed_time = 0
+            cumulative_group_count = 0
+            cumulative_command_count = 0
+            logger.debug("Loading command modules:")
+            logger.debug(self.header_mod)
 
-            count, cumulative_elapsed_time, cumulative_group_count, cumulative_command_count = \
-                self._process_results_with_timing(results)
+            for mod in [m for m in command_modules if m not in BLOCKED_MODS]:
+                try:
+                    start_time = timeit.default_timer()
+                    module_command_table, module_group_table = _load_module_command_loader(self, args, mod)
+                    import_module_breaking_changes(mod)
+                    for cmd in module_command_table.values():
+                        cmd.command_source = mod
+                    self.command_table.update(module_command_table)
+                    self.command_group_table.update(module_group_table)
+
+                    elapsed_time = timeit.default_timer() - start_time
+                    logger.debug(self.item_format_string, mod, elapsed_time,
+                                 len(module_group_table), len(module_command_table))
+                    count += 1
+                    cumulative_elapsed_time += elapsed_time
+                    cumulative_group_count += len(module_group_table)
+                    cumulative_command_count += len(module_command_table)
+                except Exception as ex:  # pylint: disable=broad-except
+                    # Changing this error message requires updating CI script that checks for failed
+                    # module loading.
+                    from azure.cli.core import telemetry
+                    logger.error("Error loading command module '%s': %s", mod, ex)
+                    telemetry.set_exception(exception=ex, fault_type='module-load-error-' + mod,
+                                            summary='Error loading module: {}'.format(mod))
+                    logger.debug(traceback.format_exc())
             # Summary line
             logger.debug(self.item_format_string,
                          "Total ({})".format(count), cumulative_elapsed_time,
@@ -367,7 +378,7 @@ def _filter_modname(extensions):
                         # from an extension requires this map to be up-to-date.
                         # self._mod_to_ext_map[ext_mod] = ext_name
                         start_time = timeit.default_timer()
-                        extension_command_table, extension_group_table, _ = \
+                        extension_command_table, extension_group_table = \
                             _load_extension_command_loader(self, args, ext_mod)
                         import_extension_breaking_changes(ext_mod)
 
@@ -576,99 +587,6 @@ def load_arguments(self, command=None):
                 self.extra_argument_registry.update(loader.extra_argument_registry)
                 loader._update_command_definitions()  # pylint: disable=protected-access
 
-    def _load_modules(self, args, command_modules):
-        """Load command modules using ThreadPoolExecutor with timeout protection."""
-        from azure.cli.core.commands import BLOCKED_MODS
-
-        results = []
-        with ThreadPoolExecutor(max_workers=MAX_WORKER_THREAD_COUNT) as executor:
-            future_to_module = {executor.submit(self._load_single_module, mod, args): mod
-                                for mod in command_modules if mod not in BLOCKED_MODS}
-
-            for future in concurrent.futures.as_completed(future_to_module):
-                try:
-                    result = future.result(timeout=MODULE_LOAD_TIMEOUT_SECONDS)
-                    results.append(result)
-                except concurrent.futures.TimeoutError:
-                    mod = future_to_module[future]
-                    logger.warning("Module '%s' load timeout after %s seconds", mod, MODULE_LOAD_TIMEOUT_SECONDS)
-                    results.append(ModuleLoadResult(mod, {}, {}, 0,
-                                                    Exception(f"Module '{mod}' load timeout")))
-                except (ImportError, AttributeError, TypeError, ValueError) as ex:
-                    mod = future_to_module[future]
-                    logger.warning("Module '%s' load failed: %s", mod, ex)
-                    results.append(ModuleLoadResult(mod, {}, {}, 0, ex))
-                except Exception as ex:  # pylint: disable=broad-exception-caught
-                    mod = future_to_module[future]
-                    logger.warning("Module '%s' load failed with unexpected exception: %s", mod, ex)
-                    results.append(ModuleLoadResult(mod, {}, {}, 0, ex))
-
-        return results
-
-    def _load_single_module(self, mod, args):
-        from azure.cli.core.breaking_change import import_module_breaking_changes
-        from azure.cli.core.commands import _load_module_command_loader
-        import traceback
-        try:
-            start_time = timeit.default_timer()
-            module_command_table, module_group_table, command_loader = _load_module_command_loader(self, args, mod)
-            import_module_breaking_changes(mod)
-            elapsed_time = timeit.default_timer() - start_time
-            return ModuleLoadResult(mod, module_command_table, module_group_table, elapsed_time, command_loader=command_loader)
-        except Exception as ex:  # pylint: disable=broad-except
-            tb_str = traceback.format_exc()
-            return ModuleLoadResult(mod, {}, {}, 0, ex, tb_str)
-
-    def _handle_module_load_error(self, result):
-        """Handle errors that occurred during module loading."""
-        from azure.cli.core import telemetry
-
-        logger.error("Error loading command module '%s': %s", result.module_name, result.error)
-        telemetry.set_exception(exception=result.error,
-                                fault_type='module-load-error-' + result.module_name,
-                                summary='Error loading module: {}'.format(result.module_name))
-        if result.traceback_str:
-            logger.debug(result.traceback_str)
-
-    def _process_successful_load(self, result):
-        """Process successfully loaded module results."""
-        if result.command_loader:
-            self.loaders.append(result.command_loader)
-
-            for cmd in result.command_table:
-                self.cmd_to_loader_map[cmd] = [result.command_loader]
-
-        for cmd in result.command_table.values():
-            cmd.command_source = result.module_name
-
-        self.command_table.update(result.command_table)
-        self.command_group_table.update(result.group_table)
-
-        logger.debug(self.item_format_string, result.module_name, result.elapsed_time,
-                     len(result.group_table), len(result.command_table))
-
-    def _process_results_with_timing(self, results):
-        """Process pre-loaded module results with timing and progress reporting."""
-        logger.debug("Loading command modules:")
-        logger.debug(self.header_mod)
-
-        count = 0
-        cumulative_elapsed_time = 0
-        cumulative_group_count = 0
-        cumulative_command_count = 0
-
-        for result in results:
-            if result.error:
-                self._handle_module_load_error(result)
-            else:
-                self._process_successful_load(result)
-                count += 1
-                cumulative_elapsed_time += result.elapsed_time
-                cumulative_group_count += len(result.group_table)
-                cumulative_command_count += len(result.command_table)
-
-        return count, cumulative_elapsed_time, cumulative_group_count, cumulative_command_count
-
 
 class CommandIndex:
 

src/azure-cli-core/azure/cli/core/commands/__init__.py

@@ -1134,17 +1134,22 @@ def _load_command_loader(loader, args, name, prefix):
             logger.debug("Module '%s' is missing `get_command_loader` entry.", name)
 
     command_table = {}
-    command_loader = None
 
     if loader_cls:
         command_loader = loader_cls(cli_ctx=loader.cli_ctx)
+        loader.loaders.append(command_loader)  # This will be used by interactive
         if command_loader.supported_resource_type():
             command_table = command_loader.load_command_table(args)
+            if command_table:
+                for cmd in list(command_table.keys()):
+                    # TODO: If desired to for extension to patch module, this can be uncommented
+                    # if loader.cmd_to_loader_map.get(cmd):
+                    #    loader.cmd_to_loader_map[cmd].append(command_loader)
+                    # else:
+                    loader.cmd_to_loader_map[cmd] = [command_loader]
     else:
         logger.debug("Module '%s' is missing `COMMAND_LOADER_CLS` entry.", name)
-
-    group_table = command_loader.command_group_table if command_loader else {}
-    return command_table, group_table, command_loader
+    return command_table, command_loader.command_group_table
 
 
 def _load_extension_command_loader(loader, args, ext):

src/azure-cli-core/azure/cli/core/tests/test_command_registration.py

@@ -230,7 +230,7 @@ def load_command_table(self, args):
             if command_table:
                 module_command_table.update(command_table)
                 loader.loaders.append(command_loader)  # this will be used later by the load_arguments method
-        return module_command_table, command_loader.command_group_table, command_loader
+        return module_command_table, command_loader.command_group_table
 
     expected_command_index = {'hello': ['azure.cli.command_modules.hello', 'azext_hello2', 'azext_hello1'],
                               'extra': ['azure.cli.command_modules.extra']}
@@ -432,12 +432,12 @@ def test_command_index_positional_argument(self):
         # Test command index is built for command with positional argument
         cmd_tbl = loader.load_command_table(["extra", "extra", "positional_argument"])
         self.assertDictEqual(INDEX[CommandIndex._COMMAND_INDEX], self.expected_command_index)
-        self.assertSetEqual(set(cmd_tbl), {'hello mod-only', 'hello overridden', 'extra final', 'hello ext-only'})
+        self.assertEqual(list(cmd_tbl), ['hello mod-only', 'hello overridden', 'extra final', 'hello ext-only'])
 
         # Test command index is used by command with positional argument
         cmd_tbl = loader.load_command_table(["hello", "mod-only", "positional_argument"])
         self.assertDictEqual(INDEX[CommandIndex._COMMAND_INDEX], self.expected_command_index)
-        self.assertSetEqual(set(cmd_tbl), {'hello mod-only', 'hello overridden', 'hello ext-only'})
+        self.assertEqual(list(cmd_tbl), ['hello mod-only', 'hello overridden', 'hello ext-only'])
 
         # Test command index is used by command with positional argument
         cmd_tbl = loader.load_command_table(["extra", "final", "positional_argument2"])

src/azure-cli-core/azure/cli/core/tests/test_parser.py

@@ -188,7 +188,7 @@ def load_command_table(self, args):
             if command_table:
                 module_command_table.update(command_table)
                 loader.loaders.append(command_loader)  # this will be used later by the load_arguments method
-        return module_command_table, command_loader.command_group_table, command_loader
+        return module_command_table, command_loader.command_group_table
 
     @mock.patch('importlib.import_module', _mock_import_lib)
     @mock.patch('pkgutil.iter_modules', _mock_iter_modules)