Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
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
44 changes: 6 additions & 38 deletions src/azure-cli-core/azure/cli/core/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,10 +37,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 = 60
# Maximum number of worker threads for parallel module loading.
MAX_WORKER_THREAD_COUNT = 4


def _get_top_level_command(args):
Expand Down Expand Up @@ -664,42 +660,14 @@ def load_arguments(self, command=None):
loader._update_command_definitions() # pylint: disable=protected-access

def _load_modules(self, args, command_modules):
"""Load command modules using ThreadPoolExecutor with timeout protection."""
import concurrent.futures
from concurrent.futures import ThreadPoolExecutor
"""Load command modules sequentially."""
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}

try:
for future in concurrent.futures.as_completed(future_to_module, timeout=MODULE_LOAD_TIMEOUT_SECONDS):
try:
result = future.result()
results.append(result)
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))
except concurrent.futures.TimeoutError:
for future, mod in future_to_module.items():
if future.done():
try:
result = future.result()
results.append(result)
except Exception as ex: # pylint: disable=broad-exception-caught
logger.warning("Module '%s' load failed: %s", mod, ex)
results.append(ModuleLoadResult(mod, {}, {}, 0, ex))
else:
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")))
for mod in command_modules:
if mod in BLOCKED_MODS:
continue
results.append(self._load_single_module(mod, args))

return results

Expand Down Expand Up @@ -749,7 +717,7 @@ def _process_successful_load(self, result):

def _process_results_with_timing(self, results):
"""Process pre-loaded module results with timing and progress reporting."""
logger.debug("Loaded command modules in parallel:")
logger.debug("Loaded command modules:")
logger.debug(self.header_mod)

count = 0
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -275,14 +275,14 @@ def test_register_command_from_extension(self):
@mock.patch('azure.cli.core.commands._load_command_loader', _mock_load_command_loader)
@mock.patch('azure.cli.core.extension.get_extension_modname', _mock_get_extension_modname)
@mock.patch('azure.cli.core.extension.get_extensions', _mock_get_extensions)
def test_cmd_to_loader_map_populated_after_parallel_loading(self):
def test_cmd_to_loader_map_populated_after_command_loading(self):
"""
Validates that all commands in command_table have corresponding entries in cmd_to_loader_map.
"""
cli = DummyCli()
loader = cli.commands_loader

# Load all commands (triggers parallel module loading)
# Load all commands.
cmd_tbl = loader.load_command_table(None)

Comment thread
DanielMicrosoft marked this conversation as resolved.
Outdated
# Verify EVERY command in command_table has an entry in cmd_to_loader_map
Expand Down
Loading