Benchmark: Llama 3.1 70B vs Mistral 8x22B for Code Generation Accuracy

"""
HumanEval+ Benchmark Harness for Llama 3.1 70B vs Mistral 8x22B
Requires: vllm==0.4.2, datasets==2.20.0, python-dotenv==1.0.0
Environment variables:
  - HUGGING_FACE_HUB_TOKEN: Token for downloading models from HF Hub
  - CUDA_VISIBLE_DEVICES: GPU IDs to use (e.g., "0,1" for 2 A100s)
"""
import os
import json
import time
import dotenv
import argparse
from vllm import LLM, SamplingParams
from datasets import load_dataset
from typing import List, Dict, Any

# Load environment variables
dotenv.load_dotenv()

# Argument parser for configurable benchmark runs
parser = argparse.ArgumentParser(description="Run HumanEval+ benchmark for Llama 3.1 70B and Mistral 8x22B")
parser.add_argument("--model", type=str, required=True, help="Model ID (e.g., meta-llama/Llama-3.1-70B-Instruct)")
parser.add_argument("--dataset", type=str, default="bigcode/humanevalplus", help="Dataset to use")
parser.add_argument("--num-samples", type=int, default=164, help="Number of tasks to run (164 for full HumanEval+)")
parser.add_argument("--output-file", type=str, default="benchmark_results.json", help="Output JSON file path")
args = parser.parse_args()

# Validate environment variables
if not os.getenv("HUGGING_FACE_HUB_TOKEN"):
    raise ValueError("Missing HUGGING_FACE_HUB_TOKEN environment variable")

# Initialize vLLM with model-specific config
def init_llm(model_id: str) -> LLM:
    """Initialize vLLM instance with optimal params for the target model"""
    try:
        # Llama 3.1 uses group query attention, Mistral uses MoE - adjust tensor parallelism
        tensor_parallel_size = 2 if "Llama-3.1-70B" in model_id else 1
        return LLM(
            model=model_id,
            tensor_parallel_size=tensor_parallel_size,
            max_model_len=131072 if "Llama" in model_id else 65536,
            gpu_memory_utilization=0.95,
            trust_remote_code=True
        )
    except Exception as e:
        raise RuntimeError(f"Failed to initialize LLM for {model_id}: {str(e)}")

