Improved serving script

Added gemma supertoken model
Small bug persists with reading HF_TOKEN

app.py

@@ -9,11 +9,13 @@ import re
 import logging
 from typing import List, Dict, Any
 import gc
+import os
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+print("hf_toke_fromglobal", os.environ.get("HF_TOKEN"))
 # Model configurations - maps display names to HF model paths
 PREDEFINED_MODELS = [
      "meta-llama/Llama-3.2-1B",
@@ -25,8 +27,8 @@ PREDEFINED_MODELS = [
      "CohereForAI/aya-expanse-8b",
      "common-pile/comma-v0.1-2t",
      "google/byt5-small",
-     "google/byt5-small",
      "gsaltintas/supertoken_models-llama_gpt2",
+     "gsaltintas/supertoken_models-llama_google-gemma-2-2b"
 ]
 # Global cache for loaded models
 model_cache = {}
@@ -104,10 +106,10 @@ def load_model_and_tokenizer(model_path, use_cache=True, progress_callback=None)
         if progress_callback:
             progress_callback(0.1, f"🔄 Starting to load model: {model_path}")
         
-        logger.info(f"Loading model: {model_path}")
         
         # Check if CUDA is available
         device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Loading model: {model_path} using device: {device}")
         
         if progress_callback:
             progress_callback(0.2, f"📥 Loading tokenizer for {model_path}...")
@@ -122,6 +124,8 @@ def load_model_and_tokenizer(model_path, use_cache=True, progress_callback=None)
         if progress_callback:
             progress_callback(0.5, f"🧠 Loading model weights for {model_path}... (this may take a while)")
         
+        logger.info(os.getcwd())
+        logger.info("hf token", os.environ.get("HF_TOKEN"))
         # Load model with appropriate settings
         model = AutoModelForCausalLM.from_pretrained(
             model_path,

serve_on_killarney.sh

@@ -5,21 +5,24 @@ CLUSTER_HOST="killarney"
 CLUSTER_USER="gsa"  
 
 # Job configuration
+ACCOUNT="aip-craffel"
 SCRIPT_NAME="gradio_job.slurm"
-APP_PATH="/home/gsa/quick-tokenizer-accuracy/app.py"
+APP_DIR="/project/$ACCOUNT/$CLUSTER_USER/quick-tokenizer-accuracy"
+APP_PATH="app.py"
 JOB_NAME="gradio-app"
-PARTITION="l40s"
+GPU_TYPE="l40s"
+NUM_GPUS=1
 NODES=1
 NTASKS_PER_NODE=1
 CPUS_PER_TASK=4
 MEM="8G"
 TIME="02:00:00"
-GRADIO_PORT=7860
-ACCOUNT="aip-craffel"
-script_location="/project/aip-craffel/gsa/$SCRIPT_NAME"
+GRADIO_PORT=7861
+LOCAL_PORT=7861
+script_location="$APP_DIR/$SCRIPT_NAME"
 
-ENV_PATH="/home/gsa/tokenizers/.venv/bin/activate"
-OUTPUT_PATH="/project/aip-craffel/gsa/.slurm"
+ENV_PATH="/home/$CLUSTER_USER/tokenizers/.venv/bin/activate"
+OUTPUT_DIR="/project/$ACCOUNT/$CLUSTER_USER/.slurm"
 
 # Function to cleanup temporary files
 cleanup() {
@@ -37,14 +40,14 @@ trap cleanup EXIT INT TERM
 cat > "$SCRIPT_NAME" << EOF
 #!/bin/bash
 #SBATCH --job-name=$JOB_NAME
-#SBATCH --partition=$PARTITION
+#SBATCH --gres=gpu:$GPU_TYPE:$NUM_GPUS
 #SBATCH --nodes=$NODES
 #SBATCH --ntasks-per-node=$NTASKS_PER_NODE
 #SBATCH --cpus-per-task=$CPUS_PER_TASK
 #SBATCH --mem=$MEM
 #SBATCH --time=$TIME
 #SBATCH --account=$ACCOUNT
-#SBATCH --output=$OUTPUT_PATH/%j.out
+#SBATCH --output=$OUTPUT_DIR/%j.out
 
 # Print job info
 echo "Job started on node: \$(hostname)"
@@ -57,15 +60,15 @@ echo "Starting time: \$(date)"
 module load slurm/killarney/24.05.7 StdEnv/2023 gcc/13.3 openmpi/5.0.3 cuda/12.6 python/3.10.13
 
 # Activate virtual environment
-source $ENV_PATH
+source "${ENV_PATH}"
 
 # Set up environment
 export GRADIO_SERVER_NAME="0.0.0.0"
 export GRADIO_SERVER_PORT=$GRADIO_PORT
 
 # Start Gradio app
-echo "Starting Gradio app on port $GRADIO_PORT..."
-python $APP_PATH --no-browser
+echo "Starting Gradio app on port ${GRADIO_PORT}..."
+gradio "${APP_PATH}" --watch-dirs "${APP_DIR}"
 
 # Keep the job alive
 echo "Gradio app finished at: \$(date)"
@@ -81,7 +84,7 @@ if [ $? -ne 0 ]; then
 fi
 
 echo "Submitting job to cluster..."
-JOB_ID=$(ssh -t "$CLUSTER_USER@$CLUSTER_HOST" "bash -l -c 'cd /project && sbatch --parsable $script_location'")
+JOB_ID=$(ssh -t "$CLUSTER_USER@$CLUSTER_HOST" "bash -l -c 'cd $APP_DIR && sbatch --parsable $script_location'")
 
 if [ $? -ne 0 ]; then
     echo "Error: Failed to submit job to cluster"
@@ -122,7 +125,7 @@ done
 
 # Get the allocated node
 NODE=$(ssh "$CLUSTER_USER@$CLUSTER_HOST" "bash -l -c 'squeue -j $JOB_ID -h -o \"%N\"'")
-echo "Job is running on node: $NODE"
+echo "Job (${JOB_ID}) is running on node: ${NODE}"
 
 # Wait a moment for the Gradio app to start
 echo "Waiting for Gradio app to initialize..."
@@ -147,10 +150,50 @@ if [ -n "$GRADIO_CHECK" ]; then
 else
     echo "⚠ Warning: Gradio app may not have started properly"
     echo "Check the job output:"
-    ssh "$CLUSTER_USER@$CLUSTER_HOST" \
-        "bash -l -c 'tail ${JOB_ID}.out'"
+    ssh "$CLUSTER_USER@$CLUSTER_HOST" "bash -l -c 'tail \"${OUTPUT_DIR}/${JOB_ID}.out\"'"
 fi
 
+
+cancel_job() {
+    read -p "Would you like to cancel the job? (y/n): " -n 1 -r
+    if [[ $REPLY =~ ^[Yy]$ ]]; then
+        ## job id known only remotely
+        # ssh "$CLUSTER_USER@$CLUSTER_HOST" "bash -l -c 'scancel \${JOB_ID}'"
+        ssh "$CLUSTER_USER@$CLUSTER_HOST" "bash -l -c 'scancel  ${JOB_ID} '"
+        # ssh "$CLUSTER_USER@$CLUSTER_HOST" "bash -l -c 'scancel ${JOB_ID}'"
+    fi
+}
+
+# Optional port forwarding
+read -p "Would you like to set up port forwarding now? (y/n): " -n 1 -r
+echo ""
+if [[ $REPLY =~ ^[Yy]$ ]]; then
+    # ssh -L "${GRADIO_PORT}:${NODE}:${GRADIO_PORT}" "$CLUSTER_USER@$CLUSTER_HOST" \
+    # -t "echo 'Port forwarding active: localhost:${GRADIO_PORT} -> ${NODE}:${GRADIO_PORT}'; bash
+    # If GRADIO_PORT is in use locally, pick a random free port
+    if lsof -iTCP:"$GRADIO_PORT" -sTCP:LISTEN >/dev/null 2>&1; then
+        echo "Port $GRADIO_PORT is already in use locally — selecting a free one..."
+        LOCAL_PORT=$(comm -23 \
+            <(seq 1024 65535 | sort) \
+            <(lsof -nP -iTCP -sTCP:LISTEN | awk 'NR>1 {print $9}' | awk -F: '{print $NF}' | sort -u) \
+            | awk 'BEGIN{srand()} {ports[NR]=$0} END{print ports[int(rand()*NR)+1]}')
+    else
+        LOCAL_PORT="$GRADIO_PORT"
+    fi
+
+    echo "Using local port: $LOCAL_PORT"
+
+    echo "Setting up port forwarding... Open https://localhost:${LOCAL_PORT} in your browser to access the app."
+    ssh -L "${LOCAL_PORT}:${NODE}:${GRADIO_PORT}" "$CLUSTER_USER@$CLUSTER_HOST" \
+        -t "echo 'Port forwarding active: localhost:${LOCAL_PORT} -> ${NODE}:${GRADIO_PORT}'; bash"
+
+
+    echo ""
+    echo "Port forwarding ended."
+    cancel_job
+else
+    echo "Skipping port forwarding."
+
 # Connection info
 cat <<EOF
 
@@ -168,29 +211,17 @@ Alternative direct SSH with forwarding:
   ssh -L $GRADIO_PORT:localhost:$GRADIO_PORT $CLUSTER_USER@$NODE.$CLUSTER_HOST
 
 Check job status:
-  ssh $CLUSTER_USER@$CLUSTER_HOST 'squeue -j $JOB_ID '
+  ssh $CLUSTER_USER@$CLUSTER_HOST \"'squeue -j $JOB_ID '\"
 
 Cancel job:
-  ssh $CLUSTER_USER@$CLUSTER_HOST 'scancel $JOB_ID '
+  ssh $CLUSTER_USER@$CLUSTER_HOST \"'scancel $JOB_ID '\"
 =========================================
 
 EOF
-
-# Optional port forwarding
-read -p "Would you like to set up port forwarding now? (y/n): " -n 1 -r
-echo ""
-if [[ $REPLY =~ ^[Yy]$ ]]; then
-    echo "Setting up port forwarding..."
-    ssh -L "${GRADIO_PORT}:${NODE}:${GRADIO_PORT}" "$CLUSTER_USER@$CLUSTER_HOST" \
-    -t "echo 'Port forwarding active: localhost:${GRADIO_PORT} -> ${NODE}:${GRADIO_PORT}'; bash"
-    echo ""
-    echo "Port forwarding ended."
-else
-    echo "Skipping port forwarding."
     echo "Later you can run: ssh -L $GRADIO_PORT:$NODE:$GRADIO_PORT $CLUSTER_USER@$CLUSTER_HOST"
 fi
 
 echo ""
 echo "Job $JOB_ID is still running on $CLUSTER_HOST:$NODE"
-echo "Don't forget to cancel it when done: ssh $CLUSTER_USER@$CLUSTER_HOST 'scancel $JOB_ID'"
+# echo "Don't forget to cancel it when done: ssh $CLUSTER_USER@$CLUSTER_HOST 'scancel $JOB_ID'"