Keeping Your Session Active in Google Colab

workspace_utils.py

# The workspace_utils.py module includes an iterator wrapper called keep_awake and a context manager 
# called active_session that can be used to maintain an active session during long-running processes. 
# The two functions are equivalent, so use whichever fits better in your code.
#
# EXAMPLE 1
#
#    from workspace_utils import keep_awake
#    for i in keep_awake(range(5)):  #anything that happens inside this loop will keep the workspace active
#        # do iteration with lots of work here
#
#
# EXAMPLE 2
#
#    from workspace_utils import active_session
#    with active_session():
#        # do long-running work here
#
#
# SOURCE
#
#     https://video.udacity-data.com/topher/2018/May/5b06f6c2_workspace-utils/workspace-utils.py
#

import signal

from contextlib import contextmanager

import requests


DELAY = INTERVAL = 4 * 60  # interval time in seconds
MIN_DELAY = MIN_INTERVAL = 2 * 60
KEEPALIVE_URL = "https://nebula.udacity.com/api/v1/remote/keep-alive"
TOKEN_URL = "http://metadata.google.internal/computeMetadata/v1/instance/attributes/keep_alive_token"
TOKEN_HEADERS = {"Metadata-Flavor":"Google"}


def _request_handler(headers):
    def _handler(signum, frame):
        requests.request("POST", KEEPALIVE_URL, headers=headers)
    return _handler


@contextmanager
def active_session(delay=DELAY, interval=INTERVAL):
    """
    Example:

    from workspace_utils import active session

    with active_session():
        # do long-running work here
    """
    token = requests.request("GET", TOKEN_URL, headers=TOKEN_HEADERS).text
    headers = {'Authorization': "STAR " + token}
    delay = max(delay, MIN_DELAY)
    interval = max(interval, MIN_INTERVAL)
    original_handler = signal.getsignal(signal.SIGALRM)
    try:
        signal.signal(signal.SIGALRM, _request_handler(headers))
        signal.setitimer(signal.ITIMER_REAL, delay, interval)
        yield
    finally:
        signal.signal(signal.SIGALRM, original_handler)
        signal.setitimer(signal.ITIMER_REAL, 0)


def keep_awake(iterable, delay=DELAY, interval=INTERVAL):
    """
    Example:

    from workspace_utils import keep_awake

    for i in keep_awake(range(5)):
        # do iteration with lots of work here
    """
    with active_session(delay, interval): yield from iterable