#!/bin/bash
# Ollama Service Manager Script
# Place this file in /etc/init.d/ollama and make it executable.

### BEGIN INIT INFO
# Provides:          ollama
# Required-Start:    $remote_fs $syslog
# Required-Stop:     $remote_fs $syslog
# Default-Start:     2 3 4 5
# Default-Stop:      0 1 6
# Short-Description: Manage Ollama service
# Description:       This script starts, stops, and restarts the Ollama service.
### END INIT INFO

SERVICE_NAME="Ollama"
OLLAMA_COMMAND="/bin/ollama" # Update with the path to the Ollama serve command
PID_FILE="/var/run/ollama.pid"
USER="root"
GROUP="root"

export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
export LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
export NVIDIA_DRIVER_CAPABILITIES=compute,utility
export NVIDIA_VISIBLE_DEVICES=all


export OLLAMA_HOST=0.0.0.0
export OLLAMA_MODELS=/data
export HOME=/root
export OLLAMA_KEEP_ALIVE=60m0s
export OLLAMA_LOAD_TIMEOUT=60m0s

start_service() {
    echo "Starting $SERVICE_NAME..."
    if [ -f "$PID_FILE" ]; then
        echo "$SERVICE_NAME is already running."
        return 1
    fi
    su -s /bin/bash -c "nohup $OLLAMA_COMMAND serve> /var/log/ollama.log 2>&1 & echo \$! > $PID_FILE" $USER
    echo "$SERVICE_NAME started with PID $(cat $PID_FILE)."
}

stop_service() {
    echo "Stopping $SERVICE_NAME..."
    if [ ! -f "$PID_FILE" ]; then
        echo "$SERVICE_NAME is not running."
        return 1
    fi
    kill $(cat "$PID_FILE")
    rm -f "$PID_FILE"
    echo "$SERVICE_NAME stopped."
}

restart_service() {
    echo "Restarting $SERVICE_NAME..."
    stop_service
    sleep 1
    start_service
}

case "$1" in
    start)
        start_service
        ;;
    stop)
        stop_service
        ;;
    restart)
        restart_service
        ;;
    *)
        echo "Usage: $0 {start|stop|restart}"
        exit 1
        ;;
esac

exit 0