Merge pull request #256 from flatironinstitute/restartable-server

Add 'restart' endpoint to server for deploy automation

.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 }}"

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

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

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