def test_retrieve_msg_predict_binary(self, socket_patches):
expected = [{
"requestId":
b"request_id",
"headers": [],
"parameters": [{
"name": "input_name",
"contentType": "",
"value": b"binary"
}]
}]
socket_patches.socket.recv.side_effect = [
b"I",
b"\x00\x00\x00\x0a",
b"request_id",
b"\xFF\xFF\xFF\xFF",
b"\x00\x00\x00\x0a",
b"input_name",
b"\x00\x00\x00\x00",
b"\x00\x00\x00\x06",
b"binary",
b"\xFF\xFF\xFF\xFF", # end of parameters
b"\xFF\xFF\xFF\xFF" # end of batch
]
cmd, ret = codec.retrieve_msg(socket_patches.socket)
assert cmd == b'I'
assert ret == expected
def test_retrieve_msg_predict(self, socket_patches):
expected = [{
"requestId":
b"request_id",
"headers": [],
"parameters": [{
"name": "input_name",
"contentType": "application/json",
"value": {
"data": "value"
}
}]
}]
socket_patches.socket.recv.side_effect = [
b"I",
b"\x00\x00\x00\x0a",
b"request_id",
b"\xFF\xFF\xFF\xFF",
b"\x00\x00\x00\x0a",
b"input_name",
b"\x00\x00\x00\x0F",
b"application/json",
b"\x00\x00\x00\x0F",
b'{"data":"value"}',
b"\xFF\xFF\xFF\xFF", # end of parameters
b"\xFF\xFF\xFF\xFF" # end of batch
]
cmd, ret = codec.retrieve_msg(socket_patches.socket)
assert cmd == b'I'
assert ret == expected
def test_retrieve_msg_predict_text(self, socket_patches):
expected = [{
"requestId":
b"request_id",
"headers": [],
"parameters": [{
"name": "input_name",
"contentType": "text/plain",
"value": u"text_value测试"
}]
}]
socket_patches.socket.recv.side_effect = [
b"I",
b"\x00\x00\x00\x0a",
b"request_id",
b"\xFF\xFF\xFF\xFF",
b"\x00\x00\x00\x0a",
b"input_name",
b"\x00\x00\x00\x0a",
b"text/plain",
b"\x00\x00\x00\x0a",
bytes(u"text_value测试", "utf-8"),
b"\xFF\xFF\xFF\xFF", # end of parameters
b"\xFF\xFF\xFF\xFF" # end of batch
]
cmd, ret = codec.retrieve_msg(socket_patches.socket)
assert cmd == b'I'
assert ret == expected
def handle_connection(self, cl_socket):
"""
Handle socket connection.
:param cl_socket:
:return:
"""
service = None
while True:
cmd, msg = retrieve_msg(cl_socket)
if cmd == b'I':
resp = service.predict(msg)
cl_socket.send(resp)
elif cmd == b'L':
service, result, code = self.load_model(msg)
resp = bytearray()
resp += create_load_model_response(code, result)
cl_socket.send(resp)
if code != 200:
raise RuntimeError("{} - {}".format(code, result))
else:
raise ValueError("Received unknown command: {}".format(cmd))
if service is not None and service.context is not None and service.context.metrics is not None:
emit_metrics(service.context.metrics.store)
def handle_connection(self, cl_socket):
"""
Handle socket connection.
:param cl_socket:
:return:
"""
logging.basicConfig(stream=sys.stdout,
format="%(message)s",
level=logging.INFO)
cl_socket.setblocking(True)
while True:
cmd, msg = retrieve_msg(cl_socket)
if cmd == b'I':
resp = self.service.predict(msg)
cl_socket.send(resp)
elif cmd == b'L':
result, code = self.load_model(msg)
resp = bytearray()
resp += create_load_model_response(code, result)
cl_socket.send(resp)
self._remap_io()
if code != 200:
raise RuntimeError("{} - {}".format(code, result))
else:
raise ValueError("Received unknown command: {}".format(cmd))
if self.service is not None and self.service.context is not None \
and self.service.context.metrics is not None:
emit_metrics(self.service.context.metrics.store)
def test_retrieve_msg_load_no_gpu(self, socket_patches):
expected = {"modelName": b"model_name", "modelPath": b"model_path",
"batchSize": 1, "handler": b"handler"}
socket_patches.socket.recv.side_effect = [
b"L",
b"\x00\x00\x00\x0a", b"model_name",
b"\x00\x00\x00\x0a", b"model_path",
b"\x00\x00\x00\x01",
b"\x00\x00\x00\x07", b"handler",
b"\xFF\xFF\xFF\xFF"
]
cmd, ret = codec.retrieve_msg(socket_patches.socket)
assert cmd == b"L"
assert ret == expected
def handle_connection(self, cl_socket):
"""
Handle socket connection.
:param cl_socket:
:return:
"""
service = None
while True:
cmd, msg = retrieve_msg(cl_socket)
if cmd == b'I':
resp = service.predict(msg)
cl_socket.send(resp)
elif cmd == b'L':
service, result, code = self.load_model(msg)
resp = bytearray()
resp += create_load_model_response(code, result)
cl_socket.send(resp)
else:
raise ValueError("Received unknown command: {}".format(cmd))
def test_retrieve_msg_exceed_buffer_size(self, socket_patches):
socket_patches.socket.recv.side_effect = [b"L", b"\x0F\x00\x00\x03"]
with pytest.raises(ValueError, match=r"Exceed max buffer size: .*"):
codec.retrieve_msg(socket_patches.socket)
def test_retrieve_msg_unknown(self, socket_patches):
socket_patches.socket.recv.side_effect = [b"U", b"\x00\x00\x00\x03"]
with pytest.raises(ValueError, match=r"Invalid command: .*"):
codec.retrieve_msg(socket_patches.socket)
