ADK cold start improvement

build.py

import ast
import compileall
import importlib
import json
import os
import shutil

# --- Configuration ---
# List of packages to process - add more packages here
PACKAGES_TO_PROCESS = [
    "google.adk",
    "google.cloud.aiplatform",
    "google.cloud.bigquery",      # If you use BigQuery tools
    "google.auth",                # Auth can be lazy-loaded
]

BUILD_DIR = "build"

# --- Files/directories to exclude from copying (saves space and time) ---
EXCLUDE_PATTERNS = {
    "tests", "test", "*_test.py", "test_*.py",
    "docs", "*.md", "*.rst",
    ".git", ".github", "__pycache__",
    "*.pyc", "*.pyo", "*.pyd",
}

# --- Directories requiring gentle handling (__init__.py is untouched) ---
# Format: {"package.name": ["dir1", "dir2"]}
SPECIAL_HANDLING_DIRS = {
    "google.adk": [
        "cli",           # CLI needs original structure
        "models",        # Model registry needs eager loading
        # Consider adding these if they cause issues:
        # "runners",     # Only if runner initialization fails
        # "memory",      # Only if memory backends fail
    ],
    "google.cloud.aiplatform": [],  # Add specific dirs if needed
}
# --------------------------------

# The template for our new, intelligent __init__.py files.
# It uses PEP 562's module-level __getattr__ to be both fast and stable.
LAZY_INIT_TEMPLATE = """
import importlib

# This dictionary is auto-generated by the build script.
_LAZY_MAP = {lazy_map_json}

_cache = {{}}

def __getattr__(name):
    \"\"\"Lazily imports an attribute from a submodule when it's first accessed.\"\"\"
    if name in _cache:
        return _cache[name]
    
    if name not in _LAZY_MAP:
        raise AttributeError(f"module '{{__name__}}' has no attribute '{{name}}'")

    module_path = _LAZY_MAP[name]
    
    try:
        # First, try to import it as a submodule (for "from . import module" cases)
        if module_path == '.':
            # This is a submodule import from current package
            full_module_name = f"{{__name__}}.{{name}}"
            try:
                module = importlib.import_module(full_module_name)
                _cache[name] = module
                globals()[name] = module
                return module
            except ImportError:
                # Not a submodule, fall through to attribute import
                pass
        
        # Import the module and get the attribute
        module = importlib.import_module(module_path, __package__)
        
        # Prevent infinite recursion: check if we're trying to get an attribute
        # from the same module we're currently in
        if module is globals().get('__loader__'):
            raise AttributeError(f"Circular import detected for '{{name}}'")
        
        attr = getattr(module, name)
        
        # Cache the result for future access
        _cache[name] = attr
        globals()[name] = attr
        
        return attr
    except (ImportError, AttributeError) as e:
        raise ImportError(f"Could not lazy-load '{{name}}' from '{{module_path}}'") from e

def __dir__():
    \"\"\"Exposes the lazy-loaded attributes to dir() and code completion.\"\"\"
    return list(globals().keys()) + list(_LAZY_MAP.keys())
"""

def get_package_path(package_name):
    try:
        spec = importlib.util.find_spec(package_name)
        if spec is None or spec.submodule_search_locations is None: 
            raise ImportError
        return spec.submodule_search_locations[0]
    except (ImportError, AttributeError):
        raise ImportError(f"Package '{package_name}' not found.")

def should_skip_file(init_path, tree):
    """Determine if an __init__.py file should be skipped based on its content.
    
    Only skip files that actually DEFINE classes/functions directly (rare in __init__.py).
    Most __init__.py files just re-export, which is perfect for lazy loading.
    """
    
    for node in tree.body:
        # Only skip if there are actual class or function DEFINITIONS in __init__.py
        # (This is rare - most __init__.py just import and re-export)
        if isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
            return True
    
    return False

def should_exclude(path, name):
    """Check if a file or directory should be excluded."""
    for pattern in EXCLUDE_PATTERNS:
        if pattern.startswith('*') and name.endswith(pattern[1:]):
            return True
        elif pattern.endswith('*') and name.startswith(pattern[:-1]):
            return True
        elif name == pattern:
            return True
    return False

