Skip to content

Instantly share code, notes, and snippets.

@thomd

thomd/workspace_utils.py

Last active Apr 17, 2020
Embed
What would you like to do?
Keeping Your Session Active in Google Colab
# 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment