mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
e06f537e08
* Update TensorRT-LLM --------- Co-authored-by: erenup <ping.nie@pku.edu.cn> Co-authored-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
42 lines
707 B
Python
42 lines
707 B
Python
from dataclasses import dataclass
|
|
from typing import List
|
|
|
|
from tensorrt_llm.hlapi.mpi_session import SocketListener
|
|
|
|
|
|
@dataclass
|
|
class ComplexData:
|
|
a: str
|
|
b: int
|
|
c: List[int]
|
|
|
|
|
|
def test_SocketServer():
|
|
|
|
messages = [
|
|
"hello", # str
|
|
123, # int
|
|
ComplexData("hello", 123, [1, 2, 3]) # complex
|
|
]
|
|
|
|
offset = 0
|
|
|
|
def callback(data):
|
|
nonlocal offset
|
|
print('get data', data)
|
|
assert data == messages[offset]
|
|
offset += 1
|
|
|
|
server = SocketListener(callback=callback)
|
|
|
|
client = server.get_client()
|
|
|
|
for data in messages:
|
|
client.send(data)
|
|
|
|
server.shutdown()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_SocketServer()
|