DeepInfer/tests/models/test_offline_inference.py

from __future__ import annotations

import json
import math
import os
import re
import signal
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor
from contextlib import suppress
from pathlib import Path

import pytest

from deep_cluster.hardware import accelerator_kind
from examples._config import DEFAULT_MODEL_NAME

GOLDEN_FILE = Path(__file__).with_name("offline_decode_ppl.json")

# The parallel configuration is bound to the golden: any change requires re-measuring that platform/workload's baseline.
# On Ascend the model does not fit on 4 GPUs, so 2 nodes × 8 GPUs are used (EP16/CP8); NVIDIA uses 4 GPUs (EP4/CP2). Both have dp_size 2.
EP_SIZE, CP_SIZE = (16, 8) if accelerator_kind() == "ascend" else (4, 2)
DP_SIZE = EP_SIZE // CP_SIZE

# Multiple nodes: examples/_launch.run_dp splits DP ranks across nodes by RANK/WORLD_SIZE, so this process's stdout
# only covers the ranks this node owns. Each node compares only its own ranks' golden entries (the coverage is a division of labor, not a gap).
NODE_RANK = int(os.environ.get("RANK", "0"))
NUM_NODES = int(os.environ.get("WORLD_SIZE", "1"))
assert DP_SIZE % NUM_NODES == 0, f"dp_size({DP_SIZE}) must be a multiple of the node count ({NUM_NODES}) (a _launch.run_dp constraint)"
RANKS_PER_NODE = DP_SIZE // NUM_NODES
LOCAL_DP_RANKS = list(range(NODE_RANK * RANKS_PER_NODE, (NODE_RANK + 1) * RANKS_PER_NODE))


def _load_expected_ppls(*, long_context: bool) -> dict[int, list[float]]:
    """Read the per-rank golden for one workload; the baseline is independent of the model directory passed in."""
    platform = accelerator_kind()
    workload = "long" if long_context else "short"
    data = json.loads(GOLDEN_FILE.read_text())
    expected = data.get(platform, {}).get(workload)
    assert expected, f"Missing {platform}/{workload} PPL goldens in {GOLDEN_FILE}; measure and review this workload first"
    values = {int(rank): ppls for rank, ppls in expected.items()}
    # DP ranks must be a contiguous run starting at 0; how many is decided by the runtime parallel configuration, and _run_offline_inference compares them against the measured values.
    assert set(values) == set(range(len(values))), f"PPL goldens must cover contiguous DP ranks from 0, got {sorted(values)}"
    assert all(ppls and all(math.isfinite(ppl) and ppl > 0 for ppl in ppls) for ppls in values.values()), "Invalid PPL goldens"
    return values


def _parse_decode_ppls(output: str) -> dict[int, list[float]]:
    """Parse ``DP <rank>; ... decode_ppl=<value>`` lines from offline_inference output.

    Returns a dict mapping dp_rank -> list of decode_ppl values (in output order).
    """
    pattern = re.compile(r"^(?:.*?)DP\s+(\d+);.*decode_ppl=([\d.eE+-]+)", re.MULTILINE)
    result: dict[int, list[float]] = {}
    for m in pattern.finditer(output):
        rank = int(m.group(1))
        ppl = float(m.group(2))
        result.setdefault(rank, []).append(ppl)
    return result


def _stream_output(proc: subprocess.Popen[str], label: str) -> str:
    stream = proc.stdout
    assert stream is not None
    lines: list[str] = []
    with stream:
        for line in stream:
            lines.append(line)
            print(f"[{label}] {line}", end="", flush=True)
    return "".join(lines)


def _run_subprocess(command: list[str], env: dict[str, str], label: str, timeout: int) -> str:
    env["PYTHONUNBUFFERED"] = "1"
    proc = subprocess.Popen(
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        env=env,
        cwd=Path(__file__).resolve().parents[2],
        start_new_session=True,
    )

    with ThreadPoolExecutor(max_workers=1) as output_executor:
        output_future = output_executor.submit(_stream_output, proc, label)
        try:
            proc.wait(timeout=timeout)
        except BaseException:
            with suppress(ProcessLookupError):
                os.killpg(proc.pid, signal.SIGKILL)
            proc.wait()
            raise
        output = output_future.result()

    assert proc.returncode == 0, f"{label} failed with returncode {proc.returncode}"
    return output


def _collect_offline_ppls(
    model_name: str, device_offset: int, *, long_context: bool = False, check_structured_output: bool = False
) -> dict[int, list[float]]:
    env = os.environ.copy()
    env["DEVICE_OFFSET"] = str(device_offset)
    env["MASTER_PORT"] = str(23333 + device_offset)
    # Match the native reference; alias spelling must not implicitly enable speculation.
    env["DEEPINFER_USE_NEXTN_DECODE"] = "5"
    if accelerator_kind() == "nvidia":
        env["DEEPINFER_BATCH_INVARIANT"] = "1"
    command = [
        sys.executable,
        "examples/offline_inference.py",
        f"--model_name={model_name}",
        f"--ep={EP_SIZE}",
        f"--cp={CP_SIZE}",
    ]
    if long_context:
        command.append("--long_context")
    if check_structured_output:
        command.append("--check_structured_output")
    output = _run_subprocess(
        command,
        env,
        model_name,
        1800 if long_context else 600,
    )

    decode_ppls = _parse_decode_ppls(output)
    assert set(decode_ppls) == set(LOCAL_DP_RANKS), (
        f"node (RANK={NODE_RANK}/{NUM_NODES}) expects DP ranks {LOCAL_DP_RANKS} (ep={EP_SIZE}, cp={CP_SIZE}), but the output has {sorted(decode_ppls)}"
    )
    for rank, values in decode_ppls.items():
        assert len(values) == 5 * (rank + 1), f"DP {rank}: unexpected sample count {len(values)}"
        assert all(math.isfinite(ppl) and ppl > 0 for ppl in values), f"DP {rank}: invalid PPL values"
    return decode_ppls


def _run_offline_inference(
    model_name: str, device_offset: int, *, long_context: bool = False, check_structured_output: bool = False
) -> None:
    expected_ppls = _load_expected_ppls(long_context=long_context)
    # The golden's rank set must match the current parallel configuration: extra or missing entries mean stale or unmeasured items.
    assert set(expected_ppls) == set(range(DP_SIZE)), (
        f"the golden's DP ranks are {sorted(expected_ppls)}, current dp_size={DP_SIZE} (ep={EP_SIZE}, cp={CP_SIZE});"
        "re-measure stale or missing entries with examples/refresh_offline_decode_ppl.py"
    )
    decode_ppls = _collect_offline_ppls(
        model_name, device_offset, long_context=long_context, check_structured_output=check_structured_output
    )

    assert set(decode_ppls) == set(LOCAL_DP_RANKS)
    for rank in LOCAL_DP_RANKS:
        expected_values = expected_ppls[rank]
        actual_values = decode_ppls[rank]
        assert len(actual_values) == len(expected_values), f"DP {rank}: golden sample count mismatch"
        for sample, (actual, expected) in enumerate(zip(actual_values, expected_values, strict=True)):
            assert actual == expected, f"DP {rank}, sample {sample}: expected PPL {expected}, got {actual}"


@pytest.mark.ci
@pytest.mark.ci_npu
def test_offline_inference() -> None:
    _run_offline_inference(DEFAULT_MODEL_NAME, 0, check_structured_output=True)