wehub-resource-sync
148 lines
4.5 KiB
Python
148 lines
4.5 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
Single-node vLLM baseline benchmark for Ray Data LLM batch inference.
|
|
|
|
Measures throughput and supports env-driven thresholds and
|
|
JSON artifact output.
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.llm._internal.batch.benchmark.dataset import ShareGPTDataset
|
|
from ray.llm._internal.batch.benchmark.benchmark_processor import (
|
|
Mode,
|
|
VLLM_SAMPLING_PARAMS,
|
|
benchmark,
|
|
)
|
|
|
|
|
|
# Benchmark constants
|
|
NUM_REQUESTS = 1000
|
|
MODEL_ID = "facebook/opt-1.3b"
|
|
BATCH_SIZE = 64
|
|
CONCURRENCY = 1
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def disable_vllm_compile_cache(monkeypatch):
|
|
"""Disable vLLM compile cache to avoid cache corruption."""
|
|
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def cleanup_ray_resources():
|
|
"""Cleanup Ray resources between tests."""
|
|
yield
|
|
ray.shutdown()
|
|
|
|
|
|
def _get_float_env(name: str, default: float | None = None) -> float | None:
|
|
value = os.getenv(name)
|
|
if value is None or value == "":
|
|
return default
|
|
try:
|
|
return float(value)
|
|
except ValueError:
|
|
raise AssertionError(f"Invalid float for {name}: {value}")
|
|
|
|
|
|
def test_single_node_baseline_benchmark():
|
|
"""
|
|
Single-node baseline benchmark: facebook/opt-1.3b, TP=1, PP=1, 1000 prompts.
|
|
|
|
Logs BENCHMARK_* metrics and optionally asserts perf thresholds from env:
|
|
- RAY_DATA_LLM_BENCHMARK_MIN_THROUGHPUT (req/s)
|
|
- RAY_DATA_LLM_BENCHMARK_MAX_LATENCY_S (seconds)
|
|
Writes JSON artifact to RAY_LLM_BENCHMARK_ARTIFACT_PATH if set.
|
|
"""
|
|
# Dataset setup
|
|
dataset_path = os.getenv(
|
|
"RAY_LLM_BENCHMARK_DATASET_PATH", "/tmp/ray_llm_benchmark_dataset"
|
|
)
|
|
|
|
dataset = ShareGPTDataset(
|
|
dataset_path=dataset_path,
|
|
seed=0,
|
|
hf_dataset_id="Crystalcareai/Code-feedback-sharegpt-renamed",
|
|
hf_split="train",
|
|
truncate_prompt=2048,
|
|
)
|
|
|
|
print(f"Loading {NUM_REQUESTS} prompts from ShareGPT dataset...")
|
|
prompts = dataset.sample(num_requests=NUM_REQUESTS)
|
|
print(f"Loaded {len(prompts)} prompts")
|
|
|
|
ds = ray.data.from_items(prompts)
|
|
|
|
# Benchmark config (single node, TP=1, PP=1)
|
|
print(
|
|
f"\nBenchmark: {MODEL_ID}, batch={BATCH_SIZE}, concurrency={CONCURRENCY}, TP=1, PP=1"
|
|
)
|
|
|
|
# Use benchmark processor to run a single-node vLLM benchmark
|
|
result = benchmark(
|
|
Mode.VLLM_ENGINE,
|
|
ds,
|
|
batch_size=BATCH_SIZE,
|
|
concurrency=CONCURRENCY,
|
|
model=MODEL_ID,
|
|
sampling_params=VLLM_SAMPLING_PARAMS,
|
|
pipeline_parallel_size=1,
|
|
tensor_parallel_size=1,
|
|
)
|
|
|
|
result.show()
|
|
|
|
# Assertions and metrics
|
|
assert result.samples == len(prompts)
|
|
assert result.throughput > 0
|
|
|
|
print("\n" + "=" * 60)
|
|
print("BENCHMARK METRICS")
|
|
print("=" * 60)
|
|
print(f"BENCHMARK_THROUGHPUT: {result.throughput:.4f} req/s")
|
|
print(f"BENCHMARK_LATENCY: {result.elapsed_s:.4f} s")
|
|
print(f"BENCHMARK_SAMPLES: {result.samples}")
|
|
print("=" * 60)
|
|
|
|
# Optional thresholds to fail on regressions
|
|
min_throughput = _get_float_env("RAY_DATA_LLM_BENCHMARK_MIN_THROUGHPUT", 5)
|
|
max_latency_s = _get_float_env("RAY_DATA_LLM_BENCHMARK_MAX_LATENCY_S", 150)
|
|
if min_throughput is not None:
|
|
assert (
|
|
result.throughput >= min_throughput
|
|
), f"Throughput regression: {result.throughput:.4f} < {min_throughput:.4f} req/s"
|
|
if max_latency_s is not None:
|
|
assert (
|
|
result.elapsed_s <= max_latency_s
|
|
), f"Latency regression: {result.elapsed_s:.4f} > {max_latency_s:.4f} s"
|
|
|
|
# Optional JSON artifact emission for downstream ingestion
|
|
artifact_path = os.getenv("RAY_LLM_BENCHMARK_ARTIFACT_PATH")
|
|
if artifact_path:
|
|
metrics = {
|
|
"model": MODEL_ID,
|
|
"batch_size": BATCH_SIZE,
|
|
"concurrency": CONCURRENCY,
|
|
"samples": int(result.samples),
|
|
"throughput_req_per_s": float(result.throughput),
|
|
"elapsed_s": float(result.elapsed_s),
|
|
}
|
|
try:
|
|
os.makedirs(os.path.dirname(artifact_path), exist_ok=True)
|
|
with open(artifact_path, "w", encoding="utf-8") as f:
|
|
json.dump(metrics, f, indent=2, sort_keys=True)
|
|
print(f"Wrote benchmark artifact to: {artifact_path}")
|
|
except Exception as e: # noqa: BLE001
|
|
print(
|
|
f"Warning: failed to write benchmark artifact to {artifact_path}: {e}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|