models.py

from __future__ import annotations

from pydantic import BaseModel, Field

from openenv.core.env_server.types import Action, Observation, State

from env.models import ActionLogEntry, GPUState, JobState, SchedulingMetrics


class SuggestedAssignment(BaseModel):
    job_id: str
    gpu_id: str
    priority: str
    deadline: int
    vram_fit: int
    estimated_finish_step: int
    can_finish_before_deadline: bool


class GPUSchedulerAction(Action):
    action: str = Field(..., description="place or defer")
    job_id: str | None = Field(default=None, description="Job identifier for place actions")
    gpu_id: str | None = Field(default=None, description="GPU identifier for place actions")
    rationale: str = Field(..., min_length=1, max_length=240, description="Short explanation for the decision")


class GPUSchedulerObservation(Observation):
    scenario_id: str
    current_tick: int
    tick_limit: int
    summary: str
    action_hint: str
    last_step_valid_action: bool | None = None
    last_step_details: str | None = None
    priority_alerts: list[str] = Field(default_factory=list)
    feasible_assignments: list[str] = Field(default_factory=list)
    suggested_assignments: list[SuggestedAssignment] = Field(default_factory=list)
    gpu_status_lines: list[str] = Field(default_factory=list)
    queue_status_lines: list[str] = Field(default_factory=list)
    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 GPUSchedulerState(State):
    scenario_id: str = ""
    current_tick: int = 0
    tick_limit: int = 0
    summary: str = ""
    action_hint: str = ""
    last_step_valid_action: bool | None = None
    last_step_details: str | None = None
    pending_count: int = 0
    running_count: int = 0
    completed_count: int = 0
    missed_count: int = 0
    priority_alerts: list[str] = Field(default_factory=list)
    feasible_assignments: list[str] = Field(default_factory=list)
    suggested_assignments: list[SuggestedAssignment] = Field(default_factory=list)
    gpu_status_lines: list[str] = Field(default_factory=list)
    queue_status_lines: list[str] = Field(default_factory=list)
    pending_jobs: list[JobState] = Field(default_factory=list)
    running_jobs: list[JobState] = Field(default_factory=list)
    completed_jobs: list[JobState] = Field(default_factory=list)
    missed_jobs: list[JobState] = Field(default_factory=list)
    gpus: list[GPUState] = Field(default_factory=list)
    action_log: list[ActionLogEntry] = Field(default_factory=list)
    metrics: SchedulingMetrics = Field(default_factory=SchedulingMetrics)
    last_event: str | None = None