From 4736c77d8eb20e9f617623cb1190597a0c796434 Mon Sep 17 00:00:00 2001
From: biswaroop1547 <biswaroop08@gmail.com>
Date: Fri, 20 Oct 2023 07:01:09 +0530
Subject: [PATCH 01/12] update: new bash script for managing installation and
 starting of cht-petals service without docker

---
 cht-petals/download.py      | 11 ++----
 cht-petals/main.py          | 13 ++++++-
 cht-petals/models.py        |  8 ++---
 cht-petals/requirements.txt | 18 +++++-----
 cht-petals/setup-petals.sh  | 69 +++++++++++++++++++++++++++++++++++++
 5 files changed, 95 insertions(+), 24 deletions(-)
 create mode 100644 cht-petals/setup-petals.sh

diff --git a/cht-petals/download.py b/cht-petals/download.py
index bddad6b..fbed12b 100644
--- a/cht-petals/download.py
+++ b/cht-petals/download.py
@@ -1,5 +1,4 @@
 import argparse
-from platform import machine
 
 import torch
 from petals import AutoDistributedModelForCausalLM
@@ -8,20 +7,16 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--model", help="Model to download")
+parser.add_argument("--dht-prefix", help="DHT prefix to use")
 args = parser.parse_args()
 
-print(f"Downloading model {args.model}")
+print(f"Downloading model {args.model} with DHT prefix {args.dht_prefix}")
 
 
 @retry(stop=stop_after_attempt(3), wait=wait_fixed(5))
 def download_model() -> None:
     Tokenizer = LlamaTokenizer if "llama" in args.model.lower() else AutoTokenizer
     _ = Tokenizer.from_pretrained(args.model)
-
-    kwargs = {}
-    if "x86_64" in machine():
-        kwargs["torch_dtype"] = torch.float32
-    _ = AutoDistributedModelForCausalLM.from_pretrained(args.model, **kwargs)
-
+    _ = AutoDistributedModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.float32, dht_prefix=args.dht_prefix)
 
 download_model()
diff --git a/cht-petals/main.py b/cht-petals/main.py
index 0b416bb..c297194 100644
--- a/cht-petals/main.py
+++ b/cht-petals/main.py
@@ -1,4 +1,6 @@
+import argparse
 import logging
+import os
 
 import uvicorn
 from dotenv import load_dotenv
@@ -8,6 +10,15 @@
 
 load_dotenv()
 
