connect_jupyter.sh

gistfile1.txt

#!/bin/bash

# Prompt user for the GPU name
read -p "Enter GPU name (e.g., c703i-gpu1): " GPU

# Prompt user for the virtual environment name
read -p "Enter virtual environment name: " ENV

# Function to clean up processes on exit
cleanup() {
    echo "Stopping Jupyter Lab on $GPU..."

    # Find and kill all Jupyter-related processes
    JUPYTER_PIDS=$(ssh $GPU "pgrep -f 'jupyter-lab'")
    if [[ -n "$JUPYTER_PIDS" ]]; then
        ssh $GPU "kill -9 $JUPYTER_PIDS"
        echo "Jupyter Lab stopped on $GPU."
    else
        echo "No Jupyter Lab processes found running on $GPU."
    fi

    # Close the SSH tunnel
    echo "Closing SSH tunnel..."
    kill $TUNNEL_PID 2>/dev/null
    exit 0
}

# Trap Ctrl+C to stop Jupyter properly
trap cleanup SIGINT

# Start Jupyter Lab on the remote GPU
echo "Connecting to $GPU and starting Jupyter Lab..."
ssh -T $GPU << EOF
    if pgrep -f "jupyter-lab" > /dev/null; then
        echo "Jupyter Lab is already running on $GPU."
    else
        echo "Activating environment: $ENV"
        source ~/venvs/$ENV/bin/activate

        echo "Starting Jupyter Lab..."
        nohup jupyter lab --no-browser --port=8888 --ip=0.0.0.0 > jupyter.log 2>&1 &
        echo "Jupyter Lab started on $GPU. Check jupyter.log if needed."
    fi
EOF

# Open SSH tunnel for port forwarding
echo "Setting up SSH tunnel..."
ssh -N -L 8888:127.0.0.1:8888 $GPU &
TUNNEL_PID=$!

echo "✅ Done! Open your browser and go to: http://localhost:8888"
echo "ℹ️ Use the token displayed in the remote terminal."

# Keep SSH tunnel open until user manually exits
echo "Press [CTRL+C] to stop Jupyter Lab and close the tunnel."
wait $TUNNEL_PID