env/models.py

from __future__ import annotations

from enum import Enum
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator


class PriorityTier(str, Enum):
    low = "low"
    medium = "medium"
    high = "high"


class ActionKind(str, Enum):
    place = "place"
    defer = "defer"


class GPUState(BaseModel):
    model_config = ConfigDict(extra="forbid")

    gpu_id: str
    vram_capacity: int = Field(ge=1)
    speed_multiplier: float = Field(gt=0)
    cost_per_step: float = Field(ge=0)
    current_job_id: str | None = None
    busy: bool = False


class JobState(BaseModel):
    model_config = ConfigDict(extra="forbid")

    job_id: str
    vram_requirement: int = Field(ge=1)
    estimated_runtime: float = Field(gt=0)
    remaining_runtime: float = Field(ge=0)
    priority: PriorityTier
    deadline: int = Field(ge=0)
    pending_ticks: int = Field(ge=0, default=0)
    assigned_gpu_id: str | None = None
    status: Literal["pending", "running", "completed", "missed"] = "pending"


class SchedulingMetrics(BaseModel):
    model_config = ConfigDict(extra="forbid")

    total_reward: float = 0.0
    completed_jobs: int = 0
    missed_deadlines: int = 0
    invalid_actions: int = 0
    allocation: float = 0.0
    total_cost: float = 0.0
    average_pending_ticks: float = 0.0


class ActionLogEntry(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tick_index: int
    action: str
    rationale: str
    score_delta: float
    valid: bool
    details: str


class SchedulingObservation(BaseModel):
    model_config = ConfigDict(extra="forbid")

    scenario_id: str
    current_tick: int = Field(ge=0)
    tick_limit: int = Field(ge=1)
    pending_jobs: list[JobState]
    running_jobs: list[JobState]
    completed_jobs: list[JobState]
    missed_jobs: list[JobState]
    gpus: list[GPUState]
    metrics: SchedulingMetrics
    last_event: str | None = None
    action_log: list[ActionLogEntry] = Field(default_factory=list)


class SchedulingAction(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action: ActionKind
    job_id: str | None = None
    gpu_id: str | None = None
    rationale: str = Field(min_length=1, max_length=240)

    @field_validator("rationale")
    @classmethod
    def normalize_rationale(cls, value: str) -> str:
        stripped = value.strip()
        if not stripped:
            raise ValueError("rationale must not be empty")
        return stripped


class SchedulingStepResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    observation: SchedulingObservation
    score: float
    done: bool
    info: dict[str, str | float | int | bool]


class ScenarioConfig(BaseModel):
    model_config = ConfigDict(extra="forbid")

    scenario_id: str
    tick_limit: int = Field(ge=1)
    gpus: list[GPUState]
    jobs: list[JobState]