def write_json(self, msg, encoding=None):
encoding = encoding or Encoding()
logger.debug("writing json: %r", msg)
try:
return self.write_message(serialize(msg, encoding), binary=True)
except: # noqa
logger.exception('Failed to write: %s', msg)
async def write_json(self, msg, encoding=None):
encoding = encoding or Encoding()
logger.debug("writing json: %r", msg)
try:
return await self.send(serialize(msg, encoding))
except: # noqa
logger.exception('Failed to write: %s', msg)
async def execute_async(self, df_name, spec_data):
encoding = Encoding()
specs = deserialize(spec_data, encoding)
df = self.service.df_map[df_name]
tasks = encoding.decode_list('task', specs, df=df)
results = self.service.execute(df, tasks)
return results
async def handle_message(self, websocket_msg):
try:
await self._handle_message(websocket_msg)
except BaseException as e:
encoding = Encoding()
websocket_msg = deserialize(websocket_msg, encoding)
msg_id = websocket_msg['msg_id']
msg = exception(e)
await self.write_json({'msg_id': msg_id, 'msg': msg})
logger.exception("unhandled exception")
def execute(self, df, tasks):
from vaex.encoding import Encoding
encoder = Encoding()
msg_id = str(uuid.uuid4())
self._msg_id_to_tasks[msg_id] = tuple(tasks)
task_specs = encoder.encode_list("task", tasks)
msg = {
'command': 'execute',
'df': df.name,
'state': df.state_get(),
'tasks': task_specs
}
try:
results, encoding = self._send(msg, msg_id=msg_id)
results = encoding.decode_list('vaex-task-result', results)
return results
finally:
del self._msg_id_to_tasks[msg_id]
async def _handle_message(self, websocket_msg):
if TEST_LATENCY:
await asyncio.sleep(TEST_LATENCY)
msg_id = 'invalid'
encoding = Encoding()
try:
websocket_msg = deserialize(websocket_msg, encoding)
logger.debug("websocket message: %s", websocket_msg)
msg_id, msg, auth = websocket_msg['msg_id'], websocket_msg[
'msg'], websocket_msg['auth']
token = auth['token'] # are we even allowed to execute?
token_trusted = auth[
'token-trusted'] # do we trust arbitrary code execution?
trusted = token_trusted == self.token_trusted and token_trusted
if not (
(token == self.token) or
(self.token_trusted and token_trusted == self.token_trusted)):
raise ValueError('No token provided, not authorized')
last_progress = None
ioloop = asyncio.get_event_loop()
progress_futures = []
def progress(f):
nonlocal last_progress
async def send_progress():
vaex.asyncio.check_patch_tornado(
) # during testing asyncio might be patched
nonlocal last_progress
logger.debug("progress: %r", f)
last_progress = f
# TODO: create task?
return await self.write_json({
'msg_id': msg_id,
'msg': {
'progress': f
}
})
# emit when it's the first time (None), at least 0.05 sec lasted, or and the end
# but never send old or same values
if (last_progress is None or (f - last_progress) > 0.05
or f == 1.0) and (last_progress is None
or f > last_progress):
def wrapper():
progress_futures.append(
asyncio.create_task(send_progress()))
ioloop.call_soon_threadsafe(wrapper)
return True
command = msg['command']
if command == 'list':
result = self.service.list()
await self.write_json({
'msg_id': msg_id,
'msg': {
'result': result
}
})
elif command == 'versions':
result = {
'vaex.core': vaex.core._version.__version_tuple__,
'vaex.server': vaex.server._version.__version_tuple__
}
await self.write_json({
'msg_id': msg_id,
'msg': {
'result': result
}
})
elif command == 'execute':
df = self.service[msg['df']].copy()
df.state_set(msg['state'],
use_active_range=True,
trusted=trusted)
tasks = encoding.decode_list('task', msg['tasks'], df=df)
self._msg_id_to_tasks[
msg_id] = tasks # keep a reference for cancelling
try:
# TODO: this assumes all tasks succeed, but we also support 1 failing
results = await self.service.execute(df,
tasks,
progress=progress)
finally:
del self._msg_id_to_tasks[msg_id]
await asyncio.gather(*progress_futures)
# make sure the final progress value is send, and also old values are not send
last_progress = 1.0
await self.write_json({
'msg_id': msg_id,
'msg': {
'progress': 1.0
}
})
encoding = Encoding()
results = encoding.encode_list('vaex-task-result', results)
await self.write_json(
{
'msg_id': msg_id,
'msg': {
'result': results
}
}, encoding)
elif command == 'cancel':
try:
tasks = self._msg_id_to_tasks[msg['cancel_msg_id']]
except KeyError:
pass # already done, or cancelled
else:
for task in tasks:
task.cancel()
elif command == 'call-dataframe':
df = self.service[msg['df']].copy()
df.state_set(msg['state'],
use_active_range=True,
trusted=trusted)
# TODO: yield
if msg['method'] not in vaex.server.dataframe.allowed_method_names:
raise NotImplementedError("Method is not rmi invokable")
results = self.service._rmi(df, msg['method'], msg['args'],
msg['kwargs'])
encoding = Encoding()
if msg['method'] == "_evaluate_implementation":
results = encoding.encode('vaex-evaluate-result', results)
else:
results = encoding.encode('vaex-rmi-result', results)
await self.write_json(
{
'msg_id': msg_id,
'msg': {
'result': results
}
}, encoding)
else:
raise ValueError(f'Unknown command: {command}')
except Exception as e:
logger.exception("Exception while handling msg")
msg = exception(e)
await self.write_json({'msg_id': msg_id, 'msg': msg})
def _on_message(self, websocket_msg):
msg_id = 'invalid'
encoding = Encoding()
try:
websocket_msg = deserialize(websocket_msg, encoding)
logger.debug("websocket message: %s", websocket_msg)
msg_id, msg, auth = websocket_msg['msg_id'], websocket_msg[
'msg'], websocket_msg['auth']
token = auth['token'] # are we even allowed to execute?
token_trusted = auth[
'token-trusted'] # do we trust arbitrary code execution?
trusted = token_trusted == self.webserver.token_trusted and token_trusted
if not ((token == self.webserver.token) or
(self.webserver.token_trusted
and token_trusted == self.webserver.token_trusted)):
raise ValueError('No token provided, not authorized')
last_progress = None
def progress(f):
nonlocal last_progress
def send_progress():
nonlocal last_progress
logger.debug("progress: %r", f)
last_progress = f
return self.write_json({
'msg_id': msg_id,
'msg': {
'progress': f
}
})
# emit when it's the first time (None), at least 0.05 sec lasted, or and the end
# but never send old or same values
if (last_progress is None or (f - last_progress) > 0.05
or f == 1.0) and (last_progress is None
or f > last_progress):
self.webserver.ioloop.add_callback(send_progress)
return True
command = msg['command']
if command == 'list':
result = self.service.list()
self.write_json({'msg_id': msg_id, 'msg': {'result': result}})
elif command == 'versions':
result = {
'vaex.core': vaex.core._version.__version_tuple__,
'vaex.server': vaex.server._version.__version_tuple__
}
self.write_json({'msg_id': msg_id, 'msg': {'result': result}})
elif command == 'execute':
df = self.service[msg['df']].copy()
df.state_set(msg['state'],
use_active_range=True,
trusted=trusted)
tasks = encoding.decode_list('task', msg['tasks'], df=df)
self._msg_id_to_tasks[
msg_id] = tasks # keep a reference for cancelling
try:
results = yield self.service.execute(df,
tasks,
progress=progress)
finally:
del self._msg_id_to_tasks[msg_id]
# make sure the final progress value is send, and also old values are not send
last_progress = 1.0
self.write_json({'msg_id': msg_id, 'msg': {'progress': 1.0}})
encoding = Encoding()
results = encoding.encode_list('vaex-task-result', results)
self.write_json({
'msg_id': msg_id,
'msg': {
'result': results
}
}, encoding)
elif command == 'cancel':
try:
tasks = self._msg_id_to_tasks[msg['cancel_msg_id']]
except KeyError:
pass # already done, or cancelled
else:
for task in tasks:
task.cancel()
elif command == 'call-dataframe':
df = self.service[msg['df']].copy()
df.state_set(msg['state'],
use_active_range=True,
trusted=trusted)
# TODO: yield
if msg['method'] not in vaex.server.dataframe.allowed_method_names:
raise NotImplementedError("Method is not rmi invokable")
results = self.service._rmi(df, msg['method'], msg['args'],
msg['kwargs'])
encoding = Encoding()
if msg['method'] == "_evaluate_implementation":
results = encoding.encode('vaex-evaluate-result', results)
else:
results = encoding.encode('vaex-rmi-result', results)
self.write_json({
'msg_id': msg_id,
'msg': {
'result': results
}
}, encoding)
else:
raise ValueError(f'Unknown command: {command}')
except Exception as e:
logger.exception("Exception while handling msg")
msg = exception(e)
self.write_json({'msg_id': msg_id, 'msg': msg})
def write_json(self, msg, encoding=None):
encoding = encoding or Encoding()
logger.debug("writing json: %r", msg)
return self.write_message(serialize(msg, encoding), binary=True)
def execute(self, df, tasks):
encoding = Encoding()
task_specs = encoding.encode_list("task", tasks)
return self.server.execute(df.name, serialize(task_specs, encoding))