From a13d2fbe97386cd38ce41b9564e14ed062505827 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Wed, 10 Dec 2025 16:09:18 +1100 Subject: [PATCH 01/17] wip: using threads to build command table --- .vscode/launch.json | 2 +- src/azure-cli-core/azure/cli/core/__init__.py | 56 ++++++++++++------- .../azure/cli/core/commands/__init__.py | 1 + 3 files changed, 39 insertions(+), 20 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index e7ed7a11353..88410318724 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -23,7 +23,7 @@ "program": "${workspaceRoot}/src/azure-cli/azure/cli/__main__.py", "cwd": "${workspaceRoot}", "args": [ - "--version" + "Version" ], "console": "externalTerminal", }, diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index ae0f253caee..6d9cbf163e1 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -9,6 +9,7 @@ import os import sys import timeit +from concurrent.futures import ThreadPoolExecutor from knack.cli import CLI from knack.commands import CLICommandsLoader @@ -258,31 +259,48 @@ def _update_command_table_from_modules(args, command_modules=None): logger.debug("Loading command modules:") logger.debug(self.header_mod) - for mod in [m for m in command_modules if m not in BLOCKED_MODS]: + print(f"*** Starting SUT***") + start_time_wrapper = timeit.default_timer() + + def load_module_threaded(mod): 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) + return (mod, module_command_table, module_group_table, elapsed_time, None) 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()) + return (mod, {}, {}, 0, ex) + + # Use ThreadPoolExecutor to load modules in parallel + with ThreadPoolExecutor(max_workers=4) as executor: + futures = [executor.submit(load_module_threaded, mod) + for mod in command_modules if mod not in BLOCKED_MODS] + + for future in futures: + mod, module_command_table, module_group_table, elapsed_time, error = future.result() + if error: + # 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, error) + telemetry.set_exception(exception=error, fault_type='module-load-error-' + mod, + summary='Error loading module: {}'.format(mod)) + logger.debug(traceback.format_exc()) + else: + 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) + + 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) + elapsed_time = timeit.default_timer() - start_time_wrapper + print(f"*** SUT operation *** took: {elapsed_time:.6f} seconds for {len(command_modules)} modules") # Summary line logger.debug(self.item_format_string, "Total ({})".format(count), cumulative_elapsed_time, diff --git a/src/azure-cli-core/azure/cli/core/commands/__init__.py b/src/azure-cli-core/azure/cli/core/commands/__init__.py index fb2a9a3dece..05d45cd27bf 100644 --- a/src/azure-cli-core/azure/cli/core/commands/__init__.py +++ b/src/azure-cli-core/azure/cli/core/commands/__init__.py @@ -517,6 +517,7 @@ def execute(self, args): args = _pre_command_table_create(self.cli_ctx, args) self.cli_ctx.raise_event(EVENT_INVOKER_PRE_CMD_TBL_CREATE, args=args) + # @NOTE: below takes forever: self.commands_loader.load_command_table(args) self.cli_ctx.raise_event(EVENT_INVOKER_PRE_CMD_TBL_TRUNCATE, load_cmd_tbl_func=self.commands_loader.load_command_table, args=args) From 8c25f0e432ec22b85100472d922880014a938536 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Wed, 10 Dec 2025 16:45:27 +1100 Subject: [PATCH 02/17] refactor: export to smaller methods --- src/azure-cli-core/azure/cli/core/__init__.py | 137 +++++++++++------- 1 file changed, 86 insertions(+), 51 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 6d9cbf163e1..9357fefebfd 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -10,6 +10,8 @@ import sys import timeit from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Dict, Any, Optional from knack.cli import CLI from knack.commands import CLICommandsLoader @@ -196,6 +198,15 @@ def _configure_style(self): format_styled_text.theme = theme +@dataclass +class ModuleLoadResult: + module_name: str + command_table: Dict[str, Any] + group_table: Dict[str, Any] + elapsed_time: float + error: Optional[Exception] = None + + class MainCommandsLoader(CLICommandsLoader): # Format string for pretty-print the command module table @@ -222,11 +233,11 @@ def load_command_table(self, args): import pkgutil import traceback from azure.cli.core.commands import ( - _load_module_command_loader, _load_extension_command_loader, BLOCKED_MODS, ExtensionCommandSource) + _load_extension_command_loader, 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_module_breaking_changes, import_extension_breaking_changes) + import_core_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. @@ -252,55 +263,11 @@ def _update_command_table_from_modules(args, command_modules=None): except ImportError as e: logger.warning(e) - count = 0 - cumulative_elapsed_time = 0 - cumulative_group_count = 0 - cumulative_command_count = 0 - logger.debug("Loading command modules:") - logger.debug(self.header_mod) - - print(f"*** Starting SUT***") - start_time_wrapper = timeit.default_timer() - - def load_module_threaded(mod): - try: - start_time = timeit.default_timer() - module_command_table, module_group_table = _load_module_command_loader(self, args, mod) - import_module_breaking_changes(mod) - elapsed_time = timeit.default_timer() - start_time - return (mod, module_command_table, module_group_table, elapsed_time, None) - except Exception as ex: # pylint: disable=broad-except - return (mod, {}, {}, 0, ex) - - # Use ThreadPoolExecutor to load modules in parallel - with ThreadPoolExecutor(max_workers=4) as executor: - futures = [executor.submit(load_module_threaded, mod) - for mod in command_modules if mod not in BLOCKED_MODS] - - for future in futures: - mod, module_command_table, module_group_table, elapsed_time, error = future.result() - if error: - # 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, error) - telemetry.set_exception(exception=error, fault_type='module-load-error-' + mod, - summary='Error loading module: {}'.format(mod)) - logger.debug(traceback.format_exc()) - else: - 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) - - 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) - elapsed_time = timeit.default_timer() - start_time_wrapper - print(f"*** SUT operation *** took: {elapsed_time:.6f} seconds for {len(command_modules)} modules") + results = self._load_modules_threaded(args, command_modules) + + # @TODO: export to own method: + count, cumulative_elapsed_time, cumulative_group_count, cumulative_command_count = \ + self._process_results_with_timing(results, command_modules) # Summary line logger.debug(self.item_format_string, "Total ({})".format(count), cumulative_elapsed_time, @@ -577,6 +544,74 @@ 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_threaded(self, args, command_modules): + """Load command modules using ThreadPoolExecutor for parallel processing.""" + from azure.cli.core.commands import _load_module_command_loader, BLOCKED_MODS + from azure.cli.core.breaking_change import import_module_breaking_changes + + def load_single_module(mod): + try: + start_time = timeit.default_timer() + module_command_table, module_group_table = _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) + except Exception as ex: # pylint: disable=broad-except + return ModuleLoadResult(mod, {}, {}, 0, ex) + + with ThreadPoolExecutor(max_workers=4) as executor: + futures = [executor.submit(load_single_module, mod) + for mod in command_modules if mod not in BLOCKED_MODS] + return [future.result() for future in futures] + + def _handle_module_load_error(self, result): + """Handle errors that occurred during module loading.""" + import traceback + from azure.cli.core import telemetry + + # Changing this error message requires updating CI script that checks for failed module loading. + 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)) + logger.debug(traceback.format_exc()) + + def _process_successful_load(self, result): + """Process successfully loaded module results.""" + # Set command source for all commands in the module + for cmd in result.command_table.values(): + cmd.command_source = result.module_name + + # Update main command and group tables + self.command_table.update(result.command_table) + self.command_group_table.update(result.group_table) + + # Log the results + 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, command_modules): + """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: From a53f78af5a75f3987f1241338693bb373b13f093 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Thu, 11 Dec 2025 14:17:56 +1100 Subject: [PATCH 03/17] refactor: refactor smaller methods and add timeout for deadlocks --- src/azure-cli-core/azure/cli/core/__init__.py | 71 ++++++++++++------- 1 file changed, 44 insertions(+), 27 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 9357fefebfd..c900311100e 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -10,7 +10,6 @@ import sys import timeit from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass from typing import Dict, Any, Optional from knack.cli import CLI @@ -198,13 +197,13 @@ def _configure_style(self): format_styled_text.theme = theme -@dataclass class ModuleLoadResult: - module_name: str - command_table: Dict[str, Any] - group_table: Dict[str, Any] - elapsed_time: float - error: Optional[Exception] = None + def __init__(self, module_name, command_table, group_table, elapsed_time, error=None): + self.module_name = module_name + self.command_table = command_table + self.group_table = group_table + self.elapsed_time = elapsed_time + self.error = error class MainCommandsLoader(CLICommandsLoader): @@ -253,7 +252,6 @@ def _update_command_table_from_modules(args, command_modules=None): if command_modules is not None: command_modules.extend(ALWAYS_LOADED_MODULES) else: - # Perform module discovery command_modules = [] try: mods_ns_pkg = import_module('azure.cli.command_modules') @@ -263,7 +261,7 @@ def _update_command_table_from_modules(args, command_modules=None): except ImportError as e: logger.warning(e) - results = self._load_modules_threaded(args, command_modules) + results = self._load_modules(args, command_modules) # @TODO: export to own method: count, cumulative_elapsed_time, cumulative_group_count, cumulative_command_count = \ @@ -544,25 +542,44 @@ 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_threaded(self, args, command_modules): - """Load command modules using ThreadPoolExecutor for parallel processing.""" - from azure.cli.core.commands import _load_module_command_loader, BLOCKED_MODS - from azure.cli.core.breaking_change import import_module_breaking_changes - - def load_single_module(mod): - try: - start_time = timeit.default_timer() - module_command_table, module_group_table = _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) - except Exception as ex: # pylint: disable=broad-except - return ModuleLoadResult(mod, {}, {}, 0, ex) - + def _load_modules(self, args, command_modules): + """Load command modules using ThreadPoolExecutor with timeout protection.""" + from azure.cli.core.commands import BLOCKED_MODS + import concurrent.futures + + results = [] with ThreadPoolExecutor(max_workers=4) as executor: - futures = [executor.submit(load_single_module, mod) - for mod in command_modules if mod not in BLOCKED_MODS] - return [future.result() for future in futures] + 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, timeout=60): + try: + # @NOTE: Timeout to counteract deadlocks, but how to test? + result = future.result(timeout=30) + results.append(result) + except concurrent.futures.TimeoutError: + mod = future_to_module[future] + logger.warning("Module '%s' load timeout", mod) + results.append(ModuleLoadResult(mod, {}, {}, 0, + Exception(f"Module '{mod}' load timeout"))) + except Exception as ex: + mod = future_to_module[future] + logger.warning("Module '%s' load failed: %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 + try: + start_time = timeit.default_timer() + module_command_table, module_group_table = _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) + except Exception as ex: # pylint: disable=broad-except + return ModuleLoadResult(mod, {}, {}, 0, ex) def _handle_module_load_error(self, result): """Handle errors that occurred during module loading.""" From f3209600d4c68ac3c66b6ede68ef8ad1dc9482b8 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Thu, 11 Dec 2025 14:58:33 +1100 Subject: [PATCH 04/17] refactor: use constants for magic numbers --- src/azure-cli-core/azure/cli/core/__init__.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index c900311100e..85b98972dc4 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -34,6 +34,8 @@ 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'] +MODULE_LOAD_TIMEOUT_SECONDS = 30 +MAX_WORKER_THREAD_COUNT = 4 def _configure_knack(): @@ -548,18 +550,18 @@ def _load_modules(self, args, command_modules): import concurrent.futures results = [] - with ThreadPoolExecutor(max_workers=4) as executor: + 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, timeout=60): + for future in future_to_module: try: - # @NOTE: Timeout to counteract deadlocks, but how to test? - result = future.result(timeout=30) + 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", mod) + logger.warning("Module '%s' load timeout after %s seconds", mod, MODULE_LOAD_TIMEOUT_SECONDS) + future.cancel() results.append(ModuleLoadResult(mod, {}, {}, 0, Exception(f"Module '{mod}' load timeout"))) except Exception as ex: From a4f239825214f3203e77223c45a220f532069cda Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Thu, 11 Dec 2025 15:27:39 +1100 Subject: [PATCH 05/17] refactor: remove comments --- src/azure-cli-core/azure/cli/core/__init__.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 85b98972dc4..a7fb7505d43 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -265,7 +265,6 @@ def _update_command_table_from_modules(args, command_modules=None): results = self._load_modules(args, command_modules) - # @TODO: export to own method: count, cumulative_elapsed_time, cumulative_group_count, cumulative_command_count = \ self._process_results_with_timing(results, command_modules) # Summary line @@ -587,25 +586,21 @@ def _handle_module_load_error(self, result): """Handle errors that occurred during module loading.""" import traceback from azure.cli.core import telemetry - - # Changing this error message requires updating CI script that checks for failed module loading. + logger.error("Error loading command module '%s': %s", result.module_name, result.error) - telemetry.set_exception(exception=result.error, + telemetry.set_exception(exception=result.error, fault_type='module-load-error-' + result.module_name, summary='Error loading module: {}'.format(result.module_name)) logger.debug(traceback.format_exc()) def _process_successful_load(self, result): """Process successfully loaded module results.""" - # Set command source for all commands in the module for cmd in result.command_table.values(): cmd.command_source = result.module_name - - # Update main command and group tables + self.command_table.update(result.command_table) self.command_group_table.update(result.group_table) - - # Log the results + logger.debug(self.item_format_string, result.module_name, result.elapsed_time, len(result.group_table), len(result.command_table)) From 430117c0a42ab137aae88c3b3e1d9712633cef77 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 15 Dec 2025 13:21:49 +1100 Subject: [PATCH 06/17] test: add test for command table integrity --- .../tests/test_command_table_integrity.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 src/azure-cli-core/azure/cli/core/tests/test_command_table_integrity.py diff --git a/src/azure-cli-core/azure/cli/core/tests/test_command_table_integrity.py b/src/azure-cli-core/azure/cli/core/tests/test_command_table_integrity.py new file mode 100644 index 00000000000..28fa52dd355 --- /dev/null +++ b/src/azure-cli-core/azure/cli/core/tests/test_command_table_integrity.py @@ -0,0 +1,57 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import unittest + +from azure.cli.core.mock import DummyCli +from azure.cli.core import MainCommandsLoader + + +class CommandTableIntegrityTest(unittest.TestCase): + + def setUp(self): + self.cli_ctx = DummyCli() + + def test_command_table_integrity(self): + """Test command table loading produces valid, complete results.""" + + # Load command table using current implementation + loader = MainCommandsLoader(self.cli_ctx) + loader.load_command_table([]) + + # Test invariants that should always hold: + + # 1. No corruption/duplicates + command_names = list(loader.command_table.keys()) + unique_command_names = set(command_names) + self.assertEqual(len(unique_command_names), len(command_names), "No duplicate commands") + + # 2. Core functionality exists (high-level groups that should always exist) + core_groups = ['vm', 'network', 'resource', 'account', 'group'] + existing_groups = {cmd.split()[0] for cmd in loader.command_table.keys() if ' ' in cmd} + missing_core = [group for group in core_groups if group not in existing_groups] + self.assertEqual(len(missing_core), 0, f"Missing core command groups: {missing_core}") + + # 3. Structural integrity + commands_without_source = [] + for cmd_name, cmd_obj in loader.command_table.items(): + if not hasattr(cmd_obj, 'command_source') or not cmd_obj.command_source: + commands_without_source.append(cmd_name) + + self.assertEqual(len(commands_without_source), 0, + f"Commands missing source: {commands_without_source[:5]}...") + + # 4. Basic sanity - we loaded SOMETHING + self.assertGreater(len(loader.command_table), 0, "Commands were loaded") + self.assertGreater(len(loader.command_group_table), 0, "Groups were loaded") + + # 5. Verify core groups are properly represented + found_core_groups = sorted(existing_groups & set(core_groups)) + self.assertGreaterEqual(len(found_core_groups), 3, + f"At least 3 core command groups should be present, found: {found_core_groups}") + + +if __name__ == '__main__': + unittest.main() From ed7147ce9debac0b750c62ec35ebcd4b79f1e7f1 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 15 Dec 2025 13:23:14 +1100 Subject: [PATCH 07/17] refactor: remove comments --- src/azure-cli-core/azure/cli/core/commands/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/azure-cli-core/azure/cli/core/commands/__init__.py b/src/azure-cli-core/azure/cli/core/commands/__init__.py index 05d45cd27bf..fb2a9a3dece 100644 --- a/src/azure-cli-core/azure/cli/core/commands/__init__.py +++ b/src/azure-cli-core/azure/cli/core/commands/__init__.py @@ -517,7 +517,6 @@ def execute(self, args): args = _pre_command_table_create(self.cli_ctx, args) self.cli_ctx.raise_event(EVENT_INVOKER_PRE_CMD_TBL_CREATE, args=args) - # @NOTE: below takes forever: self.commands_loader.load_command_table(args) self.cli_ctx.raise_event(EVENT_INVOKER_PRE_CMD_TBL_TRUNCATE, load_cmd_tbl_func=self.commands_loader.load_command_table, args=args) From 3f51c599b068d26972ac295a4c18048c8b4039cd Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 15 Dec 2025 15:17:11 +1100 Subject: [PATCH 08/17] refactor: restore launch.json --- .vscode/launch.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index 88410318724..e7ed7a11353 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -23,7 +23,7 @@ "program": "${workspaceRoot}/src/azure-cli/azure/cli/__main__.py", "cwd": "${workspaceRoot}", "args": [ - "Version" + "--version" ], "console": "externalTerminal", }, From b7ccbed4f7e0b4f8519c2385720cb5c84263a420 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 15 Dec 2025 15:57:04 +1100 Subject: [PATCH 09/17] refactor: linting issues --- src/azure-cli-core/azure/cli/core/__init__.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index a7fb7505d43..f3df38c0ca6 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -9,8 +9,8 @@ import os import sys import timeit +import concurrent.futures from concurrent.futures import ThreadPoolExecutor -from typing import Dict, Any, Optional from knack.cli import CLI from knack.commands import CLICommandsLoader @@ -199,7 +199,7 @@ def _configure_style(self): format_styled_text.theme = theme -class ModuleLoadResult: +class ModuleLoadResult: # pylint: disable=too-few-public-methods def __init__(self, module_name, command_table, group_table, elapsed_time, error=None): self.module_name = module_name self.command_table = command_table @@ -266,7 +266,7 @@ def _update_command_table_from_modules(args, command_modules=None): results = self._load_modules(args, command_modules) count, cumulative_elapsed_time, cumulative_group_count, cumulative_command_count = \ - self._process_results_with_timing(results, command_modules) + self._process_results_with_timing(results) # Summary line logger.debug(self.item_format_string, "Total ({})".format(count), cumulative_elapsed_time, @@ -551,7 +551,7 @@ def _load_modules(self, args, command_modules): 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 mod in command_modules if mod not in BLOCKED_MODS} for future in future_to_module: try: @@ -562,8 +562,8 @@ def _load_modules(self, args, command_modules): logger.warning("Module '%s' load timeout after %s seconds", mod, MODULE_LOAD_TIMEOUT_SECONDS) future.cancel() results.append(ModuleLoadResult(mod, {}, {}, 0, - Exception(f"Module '{mod}' load timeout"))) - except Exception as ex: + 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)) @@ -602,9 +602,9 @@ def _process_successful_load(self, result): 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)) + len(result.group_table), len(result.command_table)) - def _process_results_with_timing(self, results, command_modules): + 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) From a472dedbffb1f4791f90273c8d917b312a7f8203 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 15 Dec 2025 15:59:51 +1100 Subject: [PATCH 10/17] refactor: flake issues --- src/azure-cli-core/azure/cli/core/__init__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index f3df38c0ca6..381eb0fd2a8 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -234,7 +234,7 @@ def load_command_table(self, args): import pkgutil import traceback from azure.cli.core.commands import ( - _load_extension_command_loader, ExtensionCommandSource) + _load_extension_command_loader, ExtensionCommandSource) from azure.cli.core.extension import ( get_extensions, get_extension_path, get_extension_modname) from azure.cli.core.breaking_change import ( @@ -562,7 +562,7 @@ def _load_modules(self, args, command_modules): logger.warning("Module '%s' load timeout after %s seconds", mod, MODULE_LOAD_TIMEOUT_SECONDS) future.cancel() results.append(ModuleLoadResult(mod, {}, {}, 0, - Exception(f"Module '{mod}' load timeout"))) + 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) @@ -589,8 +589,8 @@ def _handle_module_load_error(self, result): 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)) + fault_type='module-load-error-' + result.module_name, + summary='Error loading module: {}'.format(result.module_name)) logger.debug(traceback.format_exc()) def _process_successful_load(self, result): From 82c2cb56b28c61d04f4b34bea73c6dcfe8e6442d Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 15 Dec 2025 16:05:40 +1100 Subject: [PATCH 11/17] refactor: flake issues --- src/azure-cli-core/azure/cli/core/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 381eb0fd2a8..1a37f3da993 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -562,7 +562,7 @@ def _load_modules(self, args, command_modules): logger.warning("Module '%s' load timeout after %s seconds", mod, MODULE_LOAD_TIMEOUT_SECONDS) future.cancel() results.append(ModuleLoadResult(mod, {}, {}, 0, - Exception(f"Module '{mod}' load timeout"))) + 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) From 89a8e32b41e2741865edd2fa65803ca92bc509e4 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 15 Dec 2025 17:10:33 +1100 Subject: [PATCH 12/17] refactor: fix lint issue (imports) --- src/azure-cli-core/azure/cli/core/__init__.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 1a37f3da993..6f14578eb05 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -9,8 +9,7 @@ import os import sys import timeit -import concurrent.futures -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor, TimeoutError from knack.cli import CLI from knack.commands import CLICommandsLoader @@ -546,7 +545,6 @@ def load_arguments(self, command=None): def _load_modules(self, args, command_modules): """Load command modules using ThreadPoolExecutor with timeout protection.""" from azure.cli.core.commands import BLOCKED_MODS - import concurrent.futures results = [] with ThreadPoolExecutor(max_workers=MAX_WORKER_THREAD_COUNT) as executor: @@ -557,7 +555,7 @@ def _load_modules(self, args, command_modules): try: result = future.result(timeout=MODULE_LOAD_TIMEOUT_SECONDS) results.append(result) - except concurrent.futures.TimeoutError: + except TimeoutError: mod = future_to_module[future] logger.warning("Module '%s' load timeout after %s seconds", mod, MODULE_LOAD_TIMEOUT_SECONDS) future.cancel() From 6580aa55d5972fc7b69abda9e9723bd49ab4f509 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Tue, 16 Dec 2025 09:42:58 +1100 Subject: [PATCH 13/17] refactor: concurrent.futures Timeout import --- src/azure-cli-core/azure/cli/core/__init__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 6f14578eb05..29eda72fc27 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -9,7 +9,8 @@ import os import sys import timeit -from concurrent.futures import ThreadPoolExecutor, TimeoutError +import concurrent.futures +from concurrent.futures import ThreadPoolExecutor from knack.cli import CLI from knack.commands import CLICommandsLoader @@ -555,7 +556,7 @@ def _load_modules(self, args, command_modules): try: result = future.result(timeout=MODULE_LOAD_TIMEOUT_SECONDS) results.append(result) - except TimeoutError: + except concurrent.futures.TimeoutError: mod = future_to_module[future] logger.warning("Module '%s' load timeout after %s seconds", mod, MODULE_LOAD_TIMEOUT_SECONDS) future.cancel() From 0f57d7a530f5b34eed872a8aa3d48a46841cceae Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Thu, 18 Dec 2025 12:59:42 +1100 Subject: [PATCH 14/17] fix: address co-pilot suggestions --- src/azure-cli-core/azure/cli/core/__init__.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 29eda72fc27..3174e6d7e00 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -34,7 +34,9 @@ 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 @@ -200,12 +202,13 @@ def _configure_style(self): class ModuleLoadResult: # pylint: disable=too-few-public-methods - def __init__(self, module_name, command_table, group_table, elapsed_time, error=None): + def __init__(self, module_name, command_table, group_table, elapsed_time, error=None, traceback_str=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 class MainCommandsLoader(CLICommandsLoader): @@ -552,26 +555,30 @@ def _load_modules(self, args, command_modules): 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 future_to_module: + 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) - future.cancel() 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: + 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 = _load_module_command_loader(self, args, mod) @@ -579,18 +586,19 @@ def _load_single_module(self, mod, args): elapsed_time = timeit.default_timer() - start_time return ModuleLoadResult(mod, module_command_table, module_group_table, elapsed_time) except Exception as ex: # pylint: disable=broad-except - return ModuleLoadResult(mod, {}, {}, 0, ex) + 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.""" - import traceback 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)) - logger.debug(traceback.format_exc()) + if result.traceback_str: + logger.debug(result.traceback_str) def _process_successful_load(self, result): """Process successfully loaded module results.""" From 25d5d6a840ac07d9d97aa0420d191cd5c3928e0a Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Thu, 18 Dec 2025 13:32:43 +1100 Subject: [PATCH 15/17] fix: address lint error --- src/azure-cli-core/azure/cli/core/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 3174e6d7e00..70a0e78b3ab 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -568,7 +568,7 @@ def _load_modules(self, args, command_modules): mod = future_to_module[future] logger.warning("Module '%s' load failed: %s", mod, ex) results.append(ModuleLoadResult(mod, {}, {}, 0, ex)) - except Exception as 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)) From d02b59f4eef6165c7a79d313a8696cc1974f5543 Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 22 Dec 2025 15:44:20 +1100 Subject: [PATCH 16/17] feature: address threading race conditions concerns --- src/azure-cli-core/azure/cli/core/__init__.py | 15 +++++++++++---- .../azure/cli/core/commands/__init__.py | 13 ++++--------- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/__init__.py b/src/azure-cli-core/azure/cli/core/__init__.py index 70a0e78b3ab..19f844c6b4a 100644 --- a/src/azure-cli-core/azure/cli/core/__init__.py +++ b/src/azure-cli-core/azure/cli/core/__init__.py @@ -202,13 +202,14 @@ def _configure_style(self): class ModuleLoadResult: # pylint: disable=too-few-public-methods - def __init__(self, module_name, command_table, group_table, elapsed_time, error=None, traceback_str=None): + 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): @@ -343,7 +344,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) @@ -581,10 +582,10 @@ def _load_single_module(self, mod, args): import traceback try: start_time = timeit.default_timer() - module_command_table, module_group_table = _load_module_command_loader(self, args, mod) + 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) + 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) @@ -602,6 +603,12 @@ def _handle_module_load_error(self, result): 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 diff --git a/src/azure-cli-core/azure/cli/core/commands/__init__.py b/src/azure-cli-core/azure/cli/core/commands/__init__.py index fb2a9a3dece..41916a4c58d 100644 --- a/src/azure-cli-core/azure/cli/core/commands/__init__.py +++ b/src/azure-cli-core/azure/cli/core/commands/__init__.py @@ -1131,22 +1131,17 @@ 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) - return command_table, command_loader.command_group_table + + group_table = command_loader.command_group_table if command_loader else {} + return command_table, group_table, command_loader def _load_extension_command_loader(loader, args, ext): From 4f890c80344e01b4b6524a0d631c8c50e1d11f4e Mon Sep 17 00:00:00 2001 From: Daniel Languiller Date: Mon, 22 Dec 2025 16:18:36 +1100 Subject: [PATCH 17/17] fix: fix test mocks to return correct params --- .../azure/cli/core/tests/test_command_registration.py | 2 +- src/azure-cli-core/azure/cli/core/tests/test_parser.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/tests/test_command_registration.py b/src/azure-cli-core/azure/cli/core/tests/test_command_registration.py index 21d03d47b5e..41cf1ea1af4 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_command_registration.py +++ b/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 + return module_command_table, command_loader.command_group_table, command_loader expected_command_index = {'hello': ['azure.cli.command_modules.hello', 'azext_hello2', 'azext_hello1'], 'extra': ['azure.cli.command_modules.extra']} diff --git a/src/azure-cli-core/azure/cli/core/tests/test_parser.py b/src/azure-cli-core/azure/cli/core/tests/test_parser.py index bff9c8ddc62..d584aafe383 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_parser.py +++ b/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 + return module_command_table, command_loader.command_group_table, command_loader @mock.patch('importlib.import_module', _mock_import_lib) @mock.patch('pkgutil.iter_modules', _mock_iter_modules)