add desktop notification (#29)

Add desktop notifications for MacOs

README.md

@@ -200,6 +200,29 @@ knockknock discord \
     sleep 10
 ```
 
+### Desktop Notification
+
+You can also get notified from a desktop notification. It is currently only available for MacOS.
+
+#### Python
+
+```python
+from knockknock import desktop_sender
+
+@desktop_sender(title="Knockknock Desktop Notifier")
+def train_your_nicest_model(your_nicest_parameters):
+    import time
+    time.sleep(10000)
+    return {"loss": 0.9}
+```
+
+#### Command Line
+```bash
+knockknock desktop \
+    --title 'Knockknock Desktop Notifier' \
+    sleep 2
+```
+
 ## 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.

knockknock/__init__.py

@@ -4,3 +4,4 @@ from knockknock.slack_sender import slack_sender
 from knockknock.sms_sender import sms_sender
 from knockknock.telegram_sender import telegram_sender
 from knockknock.teams_sender import teams_sender
+from knockknock.desktop_sender import desktop_sender

knockknock/__main__.py

@@ -1,7 +1,7 @@
 import argparse
 import subprocess
 
-from knockknock import email_sender, slack_sender, telegram_sender, teams_sender, sms_sender, discord_sender
+from knockknock import email_sender, slack_sender, telegram_sender, teams_sender, sms_sender, discord_sender, desktop_sender
 
 
 def main():
@@ -11,6 +11,14 @@ def main():
                         help="Show full command in notification.")
     subparsers = parser.add_subparsers()
 
+    ## Desktop
+    desktop_parser = subparsers.add_parser(
+        name="desktop",description="Send a desktop notification before and after function " +
+        "execution, with start and end status (successfully or crashed).")
+    desktop_parser.add_argument("--title", type=str, required=False,
+        help="The title of the notification, default to knockknock")
+    desktop_parser.set_defaults(sender_func=desktop_sender)
+    
     ## Discord
     discord_parser = subparsers.add_parser(
         name="discord", description="Send a Discord message before and after function " +

knockknock/desktop_sender.py

@@ -0,0 +1,84 @@
+import os
+import datetime
+import traceback
+import functools
+import socket
+import subprocess
+DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+def desktop_sender(title:str="knockknock"):
+    
+    def show_notification(text:str,title:str):
+        subprocess.run(["sh", "-c", "osascript -e 'display notification \"%s\" with title \"%s\"'" % (text, title)])
+
+    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)]
+                text = '\n'.join(contents)
+                show_notification(text, title)
+
+            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.")
+
+                    text = '\n'.join(contents)
+                    show_notification(text, title)
+
+                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()]
+                text = '\n'.join(contents)
+                show_notification(text, title)
+                raise ex
+
+        return wrapper_sender
+
+    return decorator_sender

test/test.py

@@ -0,0 +1,10 @@
+from knockknock.desktop_sender import desktop_sender
+
+
+@desktop_sender(title="test")
+def train():
+    import time
+    time.sleep(10)
+    return {"loss":1}
+
+train()