tests/test_simulator.py

from env.models import SchedulingAction
from env.simulator import GPUSchedulingSimulator


def test_valid_assignment_completes_job():
    simulator = GPUSchedulingSimulator.from_hand_authored("deadline_crunch")
    result = simulator.step(
        SchedulingAction(
            action="place",
            job_id="req-4",
            gpu_id="small-1",
            rationale="Short urgent job fits and should clear quickly.",
        )
    )
    assert result.score > 0
    assert simulator.jobs["req-4"].status == "completed"
    assert simulator.metrics.completed_jobs == 1


def test_invalid_assignment_is_penalized():
    simulator = GPUSchedulingSimulator.from_hand_authored("deadline_crunch")
    result = simulator.step(
        SchedulingAction(
            action="place",
            job_id="req-5",
            gpu_id="small-1",
            rationale="This should fail due to VRAM pressure.",
        )
    )
    assert result.score < 0
    assert simulator.metrics.invalid_actions == 1
    assert simulator.jobs["req-5"].status == "pending"


def test_defer_advances_time():
    simulator = GPUSchedulingSimulator.from_hand_authored("deadline_crunch")
    initial_tick = simulator.current_tick
    result = simulator.step(SchedulingAction(action="defer", rationale="Holding for a better placement."))
    assert result.observation.current_tick == initial_tick + 1
    assert simulator.metrics.total_reward == result.score


def test_wait_can_outperform_greedy_large_gpu_assignment():
    waiting_simulator = GPUSchedulingSimulator.from_hand_authored("wait_for_capacity")
    greedy_simulator = GPUSchedulingSimulator.from_hand_authored("wait_for_capacity")

    defer_result = waiting_simulator.step(
        SchedulingAction(
            action="defer",
            rationale="Cheaper capacity frees next step, so paying for the large GPU now is wasteful.",
        )
    )
    greedy_result = greedy_simulator.step(
        SchedulingAction(
            action="place",
            job_id="queued-b",
            gpu_id="large-1",
            rationale="Use the only idle GPU immediately.",
        )
    )

    assert defer_result.score > greedy_result.score
    assert waiting_simulator.jobs["queued-b"].status == "pending"
    assert greedy_simulator.jobs["queued-b"].status == "running"


def test_wait_can_preserve_large_gpu_for_visible_heavy_workload():
    waiting_simulator = GPUSchedulingSimulator.from_hand_authored("reserve_large_gpu")
    greedy_simulator = GPUSchedulingSimulator.from_hand_authored("reserve_large_gpu")

    defer_result = waiting_simulator.step(
        SchedulingAction(
            action="defer",
            rationale="Hold the expensive large GPU until the cheaper GPUs finish warmup.",
        )
    )
    greedy_result = greedy_simulator.step(
        SchedulingAction(
            action="place",
            job_id="reserve-small",
            gpu_id="large-1",
            rationale="Use the only idle GPU immediately.",
        )
    )

    assert defer_result.score > greedy_result.score
    assert waiting_simulator.jobs["reserve-heavy"].status == "pending"
    assert greedy_simulator.jobs["reserve-small"].status == "running"


def test_seeded_scenario_is_deterministic():
    sim_a = GPUSchedulingSimulator.from_seed(42)
    sim_b = GPUSchedulingSimulator.from_seed(42)
    assert sim_a.observe().model_dump() == sim_b.observe().model_dump()


def test_scenario_terminates_after_jobs_resolve():
    simulator = GPUSchedulingSimulator.from_hand_authored("deadline_crunch")
    done = False
    while not done:
        result = simulator.step(SchedulingAction(action="defer", rationale="Let the pending queue age."))
        done = result.done
    assert simulator.current_tick <= simulator.tick_limit
    assert all(job.status in {"completed", "missed"} for job in simulator.jobs.values())