vllm-omni/tests/edge/local/test_plan.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Admission: say yes with a budget, or no with a reason and a remedy."""

import pytest

from vllm_omni.edge.local.plan import (
    GiB,
    REFUSE_CAPACITY,
    REFUSE_FORMAT,
    REFUSE_NO_DEVICE,
    plan_text_session,
)

from .conftest import FP8_QUANT_CONFIG, blackwell_laptop, cpu_only_box, write_checkpoint

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


def _plan(model_dir, profile=None, platform="cuda", **kw):
    return plan_text_session(
        str(model_dir),
        profile=profile if profile is not None else blackwell_laptop(),
        platform_device_type=platform,
        **kw,
    )


def _codes(plan):
    return {r.code for r in plan.refusals}


# ------------------------------------------------------------------ the yes
def test_a_dense_checkpoint_is_admitted_to_the_gpu(dense_checkpoint):
    plan = _plan(dense_checkpoint)
    assert plan.admitted
    assert plan.selected.device_id == "cuda:0"
    assert plan.backend == "vllm:cuda"


def test_an_admitted_plan_prices_every_line_of_section_6_2(dense_checkpoint):
    """The budget is not 'the weights'. It is the sum the design specifies,
    and each term has to be visible so a reviewer can argue with one."""
    plan = _plan(dense_checkpoint)
    purposes = {r.purpose for r in plan.reservations}
    assert purposes == {
        "weights", "kv_cache", "activations", "backend_workspace",
        "load_transient", "external_reserve", "safety_margin",
    }
    assert all(r.pool == "vram" for r in plan.reservations)
    assert all(r.evidence in {"D", "E"} for r in plan.reservations)
    # The weights line is read off the files, not estimated.
    weights = next(r for r in plan.reservations if r.purpose == "weights")
    assert weights.evidence == "D"
    assert weights.bytes == plan.manifest.weight_bytes


def test_graph_capture_is_charged_only_when_graphs_are_on(dense_checkpoint):
    eager = _plan(dense_checkpoint, enforce_eager=True)
    graphs = _plan(dense_checkpoint, enforce_eager=False)
    assert "graph_capture" not in {r.purpose for r in eager.reservations}
    assert "graph_capture" in {r.purpose for r in graphs.reservations}
    assert graphs.peak_bytes > eager.peak_bytes


def test_the_peak_is_a_max_not_a_sum(dense_checkpoint):
    """The loader's transient is gone before the KV pool is sized, so adding
    both would refuse plans that fit."""
    plan = _plan(dense_checkpoint)
    everything = sum(r.bytes for r in plan.reservations)
    assert plan.peak_bytes <= everything


def test_kv_reservation_tracks_the_context(dense_checkpoint):
    short = _plan(dense_checkpoint, max_model_len=1024)
    long = _plan(dense_checkpoint, max_model_len=8192)
    kv_short = next(r for r in short.reservations if r.purpose == "kv_cache").bytes
    kv_long = next(r for r in long.reservations if r.purpose == "kv_cache").bytes
    assert kv_long > kv_short


def test_sliding_layers_are_not_charged_full_length(dense_checkpoint):
    """21 of 28 layers hold a 512 window, not the whole context. Charging them
    all at full length is the arithmetic the kv_budget module exists to fix."""
    plan = _plan(dense_checkpoint, max_model_len=8192)
    assert plan.kv_budget.sliding_layers == 21
    assert plan.kv_budget.full_layers == 7
    assert plan.kv_budget.bytes_total_hybrid < plan.kv_budget.bytes_total


# ------------------------------------------------------------------- the no
def test_int8_is_refused_on_blackwell_although_it_fits(int8_checkpoint):
    """4 MiB of weights against 24 GiB of VRAM: every capacity check passes and
    the answer is still no."""
    plan = _plan(int8_checkpoint)
    assert not plan.admitted
    assert REFUSE_FORMAT in _codes(plan)
    refusal = next(r for r in plan.refusals if r.code == REFUSE_FORMAT)
    assert refusal.device_id == "cuda:0"
    assert "int8" in refusal.message
    assert "omni-cpu" in refusal.remedy


def test_the_same_int8_checkpoint_is_admitted_on_the_cpu(int8_checkpoint):
    """The remedy the refusal names has to actually be a plan."""
    plan = _plan(int8_checkpoint, profile=cpu_only_box(), platform="cpu")
    assert plan.admitted
    assert plan.selected.device_id == "cpu"
    assert plan.backend == "vllm:cpu"


def test_fp8_is_admitted_where_int8_is_not(tmp_path):
    fp8 = write_checkpoint(tmp_path / "fp8", quantization=FP8_QUANT_CONFIG)
    assert _plan(fp8).admitted


def test_a_checkpoint_too_large_for_the_pool_is_refused_with_a_number(tmp_path):
    huge = write_checkpoint(tmp_path / "huge", weight_bytes=40 * GiB, shards=8)
    plan = _plan(huge)
    assert not plan.admitted
    assert REFUSE_CAPACITY in _codes(plan)
    refusal = next(r for r in plan.refusals if r.code == REFUSE_CAPACITY)
    assert "GiB" in refusal.message


def test_a_capacity_refusal_never_silently_shortens_the_context(tmp_path):
    """Trimming max_model_len is offered as a remedy the operator applies, not
    something the planner does on its behalf."""
    # Sized so the overrun is smaller than the KV cache it asked for: that is
    # the case where a shorter context would in fact fit, which is exactly when
    # an engine that "helpfully" truncates would do damage.
    big = write_checkpoint(tmp_path / "big", weight_bytes=16 * GiB, shards=8)
    plan = _plan(big, max_model_len=131072)
    assert not plan.admitted
    refusal = next(r for r in plan.refusals if r.code == REFUSE_CAPACITY)
    assert "max-model-len" in refusal.remedy
    assert "forbids" in refusal.remedy
    # And the plan still reports the context that was asked for, not a
    # quietly reduced one.
    assert plan.request_limits["max_model_len"] == 131072


def test_masking_the_gpu_leaves_nothing_runnable_in_a_cuda_process(dense_checkpoint):
    """The capability-masked startup of the acceptance run: a refusal here is
    the correct outcome, not a failure."""
    plan = _plan(dense_checkpoint, mask=frozenset({"gpu_discrete"}))
    assert not plan.admitted
    assert REFUSE_NO_DEVICE in _codes(plan)
    assert any("mask" in r.message for r in plan.refusals)


def test_every_unusable_device_is_explained(dense_checkpoint):
    plan = _plan(dense_checkpoint)
    explained = {r.device_id for r in plan.refusals}
    assert {"cpu", "npu:amd", "gpu_integrated:amd"} <= explained
    assert all(r.message and r.remedy for r in plan.refusals)


# ------------------------------------------------------------ engine kwargs
def test_the_memory_fraction_is_derived_not_defaulted(dense_checkpoint):
    """vLLM reads this field as host RAM on the CPU backend, so leaving it at
    its default there asks for 92% of the machine. It is always derived."""
    for platform, profile in (("cuda", blackwell_laptop()), ("cpu", cpu_only_box())):
        plan = _plan(dense_checkpoint, profile=profile, platform=platform)
        assert plan.admitted
        fraction = plan.engine_kwargs["gpu_memory_utilization"]
        assert 0 < fraction <= 0.95
        assert fraction >= plan.peak_bytes / plan.selected.memory_bytes * 0.5


def test_the_kv_pool_is_absolute_not_leftover(dense_checkpoint):
    plan = _plan(dense_checkpoint)
    assert plan.engine_kwargs["kv_cache_memory_bytes"] == plan.kv_budget.bytes_total


def test_request_limits_reach_the_backend(dense_checkpoint):
    plan = _plan(dense_checkpoint, max_model_len=2048, max_num_seqs=2)
    assert plan.engine_kwargs["max_model_len"] == 2048
    assert plan.engine_kwargs["max_num_seqs"] == 2
    assert plan.engine_kwargs["enforce_eager"] is True


# ---------------------------------------------------------------- the stage
def test_the_plan_declares_exactly_one_text_stage(dense_checkpoint):
    plan = _plan(dense_checkpoint)
    assert plan.stage["stage_id"] == 0
    assert plan.stage["execution_type"] == "LLM_AR"
    assert plan.stage["final_output_type"] == "text"
    assert plan.stage["input_sources"] == []


def test_the_record_round_trips_to_json(dense_checkpoint):
    import json

    plan = _plan(dense_checkpoint)
    text = json.dumps(plan.to_dict(), default=str)
    back = json.loads(text)
    assert back["admitted"]
    assert back["selected_device"]["device_id"] == "cuda:0"
    assert back["manifest"]["artifact_id"] == plan.manifest.artifact_id