vast_coder_spawner_ollama/vast_start.py

#!/bin/env python3
import requests
import time
import paramiko
import json
import os
import subprocess
model="qwen3-coder-next:q8_0"
vram=80000
max_bid=2.5
space=int(vram/700+20)

def load_api_key():
    path = os.path.join(
        os.environ["HOME"],
        ".config",
        "vastai",
        "vast_api_key"
    )

    with open(path, "r") as f:
        return f.read().strip()

def file_obj(path):
    with open(path, "r") as f:
        return json.load(f)

API_KEY = load_api_key()

BASE = "https://console.vast.ai/api/v0"
params = file_obj("search_params.json")
create_payload = file_obj("create_payload.json")
def vast_request(method, path, params=None, json_data=None):
    url = f"{BASE}/{path}"
    print(f"Request: {method} {url} params={params} json={json_data}")
    headers = {"Authorization": f"Bearer {API_KEY}"}
    r = requests.request(method, url, headers=headers, params=params, json=json_data)
    print(f"Response: {r.status_code} {r.text}")
    r.raise_for_status()
    return r.json()

def get_cheapest_gpus():
    params["gpu_ram"]["gt"] = vram
    r = vast_request("GET", "bundles/", params=params)
    offers = r["offers"]
    good_offers = []
    if not offers:
        raise Exception("No GPU offers found")
    for offer in offers:
        if offer["gpu_ram"] > vram \
            and offer["dph_total"] < max_bid \
            and offer["verification"] == "verified":
            good_offers.append(offer)
    if len(good_offers) == 0:
        raise Exception("No matching GPU offers found")
    return good_offers

def create_instance(offer_id):
    #create_payload['id'] = offer_id
    create_payload['env']['OLLAMA_MODEL']=model
    create_payload['disk']=space
    r = vast_request("PUT", f"asks/{offer_id}/", json_data=create_payload)

    return r["new_contract"]

def destry_instance(instance_id):
    r = requests.delete(
        f"{BASE}/instances/{instance_id}/",
        headers={"Authorization": f"Bearer {API_KEY}"}
    )

def get_instance(instance_id):
    r = requests.get(
        f"{BASE}/instances/{instance_id}/",
        headers={"Authorization": f"Bearer {API_KEY}"}
    )
    data = r.json()
    return data["instances"]


def wait_until_ready(instance_id):
    while True:
        data = get_instance(instance_id)
        for tag in [
            "actual_status", 
            "ssh_host", 
            "ssh_port", 
            "ssh_pwd", 
            "verification",
            "uptime_mins",
            "cur_state",
            "direct_port_count",
            "gpu_arch",
            "gpu_total_ram",
            "gpu_util",
            "intended_status",
            "next_state",
            "gpu_lanes",
            "gpu_name",
            "public_ipaddr",
            "ports",
            ]:
            try:
                if tag in data:
                    print(f"{tag}: {data[tag]}")
            except:
                print(f"{tag}: Error")
        print("Instance status:", data["actual_status"])
        #print("Instance data:", data)
        if data["actual_status"] == "running":
            return data

        print("Waiting for instance...\n\n")
        time.sleep(15)


def ssh_and_setup(ip, port, user, password):

    commands = [
        f"ollama pull {model}",
        "touch ~/.no_auto_tmux",
        #"ollama rm qwen3.5:35b",
        #"ollama pull lama3.1:70b",
    ]
    cmd2=["ssh","-t","-o","StrictHostKeyChecking=accept-new","-p",str(port),f"{user}@{ip}", "touch","~/.no_auto_tmux"]
    subprocess.run(cmd2)
    cmd2=["ssh","-t","-p",str(port),f"{user}@{ip}","ollama","pull",model]
    print(' '.join(cmd2))
    subprocess.run(cmd2)



def main():

    print("Searching cheapest 39GB+ GPU...")
    offers = get_cheapest_gpus()
    instance_id=0
    for offer in offers:
        offer_id = offer["id"]
        print("Offer:", offer_id, "GPU:", offer["gpu_name"], "Price:", offer["dph_total"])

        print("Creating instance...")
        try:
            for k in offer:
                v=offer[k]
                print(f"{k}: {v}")
            instance_id = create_instance(offer_id)
            break
        except:
            print(f"failed to claim offer {offer_id}")
    data=[]
    if instance_id > 0 :
        print("Instance:", instance_id)
        data = wait_until_ready(instance_id)
        ip = data["ssh_host"]
        port = data["ssh_port"]
        user = "root"
        #password = data["ssh_pwd"]

        print("Instance ready:", ip)


        pub_ip=data["public_ipaddr"]
        pub_port=data['ports']['22/tcp'][0]['HostPort']

        ssh_and_setup(ip, port, user, "")
        print(f"ssh root@{pub_ip} -p {pub_port}")
        ollama_port=data['ports']['11434/tcp'][0]['HostPort']
        print(f"ollama up at http://{pub_ip}:{ollama_port} with mode {model}")

        cmd2=["ssh","-t","-p",str(pub_port),f"{user}@{pub_ip}","ollama","run",model]
        print(' '.join(cmd2))
        subprocess.run(cmd2)


if __name__ == "__main__":
    main()