model_stats.py

import json
import asyncio
from ollama import AsyncClient, ResponseError
from typing import Dict, List, Any
import time
from datetime import datetime
import logging
from tqdm import tqdm
import os

# Configure logging with more detailed format
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('model_testing.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Ollama client configuration
OLLAMA_HOST = "http://20.7.14.25:11434"
client = AsyncClient(host=OLLAMA_HOST)

# Create results directory if it doesn't exist
RESULTS_DIR = "model_test_results"
os.makedirs(RESULTS_DIR, exist_ok=True)

async def check_model_exists(model_name: str) -> bool:
    """Check if a model is already pulled."""
    try:
        # Try to get model info - if it exists, it will return successfully
        await client.show(model_name)
        return True
    except ResponseError:
        return False
    except Exception as e:
        logger.error(f"Error checking if model {model_name} exists: {e}")
        return False

async def load_approved_models() -> List[str]:
    """Load approved models from the JSON file."""
    try:
        with open('approved_models.json', 'r') as f:
            data = json.load(f)
            models = [model[0] for model in data['approved_models']]
            logger.info(f"Successfully loaded {len(models)} models from approved_models.json")
            return models
    except Exception as e:
        logger.error(f"Error loading approved models: {e}")
        return []

async def pull_model(model_name: str) -> bool:
    """Pull a model from Ollama."""
    try:
        logger.info(f"Starting to pull model: {model_name}")
        start_time = time.time()
        await client.pull(model_name)
        end_time = time.time()
        logger.info(f"Successfully pulled {model_name} in {end_time - start_time:.2f} seconds")
        return True
    except ResponseError as e:
        logger.error(f"Error pulling model {model_name}: {e}")
        return False
    except Exception as e:
        logger.error(f"Unexpected error while pulling {model_name}: {e}")
        return False

async def check_loaded_models():
    """Check if there are any models currently loaded in memory."""
    try:
        # Use ollama ps to check loaded models
        ps_response = await client.ps()
        if ps_response and hasattr(ps_response, 'models'):
            logger.warning("Found loaded models in memory. Waiting for keep_alive to unload them...")
            # Just log the loaded models, they will be unloaded by keep_alive: 0
            for model in ps_response.models:
                if model.name:
                    logger.info(f"Model currently loaded in memory: {model.name}")
                    logger.debug(f"Model details: size={model.size}, vram={model.size_vram}, params={model.details.parameter_size}")
    except Exception as e:
        logger.error(f"Error checking loaded models: {e}")

async def test_model(model_name: str) -> Dict[str, Any]:
    """Test a model and collect performance stats."""
    stats = {
        "model_name": model_name,
        "timestamp": datetime.now().isoformat(),
        "success": False,
        "error": None,
        "performance": {},
        "model_info": {}
    }
    
    try:
        logger.info(f"Starting performance test for model: {model_name}")
        
        # Test with a comprehensive prompt that should generate a longer response
        prompt = """You are a creative writing assistant. Write a short story about a futuristic city where:
1. The city is powered by a mysterious energy source
2. The inhabitants have developed unique abilities
3. There's a hidden conflict between different factions
4. The protagonist discovers a shocking truth about the city's origins

Make the story engaging and include vivid descriptions of the city's architecture and technology."""

        # First, generate a small response to ensure the model is loaded
        await client.generate(
            model=model_name,
            prompt="test",
            stream=False,
            options={
                "max_tokens": 1,
                "keep_alive": 1  # Keep the model loaded
            }
        )
        
        # Get model info while it's loaded
        ps_response = await client.ps()
        if ps_response and hasattr(ps_response, 'models'):
            model_found = False
            for model in ps_response.models:
                if model.name == model_name:
                    model_found = True
                    stats["model_info"] = {
                        "size": model.size,
                        "size_vram": model.size_vram,
                        "parameter_size": model.details.parameter_size,
                        "quantization_level": model.details.quantization_level,
                        "format": model.details.format,
                        "family": model.details.family
                    }
                    logger.info(f"Found model info for {model_name}: {stats['model_info']}")
                    break
            if not model_found:
                logger.warning(f"Model {model_name} not found in ps response. Available models: {[m.name for m in ps_response.models]}")
        else:
            logger.warning(f"No models found in ps response")
        
        start_time = time.time()
        
        # Now generate the full response
        response = await client.generate(
            model=model_name,
            prompt=prompt,
            stream=False,
            options={
                "temperature": 0.7,
                "top_p": 0.9,
                "top_k": 40,
                "max_tokens": 1000,
                "repetition_penalty": 1.0,
                "seed": 42,
                "keep_alive": 0  # Ensure model is unloaded after generation
            }
        )
        
        end_time = time.time()
        
        # Calculate performance metrics
        total_tokens = len(response.get("response", "").split())
        total_time = end_time - start_time
        tokens_per_second = total_tokens / total_time if total_time > 0 else 0
        prompt_tokens = len(prompt.split())
        generation_tokens = total_tokens - prompt_tokens
        
        # Collect detailed performance metrics
        stats["performance"] = {
            "response_time": total_time,
            "total_tokens": total_tokens,
            "tokens_per_second": tokens_per_second,
            "prompt_tokens": prompt_tokens,
            "generation_tokens": generation_tokens,
            "generation_tokens_per_second": generation_tokens / total_time if total_time > 0 else 0,
            "response": response.get("response", ""),
            "eval_count": response.get("eval_count", 0),
            "eval_duration": response.get("eval_duration", 0),
            "prompt_eval_duration": response.get("prompt_eval_duration", 0),
            "total_duration": response.get("total_duration", 0),
        }
        
        stats["success"] = True
        logger.info(f"Successfully tested {model_name}: {tokens_per_second:.2f} tokens/second")
        
    except Exception as e:
        stats["error"] = str(e)
        logger.error(f"Error testing model {model_name}: {e}")
    
    return stats

async def save_results(results: List[Dict[str, Any]], timestamp: str):
    """Save results in multiple formats."""
    # Save detailed results
    detailed_path = os.path.join(RESULTS_DIR, f"model_stats_{timestamp}.json")
    with open(detailed_path, 'w') as f:
        json.dump(results, f, indent=2)
    logger.info(f"Saved detailed results to {detailed_path}")
    
    # Save summary results
    summary = []
    for result in results:
        if result["success"]:
            perf = result["performance"]
            model_info = result["model_info"]
            summary.append({
                "model_name": result["model_name"],
                "model_size": model_info.get("size", 0),
                "vram_size": model_info.get("size_vram", 0),
                "parameter_size": model_info.get("parameter_size", ""),
                "quantization": model_info.get("quantization_level", ""),
                "tokens_per_second": perf["tokens_per_second"],
                "generation_tokens_per_second": perf["generation_tokens_per_second"],
                "total_tokens": perf["total_tokens"],
                "response_time": perf["response_time"],
                "success": result["success"]
            })
    
    summary_path = os.path.join(RESULTS_DIR, f"model_stats_summary_{timestamp}.json")
    with open(summary_path, 'w') as f:
        json.dump(summary, f, indent=2)
    logger.info(f"Saved summary results to {summary_path}")
    
    # Log top performers
    successful_results = [r for r in results if r["success"]]
    if successful_results:
        top_performers = sorted(
            successful_results,
            key=lambda x: x["performance"]["tokens_per_second"],
            reverse=True
        )[:5]
        logger.info("\nTop 5 performers by tokens per second:")
        for r in top_performers:
            model_info = r["model_info"]
            logger.info(f"{r['model_name']}:")
            logger.info(f"  Tokens/second: {r['performance']['tokens_per_second']:.2f}")
            logger.info(f"  VRAM Usage: {model_info.get('size_vram', 0)/1024/1024/1024:.2f} GB")
            logger.info(f"  Parameter Size: {model_info.get('parameter_size', 'N/A')}")
            logger.info(f"  Quantization: {model_info.get('quantization_level', 'N/A')}")
            logger.info("  " + "-" * 30)

async def main():
    """Main function to run the model testing process."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    logger.info("Starting model testing process")
    
    # Load approved models
    models = await load_approved_models()
    if not models:
        logger.error("No models loaded. Exiting.")
        return
    
    # Check and unload any models that might be already loaded
    await check_loaded_models()
    
    # Check which models need to be pulled
    models_to_pull = []
    for model in models:
        if not await check_model_exists(model):
            models_to_pull.append(model)
    
    if models_to_pull:
        logger.info(f"Found {len(models_to_pull)} models that need to be pulled:")
        for model in models_to_pull:
            logger.info(f"- {model}")
        
        # Ask user if they want to pull missing models
        while True:
            response = input("\nDo you want to pull the missing models? (yes/no): ").lower()
            if response in ['yes', 'no']:
                break
            print("Please answer 'yes' or 'no'")
        
        if response == 'yes':
            # Pull missing models with progress bar
            logger.info("Starting model pulling phase")
            for model in tqdm(models_to_pull, desc="Pulling models"):
                await pull_model(model)
        else:
            logger.warning("Skipping model pulling. Some models may not be available for testing.")
            # Filter out models that weren't pulled
            models = [model for model in models if model not in models_to_pull]
            if not models:
                logger.error("No models available for testing. Exiting.")
                return
    else:
        logger.info("All models are already pulled. Skipping pulling phase.")
    
    # Test all models with progress bar
    logger.info("Starting model testing phase")
    results = []
    for model in tqdm(models, desc="Testing models"):
        # Check for any loaded models before testing
        await check_loaded_models()
        
        stats = await test_model(model)
        results.append(stats)
        
        # Save intermediate results after each model
        await save_results(results, timestamp)
        
        # Add sleep between model tests to ensure proper cleanup
        logger.info("Waiting 3 seconds before next model test...")
        await asyncio.sleep(3)
    
    # Save final results
    await save_results(results, timestamp)
    
    # Log summary
    successful_tests = sum(1 for r in results if r["success"])
    logger.info(f"Model testing completed. {successful_tests}/{len(models)} models tested successfully")

if __name__ == "__main__":
    asyncio.run(main())