async def _producer_handler(self, producer, source_name):
'''
Loop waiting until there is a token available. Then calls producer to
get the gabriel_pb2.InputFrame to send.
'''
await self._welcome_event.wait()
source = self._sources.get(source_name)
assert source is not None, (
"No engines consume frames from source: {}".format(source_name))
while self._running:
await source.get_token()
input_frame = await producer()
if input_frame is None:
source.return_token()
logger.info('Received None from producer')
continue
from_client = gabriel_pb2.FromClient()
from_client.frame_id = source.get_frame_id()
from_client.source_name = source_name
from_client.input_frame.CopyFrom(input_frame)
try:
await self._send_from_client(from_client)
except websockets.exceptions.ConnectionClosed:
return # stop the handler
logger.debug('num_tokens for %s is now %d', source_name,
source.get_num_tokens())
source.next_frame()
async def producer():
_, frame = self._video_capture.read()
if frame is None:
return None
frame = self._preprocess(frame)
_, jpeg_frame = cv2.imencode('.jpg', frame)
from_client = gabriel_pb2.FromClient()
from_client.payload_type = gabriel_pb2.PayloadType.IMAGE
from_client.filter_passed = self._filter_passed
from_client.payloads_for_frame.append(jpeg_frame.tostring())
extras = self._produce_extras()
if extras is not None:
from_client.extras.Pack(extras)
return from_client
async def receiver():
loop = asyncio.get_event_loop()
stream_reader = asyncio.StreamReader(loop=loop)
def protocol_factory():
return asyncio.StreamReaderProtocol(stream_reader)
transport = await loop.connect_read_pipe(
protocol_factory, os.fdopen(self._read, mode='r'))
while True:
size_bytes = await stream_reader.readexactly(
_NUM_BYTES_FOR_SIZE)
size_of_message = int.from_bytes(size_bytes, _BYTEORDER)
from_client = gabriel_pb2.FromClient()
from_client.ParseFromString(
await stream_reader.readexactly(size_of_message))
self._latest_item = from_client
self._frame_available.set()
async def _consumer(self, websocket, client):
address = websocket.remote_address
async for raw_input in websocket:
logger.debug('Received input from %s', address)
from_client = gabriel_pb2.FromClient()
from_client.ParseFromString(raw_input)
status = await self._consumer_helper(client, from_client, address)
if status == ResultWrapper.Status.SUCCESS:
client.tokens_for_source[from_client.source_name] -= 1
continue
# Send error message
to_client = gabriel_pb2.ToClient()
to_client.response.source_name = from_client.source_name
to_client.response.frame_id = from_client.frame_id
to_client.response.return_token = True
to_client.response.result_wrapper.status = status
await websocket.send(to_client.SerializeToString())
def run(engine, filter_name, server_address, timeout=TEN_SECONDS,
request_retries=REQUEST_RETRIES):
context = zmq.Context()
while request_retries > 0:
socket = context.socket(zmq.REQ)
socket.connect(server_address)
to_server_runner = gabriel_pb2.ToServerRunner()
to_server_runner.welcome.filter_name = filter_name
socket.send(to_server_runner.SerializeToString())
logger.info('Sent welcome message to server')
while True:
if socket.poll(timeout) == 0:
logger.warning('No response from server')
socket.setsockopt(zmq.LINGER, 0)
socket.close()
request_retries -= 1
break
message_from_server = socket.recv()
if message_from_server == network_engine.HEARTBEAT:
socket.send(network_engine.HEARTBEAT)
continue
from_client = gabriel_pb2.FromClient()
from_client.ParseFromString(message_from_server)
assert from_client.filter_passed == filter_name
to_server_runner = gabriel_pb2.ToServerRunner()
result_wrapper = engine.handle(from_client)
to_server_runner.result_wrapper.CopyFrom(result_wrapper)
socket.send(to_server_runner.SerializeToString())
logger.warning('Ran out of retires. Abandoning server connection.')