Secure remote Ollama API with Caddy Auth & Ngrok on Google Colab in one shot

Setup_secure_remote_ollama_server_on_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "123456789012"
      },
      "outputs": [],
      "source": [
        "# Phase 0: Ensure you have set NGROK_AUTHTOKEN in Colab \"Secrets\"\n",
        "\n",
        "# Phase 1: Installation of Ollama and Caddy\n",
        "# These commands will be executed in the Colab shell.\n",
        "print(\"Installing Ollama...\")\n",
        "!curl -fsSL https://ollama.com/install.sh | sh\n",
        "\n",
        "print(\"\\nInstalling Caddy...\")\n",
        "!curl -fsSL https://webi.sh/caddy | sh\n",
        "!source /root/.bashrc\n",
        "\n",
        "# Phase 2: Installation of Python dependencies\n",
        "print(\"\\nInstalling pyngrok...\")\n",
        "!pip install -qq pyngrok\n",
        "\n",
        "# Phase 3: Python script to configure and start services\n",
        "import os\n",
        "import subprocess\n",
        "import time\n",
        "import secrets\n",
        "import string\n",
        "import socket\n",
        "from pyngrok import ngrok\n",
        "from google.colab import userdata # To access Colab secrets\n",
        "\n",
        "# --- Configurable Variables ---\n",
        "# Allow user input, using a default if the input is empty.\n",
        "user_input_model = input(\"Input the model that runs on a single GPU (default: gemma3:4b-it-qat)... \")\n",
        "OLLAMA_MODEL_TO_PULL = user_input_model if user_input_model else \"gemma3:4b-it-qat\" # Model to download\n",
        "OLLAMA_PORT = 11434\n",
        "CADDY_PORT = 8081\n",
        "NGROK_PORT = str(CADDY_PORT)\n",
        "\n",
        "# --- Helper Functions ---\n",
        "def generate_ollama_key(length: int = 32) -> str:\n",
        "    \"\"\"Generates a secure random string suitable as an API key for Ollama.\"\"\"\n",
        "    characters = string.ascii_lowercase + string.digits\n",
        "    random_string = ''.join(secrets.choice(characters) for _ in range(length))\n",
        "    return f\"sk-ollama-{random_string}\"\n",
        "\n",
        "def wait_for_port(port: int, host: str = '127.0.0.1', timeout: int = 120, service_name: str = \"Service\") -> bool:\n",
        "    \"\"\"\n",
        "    Waits for a specific port to be open and listening on the specified host.\n",
        "    Uses socket connection attempts.\n",
        "    \"\"\"\n",
        "    print(f\"Waiting for {service_name} on port {port} to be available...\")\n",
        "    start_time = time.monotonic()\n",
        "    while time.monotonic() - start_time < timeout:\n",
        "        try:\n",
        "            with socket.create_connection((host, port), timeout=2): # Short timeout for each attempt\n",
        "                print(f\"OK: {service_name} is listening on port {port}.\")\n",
        "                return True\n",
        "        except (socket.error, ConnectionRefusedError, socket.timeout):\n",
        "            time.sleep(2)  # Wait 2 seconds before retrying\n",
        "    print(f\"ERROR: Timeout! {service_name} did not become available on port {port} within {timeout} seconds.\")\n",
        "    return False\n",
        "\n",
        "def start_ollama_server() -> subprocess.Popen:\n",
        "    \"\"\"Starts the Ollama server in the background.\"\"\"\n",
        "    print(\"Starting Ollama server...\")\n",
        "    try:\n",
        "        # Ensure Ollama is not already running or handle the error\n",
        "        # You might want to kill existing Ollama processes before starting a new one\n",
        "        # Example: !pkill -f \"ollama serve\"\n",
        "        process = subprocess.Popen(['ollama', 'serve'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
        "        print(\"Command 'ollama serve' executed.\")\n",
        "        return process\n",
        "    except FileNotFoundError:\n",
        "        print(\"ERROR: The 'ollama' command was not found. Ensure the installation was successful and it's in the PATH.\")\n",
        "        raise\n",
        "    except Exception as e:\n",
        "        print(f\"ERROR while starting Ollama: {e}\")\n",
        "        raise\n",
        "\n",
        "def create_caddy_log_directory():\n",
        "    \"\"\"Creates the directory for Caddy logs if it doesn't exist.\"\"\"\n",
        "    log_dir = \"/content/var/log/caddy\"\n",
        "    if not os.path.exists(log_dir):\n",
        "        print(f\"Creating directory for Caddy logs: {log_dir}\")\n",
        "        os.makedirs(log_dir)\n",
        "    else:\n",
        "        print(f\"Directory for Caddy logs already exists: {log_dir}\")\n",
        "\n",
        "def write_caddyfile(api_key: str) -> None:\n",
        "    \"\"\"Writes the Caddyfile with the provided API key.\"\"\"\n",
        "    caddyfile_content = f\"\"\":{CADDY_PORT} {{\n",
        "    # Define a matcher for authorized API access\n",
        "    @apiAuth {{\n",
        "        header Authorization \"Bearer {api_key}\"\n",
        "    }}\n",
        "\n",
        "    # Proxy for authorized requests to Ollama\n",
        "    reverse_proxy @apiAuth http://localhost:{OLLAMA_PORT} {{\n",
        "        header_up Host {{http.reverse_proxy.upstream.hostport}}\n",
        "        header_up X-Forwarded-Host {{host}}\n",
        "    }}\n",
        "\n",
        "    # Define a matcher for unauthorized access\n",
        "    @unauthorized {{\n",
        "        not {{\n",
        "            header Authorization \"Bearer {api_key}\"\n",
        "        }}\n",
        "    }}\n",
        "\n",
        "    # Respond to unauthorized requests\n",
        "    respond @unauthorized \"Unauthorized access. Please provide a valid API key in the 'Authorization: Bearer YOUR_API_KEY' header.\" 401 {{\n",
        "        close\n",
        "    }}\n",
        "\n",
        "    # Logging configuration\n",
        "    log {{\n",
        "        output file /content/var/log/caddy/access.log\n",
        "        level DEBUG\n",
        "    }}\n",
        "\n",
        "    # Security headers (optional, but recommended)\n",
        "    header {{\n",
        "        # Strict-Transport-Security (if using HTTPS with ngrok, even though Caddy here is HTTP)\n",
        "        Strict-Transport-Security \"max-age=31536000;\"\n",
        "        X-Content-Type-Options \"nosniff\"\n",
        "        X-Frame-Options \"DENY\"\n",
        "        Referrer-Policy \"strict-origin-when-cross-origin\"\n",
        "    }}\n",
        "}}\"\"\"\n",
        "    print(\"Writing Caddyfile to /content/Caddyfile...\")\n",
        "    with open(\"/content/Caddyfile\", \"w\") as f:\n",
        "        f.write(caddyfile_content)\n",
        "    print(\"Caddyfile written.\")\n",
        "\n",
        "def start_caddy_server() -> subprocess.Popen:\n",
        "    \"\"\"Starts the Caddy server in the background with the specified Caddyfile.\"\"\"\n",
        "    caddy_executable = \"/root/.local/bin/caddy\" # Installation path for webi.sh\n",
        "    if not os.path.exists(caddy_executable):\n",
        "        print(f\"ERROR: Caddy executable not found at {caddy_executable}. Verify installation.\")\n",
        "        # It might be in /root/.local/bin/caddy or another path if installed differently\n",
        "        # Try to find it with !which caddy\n",
        "        result = subprocess.run(['which', 'caddy'], capture_output=True, text=True)\n",
        "        if result.returncode == 0 and result.stdout.strip():\n",
        "            caddy_executable = result.stdout.strip()\n",
        "            print(f\"Found Caddy at: {caddy_executable}\")\n",
        "        else:\n",
        "            raise FileNotFoundError(f\"Caddy executable not found. Searched path: {caddy_executable}\")\n",
        "\n",
        "    print(f\"Starting Caddy server with {caddy_executable}...\")\n",
        "    try:\n",
        "        # Ensure Caddy is not already running or handle the error\n",
        "        # Example: !pkill -f \"caddy run\"\n",
        "        process = subprocess.Popen([caddy_executable, 'run', '--config', '/content/Caddyfile', '--adapter', 'caddyfile'],\n",
        "                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
        "        print(\"Command 'caddy run' executed.\")\n",
        "        return process\n",
        "    except Exception as e:\n",
        "        print(f\"ERROR while starting Caddy: {e}\")\n",
        "        raise\n",
        "\n",
        "def setup_ngrok_tunnel(port_to_forward: str, ngrok_auth_token: str) -> ngrok.NgrokTunnel:\n",
        "    \"\"\"Configures and starts an ngrok tunnel for the specified port.\"\"\"\n",
        "    print(f\"Configuring ngrok tunnel for port {port_to_forward}...\")\n",
        "    if not ngrok_auth_token:\n",
        "        raise ValueError(\"NGROK_AUTHTOKEN is not set. Set it in Colab Secrets.\")\n",
        "\n",
        "    ngrok.set_auth_token(ngrok_auth_token)\n",
        "\n",
        "    # Terminate any active ngrok tunnels on the same port or all tunnels\n",
        "    # to avoid conflicts if the cell is run multiple times.\n",
        "    try:\n",
        "        for t in ngrok.get_tunnels():\n",
        "            if t.config['addr'].endswith(f\":{port_to_forward}\"):\n",
        "                print(f\"Closing existing ngrok tunnel: {t.public_url}\")\n",
        "                ngrok.disconnect(t.public_url)\n",
        "        # ngrok.kill() # To close all ngrok processes, more aggressive\n",
        "    except Exception as e:\n",
        "        print(f\"Note: Error while closing existing ngrok tunnels (there might be none): {e}\")\n",
        "\n",
        "    try:\n",
        "        tunnel = ngrok.connect(port_to_forward, proto=\"http\", host_header=f\"localhost:{port_to_forward}\")\n",
        "        print(f\"ngrok tunnel created successfully!\")\n",
        "        return tunnel\n",
        "    except Exception as e:\n",
        "        print(f\"ERROR while creating ngrok tunnel: {e}\")\n",
        "        if \"authentication failed\" in str(e).lower():\n",
        "            print(\"Hint: Verify that your NGROK_AUTHTOKEN is correct and valid.\")\n",
        "        elif \"address already in use\" in str(e).lower() or \"port already bound\" in str(e).lower():\n",
        "            print(\"Hint: Another ngrok process might be running or the port is blocked. Try restarting the runtime.\")\n",
        "        raise\n",
        "\n",
        "# --- Main Flow ---\n",
        "if __name__ == \"__main__\":\n",
        "    ollama_api_key = \"\"\n",
        "    ngrok_public_url = \"\"\n",
        "    ollama_process = None\n",
        "    caddy_process = None\n",
        "\n",
        "    try:\n",
        "        # 1. Generate the API key for Ollama\n",
        "        ollama_api_key = generate_ollama_key()\n",
        "        print(f\"Ollama API key generated: {ollama_api_key}\")\n",
        "        # Export the API key as an environment variable for Caddy\n",
        "        os.environ['OLLAMA_API_KEY'] = ollama_api_key\n",
        "\n",
        "        # 2. Create the directory for Caddy logs\n",
        "        create_caddy_log_directory()\n",
        "\n",
        "        # 3. Write the Caddyfile\n",
        "        write_caddyfile(ollama_api_key)\n",
        "\n",
        "        # 4. Start the Ollama server and wait for it to be ready\n",
        "        ollama_process = start_ollama_server()\n",
        "        if not wait_for_port(OLLAMA_PORT, timeout=60, service_name=\"Ollama\"):\n",
        "            raise RuntimeError(f\"Ollama server did not start on port {OLLAMA_PORT}.\")\n",
        "\n",
        "        # 5. Start the Caddy server and wait for it to be ready\n",
        "        caddy_process = start_caddy_server()\n",
        "        if not wait_for_port(CADDY_PORT, timeout=30, service_name=\"Caddy\"):\n",
        "            raise RuntimeError(f\"Caddy server did not start on port {CADDY_PORT}.\")\n",
        "\n",
        "        # 6. Configure the ngrok tunnel\n",
        "        # Ensure you have set NGROK_AUTHTOKEN in Colab \"Secrets\"\n",
        "        # (left panel, key-shaped icon)\n",
        "        ngrok_auth_token = userdata.get('NGROK_AUTHTOKEN')\n",
        "        if not ngrok_auth_token:\n",
        "            print(\"ERROR: NGROK_AUTHTOKEN not found in Colab Secrets.\")\n",
        "            print(\"Please add your ngrok authentication token to Colab Secrets.\")\n",
        "            print(\"You can get one from https://dashboard.ngrok.com/get-started/your-authtoken\")\n",
        "            raise ValueError(\"NGROK_AUTHTOKEN missing.\")\n",
        "\n",
        "        ngrok_tunnel = setup_ngrok_tunnel(NGROK_PORT, ngrok_auth_token)\n",
        "        ngrok_public_url = ngrok_tunnel.public_url\n",
        "\n",
        "        print(\"\\n--- Configuration Completed Successfully! ---\")\n",
        "        print(f\"Public Ngrok URL (forwards to Caddy on port {CADDY_PORT}): {ngrok_public_url}\")\n",
        "        print(f\"Your Ollama API Key (to use in 'Authorization: Bearer ...' header): {ollama_api_key}\")\n",
        "        print(f\"Ollama is listening on: http://localhost:{OLLAMA_PORT} (accessible via Caddy/ngrok with the API key)\")\n",
        "        print(f\"Caddy is listening on: http://localhost:{CADDY_PORT} (exposed via ngrok)\")\n",
        "\n",
        "        # 7. Download an Ollama model (after the server is active)\n",
        "        print(f\"\\nDownloading Ollama model: {OLLAMA_MODEL_TO_PULL}...\")\n",
        "        # We use subprocess.run to wait for the pull to complete\n",
        "        pull_process = subprocess.run(['ollama', 'pull', OLLAMA_MODEL_TO_PULL], capture_output=True, text=True)\n",
        "        if pull_process.returncode == 0:\n",
        "            print(f\"Model {OLLAMA_MODEL_TO_PULL} downloaded successfully.\")\n",
        "            print(pull_process.stdout)\n",
        "        else:\n",
        "            print(f\"ERROR while downloading model {OLLAMA_MODEL_TO_PULL}:\")\n",
        "            print(pull_process.stderr)\n",
        "\n",
        "        print(\"\\nOllama and Caddy servers are running in the background.\")\n",
        "        print(\"You can interact with the Ollama API via the provided ngrok URL, using the API key.\")\n",
        "        print(\"To stop the servers, you will need to interrupt the execution of this cell or restart the runtime.\")\n",
        "        print(\"Caddy logs: /content/var/log/caddy/access.log\")\n",
        "        print(\"Ollama logs: Check the cell output or Ollama system logs if configured differently.\")\n",
        "        print(f\"\\nExample curl request to test (replace <YOUR_MODEL> if different from {OLLAMA_MODEL_TO_PULL}):\")\n",
        "        print(f\"curl -X POST {ngrok_public_url}/api/generate -H \\\"Authorization: Bearer {ollama_api_key}\\\" -d '{{\\\"model\\\": \\\"{OLLAMA_MODEL_TO_PULL}\\\", \\\"prompt\\\":\\\"Why is the sky blue?\\\", \\\"stream\\\": false}}'\")\n",
        "\n",
        "        #Remove comment if you open Colab Terminal\n",
        "        #print(f\"\\nOpen terminal to test example curl request and tail Caddy logs (/content/var/log/caddy/access.log)\")\n",
        "        #!pip install colab-xterm\n",
        "        #%load_ext colabxterm\n",
        "        #%xterm\n",
        "\n",
        "    except FileNotFoundError as e:\n",
        "        print(f\"CRITICAL ERROR - File not found: {e}. Installation might have failed.\")\n",
        "    except ValueError as e: # Specific for missing NGROK_AUTHTOKEN\n",
        "        print(f\"CONFIGURATION ERROR: {e}\")\n",
        "    except RuntimeError as e:\n",
        "        print(f\"RUNTIME ERROR: {e}\")\n",
        "    except Exception as e:\n",
        "        print(f\"UNEXPECTED ERROR: {e}\")"
      ]
    }
  ]
}