Send Activity Notifications to JupyterHub

Author: jwindgassen

jupyter_server_proxy/standalone/__init__.py

@@ -6,13 +6,8 @@
 from tornado.httpserver import HTTPServer
 from tornado.log import app_log
 
-from .proxy import (
-    configure_http_client,
-    get_port_from_env,
-    get_ssl_options,
-    make_app,
-    start_keep_alive,
-)
+from .activity import start_activity_update
+from .proxy import configure_http_client, get_port_from_env, get_ssl_options, make_app
 
 
 def run(
@@ -24,8 +19,7 @@ def run(
     logs=True,
     authtype="oauth",
     timeout=60,
-    last_activity_interval=300,
-    force_alive=True,
+    activity_interval=300,
     progressive=False,
     websocket_max_message_size=0,
 ):
@@ -69,8 +63,8 @@ def run(
     print(f"Auth Type: {authtype}")
     print(f"Command: {command}")
 
-    if last_activity_interval > 0:
-        start_keep_alive(last_activity_interval, force_alive, app.settings)
+    if activity_interval > 0:
+        start_activity_update(activity_interval)
 
     ioloop.IOLoop.current().start()
 
@@ -117,16 +111,10 @@ def main():
         help="Timeout to wait until the subprocess has started and can be addressed.",
     )
     parser.add_argument(
-        "--last-activity-interval",
+        "--activity-interval",
         default=300,
         type=int,
-        help="Frequency to notify Hub that the WebApp is still running in seconds. 0 for never.",
-    )
-    parser.add_argument(
-        "--force-alive",
-        action="store_true",
-        default=True,
-        help="Always report, that there has been activity (force keep alive) - only if last-activity-interval > 0.",
+        help="Frequency to notify Hub that the WebApp is still running (In seconds, 0 for never).",
     )
     parser.add_argument(
         "--progressive",

jupyter_server_proxy/standalone/activity.py

@@ -0,0 +1,75 @@
+import json
+import os
+from datetime import datetime
+
+from jupyterhub.utils import exponential_backoff, isoformat
+from tornado import httpclient, ioloop
+from tornado.log import app_log as log
+
+
+async def notify_activity():
+    """
+    Regularly notify JupyterHub of activity.
+    See `jupyrehub/singleuser/extensions#L396`
+    """
+
+    client = httpclient.AsyncHTTPClient()
+    last_activity_timestamp = isoformat(datetime.utcnow())
+    failure_count = 0
+
+    activity_url = os.environ.get("JUPYTERHUB_ACTIVITY_URL")
+    server_name = os.environ.get("JUPYTERHUB_SERVER_NAME")
+    api_token = os.environ.get("JUPYTERHUB_API_TOKEN")
+
+    if not (activity_url and server_name and api_token):
+        log.error(
+            "Could not find environment variables to send notification to JupyterHub"
+        )
+        return
+
+    async def notify():
+        """Send Notification, return if successful"""
+        nonlocal failure_count
+        log.debug(f"Notifying Hub of activity {last_activity_timestamp}")
+
+        req = httpclient.HTTPRequest(
+            url=activity_url,
+            method="POST",
+            headers={
+                "Authorization": f"token {api_token}",
+                "Content-Type": "application/json",
+            },
+            body=json.dumps(
+                {
+                    "servers": {
+                        server_name: {"last_activity": last_activity_timestamp}
+                    },
+                    "last_activity": last_activity_timestamp,
+                }
+            ),
+        )
+
+        try:
+            await client.fetch(req)
+            return True
+        except httpclient.HTTPError as e:
+            failure_count += 1
+            log.error(f"Error notifying Hub of activity: {e}")
+            return False
+
+    # Try sending notification for 1 minute
+    await exponential_backoff(
+        notify,
+        fail_message="Failed to notify Hub of activity",
+        start_wait=1,
+        max_wait=15,
+        timeout=60,
+    )
+
+    if failure_count > 0:
+        log.info(f"Sent hub activity after {failure_count} retries")
+
+
+def start_activity_update(interval):
+    pc = ioloop.PeriodicCallback(notify_activity, 1e3 * interval, 0.1)
+    pc.start()

jupyter_server_proxy/standalone/proxy.py

@@ -1,13 +1,11 @@
-import json
 import os
 import re
-from datetime import datetime
 from urllib.parse import urlparse
 
 from jupyterhub import __version__ as __jh_version__
 from jupyterhub.services.auth import HubOAuthCallbackHandler
-from jupyterhub.utils import exponential_backoff, isoformat, make_ssl_context
-from tornado import httpclient, ioloop
+from jupyterhub.utils import make_ssl_context
+from tornado import httpclient
 from tornado.web import Application, RedirectHandler, RequestHandler
 from tornado.websocket import WebSocketHandler
 
@@ -183,70 +181,3 @@ def get_port_from_env():
         elif url.scheme == "https":
             return 443
     return 8888
-
-
-def start_keep_alive(last_activity_interval, force_alive, settings):
-    client = httpclient.AsyncHTTPClient()
-
-    hub_activity_url = os.environ.get("JUPYTERHUB_ACTIVITY_URL", "")
-    server_name = os.environ.get("JUPYTERHUB_SERVER_NAME", "")
-    api_token = os.environ.get("JUPYTERHUB_API_TOKEN", "")
-
-    if api_token == "" or server_name == "" or hub_activity_url == "":
-        print(
-            "The following env vars are required to report activity back to the hub for keep alive: "
-            "JUPYTERHUB_ACTIVITY_URL ({}), JUPYTERHUB_SERVER_NAME({})".format(
-                hub_activity_url, server_name
-            )
-        )
-        return
-
-    async def send_activity():
-        async def notify():
-            print("About to notify Hub of activity")
-
-            last_activity_timestamp = None
-
-            if force_alive:
-                last_activity_timestamp = datetime.utcnow()
-            else:
-                last_activity_timestamp = settings.get("api_last_activity", None)
-
-            if last_activity_timestamp:
-                last_activity_timestamp = isoformat(last_activity_timestamp)
-                req = httpclient.HTTPRequest(
-                    url=hub_activity_url,
-                    method="POST",
-                    headers={
-                        "Authorization": f"token {api_token}",
-                        "Content-Type": "application/json",
-                    },
-                    body=json.dumps(
-                        {
-                            "servers": {
-                                server_name: {"last_activity": last_activity_timestamp}
-                            },
-                            "last_activity": last_activity_timestamp,
-                        }
-                    ),
-                )
-                try:
-                    await client.fetch(req)
-                except Exception as e:
-                    print(f"Error notifying Hub of activity: {e}")
-                    return False
-                else:
-                    return True
-
-            return True  # Nothing to report, so really it worked
-
-        await exponential_backoff(
-            notify,
-            fail_message="Failed to notify Hub of activity",
-            start_wait=1,
-            max_wait=15,
-            timeout=60,
-        )
-
-    pc = ioloop.PeriodicCallback(send_activity, 1e3 * last_activity_interval, 0.1)
-    pc.start()