env/environment.py

from __future__ import annotations

from .episodes import generate_seeded_scenario, get_hand_authored_scenario
from .models import SchedulingAction, SchedulingMetrics, SchedulingObservation, SchedulingStepResult
from .simulator import GPUSchedulingSimulator


class GPUInferenceSchedulingEnv:
    """Thin environment wrapper around the plain Python simulator."""

    def __init__(self, scenario_id: str = "deadline_crunch", seed: int | None = None):
        self.scenario_id = scenario_id
        self.seed = seed
        if seed is not None:
            self.simulator = GPUSchedulingSimulator(generate_seeded_scenario(seed, scenario_id=f"seeded_{seed}"))
        else:
            self.simulator = GPUSchedulingSimulator(get_hand_authored_scenario(scenario_id))

    def reset(self, seed: int | None = None, scenario_id: str | None = None) -> SchedulingObservation:
        if seed is not None:
            self.seed = seed
            self.scenario_id = scenario_id or f"seeded_{seed}"
            self.simulator = GPUSchedulingSimulator(generate_seeded_scenario(seed, scenario_id=self.scenario_id))
        elif scenario_id is not None:
            self.seed = None
            self.scenario_id = scenario_id
            self.simulator = GPUSchedulingSimulator(get_hand_authored_scenario(scenario_id))
        else:
            if self.seed is not None:
                self.simulator = GPUSchedulingSimulator(generate_seeded_scenario(self.seed, scenario_id=self.scenario_id))
            else:
                self.simulator = GPUSchedulingSimulator(get_hand_authored_scenario(self.scenario_id))
        return self.simulator.observe()

    def step(self, action: SchedulingAction) -> SchedulingStepResult:
        return self.simulator.step(action)

    def state(self) -> SchedulingObservation:
        return self.simulator.observe()

    def metrics(self) -> SchedulingMetrics:
        return self.simulator.metrics.model_copy(deep=True)