+MODEL_PATH = "./models"
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", help="Path to Model files directory", default=MODEL_PATH)
+    parser.add_argument("--dht_prefix", help="DHT prefix to use")
+    args = parser.parse_args()
+    MODEL_PATH = args.model_path
+    DHT_PREFIX = args.dht_prefix
+
 logging.basicConfig(
     format="%(asctime)s %(levelname)-8s %(message)s",
     level=logging.INFO,
@@ -19,7 +30,7 @@ def create_start_app_handler(app: FastAPI):
     def start_app() -> None:
         from models import PetalsBasedModel
 
-        PetalsBasedModel.get_model()
+        PetalsBasedModel.get_model(MODEL_PATH, DHT_PREFIX)
 
     return start_app
 
diff --git a/cht-petals/models.py b/cht-petals/models.py
index 8e83a15..1ca5ccb 100644
--- a/cht-petals/models.py
+++ b/cht-petals/models.py
@@ -56,13 +56,9 @@ def generate(
         return [cls.tokenizer.decode(outputs[0])]
 
     @classmethod
-    def get_model(cls):
+    def get_model(cls, model_path: str, dht_prefix: str):
         if cls.model is None:
             Tokenizer = LlamaTokenizer if "llama" in os.getenv("MODEL_ID").lower() else AutoTokenizer
             cls.tokenizer = Tokenizer.from_pretrained(os.getenv("MODEL_ID"))
-
-            kwargs = {}
-            if "x86_64" in machine():
-                kwargs["torch_dtype"] = torch.float32
-            cls.model = AutoDistributedModelForCausalLM.from_pretrained(os.getenv("MODEL_ID"), **kwargs)
+            cls.model = AutoDistributedModelForCausalLM.from_pretrained(os.getenv("MODEL_ID", model_path), torch_dtype=torch.float32, dht_prefix=os.getenv("DHT_PREFIX", dht_prefix))
         return cls.model
diff --git a/cht-petals/requirements.txt b/cht-petals/requirements.txt
index 205b6a1..5986130 100644
--- a/cht-petals/requirements.txt
+++ b/cht-petals/requirements.txt
@@ -1,9 +1,9 @@
-fastapi==0.95.0
-uvicorn==0.21.1
-pytest==7.2.2
-requests==2.28.2
-tqdm==4.65.0
-httpx==0.23.3
-python-dotenv==1.0.0
-tenacity==8.2.2
-petals==2.2.0
+fastapi==0.*
+uvicorn==0.*
+pytest==7.*
+requests==2.*
+tqdm==4.*
+httpx==0.*
+python-dotenv==1.*
+tenacity==8.*
+petals==2.*
\ No newline at end of file
diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
new file mode 100644
index 0000000..509b4e2
--- /dev/null
+++ b/cht-petals/setup-petals.sh
@@ -0,0 +1,69 @@
+#!/bin/bash
+# ./setup-petals.sh --model_path ./models/models--petals-team--StableBeluga2 --dht_prefix StableBeluga2-hf
+# TODO: replace appdir with ~/.prem/appdir (?)
+
+
+# Function to install Miniconda
+install_miniconda() {
+    local arch
+    if [ "$1" == "x86_64" ]; then
+        arch="x86_64"
+    elif [ "$1" == "arm64" ]; then
+        arch="arm64"
+    else
+        echo "Unsupported architecture: $1"
+        exit 1
+    fi
+
+    # Create the 'appdir' directory if it doesn't exist
+    if [ ! -d "appdir" ]; then
+        mkdir appdir
+    fi
+
+    # Download and install Miniconda
+    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-"$arch".sh -O miniconda_installer.sh
+    bash miniconda_installer.sh -b -u -p appdir
+    rm miniconda_installer.sh
+}
+
+# Check if 'appdir' directory exists, if not, create it
+if [ ! -d "appdir" ]; then
+    install_miniconda "$(arch | cut -d_ -f2)"
+else
+    echo "Using existing 'appdir' directory."
+fi
+
+# Activate the Miniconda environment and create 'prem_env'
+appdir/bin/conda create -n prem_env python=3.11 -y
+
+# Clone the GitHub repository and install requirements
+git clone -n --depth=1 --filter=tree:0 https://github.com/premAI-io/prem-services.git
+git -C prem-services sparse-checkout set --no-cone cht-petals
+git -C prem-services checkout
+
+appdir/envs/prem_env/bin/pip install -r prem-services/cht-petals/requirements.txt
+
+# Check for the --model_path and --dht_prefix parameters and run main.py
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --model_path)
+            model_path="$2"
+            shift 2
+            ;;
+        --dht_prefix)
+            dht_prefix="$2"
+            shift 2
+            ;;
+        *)
+            shift
+            ;;
+    esac
+done
+
+if [ -n "$model_path" ] && [ -n "$dht_prefix" ]; then
+    export PYTHONPATH="$(pwd)/prem-services"
+    appdir/envs/prem_env/bin/python prem-services/cht-petals/main.py --model_path "$model_path" --dht_prefix "$dht_prefix"
+else
+    echo "Please provide the --model_path parameter path to model directory and the --dht_prefix parameter for the DHT prefix."
+    exit 1
+fi

From ed7501043012ed6d87170a8353a77bbae464f816 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 20 Oct 2023 01:39:58 +0000
Subject: [PATCH 02/12] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 cht-petals/download.py      | 5 ++++-
 cht-petals/models.py        | 6 +++++-
 cht-petals/requirements.txt | 2 +-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/cht-petals/download.py b/cht-petals/download.py
index fbed12b..52af9f4 100644
--- a/cht-petals/download.py
+++ b/cht-petals/download.py
@@ -17,6 +17,9 @@
 def download_model() -> None:
     Tokenizer = LlamaTokenizer if "llama" in args.model.lower() else AutoTokenizer
     _ = Tokenizer.from_pretrained(args.model)
-    _ = AutoDistributedModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.float32, dht_prefix=args.dht_prefix)
+    _ = AutoDistributedModelForCausalLM.from_pretrained(
+        args.model, torch_dtype=torch.float32, dht_prefix=args.dht_prefix
+    )
+
 
 download_model()
diff --git a/cht-petals/models.py b/cht-petals/models.py
index 1ca5ccb..c25530b 100644
--- a/cht-petals/models.py
+++ b/cht-petals/models.py
@@ -60,5 +60,9 @@ def get_model(cls, model_path: str, dht_prefix: str):
         if cls.model is None:
             Tokenizer = LlamaTokenizer if "llama" in os.getenv("MODEL_ID").lower() else AutoTokenizer
             cls.tokenizer = Tokenizer.from_pretrained(os.getenv("MODEL_ID"))
-            cls.model = AutoDistributedModelForCausalLM.from_pretrained(os.getenv("MODEL_ID", model_path), torch_dtype=torch.float32, dht_prefix=os.getenv("DHT_PREFIX", dht_prefix))
+            cls.model = AutoDistributedModelForCausalLM.from_pretrained(
+                os.getenv("MODEL_ID", model_path),
+                torch_dtype=torch.float32,
+                dht_prefix=os.getenv("DHT_PREFIX", dht_prefix),
+            )
         return cls.model
