backend/queue/job_schema.py

"""
Job Schema for Evaluation Queue

Defines the job schema for asynchronous evaluation processing.
Supports job lifecycle: PENDING -> QUEUED -> RUNNING -> COMPLETED/FAILED/CANCELLED
"""

import uuid
from datetime import datetime
from enum import Enum
from typing import Any, Dict, Optional

from pydantic import BaseModel, Field


class JobStatus(str, Enum):
    """Job status enumeration matching DB schema."""
    PENDING = "pending"
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class JobType(str, Enum):
    """Type of evaluation job."""
    BENCHMARK = "benchmark"
    SINGLE_EVAL = "single_eval"
    ADAPTIVE_EVAL = "adaptive_eval"


class JobPriority(str, Enum):
    """Job priority levels."""
    LOW = "low"
    NORMAL = "normal"
    HIGH = "high"
    CRITICAL = "critical"


class WorkerStatus(str, Enum):
    """Worker status enumeration."""
    REGISTERED = "registered"
    ACTIVE = "active"
    DEGRADED = "degraded"
    OFFLINE = "offline"


class GPURequirement(int, Enum):
    """GPU requirement levels for jobs."""
    CPU_ONLY = 0  # CPU-only job, no GPU needed
    SINGLE_GPU = 1  # Requires 1 GPU
    MULTI_GPU = 2  # Requires multiple GPUs (for future use)


class EvaluationJob(BaseModel):
    """
    Evaluation job schema for async processing.
    
    This is the primary job schema that gets submitted to the queue
    and processed by workers.
    """
    job_id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        description="Unique job identifier"
    )
    job_type: JobType = Field(
        default=JobType.BENCHMARK,
        description="Type of evaluation job"
    )
    model_name: str = Field(
        description="Name of the model to evaluate"
    )
    model_version: str = Field(
        default="latest",
        description="Model version identifier"
    )
    dataset_name: str = Field(
        default="default",
        description="Dataset name"
    )
    dataset_version: str = Field(
        description="Dataset version to use"
    )
    config_hash: str = Field(
        description="SHA256 hash of configuration for reproducibility"
    )
    priority: JobPriority = Field(
        default=JobPriority.NORMAL,
        description="Job priority level"
    )
    submitted_by: Optional[str] = Field(
        default=None,
        description="API key owner who submitted the job"
    )
    status: JobStatus = Field(
        default=JobStatus.PENDING,
        description="Current job status"
    )
    progress: float = Field(
        default=0.0,
        ge=0.0,
        le=100.0,
        description="Job progress percentage (0-100)"
    )
    total_samples: int = Field(
        default=0,
        ge=0,
        description="Total number of samples to process"
    )
    completed_samples: int = Field(
        default=0,
        ge=0,
        description="Number of samples completed"
    )
    failed_samples: int = Field(
        default=0,
        ge=0,
        description="Number of samples failed"
    )
    # Results (populated when job completes)
    composite_score: Optional[float] = Field(
        default=None,
        description="Final composite robustness score (0-1)"
    )
    metrics: Optional[Dict[str, float]] = Field(
        default=None,
        description="Individual metric scores"
    )
    # Timestamps
    created_at: datetime = Field(
        default_factory=datetime.utcnow,
        description="When the job was created"
    )
    queued_at: Optional[datetime] = Field(
        default=None,
        description="When the job was queued"
    )
    started_at: Optional[datetime] = Field(
        default=None,
        description="When the job started processing"
    )
    completed_at: Optional[datetime] = Field(
        default=None,
        description="When the job completed"
    )
    # Error information
    error: Optional[str] = Field(
        default=None,
        description="Error message if job failed"
    )
    error_details: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Detailed error information"
    )
    # Checkpoint information
    last_checkpoint_at: Optional[datetime] = Field(
        default=None,
        description="When the last checkpoint was saved"
    )
    checkpoint_interval: int = Field(
        default=10,
        ge=1,
        description="Save checkpoint every N samples"
    )
    # Worker information
    worker_id: Optional[str] = Field(
        default=None,
        description="ID of the worker processing this job"
    )
    # Metadata
    metadata: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Additional job metadata"
    )
    
    # Cost tracking (Week 7 Day 4 - Intelligent Scheduling)
    estimated_gpu_hours: Optional[float] = Field(
        default=None,
        ge=0.0,
        description="Estimated GPU hours required for this job"
    )
    estimated_cost: Optional[float] = Field(
        default=None,
        ge=0.0,
        description="Estimated cost for this job in USD"
    )
    actual_gpu_hours: Optional[float] = Field(
        default=None,
        ge=0.0,
        description="Actual GPU hours consumed when job completes"
    )
    actual_cost: Optional[float] = Field(
        default=None,
        ge=0.0,
        description="Actual cost when job completes in USD"
    )
    
    # SLA enforcement (Week 7 Day 4 - Intelligent Scheduling)
    deadline_timestamp: Optional[datetime] = Field(
        default=None,
        description="Job deadline for SLA enforcement"
    )
    priority_score: Optional[float] = Field(
        default=None,
        ge=0.0,
        le=1.0,
        description="Computed priority score (0-1) for scheduling"
    )
    queue_type: Optional[str] = Field(
        default=None,
        description="Assigned queue type: high, medium, or low"
    )

    class Config:
        use_enum_values = True


