Telegram notifications from Jupyter Notebooks

When running long running code in Jupyter, I want to get notified when it finished so that I can get back to it. There is an extension to do that with browser notifications, but there are times when I leave the computer while waiting for an ML training to finish.

For long running CLI commands there is the ntfy, a command line tool that allows you to send notifications through a lot of channels.

So I hacked the two together to get some code that automatically messages me on Telegram when a cell finished more than 60 seconds after it started. This extension is registered automatically on the startup of any IPython and Jupyter notebook (even if they are installed in random virtual environments). Why Telegram? Because I already have it installed and it seemed like the easiest integration to set up.

The code has to be placed in ~\.ipython\profile_default\startup\01notification.py. You can place multiple files in this folder and they are loaded in lexicographic order, so you should prepend a number if you care about order. First, a couple of magic imports:

import time
import subprocess

from IPython.core.getipython import get_ipython
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
from IPython.core.magic import (register_line_magic)

To send the notification using ntfy, I'm simply calling the program using subprocess. The path resolution used by subprocess is not clear to me, so I had to use the full path to the executable.

def display_notification(message):
    subprocess.run([r"C:\Users\Roland\AppData\Local\Programs\Python\Python310\Scripts\ntfy.exe", "-b", "telegram", "send", message]) 

Then we define some variables, one for the threshold for notification and the other one for remembering the start of the execution:

autonotify_after = 60
run_start_time = None

And now we define the magic function (that's what these things are called in IPython). It has two arguments, one to override the duration of the threshold for notifications and the other for the default message. I copy pasted the decorators as you see them. After parsing the arguments (which come as a string), we register two event handlers: one to run before a cell is executed and one after a cell is executed.

@magic_arguments()
@argument(
    "-a", "--after", default=None,
    help="Send notification if cell execution is longer than x seconds"
)
@argument(
    "-m",
    "--message",
    default="Cell Execution Has Finished!",
    help="Custom notification message"
)
@register_line_magic
def autonotify(line):
    # Record options
    args = parse_argstring(autonotify, line)
    message = args.message.lstrip("\'\"").rstrip("\'\"")
    if args.after:
        global autonotify_after
        autonotify_after = args.after
    ### Register events
    ip = get_ipython()

    # Register new events
    ip.events.register('pre_run_cell', pre_run_cell)
    ip.events.register('post_run_cell', lambda: post_run_cell(message))

The handler to run before a cell is simple: we just record the start time of the run.

def pre_run_cell():
    global run_start_time
    run_start_time = time.time()

The second handler is slightly more complex. We look at the output of the last cell and append it to the message if it's not an "empty" value. We check how long has elapsed to know whether to show a notification or not:

def post_run_cell(message):
    # Set last output as notification message
    last_output = get_ipython().user_global_ns['_']
    # Don't use output if it's None or empty (but still allow False, 0, etc.)
    try:
        if last_output is not None and len(str(last_output)):
            message = message + "\n" + str(last_output)
    except ValueError:
        pass # can't convert to string. Use default message
    
    # Check autonotify options and perform checks
    if not check_after(): 
        return
    display_notification(message)


def check_after():
    # Check if the time elapsed is over the specified time.
    now, start = time.time(), run_start_time
    return autonotify_after >= 0 and start and (now - start) >= autonotify_after

The last piece of magic is to run this function. The other blog post I was inspired by said you should delete the function for things to work properly, so I did that as well:

ipython = get_ipython()
ipython.magic('autonotify')
del autonotify

And voila, now you will get Telegram notifications automatically when your model  finishes training! Setting up a Telegram bot is left as an exercise for the reader.