Using tensorboard with ngrok tunneled Jupyter notebook

ngrok_tboard.py

from collections import namedtuple
import json
import psutil
import requests
import re
import subprocess


class ngrok_tboard_manager:
    def __init__(self, port=None):
        self.port = port or self.get_running_ngrok()

    def get_running_ngrok(self):
        lines = subprocess.getoutput('netstat -lntp | grep ngrok | grep LISTEN').split('\n')[2:]
        for line in lines:
            pid = re.search('(\d+)/ngrok', line).group(1)
            port = re.search('127.0.0.1:(\d+)', line).group(1)
            print(f'Found ngrok with PID {pid} listening port {port}')
            return port

    def list_tunnels(self):
        tunnel = namedtuple('ngrok_tunnel', ['name', 'public_url', 'proto', 'addr'])
        url = f'http://localhost:{self.port}/api/tunnels'
        r = json.loads(requests.get(url).content.decode('utf-8'))['tunnels']
        return [tunnel(t['name'], t['public_url'], t['proto'], t['config']['addr']) for t in r]

    def start_tunnel(self, name, addr, proto='http', bind_tls=False):
        if any(name == t.name for t in self.list_tunnels()):
            print(f'tunnels named {name} already exists')
            return
        url = f'http://localhost:{self.port}/api/tunnels'
        headers = {'content-type': 'application/json'}
        data = {'name': name, 'proto': proto, 'addr': addr, 'bind_tls': bind_tls}
        r = requests.post(url, data=json.dumps(data), headers=headers)
        if r.status_code == 201:
            public_url = json.loads(r.content.decode('utf-8'))['public_url']
            print(f'tunnel {name} has been started at attr={addr} with url {public_url}')
        else:
            raise Exception(r.content.decode('utf-8'))

    def stop_tunnel(self, name):
        if all(name != t.name for t in self.list_tunnels()):
            print(f'no tunnels named {name} found')
            return
        url = f'http://localhost:{self.port}/api/tunnels/{name}'
        r = requests.delete(url)
        if r.status_code == 204:
            print(f'tunnel {name} has been stopped properly')
        else:
            raise Exception(r.content.decode('utf-8'))

    def list_tboards(self):
        res, tboard = [], namedtuple('tboard', ['pid', 'logdir', 'port'])
        for pid in subprocess.getoutput('pgrep tensorboard').split('\n'):
            line = subprocess.getoutput(f'ss -l -p -n | grep "pid={pid}"')
            if line:
                port = int(line.split()[4][3:])
                cmdline = subprocess.getoutput(f'ps -fp {pid}')
                logdir = re.search('--logdir=(.*) ', cmdline).group(1)
                res.append(tboard(int(pid), logdir, port))
        return res

    def start_tboard(self, logdir):
        def get_available_port(port=6666):
            while any(conn.laddr.port == port for conn in psutil.net_connections()):
                print(f'port {port} is not available, trying port {port + 1}')
                port += 1
            return port
        for tboard in self.list_tboards():
            if tboard.logdir != logdir:
                continue
            print(f'tensorboard with {logdir} is already in tunnel tboard{tboard.port}')
            return
        port = get_available_port()
        cmd = f'tensorboard --logdir={logdir} --port={port}'
        subprocess.Popen(cmd.split())
        self.start_tunnel(f'tboard{port}', port)

    def stop_tboard(self, logdir):
        for tboard in self.list_tboards():
            if tboard.logdir != logdir:
                continue
            subprocess.Popen(['kill', str(tboard.pid)])
            self.stop_tunnel(f'tboard{tboard.port}')

# Usage:
# ng = ngrok_tboard_manager()
# ng.start_tboard(<logdir>)
# ng.stop_tboard(<logdir>)