# Sampling params for code generation (deterministic for pass@1)
SAMPLING_PARAMS = SamplingParams(
    temperature=0.0,
    top_p=1.0,
    max_tokens=512,
    stop=[\"\n\n\", \"def \", \"class \"]
)

def run_benchmark(llm: LLM, dataset: List[Dict[str, Any]]) -> Dict[str, Any]:
    """Run pass@1 benchmark on target dataset"""
    results = {
        "total_tasks": len(dataset),
        "passed_tasks": 0,
        "failed_tasks": 0,
        "latency_ms": 0.0,
        "task_results": []
    }

    for task in dataset:
        task_id = task["task_id"]
        prompt = task["prompt"]
        test_cases = task["test"]

        # Generate code
        start_time = time.time()
        try:
            outputs = llm.generate([prompt], SAMPLING_PARAMS)
            generated_code = outputs[0].outputs[0].text
        except Exception as e:
            print(f"Generation failed for {task_id}: {str(e)}")
            results["failed_tasks"] += 1
            results["task_results"].append({"task_id": task_id, "passed": False, "error": str(e)})
            continue
        latency = (time.time() - start_time) * 1000
        results["latency_ms"] += latency

        # Execute test cases (sandboxed - simplified for example)
        # In production, use a Docker sandbox to avoid arbitrary code execution
        full_code = prompt + generated_code + \"\n\" + test_cases
        try:
            # WARNING: Never run untrusted code without sandboxing in production
            exec(full_code, {"__builtins__": {}})  # Restricted builtins for safety
            passed = True
        except AssertionError:
            passed = False
        except Exception as e:
            passed = False

        # Update results
        if passed:
            results["passed_tasks"] += 1
        else:
            results["failed_tasks"] += 1
        results["task_results"].append({
            "task_id": task_id,
            "passed": passed,
            "latency_ms": latency,
            "generated_code": generated_code
        })

    # Calculate pass@1
    results["pass_at_1"] = (results["passed_tasks"] / results["total_tasks"]) * 100
    results["avg_latency_ms"] = results["latency_ms"] / results["total_tasks"]
    return results

if __name__ == "__main__":
    # Load dataset
    print(f"Loading dataset {args.dataset}...")
    try:
        dataset = load_dataset(args.dataset, split="test")
        dataset = dataset.select(range(min(args.num_samples, len(dataset))))
    except Exception as e:
        raise RuntimeError(f"Failed to load dataset: {str(e)}")

    # Initialize model
    print(f"Initializing model {args.model}...")
    llm = init_llm(args.model)

    # Run benchmark
    print(f"Running benchmark on {len(dataset)} tasks...")
    results = run_benchmark(llm, dataset)

    # Save results
    with open(args.output_file, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Benchmark complete. Pass@1: {results['pass_at_1']:.1f}%. Results saved to {args.output_file}")

"""
Mistral 8x22B MoE Inference Optimizer
Implements expert offloading and dynamic routing for cost-sensitive workloads
Requires: vllm==0.4.2, torch==2.1.0, psutil==5.9.0
"""
import torch
import psutil
import argparse
from vllm import LLM, SamplingParams
from vllm.model_executor.layers.moe import FusedMoEMethod
from typing import List, Optional

class MoEOptimizer:
    """Optimizes Mistral 8x22B MoE inference for low-cost, high-throughput code gen"""
    def __init__(self, model_id: str = "mistralai/Mixtral-8x22B-Instruct-v0.3"):
        self.model_id = model_id
        self.llm = None
        self.active_experts = 2  # Mistral 8x22B uses 2 active experts per token
        self.total_experts = 8
        self.expert_offload_threshold = 0.7  # Offload experts if VRAM usage > 70%

    def _get_available_vram(self) -> float:
        """Get available GPU VRAM in GB across all visible devices"""
        if not torch.cuda.is_available():
            return 0.0
        total_available = 0.0
        for i in range(torch.cuda.device_count()):
            total_available += torch.cuda.get_device_properties(i).total_mem - torch.cuda.memory_allocated(i)
        return total_available / (1024 ** 3)  # Convert to GB

    def _offload_experts(self, experts_to_offload: List[int]) -> None:
        """Offload specified experts to CPU to free VRAM"""
        if not hasattr(self.llm, "model"):
            raise RuntimeError("LLM not initialized")
        model = self.llm.model
        for expert_idx in experts_to_offload:
            # Access MoE layers (simplified - actual vLLM internals may vary)
            for layer in model.layers:
                if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
                    expert = layer.mlp.experts[expert_idx]
                    expert.to("cpu")
        print(f"Offloaded {len(experts_to_offload)} experts to CPU")

    def _route_experts_dynamically(self, prompt: str) -> List[int]:
        """Predict which experts will be used for a given prompt to pre-load them"""
        # Simplified routing: for code prompts, pre-load experts 0,1 (most used for code per Mistral docs)
        if "def " in prompt or "class " in prompt:
            return [0, 1]
        return list(range(self.active_experts))  # Default to first 2 experts

    def init_llm(self, max_model_len: int = 65536) -> None:
        """Initialize LLM with MoE-specific optimizations"""
        available_vram = self._get_available_vram()
        print(f"Available VRAM: {available_vram:.1f} GB")

        # Adjust tensor parallelism based on available GPUs
        tensor_parallel_size = min(torch.cuda.device_count(), 2)

        try:
            self.llm = LLM(
                model=self.model_id,
                tensor_parallel_size=tensor_parallel_size,
                max_model_len=max_model_len,
                gpu_memory_utilization=0.9,
                trust_remote_code=True,
                # Enable MoE expert offloading if VRAM is insufficient
                enable_expert_offload=available_vram < 88.0  # 88GB is min for 8x22B FP16
            )
        except Exception as e:
            raise RuntimeError(f"Failed to initialize MoE LLM: {str(e)}")

        # Offload unused experts if VRAM usage is high
        if psutil.virtual_memory().percent > self.expert_offload_threshold * 100:
            unused_experts = [2,3,4,5,6,7]  # Keep experts 0,1 (code-optimized) on GPU
            self._offload_experts(unused_experts)

    def generate_code(self, prompt: str, max_tokens: int = 512) -> str:
        """Generate code with dynamic expert routing"""
        if not self.llm:
            raise RuntimeError("LLM not initialized. Call init_llm first.")

        # Pre-route experts for the prompt
        target_experts = self._route_experts_dynamically(prompt)
        print(f"Routing to experts: {target_experts}")

        # Sampling params for code generation
        sampling_params = SamplingParams(
            temperature=0.2,
            top_p=0.95,
            max_tokens=max_tokens,
            stop=[\"\n\n\", \"def \", \"import \"]
        )

        try:
            outputs = self.llm.generate([prompt], sampling_params)
            return outputs[0].outputs[0].text
        except Exception as e:
            raise RuntimeError(f"Code generation failed: {str(e)}")

    def get_cost_metrics(self, num_tokens: int) -> Dict[str, float]:
        """Calculate cost metrics for generated tokens"""
        power_draw_w = 400.0  # Average A100 power draw per GPU
        num_gpus = torch.cuda.device_count()
        total_power = power_draw_w * num_gpus
        cost_per_kwh = 0.12
        total_kwh = (total_power * (num_tokens / 1000)) / (3.6e6)  # Convert W*s to kWh
        return {
            "total_power_w": total_power,
            "cost_usd": total_kwh * cost_per_kwh,
            "tokens_per_second": 25.0  # Average for 8x22B on A100
        }

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run optimized Mistral 8x22B inference")
    parser.add_argument("--prompt-file", type=str, help="Path to file containing code prompt")
    args = parser.parse_args()

    optimizer = MoEOptimizer()
    optimizer.init_llm()

    # Load prompt
    if args.prompt_file:
        with open(args.prompt_file, "r") as f:
            prompt = f.read()
    else:
        prompt = "Write a Python function to calculate the Fibonacci sequence iteratively."

    # Generate code
    print(f"Generating code for prompt: {prompt[:50]}...")
    generated_code = optimizer.generate_code(prompt)
    print(f"Generated Code:\n{generated_code}")

    # Print cost metrics
    num_tokens = len(generated_code.split()) * 1.3  # Approx tokens per word
    metrics = optimizer.get_cost_metrics(int(num_tokens))
    print(f"Cost Metrics: {metrics}")

"""
Llama 3.1 70B Context Window Stress Test
Tests long-range dependency tracking for large codebases (128k token context)
Requires: vllm==0.4.2, torch==2.1.0, datasets==2.20.0
"""
import os
import json
import time
import argparse
from vllm import LLM, SamplingParams
from datasets import load_dataset
from typing import List, Dict, Any, Tuple

class ContextStressTest:
    """Runs context window stress tests for Llama 3.1 70B on code tasks"""
    def __init__(self, model_id: str = "meta-llama/Llama-3.1-70B-Instruct"):
        self.model_id = model_id
        self.llm = None
        self.max_context = 131072  # 131k tokens for Llama 3.1 70B
        self.sampling_params = SamplingParams(
            temperature=0.0,
            top_p=1.0,
            max_tokens=1024,
            stop=[\"\n\n\"]
        )

    def init_llm(self) -> None:
        """Initialize LLM with prefix caching enabled for long context"""
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA required for Llama 3.1 70B inference")
        tensor_parallel_size = min(torch.cuda.device_count(), 2)  # 2 A100s for 70B
        try:
            self.llm = LLM(
                model=self.model_id,
                tensor_parallel_size=tensor_parallel_size,
                max_model_len=self.max_context,
                gpu_memory_utilization=0.95,
                enable_prefix_caching=True,  # Critical for long context reuse
                trust_remote_code=True
            )
        except Exception as e:
            raise RuntimeError(f"Failed to initialize Llama 3.1 70B: {str(e)}")

    def _generate_long_context(self, base_code: str, target_length: int) -> str:
        """Generate a long context prompt by repeating and extending base code"""
        # Repeat base code until we reach ~target_length tokens (approx 4 chars per token)
        current_length = len(base_code)
        long_context = base_code
        while current_length < target_length * 4:
            long_context += f"\n# Extended context line {current_length}\n{base_code[:100]}"
            current_length += len(base_code) + 50
        return long_context[:target_length * 4]  # Truncate to exact length

    def run_test(self, context_length: int, num_tests: int = 5) -> Dict[str, Any]:
        """Run stress test for a given context length"""
        if not self.llm:
            raise RuntimeError("LLM not initialized")

        # Load base code (Django model example)
        base_code = """
from django.db import models

class User(models.Model):
    username = models.CharField(max_length=150, unique=True)
    email = models.EmailField(unique=True)
    created_at = models.DateTimeField(auto_now_add=True)

    def __str__(self):
        return self.username
"""
        # Generate long context
        long_context = self._generate_long_context(base_code, context_length)
        prompt = f"{long_context}\n\n# Write a Django view to list all users\n"

        results = {
            "context_length": context_length,
            "num_tests": num_tests,
            "passed": 0,
            "failed": 0,
            "avg_latency_ms": 0.0,
            "correct_dependencies": 0
        }

        for i in range(num_tests):
            start_time = time.time()
            try:
                outputs = self.llm.generate([prompt], self.sampling_params)
                generated_code = outputs[0].outputs[0].text
            except Exception as e:
                print(f"Test {i} failed: {str(e)}")
                results["failed"] += 1
                continue
            latency = (time.time() - start_time) * 1000
            results["avg_latency_ms"] += latency

            # Check if generated code uses the User model correctly (dependency tracking)
            if "User.objects.all()" in generated_code and "from django.shortcuts import render" in generated_code:
                results["correct_dependencies"] += 1
                results["passed"] += 1
            else:
                results["failed"] += 1

        # Calculate averages
        if results["num_tests"] > 0:
            results["avg_latency_ms"] /= results["num_tests"]
            results["dependency_accuracy"] = (results["correct_dependencies"] / results["num_tests"]) * 100
        return results

    def run_full_suite(self) -> List[Dict[str, Any]]:
        """Run tests across 16k, 32k, 64k, 128k context lengths"""
        context_lengths = [16384, 32768, 65536, 131072]
        all_results = []
        for length in context_lengths:
            print(f"Running test for {length} token context...")
            test_results = self.run_test(length)
            all_results.append(test_results)
            print(f"Context {length}: Dependency Accuracy {test_results['dependency_accuracy']:.1f}%, Avg Latency {test_results['avg_latency_ms']:.0f}ms")
        return all_results

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run Llama 3.1 70B context stress test")
    parser.add_argument("--output-file", type=str, default="context_test_results.json")
    args = parser.parse_args()

    # Initialize test harness
    test_harness = ContextStressTest()
    test_harness.init_llm()

    # Run full test suite
    results = test_harness.run_full_suite()

    # Save results
    with open(args.output_file, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Context stress test complete. Results saved to {args.output_file}")

from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x22B-Instruct-v0.3")
# MoE-LoRA config targeting only active experts
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=[\"experts.0\", \"experts.1\"],  # Only active experts
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=2,
    max_model_len=131072,
    enable_prefix_caching=True,  # Enable prefix caching
    gpu_memory_utilization=0.95
)

import subprocess
import json
from typing import Tuple

def filter_hallucinations(generated_code: str) -> Tuple[bool, str]:
    """Return (is_safe, reason) for generated Python code"""
    # Write code to temp file
    with open("temp_generated.py", "w") as f:
        f.write(generated_code)
    # Run Bandit static analysis
    result = subprocess.run(
        ["bandit", "-r", "temp_generated.py", "-f", "json"],
        capture_output=True,
        text=True
    )
    # Parse Bandit output
    if result.returncode == 0:
        return (True, "No security issues found")
    # Extract high-severity issues
    issues = json.loads(result.stdout).get("results", [])
    high_severity = [i for i in issues if i["issue_severity"] == "high"]
    if high_severity:
        return (False, f"High severity issues: {[i['issue_text'] for i in high_severity]}")
    return (True, "Low/medium severity issues only")