[Misc] Improve conftest (#5681)

This commit is contained in:
Cyrus Leung 2024-06-20 10:09:21 +08:00 committed by GitHub
parent 949e49a685
commit 3730a1c832
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -365,7 +365,7 @@ class HfRunner:
cleanup() cleanup()
@pytest.fixture @pytest.fixture(scope="session")
def hf_runner(): def hf_runner():
return HfRunner return HfRunner
@ -385,6 +385,7 @@ class VllmRunner:
block_size: int = 16, block_size: int = 16,
enable_chunked_prefill: bool = False, enable_chunked_prefill: bool = False,
swap_space: int = 4, swap_space: int = 4,
enforce_eager: bool = False,
**kwargs, **kwargs,
) -> None: ) -> None:
self.model = LLM( self.model = LLM(
@ -393,6 +394,7 @@ class VllmRunner:
trust_remote_code=True, trust_remote_code=True,
dtype=dtype, dtype=dtype,
swap_space=swap_space, swap_space=swap_space,
enforce_eager=enforce_eager,
disable_log_stats=disable_log_stats, disable_log_stats=disable_log_stats,
tensor_parallel_size=tensor_parallel_size, tensor_parallel_size=tensor_parallel_size,
max_model_len=max_model_len, max_model_len=max_model_len,