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