add feishu support

README.md

@@ -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)                              |
+|           [Feishu](#feishu)           |                             [@suxnju](https://github.com/suxnju)                           |
 
 
 ### Email
@@ -397,6 +398,28 @@ knockknock wechat \
 
 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>"]`.
 
+### Feishu
+
+#### Python
+
+```python
+from knockknock import feishu_sender
+
+webhook_url = "<webhook_url_to_your_feishu_chatroom_robot>"
+@dingtalk_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 feishu \
+    --webhook-url <webhook_url_to_your_feishu_chatroom_robot> \
+    sleep 10
+```
 
 ## Note on distributed training
 

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.feishu_sender import feishu_sender

knockknock/__main__.py

@@ -12,7 +12,8 @@
                         sms_sender,
                         teams_sender,
                         telegram_sender,
-                        wechat_sender,)
+                        wechat_sender,
+                        feishu_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)
 
+    # WeChat Work
+    feishu_parser = subparsers.add_parser(
+        name="feishu", description="Send a Feishu message before and after function " +
+        "execution, with start and end status (sucessfully or crashed).")
+    feishu_parser.add_argument(
+        "--webhook-url", type=str, required=True,
+        help="The webhook URL to access your Feishu chatroom")
+    feishu_parser.set_defaults(sender_func=feishu_sender)
+
     args, remaining_args = parser.parse_known_args()
     args = vars(args)
 

knockknock/feishu_sender.py

@@ -0,0 +1,122 @@
+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 feishu_sender(webhook_url: str):
+    """
+    Feishu sender wrapper: execute func, send a Feishu message with the end status
+    (sucessfully finished or crashed) at the end. Also send a Feishu message before
+    executing func. To obtain the webhook, add a Group Robot in your Feishu Group. Visit
+    https://open.feishu.cn/document/ukTMukTMukTM/ucTM5YjL3ETO24yNxkjN for more details.
+
+    `webhook_url`: str
+        The webhook URL to access your Feishu chatroom.
+        Visit https://open.feishu.cn/document/ukTMukTMukTM/ucTM5YjL3ETO24yNxkjN for more details.
+    """
+
+    msg_template = {
+        "msg_type": "text",
+        "content" : {"text":""},
+        "timestamp" : "",
+    }
+
+    def decorator_sender(func):
+        @functools.wraps(func)
+        def wrapper_sender(*args, **kwargs):
+
+            timestamp = int(datetime.datetime.now().timestamp())
+            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['content']['text'] = '\n'.join(contents)
+                msg_template['timestamp'] = timestamp
+                resp = requests.post(webhook_url, json=msg_template)
+                resp.raise_for_status()
+                result = resp.json()
+                if result.get("code") and result.get("code") != 0:
+                    print(f"error：{result['msg']}")
+                    return
+            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['content']['text'] = '\n'.join(contents)
+                    msg_template['timestamp'] = timestamp
+                    resp = requests.post(webhook_url, json=msg_template)
+                    resp.raise_for_status()
+                    result = resp.json()
+                    if result.get("code") and result.get("code") != 0:
+                        print(f"error：{result['msg']}")
+                        return
+
+                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['content']['text'] = '\n'.join(contents)
+                msg_template['timestamp'] = timestamp
+                resp = requests.post(webhook_url, json=msg_template)
+                resp.raise_for_status()
+                result = resp.json()
+                if result.get("code") and result.get("code") != 0:
+                    print(f"error：{result['msg']}")
+                    return
+
+                raise ex
+
+        return wrapper_sender
+
+    return decorator_sender