def __call__(self, img, available_at):
# Choose the return value
if len(self.deque) > 0 and self.deque[0].ready():
last = self.deque.popleft()
res = last.get()
if res is not None:
pyprofile.timing('vnc_env.diagnostics.async_decode.latency',
time.time() - res['available_at'])
else:
res = False
pyprofile.gauge('vnc_env.diagnostics.async_decode.queue_depth',
len(self.deque))
# Just grayscale it by keeping only one component. Should be
# good enough as this region is black and white anyway.
grayscale = img[self.y:self.y + self.height,
self.x:self.x + self.width, 0]
# Apply processing if needed
match = np.array_equal(self._last_img, grayscale)
if not match:
pyprofile.incr('vnc_env.diagnostics.async_decode.schedule')
# sneakily copy if numpy hasn't, so it can be cached
self._last_img = np.ascontiguousarray(grayscale)
async = self.qr_pool.apply_async(
self.method, (self._last_img, time.time(), available_at))
self.deque.append(async)
else:
pyprofile.incr('vnc_env.diagnostics.async_decode.cache_hit')
return res
def sendMessage(self, data):
pyprofile.incr('vnc_client.data.sent.messages')
pyprofile.incr('vnc_client.data.sent.bytes',
len(data),
unit=pyprofile.BYTES)
if self.transport:
self.transport.write(data)
def dataReceived(self, data):
pyprofile.incr('vnc_proxy_server.data.sent.messages')
pyprofile.incr('vnc_proxy_server.data.sent.bytes', len(data))
self.buf.append(data)
self.buf_len += len(data)
self.flush()
def send_env_reward(self, reward, done, info, episode_id):
pyprofile.incr('agent_conn.reward', reward)
if done:
pyprofile.incr('agent_conn.done')
reactor.callFromThread(self._send_env_reward, reward, done, info,
episode_id)
def recv_DecodeZlib_value(self, block, x, y, width, height):
pyprofile.incr('vncdriver.recv_rectangle.zlib_encoding.bytes',
len(block),
unit=pyprofile.BYTES)
rectangle = server_messages.ZlibEncoding.parse_rectangle(
self, x, y, width, height, block)
self._rectangles.append(rectangle)
self._process_rectangles()
def recv_DecodeZlib(self, block, x, y, width, height):
pyprofile.incr('vncdriver.recv_rectangle.zlib_encoding')
pyprofile.incr('vncdriver.recv_rectangle.zlib_encoding.bytes',
len(block),
unit=pyprofile.BYTES)
(length, ) = struct.unpack('!I', block)
self.expect(self.recv_DecodeZlib_value, length, x, y, width, height)
def update_rectangle(self, x, y, width, height, data):
bytes = data.tobytes()
pyprofile.incr('vncdriver.pyglet_screen.blit')
pyprofile.incr('vncdriver.pyglet_screen.blit.bytes', len(bytes), unit=pyprofile.BYTES)
import pyglet
image = pyglet.image.ImageData(width, height, 'RGB', bytes, pitch=width * -3)
self.texture.blit_into(image, x, self._height-height-y, 0)
self._is_updated = True
def recv_FramebufferUpdate(self, block):
(number_of_rectangles, ) = struct.unpack('!xH', block)
logger.debug('Receiving %d rectangles', number_of_rectangles)
pyprofile.incr('vncdriver.framebuffer_update')
pyprofile.incr('vncdriver.framebuffer_update.number_of_rectangles',
number_of_rectangles)
self._remaining_rectangles = number_of_rectangles
self._rectangles = []
self._process_rectangles()
def dataReceived(self, data):
pyprofile.incr('vnc_client.data.received.messages')
pyprofile.incr('vnc_client.data.received.bytes',
len(data),
unit=pyprofile.BYTES)
self.buf.append(data)
self.buf_len += len(data)
logger.debug('Received data: %s bytes (brings us to %s total)',
len(data), self.buf_len)
self.flush()
def dataReceived(self, data):
if self.expected is None:
# We're in direct proxy mode
self.proxyData(data)
return
pyprofile.incr('vnc_proxy_server.data.sent.messages')
pyprofile.incr('vnc_proxy_server.data.sent.bytes', len(data))
self.buf.append(data)
self.buf_len += len(data)
self.flush()
def parse_rectangle(cls, client, x, y, width, height, data):
decompressed = client.zlib_decompressor.decompress(data)
logger.debug('[zlib] Decompressed from %s bytes -> %s bytes',
len(data), len(decompressed))
pyprofile.incr(
'vncdriver.recv_rectangle.zlib_encoding.decompressed_bytes',
len(decompressed),
unit=pyprofile.BYTES)
data = np.frombuffer(decompressed, np.uint8).reshape(
(height, width, 4))[:, :, [0, 1, 2]]
encoding = cls(data)
return Rectangle(x, y, width, height, encoding)
def parse_rectangle(cls, client, x, y, width, height, data):
decompressed = client.zlib_decompressor.decompress(data)
logger.debug('[zrle] Decompressed from %s bytes -> %s bytes',
len(data), len(decompressed))
pyprofile.incr(
'vncdriver.recv_rectangle.zrle_encoding.decompressed_bytes',
len(decompressed),
unit=pyprofile.BYTES)
if client.framebuffer.bpp > 24:
bytes_per_pixel = 3
else:
bytes_per_pixel = client.framebuffer.bypp
buf = BytesIO(decompressed)
data = cls._read(x, y, width, height, buf, bytes_per_pixel)
encoding = cls(data)
return Rectangle(x, y, width, height, encoding)
def onMessage(self, payload, isBinary):
if not self.factory.agent_conn.check_message(self):
return
assert not isBinary, "Binary websocket not supported"
payload = ujson.loads(payload)
context = {
'start': time.time(),
'conn': self,
}
latency = context['start'] - payload['headers']['sent_at']
pyprofile.incr('rewarder_protocol.messages')
pyprofile.incr('rewarder_protocol.messages.{}'.format(
payload['method']))
pyprofile.timing('rewarder_protocol.latency.rtt.skew_unadjusted',
2 * latency)
if latency < 0:
pyprofile.incr(
'rewarder_protocol.latency.rtt.skew_unadjusted.negative')
if payload['method'] == 'v0.env.reset':
logger.info('Received reset message: %s', payload)
self.factory.agent_conn.control_buffer.recv_rpc(context, payload)
elif payload['method'] == 'v0.control.ping':
logger.debug('Received ping message: %s', payload)
parent_message_id = payload['headers']['message_id']
headers = {'parent_message_id': parent_message_id}
self.send_message('v0.reply.control.ping', {}, headers)
else:
logger.warn('Received unsupported message: %s', payload)
def onMessage(self, payload, isBinary):
extra_logger.debug('[%s] Received payload: %s', self.factory.label,
payload)
assert not isBinary
payload = ujson.loads(payload)
context = self._make_context()
latency = context['start'] - payload['headers']['sent_at']
pyprofile.incr('rewarder_protocol.messages')
pyprofile.incr('rewarder_protocol.messages.{}'.format(
payload['method']))
# Double latency to model RTT
pyprofile.timing('rewarder_protocol.latency.rtt.skew_unadjusted',
2 * latency)
if latency < 0:
pyprofile.incr(
'rewarder_protocol.latency.rtt.skew_unadjusted.negative')
self.recv(context, payload)
def recv(self, context, response):
method = response['method']
body = response['body']
headers = response['headers']
remote_time = headers['sent_at']
local_time = context['start']
episode_id = headers.get('episode_id')
if episode_id is not None:
self.reward_buffer.push_time(episode_id, remote_time, local_time)
# Gets called by RewarderClient
if method == 'v0.env.reward':
episode_id = headers['episode_id']
reward = body['reward']
done = body['done']
info = body['info']
extra_logger.debug(
'[%s] Received %s: reward=%s done=%s info=%s episode_id=%s',
self.factory.label, method, reward, done, info, episode_id)
pyprofile.incr('rewarder_client.reward', reward)
if done:
pyprofile.incr('rewarder_client.done')
self.reward_buffer.push(episode_id, reward, done, info)
elif method == 'v0.env.text':
episode_id = headers['episode_id']
text = body['text']
extra_logger.debug('[%s] Received %s: text=%s episode_id=%s',
self.factory.label, method, text, episode_id)
self.reward_buffer.push_text(episode_id, text)
elif method == 'v0.env.observation':
episode_id = headers['episode_id']
jsonable = body['observation']
extra_logger.debug(
'[%s] Received %s: observation=%s episode_id=%s',
self.factory.label, method, jsonable, episode_id)
self.reward_buffer.set_observation(episode_id=episode_id,
observation=jsonable)
elif method == 'v0.env.describe':
episode_id = headers['episode_id']
env_id = body['env_id']
env_state = body['env_state']
fps = body['fps']
extra_logger.info(
'[%s] Received %s: env_id=%s env_state=%s episode_id=%s',
self.factory.label, method, env_id, env_state, episode_id)
self.reward_buffer.set_env_info(env_state,
env_id=env_id,
episode_id=episode_id,
fps=fps)
elif method == 'v0.reply.env.reset':
episode_id = headers['episode_id']
self._finish_reset(episode_id)
elif method in ['v0.reply.error', 'v0.reply.control.ping']:
assert headers.get('parent_message_id') is not None
elif method == 'v0.connection.close':
assert headers.get('parent_message_id') is None
logger.debug('Server hanging up: %s', body['message'])
self._close_message = body['message']
e = error.Error(body['message'])
self.factory.record_error(e)
else:
logger.error(
'Unrecognized websocket method: method=%s body=%s headers=%s (consider adding to rewarder_state.py)',
method, body, headers)
return
parent_id = headers.get('parent_message_id')
if parent_id is not None:
try:
spec = self._requests.pop(parent_id)
except KeyError:
logger.error(
'[%s] Received extra reply to %d; ignoring: method=%s body=%s headers=%s ',
self.factory.label, parent_id, method, body, headers)
else:
request, d = spec
if method != 'v0.reply.error':
d.callback((context, request, response))
else:
e = RemoteError('[{}] Remote error: {}'.format(
self.factory.label, body['message']))
d.errback(e)