diff --git a/cht-petals/requirements.txt b/cht-petals/requirements.txt
index 5986130..77f5641 100644
--- a/cht-petals/requirements.txt
+++ b/cht-petals/requirements.txt
@@ -6,4 +6,4 @@ tqdm==4.*
 httpx==0.*
 python-dotenv==1.*
 tenacity==8.*
-petals==2.*
\ No newline at end of file
+petals==2.*

From 919962c4f0e7195c6917854c5992be92c685d0dd Mon Sep 17 00:00:00 2001
From: biswaroop1547 <biswaroop08@gmail.com>
Date: Fri, 20 Oct 2023 08:41:52 +0530
Subject: [PATCH 03/12] update: more resilient script + tested locally

---
 cht-petals/main.py         |  6 +--
 cht-petals/models.py       |  5 +--
 cht-petals/setup-petals.sh | 87 ++++++++++++++++++++++++++------------
 3 files changed, 66 insertions(+), 32 deletions(-)
 mode change 100644 => 100755 cht-petals/setup-petals.sh

diff --git a/cht-petals/main.py b/cht-petals/main.py
index c297194..b34e15d 100644
--- a/cht-petals/main.py
+++ b/cht-petals/main.py
@@ -1,6 +1,5 @@
 import argparse
 import logging
-import os
 
 import uvicorn
 from dotenv import load_dotenv
@@ -11,6 +10,7 @@
 load_dotenv()
 
 MODEL_PATH = "./models"
+DHT_PREFIX = "StableBeluga2"
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_path", help="Path to Model files directory", default=MODEL_PATH)
@@ -30,6 +30,7 @@ def create_start_app_handler(app: FastAPI):
     def start_app() -> None:
         from models import PetalsBasedModel
 
+        print(f"Using {MODEL_PATH=} and {DHT_PREFIX=}")
         PetalsBasedModel.get_model(MODEL_PATH, DHT_PREFIX)
 
     return start_app
@@ -51,6 +52,5 @@ def get_application() -> FastAPI:
 
 app = get_application()
 
-
 if __name__ == "__main__":
-    uvicorn.run("main:app", host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)
diff --git a/cht-petals/models.py b/cht-petals/models.py
index c25530b..f80849a 100644
--- a/cht-petals/models.py
+++ b/cht-petals/models.py
@@ -1,6 +1,5 @@
 import os
 from abc import ABC, abstractmethod
-from platform import machine
 from typing import List
 
 import torch
