ComputerAssignments/CA19_gpu_memory_optimization/scripts/main.py

import sys
import os
import argparse
import logging
from typing import Dict, Any


sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))

from utils.config_utils import ConfigManager
from utils.data_utils import create_directory_structure
from utils.plot_utils import create_all_plots
from evaluation.benchmark import run_comprehensive_benchmark
from visualization import create_all_visualizations

def setup_logging(log_level: str = "INFO", log_file: str = "logs/experiment.log"):
    os.makedirs(os.path.dirname(log_file), exist_ok=True)

    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        handlers=[logging.FileHandler(log_file), logging.StreamHandler(sys.stdout)],
    )

    return logging.getLogger(__name__)

def run_memory_profiling_analysis(config: ConfigManager, logger: logging.Logger):
    logger.info("Starting memory profiling analysis...")

    try:
        from core.memory_profiler import MemoryProfiler
        from core.gpu_profiler import GPUMemoryProfiler


        cpu_profiler = MemoryProfiler()
        gpu_profiler = GPUMemoryProfiler()

        logger.info("Memory profilers initialized successfully")


        import numpy as np

        data_sizes = [100, 500, 1000]
        for size_mb in data_sizes:
            logger.info(f"Testing {size_mb}MB memory allocation...")

            cpu_profiler.reset()
            cpu_profiler.measure("Before allocation")


            data = np.random.rand(size_mb * 1024 * 1024 // 8).astype(np.float64)
            cpu_profiler.measure(f"After {size_mb}MB allocation")


            result = np.sum(data)
            cpu_profiler.measure(f"After processing {size_mb}MB")

            del data
            cpu_profiler.measure(f"After cleanup {size_mb}MB")

            logger.info(f"✅ {size_mb}MB test completed")

        logger.info("Memory profiling analysis completed successfully")

    except Exception as e:
        logger.error(f"Error in memory profiling analysis: {e}")
        raise

def run_cache_optimization_analysis(config: ConfigManager, logger: logging.Logger):
    logger.info("Starting cache optimization analysis...")

    try:
        from algorithms.cache_optimized import *
        import numpy as np
        import time


        sizes = [128, 256, 512]
        algorithms = {
            "Naive": naive_matrix_multiply,
            "Blocked": lambda A, B: blocked_matrix_multiply(A, B, 64),
            "Cache-oblivious": lambda A, B: cache_oblivious_multiply(A, B, 64),
            "NumPy": np.dot,
        }

        results = {}

        for size in sizes:
            logger.info(f"Testing {size}x{size} matrices...")
            A = np.random.rand(size, size).astype(np.float32)
            B = np.random.rand(size, size).astype(np.float32)

            size_results = []
            for name, func in algorithms.items():
                if name == "Naive" and size > 256:
                    continue

                start_time = time.time()
                result = func(A, B)
                end_time = time.time()

                execution_time = end_time - start_time
                size_results.append((size, execution_time, 0))

                logger.info(f"  {name}: {execution_time:.4f}s")

            results[f"size_{size}"] = size_results

        logger.info("Cache optimization analysis completed successfully")
        return results

    except Exception as e:
        logger.error(f"Error in cache optimization analysis: {e}")
        raise

def run_memory_efficient_nets_analysis(config: ConfigManager, logger: logging.Logger):
    logger.info("Starting memory-efficient neural networks analysis...")

    try:
        from memory.efficient_nets import *
        import torch


        batch_size = 128
        input_size = 784
        input_tensor = torch.randn(batch_size, input_size)

        models = {
            "Standard": MemoryEfficientNet(use_checkpointing=False),
            "Checkpointed": MemoryEfficientNet(use_checkpointing=True),
            "Mixed Precision": MixedPrecisionNet(),
        }

        results = {}

        for name, model in models.items():
            logger.info(f"Testing {name} model...")


            output = model(input_tensor)
            loss = output.sum()


            loss.backward()

            results[name] = {"output_shape": output.shape, "loss_value": loss.item()}

            logger.info(f"  ✅ {name} model test completed")

        logger.info("Memory-efficient neural networks analysis completed successfully")
        return results

    except Exception as e:
        logger.error(f"Error in memory-efficient neural networks analysis: {e}")
        raise

def run_distributed_systems_analysis(config: ConfigManager, logger: logging.Logger):
    logger.info("Starting distributed memory systems analysis...")

    try:
        from distributed.zero_optimizer import *


        simulator = DistributedMemorySimulator(
            model_size_mb=4000, batch_size_mb=200, num_devices=4
        )

        logger.info("Testing distributed memory strategies...")


        dp_mem, dp_total = simulator.data_parallel_memory()
        logger.info(f"Data Parallel: {dp_mem['total']/1024:.2f} GB per device")


        mp_mem, mp_total = simulator.model_parallel_memory()
        logger.info(f"Model Parallel: {mp_mem['total']/1024:.2f} GB per device")


        zero_mem, zero_total = simulator.zero_stage3_memory()
        logger.info(f"ZeRO Stage 3: {zero_mem['total']/1024:.2f} GB per device")

        results = {
            "data_parallel": dp_mem,
            "model_parallel": mp_mem,
            "zero_stage3": zero_mem,
        }

        logger.info("Distributed memory systems analysis completed successfully")
        return results

    except Exception as e:
        logger.error(f"Error in distributed memory systems analysis: {e}")
        raise

def run_comprehensive_analysis(config: ConfigManager, logger: logging.Logger):
    logger.info("Starting comprehensive analysis...")

    try:

        logger.info("Generating visualizations...")
        create_all_visualizations("visualizations")


        logger.info("Running benchmark suite...")
        benchmark_suite = run_comprehensive_benchmark("results")

        logger.info("Comprehensive analysis completed successfully")
        return benchmark_suite

    except Exception as e:
        logger.error(f"Error in comprehensive analysis: {e}")
        raise

def main():
    parser = argparse.ArgumentParser(description="Systems Memory in AI - CA19 Analysis")
    parser.add_argument(
        "--config",
        type=str,
        default="configs/config.yaml",
        help="Configuration file path",
    )
    parser.add_argument(
        "--log-level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="Logging level",
    )
    parser.add_argument(
        "--output-dir", type=str, default="results", help="Output directory for results"
    )
    parser.add_argument(
        "--skip-profiling", action="store_true", help="Skip memory profiling analysis"
    )
    parser.add_argument(
        "--skip-cache", action="store_true", help="Skip cache optimization analysis"
    )
    parser.add_argument(
        "--skip-nets",
        action="store_true",
        help="Skip memory-efficient neural networks analysis",
    )
    parser.add_argument(
        "--skip-distributed",
        action="store_true",
        help="Skip distributed systems analysis",
    )

    args = parser.parse_args()


    logger = setup_logging(args.log_level)
    logger.info("Starting Systems Memory in AI - CA19 Analysis")


    create_directory_structure()


    config = ConfigManager(args.config)
    exp_config = config.get_experiment_config()

    logger.info(f"Experiment: {exp_config.name} v{exp_config.version}")
    logger.info(f"Description: {exp_config.description}")


    results = {}

    try:
        if not args.skip_profiling:
            results["profiling"] = run_memory_profiling_analysis(config, logger)

        if not args.skip_cache:
            results["cache_optimization"] = run_cache_optimization_analysis(
                config, logger
            )

        if not args.skip_nets:
            results["memory_efficient_nets"] = run_memory_efficient_nets_analysis(
                config, logger
            )

        if not args.skip_distributed:
            results["distributed_systems"] = run_distributed_systems_analysis(
                config, logger
            )


        results["comprehensive"] = run_comprehensive_analysis(config, logger)

        logger.info("All analyses completed successfully!")
        logger.info(f"Results saved to: {args.output_dir}")
        logger.info("Visualizations saved to: visualizations")

    except Exception as e:
        logger.error(f"Analysis failed: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()