Convert formatting to use ruff instead of yapf + isort (#26247)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
@@ -32,7 +32,6 @@ def cleanup():
|
||||
|
||||
@pytest.fixture()
|
||||
def just_serialize_model_tensors(model_ref, monkeypatch, tmp_path):
|
||||
|
||||
def noop(*args, **kwargs):
|
||||
return None
|
||||
|
||||
@@ -56,8 +55,7 @@ def model_path(model_ref, tmp_path):
|
||||
yield tmp_path / model_ref / "model.tensors"
|
||||
|
||||
|
||||
def assert_from_collective_rpc(engine: LLM, closure: Callable,
|
||||
closure_kwargs: dict):
|
||||
def assert_from_collective_rpc(engine: LLM, closure: Callable, closure_kwargs: dict):
|
||||
res = engine.collective_rpc(method=closure, kwargs=closure_kwargs)
|
||||
return all(res)
|
||||
|
||||
@@ -67,18 +65,13 @@ def assert_from_collective_rpc(engine: LLM, closure: Callable,
|
||||
# method. It's purely used as a dummy utility to run methods that test
|
||||
# Tensorizer functionality
|
||||
class DummyExecutor(UniProcExecutor):
|
||||
|
||||
def _init_executor(self) -> None:
|
||||
"""Initialize the worker and load the model.
|
||||
"""
|
||||
self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config,
|
||||
rpc_rank=0)
|
||||
distributed_init_method = get_distributed_init_method(
|
||||
get_ip(), get_open_port())
|
||||
"""Initialize the worker and load the model."""
|
||||
self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config, rpc_rank=0)
|
||||
distributed_init_method = get_distributed_init_method(get_ip(), get_open_port())
|
||||
local_rank = 0
|
||||
# set local rank as the device index if specified
|
||||
device_info = self.vllm_config.device_config.device.__str__().split(
|
||||
":")
|
||||
device_info = self.vllm_config.device_config.device.__str__().split(":")
|
||||
if len(device_info) > 1:
|
||||
local_rank = int(device_info[1])
|
||||
rank = 0
|
||||
@@ -91,7 +84,7 @@ class DummyExecutor(UniProcExecutor):
|
||||
is_driver_worker=is_driver_worker,
|
||||
)
|
||||
self.mm_receiver_cache = None
|
||||
self.collective_rpc("init_worker", args=([kwargs], ))
|
||||
self.collective_rpc("init_worker", args=([kwargs],))
|
||||
self.collective_rpc("init_device")
|
||||
|
||||
@property
|
||||
@@ -99,5 +92,5 @@ class DummyExecutor(UniProcExecutor):
|
||||
return 2
|
||||
|
||||
def shutdown(self):
|
||||
if hasattr(self, 'thread_pool'):
|
||||
if hasattr(self, "thread_pool"):
|
||||
self.thread_pool.shutdown(wait=False)
|
||||
|
||||
@@ -17,14 +17,19 @@ import vllm.model_executor.model_loader.tensorizer
|
||||
from tests.utils import VLLM_PATH, RemoteOpenAIServer
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
|
||||
# yapf: disable
|
||||
from vllm.model_executor.model_loader.tensorizer import (TensorizerConfig,
|
||||
TensorSerializer,
|
||||
is_vllm_tensorized,
|
||||
open_stream,
|
||||
tensorize_vllm_model)
|
||||
from vllm.model_executor.model_loader.tensorizer import (
|
||||
TensorizerConfig,
|
||||
TensorSerializer,
|
||||
is_vllm_tensorized,
|
||||
open_stream,
|
||||
tensorize_vllm_model,
|
||||
)
|
||||
from vllm.model_executor.model_loader.tensorizer_loader import (
|
||||
BLACKLISTED_TENSORIZER_ARGS)
|
||||
BLACKLISTED_TENSORIZER_ARGS,
|
||||
)
|
||||
|
||||
# yapf: enable
|
||||
from vllm.utils import PlaceholderModule
|
||||
|
||||
@@ -44,7 +49,7 @@ class TensorizerCaughtError(Exception):
|
||||
|
||||
EXAMPLES_PATH = VLLM_PATH / "examples"
|
||||
|
||||
pytest_plugins = "pytest_asyncio",
|
||||
pytest_plugins = ("pytest_asyncio",)
|
||||
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
@@ -56,8 +61,7 @@ prompts = [
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
|
||||
|
||||
|
||||
def patch_init_and_catch_error(self, obj, method_name,
|
||||
expected_error: type[Exception]):
|
||||
def patch_init_and_catch_error(self, obj, method_name, expected_error: type[Exception]):
|
||||
original = getattr(obj, method_name, None)
|
||||
if original is None:
|
||||
raise ValueError("Method '{}' not found.".format(method_name))
|
||||
@@ -80,17 +84,19 @@ def assert_specific_tensorizer_error_is_raised(
|
||||
expected_error: type[Exception],
|
||||
):
|
||||
with pytest.raises(TensorizerCaughtError):
|
||||
executor.collective_rpc(patch_init_and_catch_error,
|
||||
args=(
|
||||
obj,
|
||||
method_name,
|
||||
expected_error,
|
||||
))
|
||||
executor.collective_rpc(
|
||||
patch_init_and_catch_error,
|
||||
args=(
|
||||
obj,
|
||||
method_name,
|
||||
expected_error,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def is_curl_installed():
|
||||
try:
|
||||
subprocess.check_call(['curl', '--version'])
|
||||
subprocess.check_call(["curl", "--version"])
|
||||
return True
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
return False
|
||||
@@ -99,13 +105,14 @@ def is_curl_installed():
|
||||
def write_keyfile(keyfile_path: str):
|
||||
encryption_params = EncryptionParams.random()
|
||||
pathlib.Path(keyfile_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(keyfile_path, 'wb') as f:
|
||||
with open(keyfile_path, "wb") as f:
|
||||
f.write(encryption_params.key)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
|
||||
def test_deserialized_encrypted_vllm_model_has_same_outputs(
|
||||
model_ref, vllm_runner, tmp_path, model_path):
|
||||
model_ref, vllm_runner, tmp_path, model_path
|
||||
):
|
||||
args = EngineArgs(model=model_ref)
|
||||
with vllm_runner(model_ref) as vllm_model:
|
||||
key_path = tmp_path / model_ref / "model.key"
|
||||
@@ -113,29 +120,30 @@ def test_deserialized_encrypted_vllm_model_has_same_outputs(
|
||||
|
||||
outputs = vllm_model.generate(prompts, sampling_params)
|
||||
|
||||
config_for_serializing = TensorizerConfig(tensorizer_uri=str(model_path),
|
||||
encryption_keyfile=str(key_path))
|
||||
config_for_serializing = TensorizerConfig(
|
||||
tensorizer_uri=str(model_path), encryption_keyfile=str(key_path)
|
||||
)
|
||||
|
||||
tensorize_vllm_model(args, config_for_serializing)
|
||||
|
||||
config_for_deserializing = TensorizerConfig(
|
||||
tensorizer_uri=str(model_path), encryption_keyfile=str(key_path))
|
||||
tensorizer_uri=str(model_path), encryption_keyfile=str(key_path)
|
||||
)
|
||||
|
||||
with vllm_runner(model_ref,
|
||||
load_format="tensorizer",
|
||||
model_loader_extra_config=config_for_deserializing
|
||||
) as loaded_vllm_model: # noqa: E501
|
||||
|
||||
deserialized_outputs = loaded_vllm_model.generate(
|
||||
prompts, sampling_params)
|
||||
with vllm_runner(
|
||||
model_ref,
|
||||
load_format="tensorizer",
|
||||
model_loader_extra_config=config_for_deserializing,
|
||||
) as loaded_vllm_model: # noqa: E501
|
||||
deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
|
||||
# noqa: E501
|
||||
|
||||
assert outputs == deserialized_outputs
|
||||
|
||||
|
||||
def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
|
||||
tmp_path, model_ref,
|
||||
model_path):
|
||||
def test_deserialized_hf_model_has_same_outputs(
|
||||
hf_runner, vllm_runner, tmp_path, model_ref, model_path
|
||||
):
|
||||
with hf_runner(model_ref) as hf_model:
|
||||
max_tokens = 50
|
||||
outputs = hf_model.generate_greedy(prompts, max_tokens=max_tokens)
|
||||
@@ -143,14 +151,17 @@ def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
|
||||
serializer = TensorSerializer(stream)
|
||||
serializer.write_module(hf_model.model)
|
||||
|
||||
with vllm_runner(model_ref,
|
||||
load_format="tensorizer",
|
||||
model_loader_extra_config=TensorizerConfig(
|
||||
tensorizer_uri=str(model_path),
|
||||
num_readers=1,
|
||||
)) as loaded_hf_model:
|
||||
with vllm_runner(
|
||||
model_ref,
|
||||
load_format="tensorizer",
|
||||
model_loader_extra_config=TensorizerConfig(
|
||||
tensorizer_uri=str(model_path),
|
||||
num_readers=1,
|
||||
),
|
||||
) as loaded_hf_model:
|
||||
deserialized_outputs = loaded_hf_model.generate_greedy(
|
||||
prompts, max_tokens=max_tokens)
|
||||
prompts, max_tokens=max_tokens
|
||||
)
|
||||
|
||||
assert outputs == deserialized_outputs
|
||||
|
||||
@@ -159,35 +170,37 @@ def test_load_without_tensorizer_load_format(vllm_runner, capfd, model_ref):
|
||||
model = None
|
||||
try:
|
||||
model = vllm_runner(
|
||||
model_ref,
|
||||
model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"))
|
||||
model_ref, model_loader_extra_config=TensorizerConfig(tensorizer_uri="test")
|
||||
)
|
||||
pytest.fail("Expected RuntimeError for extra config keys")
|
||||
except RuntimeError:
|
||||
out, err = capfd.readouterr()
|
||||
combined_output = out + err
|
||||
assert ("ValueError: Unexpected extra config keys for load "
|
||||
"format auto") in combined_output
|
||||
assert (
|
||||
"ValueError: Unexpected extra config keys for load format auto"
|
||||
) in combined_output
|
||||
finally:
|
||||
del model
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd,
|
||||
model_ref):
|
||||
def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd, model_ref):
|
||||
model = None
|
||||
try:
|
||||
model = vllm_runner(
|
||||
model_ref,
|
||||
load_format="safetensors",
|
||||
model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"))
|
||||
model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"),
|
||||
)
|
||||
pytest.fail("Expected RuntimeError for extra config keys")
|
||||
except RuntimeError:
|
||||
out, err = capfd.readouterr()
|
||||
|
||||
combined_output = out + err
|
||||
assert ("ValueError: Unexpected extra config keys "
|
||||
"for load format safetensors") in combined_output
|
||||
assert (
|
||||
"ValueError: Unexpected extra config keys for load format safetensors"
|
||||
) in combined_output
|
||||
finally:
|
||||
del model
|
||||
gc.collect()
|
||||
@@ -214,21 +227,24 @@ def test_tensorizer_with_tp_path_without_template(vllm_runner, capfd):
|
||||
except RuntimeError:
|
||||
out, err = capfd.readouterr()
|
||||
combined_output = out + err
|
||||
assert ("ValueError: For a sharded model, tensorizer_uri "
|
||||
"should include a string format template like '%04d' "
|
||||
"to be formatted with the rank "
|
||||
"of the shard") in combined_output
|
||||
assert (
|
||||
"ValueError: For a sharded model, tensorizer_uri "
|
||||
"should include a string format template like '%04d' "
|
||||
"to be formatted with the rank "
|
||||
"of the shard"
|
||||
) in combined_output
|
||||
|
||||
|
||||
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")
|
||||
def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
|
||||
vllm_runner, tmp_path):
|
||||
vllm_runner, tmp_path
|
||||
):
|
||||
model_ref = "EleutherAI/pythia-1.4b"
|
||||
# record outputs from un-sharded un-tensorized model
|
||||
with vllm_runner(
|
||||
model_ref,
|
||||
disable_custom_all_reduce=True,
|
||||
enforce_eager=True,
|
||||
model_ref,
|
||||
disable_custom_all_reduce=True,
|
||||
enforce_eager=True,
|
||||
) as base_model:
|
||||
outputs = base_model.generate(prompts, sampling_params)
|
||||
|
||||
@@ -254,21 +270,22 @@ def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
|
||||
assert os.path.isfile(model_path % 1), "Serialization subprocess failed"
|
||||
|
||||
with vllm_runner(
|
||||
model_ref,
|
||||
tensor_parallel_size=2,
|
||||
load_format="tensorizer",
|
||||
disable_custom_all_reduce=True,
|
||||
enforce_eager=True,
|
||||
model_loader_extra_config=tensorizer_config) as loaded_vllm_model:
|
||||
deserialized_outputs = loaded_vllm_model.generate(
|
||||
prompts, sampling_params)
|
||||
model_ref,
|
||||
tensor_parallel_size=2,
|
||||
load_format="tensorizer",
|
||||
disable_custom_all_reduce=True,
|
||||
enforce_eager=True,
|
||||
model_loader_extra_config=tensorizer_config,
|
||||
) as loaded_vllm_model:
|
||||
deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
|
||||
|
||||
assert outputs == deserialized_outputs
|
||||
|
||||
|
||||
@pytest.mark.flaky(reruns=3)
|
||||
def test_vllm_tensorized_model_has_same_outputs(model_ref, vllm_runner,
|
||||
tmp_path, model_path):
|
||||
def test_vllm_tensorized_model_has_same_outputs(
|
||||
model_ref, vllm_runner, tmp_path, model_path
|
||||
):
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
config = TensorizerConfig(tensorizer_uri=str(model_path))
|
||||
@@ -280,11 +297,10 @@ def test_vllm_tensorized_model_has_same_outputs(model_ref, vllm_runner,
|
||||
tensorize_vllm_model(args, config)
|
||||
assert is_vllm_tensorized(config)
|
||||
|
||||
with vllm_runner(model_ref,
|
||||
load_format="tensorizer",
|
||||
model_loader_extra_config=config) as loaded_vllm_model:
|
||||
deserialized_outputs = loaded_vllm_model.generate(
|
||||
prompts, sampling_params)
|
||||
with vllm_runner(
|
||||
model_ref, load_format="tensorizer", model_loader_extra_config=config
|
||||
) as loaded_vllm_model:
|
||||
deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
|
||||
# noqa: E501
|
||||
|
||||
assert outputs == deserialized_outputs
|
||||
@@ -314,15 +330,17 @@ def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref):
|
||||
|
||||
|
||||
def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
|
||||
|
||||
serialization_params = {
|
||||
"limit_cpu_concurrency": 2,
|
||||
}
|
||||
model_ref = "facebook/opt-125m"
|
||||
model_path = tmp_path / (model_ref + ".tensors")
|
||||
config = TensorizerConfig(tensorizer_uri=str(model_path),
|
||||
serialization_kwargs=serialization_params)
|
||||
llm = LLM(model=model_ref, )
|
||||
config = TensorizerConfig(
|
||||
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
|
||||
)
|
||||
llm = LLM(
|
||||
model=model_ref,
|
||||
)
|
||||
|
||||
def serialization_test(self, *args, **kwargs):
|
||||
# This is performed in the ephemeral worker process, so monkey-patching
|
||||
@@ -340,10 +358,13 @@ def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
|
||||
return original(self, *args, **kwargs)
|
||||
|
||||
tensorizer.serialization.TensorSerializer.__init__ = (
|
||||
tensorizer_serializer_wrapper)
|
||||
tensorizer_serializer_wrapper
|
||||
)
|
||||
|
||||
tensorizer_config = TensorizerConfig(**kwargs["tensorizer_config"])
|
||||
self.save_tensorized_model(tensorizer_config=tensorizer_config, )
|
||||
self.save_tensorized_model(
|
||||
tensorizer_config=tensorizer_config,
|
||||
)
|
||||
return to_compare | original_dict == to_compare
|
||||
|
||||
kwargs = {"tensorizer_config": config.to_serializable()}
|
||||
@@ -351,9 +372,7 @@ def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
|
||||
assert assert_from_collective_rpc(llm, serialization_test, kwargs)
|
||||
|
||||
|
||||
def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(
|
||||
tmp_path, capfd):
|
||||
|
||||
def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
|
||||
deserialization_kwargs = {
|
||||
"num_readers": "bar", # illegal value
|
||||
}
|
||||
@@ -364,8 +383,9 @@ def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(
|
||||
|
||||
model_ref = "facebook/opt-125m"
|
||||
model_path = tmp_path / (model_ref + ".tensors")
|
||||
config = TensorizerConfig(tensorizer_uri=str(model_path),
|
||||
serialization_kwargs=serialization_params)
|
||||
config = TensorizerConfig(
|
||||
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
|
||||
)
|
||||
|
||||
args = EngineArgs(model=model_ref)
|
||||
tensorize_vllm_model(args, config)
|
||||
@@ -393,7 +413,6 @@ def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(
|
||||
|
||||
|
||||
def test_assert_stream_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
|
||||
|
||||
deserialization_kwargs = {
|
||||
"num_readers": 1,
|
||||
}
|
||||
@@ -404,8 +423,9 @@ def test_assert_stream_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
|
||||
|
||||
model_ref = "facebook/opt-125m"
|
||||
model_path = tmp_path / (model_ref + ".tensors")
|
||||
config = TensorizerConfig(tensorizer_uri=str(model_path),
|
||||
serialization_kwargs=serialization_params)
|
||||
config = TensorizerConfig(
|
||||
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
|
||||
)
|
||||
|
||||
args = EngineArgs(model=model_ref)
|
||||
tensorize_vllm_model(args, config)
|
||||
@@ -441,16 +461,24 @@ async def test_serialize_and_serve_entrypoints(tmp_path):
|
||||
|
||||
suffix = "test"
|
||||
try:
|
||||
result = subprocess.run([
|
||||
sys.executable,
|
||||
f"{VLLM_PATH}/examples/others/tensorize_vllm_model.py", "--model",
|
||||
model_ref, "serialize", "--serialized-directory",
|
||||
str(tmp_path), "--suffix", suffix, "--serialization-kwargs",
|
||||
'{"limit_cpu_concurrency": 4}'
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True)
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
f"{VLLM_PATH}/examples/others/tensorize_vllm_model.py",
|
||||
"--model",
|
||||
model_ref,
|
||||
"serialize",
|
||||
"--serialized-directory",
|
||||
str(tmp_path),
|
||||
"--suffix",
|
||||
suffix,
|
||||
"--serialization-kwargs",
|
||||
'{"limit_cpu_concurrency": 4}',
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print("Tensorizing failed.")
|
||||
print("STDOUT:\n", e.stdout)
|
||||
@@ -470,14 +498,20 @@ async def test_serialize_and_serve_entrypoints(tmp_path):
|
||||
"deserialization_kwargs": {
|
||||
"verify_hash": True,
|
||||
"num_readers": 8,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
cmd = [
|
||||
"-m", "vllm.entrypoints.cli.main", "serve", "--host", "localhost",
|
||||
"--load-format", "tensorizer", model_ref,
|
||||
"-m",
|
||||
"vllm.entrypoints.cli.main",
|
||||
"serve",
|
||||
"--host",
|
||||
"localhost",
|
||||
"--load-format",
|
||||
"tensorizer",
|
||||
model_ref,
|
||||
"--model-loader-extra-config",
|
||||
json.dumps(model_loader_extra_config, indent=2)
|
||||
json.dumps(model_loader_extra_config, indent=2),
|
||||
]
|
||||
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
@@ -500,17 +534,16 @@ async def test_serialize_and_serve_entrypoints(tmp_path):
|
||||
|
||||
|
||||
@pytest.mark.parametrize("illegal_value", BLACKLISTED_TENSORIZER_ARGS)
|
||||
def test_blacklisted_parameter_for_loading(tmp_path, vllm_runner, capfd,
|
||||
illegal_value):
|
||||
|
||||
def test_blacklisted_parameter_for_loading(tmp_path, vllm_runner, capfd, illegal_value):
|
||||
serialization_params = {
|
||||
"limit_cpu_concurrency": 2,
|
||||
}
|
||||
|
||||
model_ref = "facebook/opt-125m"
|
||||
model_path = tmp_path / (model_ref + ".tensors")
|
||||
config = TensorizerConfig(tensorizer_uri=str(model_path),
|
||||
serialization_kwargs=serialization_params)
|
||||
config = TensorizerConfig(
|
||||
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
|
||||
)
|
||||
|
||||
args = EngineArgs(model=model_ref)
|
||||
tensorize_vllm_model(args, config)
|
||||
@@ -526,5 +559,6 @@ def test_blacklisted_parameter_for_loading(tmp_path, vllm_runner, capfd,
|
||||
except RuntimeError:
|
||||
out, err = capfd.readouterr()
|
||||
combined_output = out + err
|
||||
assert (f"ValueError: {illegal_value} is not an allowed "
|
||||
f"Tensorizer argument.") in combined_output
|
||||
assert (
|
||||
f"ValueError: {illegal_value} is not an allowed Tensorizer argument."
|
||||
) in combined_output
|
||||
|
||||
Reference in New Issue
Block a user