Deferring Imports to improve hyp --help command

[mohamedzeidan2021]

These code changes implement lazy loading and deferred imports to significantly improve CLI startup performance. The changes reduce initial import time by deferring heavy dependencies (k7s, boto3, telemetry) until they are actually needed during command execution

Problem Statement

the CLI had slow startup times due to eager loading of heavy dependencies at import time

• kubernetes library (~300ms import time)
• boto3 and AWS SDK components (~318ms import time)
• telemetry and other utility modules

Total startup time of hyp --help was ~920ms, impacting user experience (>2 seconds for some users)

Solution Overview

Implemented lazy loading architecture that defers imports until command execution:

• Package-level lazy loading - Removed star imports from __init__.py
• Command module lazy loading - Deferred heavy imports in command modules
• Dynamic CLI registration - Load and register commands only when accessed
• Optimized help generation - Static help text without module loading

Performance Impact

• CLI Startup time: Reduced from ~920ms to ~36ms (25x faster)
• Help Command: Maintains functionality while avoiding heavy imports
• Command Execution: No performance impact because imports happen on first use

Detailed Changes

1. Package-Level Lazy Loading __init__.py

Before
from .common.utils import * # Loaded kubernetes, boto3 at import time

After

# Lazy-loading wrapper functions for public API
def list_clusters(*args, **kwargs):
    from .common.utils import list_clusters as _list_clusters
    return _list_clusters(*args, **kwargs)

def set_cluster_context(*args, **kwargs):
    from .common.utils import set_cluster_context as _set_cluster_context
    return _set_cluster_context(*args, **kwargs)

def get_cluster_context(*args, **kwargs):
    from .common.utils import get_cluster_context as _get_cluster_context
    return _get_cluster_context(*args, **kwargs)

Impact - Eliminates ~320 ms of import time

2. Command Module Lazy Loading

training.py
Before

from sagemaker.hyperpod.training.hyperpod_pytorch_job import HyperPodPytorchJob
from sagemaker.hyperpod.common.telemetry.telemetry_logging import _hyperpod_telemetry_emitter
# ... other heavy imports at module level

After

def _get_training_dependencies():
    """Lazy load ALL heavy training dependencies"""
    from sagemaker.hyperpod.training.hyperpod_pytorch_job import HyperPodPytorchJob
    from sagemaker.hyperpod.common.telemetry.telemetry_logging import _hyperpod_telemetry_emitter
    # ... other imports
    return (HyperPodPytorchJob, _hyperpod_telemetry_emitter, ...)

def _ensure_training_deps():
    global _training_deps
    if _training_deps is None:
        _training_deps = _get_training_dependencies()
    return _training_deps

Similar pattern applied in inference.py and cluster.py

3. Dynamic CLI Registration (hyp_cli.py)

LazyGroup Class

class LazyGroup(click.Group):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.modules_registered = set()
        
        # Command-to-module mapping for selective loading
        self._command_module_map = {
            'list-cluster': 'cluster',
            'set-cluster-context': 'cluster', 
            # ...
        }

Dynamic Module Registration

def _register_module(self, module_name):
    """Register commands from a specific module"""
    if module_name == 'training':
        from sagemaker.hyperpod.cli.commands.training import (
            pytorch_create, list_jobs, pytorch_describe, # ...
        )
        self.commands['create'].add_command(pytorch_create())
        self.commands['list'].add_command(list_jobs)
        # ...

• Commands loaded only when accessed

Impact on Commands

hyp --help:

• Before: 1.3 s ± 22.3 ms
• After: 36.1 ms ± 1.6 ms

PR Approval Steps

For Requester

1. Description
   • Check the PR title and description for clarity. It should describe the changes made and the reason behind them.
   • Ensure that the PR follows the contribution guidelines, if applicable.
2. Security requirements
   • Ensure that a Pull Request (PR) does not expose passwords and other sensitive information by using git-secrets and upload relevant evidence: https://github.com/awslabs/git-secrets
   • Ensure commit has GitHub Commit Signature
3. Manual review
   1. Click on the Files changed tab to see the code changes. Review the changes thoroughly:
      • Code Quality: Check for coding standards, naming conventions, and readability.
      • Functionality: Ensure that the changes meet the requirements and that all necessary code paths are tested.
      • Security: Check for any security issues or vulnerabilities.
      • Documentation: Confirm that any necessary documentation (code comments, README updates, etc.) has been updated.
4. Check for Merge Conflicts:
   • Verify if there are any merge conflicts with the base branch. GitHub will usually highlight this. If there are conflicts, you should resolve them.

For Reviewer

1. Go through For Requester section to double check each item.
2. Request Changes or Approve the PR:
   1. If the PR is ready to be merged, click Review changes and select Approve.
   2. If changes are required, select Request changes and provide feedback. Be constructive and clear in your feedback.
3. Merging the PR
   1. Check the Merge Method:
      1. Decide on the appropriate merge method based on your repository's guidelines (e.g., Squash and merge, Rebase and merge, or Merge).
   2. Merge the PR:
      1. Click the Merge pull request button.
      2. Confirm the merge by clicking Confirm merge.