add Wechat Work support (#47)

* add Wechat Work support add Wechat Work chatroom robot notification support. Wechat Work is one of the most popular business communication and office collaboration tool, developed by Tencent. * updates * uniformization + bug * uniformization * uniformization * uniformization * tested * Update README.md * Update __main__.py Co-authored-by: Victor SANH <victorsanh@gmail.com>
2021-08-28 00:30:42 +03:00 · 2020-03-16 12:26:47 +08:00
parent b59a7be6e8
commit 1160ab756d
4 changed files with 174 additions and 2 deletions
--- 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 *eleven* ways to setup notifications:
+There are currently *twelve* ways to setup notifications:

 |               Platform                |                                   External Contributors                                   |
 | :-----------------------------------: | :---------------------------------------------------------------------------------------: |
@@ -35,6 +35,7 @@ There are currently *eleven* ways to setup notifications:
 |     [Amazon Chime](#amazon-chime)     |                     [@prabhakar267](https://github.com/prabhakar267)                      |
 |         [DingTalk](#dingtalk)         |                         [@wuutiing](https://github.com/wuutiing)                          |
 |       [RocketChat](#rocketchat)       |                            [@radao](https://github.com/radao)                             |
+|      [WeChat Work](#wechat-work)      |                             [@jcyk](https://github.com/jcyk)                              |


 ### Email
@@ -66,6 +67,7 @@ If `sender_email` is not specified, then the first email in `recipient_emails` w

 Note that launching this will asks you for the sender's email password. It will be safely stored in the system keyring service through the [`keyring` Python library](https://pypi.org/project/keyring/).

+
 ### Slack

 Similarly, you can also use Slack to get notifications. You'll have to get your Slack room [webhook URL](https://api.slack.com/incoming-webhooks#create_a_webhook) and optionally your [user id](https://api.slack.com/methods/users.identity) (if you want to tag yourself or someone else).
@@ -96,6 +98,7 @@ knockknock slack \

 You can also specify an optional argument to tag specific people: `--user-mentions <your_slack_id>,<grandma's_slack_id>`.

+
 ### Telegram

 You can also use Telegram Messenger to get notifications. You'll first have to create your own notification bot by following the three steps provided by Telegram [here](https://core.telegram.org/bots#6-botfather) and save your API access `TOKEN`.
@@ -124,6 +127,7 @@ knockknock telegram \
    sleep 10
 ```

+
 ### Microsoft Teams

 Thanks to [@noklam](https://github.com/noklam), you can also use Microsoft Teams to get notifications. You'll have to get your Team Channel [webhook URL](https://docs.microsoft.com/en-us/microsoftteams/platform/concepts/connectors/connectors-using).
@@ -150,6 +154,7 @@ knockknock teams \

 You can also specify an optional argument to tag specific people: `user_mentions=[<your_teams_id>, <grandma's_teams_id>]`.

+
 ### Text Message (SMS)

 Thanks to [@abhishekkrthakur](https://github.com/abhishekkrthakur), you can use Twilio to send text message notifications. You'll have to setup a [Twilio](www.twilio.com) account [here](https://www.twilio.com/try-twilio), which is paid service with competitive prices: for instance in the US, getting a new number and sending one text message through this service respectively cost $1.00 and $0.0075. You'll need to get (a) a phone number, (b) your [account SID](https://www.twilio.com/docs/glossary/what-is-a-sid) and (c) your [authentification token](https://www.twilio.com/docs/iam/access-tokens). Some detail [here](https://www.twilio.com/docs/iam/api/account).
@@ -342,6 +347,10 @@ from knockknock import rocketchat_sender
    rocketchat_user_id="<your_rocketchat_user_id>",
    rocketchat_auth_token="<your_rocketchat_auth_token>",
    channel="<channel_name>")
+def train_your_nicest_model(your_nicest_parameters):
+    import time
+    time.sleep(10000)
+    return {'loss': 0.9} # Optional return value
 ```

 You can also specify two optional arguments:
@@ -361,6 +370,34 @@ knockknock rocketchat \
 ```


+### WeChat Work
+
+WeChat Work is now supported thanks to [@jcyk](https://github.com/jcyk). Given WeChat Work chatroom robot's webhook url, your notifications will be sent to reach anyone in that chatroom.
+
+#### Python
+
+```python
+from knockknock import wechat_sender
+
+webhook_url = "<webhook_url_to_your_wechat_work_chatroom_robot>"
+@wechat_sender(webhook_url=webhook_url)
+def train_your_nicest_model(your_nicest_parameters):
+    import time
+    time.sleep(10000)
+    return {'loss': 0.9} # Optional return value
+```
+
+#### Command-line
+
+```bash
+knockknock wechat \
+    --webhook-url <webhook_url_to_your_wechat_work_chatroom_robot> \
+    sleep 10
+```
+
+You can also specify an optional argument to tag specific people: `user-mentions=["<list_of_userids_you_want_to_tag>"]` and/or `user-mentions-mobile=["<list_of_phonenumbers_you_want_to_tag>"]`.
+
+
 ## 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.
--- a/knockknock/init.py
+++ b/knockknock/init.py
@@ -8,4 +8,5 @@ from knockknock.teams_sender import teams_sender
 from knockknock.desktop_sender import desktop_sender
 from knockknock.matrix_sender import matrix_sender
 from knockknock.dingtalk_sender import dingtalk_sender
+from knockknock.wechat_sender import wechat_sender
 from knockknock.rocketchat_sender import rocketchat_sender
--- a/knockknock/main.py
+++ b/knockknock/main.py
@@ -11,7 +11,8 @@ from knockknock import (chime_sender,
                        slack_sender,
                        sms_sender,
                        teams_sender,
-                        telegram_sender,)
+                        telegram_sender,
+                        wechat_sender,)

 def main():
    parser = argparse.ArgumentParser(
@@ -173,6 +174,21 @@ def main():
        "--alias", type=str, required=False, default="", help="Optional alias to use for the notification.")
    rocketchat_parser.set_defaults(sender_func=rocketchat_sender)

+    # WeChat Work
+    wechat_parser = subparsers.add_parser(
+        name="wechat", description="Send a WeChat Work message before and after function " +
+        "execution, with start and end status (sucessfully or crashed).")
+    wechat_parser.add_argument(
+        "--webhook-url", type=str, required=True,
+        help="The webhook URL to access your wechat_work chatroom")
+    wechat_parser.add_argument(
+        "--user-mentions", type=lambda s: s.split(","), required=False, default=[],
+        help="Optional userids to notify (use '@all' for all group members), as comma seperated list.")
+    wechat_parser.add_argument(
+        "--user-mentions-mobile", type=lambda s: s.split(","), required=False, default=[],
+        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)
+
    args, remaining_args = parser.parse_known_args()
    args = vars(args)

--- a/knockknock/wechat_sender.py
+++ b/knockknock/wechat_sender.py
@@ -0,0 +1,118 @@
+from typing import List
+import os
+import datetime
+import traceback
+import functools
+import json
+import socket
+import requests
+import time
+
+DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+def wechat_sender(webhook_url: str,
+                  user_mentions: List[str] = [],
+                  user_mentions_mobile: List[str] = []):
+    """
+    WeChat Work sender wrapper: execute func, send a WeChat Work notification with the end status
+    (sucessfully finished or crashed) at the end. Also send a WeChat Work notification before
+    executing func. To obtain the webhook, add a Group Robot in your WeChat Work Group. Visit
+    https://work.weixin.qq.com/api/doc/90000/90136/91770 for more details.
+
+    `webhook_url`: str
+        The webhook URL to access your WeChat Work chatroom.
+        Visit https://work.weixin.qq.com/api/doc/90000/90136/91770 for more details.
+    `user_mentions`: List[str] (default=[])
+        Optional userids to notify (use '@all' for all group members).
+        Visit https://work.weixin.qq.com/api/doc/90000/90136/91770 for more details.
+    `user_mentions_mobile`: List[str] (default=[])
+        Optional user's phone numbers to notify (use '@all' for all group members).
+        Visit https://work.weixin.qq.com/api/doc/90000/90136/91770 for more details.
+    """
+    
+    msg_template = {
+        "msgtype": "text", 
+        "text": {
+            "content": "",
+            "mentioned_list":user_mentions,
+            "mentioned_mobile_list":user_mentions_mobile
+        }
+    }
+
+    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)]
+                
+                msg_template['text']['content'] = '\n'.join(contents)
+                requests.post(webhook_url, json=msg_template)
+
+            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.")
+
+                    
+                    msg_template['text']['content'] = '\n'.join(contents)
+                    requests.post(webhook_url, json=msg_template)
+                    print(msg_template)
+
+                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()]
+                
+                msg_template['text']['content'] = '\n'.join(contents)
+                requests.post(webhook_url, json=msg_template)
+                print(msg_template)
+
+                raise ex
+
+        return wrapper_sender
+
+    return decorator_sender