Fetch container uptime in timezone-independent way

Author: edan-bainglass

stack/lab/Dockerfile

@@ -45,14 +45,25 @@ RUN jupyter nbextension enable --py --sys-prefix appmode && \
      jupyter serverextension enable --py --sys-prefix appmode
 
 # Swap appmode icon for AiiDAlab gears icon, shown during app load
-COPY --chown=${NB_UID}:${NB_GID} gears.svg ${CONDA_DIR}/share/jupyter/nbextensions/appmode/gears.svg
+COPY gears.svg ${CONDA_DIR}/share/jupyter/nbextensions/appmode/gears.svg
 
-# Set up opt-in countdown feature
 ARG PYTHON_MINOR_VERSION
-ENV PYTHON_MINOR_VERSION=${PYTHON_MINOR_VERSION}
-COPY --chown=${NB_UID}:${NB_GID} countdown/ ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/custom/
+
+# Set up the logo of notebook interface
+COPY aiidalab-wide-logo.png ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/base/images/logo.png
+
+# Copy custom CSS and JS files for AiiDAlab container countdown feature
+COPY countdown/ ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/custom/
+
+# Add endpoint to fetch container uptime
+COPY aiidalab_uptime.py ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/
+COPY aiidalab_uptime.json ${CONDA_DIR}/etc/jupyter/jupyter_notebook_config.d/
 
 # Copy start-up scripts for AiiDAlab.
