env/simulator.py

from __future__ import annotations

from dataclasses import dataclass

from .episodes import generate_seeded_scenario, get_hand_authored_scenario
from .models import (
    ActionKind,
    ActionLogEntry,
    PriorityTier,
    ScenarioConfig,
    SchedulingAction,
    SchedulingMetrics,
    SchedulingObservation,
    SchedulingStepResult,
)


PRIORITY_BONUS = {
    PriorityTier.low: 1.0,
    PriorityTier.medium: 2.0,
    PriorityTier.high: 4.0,
}


@dataclass
class RewardConfig:
    completion_reward: float = 10.0
    missed_deadline_penalty: float = -8.0
    invalid_action_penalty: float = -5.0
    idle_capacity_penalty: float = -0.5


class GPUSchedulingSimulator:
    def __init__(self, scenario: ScenarioConfig, reward_config: RewardConfig | None = None):
        self.reward_config = reward_config or RewardConfig()
        self._base_scenario = scenario.model_copy(deep=True)
        self._load_scenario(self._base_scenario)

    @classmethod
    def from_hand_authored(cls, scenario_id: str) -> "GPUSchedulingSimulator":
        return cls(get_hand_authored_scenario(scenario_id))

    @classmethod
    def from_seed(cls, seed: int) -> "GPUSchedulingSimulator":
        return cls(generate_seeded_scenario(seed))

    def _load_scenario(self, scenario: ScenarioConfig) -> None:
        self.scenario_id = scenario.scenario_id
        self.tick_limit = scenario.tick_limit
        self.current_tick = 0
        self.gpus = {gpu.gpu_id: gpu.model_copy(deep=True) for gpu in scenario.gpus}
        self.jobs = {job.job_id: job.model_copy(deep=True) for job in scenario.jobs}
        self.metrics = SchedulingMetrics()
        self.action_log: list[ActionLogEntry] = []
        self.last_event = "scenario_reset"
        self._allocation_sum = 0.0

    def reset(self, scenario: ScenarioConfig | None = None) -> SchedulingObservation:
        if scenario is not None:
            self._base_scenario = scenario.model_copy(deep=True)
        self._load_scenario(self._base_scenario)
        return self.observe()

    def observe(self) -> SchedulingObservation:
        pending = [job.model_copy(deep=True) for job in self.jobs.values() if job.status == "pending"]
        running = [job.model_copy(deep=True) for job in self.jobs.values() if job.status == "running"]
        completed = [job.model_copy(deep=True) for job in self.jobs.values() if job.status == "completed"]
        missed = [job.model_copy(deep=True) for job in self.jobs.values() if job.status == "missed"]
        gpu_list = [gpu.model_copy(deep=True) for gpu in self.gpus.values()]
        return SchedulingObservation(
            scenario_id=self.scenario_id,
            current_tick=self.current_tick,
            tick_limit=self.tick_limit,
            pending_jobs=sorted(pending, key=lambda job: job.job_id),
            running_jobs=sorted(running, key=lambda job: job.job_id),
            completed_jobs=sorted(completed, key=lambda job: job.job_id),
            missed_jobs=sorted(missed, key=lambda job: job.job_id),
            gpus=sorted(gpu_list, key=lambda gpu: gpu.gpu_id),
            metrics=self.metrics.model_copy(deep=True),
            last_event=self.last_event,
            action_log=[entry.model_copy(deep=True) for entry in self.action_log],
        )

    def step(self, action: SchedulingAction) -> SchedulingStepResult:
        score = 0.0
        valid, details = self._apply_action(action)
        if not valid:
            score += self.reward_config.invalid_action_penalty
            self.metrics.invalid_actions += 1
        progress_score, event = self._advance_time()
        score += progress_score
        self.current_tick += 1
        self.metrics.total_reward += score
        self._update_metrics()
        done = self._is_done()
        self.action_log.append(
            ActionLogEntry(
                tick_index=self.current_tick,
                action=action.model_dump_json(),
                rationale=action.rationale,
                score_delta=score,
                valid=valid,
                details=f"{details}; {event}",
            )
        )
        return SchedulingStepResult(
            observation=self.observe(),
            score=score,
            done=done,
            info={"valid_action": valid, "details": f"{details}; {event}", "scenario_id": self.scenario_id},
        )

    def _apply_action(self, action: SchedulingAction) -> tuple[bool, str]:
        if action.action == ActionKind.defer:
            self.last_event = "deferred"
            return True, self.last_event
        if action.job_id is None or action.gpu_id is None:
            return False, "place requires job_id and gpu_id"
        if action.job_id not in self.jobs:
            return False, f"unknown job {action.job_id}"
        if action.gpu_id not in self.gpus:
            return False, f"unknown gpu {action.gpu_id}"
        job = self.jobs[action.job_id]
        gpu = self.gpus[action.gpu_id]
        if job.status != "pending":
            return False, f"job {job.job_id} is not pending"
        if gpu.busy:
            return False, f"gpu {gpu.gpu_id} is allocated"
        if job.vram_requirement > gpu.vram_capacity:
            return False, f"job {job.job_id} exceeds {gpu.gpu_id} VRAM"
        job.status = "running"
        job.assigned_gpu_id = gpu.gpu_id
        gpu.busy = True
        gpu.current_job_id = job.job_id
        self.last_event = f"placed {job.job_id} on {gpu.gpu_id}"
        return True, self.last_event

    def _advance_time(self) -> tuple[float, str]:
        score = 0.0
        events: list[str] = []
        for gpu in self.gpus.values():
            if not gpu.busy or gpu.current_job_id is None:
                continue
            job = self.jobs[gpu.current_job_id]
            job.remaining_runtime = max(0.0, job.remaining_runtime - gpu.speed_multiplier)
            score -= gpu.cost_per_step
            self.metrics.total_cost += gpu.cost_per_step
            if job.remaining_runtime <= 0:
                job.status = "completed"
                gpu.busy = False
                gpu.current_job_id = None
                score += self.reward_config.completion_reward + PRIORITY_BONUS[job.priority]
                self.metrics.completed_jobs += 1
                events.append(f"completed {job.job_id}")

        for job in self.jobs.values():
            if job.status == "pending":
                job.pending_ticks += 1

        next_tick = self.current_tick + 1
        for job in self.jobs.values():
            if job.status in {"completed", "missed"}:
                continue
            if next_tick > job.deadline:
                if job.assigned_gpu_id:
                    gpu = self.gpus[job.assigned_gpu_id]
                    gpu.busy = False
                    gpu.current_job_id = None
                job.status = "missed"
                score += self.reward_config.missed_deadline_penalty
                self.metrics.missed_deadlines += 1
                events.append(f"missed deadline {job.job_id}")

        pending_jobs = [job for job in self.jobs.values() if job.status == "pending"]
        idle_gpus = [gpu for gpu in self.gpus.values() if not gpu.busy]
        if pending_jobs and idle_gpus:
            score += self.reward_config.idle_capacity_penalty * len(idle_gpus)
            events.append("idle capacity penalty")

        busy_gpu_count = sum(1 for gpu in self.gpus.values() if gpu.busy)
        tick_allocation = busy_gpu_count / max(1, len(self.gpus))
        self._allocation_sum += tick_allocation
        self.metrics.allocation = self._allocation_sum / max(1, self.current_tick + 1)
        self.last_event = " | ".join(events or ["tick advanced"])
        return score, self.last_event

    def _update_metrics(self) -> None:
        pending_ticks = [job.pending_ticks for job in self.jobs.values()]
        self.metrics.average_pending_ticks = sum(pending_ticks) / len(pending_ticks) if pending_ticks else 0.0

    def _is_done(self) -> bool:
        unfinished = [job for job in self.jobs.values() if job.status in {"pending", "running"}]
        return self.current_tick >= self.tick_limit or not unfinished