def copy_without_excluded(src, dst):
    """Copy directory tree excluding certain patterns."""
    def ignore_patterns(directory, contents):
        return [name for name in contents if should_exclude(directory, name)]
    
    shutil.copytree(src, dst, ignore=ignore_patterns)

def process_package(package_name):
    """Process a single package for lazy loading optimization."""
    
    print(f"\n{'='*60}")
    print(f"Processing package: {package_name}")
    print(f"{'='*60}")
    
    try:
        source_dir = get_package_path(package_name)
        print(f"Found '{package_name}' at: {source_dir}")
    except ImportError as e:
        print(f"Warning: {e} - Skipping this package.")
        return False

    # Create target directory
    target_dir = os.path.join(BUILD_DIR, *package_name.split('.'))
    
    # If target already exists from a parent package, skip
    if os.path.exists(target_dir):
        print(f"Target directory already exists: {target_dir} - Using existing copy")
    else:
        copy_without_excluded(source_dir, target_dir)
        print(f"Copied original package to '{target_dir}' (excluded tests/docs)")

    target_dir_abs = os.path.abspath(target_dir)
    
    # Get special handling directories for this package
    special_dirs = SPECIAL_HANDLING_DIRS.get(package_name, [])
    special_paths = {os.path.join(target_dir_abs, d) for d in special_dirs}
    special_paths.add(target_dir_abs)  # Always skip root __init__.py

    for root, _, files in os.walk(target_dir):
        root_abs = os.path.abspath(root)

        if root_abs in special_paths:
            print(f"Leaving sensitive directory untouched: {root}")
            continue
        
        if '__init__.py' in files:
            init_path = os.path.join(root, '__init__.py')
            
            print(f"Applying AST rewrite to __init__.py in: {root}")
            try:
                with open(init_path, "r", encoding="utf-8") as f:
                    source_code = f.read()
                    tree = ast.parse(source_code)

                lazy_map = {}
                # Find all relative imports and build our map
                for node in tree.body:
                    if isinstance(node, ast.ImportFrom) and node.level > 0:
                        module_path = "." * node.level + (node.module or "")
                        for alias in node.names:
                            if alias.name != "*":
                                lazy_map[alias.asname or alias.name] = module_path
                
                if lazy_map:
                    # If we found lazy-loadable imports, rewrite the file
                    lazy_map_json = json.dumps(lazy_map, indent=4)
                    with open(init_path, "w", encoding="utf-8") as f:
                        f.write(LAZY_INIT_TEMPLATE.format(lazy_map_json=lazy_map_json))
                    print(f"  ✓ Rewrote with {len(lazy_map)} lazy imports")
                else:
                    print(f"  - No relative imports to rewrite in {init_path}. Leaving as-is.")

            except Exception as e:
                print(f"  - WARNING: Could not parse/rewrite {init_path}. Leaving as-is. Error: {e}")
    
    return True

def create_minified_packages():
    """Process all configured packages."""
    
    print(f"\nStarting multi-package lazy loading optimization")
    print(f"Packages to process: {', '.join(PACKAGES_TO_PROCESS)}")
    print(f"Build directory: {BUILD_DIR}")
    
    # Clean build directory
    if os.path.exists(BUILD_DIR): 
        shutil.rmtree(BUILD_DIR)
        print(f"\nCleaned existing build directory")
    
    os.makedirs(BUILD_DIR, exist_ok=True)
    
    # Process each package
    successful = 0
    failed = 0
    
    for package_name in PACKAGES_TO_PROCESS:
        if process_package(package_name):
            successful += 1
        else:
            failed += 1
    
    # Precompile all Python files for faster loading
    print(f"\n{'='*60}")
    print("Precompiling Python bytecode...")
    print(f"{'='*60}")
    compileall.compile_dir(BUILD_DIR, force=True, quiet=1, legacy=True)
    print("✓ Bytecode compilation complete")
    
    print(f"\n{'='*60}")
    print(f"Build process complete!")
    print(f"{'='*60}")
    print(f"Successfully processed: {successful} package(s)")
    print(f"Failed/Skipped: {failed} package(s)")
    print(f"\nAttribute-level lazy-loading structure created in: {BUILD_DIR}")

if __name__ == "__main__":
    create_minified_packages()