mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-29 04:05:35 +00:00
55cc7f3df7
* clean requirements * modify example inference struct * add test ci scripts * mark test_infer as submodule * rm deprecated cls & deps * import of HAS_FLASH_ATTN * prune inference tests to be run * prune triton kernel tests * increment pytest timeout mins * revert import path in openmoe
47 lines
1.2 KiB
Python
Executable File
47 lines
1.2 KiB
Python
Executable File
import pytest
|
|
|
|
import colossalai
|
|
from colossalai.inference.config import InferenceConfig
|
|
from colossalai.inference.struct import RequestStatus, Sequence
|
|
from colossalai.testing import rerun_if_address_is_in_use, spawn
|
|
|
|
|
|
def check_config_and_inference():
|
|
config = InferenceConfig()
|
|
assert config.max_batch_size == 8
|
|
sequence = Sequence(
|
|
request_id=1,
|
|
prompt="abc",
|
|
input_token_id=[1, 2, 3],
|
|
block_size=16,
|
|
sample_params=None,
|
|
eos_token_id=2,
|
|
pad_token_id=2,
|
|
max_output_len=256,
|
|
)
|
|
|
|
sequence.mark_running()
|
|
assert sequence.status == RequestStatus.RUNNING
|
|
sequence.recycle()
|
|
assert sequence.status == RequestStatus.RECYCLED
|
|
|
|
assert sequence.sentence_len == 3
|
|
assert sequence.input_len == 3
|
|
assert sequence.output_len == 0
|
|
assert sequence.check_finish() == False
|
|
|
|
|
|
def run_dist(rank, world_size, port):
|
|
colossalai.launch(rank=rank, world_size=world_size, port=port, host="localhost")
|
|
check_config_and_inference()
|
|
|
|
|
|
@pytest.mark.dist
|
|
@rerun_if_address_is_in_use()
|
|
def test_config_and_inference():
|
|
spawn(run_dist, 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_config_and_inference()
|