wip: add spot pod creation

- Add support for creating spot pods using the RunPod API.
- Implement `create_spot_pod` function in the `RunPod` client.
- Refactor RunPod client to handle different query types and response structures.
- Improve error handling and logging for GraphQL requests.
- Update example to demonstrate spot pod creation.

Co-authored-by: mahmmoud.hassanein <mahmmoud.hassanein@gmail.com>
This commit is contained in:
2025-01-20 21:34:16 +02:00
parent d54a1e5a34
commit 4422d67701
4 changed files with 180 additions and 138 deletions

View File

@@ -10,12 +10,40 @@ mut rp := runpod.get_or_create(
)! )!
// Create a new pod // Create a new pod
pod_response := rp.create_pod( // pod_response := rp.create_on_demand_pod(
name: 'RunPod Tensorflow' // name: 'RunPod Tensorflow'
image_name: 'runpod/tensorflow' // image_name: 'runpod/tensorflow'
env: [ // env: [
{"JUPYTER_PASSWORD": "rn51hunbpgtltcpac3ol"} // runpod.EnvironmentVariableInput{
// key: 'JUPYTER_PASSWORD'
// value: 'rn51hunbpgtltcpac3ol'
// },
// ]
// )!
// println('Created pod with ID: ${pod_response.id}')
// create spot pod
// "mutation { podRentInterruptable( input: { bidPerGpu: 0.2, cloudType: SECURE, gpuCount: 1, volumeInGb: 40, containerDiskInGb: 40, minVcpuCount: 2, minMemoryInGb: 15, gpuTypeId: \"NVIDIA RTX A6000\", name: \"RunPod Pytorch\", imageName: \"runpod/pytorch\", dockerArgs: \"\", ports: \"8888/http\", volumeMountPath: \"/workspace\", env: [{ key: \"JUPYTER_PASSWORD\", value: \"vunw9ybnzqwpia2795p2\" }] } ) { id imageName env machineId machine { podHostId } } }"
spot_pod_resp := rp.create_spot_pod(
port: 1826
bid_per_gpu: 0.2
cloud_type: .secure
gpu_count: 1
volume_in_gb: 40
container_disk_in_gb: 40
min_vcpu_count: 2
min_memory_in_gb: 15
gpu_type_id: 'NVIDIA RTX A6000'
name: 'RunPod Pytorch'
image_name: 'runpod/pytorch'
docker_args: ''
ports: '8888/http'
volume_mount_path: '/workspace'
env: [
runpod.EnvironmentVariableInput{
key: 'JUPYTER_PASSWORD'
value: 'rn51hunbpgtltcpac3ol'
},
] ]
)! )!
println('Created spot pod with ID: ${spot_pod_resp.id}')
println('Created pod with ID: ${pod_response.id}')

View File

@@ -28,19 +28,25 @@ struct GqlQuery {
} }
// GraphQL response wrapper // GraphQL response wrapper
struct GqlResponse { struct GqlResponse[T] {
data GqlResponseData data map[string]T
} }
struct GqlResponseData { // struct GqlResponseData[T] {
pod_find_and_deploy_on_demand PodFindAndDeployOnDemandResponse @[json: 'podFindAndDeployOnDemand'] // pod_find_and_deploy_on_demand T @[json: 'podFindAndDeployOnDemand']
} // }
fn (mut rp RunPod) create_pod_request[T, R](request T, response R) !R { fn (mut rp RunPod) create_pop_find_and_deploy_on_demand_request(request PodFindAndDeployOnDemandRequest) !PodFindAndDeployOnDemandResponse {
gql := rp.build_query[T, R](request, response) gql := build_query(BuildQueryArgs{
query_type: .mutation
method_name: 'podFindAndDeployOnDemand'
}, request, PodFindAndDeployOnDemandResponse{})
println('gql: ${gql}') println('gql: ${gql}')
response_ := rp.make_request[GqlResponse](.post, '/graphql', gql)! response_ := rp.make_request[GqlResponse[PodFindAndDeployOnDemandResponse]](.post,
'/graphql', gql)!
println('response: ${json.encode(response_)}') println('response: ${json.encode(response_)}')
return response return response_.data['podFindAndDeployOnDemand'] or {
return error('Could not find podFindAndDeployOnDemand in response data: ${response_.data}')
}
// return response.data.pod_find_and_deploy_on_demand // return response.data.pod_find_and_deploy_on_demand
} }

View File

