debug_memory.py

#!/usr/bin/env python3
"""
Debug script to monitor memory usage during model loading.
Run this to identify exactly where the memory issues occur.
"""

import gc
import logging
import os
import sys
from typing import Optional

import psutil
import torch

# Set up logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def get_memory_info():
    """Get current memory usage information."""
    process = psutil.Process(os.getpid())
    memory_info = process.memory_info()
    virtual_memory = psutil.virtual_memory()

    return {
        "process_rss_gb": memory_info.rss / (1024**3),  # Resident Set Size
        "process_vms_gb": memory_info.vms / (1024**3),  # Virtual Memory Size
        "system_total_gb": virtual_memory.total / (1024**3),
        "system_available_gb": virtual_memory.available / (1024**3),
        "system_used_gb": virtual_memory.used / (1024**3),
        "system_percent": virtual_memory.percent,
    }


def log_memory_usage(step: str):
    """Log current memory usage with a step description."""
    mem_info = get_memory_info()
    logger.info(f"=== {step} ===")
    logger.info(f"Process RSS: {mem_info['process_rss_gb']:.2f} GB")
    logger.info(f"Process VMS: {mem_info['process_vms_gb']:.2f} GB")
    logger.info(f"System Total: {mem_info['system_total_gb']:.2f} GB")
    logger.info(f"System Available: {mem_info['system_available_gb']:.2f} GB")
    logger.info(
        f"System Used: {mem_info['system_used_gb']:.2f} GB ({mem_info['system_percent']:.1f}%)"
    )

    if torch.cuda.is_available():
        logger.info(
            f"CUDA Memory Allocated: {torch.cuda.memory_allocated() / (1024**3):.2f} GB"
        )
        logger.info(
            f"CUDA Memory Cached: {torch.cuda.memory_reserved() / (1024**3):.2f} GB"
        )

    logger.info("")


