safer interrupt handling

cht-petals/setup-petals.sh

@@ -1,15 +1,26 @@
-#!/bin/bash
-# ./setup-petals.sh --model-path ./models/models--petals-team--StableBeluga2 --dht-prefix StableBeluga2-hf --port 8794
+#!/usr/bin/env bash
+# Usage: setup-petals.sh [--model-path=<DIR>] [--dht-prefix=<PREFIX>] [--port=<INT>]
 
 tmpdir="${PREM_APPDIR:-.}/$(uuid)"
 
+cleanup(){
+  for i in $(jobs -p); do
+    kill -n 9 $i
+  done
+  rm -rf "$tmpdir"
+  exit 0
+}
+
+trap "cleanup" SIGTERM
+trap "cleanup" SIGINT
+
 # clone source
 git clone -n --depth=1 --filter=tree:0 https://github.com/premAI-io/prem-services.git "$tmpdir"
 git -C "$tmpdir" sparse-checkout set --no-cone cht-petals
 git -C "$tmpdir" checkout
 # install deps
-"${PREM_PYTHON:-python}" -m pip install -r "$tmpdir/cht-petals/requirements.txt" || :
+"${PREM_PYTHON:-python}" -m pip install -r "$tmpdir/cht-petals/requirements.txt"
 # run server
-PYTHONPATH="$tmpdir/cht-petals" "${PREM_PYTHON:-python}" "$tmpdir/cht-petals/main.py" "$@" || :
+PYTHONPATH="$tmpdir/cht-petals" "${PREM_PYTHON:-python}" "$tmpdir/cht-petals/main.py" "$@" &
 
-rm -rf "$tmpdir"
+wait