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 ; ... decode_ppl=`` 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)