diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
index e030c3d..4508465 100755
--- a/cht-petals/setup-petals.sh
+++ b/cht-petals/setup-petals.sh
@@ -1,15 +1,27 @@
-#!/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=
] [--dht-prefix=] [--port=]
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