runpod_update_endpoint_image.py

import os
import json

import runpod
import requests

runpod.api_key = os.environ['RUNPOD_API_KEY']

def update_endpoint_image(endpoint_id, new_image_name):
  endpoint_with_template_response = run_graphql_query(
    fetch_endpoint_with_template_query,
    {'id': endpoint_id})

  template = endpoint_with_template_response['data']['myself']['endpoint']['template']

  template["imageName"] = new_image_name

  run_graphql_query(save_template_query, { 'input': template })

fetch_endpoint_with_template_query = """
query($id: String!) {
  myself {
    endpoint(id: $id) {
      template {
        advancedStart
        containerDiskInGb
        containerRegistryAuthId
        dockerArgs
        env {
          key
          value
        }
        id
        imageName
        isPublic
        isServerless
        name
        ports
        readme
        startJupyter
        startScript
        startSsh
        volumeInGb
        volumeMountPath
        config
        category
      }
    }
  }
}
"""

save_template_query = """
mutation saveTemplate($input: SaveTemplateInput) {
  saveTemplate(input: $input) {
    id
    __typename
  }
}
"""

# modified version of runpod.api.graphql.run_graphql_query, which sadly sucks.
# it doesn't have a parameter to pass variables through the API. instead they
# use string concatenation to inject variables.
# seems like a horrific idea.
def run_graphql_query(query, variables, api_key = None):
    from runpod import api_key as global_api_key
    effective_api_key = api_key or global_api_key
    api_url_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io")
    url = f"{api_url_base}/graphql"
    headers = {
        "Content-Type": "application/json",
        "User-Agent": runpod.user_agent.USER_AGENT,
        "Authorization": f"Bearer {effective_api_key}",
    }
    data = json.dumps({"query": query, "variables": variables})
    response = requests.post(url, headers=headers, data=data, timeout=30)
    if "errors" in response.json():
        raise Exception(response.json()["errors"][0]["message"])
    return response.json()

image_name = "${{ inputs.image_name }}"
endpoint_id = "${{ inputs.endpoint_id }}"
update_endpoint_image(endpoint_id, image_name)

EOF