wehub-resource-sync
61 lines
1.5 KiB
Python
61 lines
1.5 KiB
Python
import pytest
|
|
|
|
import ray
|
|
from ray.data.llm import build_processor, vLLMEngineProcessorConfig
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def cleanup_ray_resources():
|
|
"""Automatically cleanup Ray resources between tests to prevent conflicts."""
|
|
yield
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"tp_size,pp_size",
|
|
[
|
|
(2, 4),
|
|
(4, 2),
|
|
],
|
|
)
|
|
def test_vllm_multi_node(tp_size, pp_size):
|
|
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=pp_size,
|
|
tensor_parallel_size=tp_size,
|
|
distributed_executor_backend="ray",
|
|
),
|
|
tokenize_stage=False,
|
|
detokenize_stage=False,
|
|
concurrency=1,
|
|
batch_size=64,
|
|
chat_template_stage=False,
|
|
)
|
|
|
|
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)
|