chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,300 @@
|
||||
import sys
|
||||
from typing import Any, Dict
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray import serve
|
||||
from ray.data import ActorPoolStrategy
|
||||
from ray.data.llm import ServeDeploymentProcessorConfig, build_processor
|
||||
from ray.llm._internal.batch.processor import ProcessorBuilder
|
||||
from ray.serve.llm.openai_api_models import ChatCompletionRequest, CompletionRequest
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype_mapping", [None, {"CompletionRequest": CompletionRequest}]
|
||||
)
|
||||
def test_serve_deployment_processor(dtype_mapping):
|
||||
app_name = "test_serve_deployment_processor_app"
|
||||
deployment_name = "test_serve_deployment_name"
|
||||
|
||||
config_kwargs = dict(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
batch_size=16,
|
||||
concurrency=1,
|
||||
)
|
||||
if dtype_mapping is not None:
|
||||
config_kwargs["dtype_mapping"] = dtype_mapping
|
||||
config = ServeDeploymentProcessorConfig(**config_kwargs)
|
||||
|
||||
processor = ProcessorBuilder.build(config)
|
||||
assert processor.list_stage_names() == [
|
||||
"ServeDeploymentStage",
|
||||
]
|
||||
|
||||
stage = processor.get_stage_by_name("ServeDeploymentStage")
|
||||
assert stage.fn_constructor_kwargs == {
|
||||
"deployment_name": deployment_name,
|
||||
"app_name": app_name,
|
||||
"dtype_mapping": dtype_mapping,
|
||||
"should_continue_on_error": False,
|
||||
"request_timeout_s": None,
|
||||
}
|
||||
|
||||
assert "compute" in stage.map_batches_kwargs
|
||||
assert isinstance(stage.map_batches_kwargs["compute"], ActorPoolStrategy)
|
||||
assert stage.map_batches_kwargs["compute"].min_size == 1
|
||||
assert stage.map_batches_kwargs["compute"].max_size == 1
|
||||
|
||||
|
||||
def test_simple_serve_deployment(serve_cleanup):
|
||||
@serve.deployment
|
||||
class SimpleServeDeployment:
|
||||
# ServeDeploymentStageUDF expects an async generator.
|
||||
async def add(self, request: Dict[str, Any]):
|
||||
yield {"result": request["x"] + 1}
|
||||
|
||||
app_name = "simple_serve_deployment_app"
|
||||
deployment_name = "SimpleServeDeployment"
|
||||
|
||||
serve.run(SimpleServeDeployment.bind(), name=app_name)
|
||||
|
||||
config = ServeDeploymentProcessorConfig(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
batch_size=16,
|
||||
concurrency=1,
|
||||
)
|
||||
|
||||
processor = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="add",
|
||||
dtype=None, # Empty dtype since output is already dict format
|
||||
request_kwargs=dict(x=row["id"]),
|
||||
),
|
||||
postprocess=lambda row: dict(
|
||||
resp=row["result"],
|
||||
id=row["id"],
|
||||
),
|
||||
)
|
||||
|
||||
ds = ray.data.range(60)
|
||||
ds = ds.map(lambda x: {"id": x["id"]})
|
||||
ds = processor(ds)
|
||||
|
||||
outs = ds.take_all()
|
||||
assert len(outs) == 60
|
||||
assert all("resp" in out for out in outs)
|
||||
assert all(out["resp"] == out["id"] + 1 for out in outs)
|
||||
|
||||
|
||||
def test_serve_deployment_continue_on_error(serve_cleanup):
|
||||
@serve.deployment
|
||||
class FailingServeDeployment:
|
||||
async def process(self, request: Dict[str, Any]):
|
||||
x = request["x"]
|
||||
if x % 10 == 0: # Fail every 10th row
|
||||
raise ValueError(f"Intentional failure for x={x}")
|
||||
yield {"result": x * 2}
|
||||
|
||||
app_name = "failing_serve_deployment_app"
|
||||
deployment_name = "FailingServeDeployment"
|
||||
|
||||
serve.run(FailingServeDeployment.bind(), name=app_name)
|
||||
|
||||
config = ServeDeploymentProcessorConfig(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
batch_size=16,
|
||||
concurrency=1,
|
||||
should_continue_on_error=True,
|
||||
)
|
||||
|
||||
processor = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="process",
|
||||
dtype=None,
|
||||
request_kwargs=dict(x=row["id"]),
|
||||
),
|
||||
# Error rows will bypass this postprocess and return raw data with
|
||||
# __inference_error__ set. Only success rows get resp/id keys.
|
||||
postprocess=lambda row: dict(
|
||||
resp=row.get("result"),
|
||||
id=row.get("id"),
|
||||
),
|
||||
)
|
||||
|
||||
ds = ray.data.range(60)
|
||||
ds = ds.map(lambda x: {"id": x["id"]})
|
||||
ds = processor(ds)
|
||||
|
||||
outs = ds.take_all()
|
||||
assert len(outs) == 60
|
||||
|
||||
# Check __inference_error__ directly
|
||||
errors = [o for o in outs if o.get("__inference_error__", "")]
|
||||
successes = [o for o in outs if not o.get("__inference_error__", "")]
|
||||
|
||||
assert len(errors) == 6, f"Expected 6 errors, got {len(errors)}: {errors[:3]}..."
|
||||
assert len(successes) == 54
|
||||
|
||||
for e in errors:
|
||||
error_msg = e["__inference_error__"]
|
||||
assert "ValueError" in error_msg, f"Expected ValueError in: {error_msg}"
|
||||
assert (
|
||||
"Intentional failure" in error_msg
|
||||
), f"Expected 'Intentional failure' in: {error_msg}"
|
||||
|
||||
for s in successes:
|
||||
assert s.get("resp") is not None, f"Missing resp in success row: {s}"
|
||||
|
||||
|
||||
def test_completion_model(model_opt_125m, create_model_opt_125m_deployment):
|
||||
deployment_name, app_name = create_model_opt_125m_deployment
|
||||
config = ServeDeploymentProcessorConfig(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
dtype_mapping={
|
||||
"CompletionRequest": CompletionRequest,
|
||||
},
|
||||
batch_size=16,
|
||||
concurrency=1,
|
||||
)
|
||||
|
||||
processor = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="completions",
|
||||
dtype="CompletionRequest",
|
||||
request_kwargs=dict(
|
||||
model=model_opt_125m,
|
||||
prompt=row["prompt"],
|
||||
stream=False,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: dict(
|
||||
resp=row["choices"][0]["text"],
|
||||
),
|
||||
)
|
||||
|
||||
ds = ray.data.range(60)
|
||||
ds = ds.map(lambda x: {"prompt": f"Hello {x['id']}"})
|
||||
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_multi_turn_completion_model(model_opt_125m, create_model_opt_125m_deployment):
|
||||
deployment_name, app_name = create_model_opt_125m_deployment
|
||||
|
||||
config1 = ServeDeploymentProcessorConfig(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
dtype_mapping={
|
||||
"CompletionRequest": CompletionRequest,
|
||||
},
|
||||
# Use lower batch size to reduce resource usage as there are multiple processors
|
||||
batch_size=4,
|
||||
concurrency=1,
|
||||
)
|
||||
|
||||
processor1 = build_processor(
|
||||
config1,
|
||||
preprocess=lambda row: dict(
|
||||
dtype="CompletionRequest",
|
||||
method="completions",
|
||||
request_kwargs=dict(
|
||||
model=model_opt_125m,
|
||||
prompt=row["prompt"],
|
||||
stream=False,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: dict(
|
||||
prompt=row["choices"][0]["text"],
|
||||
),
|
||||
)
|
||||
|
||||
config2 = ServeDeploymentProcessorConfig(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
dtype_mapping={
|
||||
"CompletionRequest": CompletionRequest,
|
||||
},
|
||||
batch_size=4,
|
||||
concurrency=1,
|
||||
)
|
||||
|
||||
processor2 = build_processor(
|
||||
config2,
|
||||
preprocess=lambda row: dict(
|
||||
dtype="CompletionRequest",
|
||||
method="completions",
|
||||
request_kwargs=dict(
|
||||
model=model_opt_125m,
|
||||
prompt=row["prompt"],
|
||||
stream=False,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: dict(
|
||||
resp=row["choices"][0]["text"],
|
||||
),
|
||||
)
|
||||
|
||||
ds = ray.data.range(60)
|
||||
ds = ds.map(lambda x: {"prompt": f"Hello {x['id']}"})
|
||||
ds = processor1(ds)
|
||||
ds = processor2(ds)
|
||||
|
||||
ds = ds.materialize()
|
||||
outs = ds.take_all()
|
||||
assert len(outs) == 60
|
||||
assert all("resp" in out for out in outs)
|
||||
|
||||
|
||||
def test_chat_model(model_opt_125m, create_model_opt_125m_deployment):
|
||||
deployment_name, app_name = create_model_opt_125m_deployment
|
||||
config = ServeDeploymentProcessorConfig(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
dtype_mapping={
|
||||
"ChatCompletionRequest": ChatCompletionRequest,
|
||||
},
|
||||
batch_size=16,
|
||||
concurrency=1,
|
||||
)
|
||||
|
||||
processor = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
dtype="ChatCompletionRequest",
|
||||
method="chat",
|
||||
request_kwargs=dict(
|
||||
model=model_opt_125m,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": f"Hello {row['id']}"},
|
||||
],
|
||||
stream=False,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: dict(
|
||||
resp=row["choices"][0]["message"]["content"],
|
||||
),
|
||||
)
|
||||
|
||||
ds = ray.data.range(60)
|
||||
ds = ds.map(lambda x: {"id": x["id"]})
|
||||
ds = processor(ds)
|
||||
ds = ds.materialize()
|
||||
outs = ds.take_all()
|
||||
assert len(outs) == 60
|
||||
assert all("resp" in out for out in outs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,126 @@
|
||||
"""This test suite does not need sglang to be installed."""
|
||||
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.llm import SGLangEngineProcessorConfig
|
||||
from ray.llm._internal.batch.constants import SGLangTaskType
|
||||
from ray.llm._internal.batch.processor import ProcessorBuilder
|
||||
from ray.llm._internal.batch.processor.sglang_engine_proc import (
|
||||
build_sglang_engine_processor,
|
||||
)
|
||||
|
||||
|
||||
def test_sglang_engine_processor(gpu_type, model_llama_3_2_216M):
|
||||
config = SGLangEngineProcessorConfig(
|
||||
model_source=model_llama_3_2_216M,
|
||||
engine_kwargs=dict(
|
||||
context_length=8192,
|
||||
tp_size=2,
|
||||
dp_size=2,
|
||||
disable_cuda_graph=True,
|
||||
dtype="half", # Older GPUs (e.g. T4) don't support bfloat16
|
||||
),
|
||||
runtime_env=dict(
|
||||
env_vars=dict(
|
||||
RANDOM_ENV_VAR="12345",
|
||||
),
|
||||
),
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=4,
|
||||
batch_size=64,
|
||||
max_concurrent_batches=4,
|
||||
max_pending_requests=111,
|
||||
chat_template_stage=True,
|
||||
tokenize_stage=True,
|
||||
detokenize_stage=True,
|
||||
)
|
||||
processor = ProcessorBuilder.build(config)
|
||||
assert processor.list_stage_names() == [
|
||||
"ChatTemplateStage",
|
||||
"TokenizeStage",
|
||||
"SGLangEngineStage",
|
||||
"DetokenizeStage",
|
||||
]
|
||||
|
||||
stage = processor.get_stage_by_name("SGLangEngineStage")
|
||||
assert stage.fn_constructor_kwargs == {
|
||||
"model": model_llama_3_2_216M,
|
||||
"engine_kwargs": {
|
||||
"context_length": 8192,
|
||||
"tp_size": 2,
|
||||
"dp_size": 2,
|
||||
"disable_cuda_graph": True,
|
||||
"dtype": "half",
|
||||
"task": SGLangTaskType.GENERATE,
|
||||
},
|
||||
"task_type": SGLangTaskType.GENERATE,
|
||||
"max_pending_requests": 111,
|
||||
}
|
||||
|
||||
runtime_env = stage.map_batches_kwargs.pop("runtime_env")
|
||||
assert "env_vars" in runtime_env
|
||||
assert runtime_env["env_vars"]["RANDOM_ENV_VAR"] == "12345"
|
||||
compute = stage.map_batches_kwargs.pop("compute")
|
||||
assert isinstance(compute, ray.data._internal.compute.ActorPoolStrategy)
|
||||
assert stage.map_batches_kwargs == {
|
||||
"zero_copy_batch": True,
|
||||
"max_concurrency": 4,
|
||||
"accelerator_type": gpu_type,
|
||||
"num_gpus": 4, # Based on tp_size=2, dp_size=2 in engine_kwargs
|
||||
}
|
||||
|
||||
|
||||
class TestSGLangEngineProcessorConfig:
|
||||
def test_build_processor_autoconfig_failure_with_trust_remote_code(self):
|
||||
config = SGLangEngineProcessorConfig(
|
||||
model_source="nonexistent-org/nonexistent-model",
|
||||
engine_kwargs={"trust_remote_code": True},
|
||||
)
|
||||
|
||||
processor = build_sglang_engine_processor(config)
|
||||
assert processor is not None
|
||||
|
||||
def test_build_processor_import_error_with_trust_remote_code(self):
|
||||
config = SGLangEngineProcessorConfig(
|
||||
model_source="org/model-with-custom-code",
|
||||
engine_kwargs={"trust_remote_code": True},
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"ray.llm._internal.batch.processor.sglang_engine_proc."
|
||||
"download_model_files",
|
||||
return_value="/tmp/fake_model_dir",
|
||||
),
|
||||
patch(
|
||||
"ray.llm._internal.batch.processor.sglang_engine_proc."
|
||||
"transformers.AutoConfig.from_pretrained",
|
||||
side_effect=ModuleNotFoundError("custom modeling module missing"),
|
||||
),
|
||||
):
|
||||
processor = build_sglang_engine_processor(config)
|
||||
|
||||
assert processor is not None
|
||||
|
||||
def test_build_processor_download_error_with_trust_remote_code(self):
|
||||
config = SGLangEngineProcessorConfig(
|
||||
model_source="org/model-with-custom-code",
|
||||
engine_kwargs={"trust_remote_code": True},
|
||||
)
|
||||
|
||||
with patch(
|
||||
"ray.llm._internal.batch.processor.sglang_engine_proc."
|
||||
"download_model_files",
|
||||
side_effect=RuntimeError("download failed"),
|
||||
):
|
||||
processor = build_sglang_engine_processor(config)
|
||||
|
||||
assert processor is not None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,621 @@
|
||||
import sys
|
||||
|
||||
import pydantic
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
import ray
|
||||
from ray.data.llm import build_processor, vLLMEngineProcessorConfig
|
||||
from ray.llm._internal.batch.constants import vLLMTaskType
|
||||
from ray.llm._internal.batch.processor import ProcessorBuilder
|
||||
from ray.llm._internal.batch.stages.configs import (
|
||||
ChatTemplateStageConfig,
|
||||
DetokenizeStageConfig,
|
||||
PrepareMultimodalStageConfig,
|
||||
TokenizerStageConfig,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tensor_parallel_size, expected_distributed_executor_backend",
|
||||
[(1, "uni"), (2, "ray")],
|
||||
)
|
||||
def test_vllm_engine_processor(
|
||||
gpu_type,
|
||||
model_opt_125m,
|
||||
tensor_parallel_size,
|
||||
expected_distributed_executor_backend,
|
||||
):
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model_opt_125m,
|
||||
engine_kwargs=dict(
|
||||
max_model_len=8192,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
),
|
||||
runtime_env=dict(
|
||||
env_vars=dict(
|
||||
RANDOM_ENV_VAR="12345",
|
||||
),
|
||||
),
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=4,
|
||||
batch_size=64,
|
||||
max_pending_requests=111,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=True),
|
||||
prepare_multimodal_stage=PrepareMultimodalStageConfig(enabled=True),
|
||||
)
|
||||
processor = ProcessorBuilder.build(config)
|
||||
assert processor.list_stage_names() == [
|
||||
"PrepareMultimodalStage",
|
||||
"ChatTemplateStage",
|
||||
"TokenizeStage",
|
||||
"vLLMEngineStage",
|
||||
"DetokenizeStage",
|
||||
]
|
||||
|
||||
stage = processor.get_stage_by_name("vLLMEngineStage")
|
||||
assert stage.fn_constructor_kwargs == {
|
||||
"model": model_opt_125m,
|
||||
"engine_kwargs": {
|
||||
"max_model_len": 8192,
|
||||
"distributed_executor_backend": expected_distributed_executor_backend,
|
||||
"tensor_parallel_size": tensor_parallel_size,
|
||||
"task_type": vLLMTaskType.GENERATE,
|
||||
},
|
||||
"task_type": vLLMTaskType.GENERATE,
|
||||
"max_pending_requests": 111,
|
||||
"dynamic_lora_loading_path": None,
|
||||
"max_concurrent_batches": 8,
|
||||
"batch_size": 64,
|
||||
"should_continue_on_error": False,
|
||||
"log_engine_metrics": True,
|
||||
}
|
||||
|
||||
runtime_env = stage.map_batches_kwargs.pop("runtime_env")
|
||||
assert "env_vars" in runtime_env
|
||||
assert runtime_env["env_vars"]["RANDOM_ENV_VAR"] == "12345"
|
||||
compute = stage.map_batches_kwargs.pop("compute")
|
||||
assert isinstance(compute, ray.data._internal.compute.ActorPoolStrategy)
|
||||
|
||||
if expected_distributed_executor_backend == "ray":
|
||||
ray_remote_args_fn = stage.map_batches_kwargs.pop("ray_remote_args_fn")
|
||||
assert ray_remote_args_fn is not None
|
||||
assert stage.map_batches_kwargs == {
|
||||
"zero_copy_batch": True,
|
||||
"max_concurrency": 8,
|
||||
"accelerator_type": gpu_type,
|
||||
"num_gpus": 0,
|
||||
}
|
||||
else:
|
||||
assert "ray_remote_args_fn" not in stage.map_batches_kwargs
|
||||
assert stage.map_batches_kwargs == {
|
||||
"zero_copy_batch": True,
|
||||
"max_concurrency": 8,
|
||||
"accelerator_type": gpu_type,
|
||||
"num_gpus": 1,
|
||||
}
|
||||
|
||||
|
||||
def test_vllm_engine_processor_task_override(model_opt_125m):
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model_opt_125m,
|
||||
engine_kwargs=dict(
|
||||
task_type=vLLMTaskType.EMBED,
|
||||
),
|
||||
task_type=vLLMTaskType.GENERATE,
|
||||
concurrency=4,
|
||||
batch_size=64,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=True),
|
||||
prepare_multimodal_stage=PrepareMultimodalStageConfig(enabled=True),
|
||||
)
|
||||
processor = ProcessorBuilder.build(config)
|
||||
stage = processor.get_stage_by_name("vLLMEngineStage")
|
||||
|
||||
assert stage.fn_constructor_kwargs["task_type"] == vLLMTaskType.GENERATE
|
||||
assert (
|
||||
stage.fn_constructor_kwargs["engine_kwargs"]["task_type"]
|
||||
== vLLMTaskType.GENERATE
|
||||
)
|
||||
|
||||
|
||||
def test_vllm_engine_processor_invalid_task(model_opt_125m):
|
||||
with pytest.raises(
|
||||
pydantic.ValidationError, match="Invalid task type: invalid_task"
|
||||
):
|
||||
vLLMEngineProcessorConfig(
|
||||
model_source=model_opt_125m,
|
||||
engine_kwargs=dict(
|
||||
task_type=vLLMTaskType.EMBED,
|
||||
),
|
||||
task_type="invalid_task",
|
||||
concurrency=4,
|
||||
batch_size=64,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=True),
|
||||
prepare_multimodal_stage=PrepareMultimodalStageConfig(enabled=True),
|
||||
)
|
||||
|
||||
|
||||
def test_vllm_engine_processor_placement_group(gpu_type, model_opt_125m):
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model_opt_125m,
|
||||
engine_kwargs=dict(
|
||||
max_model_len=8192,
|
||||
),
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=4,
|
||||
batch_size=64,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
placement_group_config=dict(bundles=[{"CPU": 1, "GPU": 1}]),
|
||||
)
|
||||
processor = ProcessorBuilder.build(config)
|
||||
stage = processor.get_stage_by_name("vLLMEngineStage")
|
||||
|
||||
stage.map_batches_kwargs.pop("runtime_env")
|
||||
stage.map_batches_kwargs.pop("compute")
|
||||
|
||||
assert stage.map_batches_kwargs == {
|
||||
"zero_copy_batch": True,
|
||||
"max_concurrency": 8,
|
||||
"accelerator_type": gpu_type,
|
||||
"num_cpus": 1,
|
||||
"num_gpus": 1,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"engine_kwargs_extra,expected_num_bundles",
|
||||
[
|
||||
({"tensor_parallel_size": 2}, 2),
|
||||
(
|
||||
{"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
|
||||
4,
|
||||
),
|
||||
({}, 1), # Default case: tp=1, pp=1 → executor_backend="uni"
|
||||
],
|
||||
)
|
||||
def test_vllm_engine_processor_bundle_per_worker(
|
||||
gpu_type, model_opt_125m, engine_kwargs_extra, expected_num_bundles
|
||||
):
|
||||
"""Test bundle_per_worker auto-expands based on tp*pp."""
|
||||
engine_kwargs = dict(max_model_len=8192)
|
||||
engine_kwargs.update(engine_kwargs_extra)
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model_opt_125m,
|
||||
engine_kwargs=engine_kwargs,
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=4,
|
||||
batch_size=64,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
placement_group_config={"bundle_per_worker": {"CPU": 1, "GPU": 1}},
|
||||
)
|
||||
processor = ProcessorBuilder.build(config)
|
||||
stage = processor.get_stage_by_name("vLLMEngineStage")
|
||||
|
||||
stage.map_batches_kwargs.pop("runtime_env")
|
||||
stage.map_batches_kwargs.pop("compute")
|
||||
|
||||
expected_kwargs = {
|
||||
"zero_copy_batch": True,
|
||||
"max_concurrency": 8,
|
||||
"accelerator_type": gpu_type,
|
||||
}
|
||||
|
||||
if expected_num_bundles > 1:
|
||||
# TP/PP > 1 -> executor_backend="ray"
|
||||
ray_remote_args_fn = stage.map_batches_kwargs.pop("ray_remote_args_fn")
|
||||
assert ray_remote_args_fn.args[0] == expected_num_bundles
|
||||
assert ray_remote_args_fn.args[1] == gpu_type
|
||||
expected_bundles = [{"CPU": 1, "GPU": 1}] * expected_num_bundles
|
||||
assert ray_remote_args_fn.args[2]["bundles"] == expected_bundles
|
||||
assert ray_remote_args_fn.args[2]["strategy"] == "PACK"
|
||||
expected_kwargs["num_gpus"] = 0
|
||||
else:
|
||||
# TP=1, PP=1 -> executor_backend="uni"
|
||||
expected_kwargs["num_cpus"] = 1
|
||||
expected_kwargs["num_gpus"] = 1
|
||||
|
||||
assert stage.map_batches_kwargs == expected_kwargs
|
||||
|
||||
|
||||
def test_vllm_engine_processor_bundle_per_worker_conflict(gpu_type, model_opt_125m):
|
||||
"""Test that specifying both bundle_per_worker and bundles raises error."""
|
||||
with pytest.raises(ValueError, match="Cannot specify both"):
|
||||
vLLMEngineProcessorConfig(
|
||||
model_source=model_opt_125m,
|
||||
engine_kwargs=dict(max_model_len=8192),
|
||||
accelerator_type=gpu_type,
|
||||
placement_group_config={
|
||||
"bundle_per_worker": {"CPU": 1, "GPU": 1},
|
||||
"bundles": [{"CPU": 1, "GPU": 1}],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_multimodal_stage_vllm_engine_processor(gpu_type, model_smolvlm_256m):
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model_smolvlm_256m,
|
||||
engine_kwargs=dict(
|
||||
max_model_len=8192,
|
||||
),
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=1,
|
||||
batch_size=16,
|
||||
prepare_multimodal_stage=PrepareMultimodalStageConfig(
|
||||
enabled=True,
|
||||
model_config_kwargs=dict(
|
||||
allowed_local_media_path="/tmp",
|
||||
),
|
||||
),
|
||||
)
|
||||
processor = ProcessorBuilder.build(config)
|
||||
|
||||
assert "PrepareMultimodalStage" in processor.list_stage_names()
|
||||
stage = processor.get_stage_by_name("PrepareMultimodalStage")
|
||||
fn_kwargs = stage.fn_constructor_kwargs
|
||||
|
||||
assert "model_config_kwargs" in fn_kwargs
|
||||
model_config_kwargs = fn_kwargs["model_config_kwargs"]
|
||||
assert model_config_kwargs["allowed_local_media_path"] == "/tmp"
|
||||
assert model_config_kwargs["model"] == model_smolvlm_256m
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["uni", "mp", "ray"])
|
||||
def test_generation_model(gpu_type, model_opt_125m, backend):
|
||||
# OPT models don't have chat template, so we use ChatML template
|
||||
# here to demonstrate the usage of custom chat template.
|
||||
chat_template = """
|
||||
{% if messages[0]['role'] == 'system' %}
|
||||
{% set offset = 1 %}
|
||||
{% else %}
|
||||
{% set offset = 0 %}
|
||||
{% endif %}
|
||||
|
||||
{{ bos_token }}
|
||||
{% for message in messages %}
|
||||
{% if (message['role'] == 'user') != (loop.index0 % 2 == offset) %}
|
||||
{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
|
||||
{% endif %}
|
||||
|
||||
{{ '<|im_start|>' + message['role'] + '\n' + message['content'] | trim + '<|im_end|>\n' }}
|
||||
{% endfor %}
|
||||
|
||||
{% if add_generation_prompt %}
|
||||
{{ '<|im_start|>assistant\n' }}
|
||||
{% endif %}
|
||||
"""
|
||||
|
||||
processor_config = vLLMEngineProcessorConfig(
|
||||
model_source=model_opt_125m,
|
||||
engine_kwargs=dict(
|
||||
enable_prefix_caching=False,
|
||||
enable_chunked_prefill=True,
|
||||
max_num_batched_tokens=2048,
|
||||
max_model_len=2048,
|
||||
# Skip CUDA graph capturing to reduce startup time.
|
||||
enforce_eager=True,
|
||||
distributed_executor_backend=backend,
|
||||
),
|
||||
batch_size=16,
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=1,
|
||||
chat_template_stage=ChatTemplateStageConfig(
|
||||
enabled=True, chat_template=chat_template
|
||||
),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=True),
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def test_generation_model_tokenized_prompt(gpu_type, model_opt_125m):
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_opt_125m, trust_remote_code=True)
|
||||
|
||||
processor_config = vLLMEngineProcessorConfig(
|
||||
model_source=model_opt_125m,
|
||||
engine_kwargs=dict(
|
||||
enable_prefix_caching=False,
|
||||
enable_chunked_prefill=True,
|
||||
max_num_batched_tokens=2048,
|
||||
max_model_len=2048,
|
||||
enforce_eager=True,
|
||||
),
|
||||
batch_size=16,
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=1,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=False),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=False),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=False),
|
||||
)
|
||||
|
||||
def preprocess(row):
|
||||
prompt_text = f"Calculate {row['id']} ** 3"
|
||||
|
||||
return dict(
|
||||
tokenized_prompt=tokenizer(prompt_text)["input_ids"],
|
||||
sampling_params=dict(
|
||||
temperature=0.3,
|
||||
max_tokens=50,
|
||||
),
|
||||
)
|
||||
|
||||
processor = build_processor(
|
||||
processor_config,
|
||||
preprocess=preprocess,
|
||||
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)
|
||||
|
||||
|
||||
def test_embedding_model(gpu_type, model_smolvlm_256m):
|
||||
processor_config = vLLMEngineProcessorConfig(
|
||||
model_source=model_smolvlm_256m,
|
||||
task_type="embed",
|
||||
engine_kwargs=dict(
|
||||
enable_prefix_caching=False,
|
||||
enable_chunked_prefill=False,
|
||||
max_model_len=2048,
|
||||
# Skip CUDA graph capturing to reduce startup time.
|
||||
enforce_eager=True,
|
||||
),
|
||||
batch_size=16,
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=1,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=False),
|
||||
)
|
||||
|
||||
processor = build_processor(
|
||||
processor_config,
|
||||
preprocess=lambda row: dict(
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a calculator"},
|
||||
{"role": "user", "content": f"{row['id']} ** 3 = ?"},
|
||||
],
|
||||
),
|
||||
postprocess=lambda row: {
|
||||
"resp": row["embeddings"],
|
||||
"prompt": row["prompt"],
|
||||
},
|
||||
)
|
||||
|
||||
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)
|
||||
assert all("prompt" in out for out in outs)
|
||||
|
||||
|
||||
def test_classification_model(gpu_type):
|
||||
processor_config = vLLMEngineProcessorConfig(
|
||||
model_source="HuggingFaceTB/fineweb-edu-classifier",
|
||||
task_type="classify",
|
||||
engine_kwargs=dict(
|
||||
max_model_len=512, # Model only supports up to 512 tokens
|
||||
),
|
||||
batch_size=16,
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=1,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=False),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=False),
|
||||
)
|
||||
|
||||
processor = build_processor(
|
||||
processor_config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt="This is a great educational content.",
|
||||
tokenization_kwargs={"truncation": True, "max_length": 512},
|
||||
),
|
||||
postprocess=lambda row: {
|
||||
"probs": float(row["embeddings"][0])
|
||||
if row.get("embeddings") is not None and len(row["embeddings"]) > 0
|
||||
else None,
|
||||
},
|
||||
)
|
||||
|
||||
ds = ray.data.range(60)
|
||||
ds = processor(ds)
|
||||
ds = ds.materialize()
|
||||
outs = ds.take_all()
|
||||
assert len(outs) == 60
|
||||
assert all("probs" in out for out in outs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("input_raw_image_data", [True, False])
|
||||
@pytest.mark.parametrize("decouple_tokenizer", [True, False])
|
||||
def test_vision_model(
|
||||
gpu_type, model_smolvlm_256m, image_asset, input_raw_image_data, decouple_tokenizer
|
||||
):
|
||||
image_url, image_pil = image_asset
|
||||
llm_processor_config = vLLMEngineProcessorConfig(
|
||||
model_source=model_smolvlm_256m,
|
||||
task_type="generate",
|
||||
engine_kwargs=dict(
|
||||
# Skip CUDA graph capturing to reduce startup time.
|
||||
enforce_eager=True,
|
||||
# CI uses T4 GPU which does not support bfloat16.
|
||||
dtype="half",
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
),
|
||||
batch_size=16,
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=1,
|
||||
prepare_multimodal_stage=PrepareMultimodalStageConfig(
|
||||
enabled=True,
|
||||
chat_template_content_format="openai",
|
||||
),
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=decouple_tokenizer),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=decouple_tokenizer),
|
||||
)
|
||||
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": image_url},
|
||||
"uuid": "image-1-id", # UUID will not be included in the output as it's only used for internal caching
|
||||
}
|
||||
if input_raw_image_data
|
||||
else {
|
||||
"type": "image_pil",
|
||||
"image_pil": image_pil,
|
||||
"uuid": "image-2-id",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
sampling_params=dict(
|
||||
temperature=0.3,
|
||||
max_tokens=50,
|
||||
detokenize=not decouple_tokenizer,
|
||||
),
|
||||
),
|
||||
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("input_raw_audio_data", [True, False])
|
||||
def test_audio_model(
|
||||
gpu_type, model_qwen_2_5_omni_3b, audio_asset, input_raw_audio_data
|
||||
):
|
||||
audio_url, audio_data = audio_asset
|
||||
llm_processor_config = vLLMEngineProcessorConfig(
|
||||
model_source=model_qwen_2_5_omni_3b,
|
||||
task_type="generate",
|
||||
engine_kwargs=dict(
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
),
|
||||
batch_size=16,
|
||||
accelerator_type=gpu_type,
|
||||
concurrency=1,
|
||||
prepare_multimodal_stage=PrepareMultimodalStageConfig(
|
||||
enabled=True,
|
||||
chat_template_content_format="openai",
|
||||
),
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=True),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=False),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=False),
|
||||
)
|
||||
|
||||
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"Describe this audio in {row['val']} words.",
|
||||
},
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {"data": audio_data, "format": "wav"},
|
||||
}
|
||||
if input_raw_audio_data
|
||||
else {
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": audio_url},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
sampling_params=dict(
|
||||
temperature=0.3,
|
||||
max_tokens=50,
|
||||
),
|
||||
),
|
||||
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)
|
||||
|
||||
|
||||
class TestVLLMEngineProcessorConfig:
|
||||
def test_build_processor_autoconfig_failure(self):
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source="nonexistent-org/nonexistent-model",
|
||||
)
|
||||
|
||||
processor = build_processor(config)
|
||||
assert processor is not None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user