+# We expose PYTHON_MINOR_VERSION in the container, as it is used by the
+# before-notebook.d/70_prepare_countdown_config.sh script to write the
+# environment-variable-dependent (opt-in) configuration file.
+ENV PYTHON_MINOR_VERSION=${PYTHON_MINOR_VERSION}
 COPY before-notebook.d/* /usr/local/bin/before-notebook.d/
 
 # Configure AiiDAlab environment.
@@ -76,12 +87,6 @@ ENV AIIDALAB_DEFAULT_APPS=""
 # Specify default factory reset (not set):
 ENV AIIDALAB_FACTORY_RESET=""
 
-USER ${NB_USER}
-
-WORKDIR "/home/${NB_USER}"
-
-RUN  mkdir -p /home/${NB_USER}/apps
-
 # When a Jupyter notebook server looses a connection to the frontend,
 # it keeps the messages in a buffer. If there is a background thread running
 # and trying to update the frontend, the buffer grows indefinitely,
@@ -103,5 +108,8 @@ ENV NOTEBOOK_ARGS=\
 "--TerminalManager.cull_inactive_timeout=3600 "\
 "--TerminalManager.cull_interval=300"
 
-# Set up the logo of notebook interface
-COPY --chown=${NB_UID}:${NB_GID} aiidalab-wide-logo.png ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/base/images/logo.png
+USER ${NB_USER}
+
+WORKDIR "/home/${NB_USER}"
+
+RUN  mkdir -p /home/${NB_USER}/apps

stack/lab/aiidalab_uptime.json

@@ -0,0 +1,7 @@
+{
+  "NotebookApp": {
+    "nbserver_extensions": {
+      "aiidalab_uptime": true
+    }
+  }
+}

stack/lab/aiidalab_uptime.py

@@ -0,0 +1,32 @@
+import re
+import subprocess
+
+from notebook.base.handlers import IPythonHandler
+from notebook.utils import url_path_join
+
+
+class UptimeHandler(IPythonHandler):
+    def get(self):
+        try:
+            output = subprocess.check_output(["ps", "-p", "1", "-o", "etime="])
+            etime = output.decode("utf-8").strip()
+            seconds = self._parse_etime_to_seconds(etime)
+            self.finish({"uptime": seconds})
+        except Exception as e:
+            self.set_status(500)
+            self.finish({"error": str(e)})
+
+    def _parse_etime_to_seconds(self, etime):
+        # Supports MM:SS, HH:MM:SS, or D-HH:MM:SS formats
+        match = re.match(r"(?:(\d+)-)?(?:(\d+):)?(\d+):(\d+)", etime)
+        if not match:
+            raise ValueError(f"Unrecognized etime format: {etime}")
+
+        days, hours, minutes, seconds = match.groups(default="0")
+        return int(days) * 86400 + int(hours) * 3600 + int(minutes) * 60 + int(seconds)
+
+
+def load_jupyter_server_extension(nb_server_app):
+    web_app = nb_server_app.web_app
+    route = url_path_join(web_app.settings["base_url"], "/uptime")
+    web_app.add_handlers(".*", [(route, UptimeHandler)])

stack/lab/before-notebook.d/70_prepare_countdown_config.sh

@@ -3,20 +3,9 @@ set -e
 
 CUSTOM_DIR="${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/custom"
 
-if [ "$LIFETIME" ]; then
-    EPHEMERAL=true
-
-    # Convert LIFETIME from HH:MM:SS to seconds
-    IFS=: read -r H M S <<<"$LIFETIME"
-    LIFETIME_SEC=$((10#$H * 3600 + 10#$M * 60 + 10#$S))
-
-    # Calculate expiry timestamp in UTC
-    EXPIRY=$(date -u -d "+${LIFETIME_SEC} seconds" +"%Y-%m-%dT%H:%M:%SZ")
-    export EXPIRY
-else
-    EPHEMERAL=false
-fi
-
-export EPHEMERAL
-envsubst <"${CUSTOM_DIR}/config.json.template" >"$CUSTOM_DIR/config.json"
-rm "${CUSTOM_DIR}/config.json.template"
+cat <<EOF >"${CUSTOM_DIR}/config.json"
+{
+  "ephemeral": $([ -n "$LIFETIME" ] && echo 1 || echo 0),
+  "lifetime": "$LIFETIME"
+}
+EOF

stack/lab/countdown/config.json.template

@@ -1,11 +1,18 @@
 require(["base/js/namespace", "base/js/events"], (Jupyter, events) => {
   const parseLifetimeToMs = (str) => {
-    const parts = str.split(":").map(Number);
-    if (parts.length !== 3 || parts.some(isNaN)) {
+    // Supports MM:SS, HH:MM:SS, or D-HH:MM:SS formats
+    const regex = /^(?:(\d+)-)?(?:(\d+):)?(\d+):(\d+)$/;
+    const match = regex.exec(str.trim());
+
+    if (!match) {
+      console.warn("Invalid lifetime format:", str);
       return null;
     }
-    const [h, m, s] = parts;
-    return ((h * 60 + m) * 60 + s) * 1000;
+
+    const [_, days, hours, minutes, seconds] = match.map((v) => Number(v) || 0);
+
+    const totalSeconds = days * 86400 + hours * 3600 + minutes * 60 + seconds;
+    return totalSeconds * 1000;
   };
 
   const insertCountdown = (remainingMs) => {
@@ -37,18 +44,19 @@ require(["base/js/namespace", "base/js/events"], (Jupyter, events) => {
     `;
     banner.appendChild(saveInfo);
 
-    const endTime = new Date(Date.now() + remainingMs);
+    const startTime = Date.now();
+    const endTime = startTime + remainingMs;
 
     const formatTime = (seconds) => {
       const hrs = `${Math.floor(seconds / 3600)}`.padStart(2, "0");
       const mins = `${Math.floor((seconds % 3600) / 60)}`.padStart(2, "0");
       const secs = `${Math.floor(seconds % 60)}`.padStart(2, "0");
+      // Format as HH:MM:SS, even if > 1 day (for now - rare?)
       return `${hrs}:${mins}:${secs}`;
     };
 
     const updateTimer = () => {
-      const now = new Date();
-      const timeLeft = (endTime - now) / 1000;
+      const timeLeft = (endTime - Date.now()) / 1000;
       if (timeLeft < 0) {
         clearInterval(interval);
         return;
@@ -74,23 +82,28 @@ require(["base/js/namespace", "base/js/events"], (Jupyter, events) => {
     }
   };
 
-  loadCountdown = async () => {
+  const loadCountdown = async () => {
     try {
-      const response = await fetch("/static/custom/config.json");
-      const config = await response.json();
+      const [configResponse, uptimeResponse] = await Promise.all([
+        fetch("/static/custom/config.json"),
+        fetch("/uptime"),
+      ]);
+
+      const config = await configResponse.json();
+      const uptimeData = await uptimeResponse.json();
 
-      // Opt-in point for deployments
       if (!config.ephemeral) {
         return;
       }
 
-      if (!config.expiry) {
-        console.warn("Missing `expiry` in config file");
+      if (!config.lifetime) {
+        console.warn("Missing `lifetime` in config file");
         return;
       }
 
-      const expiry = new Date(config.expiry).getTime();
-      const remaining = expiry - Date.now();
+      const lifetimeMs = parseLifetimeToMs(config.lifetime);
+      const uptimeMs = uptimeData.uptime * 1000;
+      const remaining = lifetimeMs - uptimeMs;
 
       insertCountdown(Math.max(0, remaining));
     } catch (err) {