WIP: refactor RunPod client

- Refactor RunPod client to use a new GraphQL builder.
- This improves the readability and maintainability of the code.
- The old `build_query` function was removed, and the new
- `QueryBuilder` struct is now used.  This allows for a more
- flexible and extensible approach to constructing GraphQL
- queries.  The example in `runpod_example.vsh` is now
- commented out until the new GraphQL builder is fully
- implemented.

Co-authored-by: mariobassem12 <mariobassem12@gmail.com>
This commit is contained in:
Mahmoud Emad
2025-01-22 20:35:45 +02:00
parent 7486d561ec
commit 6f9d570a93
5 changed files with 340 additions and 188 deletions

View File

@@ -32,51 +32,51 @@ on_demand_pod_response := rp.create_on_demand_pod(
println('Created pod with ID: ${on_demand_pod_response.id}')
// create a spot pod
spot_pod_response := rp.create_spot_pod(
port: 1826
bid_per_gpu: 0.2
cloud_type: .secure
gpu_count: 1
volume_in_gb: 5
container_disk_in_gb: 5
min_vcpu_count: 1
min_memory_in_gb: 4
gpu_type_id: 'NVIDIA RTX 2000 Ada'
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_response.id}')
// // create a spot pod
// spot_pod_response := rp.create_spot_pod(
// port: 1826
// bid_per_gpu: 0.2
// cloud_type: .secure
// gpu_count: 1
// volume_in_gb: 5
// container_disk_in_gb: 5
// min_vcpu_count: 1
// min_memory_in_gb: 4
// gpu_type_id: 'NVIDIA RTX 2000 Ada'
// 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_response.id}')
// stop on-demand pod
stop_on_demand_pod := rp.stop_pod(
pod_id: '${on_demand_pod_response.id}'
)!
println('Stopped on-demand pod with ID: ${stop_on_demand_pod.id}')
// // stop on-demand pod
// stop_on_demand_pod := rp.stop_pod(
// pod_id: '${on_demand_pod_response.id}'
// )!
// println('Stopped on-demand pod with ID: ${stop_on_demand_pod.id}')
// stop spot pod
stop_spot_pod := rp.stop_pod(
pod_id: '${spot_pod_response.id}'
)!
println('Stopped spot pod with ID: ${stop_spot_pod.id}')
// // stop spot pod
// stop_spot_pod := rp.stop_pod(
// pod_id: '${spot_pod_response.id}'
// )!
// println('Stopped spot pod with ID: ${stop_spot_pod.id}')
// start on-demand pod
start_on_demand_pod := rp.start_on_demand_pod(pod_id: '${on_demand_pod_response.id}', gpu_count: 1)!
println('Started on demand pod with ID: ${start_on_demand_pod.id}')
// // start on-demand pod
// start_on_demand_pod := rp.start_on_demand_pod(pod_id: '${on_demand_pod_response.id}', gpu_count: 1)!
// println('Started on demand pod with ID: ${start_on_demand_pod.id}')
// start spot pod
start_spot_pod := rp.start_spot_pod(
pod_id: '${spot_pod_response.id}'
gpu_count: 1
bid_per_gpu: 0.2
)!
println('Started spot pod with ID: ${start_on_demand_pod.id}')
// // start spot pod
// start_spot_pod := rp.start_spot_pod(
// pod_id: '${spot_pod_response.id}'
// gpu_count: 1
// bid_per_gpu: 0.2
// )!
// println('Started spot pod with ID: ${start_on_demand_pod.id}')

View File

@@ -0,0 +1,82 @@
module runpod
enum OperationType {
query
mutation
}
struct QueryBuilder {
pub mut:
operation OperationType
fields []Field
variables map[string]string
}
struct Field {
name string
arguments map[string]string
sub_fields []Field
}
fn new_field(name string, args map[string]string, sub_fields []Field) Field {
return Field{
name: name
arguments: args
sub_fields: sub_fields
}
}
fn build_arguments(args map[string]string) string {
if args.len == 0 {
return ''
}
mut sb := ''
sb += '('
for key, value in args {
if value.len == 0 {
continue
}
sb += '${key}: ${value}, '
}
return sb.trim_right(', ') + ')'
}
fn build_fields(fields []Field) string {
mut sb := ' { '
for field in fields {
sb += field.name
if field.arguments.len > 0 {
sb += build_arguments(field.arguments)
}
if field.sub_fields.len > 0 {
sb += build_fields(field.sub_fields)
}
sb += ' '
}
sb += ' } '
return sb
}
fn (mut q QueryBuilder) add_operation(operation OperationType, fields []Field, variables map[string]string) {
q.operation = operation
q.fields = fields
q.variables = variables.clone()
}
fn (q QueryBuilder) build_query() string {
mut sb := ''
sb += '${q.operation}' + ' myOperation'
if q.variables.len > 0 {
sb += build_arguments(q.variables)
}
sb += build_fields(q.fields)
return sb
}

View File

@@ -1,21 +1,87 @@
module runpod
import x.json2
import json
// fn main() {
// mut fields := []Field{}
// fields << new_field('gpuTypes', {
// 'id': '"NVIDIA GeForce RTX 3090"'
// }, [
// new_field('displayName', {}, []),
// new_field('d', {}, []),
// new_field('communmemoryInGb', {}, []),
// new_field('secureClouityCloud', {}, []),
// new_field('lowestPrice', {
// 'gpuCount': '1'
// }, [
// new_field('minimumBidPrice', {}, []),
// new_field('uninterruptablePrice', {}, []),
// ]),
// ])
// // Create Query Builder
// mut builder := QueryBuilder{}
// builder.add_operation(.query, fields, {})
// // Build and print the query
// query := builder.build_query()
// println('query: ${query}')
// }
// #### Internally method doing a network call to create a new on-demand pod.
// - Build the required query based pn the input sent by the user and send the request.
// - Decode the response received from the API into two objects `Data` and `Error`.
// - The data field should contains the pod details same as `PodResult` struct.
// - The error field should contain the error message.
fn (mut rp RunPod) create_on_demand_pod_request(request PodFindAndDeployOnDemandRequest) !PodResult {
gql := build_query(
query_type: .mutation
method_name: 'podFindAndDeployOnDemand'
request_model: request
response_model: PodResult{}
)
response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
fn (mut rp RunPod) create_on_demand_pod_request(input PodFindAndDeployOnDemandRequest) !PodResult {
mut fields := []Field{}
mut machine_fields := []Field{}
mut output_fields := []Field{}
mut arguments := map[string]string{}
mut builder := QueryBuilder{}
// $for field in input.fields {
// // TODO: Handle option fields
// // TODO: Handle the skip chars \
// item := input.$(field.name)
// arguments[get_field_name(field)] = '${item}'
// }
machine_fields << new_field('podHostId', {}, [])
output_fields << new_field('id', {}, [])
output_fields << new_field('imageName', {}, [])
output_fields << new_field('env', {}, [])
output_fields << new_field('machineId', {}, [])
output_fields << new_field('desiredStatus', {}, [])
output_fields << new_field('machine', {}, machine_fields)
fields << new_field('podFindAndDeployOnDemand', {
'input': '\$arguments'
}, output_fields)
builder.add_operation(.mutation, fields, {
'\$arguments': 'PodFindAndDeployOnDemandInput'
})
query := builder.build_query()
encoded_input := json.encode(input)
decoded_input := json2.raw_decode(encoded_input)!.as_map()
mut q_map := map[string]json2.Any{}
mut variables := map[string]json2.Any{}
variables['arguments'] = decoded_input
q_map['query'] = json2.Any(query)
q_map['variables'] = json2.Any(variables)
q := json2.encode(q_map)
println('query: ${q}')
response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', q)!
return response.data['podFindAndDeployOnDemand'] or {
return error('Could not find "podFindAndDeployOnDemand" in response data: ${response.data}')
}
// return PodResult{}
}
// #### Internally method doing a network call to create a new spot pod.
@@ -24,16 +90,17 @@ fn (mut rp RunPod) create_on_demand_pod_request(request PodFindAndDeployOnDemand
// - The data field should contains the pod details same as `PodResult` struct.
// - The error field should contain the error message.
fn (mut rp RunPod) create_spot_pod_request(input PodRentInterruptableInput) !PodResult {
gql := build_query(
query_type: .mutation
method_name: 'podRentInterruptable'
request_model: input
response_model: PodResult{}
)
response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
return response.data['podRentInterruptable'] or {
return error('Could not find "podRentInterruptable" in response data: ${response.data}')
}
// gql := build_query(
// query_type: .mutation
// method_name: 'podRentInterruptable'
// request_model: input
// response_model: PodResult{}
// )
// response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
// return response.data['podRentInterruptable'] or {
// return error('Could not find "podRentInterruptable" in response data: ${response.data}')
// }
return PodResult{}
}
// #### Internally method doing a network call to start on demand pod.
@@ -42,16 +109,17 @@ fn (mut rp RunPod) create_spot_pod_request(input PodRentInterruptableInput) !Pod
// - The data field should contains the pod details same as `PodResult` struct.
// - The error field should contain the error message.
fn (mut rp RunPod) start_on_demand_pod_request(input PodResume) !PodResult {
gql := build_query(
query_type: .mutation
method_name: 'podResume'
request_model: input
response_model: PodResult{}
)
response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
return response.data['podResume'] or {
return error('Could not find "podResume" in response data: ${response.data}')
}
// gql := build_query(
// query_type: .mutation
// method_name: 'podResume'
// request_model: input
// response_model: PodResult{}
// )
// response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
// return response.data['podResume'] or {
// return error('Could not find "podResume" in response data: ${response.data}')
// }
return PodResult{}
}
// #### Internally method doing a network call to start spot pod.
@@ -60,16 +128,17 @@ fn (mut rp RunPod) start_on_demand_pod_request(input PodResume) !PodResult {
// - The data field should contains the pod details same as `PodResult` struct.
// - The error field should contain the error message.
fn (mut rp RunPod) start_spot_pod_request(input PodBidResume) !PodResult {
gql := build_query(
query_type: .mutation
method_name: 'podBidResume'
request_model: input
response_model: PodResult{}
)
response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
return response.data['podBidResume'] or {
return error('Could not find "podBidResume" in response data: ${response.data}')
}
// gql := build_query(
// query_type: .mutation
// method_name: 'podBidResume'
// request_model: input
// response_model: PodResult{}
// )
// response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
// return response.data['podBidResume'] or {
// return error('Could not find "podBidResume" in response data: ${response.data}')
// }
return PodResult{}
}
// #### Internally method doing a network call to stop a pod.
@@ -78,14 +147,15 @@ fn (mut rp RunPod) start_spot_pod_request(input PodBidResume) !PodResult {
// - The data field should contains the pod details same as `PodResult` struct.
// - The error field should contain the error message.
fn (mut rp RunPod) stop_pod_request(input PodResume) !PodResult {
gql := build_query(
query_type: .mutation
method_name: 'podStop'
request_model: input
response_model: PodResult{}
)
response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
return response.data['podStop'] or {
return error('Could not find "podStop" in response data: ${response.data}')
}
// gql := build_query(
// query_type: .mutation
// method_name: 'podStop'
// request_model: input
// response_model: PodResult{}
// )
// response := rp.make_request[GqlResponse[PodResult]](.post, '/graphql', gql)!
// return response.data['podStop'] or {
// return error('Could not find "podStop" in response data: ${response.data}')
// }
return PodResult{}
}

View File

@@ -29,7 +29,7 @@ pub enum CloudType {
community
}
fn (ct CloudType) to_string() string {
fn (ct CloudType) str() string {
return match ct {
.all {
'ALL'
@@ -44,7 +44,7 @@ fn (ct CloudType) to_string() string {
}
pub struct EnvironmentVariableInput {
pub:
pub mut:
key string
value string
}

View File

@@ -33,117 +33,117 @@ fn get_field_name(field FieldData) string {
return field_name
}
// Constructs JSON-like request fields from a struct.
fn get_request_fields[T](struct_ T) string {
mut body_ := '{ '
mut fields := []string{}
// // Constructs JSON-like request fields from a struct.
// fn get_request_fields[T](struct_ T) string {
// mut body_ := '{ '
// mut fields := []string{}
$for field in T.fields {
mut string_ := ''
omit := 'omitempty' in field.attrs
mut empty_string := false
$if field.typ is string {
empty_string = struct_.$(field.name) == ''
}
if !omit || !empty_string {
string_ += get_field_name(field)
string_ += ': '
// $for field in T.fields {
// mut string_ := ''
// omit := 'omitempty' in field.attrs
// mut empty_string := false
// $if field.typ is string {
// empty_string = struct_.$(field.name) == ''
// }
// if !omit || !empty_string {
// string_ += get_field_name(field)
// string_ += ': '
$if field.is_enum {
string_ += struct_.$(field.name).to_string()
} $else $if field.typ is string {
item := struct_.$(field.name)
string_ += "\"${item}\""
} $else $if field.is_array {
string_ += '['
for i in struct_.$(field.name) {
string_ += get_request_fields(i)
}
string_ += ']'
} $else $if field.is_struct {
string_ += get_request_fields(struct_.$(field.name))
} $else {
item := struct_.$(field.name)
string_ += '${item}'
}
// $if field.is_enum {
// string_ += struct_.$(field.name).to_string()
// } $else $if field.typ is string {
// item := struct_.$(field.name)
// string_ += "\"${item}\""
// } $else $if field.is_array {
// string_ += '['
// for i in struct_.$(field.name) {
// string_ += get_request_fields(i)
// }
// string_ += ']'
// } $else $if field.is_struct {
// string_ += get_request_fields(struct_.$(field.name))
// } $else {
// item := struct_.$(field.name)
// string_ += '${item}'
// }
fields << string_
}
}
body_ += fields.join(', ')
body_ += ' }'
return body_
}
// fields << string_
// }
// }
// body_ += fields.join(', ')
// body_ += ' }'
// return body_
// }
// Constructs JSON-like response fields for a given struct.
fn get_response_fields[R](struct_ R) string {
mut body_ := '{ '
// // Constructs JSON-like response fields for a given struct.
// fn get_response_fields[R](struct_ R) string {
// mut body_ := '{ '
$for field in R.fields {
$if field.is_struct {
// Recursively process nested structs
body_ += '${field.name} '
body_ += get_response_fields(struct_.$(field.name))
} $else {
body_ += get_field_name(field)
body_ += ' '
}
}
body_ += ' }'
return body_
}
// $for field in R.fields {
// $if field.is_struct {
// // Recursively process nested structs
// body_ += '${field.name} '
// body_ += get_response_fields(struct_.$(field.name))
// } $else {
// body_ += get_field_name(field)
// body_ += ' '
// }
// }
// body_ += ' }'
// return body_
// }
// Enum representing the type of GraphQL operation.
pub enum QueryType {
query
mutation
}
// // Enum representing the type of GraphQL operation.
// pub enum QueryType {
// query
// mutation
// }
// Struct for building GraphQL queries with request and response models.
@[params]
pub struct BuildQueryArgs[T, R] {
pub:
query_type QueryType // query or mutation
method_name string
request_model T @[required]
response_model R @[required]
}
// // Struct for building GraphQL queries with request and response models.
// @[params]
// pub struct BuildQueryArgs[T, R] {
// pub:
// query_type QueryType // query or mutation
// method_name string
// request_model T @[required]
// response_model R @[required]
// }
// Builds a GraphQL query or mutation string from provided arguments.
fn build_query[T, R](args BuildQueryArgs[T, R]) string {
mut request_fields := T{}
mut response_fields := R{}
// // Builds a GraphQL query or mutation string from provided arguments.
// fn build_query[T, R](args BuildQueryArgs[T, R]) string {
// mut request_fields := T{}
// mut response_fields := R{}
if args.request_model {
request_fields = get_request_fields(args.request_model)
}
// if args.request_model {
// request_fields = get_request_fields(args.request_model)
// }
if args.response_model {
response_fields = get_response_fields(args.response_model)
}
// if args.response_model {
// response_fields = get_response_fields(args.response_model)
// }
mut query := ''
// mut query := ''
if args.request_model && args.response_model{
query := '${args.query_type.to_string()} { ${args.method_name}(input: ${request_fields}) ${response_fields} }'
}
// if args.request_model && args.response_model{
// query := '${args.query_type.to_string()} { ${args.method_name}(input: ${request_fields}) ${response_fields} }'
// }
if args.response_model && !args.request_model{
query := '${args.query_type.to_string()} { ${response_fields} }'
}
// if args.response_model && !args.request_model{
// query := '${args.query_type.to_string()} { ${response_fields} }'
// }
// Wrap in the final structure
gql := GqlQuery{
query: query
}
// // Wrap in the final structure
// gql := GqlQuery{
// query: query
// }
// Return the final GraphQL query as a JSON string
return json.encode(gql)
}
// // Return the final GraphQL query as a JSON string
// return json.encode(gql)
// }
// Converts the `QueryType` enum to its string representation.
fn (q QueryType) to_string() string {
return match q {
fn (op OperationType) to_string() string {
return match op {
.query {
'query'
}