diff --git a/README.md b/README.md index 660ce97..d6fb296 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ This code has only been tested with Python >= 3.6. The library is designed to be used in a seamless way, with minimal code modification: you only need to add a decorator on top your main function call. The return value (if there is one) is also reported in the notification. -There are currently *twelve* ways to setup notifications: +There are currently *thirteen* ways to setup notifications: | Platform | External Contributors | | :-----------------------------------: | :---------------------------------------------------------------------------------------: | @@ -36,6 +36,7 @@ There are currently *twelve* ways to setup notifications: | [DingTalk](#dingtalk) | [@wuutiing](https://github.com/wuutiing) | | [RocketChat](#rocketchat) | [@radao](https://github.com/radao) | | [WeChat Work](#wechat-work) | [@jcyk](https://github.com/jcyk) | +| [LINE](#line) | [@phiradet](https://github.com/phiradet) | ### Email @@ -398,6 +399,31 @@ knockknock wechat \ You can also specify an optional argument to tag specific people: `user-mentions=[""]` and/or `user-mentions-mobile=[""]`. +### LINE +LINE is now supported thanks to [@phiradet](https://github.com/phiradet). You will first have to generate your LINE Notify personal access token from [My page](https://notify-bot.line.me/my/). Please follow the steps explained in the "Generating personal access tokens" section of this [LINE Engineering blog](https://engineering.linecorp.com/en/blog/using-line-notify-to-send-messages-to-line-from-the-command-line/). Your notifications will be sent to you or your selected group. + +#### Python + +```python +from knockknock import line_sender + +token = "" + +@line_sender(token=token) +def train_your_nicest_model(your_nicest_parameters): + import time + time.sleep(10000) + return {'loss': 0.9} # Optional return value +``` + +#### Command-line + +```bash +knockknock line \ + --token \ + sleep 10 +``` + ## Note on distributed training When using distributed training, a GPU is bound to its process using the local rank variable. Since knockknock works at the process level, if you are using 8 GPUs, you would get 8 notifications at the beginning and 8 notifications at the end... To circumvent that, except for errors, only the master process is allowed to send notifications so that you receive only one notification at the beginning and one notification at the end. diff --git a/knockknock/__init__.py b/knockknock/__init__.py index a5c7f51..9995843 100644 --- a/knockknock/__init__.py +++ b/knockknock/__init__.py @@ -10,3 +10,4 @@ from knockknock.dingtalk_sender import dingtalk_sender from knockknock.wechat_sender import wechat_sender from knockknock.rocketchat_sender import rocketchat_sender +from knockknock.line_sender import line_sender diff --git a/knockknock/__main__.py b/knockknock/__main__.py index 215939d..83b30ce 100644 --- a/knockknock/__main__.py +++ b/knockknock/__main__.py @@ -12,7 +12,8 @@ sms_sender, teams_sender, telegram_sender, - wechat_sender,) + wechat_sender, + line_sender,) def main(): parser = argparse.ArgumentParser( @@ -189,6 +190,15 @@ def main(): help="Optional user phone numbers to notify (use '@all' for all group members), as comma seperated list.") wechat_parser.set_defaults(sender_func=wechat_sender) + # LINE + line_parser = subparsers.add_parser( + name="line", description="Send a LINE message before and after function " + + "execution, with start and end status (successfully or crashed).") + line_parser.add_argument( + "--token", type=str, required=True, + help="The personal access token required to use the LINE Notify API.") + line_parser.set_defaults(sender_func=line_sender) + args, remaining_args = parser.parse_known_args() args = vars(args) diff --git a/knockknock/line_sender.py b/knockknock/line_sender.py new file mode 100644 index 0000000..73c7ab6 --- /dev/null +++ b/knockknock/line_sender.py @@ -0,0 +1,103 @@ +import os +import datetime +import traceback +import functools +import socket +import requests + + +DATE_FORMAT = "%Y-%m-%d %H:%M:%S" +NOTIFY_API_URL = "https://notify-api.line.me/api/notify" + + +def line_sender(token: str): + """ + LINE sender wrapper: execute func, send a LINE notification with the end status + (sucessfully finished or crashed) at the end. Also send a Line notification before + executing func. + + `token`: str + The personal access token required to use the LINE Notify API + Visit https://engineering.linecorp.com/en/blog/using-line-notify-to-send-messages-to-line-from-the-command-line/ + for more details. + """ + headers = {"Authorization": f"Bearer {token}", + "Content-Type": "application/x-www-form-urlencoded"} + dump = {} + + def decorator_sender(func): + @functools.wraps(func) + def wrapper_sender(*args, **kwargs): + + start_time = datetime.datetime.now() + host_name = socket.gethostname() + func_name = func.__name__ + + # Handling distributed training edge case. + # In PyTorch, the launch of `torch.distributed.launch` sets up a RANK environment variable for each process. + # This can be used to detect the master process. + # See https://github.com/pytorch/pytorch/blob/master/torch/distributed/launch.py#L211 + # Except for errors, only the master process will send notifications. + if 'RANK' in os.environ: + master_process = (int(os.environ['RANK']) == 0) + host_name += ' - RANK: %s' % os.environ['RANK'] + else: + 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) + ] + dump['message'] = '\n'.join(contents) + requests.post(url=NOTIFY_API_URL, headers=headers, params=dump) + + 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.") + + dump['message'] = '\n'.join(contents) + requests.post(url=NOTIFY_API_URL, headers=headers, params=dump) + + 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(), + ] + dump['message'] = '\n'.join(contents) + requests.post(url=NOTIFY_API_URL, headers=headers, params=dump) + raise ex + + return wrapper_sender + + return decorator_sender