chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,560 @@
|
||||
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__]))
|
||||
Reference in New Issue
Block a user