def _ws_send_recv(self):
to_server = SimulatorToServer()
self._on_send(to_server)
log.pb("to_server: {}".format(MessageToJson(to_server)))
if (to_server.message_type):
out_bytes = to_server.SerializeToString()
try:
yield self._ws.write_message(out_bytes, binary=True)
except (StreamClosedError, WebSocketClosedError) as e:
raise BonsaiServerError(
"Websocket connection closed. Code: {}, Reason: {}".format(
self._ws.close_code, self._ws.close_reason))
# read response from server
in_bytes = yield self._ws.read_message()
if in_bytes is None:
raise BonsaiServerError(
"Websocket connection closed. Code: {}, Reason: {}".format(
self._ws.close_code, self._ws.close_reason))
from_server = ServerToSimulator()
from_server.ParseFromString(in_bytes)
log.pb("from_server: {}".format(MessageToJson(from_server)))
self._on_recv(from_server)
def on_message(self, in_bytes):
from_sim = SimulatorToServer()
from_sim.ParseFromString(in_bytes)
self._validate_message(from_sim)
mtype = self._dispatch_mtype(from_sim.message_type)
# dict->json->Message->binary
msg_dict = self._message_data[mtype]
json_msg = json.dumps(msg_dict)
msg = ServerToSimulator()
Parse(json_msg, msg)
self.write_message(msg.SerializeToString(), binary=True)
self._prev = mtype
def run(self):
""" Run loop called from Simulator. Encapsulates one round trip
to the backend, which might include a simulation loop.
"""
# Grab a web socket connection if needed
if self._ws is None:
message = yield self._connect()
# If the connection failed, report
if message is not None:
raise BonsaiServerError(
"Error while connecting to websocket: {}".format(message))
# If there is a batch of predictions cued up, step through it
if self._prev_message_type == ServerToSimulator.PREDICTION:
for step in self._sim_steps:
self._advance(step)
# send message to server
to_server = SimulatorToServer()
self._on_send(to_server)
if (to_server.message_type):
out_bytes = to_server.SerializeToString()
try:
yield self._ws.write_message(out_bytes, binary=True)
except (StreamClosedError, WebSocketClosedError) as e:
raise BonsaiServerError(
"Websocket connection closed. Code: {}, Reason: {}".format(
self._ws.close_code, self._ws.close_reason))
# read response from server
in_bytes = yield self._ws.read_message()
if in_bytes is None:
raise BonsaiServerError(
"Websocket connection closed. Code: {}, Reason: {}".format(
self._ws.close_code, self._ws.close_reason))
from_server = ServerToSimulator()
from_server.ParseFromString(in_bytes)
self._on_recv(from_server)
if self._prev_message_type == ServerToSimulator.FINISHED:
yield self._ws.close()
raise gen.Return(False)
# You've come this far, celebrate!
raise gen.Return(True)
def on_message(self, in_bytes):
# dummy message and response for testing libbonsai's custom
# websocket client
if in_bytes == b"foobar":
self.write_message("bazqux", binary=True)
return
from_sim = SimulatorToServer()
from_sim.ParseFromString(in_bytes)
self._validate_message(from_sim)
sim_name = from_sim.register_data.simulator_name
mtype = self._dispatch_mtype(from_sim.message_type)
if self._FLAKY and \
self._count > self._fail_point and \
self._count < self._fail_point + self._fail_duration:
self.close(code=1008, reason=None)
return
elif BRAIN_STATUS['state'] == "Stopped":
self.close(
# code 1001 means: brain has finished training
code=1001,
reason="Brain no longer training"
)
return
elif (from_sim.message_type == SimulatorToServer.REGISTER and
sim_name not in SIMS.keys()):
self.close(
# code 4043 means: simulator does not exist
code=4043,
reason="Simulator {} does not exist.".format(sim_name)
)
return
# dict->json->Message->binary
msg_dict = self._message_data[mtype]
json_msg = json.dumps(msg_dict)
msg = ServerToSimulator()
Parse(json_msg, msg)
self.write_message(msg.SerializeToString(), binary=True)
self._prev = mtype
async def _ws_send_recv(self):
to_server = SimulatorToServer()
self._on_send(to_server)
log.pb("to_server: {}".format(MessageToJson(to_server)))
if to_server.message_type:
out_bytes = to_server.SerializeToString()
try:
with self._sim_connection.lock:
if self._sim_connection.client.closed:
await self._handle_disconnect(
"Attempted write to closed web socket")
return
await self._sim_connection.client.send_bytes(out_bytes)
except ClientError as e:
await self._handle_disconnect(e)
return
try:
with self._sim_connection.lock:
log.network('Reading response from server')
msg = await self._sim_connection.client.receive()
log.network('Received response from server')
except TimeoutError as e:
log.error('WS read took longer than {} seconds. '
'Sim will be disconnected.'.format(
self._sim_connection.read_timeout_seconds))
await self._handle_disconnect()
return
if msg.type == WSMsgType.CLOSE or msg.type == WSMsgType.CLOSED \
or msg.type == WSMsgType.ERROR or isinstance(msg.data, EofStream):
await self._handle_disconnect(msg.extra)
return
from_server = ServerToSimulator()
from_server.ParseFromString(msg.data)
log.pb("from_server: {}".format(MessageToJson(from_server)))
self._on_recv(from_server)
def _ws_send_recv(self):
to_server = SimulatorToServer()
self._on_send(to_server)
log.pb("to_server: {}".format(MessageToJson(to_server)))
if to_server.message_type:
out_bytes = to_server.SerializeToString()
try:
with self._sim_connection.lock:
log.network('Writing message to server')
yield self._sim_connection.client.write_message(
out_bytes, binary=True)
log.network('Wrote message to server')
except (StreamClosedError, WebSocketClosedError) as e:
self._handle_disconnect()
return
try:
with self._sim_connection.lock:
log.network('Reading response from server')
in_bytes = yield gen.with_timeout(
timedelta(
seconds=self._sim_connection.read_timeout_seconds),
self._sim_connection.client.read_message())
log.network('Received response from server')
except gen.TimeoutError as e:
log.error(
'WS read took longer than {} seconds. '
'Sim will be disconnected.'.format(
self._sim_connection.read_timeout_seconds))
self._handle_disconnect()
return
if in_bytes is None:
self._handle_disconnect()
return
from_server = ServerToSimulator()
from_server.ParseFromString(in_bytes)
log.pb("from_server: {}".format(MessageToJson(from_server)))
self._on_recv(from_server)
async def _ws_send_recv(self):
to_server = SimulatorToServer()
self._on_send(to_server)
log.pb("to_server: {}".format(MessageToJson(to_server)))
if to_server.message_type:
out_bytes = to_server.SerializeToString()
try:
with self._sim_connection.lock:
if self._sim_connection.client.closed:
await self._handle_disconnect(
"Attempted write to closed web socket")
return
log.network('Attempting to send message to server.')
await self._sim_connection.client.send_bytes(out_bytes)
log.network('Message sent to server.')
except ClientError as e:
await self._handle_disconnect(e)
return
with self._sim_connection.lock:
log.network('Waiting for message from server.')
self._receive_handle = ensure_future(
self._sim_connection.client.receive(), loop=self._ioloop)
msg = await self._receive_handle
log.network('Received message from server.')
if msg.type == WSMsgType.CLOSE or msg.type == WSMsgType.CLOSED \
or msg.type == WSMsgType.ERROR or isinstance(msg.data, EofStream):
await self._handle_disconnect(msg.extra)
return
from_server = ServerToSimulator()
from_server.ParseFromString(msg.data)
log.pb("from_server: {}".format(MessageToJson(from_server)))
self._on_recv(from_server)
async def handle_msg(self, request):
self._prev = ServerToSimulator.UNKNOWN
if self._UNAUTHORIZED:
return web.Response(status=401, text="Unauthorized")
if self._FORBIDDEN:
return web.Response(status=403, text="Forbidden")
if self._FLAKY and \
self._count > self._fail_point and \
self._count < self._fail_point + self._fail_duration:
return web.Response(status=503, text="Service Unavailable")
ws = web.WebSocketResponse(protocols=['', 'bonsaiauth'])
await ws.prepare(request)
ws.force_close()
if self._PONG:
if os.path.exists('pong.json'):
os.remove('pong.json')
with async_timeout.timeout(1, loop=ws._loop):
msg = await ws._reader.read()
if msg.type == WSMsgType.PONG:
pong_json = {'PONG': 1}
with open('pong.json', 'w') as outfile:
json.dump(pong_json, outfile)
await ws.close()
return ws
request.app['websockets'].add(ws)
try:
async for msg in ws:
if msg.type == WSMsgType.CLOSE:
await ws.close()
elif msg.data == b'foobar':
await ws.send_bytes(b'bazqux')
else:
self._count += 1
from_sim = SimulatorToServer()
from_sim.ParseFromString(msg.data)
self._validate_message(from_sim)
sim_name = from_sim.register_data.simulator_name
mtype = self._dispatch_mtype(from_sim.message_type)
if self._FLAKY and \
self._count > self._fail_point and \
self._count < self._fail_point + self._fail_duration:
await ws.close(code=1008, message=b'')
return ws
elif BRAIN_STATUS['state'] == "Stopped":
await ws.close(
# code 1001 means: brain has finished training
code=1001,
message=b"Brain no longer training"
)
return ws
elif (from_sim.message_type == SimulatorToServer.REGISTER and
sim_name not in SIMS.keys()):
msg = "Simulator {} does not exist.".format(sim_name)
await ws.close(
# code 4043 means: simulator does not exist
code=4043, message=bytes(msg, 'utf-8'))
return ws
if self._EOFSTREAM:
# Intentionally send a bad parameter
await ws.send_bytes(cast(Any, EofStream()))
elif self._ERROR_MSG:
# Intentionally send a bad parameter
await ws.send_bytes(cast(Any, 'foo'))
else:
# dict->json->Message->binary
msg_dict = self._message_data[mtype]
json_msg = json.dumps(msg_dict)
msg = ServerToSimulator()
Parse(json_msg, msg)
await ws.send_bytes(msg.SerializeToString())
self._prev = mtype
finally:
request.app['websockets'].discard(ws)
await ws.close()
return ws
_msg_files = {
ServerToSimulator.ACKNOWLEDGE_REGISTER: "_on_acknowledge_register",
ServerToSimulator.SET_PROPERTIES: "_on_set_properties",
ServerToSimulator.START: "_on_start",
ServerToSimulator.PREDICTION: "_on_prediction",
ServerToSimulator.STOP: "_on_stop",
ServerToSimulator.RESET: "_on_reset"
}
_other_msg_types = [ServerToSimulator.UNKNOWN, ServerToSimulator.FINISHED]
blob = [None] * N_MESSAGES
for mtype, fname in _msg_files.items():
with open(fname, 'rb') as f:
msg = ServerToSimulator()
b = f.read()
msg.ParseFromString(b)
j = MessageToJson(msg)
data = json.loads(j)
blob[mtype] = data
for mtype in _other_msg_types:
msg = ServerToSimulator()
msg.message_type = mtype
j = MessageToJson(msg)
data = json.loads(j)
blob[mtype] = data
with open(args.fname[0], 'w') as f:
blob_str = json.dumps(blob, indent=2)