@@ -58,8 +57,8 @@ def generate(
     @classmethod
     def get_model(cls, model_path: str, dht_prefix: str):
         if cls.model is None:
-            Tokenizer = LlamaTokenizer if "llama" in os.getenv("MODEL_ID").lower() else AutoTokenizer
-            cls.tokenizer = Tokenizer.from_pretrained(os.getenv("MODEL_ID"))
+            Tokenizer = LlamaTokenizer if "llama" in os.getenv("MODEL_ID", model_path).lower() else AutoTokenizer
+            cls.tokenizer = Tokenizer.from_pretrained(os.getenv("MODEL_ID", model_path))
             cls.model = AutoDistributedModelForCausalLM.from_pretrained(
                 os.getenv("MODEL_ID", model_path),
                 torch_dtype=torch.float32,
diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
old mode 100644
new mode 100755
index 509b4e2..1955da7
--- a/cht-petals/setup-petals.sh
+++ b/cht-petals/setup-petals.sh
@@ -2,46 +2,81 @@
 # ./setup-petals.sh --model_path ./models/models--petals-team--StableBeluga2 --dht_prefix StableBeluga2-hf
 # TODO: replace appdir with ~/.prem/appdir (?)
 
+# Define the Miniconda installation path and environment name
+miniconda_path="appdir"
+conda_env="prem_env"
 
-# Function to install Miniconda
+# Function to install Miniconda if not already present
 install_miniconda() {
     local arch
-    if [ "$1" == "x86_64" ]; then
-        arch="x86_64"
-    elif [ "$1" == "arm64" ]; then
-        arch="arm64"
+
+    if [ "$(uname -s)" == "Darwin" ]; then
+        # macOS
+        if [ "$1" == "x86_64" ]; then
+            arch="x86_64"
+        elif [ "$1" == "arm64" ]; then
+            arch="arm64"
+        else
+            echo "Unsupported architecture: $1"
+            exit 1
+        fi
+    elif [ "$(uname -s)" == "Linux" ]; then
+        # Linux
+        if [ "$1" == "x86_64" ]; then
+            arch="x86_64"
+        elif [ "$1" == "aarch64" ]; then
+            arch="aarch64"
+        elif [ "$1" == "armv7l" ]; then
+            arch="armv7l"
+        else
+            echo "Unsupported architecture: $1"
+            exit 1
+        fi
     else
-        echo "Unsupported architecture: $1"
+        echo "Unsupported operating system: $(uname -s)"
         exit 1
     fi
 
     # Create the 'appdir' directory if it doesn't exist
-    if [ ! -d "appdir" ]; then
-        mkdir appdir
-    fi
+    if [ ! -d "$miniconda_path" ]; then
+        mkdir "$miniconda_path"
+
+        # Download and install Miniconda based on OS and architecture
+        if [ "$(uname -s)" == "Darwin" ]; then
+            wget "https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-$arch.sh" -O miniconda_installer.sh
+        elif [ "$(uname -s)" == "Linux" ]; then
+            wget "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-$arch.sh" -O miniconda_installer.sh
+        fi
 
-    # Download and install Miniconda
-    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-"$arch".sh -O miniconda_installer.sh
-    bash miniconda_installer.sh -b -u -p appdir
-    rm miniconda_installer.sh
+        bash miniconda_installer.sh -b -u -p "$miniconda_path"
+        rm miniconda_installer.sh
+    else
+        echo "Using existing Miniconda in '$miniconda_path'."
+    fi
 }
 
-# Check if 'appdir' directory exists, if not, create it
-if [ ! -d "appdir" ]; then
+# Check if Miniconda is already installed
+if ! command -v "$miniconda_path/bin/conda" &>/dev/null; then
     install_miniconda "$(arch | cut -d_ -f2)"
-else
-    echo "Using existing 'appdir' directory."
 fi
 
-# Activate the Miniconda environment and create 'prem_env'
-appdir/bin/conda create -n prem_env python=3.11 -y
+# Check if the Miniconda environment 'prem_env' exists
+if [ ! -d "$miniconda_path/envs/$conda_env" ]; then
+    "$miniconda_path/bin/conda" create -n "$conda_env" python=3.11 -y
+fi
 
-# Clone the GitHub repository and install requirements
-git clone -n --depth=1 --filter=tree:0 https://github.com/premAI-io/prem-services.git
-git -C prem-services sparse-checkout set --no-cone cht-petals
-git -C prem-services checkout
+# Clone the GitHub repository if not already present
+if [ ! -d "prem-services" ]; then
+    # only clone the required directory - https://stackoverflow.com/a/52269934
+    git clone -n --depth=1 --filter=tree:0 https://github.com/premAI-io/prem-services.git
+    git -C prem-services sparse-checkout set --no-cone cht-petals
+    git -C prem-services checkout
+else
+    echo "Using existing 'prem-services' directory."
+fi
 
-appdir/envs/prem_env/bin/pip install -r prem-services/cht-petals/requirements.txt
+# Install requirements in the Miniconda environment
+"$miniconda_path/envs/$conda_env/bin/pip" install -r prem-services/cht-petals/requirements.txt
 
 # Check for the --model_path and --dht_prefix parameters and run main.py
 while [[ $# -gt 0 ]]; do
@@ -62,8 +97,8 @@ done
 
 if [ -n "$model_path" ] && [ -n "$dht_prefix" ]; then
     export PYTHONPATH="$(pwd)/prem-services"
-    appdir/envs/prem_env/bin/python prem-services/cht-petals/main.py --model_path "$model_path" --dht_prefix "$dht_prefix"
+    "$miniconda_path/envs/$conda_env/bin/python" prem-services/cht-petals/main.py --model_path "$model_path" --dht_prefix "$dht_prefix"
 else
-    echo "Please provide the --model_path parameter path to model directory and the --dht_prefix parameter for the DHT prefix."
+    echo "Please provide the --model_path parameter with the path to model directory and the --dht_prefix parameter for the DHT prefix."
     exit 1
 fi

From d700ee213c7d8004843678a43c6b756fcc3180a2 Mon Sep 17 00:00:00 2001
From: Biswaroop Bhattacharjee <biswaroop08@gmail.com>
Date: Fri, 20 Oct 2023 19:09:07 +0530
Subject: [PATCH 04/12] update: requirements semver

Co-authored-by: Casper da Costa-Luis <casper.dcl@physics.org>
---
 cht-petals/requirements.txt | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cht-petals/requirements.txt b/cht-petals/requirements.txt
index 77f5641..0863be0 100644
--- a/cht-petals/requirements.txt
+++ b/cht-petals/requirements.txt
@@ -1,9 +1,9 @@
-fastapi==0.*
-uvicorn==0.*
+fastapi
+uvicorn
 pytest==7.*
 requests==2.*
 tqdm==4.*
-httpx==0.*
+httpx
 python-dotenv==1.*
 tenacity==8.*
 petals==2.*

From 02dc713106a71149d538faeae1a2cc1e1d44331b Mon Sep 17 00:00:00 2001
From: biswaroop1547 <biswaroop08@gmail.com>
Date: Fri, 20 Oct 2023 22:46:40 +0530
Subject: [PATCH 05/12] add: configurable --port for server and script trimdown
 to only handle cht-petals

---
 cht-petals/main.py         |  3 +-
 cht-petals/setup-petals.sh | 88 ++++++++------------------------------
 2 files changed, 21 insertions(+), 70 deletions(-)

diff --git a/cht-petals/main.py b/cht-petals/main.py
index b34e15d..7c520ae 100644
--- a/cht-petals/main.py
+++ b/cht-petals/main.py
@@ -15,6 +15,7 @@
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_path", help="Path to Model files directory", default=MODEL_PATH)
     parser.add_argument("--dht_prefix", help="DHT prefix to use")
+    parser.add_argument("--port", help="Port to run model server on", type=int)
     args = parser.parse_args()
     MODEL_PATH = args.model_path
     DHT_PREFIX = args.dht_prefix
@@ -53,4 +54,4 @@ def get_application() -> FastAPI:
 app = get_application()
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=args.port)
diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
index 1955da7..38336bb 100755
--- a/cht-petals/setup-petals.sh
+++ b/cht-petals/setup-petals.sh
@@ -1,69 +1,15 @@
 #!/bin/bash
-# ./setup-petals.sh --model_path ./models/models--petals-team--StableBeluga2 --dht_prefix StableBeluga2-hf
-# TODO: replace appdir with ~/.prem/appdir (?)
+# ./setup-petals.sh --model_path ./models/models--petals-team--StableBeluga2 --dht_prefix StableBeluga2-hf --port 8794
 
-# Define the Miniconda installation path and environment name
-miniconda_path="appdir"
-conda_env="prem_env"
-
-# Function to install Miniconda if not already present
-install_miniconda() {
-    local arch
-
-    if [ "$(uname -s)" == "Darwin" ]; then
-        # macOS
-        if [ "$1" == "x86_64" ]; then
-            arch="x86_64"
-        elif [ "$1" == "arm64" ]; then
-            arch="arm64"
-        else
-            echo "Unsupported architecture: $1"
-            exit 1
-        fi
-    elif [ "$(uname -s)" == "Linux" ]; then
-        # Linux
-        if [ "$1" == "x86_64" ]; then
-            arch="x86_64"
-        elif [ "$1" == "aarch64" ]; then
-            arch="aarch64"
-        elif [ "$1" == "armv7l" ]; then
-            arch="armv7l"
-        else
-            echo "Unsupported architecture: $1"
-            exit 1
-        fi
-    else
-        echo "Unsupported operating system: $(uname -s)"
-        exit 1
-    fi
-
-    # Create the 'appdir' directory if it doesn't exist
-    if [ ! -d "$miniconda_path" ]; then
-        mkdir "$miniconda_path"
-
-        # Download and install Miniconda based on OS and architecture
-        if [ "$(uname -s)" == "Darwin" ]; then
-            wget "https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-$arch.sh" -O miniconda_installer.sh
-        elif [ "$(uname -s)" == "Linux" ]; then
-            wget "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-$arch.sh" -O miniconda_installer.sh
-        fi
-
-        bash miniconda_installer.sh -b -u -p "$miniconda_path"
-        rm miniconda_installer.sh
-    else
-        echo "Using existing Miniconda in '$miniconda_path'."
-    fi
-}
-
-# Check if Miniconda is already installed
-if ! command -v "$miniconda_path/bin/conda" &>/dev/null; then
-    install_miniconda "$(arch | cut -d_ -f2)"
+# Check if the required environment variables are set
+if [ -z "$PREM_PYTHON" ]; then
+    echo "Please set the required PREM_PYTHON environment variable."
+    echo "Example: export PREM_PYTHON=appdir/envs/prem_env/bin/python"
+    exit 1
 fi
 
-# Check if the Miniconda environment 'prem_env' exists
-if [ ! -d "$miniconda_path/envs/$conda_env" ]; then
-    "$miniconda_path/bin/conda" create -n "$conda_env" python=3.11 -y
-fi
+# Define the paths based on environment variables
+python_exec="${PREM_PYTHON:-python}"
 
 # Clone the GitHub repository if not already present
 if [ ! -d "prem-services" ]; then
@@ -72,13 +18,13 @@ if [ ! -d "prem-services" ]; then
     git -C prem-services sparse-checkout set --no-cone cht-petals
     git -C prem-services checkout
 else
-    echo "Using existing 'prem-services' directory."
+    echo "Using the existing 'prem-services' directory."
 fi
 
-# Install requirements in the Miniconda environment
-"$miniconda_path/envs/$conda_env/bin/pip" install -r prem-services/cht-petals/requirements.txt
+# Install requirements using the specified Python binary
+"$python_exec" -m pip install -r prem-services/cht-petals/requirements.txt
 
-# Check for the --model_path and --dht_prefix parameters and run main.py
+# Check for the --model_path, --dht_prefix, and --port parameters and run main.py
 while [[ $# -gt 0 ]]; do
     case "$1" in
         --model_path)
@@ -89,16 +35,20 @@ while [[ $# -gt 0 ]]; do
             dht_prefix="$2"
             shift 2
             ;;
+        --port)
+            port="$2"
+            shift 2
+            ;;
         *)
             shift
             ;;
     esac
 done
 
-if [ -n "$model_path" ] && [ -n "$dht_prefix" ]; then
+if [ -n "$model_path" ] && [ -n "$dht_prefix" ] && [ -n "$port" ]; then
     export PYTHONPATH="$(pwd)/prem-services"
-    "$miniconda_path/envs/$conda_env/bin/python" prem-services/cht-petals/main.py --model_path "$model_path" --dht_prefix "$dht_prefix"
+    "$python_exec" prem-services/cht-petals/main.py --model_path "$model_path" --dht_prefix "$dht_prefix" --port $port
 else
-    echo "Please provide the --model_path parameter with the path to model directory and the --dht_prefix parameter for the DHT prefix."
+    echo "Please provide the --model_path parameter with the path to the model directory, the --dht_prefix parameter for the DHT prefix, and the --port parameter for the port number."
     exit 1
 fi

From 771920548697ce7a67351b5d9c5fbd9d346f3f52 Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Mon, 23 Oct 2023 19:27:47 +0100
Subject: [PATCH 06/12] misc tidy

---
 cht-petals/main.py         |  9 +++---
 cht-petals/models.py       |  6 ++--
 cht-petals/setup-petals.sh | 61 +++++++-------------------------------
 3 files changed, 19 insertions(+), 57 deletions(-)

diff --git a/cht-petals/main.py b/cht-petals/main.py
index 7c520ae..126202f 100644
--- a/cht-petals/main.py
+++ b/cht-petals/main.py
@@ -1,5 +1,6 @@
 import argparse
 import logging
+import os
 
 import uvicorn
 from dotenv import load_dotenv
@@ -9,12 +10,12 @@
 
 load_dotenv()
 
-MODEL_PATH = "./models"
-DHT_PREFIX = "StableBeluga2"
+MODEL_PATH = os.getenv("MODEL_PATH", "./models")
+DHT_PREFIX = os.getenv("DHT_PREFIX", "StableBeluga2")
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", help="Path to Model files directory", default=MODEL_PATH)
-    parser.add_argument("--dht_prefix", help="DHT prefix to use")
+    parser.add_argument("--model-path", help="Path to Model files directory", default=MODEL_PATH)
+    parser.add_argument("--dht-prefix", help="DHT prefix to use")
     parser.add_argument("--port", help="Port to run model server on", type=int)
     args = parser.parse_args()
     MODEL_PATH = args.model_path
diff --git a/cht-petals/models.py b/cht-petals/models.py
index f80849a..8d004de 100644
--- a/cht-petals/models.py
+++ b/cht-petals/models.py
@@ -55,13 +55,13 @@ def generate(
         return [cls.tokenizer.decode(outputs[0])]
 
     @classmethod
-    def get_model(cls, model_path: str, dht_prefix: str):
+    def get_model(cls, model_path: str = "./models", dht_prefix: str = "StableBeluga2"):
         if cls.model is None:
-            Tokenizer = LlamaTokenizer if "llama" in os.getenv("MODEL_ID", model_path).lower() else AutoTokenizer
+            Tokenizer = LlamaTokenizer if "llama" in model_path.lower() else AutoTokenizer
             cls.tokenizer = Tokenizer.from_pretrained(os.getenv("MODEL_ID", model_path))
             cls.model = AutoDistributedModelForCausalLM.from_pretrained(
                 os.getenv("MODEL_ID", model_path),
                 torch_dtype=torch.float32,
-                dht_prefix=os.getenv("DHT_PREFIX", dht_prefix),
+                dht_prefix=dht_prefix,
             )
         return cls.model
diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
index 38336bb..51314af 100755
--- a/cht-petals/setup-petals.sh
+++ b/cht-petals/setup-petals.sh
@@ -1,54 +1,15 @@
 #!/bin/bash
-# ./setup-petals.sh --model_path ./models/models--petals-team--StableBeluga2 --dht_prefix StableBeluga2-hf --port 8794
+# ./setup-petals.sh --model-path ./models/models--petals-team--StableBeluga2 --dht-prefix StableBeluga2-hf --port 8794
 
-# Check if the required environment variables are set
-if [ -z "$PREM_PYTHON" ]; then
-    echo "Please set the required PREM_PYTHON environment variable."
-    echo "Example: export PREM_PYTHON=appdir/envs/prem_env/bin/python"
-    exit 1
-fi
+tmpdir="$(mktemp)"
 
-# Define the paths based on environment variables
-python_exec="${PREM_PYTHON:-python}"
+# 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"
+# run server
+PYTHONPATH="$tmpdir/cht-petals" "${PREM_PYTHON:-python}" "$tmpdir/cht-petals/main.py" "$@"
 
-# Clone the GitHub repository if not already present
-if [ ! -d "prem-services" ]; then
-    # only clone the required directory - https://stackoverflow.com/a/52269934
-    git clone -n --depth=1 --filter=tree:0 https://github.com/premAI-io/prem-services.git
-    git -C prem-services sparse-checkout set --no-cone cht-petals
-    git -C prem-services checkout
-else
-    echo "Using the existing 'prem-services' directory."
-fi
-
-# Install requirements using the specified Python binary
-"$python_exec" -m pip install -r prem-services/cht-petals/requirements.txt
-
-# Check for the --model_path, --dht_prefix, and --port parameters and run main.py
-while [[ $# -gt 0 ]]; do
-    case "$1" in
-        --model_path)
-            model_path="$2"
-            shift 2
-            ;;
-        --dht_prefix)
-            dht_prefix="$2"
-            shift 2
-            ;;
-        --port)
-            port="$2"
-            shift 2
-            ;;
-        *)
-            shift
-            ;;
-    esac
-done
-
-if [ -n "$model_path" ] && [ -n "$dht_prefix" ] && [ -n "$port" ]; then
-    export PYTHONPATH="$(pwd)/prem-services"
-    "$python_exec" prem-services/cht-petals/main.py --model_path "$model_path" --dht_prefix "$dht_prefix" --port $port
-else
-    echo "Please provide the --model_path parameter with the path to the model directory, the --dht_prefix parameter for the DHT prefix, and the --port parameter for the port number."
-    exit 1
-fi
+rm -r "$tmpdir"

From 20edd1f4c855cc4df94c5b06e873dcba891bbef7 Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Mon, 23 Oct 2023 19:47:39 +0100
Subject: [PATCH 07/12] fix default --port

---
 cht-petals/download.py | 1 -
 cht-petals/main.py     | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/cht-petals/download.py b/cht-petals/download.py
index 52af9f4..64c07bf 100644
--- a/cht-petals/download.py
+++ b/cht-petals/download.py
@@ -9,7 +9,6 @@
 parser.add_argument("--model", help="Model to download")
 parser.add_argument("--dht-prefix", help="DHT prefix to use")
 args = parser.parse_args()
-
 print(f"Downloading model {args.model} with DHT prefix {args.dht_prefix}")
 
 
diff --git a/cht-petals/main.py b/cht-petals/main.py
index 126202f..da31143 100644
--- a/cht-petals/main.py
+++ b/cht-petals/main.py
@@ -16,7 +16,7 @@
     parser = argparse.ArgumentParser()
     parser.add_argument("--model-path", help="Path to Model files directory", default=MODEL_PATH)
     parser.add_argument("--dht-prefix", help="DHT prefix to use")
-    parser.add_argument("--port", help="Port to run model server on", type=int)
+    parser.add_argument("--port", help="Port to run model server on", type=int, default=8000)
     args = parser.parse_args()
     MODEL_PATH = args.model_path
     DHT_PREFIX = args.dht_prefix

From 9bec015737df12c231acaee646c6e23118aa5e10 Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Mon, 23 Oct 2023 19:53:33 +0100
Subject: [PATCH 08/12] fix git clone tmpdir

---
 cht-petals/setup-petals.sh | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
index 51314af..55ec8d8 100755
--- a/cht-petals/setup-petals.sh
+++ b/cht-petals/setup-petals.sh
@@ -1,15 +1,15 @@
 #!/bin/bash
 # ./setup-petals.sh --model-path ./models/models--petals-team--StableBeluga2 --dht-prefix StableBeluga2-hf --port 8794
 
-tmpdir="$(mktemp)"
+tmpdir="${PREM_APPDIR:-.}/$(uuid)"
 
 # 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 -r "$tmpdir"

From 003e668f81841398de2aea5ade29da03b88cd07c Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Mon, 23 Oct 2023 21:46:53 +0100
Subject: [PATCH 09/12] fix & expose --model-id

---
 cht-petals/main.py   |  4 +++-
 cht-petals/models.py | 14 ++++++++------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/cht-petals/main.py b/cht-petals/main.py
index da31143..795ddd8 100644
--- a/cht-petals/main.py
+++ b/cht-petals/main.py
@@ -10,12 +10,14 @@
 
 load_dotenv()
 
+MODEL_ID = os.getenv("MODEL_ID", "petals-team/StableBeluga2")
 MODEL_PATH = os.getenv("MODEL_PATH", "./models")
 DHT_PREFIX = os.getenv("DHT_PREFIX", "StableBeluga2")
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
+    parser.add_argument("--model-id", help="HuggingFace Model", default=MODEL_ID)
     parser.add_argument("--model-path", help="Path to Model files directory", default=MODEL_PATH)
-    parser.add_argument("--dht-prefix", help="DHT prefix to use")
+    parser.add_argument("--dht-prefix", help="DHT prefix to use", default=DHT_PREFIX)
     parser.add_argument("--port", help="Port to run model server on", type=int, default=8000)
     args = parser.parse_args()
     MODEL_PATH = args.model_path
diff --git a/cht-petals/models.py b/cht-petals/models.py
index 8d004de..7b54bad 100644
--- a/cht-petals/models.py
+++ b/cht-petals/models.py
@@ -1,4 +1,3 @@
-import os
 from abc import ABC, abstractmethod
 from typing import List
 
@@ -55,13 +54,16 @@ def generate(
         return [cls.tokenizer.decode(outputs[0])]
 
     @classmethod
-    def get_model(cls, model_path: str = "./models", dht_prefix: str = "StableBeluga2"):
+    def get_model(
+        cls,
+        model_path: str = "./models",
+        dht_prefix: str = "StableBeluga2",
+        model_id: str = "petals-team/StableBeluga2",
+    ):
         if cls.model is None:
             Tokenizer = LlamaTokenizer if "llama" in model_path.lower() else AutoTokenizer
-            cls.tokenizer = Tokenizer.from_pretrained(os.getenv("MODEL_ID", model_path))
+            cls.tokenizer = Tokenizer.from_pretrained(model_id)
             cls.model = AutoDistributedModelForCausalLM.from_pretrained(
-                os.getenv("MODEL_ID", model_path),
-                torch_dtype=torch.float32,
-                dht_prefix=dht_prefix,
+                model_id, torch_dtype=torch.float32, dht_prefix=dht_prefix
             )
         return cls.model

From f3b18308ef87a64b2f6e1eb17a4300d11849e349 Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Mon, 23 Oct 2023 21:51:26 +0100
Subject: [PATCH 10/12] force remove tmpdir

---
 cht-petals/setup-petals.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
index 55ec8d8..e030c3d 100755
--- a/cht-petals/setup-petals.sh
+++ b/cht-petals/setup-petals.sh
@@ -12,4 +12,4 @@ git -C "$tmpdir" checkout
 # run server
 PYTHONPATH="$tmpdir/cht-petals" "${PREM_PYTHON:-python}" "$tmpdir/cht-petals/main.py" "$@" || :
 
-rm -r "$tmpdir"
+rm -rf "$tmpdir"

From e7b98ee355e656e594f19ae1d1bde8fb56e4b296 Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Mon, 23 Oct 2023 22:03:39 +0100
Subject: [PATCH 11/12] safer interrupt handling

---
 cht-petals/setup-petals.sh | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
index e030c3d..f467df5 100755
--- a/cht-petals/setup-petals.sh
+++ b/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

From fc96cf8ab5e08aa3e4b74f96dd23c7a4246146c5 Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Mon, 23 Oct 2023 22:16:13 +0100
Subject: [PATCH 12/12] even safer traps

---
 cht-petals/setup-petals.sh | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/cht-petals/setup-petals.sh b/cht-petals/setup-petals.sh
index f467df5..f62aabb 100755
--- a/cht-petals/setup-petals.sh
+++ b/cht-petals/setup-petals.sh
@@ -1,11 +1,12 @@
 #!/usr/bin/env bash
 # Usage: setup-petals.sh [--model-path=<DIR>] [--dht-prefix=<PREFIX>] [--port=<INT>]
+set -eEuo pipefail
 
-tmpdir="${PREM_APPDIR:-.}/$(uuid)"
+tmpdir="${PREM_APPDIR:-.}/petals-$(uuid)"
 
 cleanup(){
   for i in $(jobs -p); do
-    kill -n 9 $i
+    kill -n 9 $i || :
   done
   rm -rf "$tmpdir"
   exit 0
@@ -13,6 +14,7 @@ cleanup(){
 
 trap "cleanup" SIGTERM
 trap "cleanup" SIGINT
+trap "cleanup" ERR
 
 # clone source
 git clone -n --depth=1 --filter=tree:0 https://github.com/premAI-io/prem-services.git "$tmpdir"