Implement and utilize countdown widget (for temporary deployments)

home/widgets.py

@@ -1,5 +1,7 @@
 """AiiDAlab basic widgets."""
 
+import subprocess
+from datetime import datetime
 from threading import Timer
 
 import ipywidgets as ipw
@@ -187,3 +189,53 @@ def _refresh_output(self, _=None):
             self._output.value = (
                 self.template.format(text=self.value) if self.value else ""
             )
+
+
+class CountDownWidget(ipw.VBox):
+    """A countdown w.r.t the AiiDAlab container creation time."""
+
+    TEMPLATE = "<h1 style='color: {color}'>Time remaining: <b>{countdown}</b></h1>"
+    FINAL_WARNING = "<h1 style='color: red'>AiiDAlab container shutdown imminent - please save your work!</h1>"
+
+    def __init__(self, lifetime: str, **kwargs):
+        self.reference_time = datetime.strptime(lifetime, "%H:%M:%S")
+
+        self.countdown = ipw.HTML()
+
+        super().__init__(
+            children=[
+                self.countdown,
+            ],
+            **kwargs,
+        )
+
+    def start(self):
+        Timer(1, self._update_countdown).start()
+
+    def _update_countdown(self):
+        try:
+            remaining, time_running_out = self._remaining_time()
+            message = self.TEMPLATE.format(
+                color="red" if time_running_out else "black",
+                countdown=remaining,
+            )
+        except Exception:
+            remaining = None
+            message = "Failed to obtain remaining time."
+        self.countdown.value = message
+        if remaining != "00:00:00":
+            self.start()
+        else:
+            self.countdown.value = self.FINAL_WARNING
+
+    def _remaining_time(self) -> tuple[str, bool]:
+        process = subprocess.check_output(["ps", "-o", "etime=", "-p", "1"])
+        elapsed_str = process.decode().strip()
+        if len(elapsed_str) < 6:
+            elapsed_str = f"00:{elapsed_str}"
+        elapsed_time = datetime.strptime(elapsed_str, "%H:%M:%S")
+        remaining = self.reference_time - elapsed_time
+        if remaining.total_seconds() >= 0:
+            time_running_out = remaining.total_seconds() < 60 * 5
+            return (datetime.min + remaining).strftime("%H:%M:%S"), time_running_out
+        return "00:00:00", True

start.ipynb

@@ -46,6 +46,27 @@
     "    parsed_url = urlparse.parse_qs(url.query)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from yaml import safe_load\n",
+    "\n",
+    "from home.widgets import CountDownWidget\n",
+    "\n",
+    "# For temporary deployments (e.g. demo server), duration is read from a local\n",
+    "# config file. This will initiate the countdown widget.\n",
+    "demo_server_config_file = Path.home() / \".aiidalab\" / \"demo-server-config.yml\"\n",
+    "if demo_server_config_file.exists():\n",
+    "    demo_server_config: dict = safe_load(demo_server_config_file.read_text())  # type: ignore\n",
+    "    if \"lifetime\" in demo_server_config:\n",
+    "        countdown = CountDownWidget(lifetime=demo_server_config[\"lifetime\"])\n",
+    "        countdown.start()\n",
+    "        display(countdown)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,