@@ -23,7 +23,7 @@ pub mut:
base_url string = 'https://api.runpod.io/v1' base_url string = 'https://api.runpod.io/v1'
} }
enum CloudType { pub enum CloudType {
all all
secure secure
community community
@@ -47,34 +47,41 @@ fn (ct CloudType) to_string() string {
@[params] @[params]
pub struct PodFindAndDeployOnDemandRequest { pub struct PodFindAndDeployOnDemandRequest {
pub mut: pub mut:
cloud_type CloudType = .all @[json: 'cloudType'] cloud_type CloudType = .all @[json: 'cloudType']
gpu_count int = 1 @[json: 'gpuCount'] gpu_count int = 1 @[json: 'gpuCount']
volume_in_gb int = 40 @[json: 'volumeInGb'] volume_in_gb int = 40 @[json: 'volumeInGb']
container_disk_in_gb int = 40 @[json: 'containerDiskInGb'] container_disk_in_gb int = 40 @[json: 'containerDiskInGb']
min_vcpu_count int = 2 @[json: 'minVcpuCount'] min_vcpu_count int = 2 @[json: 'minVcpuCount']
min_memory_in_gb int = 15 @[json: 'minMemoryInGb'] min_memory_in_gb int = 15 @[json: 'minMemoryInGb']
gpu_type_id string = 'NVIDIA RTX A6000' @[json: 'gpuTypeId'] gpu_type_id string = 'NVIDIA RTX A6000' @[json: 'gpuTypeId']
name string = 'RunPod Tensorflow' @[json: 'name'] name string = 'RunPod Tensorflow' @[json: 'name']
image_name string = 'runpod/tensorflow' @[json: 'imageName'] image_name string = 'runpod/tensorflow' @[json: 'imageName']
docker_args string = '' @[json: 'dockerArgs'] docker_args string = '' @[json: 'dockerArgs']
ports string = '8888/http' @[json: 'ports'] ports string = '8888/http' @[json: 'ports']
volume_mount_path string = '/workspace' @[json: 'volumeMountPath'] volume_mount_path string = '/workspace' @[json: 'volumeMountPath']
env []map[string]string = [] @[json: 'env'] env []EnvironmentVariableInput @[json: 'env']
}
pub struct EnvironmentVariableInput {
pub:
key string
value string
} }
// Represents the nested machine structure in the response // Represents the nested machine structure in the response
struct Machine { pub struct Machine {
pub:
pod_host_id string @[json: 'podHostId'] pod_host_id string @[json: 'podHostId']
} }
// Response structure for the mutation // Response structure for the mutation
pub struct PodFindAndDeployOnDemandResponse { pub struct PodFindAndDeployOnDemandResponse {
pub: pub:
id string @[json: 'id'] id string @[json: 'id']
image_name string @[json: 'imageName'] image_name string @[json: 'imageName']
env []map[string]string @[json: 'env'] env []string @[json: 'env']
machine_id int @[json: 'machineId'] machine_id int @[json: 'machineId']
machine Machine @[json: 'machine'] machine Machine @[json: 'machine']
} }
// new creates a new RunPod client // new creates a new RunPod client
@@ -88,76 +95,58 @@ pub fn new(api_key string) !&RunPod {
} }
// create_endpoint creates a new endpoint // create_endpoint creates a new endpoint
pub fn (mut rp RunPod) create_pod(pod PodFindAndDeployOnDemandRequest) !PodFindAndDeployOnDemandResponse { pub fn (mut rp RunPod) create_on_demand_pod(pod PodFindAndDeployOnDemandRequest) !PodFindAndDeployOnDemandResponse {
response_type := PodFindAndDeployOnDemandResponse{} response_type := PodFindAndDeployOnDemandResponse{}
request_type := pod request_type := pod
response := rp.create_pod_request[PodFindAndDeployOnDemandRequest, PodFindAndDeployOnDemandResponse](request_type, response := rp.create_pop_find_and_deploy_on_demand_request(request_type)!
response_type)!
return response return response
} }
// // list_endpoints lists all endpoints @[params]
// pub fn (mut rp RunPod) list_endpoints() ![]Endpoint { pub struct PodRentInterruptableInput {
// response := rp.list_endpoints_request()! pub mut:
// endpoints := json.decode([]Endpoint, response) or { port int @[json: 'port']
// return error('Failed to parse endpoints from response: ${response}') network_volume_id string @[json: 'networkVolumeId'; omitempty]
// } start_jupyter bool @[json: 'startJupyter']
// return endpoints start_ssh bool @[json: 'startSsh']
// } bid_per_gpu f32 @[json: 'bidPerGpu']
cloud_type CloudType @[json: 'cloudType']
container_disk_in_gb int @[json: 'containerDiskInGb']
country_code string @[json: 'countryCode'; omitempty]
docker_args string @[json: 'dockerArgs'; omitempty]
env []EnvironmentVariableInput @[json: 'env']
gpu_count int @[json: 'gpuCount']
gpu_type_id string @[json: 'gpuTypeId'; omitempty]
image_name string @[json: 'imageName'; omitempty]
min_disk int @[json: 'minDisk']
min_download int @[json: 'minDownload']
min_memory_in_gb int @[json: 'minMemoryInGb']
min_upload int @[json: 'minUpload']
min_vcpu_count int @[json: 'minVcpuCount']
name string @[json: 'name'; omitempty]
ports string @[json: 'ports'; omitempty]
stop_after string @[json: 'stopAfter'; omitempty]
support_public_ip bool @[json: 'supportPublicIp']
template_id string @[json: 'templateId'; omitempty]
terminate_after string @[json: 'terminateAfter'; omitempty]
volume_in_gb int @[json: 'volumeInGb']
volume_key string @[json: 'volumeKey'; omitempty]
volume_mount_path string @[json: 'volumeMountPath'; omitempty]
data_center_id string @[json: 'dataCenterId'; omitempty]
cuda_version string @[json: 'cudeVersion'; omitempty]
allowed_cuda_versions []string @[json: 'allowedCudaVersions']
}
// // get_endpoint gets an endpoint by ID pub fn (mut rp RunPod) create_spot_pod(input PodRentInterruptableInput) !PodFindAndDeployOnDemandResponse {
// pub fn (mut rp RunPod) get_endpoint(id string) !Endpoint { gql := build_query(BuildQueryArgs{
// response := rp.get_endpoint_request(id)! query_type: .mutation
// endpoint := json.decode(Endpoint, response) or { method_name: 'podRentInterruptable'
// return error('Failed to parse endpoint from response: ${response}') }, input, PodFindAndDeployOnDemandResponse{})
// } println('gql: ${gql}')
// return endpoint response_ := rp.make_request[GqlResponse[PodFindAndDeployOnDemandResponse]](.post,
// } '/graphql', gql)!
println('response: ${response_}')
// // delete_endpoint deletes an endpoint return response_.data['podRentInterruptable'] or {
// pub fn (mut rp RunPod) delete_endpoint(id string) ! { return error('Could not find podRentInterruptable in response data: ${response_.data}')
// rp.delete_endpoint_request(id)! }
// } }
// // list_gpu_instances lists available GPU instances
// pub fn (mut rp RunPod) list_gpu_instances() ![]GPUInstance {
// response := rp.list_gpu_instances_request()!
// instances := json.decode([]GPUInstance, response) or {
// return error('Failed to parse GPU instances from response: ${response}')
// }
// return instances
// }
// // run_on_endpoint runs a request on an endpoint
// pub fn (mut rp RunPod) run_on_endpoint(endpoint_id string, request RunRequest) !RunResponse {
// response := rp.run_on_endpoint_request(endpoint_id, request)!
// run_response := json.decode(RunResponse, response) or {
// return error('Failed to parse run response: ${response}')
// }
// return run_response
// }
// // get_run_status gets the status of a run
// pub fn (mut rp RunPod) get_run_status(endpoint_id string, run_id string) !RunResponse {
// response := rp.get_run_status_request(endpoint_id, run_id)!
// run_response := json.decode(RunResponse, response) or {
// return error('Failed to parse run status response: ${response}')
// }
// return run_response
// }
// // cfg_play configures a RunPod instance from heroscript parameters
// fn cfg_play(p paramsparser.Params) ! {
// mut rp := RunPod{
// name: p.get_default('name', 'default')!
// api_key: p.get('api_key')!
// base_url: p.get_default('base_url', 'https://api.runpod.io/v1')!
// }
// set(rp)!
// }
// fn obj_init(obj_ RunPod) !RunPod {
// // never call get here, only thing we can do here is work on object itself
// mut obj := obj_
// return obj
// }

View File

@@ -3,7 +3,7 @@ module runpod
import freeflowuniverse.herolib.core.httpconnection import freeflowuniverse.herolib.core.httpconnection
import json import json
fn (mut rp RunPod) get_field_name(field FieldData) string { fn get_field_name(field FieldData) string {
mut field_name := '' mut field_name := ''
// Process attributes to fetch the JSON field name or fallback to field name // Process attributes to fetch the JSON field name or fallback to field name
if field.attrs.len > 0 { if field.attrs.len > 0 {
@@ -20,53 +20,49 @@ fn (mut rp RunPod) get_field_name(field FieldData) string {
return field_name return field_name
} }
fn (mut rp RunPod) get_request_fields[T](struct_ T) string { fn get_request_fields[T](struct_ T) string {
// Start the current level // Start the current level
mut body_ := '{ ' mut body_ := '{ '
mut fields := []string{} mut fields := []string{}
$for field in T.fields { $for field in T.fields {
mut string_ := '' mut string_ := ''
string_ += rp.get_field_name(field) omit := 'omitempty' in field.attrs
string_ += ': ' mut empty_string := false
$if field.is_enum {
string_ += struct_.$(field.name).to_string()
}
$if field.typ is string { $if field.typ is string {
item := struct_.$(field.name) empty_string = struct_.$(field.name) == ''
string_ += "\"${item}\""
} }
if !omit || !empty_string {
string_ += get_field_name(field)
string_ += ': '
$if field.typ is int { $if field.is_enum {
item := struct_.$(field.name) string_ += struct_.$(field.name).to_string()
string_ += '${item}' } $else $if field.typ is string {
} item := struct_.$(field.name)
string_ += "\"${item}\""
// TODO: Loop only on the env map } $else $if field.is_array {
$if field.is_array { string_ += '['
for i in struct_.$(field.name) { for i in struct_.$(field.name) {
for k, v in i { string_ += get_request_fields(i)
string_ += '[{ '
string_ += "key: \"${k}\", value: \"${v}\""
string_ += ' }]'
} }
string_ += ']'
} $else $if field.is_struct {
string_ += get_request_fields(struct_.$(field.name))
} $else {
item := struct_.$(field.name)
string_ += '${item}'
} }
}
$if field.is_struct { fields << string_
rp.get_request_fields(struct_.$(field.name))
} }
fields << string_
} }
body_ += fields.join(', ') body_ += fields.join(', ')
body_ += ' }' body_ += ' }'
return body_ return body_
} }
fn (mut rp RunPod) get_response_fields[R](struct_ R) string { fn get_response_fields[R](struct_ R) string {
// Start the current level // Start the current level
mut body_ := '{ ' mut body_ := '{ '
@@ -74,9 +70,9 @@ fn (mut rp RunPod) get_response_fields[R](struct_ R) string {
$if field.is_struct { $if field.is_struct {
// Recursively process nested structs // Recursively process nested structs
body_ += '${field.name} ' body_ += '${field.name} '
body_ += rp.get_response_fields(struct_.$(field.name)) body_ += get_response_fields(struct_.$(field.name))
} $else { } $else {
body_ += rp.get_field_name(field) body_ += get_field_name(field)
body_ += ' ' body_ += ' '
} }
} }
@@ -84,16 +80,28 @@ fn (mut rp RunPod) get_response_fields[R](struct_ R) string {
return body_ return body_
} }
fn (mut rp RunPod) build_query[T, R](request T, response R) string { pub enum QueryType {
query
mutation
}
@[params]
pub struct BuildQueryArgs {
pub:
query_type QueryType // query or mutation
method_name string
}
fn build_query[T, R](args BuildQueryArgs, request T, response R) string {
// Convert input to JSON // Convert input to JSON
// input_json := json.encode(request) // input_json := json.encode(request)
// Build the GraphQL mutation string // Build the GraphQL mutation string
mut request_fields := rp.get_request_fields(request) mut request_fields := get_request_fields(request)
mut response_fields := rp.get_response_fields(response) mut response_fields := get_response_fields(response)
// Wrap the query correctly // Wrap the query correctly
query := 'mutation { podFindAndDeployOnDemand(input: ${request_fields}) ${response_fields} }' query := '${args.query_type.to_string()} { ${args.method_name}(input: ${request_fields}) ${response_fields} }'
// Wrap in the final structure // Wrap in the final structure
gql := GqlQuery{ gql := GqlQuery{
@@ -104,6 +112,17 @@ fn (mut rp RunPod) build_query[T, R](request T, response R) string {
return json.encode(gql) return json.encode(gql)
} }
fn (q QueryType) to_string() string {
return match q {
.query {
'query'
}
.mutation {
'mutation'
}
}
}
enum HTTPMethod { enum HTTPMethod {
get get
post post