task.py

"""RL Training Task: LLM Inference Optimization through Quantization"""

import io
import json
import sys
import signal
from collections.abc import Callable
from contextlib import redirect_stdout, redirect_stderr
from typing import Any, TypedDict

from anthropic.types import ToolUnionParam


# Timeout handler
class TimeoutException(Exception):
    pass


def timeout_handler(signum, frame):
    raise TimeoutException("Code execution timed out after 60 seconds")


class PythonToolResult(TypedDict):
    output: str | None
    error: str | None


class SubmitAnswerToolResult(TypedDict):
    answer: Any
    submitted: bool


_PYTHON_NAMESPACE = {
    '__builtins__': __builtins__,
}


def reset_python_namespace():
    global _PYTHON_NAMESPACE
    _PYTHON_NAMESPACE = {
        '__builtins__': __builtins__,
    }


def python_tool(code: str) -> PythonToolResult:
    """Execute Python code with GPU support. Captures stdout/stderr. 60s timeout."""
    try:
        signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(60)
        
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        
        global _PYTHON_NAMESPACE
        
        try:
            with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
                exec(code, _PYTHON_NAMESPACE)
        finally:
            signal.alarm(0)
        
        output = stdout_capture.getvalue()
        errors = stderr_capture.getvalue()
        
        if errors and not output:
            return {"output": None, "error": errors}
        
        return {"output": output + (f"\n[stderr]: {errors}" if errors else ""), "error": None}
    
    except TimeoutException as e:
        signal.alarm(0)
        return {"output": None, "error": str(e)}
        
    except KeyboardInterrupt:
        signal.alarm(0)
        raise
    except Exception as e:
        signal.alarm(0)
        return {"output": None, "error": f"{type(e).__name__}: {str(e)}"}


def submit_answer_tool(answer: str) -> SubmitAnswerToolResult:
    """Submit final quantization code."""
    return {"answer": answer, "submitted": True}


# Tool definitions for Anthropic API
TOOLS: list[ToolUnionParam] = [
    {
        "name": "python",
        "description": (
            "Execute Python code. "
            "Variables and functions persist across calls."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Python code to execute"
                }
            },
            "required": ["code"],
        },
    },
    {
        "name": "submit_answer",
        "description": (
            "Submit your final quantization solution. "
            "Your code must define 'original_model' and 'quantized_model' variables."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "answer": {
                    "type": "string",
                    "description": "Complete Python code defining original_model and quantized_model"
                }
            },
            "required": ["answer"],
        },
    },
]

# Tool handlers mapping
TOOL_HANDLERS: dict[str, Callable[..., Any]] = {
    "python": python_tool,
    "submit_answer": submit_answer_tool,
}


# The challenge prompt
PROMPT = """You are tasked with optimizing LLM inference through quantization.

**OBJECTIVE:**
Quantize a pre-trained language model (e.g., GPT-2) to achieve significant size reduction while maintaining accuracy.

**REQUIREMENTS:**
1. Achieve ≥50% model size reduction (measured in bytes)
2. Maintain accuracy: quantized model perplexity must be ≤1.07x original
3. Your code must define both `original_model` and `quantized_model` variables

**EVALUATION:**
The grader will execute your code, extract both model objects, and independently verify:
- Actual size reduction in model parameters (≥50%)
- Perplexity on test sentences (≤1.07x original)

**IMPORTANT:**
Use this function to measure model size (includes parameters AND buffers):
```python
def get_model_size(model):
    total_bytes = 0
    for param in model.parameters():
        total_bytes += param.numel() * param.element_size()
    for buffer in model.buffers():
        total_bytes += buffer.numel() * buffer.element_size()
    return total_bytes
```

**HINTS:**
Consider exploring different quantization approaches (FP16, INT8, etc.) to find the best trade-off.

Submit your complete solution using the `submit_answer` tool."""



def grading_func(result: Any) -> bool:
    """Validates quantization solution by measuring actual model metrics."""
    signal.alarm(0)
    
    if result is None or not isinstance(result, str):
        return False
    
    try:
        import torch
        from transformers import GPT2LMHeadModel, GPT2Tokenizer
        import numpy as np
        
        test_sentences = [
            "The quick brown fox jumps over the lazy dog.",
            "Machine learning is revolutionizing artificial intelligence.",
            "Quantum computing promises to solve complex problems faster."
        ]
        
        def get_model_size_mb(model):
            total_bytes = 0
            for param in model.parameters():
                total_bytes += param.numel() * param.element_size()
            for buffer in model.buffers():
                total_bytes += buffer.numel() * buffer.element_size()
            return total_bytes / (1024 ** 2)
        
        def calculate_perplexity(model, tokenizer, sentences, device):
            model.eval()
            total_loss = 0
            total_tokens = 0
            
            with torch.no_grad():
                for sentence in sentences:
                    inputs = tokenizer(sentence, return_tensors="pt")
                    input_ids = inputs["input_ids"].to(device)
                    outputs = model(input_ids, labels=input_ids)
                    loss = outputs.loss
                    if loss is not None:
                        total_loss += loss.item() * input_ids.numel()
                        total_tokens += input_ids.numel()
            
            if total_tokens == 0:
                return None
            
            return np.exp(total_loss / total_tokens)
        
        namespace = {'__builtins__': __builtins__}
        
        try:
            with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()):
                exec(result, namespace)
        except Exception:
            return False
        
        if 'quantized_model' not in namespace:
            return False
        if 'original_model' not in namespace:
            if 'model' in namespace and namespace['model'] is not namespace.get('quantized_model'):
                namespace['original_model'] = namespace['model']
            else:
                return False
        
        quantized_model = namespace['quantized_model']
        original_model = namespace['original_model']
        
        model_device = next(original_model.parameters()).device
        tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
        
        original_model.eval()
        original_size_mb = get_model_size_mb(original_model)
        original_perplexity = calculate_perplexity(original_model, tokenizer, test_sentences, str(model_device))
        
        if original_perplexity is None or original_size_mb < 100:
            return False
        
        quantized_device = next(quantized_model.parameters()).device
        quantized_size_mb = get_model_size_mb(quantized_model)
        
        size_reduction = (original_size_mb - quantized_size_mb) / original_size_mb
        if size_reduction < 0.50:
            return False
        
        try:
            quantized_perplexity = calculate_perplexity(quantized_model, tokenizer, test_sentences, str(quantized_device))
        except Exception:
            try:
                quantized_model_cpu = quantized_model.to("cpu")
                quantized_perplexity = calculate_perplexity(quantized_model_cpu, tokenizer, test_sentences, "cpu")
            except Exception:
                return False
        
        if quantized_perplexity is None:
            return False
        
        if not (1 < original_perplexity < 10000) or not (1 < quantized_perplexity < 10000):
            return False
        
        accuracy_ratio = quantized_perplexity / original_perplexity
        return accuracy_ratio <= 1.07
        
    except KeyboardInterrupt:
        raise
    except Exception:
        return False