server/gpu_scheduler_openenv_environment.py

from __future__ import annotations

from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.interfaces import EnvironmentMetadata

from env.environment import GPUInferenceSchedulingEnv
from env.models import SchedulingAction, SchedulingObservation
from models import GPUSchedulerAction, GPUSchedulerObservation, GPUSchedulerState


class GPUSchedulerOpenEnvEnvironment(
    Environment[GPUSchedulerAction, GPUSchedulerObservation, GPUSchedulerState]
):
    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self):
        super().__init__()
        self._env = GPUInferenceSchedulingEnv()
        initial_scenario_id = str(uuid4())
        self._state = GPUSchedulerState(
            scenario_id=initial_scenario_id,
            episode_id=initial_scenario_id,
            step_count=0,
            current_tick=0,
        )
        self._last_step_valid_action: bool | None = None
        self._last_step_details: str | None = "Environment ready. Reset to start a new scenario."

    def reset(
        self,
        seed: int | None = None,
        scenario_id: str | None = None,
        **kwargs,
    ) -> GPUSchedulerObservation:
        observation = self._env.reset(seed=seed, scenario_id=scenario_id)
        self._last_step_valid_action = None
        self._last_step_details = "Scenario reset. Choose a placement or inspect the current state."
        openenv_observation = self._build_observation(observation, reward=0.0, done=False)
        self._state = self._build_state(observation)
        return openenv_observation

    def step(
        self,
        action: GPUSchedulerAction,
        timeout_s: float | None = None,
        **kwargs,
    ) -> GPUSchedulerObservation:
        internal_action = SchedulingAction(
            action=action.action,
            job_id=action.job_id,
            gpu_id=action.gpu_id,
            rationale=action.rationale,
        )
        result = self._env.step(internal_action)
        observation = result.observation
        self._last_step_valid_action = bool(result.info.get("valid_action", True))
        self._last_step_details = str(result.info.get("details", ""))
        self._state = self._build_state(observation)
        return self._build_observation(
            observation,
            reward=result.score,
            done=result.done,
        )

    @property
    def state(self) -> GPUSchedulerState:
        return self._state

    def get_metadata(self) -> EnvironmentMetadata:
        return EnvironmentMetadata(
            name="GPU Scheduler OpenEnv",
            description="Schedule pending inference jobs onto a small GPU cluster under VRAM, deadline, and cost pressure.",
            version="0.1.0",
            author="OpenEnv hackathon project",
            documentation_url="https://meta-pytorch.org/OpenEnv/",
        )

    def _build_state(self, observation: SchedulingObservation) -> GPUSchedulerState:
        return GPUSchedulerState(
            scenario_id=observation.scenario_id,
            episode_id=observation.scenario_id,
            step_count=observation.current_tick,
            current_tick=observation.current_tick,
            tick_limit=observation.tick_limit,
            summary=self._build_summary(observation),
            action_hint=self._build_action_hint(observation),
            last_step_valid_action=self._last_step_valid_action,
            last_step_details=self._last_step_details,
            pending_count=len(observation.pending_jobs),
            running_count=len(observation.running_jobs),
            completed_count=len(observation.completed_jobs),
            missed_count=len(observation.missed_jobs),
            priority_alerts=[],
            feasible_assignments=[],
            suggested_assignments=[],
            gpu_status_lines=self._build_gpu_status_lines(observation),
            queue_status_lines=self._build_queue_status_lines(observation),
            pending_jobs=observation.pending_jobs,
            running_jobs=observation.running_jobs,
            completed_jobs=observation.completed_jobs,
            missed_jobs=observation.missed_jobs,
            gpus=observation.gpus,
            action_log=observation.action_log,
            metrics=observation.metrics,
            last_event=observation.last_event,
        )

    def _build_observation(
        self,
        observation: SchedulingObservation,
        reward: float,
        done: bool,
    ) -> GPUSchedulerObservation:
        return GPUSchedulerObservation(
            **observation.model_dump(mode="python"),
            summary=self._build_summary(observation),
            action_hint=self._build_action_hint(observation),
            last_step_valid_action=self._last_step_valid_action,
            last_step_details=self._last_step_details,
            priority_alerts=[],
            feasible_assignments=[],
            suggested_assignments=[],
            gpu_status_lines=self._build_gpu_status_lines(observation),
            queue_status_lines=self._build_queue_status_lines(observation),
            done=done,
            reward=reward,
        )

    def _build_summary(self, observation: SchedulingObservation) -> str:
        allocation = len(observation.running_jobs) / max(1, len(observation.gpus))
        return (
            f"Tick {observation.current_tick}/{observation.tick_limit}. "
            f"Pending {len(observation.pending_jobs)}, running {len(observation.running_jobs)}, "
            f"completed {len(observation.completed_jobs)}, missed {len(observation.missed_jobs)}. "
            f"GPU allocation {allocation:.2f}."
        )

    def _build_action_hint(self, observation: SchedulingObservation) -> str:
        if not observation.pending_jobs:
            return "No pending jobs remain. Use defer to confirm run completion."
        return "Inspect pending jobs and eligible GPUs, then choose a placement or defer."

    def _build_gpu_status_lines(self, observation: SchedulingObservation) -> list[str]:
        return [
            (
                f"{gpu.gpu_id}: {'busy' if gpu.busy else 'idle'}, "
                f"vram={gpu.vram_capacity}, speed={gpu.speed_multiplier}, cost={gpu.cost_per_step}, "
                f"job={gpu.current_job_id or 'none'}"
            )
            for gpu in observation.gpus
        ]

    def _build_queue_status_lines(self, observation: SchedulingObservation) -> list[str]:
        lines: list[str] = []
        for job in sorted(observation.pending_jobs, key=lambda item: (item.deadline, item.job_id))[:6]:
            lines.append(
                f"{job.job_id}: priority={job.priority.value}, deadline={job.deadline}, "
                f"runtime={job.remaining_runtime}, vram={job.vram_requirement}, pending_ticks={job.pending_ticks}"
            )
        return lines