From 2ed38d3ed06ed9779079367b2aa461303cff32cb Mon Sep 17 00:00:00 2001
From: Brian Ward <bward@flatironinstitute.org>
Date: Wed, 18 Dec 2024 20:32:51 +0000
Subject: [PATCH 1/2] Add 'restart' endpoint to server for deploy automation

---
 .github/workflows/docker_deploy.yml        |  4 ++++
 backend/stan-wasm-server/src/app/config.py |  3 ++-
 backend/stan-wasm-server/src/app/main.py   | 17 +++++++++++++++++
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/docker_deploy.yml b/.github/workflows/docker_deploy.yml
index bb9353c7..dd529e52 100644
--- a/.github/workflows/docker_deploy.yml
+++ b/.github/workflows/docker_deploy.yml
@@ -37,3 +37,7 @@ jobs:
           TAG=${INPUT_TAG:-"latest"}
           docker build -t ghcr.io/flatironinstitute/stan-wasm-server:$TAG .
           docker push ghcr.io/flatironinstitute/stan-wasm-server:$TAG
+
+      - Name: Ping public server to restart
+        run: |
+          curl -X POST https://stan-wasm.flatironinstitute.org/restart -H "Authorization: Bearer ${{ secrets.PUBLIC_SERVER_RESTART_TOKEN }}"
diff --git a/backend/stan-wasm-server/src/app/config.py b/backend/stan-wasm-server/src/app/config.py
index cfdaf720..1e6c7c2c 100644
--- a/backend/stan-wasm-server/src/app/config.py
+++ b/backend/stan-wasm-server/src/app/config.py
@@ -1,7 +1,7 @@
 import logging
 from functools import lru_cache
 from pathlib import Path
-from typing import Annotated, Literal
+from typing import Annotated, Literal, Optional
 
 from pydantic import (
     AliasChoices,
@@ -23,6 +23,7 @@ class StanWasmServerSettings(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="SWS_")
 
     passcode: SecretStr
+    restart_token: Optional[SecretStr] = None
     job_dir: Path = Path("/jobs")
     built_model_dir: Path = Path("/compiled_models")
     compilation_timeout: PositiveInt = 60 * 5
diff --git a/backend/stan-wasm-server/src/app/main.py b/backend/stan-wasm-server/src/app/main.py
index 4cbff4ec..c137fe7d 100644
--- a/backend/stan-wasm-server/src/app/main.py
+++ b/backend/stan-wasm-server/src/app/main.py
@@ -154,3 +154,20 @@ async def run_job(job_id: str, settings: DependsOnSettings) -> DictResponse:
     )
 
     return {"success": True}
+
+
+@app.post("/restart")
+async def restart(
+    settings: DependsOnSettings, authorization: str = Header(None)
+) -> None:
+    if settings.restart_token is None:
+        raise StanPlaygroundAuthenticationException("Restart token not set at startup")
+    check_authorization(authorization, settings.restart_token)
+
+    import os
+    import signal
+
+    # send an interrupt signal to the parent process
+    # uvicorn interprets this like Ctrl-C, and gracefully shuts down
+    os.kill(os.getppid(), signal.SIGINT)
+    # actual restart is handled by the orchestrator

From 1e724b31b9d9b847afa0eb032492d8821c881238 Mon Sep 17 00:00:00 2001
From: Brian Ward <bward@flatironinstitute.org>
Date: Wed, 18 Dec 2024 21:02:02 +0000
Subject: [PATCH 2/2] Set max shutdown time to 20 seconds

---
 backend/stan-wasm-server/run.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/backend/stan-wasm-server/run.sh b/backend/stan-wasm-server/run.sh
index 15e983c6..78610bbd 100644
--- a/backend/stan-wasm-server/run.sh
+++ b/backend/stan-wasm-server/run.sh
@@ -2,4 +2,4 @@
 
 set -ex
 
-uvicorn --app-dir ./src/app main:app --host 0.0.0.0 --port 8080 --workers 4
+uvicorn --app-dir ./src/app main:app --host 0.0.0.0 --port 8080 --workers 4 --timeout-graceful-shutdown 20