diff --git a/README.md b/README.md index 2273e28..599a4fe 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ Similarly, you can also use Slack to get notifications. You'll have to get your from knockknock import slack_sender webhook_url = "" -@slack_sender(webhook_url=webhook_url, channel="") +@slack_sender(webhook_url=webhook_url) def train_your_nicest_model(your_nicest_parameters): import time time.sleep(10000) diff --git a/knockknock/slack_sender.py b/knockknock/slack_sender.py index 7ca63f3..763c5d3 100644 --- a/knockknock/slack_sender.py +++ b/knockknock/slack_sender.py @@ -9,7 +9,8 @@ DATE_FORMAT = "%Y-%m-%d %H:%M:%S" -def slack_sender(webhook_url: str, channel: str, user_mentions: List[str] = []): + +def slack_sender(webhook_url: str, user_mentions: List[str] = []): """ Slack sender wrapper: execute func, send a Slack notification with the end status (sucessfully finished or crashed) at the end. Also send a Slack notification before @@ -18,18 +19,14 @@ def slack_sender(webhook_url: str, channel: str, user_mentions: List[str] = []): `webhook_url`: str The webhook URL to access your slack room. Visit https://api.slack.com/incoming-webhooks#create_a_webhook for more details. - `channel`: str - The slack room to log. `user_mentions`: List[str] (default=[]) - Optional users ids to notify. + Optional user ids to notify. Visit https://api.slack.com/methods/users.identity for more details. """ - dump = { - "username": "Knock Knock", - "channel": channel, - "icon_emoji": ":clapper:", - } + if user_mentions: + user_mentions = ["<@{}>".format(user) for user in user_mentions] + def decorator_sender(func): @functools.wraps(func) def wrapper_sender(*args, **kwargs): @@ -50,60 +47,179 @@ def wrapper_sender(*args, **kwargs): master_process = True if master_process: - contents = ['Your training has started 🎬', - 'Machine name: %s' % host_name, - 'Main call: %s' % func_name, - 'Starting date: %s' % start_time.strftime(DATE_FORMAT)] - contents.append(' '.join(user_mentions)) - dump['text'] = '\n'.join(contents) - dump['icon_emoji'] = ':clapper:' - requests.post(webhook_url, json.dumps(dump)) + notification = "Your training has started! 🎬" + if user_mentions: + notification = _add_mentions(notification) + + payload = { + "blocks": _starting_message( + func_name, host_name, notification, start_time + ), + "text": notification, + } + + requests.post(webhook_url, json.dumps(payload)) try: value = func(*args, **kwargs) if master_process: - end_time = datetime.datetime.now() - elapsed_time = end_time - start_time - contents = ["Your training is complete 🎉", - 'Machine name: %s' % host_name, - 'Main call: %s' % func_name, - 'Starting date: %s' % start_time.strftime(DATE_FORMAT), - 'End date: %s' % end_time.strftime(DATE_FORMAT), - 'Training duration: %s' % str(elapsed_time)] - - try: - str_value = str(value) - contents.append('Main call returned value: %s'% str_value) - except: - contents.append('Main call returned value: %s'% "ERROR - Couldn't str the returned value.") - - contents.append(' '.join(user_mentions)) - dump['text'] = '\n'.join(contents) - dump['icon_emoji'] = ':tada:' - requests.post(webhook_url, json.dumps(dump)) + notification = "Your training is complete 🎉" + if user_mentions: + notification = _add_mentions(notification) + + payload = { + "blocks": _successful_message( + func_name, host_name, notification, start_time, value + ), + "text": notification, + } + requests.post(webhook_url, json.dumps(payload)) return value except Exception as ex: - end_time = datetime.datetime.now() - elapsed_time = end_time - start_time - contents = ["Your training has crashed ☠️", - 'Machine name: %s' % host_name, - 'Main call: %s' % func_name, - 'Starting date: %s' % start_time.strftime(DATE_FORMAT), - 'Crash date: %s' % end_time.strftime(DATE_FORMAT), - 'Crashed training duration: %s\n\n' % str(elapsed_time), - "Here's the error:", - '%s\n\n' % ex, - "Traceback:", - '%s' % traceback.format_exc()] - contents.append(' '.join(user_mentions)) - dump['text'] = '\n'.join(contents) - dump['icon_emoji'] = ':skull_and_crossbones:' - requests.post(webhook_url, json.dumps(dump)) + notification = "Your training has crashed ☠️" + if user_mentions: + notification = _add_mentions(notification) + + payload = { + "blocks": _error_message( + ex, func_name, host_name, notification, start_time + ), + "text": notification, + } + requests.post(webhook_url, json.dumps(payload)) raise ex + def _error_message(ex, func_name, host_name, notification, start_time): + """Uses Slack blocks to create a formatted report of exception 'ex'.""" + end_time = datetime.datetime.now() + training_time = _format_train_time(end_time, start_time) + return [ + {"type": "section", "text": {"type": "mrkdwn", "text": notification}}, + {"type": "divider"}, + { + "type": "context", + "elements": [ + { + "type": "mrkdwn", + "text": "*Machine name:* {}\n" + "*Main call:* `{}`\n" + "*Starting date:* {}\n" + "*Crash date:* {}\n" + "*Time elapsed before crash:* {}".format( + host_name, + func_name, + start_time.strftime(DATE_FORMAT), + end_time.strftime(DATE_FORMAT), + training_time, + ), + } + ], + }, + {"type": "divider"}, + { + "type": "section", + "text": {"type": "mrkdwn", "text": "*Error:* `{}`".format(ex)}, + }, + { + "type": "section", + "text": { + "type": "mrkdwn", + "text": "*Traceback:*\n```{}```".format(traceback.format_exc()), + }, + }, + ] + + def _starting_message(func_name, host_name, notification, start_time): + """Uses Slack blocks to create an initial report of training.""" + return [ + {"type": "section", "text": {"type": "mrkdwn", "text": notification}}, + {"type": "divider"}, + { + "type": "context", + "elements": [ + { + "type": "mrkdwn", + "text": "*Machine name:* {}\n" + "*Main call:* `{}`\n" + "*Starting date:* {}\n".format( + host_name, func_name, start_time.strftime(DATE_FORMAT) + ), + } + ], + }, + ] + + def _successful_message(func_name, host_name, notification, start_time, value): + """Uses Slack blocks to report a successful training run with statistics.""" + end_time = datetime.datetime.now() + training_time = _format_train_time(end_time, start_time) + blocks = [ + {"type": "section", "text": {"type": "mrkdwn", "text": notification}}, + {"type": "divider"}, + { + "type": "context", + "elements": [ + { + "type": "mrkdwn", + "text": "*Machine name:* {}\n" + "*Main call:* `{}`\n" + "*Starting date:* {}\n" + "*End date:* {}\n" + "*Training Duration:* {}".format( + host_name, + func_name, + start_time.strftime(DATE_FORMAT), + end_time.strftime(DATE_FORMAT), + training_time, + ), + } + ], + }, + ] + + if value is not None: + blocks.append({"type": "divider"}) + try: + str_value = str(value) + blocks.append( + { + "type": "section", + "text": { + "type": "mrkdwn", + "text": "*Main call returned value:* `{}`".format( + str_value + ), + }, + } + ) + except Exception as e: + blocks.append( + "Couldn't str the returned value due to the following error: \n`{}`".format( + e + ) + ) + + return blocks + + def _format_train_time(end_time, start_time): + """Returns a time delta as a string in the format '%d %H:%M:%S'""" + elapsed_time = end_time - start_time + days, remainder = divmod(elapsed_time.seconds, 86400) + hours, remainder = divmod(remainder, 24) + minutes, seconds = divmod(remainder, 60) + training_time = "{:2d}:{:02d}:{:02d}".format(hours, minutes, seconds) + + if days: + training_time = "{}d ".format(days) + training_time + return training_time + + def _add_mentions(notification): + notification = notification + " " + " ".join(user_mentions) + return notification + return wrapper_sender return decorator_sender