Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
DanielMicrosoft wants to merge 1 commit into from
+47 −181
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
150 changes: 34 additions & 116 deletions src/azure-cli-core/azure/cli/core/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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():
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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:

Expand Down
13 changes: 9 additions & 4 deletions src/azure-cli-core/azure/cli/core/commands/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link

Copilot AI Jan 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Grammar issue in TODO comment. The phrase "If desired to for extension to patch module" should be corrected to either "If desired for extension to patch module" or "If it's desired to allow extension to patch module".

Suggested change
# TODO: If desired to for extension to patch module, this can be uncommented
# TODO: If desired for extension to patch module, this can be uncommented

Copilot uses AI. Check for mistakes.
# 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
Copy link

Copilot AI Jan 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The variable command_loader is only defined inside the if loader_cls: block (line 1139), but is referenced unconditionally in the return statement at line 1152. If loader_cls is None (when the module doesn't have COMMAND_LOADER_CLS or get_command_loader), this will raise an UnboundLocalError.

The correct fix should initialize command_loader = None before the if statement, and then return an empty dictionary for the command_group_table when command_loader is None, similar to what was removed:

group_table = command_loader.command_group_table if command_loader else {}
return command_table, group_table

Copilot uses AI. Check for mistakes.


def _load_extension_command_loader(loader, args, ext):
Expand Down
6 changes: 3 additions & 3 deletions src/azure-cli-core/azure/cli/core/tests/test_command_registration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link

Copilot AI Jan 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The variable command_loader may be undefined if the for loop doesn't execute (e.g., if no loader matches the name filter or if command_loaders is empty). While the current test setup ensures at least one loader is processed, this is fragile. Consider initializing command_loader = None before the loop and returning an empty dict for command_group_table when it's None, similar to the fix needed in the main implementation.

Copilot uses AI. Check for mistakes.

expected_command_index = {'hello': ['azure.cli.command_modules.hello', 'azext_hello2', 'azext_hello1'],
'extra': ['azure.cli.command_modules.extra']}
Expand Down Expand Up @@ -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'])
Copy link

Copilot AI Jan 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changing from assertSetEqual(set(cmd_tbl), ...) to assertEqual(list(cmd_tbl), ...) makes the test depend on the insertion order of commands in the dictionary. While Python 3.7+ guarantees dict order preservation, this change makes the test more brittle and less robust to implementation changes that might affect the loading order of commands. Unless there's a specific requirement to test the ordering, using assertSetEqual (or assertEqual(set(cmd_tbl), set(...))) would be more appropriate for testing command presence regardless of order.

Copilot uses AI. Check for mistakes.

# 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'])
Copy link

Copilot AI Jan 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changing from assertSetEqual(set(cmd_tbl), ...) to assertEqual(list(cmd_tbl), ...) makes the test depend on the insertion order of commands in the dictionary. While Python 3.7+ guarantees dict order preservation, this change makes the test more brittle and less robust to implementation changes that might affect the loading order of commands. Unless there's a specific requirement to test the ordering, using assertSetEqual (or assertEqual(set(cmd_tbl), set(...))) would be more appropriate for testing command presence regardless of order.

Copilot uses AI. Check for mistakes.

# Test command index is used by command with positional argument
cmd_tbl = loader.load_command_table(["extra", "final", "positional_argument2"])
Expand Down
57 changes: 0 additions & 57 deletions src/azure-cli-core/azure/cli/core/tests/test_command_table_integrity.py
View file
Open in desktop

This file was deleted.

2 changes: 1 addition & 1 deletion src/azure-cli-core/azure/cli/core/tests/test_parser.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link

Copilot AI Jan 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The variable command_loader may be undefined if the for loop doesn't execute (e.g., if command_loaders is empty). While the current test setup ensures the loop always executes, this is fragile. Consider initializing command_loader = None before the loop and returning an empty dict for command_group_table when it's None, similar to the fix needed in the main implementation.

Copilot uses AI. Check for mistakes.

@mock.patch('importlib.import_module', _mock_import_lib)
@mock.patch('pkgutil.iter_modules', _mock_iter_modules)
Expand Down
Loading