Files
2026-07-13 13:17:40 +08:00

561 lines
23 KiB
Python

import sys
from pathlib import Path
from unittest.mock import MagicMock, patch
import pydantic
import pytest
from ray.llm._internal.common.utils.download_utils import NodeModelDownloadable
from ray.llm._internal.serve.core.configs.accelerators import (
CPUAccelerator,
CPUConfig,
GPUAccelerator,
GPUConfig,
TPUAccelerator,
TPUConfig,
)
from ray.llm._internal.serve.core.configs.llm_config import (
LLMConfig,
LoraConfig,
ModelLoadingConfig,
)
from ray.llm._internal.serve.engines.vllm.vllm_models import VLLMEngineConfig
CONFIG_DIRS_PATH = str(Path(__file__).parent / "configs")
class TestModelConfig:
def test_construction(self):
"""Test construct an LLMConfig doesn't error out and has correct attributes."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(
model_id="llm_model_id",
),
accelerator_type="A100-40G", # Dash instead of underscore when specifying accelerator type
deployment_config={
"autoscaling_config": {
"min_replicas": 3,
"max_replicas": 7,
}
},
)
assert llm_config.deployment_config["autoscaling_config"]["min_replicas"] == 3
assert llm_config.deployment_config["autoscaling_config"]["max_replicas"] == 7
assert llm_config.model_loading_config.model_id == "llm_model_id"
assert llm_config.accelerator_type == "A100-40G"
def test_construction_requires_model_loading_config(self):
"""Test that constructing an LLMConfig without model_loading_config errors out"""
with pytest.raises(
pydantic.ValidationError,
):
LLMConfig(
accelerator_type="L4",
)
def test_accelerator_type_optional(self):
"""Test that accelerator_type is optional when initializing LLMConfig."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model")
)
assert llm_config.model_loading_config.model_id == "test_model"
assert llm_config.accelerator_type is None
def test_invalid_accelerator_type(self):
"""Test that invalid accelerator types raise validation errors."""
with pytest.raises(pydantic.ValidationError):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="INVALID_GPU", # Invalid string value
)
# Test invalid numeric value
with pytest.raises(pydantic.ValidationError):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type=123, # Must be a string
)
# Test that underscore is not supported in accelerator type
with pytest.raises(pydantic.ValidationError):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="A100_40G", # Should use A100-40G instead
)
def test_model_loading_config_forbids_extra_fields(self):
"""Test that ModelLoadingConfig rejects extra fields."""
with pytest.raises(pydantic.ValidationError, match="engine_kwargs"):
ModelLoadingConfig(
model_id="test_model",
model_source="test_source",
engine_kwargs={"max_model_len": 8000}, # This should be rejected
)
valid_config = ModelLoadingConfig(
model_id="test_model", model_source="test_source"
)
assert valid_config.model_id == "test_model"
assert valid_config.model_source == "test_source"
def test_invalid_generation_config(self, disable_placement_bundles):
"""Test that passing an invalid generation_config raises an error."""
with pytest.raises(
pydantic.ValidationError,
):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="L4",
generation_config="invalid_config", # Should be a dictionary, not a string
)
def test_deployment_type_checking(self, disable_placement_bundles):
"""Test that deployment config type checking works."""
with pytest.raises(
pydantic.ValidationError,
):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
deployment_config={
"max_ongoing_requests": -1,
},
accelerator_type="L4",
)
def test_autoscaling_type_checking(self, disable_placement_bundles):
"""Test that autoscaling config type checking works."""
with pytest.raises(
pydantic.ValidationError,
):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
deployment_config={
"autoscaling_config": {
"min_replicas": -1,
},
},
accelerator_type="L4",
)
def test_deployment_unset_fields_are_not_included(self, disable_placement_bundles):
"""Test that unset fields are not included in the deployment config."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="L4",
)
assert "max_ongoing_requests" not in llm_config.deployment_config
assert "graceful_shutdown_timeout_s" not in llm_config.deployment_config
def test_autoscaling_unset_fields_are_not_included(self, disable_placement_bundles):
"""Test that unset fields are not included in the autoscaling config."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
deployment_config={
"autoscaling_config": {
"min_replicas": 3,
"max_replicas": 7,
},
},
accelerator_type="L4",
)
assert (
"metrics_interval_s"
not in llm_config.deployment_config["autoscaling_config"]
)
assert (
"upscaling_factor" not in llm_config.deployment_config["autoscaling_config"]
)
def test_engine_config_cached(self):
"""Test that the engine config is cached and not recreated when calling
get_engine_config so the attributes on the engine will be persisted."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(
model_id="llm_model_id",
),
)
old_engine_config = llm_config.get_engine_config()
old_engine_config.hf_model_id = "fake_hf_model_id"
new_engine_config = llm_config.get_engine_config()
assert new_engine_config is old_engine_config
def test_remote_model_source_uses_model_id_as_hf_model_id(self):
"""A remote model_source must not leak its URI into hf_model_id.
Using the URI verbatim propagates the scheme and slashes into the HF
cache directory name (e.g. ``models--s3:----bucket--...``). The URI
should instead be carried by mirror_config while hf_model_id falls back
to the user-supplied model_id.
"""
bucket_uri = "s3://my-bucket/my-model"
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(
model_id="llm_model_id",
model_source=bucket_uri,
),
)
engine_config = llm_config.get_engine_config()
assert engine_config.hf_model_id == "llm_model_id"
assert engine_config.mirror_config is not None
assert engine_config.mirror_config.bucket_uri == bucket_uri
def test_hf_model_source_used_as_hf_model_id(self):
"""A plain HuggingFace model_source is used directly as hf_model_id."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(
model_id="llm_model_id",
model_source="facebook/opt-1.3b",
),
)
engine_config = llm_config.get_engine_config()
assert engine_config.hf_model_id == "facebook/opt-1.3b"
assert engine_config.mirror_config is None
def test_no_model_source_falls_back_to_model_id(self):
"""With no model_source, hf_model_id falls back to model_id."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(
model_id="llm_model_id",
),
)
engine_config = llm_config.get_engine_config()
assert engine_config.hf_model_id == "llm_model_id"
assert engine_config.mirror_config is None
def test_experimental_configs(self):
"""Test that `experimental_configs` can be used."""
# Test with a valid dictionary can be used.
experimental_configs = {
"experimental_feature1": "value1",
"experimental_feature2": "value2",
}
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(
model_id="llm_model_id",
),
experimental_configs=experimental_configs,
)
assert llm_config.experimental_configs == experimental_configs
# test with invalid dictionary will raise a validation error.
with pytest.raises(
pydantic.ValidationError,
):
LLMConfig(
model_loading_config=ModelLoadingConfig(
model_id="llm_model_id",
),
experimental_configs={123: "value1"},
)
def test_log_engine_metrics_disable_log_stats_validation(self):
"""Test that log_engine_metrics=True prevents disable_log_stats=True."""
with pytest.raises(
pydantic.ValidationError,
match="disable_log_stats cannot be set to True when log_engine_metrics is enabled",
):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
log_engine_metrics=True,
engine_kwargs={"disable_log_stats": True},
)
@pytest.mark.parametrize(
"load_format,expected_download_model",
[
("runai_streamer", NodeModelDownloadable.NONE),
("runai_streamer_sharded", NodeModelDownloadable.NONE),
("tensorizer", NodeModelDownloadable.NONE),
(None, NodeModelDownloadable.MODEL_AND_TOKENIZER),
],
)
def test_load_format_callback_context(self, load_format, expected_download_model):
"""Test that different load_format values set correct worker_node_download_model in callback context."""
engine_kwargs = {"load_format": load_format} if load_format is not None else {}
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
engine_kwargs=engine_kwargs,
)
# Get the callback instance which should trigger the context setup
callback = llm_config.get_or_create_callback()
# Check that the callback context has the correct worker_node_download_model value
assert hasattr(callback, "ctx"), "Callback should have ctx attribute"
assert callback.ctx.worker_node_download_model == expected_download_model
class TestFieldValidators:
"""Test the field validators for dict validation."""
def test_model_loading_config_dict_validation(self):
"""Test that model_loading_config accepts and validates dict input."""
config_dict = {"model_id": "microsoft/DialoGPT-medium"}
llm_config = LLMConfig(model_loading_config=config_dict, llm_engine="vLLM")
assert isinstance(llm_config.model_loading_config, ModelLoadingConfig)
assert llm_config.model_loading_config.model_id == "microsoft/DialoGPT-medium"
def test_model_loading_config_validation_error(self):
"""Test that invalid dict raises proper validation error."""
with pytest.raises(pydantic.ValidationError) as exc_info:
LLMConfig(
model_loading_config={"invalid_field": "value"}, llm_engine="vLLM"
)
assert "Invalid model_loading_config" in str(exc_info.value)
def test_lora_config_dict_validation(self):
"""Test that lora_config accepts and validates dict input."""
llm_config = LLMConfig(
model_loading_config={"model_id": "test"},
lora_config=None,
llm_engine="vLLM",
)
assert llm_config.lora_config is None
lora_dict = {
"dynamic_lora_loading_path": "s3://bucket/lora",
"max_num_adapters_per_replica": 8,
}
llm_config2 = LLMConfig(
model_loading_config={"model_id": "test"},
lora_config=lora_dict,
llm_engine="vLLM",
)
assert isinstance(llm_config2.lora_config, LoraConfig)
assert llm_config2.lora_config.max_num_adapters_per_replica == 8
assert llm_config2.lora_config.dynamic_lora_loading_path == "s3://bucket/lora"
def test_lora_config_validation_error(self):
"""Test that invalid lora config dict raises proper validation error."""
with pytest.raises(pydantic.ValidationError) as exc_info:
LLMConfig(
model_loading_config={"model_id": "test"},
lora_config={"max_num_adapters_per_replica": "invalid_string"},
llm_engine="vLLM",
)
assert "Invalid lora_config" in str(exc_info.value)
class TestAcceleratorConfigLogic:
"""Test the accelerator_config logic and its interaction with accelerator_type."""
def test_accelerator_config_field_basic(self):
"""Test that accelerator_config field works with basic values."""
# Test CPU config
llm_config_cpu = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_config={"kind": "cpu"},
)
assert llm_config_cpu.accelerator_config.kind == "cpu"
engine_config = llm_config_cpu.get_engine_config()
assert engine_config.accelerator_config.kind == "cpu"
# Test GPU config
llm_config_gpu = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_config={"kind": "gpu"},
)
assert llm_config_gpu.accelerator_config.kind == "gpu"
engine_config_gpu = llm_config_gpu.get_engine_config()
assert engine_config_gpu.accelerator_config.kind == "gpu"
def test_accelerator_type_with_cpu_config_raises_error(self):
"""Test that accelerator_type with CPU config raises a validation error."""
with pytest.raises(
pydantic.ValidationError,
match="accelerator_type='L4' cannot be used with CPU-only configurations",
):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_config={"kind": "cpu"},
accelerator_type="L4",
)
def test_accelerator_type_with_cpu_only_placement_group_raises_error(self):
"""Test that accelerator_type with CPU-only placement_group_config raises error."""
with pytest.raises(
pydantic.ValidationError,
match="accelerator_type='L4' cannot be used with CPU-only configurations",
):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="L4",
placement_group_config={"bundles": [{"CPU": 4}]},
)
def test_accelerator_type_with_empty_bundles_raises_error(self):
"""Test that accelerator_type with empty bundles list raises error."""
with pytest.raises(
pydantic.ValidationError,
match="accelerator_type='L4' cannot be used with CPU-only configurations",
):
LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="L4",
placement_group_config={"bundles": []},
)
def test_accelerator_type_with_gpu_placement_group_succeeds(self):
"""Test that accelerator_type with GPU-containing placement_group_config succeeds."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="L4",
placement_group_config={"bundles": [{"GPU": 1, "CPU": 4}]},
)
assert llm_config.accelerator_type == "L4"
def test_accelerator_type_with_gpu_config_succeeds(self):
"""Test that accelerator_type with GPU config succeeds."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="L4",
accelerator_config={"kind": "gpu"},
)
assert llm_config.accelerator_type == "L4"
engine_config = llm_config.get_engine_config()
assert engine_config.accelerator_type == "L4"
def test_vllm_engine_config_accelerator_type_with_cpu_config_raises_error(self):
"""Test that VLLMEngineConfig rejects accelerator_type with CPU config."""
with pytest.raises(
pydantic.ValidationError,
match="accelerator_type='L4' cannot be used with CPU-only configurations",
):
VLLMEngineConfig(
model_id="test-model",
accelerator_type="L4",
accelerator_config=CPUConfig(kind="cpu"),
)
def test_vllm_engine_config_accelerator_type_with_gpu_config_succeeds(self):
"""Test that VLLMEngineConfig accepts accelerator_type with GPU config."""
engine_config = VLLMEngineConfig(
model_id="test-model",
accelerator_type="L4",
accelerator_config=GPUConfig(kind="gpu"),
)
assert engine_config.accelerator_type == "L4"
def test_llm_config_accelerator_type_hardware_mismatch(self):
"""Test that passing a GPU accelerator_type with a TPU config raises an error."""
with pytest.raises(
pydantic.ValidationError,
match="Hardware mismatch",
):
LLMConfig(
model_loading_config={"model_id": "test_model"},
accelerator_type="L4",
accelerator_config={"kind": "tpu", "topology": "4x4"},
)
def test_engine_config_infers_tpu_from_accelerator_type_string(self):
"""Test that the engine config infers a TPU backend directly from the accelerator_type string."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model"),
accelerator_type="TPU-V6E",
)
# Validate engine correctly inferred the TPU backend
engine_config = llm_config.get_engine_config()
assert isinstance(engine_config.accelerator, TPUAccelerator)
assert engine_config.accelerator_type == "TPU-V6E"
def test_requires_deferred_placement_group(self):
"""Test that requires_deferred_placement_group correctly identifies deferred PG requirements."""
cpu_accel = CPUAccelerator()
assert cpu_accel.requires_deferred_placement_group is False
gpu_accel = GPUAccelerator()
assert gpu_accel.requires_deferred_placement_group is False
tpu_accel_no_topo = TPUAccelerator(TPUConfig(kind="tpu"))
assert tpu_accel_no_topo.requires_deferred_placement_group is False
tpu_accel_with_topo = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
assert tpu_accel_with_topo.requires_deferred_placement_group is True
@pytest.mark.parametrize(
"topology,num_devices,accelerator_type_str,expected_bundles_count,expected_chips_per_host",
[
("1x1", 1, "TPU-V6E", 1, 1),
("1x1", 1, "TPU-V7X", 1, 1),
("4x4", 16, "TPU-V6E", 4, 4),
("2x2x2", 8, "TPU-V5P", 2, 4),
("2x2", 4, "TPU-V5LITEPOD", 1, 4),
("2x2x1", 4, "TPU-V4", 1, 4),
("2x4", 8, "TPU-V6E", 1, 8),
],
)
def test_default_bundles_topology(
self,
topology,
num_devices,
accelerator_type_str,
expected_bundles_count,
expected_chips_per_host,
):
"""Test that different topologies return correct per-host bundles."""
tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology=topology))
bundles = tpu_accel.default_bundles(
num_devices=num_devices, accelerator_type_str=accelerator_type_str
)
assert len(bundles) == expected_bundles_count
for bundle in bundles:
assert bundle["TPU"] == expected_chips_per_host
assert f"accelerator_type:{accelerator_type_str}" in bundle
def test_default_bundles_topology_missing_accelerator_type_raises(self):
"""Test that ValueError is raised when topology is present but accelerator type is missing."""
tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
with pytest.raises(
ValueError,
match="`accelerator_type` must be specified when `topology` is present",
):
tpu_accel.default_bundles(num_devices=16, accelerator_type_str=None)
def test_default_bundles_topology_non_multiple_num_devices_raises(self):
"""Test that ValueError is raised when num_devices is not a multiple of chips_per_host."""
tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
with pytest.raises(ValueError, match="must be a multiple of chips_per_host"):
tpu_accel.default_bundles(num_devices=6, accelerator_type_str="TPU-V6E")
class TestCheckpointInfo:
def test_apply_checkpoint_info_uses_autoconfig_and_threads_trust_remote_code(self):
"""apply_checkpoint_info uses AutoConfig (not PretrainedConfig) and forwards
trust_remote_code to every HF config load call."""
llm_config = LLMConfig(
model_loading_config=ModelLoadingConfig(model_id="test_model")
)
mock_hf_config = MagicMock(spec=["architectures", "vision_config"])
mock_hf_config.architectures = ["LlavaForCausalLM"]
with patch(
"transformers.AutoConfig.from_pretrained", return_value=mock_hf_config
) as mock_auto:
llm_config.apply_checkpoint_info("vision/model", trust_remote_code=True)
assert all(
call.kwargs["trust_remote_code"] is True
for call in mock_auto.call_args_list
)
assert llm_config._supports_vision is True
assert llm_config._model_architecture == "LlavaForCausalLM"
if __name__ == "__main__":
sys.exit(pytest.main(["-v", __file__]))