def force_cleanup():
    """Force garbage collection and memory cleanup."""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def debug_model_loading(models_dir: str = "/home/user/app/models"):
    """Debug the model loading process step by step."""

    ckpt_path = os.path.join(models_dir, "mms_XRI.pt")
    tokenizer_path = os.path.join(models_dir, "mms_1143_langs_tokenizer_spm.model")

    logger.info("Starting memory debugging for MMS model loading...")
    logger.info(f"Checkpoint path: {ckpt_path}")
    logger.info(f"Tokenizer path: {tokenizer_path}")

    # Check file sizes
    if os.path.exists(ckpt_path):
        ckpt_size_gb = os.path.getsize(ckpt_path) / (1024**3)
        logger.info(f"Checkpoint file size: {ckpt_size_gb:.2f} GB")
    else:
        logger.error(f"Checkpoint file not found: {ckpt_path}")
        return

    log_memory_usage("Initial state")

    try:
        # Step 1: Check available memory before loading
        mem_info = get_memory_info()
        if mem_info["system_available_gb"] < ckpt_size_gb * 1.5:
            logger.warning(
                f"Available memory ({mem_info['system_available_gb']:.2f} GB) may be insufficient for checkpoint ({ckpt_size_gb:.2f} GB)"
            )

        # Step 2: Try to load checkpoint with memory mapping
        logger.info("Step 1: Loading checkpoint with memory mapping...")
        try:
            # Use mmap=True to avoid loading entire file into memory
            model_params = torch.load(ckpt_path, map_location="cpu", mmap=True)
            log_memory_usage("After loading checkpoint (mmap)")
        except Exception as e:
            logger.error(f"Memory-mapped loading failed: {e}")
            logger.info("Falling back to regular loading...")
            model_params = torch.load(ckpt_path, map_location="cpu")
            log_memory_usage("After loading checkpoint (regular)")

        # Step 3: Setup fairseq2 and configs
        logger.info("Step 2: Setting up fairseq2 and configs...")
        from fairseq2 import setup_fairseq2
        from fairseq2.context import get_runtime_context
        from fairseq2.models.llama import LLaMAConfig

        # Import the model classes
        sys.path.append("/home/user/app/server")
        from model import (
            register_wav2vec2_asr_configs,
            register_wav2vec2_configs,
            Wav2Vec2AsrConfig,
            Wav2Vec2LlamaConfig,
            Wav2Vec2LlamaFactory,
        )

        setup_fairseq2()
        context = get_runtime_context()
        register_wav2vec2_configs(context)
        register_wav2vec2_asr_configs(context)

        log_memory_usage("After fairseq2 setup")

        # Step 4: Create configs
        logger.info("Step 3: Creating model configuration...")
        w2v2_ctc_registry = context.get_config_registry(Wav2Vec2AsrConfig)
        wav2vec_ctc_config = w2v2_ctc_registry.get("7b_bib1143")

        llama_config = LLaMAConfig(
            model_dim=4096,
            max_seq_len=8192,
            vocab_info=wav2vec_ctc_config.vocab_info,
            num_layers=12,
            num_attn_heads=8,
            num_key_value_heads=8,
            ffn_inner_dim=4096,
            rope_theta=10_000.0,
            dropout_p=0.1,
        )

        config = Wav2Vec2LlamaConfig()
        config.wav2vec_ctc_config = wav2vec_ctc_config
        config.llama_config = llama_config

        log_memory_usage("After creating configs")

        # Step 5: Create model architecture (without loading weights)
        logger.info("Step 4: Creating model architecture...")
        factory = Wav2Vec2LlamaFactory(config)
        model = factory.create_model()

        log_memory_usage("After creating model architecture")

        # Step 6: Load state dict
        logger.info("Step 5: Loading model weights...")
        try:
            model.load_state_dict(model_params["model"])
            log_memory_usage("After loading model weights")
        except Exception as e:
            logger.error(f"Failed to load model weights: {e}")
            return

        # Step 7: Clean up checkpoint data
        logger.info("Step 6: Cleaning up checkpoint data...")
        del model_params
        force_cleanup()
        log_memory_usage("After cleanup")

        # Step 8: Move to device (if specified)
        device = torch.device("cpu")  # Force CPU for debugging
        logger.info(f"Step 7: Moving model to device {device}...")
        model = model.to(device).eval()
        log_memory_usage("After moving to device")

        logger.info("Model loading completed successfully!")

        # Step 9: Test a small inference to see memory usage
        logger.info("Step 8: Testing small inference...")
        try:
            # Create a small dummy input
            dummy_input = torch.randn(1, 16000).to(device)  # 1 second of audio
            with torch.no_grad():
                # Just test the encoder part to avoid full inference
                enc_out = model.encoder_frontend.extract_features(dummy_input, None)
                log_memory_usage("After small inference test")
        except Exception as e:
            logger.error(f"Small inference test failed: {e}")

        return model

    except Exception as e:
        logger.error(f"Error during model loading: {str(e)}")
        log_memory_usage("After error")
        raise


def check_docker_memory_limits():
    """Check if we're running in Docker and what the memory limits are."""
    logger.info("Checking Docker memory configuration...")

    # Check if we're in a container
    if os.path.exists("/.dockerenv"):
        logger.info("Running inside Docker container")

        # Check cgroup memory limits
        try:
            with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "r") as f:
                limit_bytes = int(f.read().strip())
                limit_gb = limit_bytes / (1024**3)
                logger.info(f"Docker memory limit: {limit_gb:.2f} GB")

                # Check if limit is reasonable (not the default huge value)
                if limit_gb > 1000:  # Probably unlimited
                    logger.warning("Docker memory limit appears to be unlimited")
                else:
                    logger.info(f"Docker memory limit is set to {limit_gb:.2f} GB")

        except Exception as e:
            logger.warning(f"Could not read Docker memory limit: {e}")

        # Check current memory usage in container
        try:
            with open("/sys/fs/cgroup/memory/memory.usage_in_bytes", "r") as f:
                usage_bytes = int(f.read().strip())
                usage_gb = usage_bytes / (1024**3)
                logger.info(f"Current Docker memory usage: {usage_gb:.2f} GB")
        except Exception as e:
            logger.warning(f"Could not read Docker memory usage: {e}")
    else:
        logger.info("Not running in Docker container")


if __name__ == "__main__":
    # Check Docker memory configuration
    check_docker_memory_limits()

    # Get models directory from environment or use default
    models_dir = os.environ.get("MODELS_DIR", "/home/user/app/models")

    # Run the debugging
    try:
        model = debug_model_loading(models_dir)
        logger.info("Memory debugging completed successfully!")
    except Exception as e:
        logger.error(f"Memory debugging failed: {e}")
        sys.exit(1)