def on_failure(self, exc, task_id, args, kwargs, einfo):
channel_name = TASK_PUBSUB_CHANNEL.format(task_id=task_id)
failure_msg = {'error': str(exc), 'args': args, 'kwargs': kwargs}
rds.publish(channel_name,
json.dumps(failure_msg, cls=VersatileEncoder))
rds.publish(channel_name, TASK_PUBSUB_EOF.format(task_id=task_id))
msg = 'Console task {}:\nargs\n```\n{}\n```\nkwargs:\n```\n{}\n```\nerror message:\n```\n{}\n```'.format(
self.name, args, kwargs, str(exc))
im_sendmsg(IM_WEBHOOK_CHANNEL, msg)
def celery_task_stream_response(celery_task_ids,
timeout=0,
exit_when_timeout=True):
if isinstance(celery_task_ids, str):
celery_task_ids = celery_task_ids,
task_progress_channels = [
TASK_PUBSUB_CHANNEL.format(task_id=id_) for id_ in celery_task_ids
]
pubsub = rds.pubsub()
pubsub.subscribe(task_progress_channels)
try:
while pubsub.subscribed:
resp = pubsub.get_message(timeout=timeout)
if resp is None:
if exit_when_timeout:
logger.warn("pubsub timeout {}".format(celery_task_ids))
return None
continue
raw_content = resp['data']
# omit the initial message where item['data'] is 1L
if not isinstance(raw_content, (bytes, str)):
continue
content = raw_content
if isinstance(content, bytes):
content = content.decode('utf-8')
logger.debug('Got pubsub message: %s', content)
# task will publish TASK_PUBSUB_EOF at success or failure
if content.startswith('CELERY_TASK_DONE'):
finished_task_id = content[content.find(':') + 1:]
finished_task_channel = TASK_PUBSUB_CHANNEL.format(
task_id=finished_task_id)
logger.debug(
'Task %s finished, break celery_task_stream_response',
finished_task_id)
pubsub.unsubscribe(finished_task_channel)
else:
yield content
finally:
logger.debug("celery stream response exit ************")
pubsub.unsubscribe()
pubsub.close()
def on_success(self, retval, task_id, args, kwargs):
channel_name = TASK_PUBSUB_CHANNEL.format(task_id=task_id)
rds.publish(channel_name, TASK_PUBSUB_EOF.format(task_id=task_id))
def stream_output(self, data, task_id=None):
channel_name = TASK_PUBSUB_CHANNEL.format(
task_id=task_id or self.request.id)
rds.publish(channel_name, json.dumps(data, cls=VersatileEncoder))