[Inference]Update inference config and fix test (#5178)

* unify the config setting

* fix test

* fix import

* fix test

* fix

* fix

* add logger

* revise log info

---------

Co-authored-by: CjhHa1 <cjh18671720497outlook.com>
This commit is contained in:
Jianghai
2023-12-12 17:22:41 +08:00
committed by FrankLeeeee
parent 3de2e62299
commit 93aeacca34
9 changed files with 61 additions and 25 deletions

View File

@@ -1,9 +1,10 @@
from colossalai.inference.core.config import InferenceConfig
from colossalai.inference.core.inference_struct import BatchHandler, Sequence
from colossalai.inference.config import InferenceConfig
from colossalai.inference.struct import BatchInfo, RequsetStatus, Sequence
def test_config_and_struct():
InferenceConfig("/llama")
def test_config_and_inferenceData():
config = InferenceConfig("/llama")
assert config.max_batch_size
sequence = Sequence(
request_id=1,
prompt="abc",
@@ -27,11 +28,16 @@ def test_config_and_struct():
assert sequence.get_output_len() == 0
assert sequence.check_finish() == False
batch = BatchHandler.init_batch([sequence])
batch = BatchInfo.init_batch([sequence])
assert batch.block_table[sequence.request_id] == sequence.block_table_index
sequence.status = RequsetStatus.COMPLETED
batch.fliter_batch()
assert batch.block_table == {}
batch.add_seqs([sequence2])
assert batch.block_table[sequence2.request_id] == sequence2.block_table_index
batch.clear_batch()
assert batch.block_table == {}
if __name__ == "__main__":
test_config_and_struct()
test_config_and_inferenceData()