wehub-resource-sync
780 lines
24 KiB
Python
780 lines
24 KiB
Python
import sys
|
|
import logging
|
|
import time
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data.llm import (
|
|
build_processor,
|
|
vLLMEngineProcessorConfig,
|
|
ChatTemplateStageConfig,
|
|
DetokenizeStageConfig,
|
|
PrepareMultimodalStageConfig,
|
|
TokenizerStageConfig,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
S3_ARTIFACT_ASSETS_URL = (
|
|
"https://air-example-data.s3.amazonaws.com/rayllm-ossci/assets/"
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def disable_vllm_compile_cache(monkeypatch):
|
|
"""Automatically disable vLLM compile cache for all tests.
|
|
|
|
Avoids AssertionError due to torch compile cache corruption caused by
|
|
running multiple engines on the same node.
|
|
See: https://github.com/vllm-project/vllm/issues/18851, fix expected with
|
|
PyTorch 2.8.0
|
|
"""
|
|
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def add_buffer_time_between_tests():
|
|
"""Add buffer time after each test to avoid resource conflicts, which cause
|
|
flakiness.
|
|
"""
|
|
import gc
|
|
|
|
gc.collect()
|
|
time.sleep(15)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def cleanup_ray_resources():
|
|
"""Automatically cleanup Ray resources between tests to prevent conflicts."""
|
|
yield
|
|
_cleanup_gpu_processes()
|
|
ray.shutdown()
|
|
|
|
|
|
def _cleanup_gpu_processes():
|
|
"""
|
|
Kill GPU processes on all nodes in the cluster. With Ray as the external orchestrator,
|
|
mp backend suffers from uncoordinated shutdown issues, leaving orphaned GPU processes.
|
|
|
|
TODO (jeffreywang): Remove this once https://github.com/vllm-project/vllm/pull/39846 lands.
|
|
"""
|
|
if not ray.is_initialized():
|
|
return
|
|
|
|
@ray.remote(num_cpus=0)
|
|
def _remote_kill_gpu_processes():
|
|
import os
|
|
import signal
|
|
|
|
import pynvml
|
|
|
|
pids = set()
|
|
try:
|
|
pynvml.nvmlInit()
|
|
device_count = pynvml.nvmlDeviceGetCount()
|
|
for i in range(device_count):
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
|
|
for proc in pynvml.nvmlDeviceGetComputeRunningProcesses(handle):
|
|
pids.add(proc.pid)
|
|
pynvml.nvmlShutdown()
|
|
except Exception:
|
|
pass
|
|
|
|
for pid in pids:
|
|
try:
|
|
os.kill(pid, signal.SIGKILL)
|
|
except (ProcessLookupError, ValueError):
|
|
pass
|
|
|
|
try:
|
|
nodes = ray.nodes()
|
|
refs = []
|
|
for node in nodes:
|
|
if not node.get("Alive", False):
|
|
continue
|
|
node_id = node["NodeID"]
|
|
refs.append(
|
|
_remote_kill_gpu_processes.options(
|
|
scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
|
|
node_id=node_id, soft=False
|
|
),
|
|
).remote()
|
|
)
|
|
if refs:
|
|
ray.get(refs, timeout=30)
|
|
except Exception as e:
|
|
logging.warning(f"Failed to kill GPU processes on remote nodes: {e}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vllm_multimodal_utils():
|
|
"""Test vLLM's multimodal utilities.
|
|
|
|
This test is adapted from https://github.com/vllm-project/vllm/blob/main/tests/entrypoints/test_chat_utils.py.
|
|
`parse_chat_messages_async` is thoroughly tested in vLLM. This test serves as an
|
|
integration test to verify that the function isn't moved to an unexpected location and its signature isn't changed.
|
|
"""
|
|
from vllm.config import ModelConfig
|
|
from vllm.entrypoints.chat_utils import parse_chat_messages_async
|
|
|
|
image_url = "https://air-example-data.s3.us-west-2.amazonaws.com/rayllm-ossci/assets/cherry_blossom.jpg"
|
|
image_uuid = str(hash(image_url))
|
|
|
|
conversation, mm_data, mm_uuids = await parse_chat_messages_async(
|
|
[
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": image_url},
|
|
"uuid": image_uuid,
|
|
},
|
|
{"type": "text", "text": "What's in the image?"},
|
|
],
|
|
}
|
|
],
|
|
ModelConfig(
|
|
"microsoft/Phi-3.5-vision-instruct",
|
|
runner="generate",
|
|
trust_remote_code=True,
|
|
limit_mm_per_prompt={"image": 2},
|
|
),
|
|
content_format="string",
|
|
)
|
|
|
|
assert conversation == [
|
|
{"role": "user", "content": "<|image_1|>\nWhat's in the image?"}
|
|
]
|
|
|
|
assert mm_data is not None
|
|
assert set(mm_data.keys()) == {"image"}
|
|
|
|
image_data = mm_data.get("image")
|
|
assert image_data is not None
|
|
|
|
assert isinstance(image_data, list) and len(image_data) == 1
|
|
|
|
assert mm_uuids is not None
|
|
assert "image" in mm_uuids
|
|
|
|
image_uuids = mm_uuids.get("image")
|
|
assert image_uuids is not None
|
|
assert isinstance(image_uuids, list) and len(image_uuids) == 1
|
|
assert image_uuids[0] == image_uuid
|
|
|
|
|
|
def test_chat_template_with_vllm():
|
|
"""Test vLLM with explicit chat template."""
|
|
|
|
processor_config = vLLMEngineProcessorConfig(
|
|
model_source="unsloth/Llama-3.2-1B-Instruct",
|
|
engine_kwargs=dict(
|
|
max_model_len=16384,
|
|
enable_chunked_prefill=True,
|
|
max_num_batched_tokens=2048,
|
|
),
|
|
tokenize_stage=True,
|
|
detokenize_stage=True,
|
|
batch_size=16,
|
|
concurrency=1,
|
|
runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
|
|
)
|
|
|
|
processor = build_processor(
|
|
processor_config,
|
|
preprocess=lambda row: dict(
|
|
messages=[
|
|
{"role": "system", "content": "You are a calculator"},
|
|
{"role": "user", "content": f"{row['id']} ** 3 = ?"},
|
|
],
|
|
sampling_params=dict(
|
|
temperature=0.3,
|
|
max_tokens=50,
|
|
detokenize=False,
|
|
),
|
|
),
|
|
postprocess=lambda row: {
|
|
"resp": row["generated_text"],
|
|
},
|
|
)
|
|
|
|
ds = ray.data.range(60)
|
|
ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
|
|
ds = processor(ds)
|
|
ds = ds.materialize()
|
|
outs = ds.take_all()
|
|
assert len(outs) == 60
|
|
assert all("resp" in out for out in outs)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"tp_size,pp_size,concurrency",
|
|
[
|
|
(2, 1, 2), # TP=2, concurrency=2
|
|
(1, 2, 2), # PP=2, concurrency=2
|
|
],
|
|
)
|
|
def test_vllm_llama_parallel(tp_size, pp_size, concurrency):
|
|
"""Test vLLM with Llama model using different parallelism configurations."""
|
|
|
|
# vLLM v1 does not support decoupled tokenizer,
|
|
# but since the tokenizer is in a separate process,
|
|
# the overhead should be moderated.
|
|
tokenize = False
|
|
detokenize = False
|
|
|
|
processor_config = vLLMEngineProcessorConfig(
|
|
model_source="unsloth/Llama-3.2-1B-Instruct",
|
|
engine_kwargs=dict(
|
|
tensor_parallel_size=tp_size,
|
|
pipeline_parallel_size=pp_size,
|
|
max_model_len=16384,
|
|
enable_chunked_prefill=True,
|
|
max_num_batched_tokens=2048,
|
|
),
|
|
tokenize_stage=tokenize,
|
|
detokenize_stage=detokenize,
|
|
batch_size=16,
|
|
accelerator_type=None,
|
|
concurrency=concurrency,
|
|
runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
|
|
)
|
|
|
|
processor = build_processor(
|
|
processor_config,
|
|
preprocess=lambda row: dict(
|
|
messages=[
|
|
{"role": "system", "content": "You are a calculator"},
|
|
{"role": "user", "content": f"{row['id']} ** 3 = ?"},
|
|
],
|
|
sampling_params=dict(
|
|
temperature=0.3,
|
|
max_tokens=50,
|
|
detokenize=False,
|
|
),
|
|
),
|
|
postprocess=lambda row: {
|
|
"resp": row["generated_text"],
|
|
},
|
|
)
|
|
|
|
ds = ray.data.range(120)
|
|
ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
|
|
ds = processor(ds)
|
|
ds = ds.materialize()
|
|
|
|
# Verify results
|
|
outs = ds.take_all()
|
|
assert len(outs) == 120
|
|
assert all("resp" in out for out in outs)
|
|
|
|
|
|
def test_vllm_llama_lora():
|
|
"""Test vLLM with Llama model and LoRA adapter support.
|
|
|
|
Validates that the LoRA adapter is actually applied during generation by
|
|
using greedy sampling (temperature=0) and short, diverse, open-ended prompts,
|
|
then asserting that a significant number of base-model vs LoRA paired outputs differ.
|
|
|
|
If the LoRA adapter is not being applied, we expect very few pairs to differ
|
|
due to non-determinism from vLLM's dynamic batching and CUDA. If it is being applied,
|
|
we expect a significant fraction of pairs to differ due to the adapter's effect on the output.
|
|
|
|
The two regimes are separated by the `min_diff_fraction` threshold defined below.
|
|
"""
|
|
# Minimum fraction of base/LoRA pairs that must differ to pass.
|
|
# Determined empirically.
|
|
min_diff_fraction = 0.5
|
|
model_source = "s3://air-example-data/llama-3.2-216M-dummy/"
|
|
lora_path = "s3://air-example-data/"
|
|
lora_name = "llama-3.2-216M-lora-dummy"
|
|
max_lora_rank = 32
|
|
|
|
# Short, diverse prompts — shorter prompts let the LoRA perturbation
|
|
# manifest earlier in generation before the base model's momentum
|
|
# dominates the output. Also prefer open-ended prompts to encourage LoRA-induced diversity.
|
|
prompts = [
|
|
"Hello world",
|
|
"The capital of France is",
|
|
"Once upon a time",
|
|
"1 + 1 =",
|
|
"def fibonacci(n):",
|
|
"The quick brown fox",
|
|
"In the beginning",
|
|
"To be or not to be",
|
|
"import numpy as np",
|
|
"The weather today is",
|
|
"My favorite color is",
|
|
"How to cook rice:",
|
|
"The meaning of life is",
|
|
"SELECT * FROM",
|
|
"Dear Sir or Madam,",
|
|
"Breaking news:",
|
|
"A long time ago in a galaxy",
|
|
"The first law of thermodynamics",
|
|
"class MyClass:",
|
|
"Roses are red,",
|
|
"According to recent studies,",
|
|
"Step 1: Preheat the oven",
|
|
"The president announced",
|
|
"In mathematics, a prime number",
|
|
"function hello() {",
|
|
"The cat sat on the",
|
|
"Water boils at",
|
|
"Happy birthday to",
|
|
"ERROR: NullPointerException",
|
|
"The mitochondria is the",
|
|
]
|
|
num_pairs = len(prompts)
|
|
|
|
# The following controls are propagated to the vLLM worker to minimize non-determinism:
|
|
# * engine_kwargs["seed"]: vLLM seeds its internal torch/np/random state.
|
|
# * PYTHONHASHSEED: stabilizes dict/set iteration order in the worker.
|
|
# * CUBLAS_WORKSPACE_CONFIG: forces deterministic cuBLAS workspace.
|
|
# Flash-attention kernels and the vLLM v1 async scheduler still introduce
|
|
# residual non-determinism we can't eliminate from the test, but in
|
|
# practice the regime separation (base/LoRA divergence signal vs noise)
|
|
# is very wide (~30/30 vs ~5/30), so half-of-num_pairs is a robust threshold.
|
|
# Note for future: VLLM_BATCH_INVARIANT=1 with attention_backend=FLASH_ATTN
|
|
# eliminates the remaining non-determinism entirely for perfect separation,
|
|
# but it requires a GPU compute capability (>=9.0) the CI runners don't meet, so
|
|
# we rely on the median threshold instead. Recommended if CI hardware is upgraded.
|
|
seed = 42
|
|
|
|
processor_config = vLLMEngineProcessorConfig(
|
|
model_source=model_source,
|
|
dynamic_lora_loading_path=lora_path,
|
|
engine_kwargs=dict(
|
|
max_model_len=4096,
|
|
enable_chunked_prefill=True,
|
|
enable_lora=True,
|
|
max_lora_rank=max_lora_rank,
|
|
seed=seed,
|
|
),
|
|
tokenize_stage=True,
|
|
detokenize_stage=True,
|
|
# minimize non-determinism from vLLM's dynamic batching by sending (1 base + 1 LoRA) per batch
|
|
batch_size=2,
|
|
concurrency=1,
|
|
runtime_env={
|
|
"env_vars": {
|
|
"VLLM_DISABLE_COMPILE_CACHE": "1",
|
|
"PYTHONHASHSEED": "0",
|
|
"CUBLAS_WORKSPACE_CONFIG": ":4096:8",
|
|
}
|
|
},
|
|
)
|
|
|
|
processor = build_processor(
|
|
processor_config,
|
|
preprocess=lambda row: dict(
|
|
model=model_source if row["id"] % 2 == 0 else lora_name,
|
|
messages=[{"role": "user", "content": row["prompt"]}],
|
|
sampling_params=dict(
|
|
temperature=0,
|
|
max_tokens=50,
|
|
detokenize=False,
|
|
),
|
|
),
|
|
postprocess=lambda row: {
|
|
"id": row["id"],
|
|
"resp": row["generated_text"],
|
|
},
|
|
)
|
|
|
|
# Build paired rows: id 2k = base, id 2k+1 = LoRA, same prompt.
|
|
rows = []
|
|
for i, prompt in enumerate(prompts):
|
|
rows.append({"id": 2 * i, "prompt": prompt})
|
|
rows.append({"id": 2 * i + 1, "prompt": prompt})
|
|
|
|
ds = ray.data.from_items(rows)
|
|
ds = processor(ds)
|
|
ds = ds.materialize()
|
|
outs = ds.take_all()
|
|
|
|
assert len(outs) == 2 * num_pairs
|
|
assert all("resp" in out for out in outs)
|
|
|
|
# LoRA exercise check: a significant fraction of pairs must differ.
|
|
by_id = {out["id"]: out["resp"] for out in outs}
|
|
diffs = sum(1 for k in range(num_pairs) if by_id[2 * k] != by_id[2 * k + 1])
|
|
min_diffs = int(num_pairs * min_diff_fraction)
|
|
assert diffs >= min_diffs, (
|
|
f"Only {diffs}/{num_pairs} base/LoRA pairs differ (need >= {min_diffs}) — "
|
|
"the LoRA adapter does not appear to be applied."
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_source,tp_size,pp_size,concurrency,sample_size,chat_template_content_format,apply_sys_msg_formatting",
|
|
[
|
|
# LLaVA model with TP=1, PP=1, concurrency=1
|
|
("llava-hf/llava-1.5-7b-hf", 1, 1, 1, 60, "openai", False),
|
|
# Pixtral model with TP=2, PP=1, concurrency=2
|
|
("mistral-community/pixtral-12b", 2, 1, 2, 60, "openai", True),
|
|
],
|
|
)
|
|
def test_vllm_vision_language_models(
|
|
model_source,
|
|
tp_size,
|
|
pp_size,
|
|
concurrency,
|
|
sample_size,
|
|
chat_template_content_format,
|
|
apply_sys_msg_formatting,
|
|
):
|
|
"""Test vLLM with vision language models using different configurations."""
|
|
|
|
# vLLM v1 does not support decoupled tokenizer,
|
|
# but since the tokenizer is in a separate process,
|
|
# the overhead should be moderated.
|
|
tokenize = False
|
|
detokenize = False
|
|
|
|
llm_processor_config = vLLMEngineProcessorConfig(
|
|
model_source=model_source,
|
|
task_type="generate",
|
|
engine_kwargs=dict(
|
|
tensor_parallel_size=tp_size,
|
|
pipeline_parallel_size=pp_size,
|
|
max_model_len=4096,
|
|
enable_chunked_prefill=True,
|
|
),
|
|
prepare_multimodal_stage=PrepareMultimodalStageConfig(
|
|
enabled=True,
|
|
chat_template_content_format=chat_template_content_format,
|
|
apply_sys_msg_formatting=apply_sys_msg_formatting,
|
|
),
|
|
chat_template_stage=True,
|
|
tokenize_stage=tokenize,
|
|
detokenize_stage=detokenize,
|
|
batch_size=16,
|
|
concurrency=concurrency,
|
|
runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
|
|
)
|
|
llm_processor = build_processor(
|
|
llm_processor_config,
|
|
preprocess=lambda row: dict(
|
|
messages=[
|
|
{"role": "system", "content": "You are an assistant"},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": f"Say {row['val']} words about this image.",
|
|
},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": S3_ARTIFACT_ASSETS_URL + "cherry_blossom.jpg"
|
|
},
|
|
},
|
|
],
|
|
},
|
|
],
|
|
sampling_params=dict(
|
|
temperature=0.3,
|
|
max_tokens=50,
|
|
),
|
|
),
|
|
postprocess=lambda row: {
|
|
"resp": row["generated_text"],
|
|
},
|
|
)
|
|
|
|
ds = ray.data.range(sample_size)
|
|
ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
|
|
ds = llm_processor(ds)
|
|
ds = ds.materialize()
|
|
outs = ds.take_all()
|
|
assert len(outs) == sample_size
|
|
assert all("resp" in out for out in outs)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"multimodal_content",
|
|
[
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": S3_ARTIFACT_ASSETS_URL + "cherry_blossom.jpg"},
|
|
},
|
|
{
|
|
"type": "video_url",
|
|
"video_url": {"url": S3_ARTIFACT_ASSETS_URL + "free-videos.mp4"},
|
|
},
|
|
],
|
|
)
|
|
def test_vllm_qwen_vl_multimodal(multimodal_content):
|
|
model_source = "Qwen/Qwen2.5-VL-3B-Instruct"
|
|
|
|
llm_processor_config = vLLMEngineProcessorConfig(
|
|
model_source=model_source,
|
|
task_type="generate",
|
|
engine_kwargs=dict(
|
|
enable_chunked_prefill=True,
|
|
distributed_executor_backend="ray",
|
|
# A single GPU won't be able to accomodate Qwen/Qwen2.5-VL-3B-Instruct's memory requirements
|
|
# due to vllm0.12.0 resource/profiling issues.
|
|
# Issue: https://github.com/vllm-project/vllm/issues/30521.
|
|
tensor_parallel_size=2,
|
|
pipeline_parallel_size=1,
|
|
),
|
|
prepare_multimodal_stage=PrepareMultimodalStageConfig(
|
|
enabled=True,
|
|
),
|
|
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
|
tokenize_stage=TokenizerStageConfig(enabled=False),
|
|
detokenize_stage=DetokenizeStageConfig(enabled=False),
|
|
batch_size=16,
|
|
concurrency=1,
|
|
)
|
|
|
|
llm_processor = build_processor(
|
|
llm_processor_config,
|
|
preprocess=lambda row: dict(
|
|
sampling_params=dict(
|
|
temperature=0.3,
|
|
max_tokens=50,
|
|
),
|
|
mm_processor_kwargs=dict(
|
|
min_pixels=28 * 28,
|
|
max_pixels=1280 * 28 * 28,
|
|
fps=1,
|
|
),
|
|
messages=[
|
|
{"role": "system", "content": "You are an assistant"},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": f"Describe this asset in {row['id']} sentences.",
|
|
},
|
|
multimodal_content,
|
|
],
|
|
},
|
|
],
|
|
),
|
|
postprocess=lambda row: {
|
|
"resp": row["generated_text"],
|
|
},
|
|
)
|
|
|
|
ds = ray.data.range(60)
|
|
ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
|
|
ds = llm_processor(ds)
|
|
ds = ds.materialize()
|
|
outs = ds.take_all()
|
|
assert len(outs) == 60
|
|
assert all("resp" in out for out in outs)
|
|
|
|
|
|
@pytest.mark.parametrize("concurrency", [1, 4])
|
|
def test_async_udf_queue_capped(concurrency):
|
|
"""
|
|
Test that the large object in input/output rows
|
|
are stored in object store and does not OOM.
|
|
"""
|
|
|
|
processor_config = vLLMEngineProcessorConfig(
|
|
model_source="unsloth/Llama-3.2-1B-Instruct",
|
|
engine_kwargs=dict(
|
|
max_model_len=16384,
|
|
enable_chunked_prefill=True,
|
|
max_num_batched_tokens=2048,
|
|
),
|
|
tokenize_stage=False,
|
|
detokenize_stage=False,
|
|
batch_size=4,
|
|
accelerator_type=None,
|
|
concurrency=concurrency,
|
|
runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
|
|
)
|
|
|
|
processor = build_processor(
|
|
processor_config,
|
|
preprocess=lambda row: dict(
|
|
# 1M emoji (4 bytes), should not leak to memory heap.
|
|
large_memory_to_carry_over="🤗" * 1_000_000,
|
|
messages=[
|
|
{"role": "system", "content": "You are a calculator"},
|
|
{"role": "user", "content": f"{row['id']} ** 3 = ?"},
|
|
],
|
|
sampling_params=dict(
|
|
temperature=0.3,
|
|
# we don't care about the actual output
|
|
max_tokens=1,
|
|
detokenize=False,
|
|
),
|
|
),
|
|
postprocess=lambda row: {
|
|
"resp": row["generated_text"],
|
|
"large_memory_still_there": "large_memory_to_carry_over" in row,
|
|
},
|
|
)
|
|
|
|
ds = ray.data.range(12000)
|
|
|
|
def map_id_to_val_in_test_no_memory_leak(x):
|
|
return {"id": x["id"], "val": x["id"] + 5}
|
|
|
|
ds = ds.map(map_id_to_val_in_test_no_memory_leak)
|
|
ds = processor(ds)
|
|
ds = ds.materialize()
|
|
|
|
outs = ds.take_all()
|
|
assert all(out["large_memory_still_there"] for out in outs)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"backend, placement_group_config",
|
|
[
|
|
# Custom placement group with STRICT_PACK strategy
|
|
(
|
|
"ray",
|
|
dict(bundles=[{"CPU": 1, "GPU": 1}] * 4, strategy="STRICT_PACK"),
|
|
),
|
|
# Strategy omitted (PACK default). Omitted GPU now validates as 0.0; explicit GPU: 1 for this GPU job.
|
|
(
|
|
"ray",
|
|
dict(bundles=[{"CPU": 1, "GPU": 1}] * 4),
|
|
),
|
|
# Empty placement group
|
|
(
|
|
"ray",
|
|
None,
|
|
),
|
|
# Custom placement group with MP backend
|
|
(
|
|
"mp",
|
|
dict(bundles=[{"GPU": 1}] * 4),
|
|
),
|
|
# Empty placement group with MP backend
|
|
(
|
|
"mp",
|
|
None,
|
|
),
|
|
],
|
|
)
|
|
def test_vllm_placement_group(backend, placement_group_config):
|
|
"""Test vLLM with different placement group configurations."""
|
|
|
|
config = vLLMEngineProcessorConfig(
|
|
model_source="facebook/opt-1.3b",
|
|
engine_kwargs=dict(
|
|
enable_prefix_caching=True,
|
|
enable_chunked_prefill=True,
|
|
max_num_batched_tokens=4096,
|
|
pipeline_parallel_size=2,
|
|
tensor_parallel_size=2,
|
|
distributed_executor_backend=backend,
|
|
),
|
|
tokenize_stage=False,
|
|
detokenize_stage=False,
|
|
concurrency=1,
|
|
batch_size=16,
|
|
chat_template_stage=False,
|
|
placement_group_config=placement_group_config,
|
|
)
|
|
|
|
processor = build_processor(
|
|
config,
|
|
preprocess=lambda row: dict(
|
|
prompt=f"You are a calculator. {row['id']} ** 3 = ?",
|
|
sampling_params=dict(
|
|
temperature=0.3,
|
|
max_tokens=20,
|
|
detokenize=True,
|
|
),
|
|
),
|
|
postprocess=lambda row: dict(
|
|
resp=row["generated_text"],
|
|
),
|
|
)
|
|
|
|
ds = ray.data.range(60)
|
|
ds = processor(ds)
|
|
ds = ds.materialize()
|
|
|
|
outs = ds.take_all()
|
|
assert len(outs) == 60
|
|
assert all("resp" in out for out in outs)
|
|
|
|
|
|
def test_vllm_autoscaling_no_starvation():
|
|
"""Test that chained vLLMEngineProcessor instances with autoscaling
|
|
concurrency can run without starving each other.
|
|
"""
|
|
processor_config_1 = vLLMEngineProcessorConfig(
|
|
model_source="facebook/opt-1.3b",
|
|
chat_template_stage=False,
|
|
tokenize_stage=False,
|
|
detokenize_stage=False,
|
|
batch_size=16,
|
|
concurrency=(1, 4),
|
|
)
|
|
|
|
processor_config_2 = vLLMEngineProcessorConfig(
|
|
model_source="facebook/opt-1.3b",
|
|
chat_template_stage=False,
|
|
tokenize_stage=False,
|
|
detokenize_stage=False,
|
|
batch_size=16,
|
|
concurrency=(3, 4),
|
|
)
|
|
|
|
processor_1 = build_processor(
|
|
processor_config_1,
|
|
preprocess=lambda row: dict(
|
|
prompt=f"Calculate {row['id']} ** 2 = ",
|
|
sampling_params=dict(
|
|
temperature=0.3,
|
|
max_tokens=30,
|
|
detokenize=True,
|
|
),
|
|
),
|
|
postprocess=lambda row: {
|
|
"resp_1": row["generated_text"],
|
|
"id": row.get("id", None),
|
|
},
|
|
)
|
|
|
|
processor_2 = build_processor(
|
|
processor_config_2,
|
|
preprocess=lambda row: dict(
|
|
prompt=f"Previous result: {row.get('resp_1', 'N/A')}. Now calculate its cube: ",
|
|
sampling_params=dict(
|
|
temperature=0.3,
|
|
max_tokens=30,
|
|
detokenize=True,
|
|
),
|
|
),
|
|
postprocess=lambda row: {
|
|
"resp_2": row["generated_text"],
|
|
"resp_1": row.get("resp_1", None),
|
|
"id": row.get("id", None),
|
|
},
|
|
)
|
|
|
|
ds = ray.data.range(60)
|
|
ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 1})
|
|
|
|
processed_ds = processor_2(processor_1(ds))
|
|
processed_ds = processed_ds.materialize()
|
|
results = processed_ds.take_all()
|
|
|
|
assert len(results) == 60
|
|
assert all("resp_1" in out for out in results)
|
|
assert all("resp_2" in out for out in results)
|
|
assert all("id" in out for out in results)
|
|
assert all(out.get("resp_1") for out in results)
|
|
assert all(out.get("resp_2") for out in results)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|