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__]))