Python ContextManager to yield an Ollama client connected to a remote GPU

ephemeral_runpod.py

import os
import sys
import time
import subprocess
import contextlib
import functools

import tqdm

import requests
import runpod
import ollama

runpod.api_key = os.environ.get('RUNPOD_API_KEY')
PORT = 11434


class LaunchFailureException(Exception):
    """Raised when a pod or instance fails to launch properly"""

    def __init__(self, message):
        self.message = message
        super().__init__(self.message)


def with_retry(f, should_retry=lambda _: True,
               max_attempts=5, delay=2, backoff=False):
    """
    should_retry is a function that accepts the exception raised by f and
    returns a bool
    """

    for attempt in range(max_attempts):
        try:
            return f()
        except Exception as e:
            if attempt == max_attempts - 1:
                raise e

            if not should_retry(e):
                raise e

            print(
                f"Attempt {attempt + 1} failed. "
                f"Retrying in {delay} seconds..."
            )
            time.sleep(delay)
            if backoff:
                delay *= 2


def wait_for_pod_reported_ready(pod, ready_timeout_seconds=450):
    print("Waiting for pod: ", end="")
    start_t = time.time()
    while (time.time() - start_t) < ready_timeout_seconds:
        pod_info = runpod.get_pod(pod['id'])

        if not pod_info:
            print("-", end="")
            continue

        if pod_info['runtime']:
            print("\n")
            assert 'ports' in pod_info['runtime']
            return pod_info

        print(".", end="")
        sys.stdout.flush()
        time.sleep(1)

    raise LaunchFailureException(
        message=f"Pod failed to become ready. \n Pod ID: [{pod['id']}]"
    )


def pull_with_progress_bar(client, model):
    current_digest = ''
    progress_bars = {}

    # Note that streaming mode is necessary, otherwise the request gets timed
    # out, I think because of the proxying through Cloudflare (100 seconds).
    # See doc here:
    #   https://docs.runpod.io/pods/configuration/expose-ports
    for response in client.pull(model, stream=True):
        digest = response.digest or ''

        if not digest:
            print(response.status)
            continue

        if digest != current_digest and current_digest in progress_bars:
            progress_bars[current_digest].close()

        if digest not in progress_bars and response.total:
            progress_bars[digest] = tqdm.tqdm(
                total=response.total,
                desc=f"Pulling {digest[7:19]}",
                unit='B',
                unit_scale=True
            )

        if response.completed and digest in progress_bars:
            bar = progress_bars[digest]
            bar.update(response.completed - bar.n)

        current_digest = digest


def load_model_on_server(client, model):

    def load_model_on_server_inner():
        print(f"Loading model...")
        pull_with_progress_bar(client, model)

    with_retry(
        load_model_on_server_inner,
        should_retry=lambda e: (
            isinstance(e, ollama._types.ResponseError)
            or
            isinstance(e, httpx.RemoteProtocolError)
        ),
        backoff=True
    )


def wait_for_ollama_generate_ready(client, model):
    """
    Sometimes /api/generate seems to not be ready immediately after the pull.
    """
    def wait_for_ollama_generate_ready_inner():
        print("Waiting for /api/generate...")
        response = client.generate(model=model, prompt="check check 123")
        assert response.response

    with_retry(
        wait_for_ollama_generate_ready_inner,
        should_retry=lambda e: (
            isinstance(e, ollama._types.ResponseError)
            or
            isinstance(e, httpx.RemoteProtocolError)
        )
    )


@contextlib.contextmanager
def runpod_ollama_client(name, podspec="NVIDIA GeForce RTX 3090",
                         load_model=None, storage_gb=128):

    def create_pod():
        return \
            runpod.create_pod(
                name,
                "ollama/ollama",
                podspec,  # use the ids from `runpod.get_gpus()`
                volume_mount_path="/root/.ollama",
                ports=f"{PORT}/http",
                container_disk_in_gb=storage_gb,  # required even with template
                volume_in_gb=storage_gb  # required even with template
            )

    pod = \
        with_retry(
            create_pod,
            should_retry=lambda e: (
                isinstance(e, runpod.error.QueryError)
                and
                (
                    "does not have the resources to deploy your pod" in str(e)
                    or
                    "no longer any instances available" in str(e)
                )
            )
        )

    try:
        pod_info = wait_for_pod_reported_ready(pod)
        ollama_host_url = f"https://{pod_info['id']}-{PORT}.proxy.runpod.net"
        client = ollama.Client(host=ollama_host_url)

        # You need to load a model before you can get completions.
        if load_model is None:
            print('Proceeding without loading a model.')
            print('Probably this is a mistake.')
        else:
            load_model_on_server(client, load_model)
            wait_for_ollama_generate_ready(client, load_model)
            print("Verified /api/generate is ready!")

        yield client

    except LaunchFailureException as lfe:
        print("Looks like a bad machine.")
        raise lfe
    finally:
        runpod.stop_pod(pod['id'])
        runpod.terminate_pod(pod['id'])


def kill_all_pods():
    pods = runpod.get_pods()
    print(f"Found [{len(pods)}] pods to kill.")
    killed = 0

    for pod in pods:
        pod_id = pod["id"]
        try:
            print(f"Stopping pod {pod_id}...")
            runpod.stop_pod(pod_id)

            print(f"Terminating pod {pod_id}...")
            runpod.terminate_pod(pod_id)
            killed += 1
            print(f"Successfully terminated pod {pod_id}")
        except Exception as e:
            print(f"Error terminating pod {pod_id}: {str(e)}")

    print(f"Terminated [{killed}] pods.")


def tell_me_a_joke():
    with runpod_ollama_client("tinyllama-test", "NVIDIA GeForce RTX 3080",
                              load_model="tinyllama") as client:

        prompt = "Tell your favorite joke?"
        response = client.generate(model='tinyllama', prompt=prompt)
        print(response.response)

def tell_me_a_more_expensive_joke__llama3_3_70b_q8():
    model = "llama3.3:70b-instruct-q8_0"
    with runpod_ollama_client("llama3_3__70b__q8_test", "NVIDIA A100 80GB PCIe",
                              load_model=model, storage_gb=256) as client:
        prompt = """
        Tell me your favorite joke, but talk like a pirate.
        To be clear, the joke should not be even remotely nautical.
        No pirate jokes whatsoever.
        Tell me a normal, non-pirate joke.
        Just, you know, tell it like as if you *also* happened to be a pirate.
        """
        response = client.generate(model=model, prompt=prompt)
        print(response.response)


if __name__ == "__main__":
    tell_me_a_joke()