class JobSubmissionRequest(BaseModel):
    """Request model for submitting a new evaluation job."""
    job_type: JobType = Field(
        default=JobType.BENCHMARK,
        description="Type of evaluation job"
    )
    model_name: str = Field(
        description="Name of the model to evaluate",
        examples=["meta-llama/Llama-2-7b-hf"]
    )
    model_version: str = Field(
        default="latest",
        description="Model version"
    )
    dataset_name: str = Field(
        default="default",
        description="Dataset name"
    )
    dataset_version: str = Field(
        description="Dataset version to use"
    )
    priority: JobPriority = Field(
        default=JobPriority.NORMAL,
        description="Job priority"
    )
    mutation_depth: int = Field(
        default=2,
        ge=0,
        le=10,
        description="Mutation depth"
    )
    attack_types: list[str] = Field(
        default_factory=lambda: ["jailbreak"],
        description="List of attack types"
    )
    max_concurrency: int = Field(
        default=4,
        ge=1,
        le=32,
        description="Maximum concurrent samples"
    )
    checkpoint_interval: int = Field(
        default=10,
        ge=1,
        description="Save checkpoint every N samples"
    )
    sampling_config: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Optional sampling configuration"
    )
    # SLA enforcement (Week 7 Day 4 - Intelligent Scheduling)
    deadline_timestamp: Optional[datetime] = Field(
        default=None,
        description="Job deadline for SLA enforcement"
    )


class JobStatusResponse(BaseModel):
    """Response model for job status queries."""
    job_id: uuid.UUID
    job_type: JobType
    status: JobStatus
    progress: float
    total_samples: int
    completed_samples: int
    failed_samples: int
    composite_score: Optional[float] = None
    metrics: Optional[Dict[str, float]] = None
    error: Optional[str] = None
    created_at: datetime
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    worker_id: Optional[str] = None
    # Cost tracking fields
    estimated_gpu_hours: Optional[float] = None
    estimated_cost: Optional[float] = None
    actual_gpu_hours: Optional[float] = None
    actual_cost: Optional[float] = None
    priority_score: Optional[float] = None
    queue_type: Optional[str] = None


class JobProgressUpdate(BaseModel):
    """Model for job progress updates during checkpointing."""
    job_id: uuid.UUID
    completed_samples: int
    failed_samples: int
    composite_score: Optional[float] = None
    metrics: Optional[Dict[str, float]] = None
    checkpoint_at: datetime = Field(default_factory=datetime.utcnow)


__all__ = [
    "JobStatus",
    "JobType",
    "JobPriority",
    "WorkerStatus",
    "GPURequirement",
    "EvaluationJob",
    "JobSubmissionRequest",
    "JobStatusResponse",
    "JobProgressUpdate",
]