Merge pull request #157 from youcefs21/main

feat and fix: add gpuCount argument to StartSpotPod
runpod · Oct 6, 2024 · 90665ad · 90665ad
2 parents ae0a532 + 7ecfae9
commit 90665ad
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/api/pod.go b/api/pod.go
@@ -355,19 +355,19 @@ func StartOnDemandPod(id string) (pod map[string]interface{}, err error) {
 	return
 }
 
-func StartSpotPod(id string, bidPerGpu float32) (podBidResume map[string]interface{}, err error) {
+func StartSpotPod(id string, bidPerGpu float32, gpuCount int) (podBidResume map[string]interface{}, err error) {
 	input := Input{
 		Query: `
 		mutation Mutation($podId: String!, $bidPerGpu: Float!) {
-			podBidResume(input: {podId: $podId, bidPerGpu: $bidPerGpu}) {
+			podBidResume(input: {podId: $podId, bidPerGpu: $bidPerGpu, gpuCount: $gpuCount}) {
 			  id
 			  costPerHr
 			  desiredStatus
 			  lastStatusChange
 			}
 		}
 		`,
-		Variables: map[string]interface{}{"podId": id, "bidPerGpu": bidPerGpu},
+		Variables: map[string]interface{}{"podId": id, "bidPerGpu": bidPerGpu, "gpuCount": gpuCount},
 	}
 	res, err := Query(input)
 	if err != nil {

diff --git a/cmd/pod/startPod.go b/cmd/pod/startPod.go
@@ -9,6 +9,7 @@ import (
 )
 
 var bidPerGpu float32
+var gpuCount int
 
 var StartPodCmd = &cobra.Command{
 	Use:   "pod [podId]",
@@ -19,7 +20,7 @@ var StartPodCmd = &cobra.Command{
 		var err error
 		var pod map[string]interface{}
 		if bidPerGpu > 0 {
-			pod, err = api.StartSpotPod(args[0], bidPerGpu)
+			pod, err = api.StartSpotPod(args[0], bidPerGpu, gpuCount)
 		} else {
 			pod, err = api.StartOnDemandPod(args[0])
 		}
@@ -36,4 +37,5 @@ var StartPodCmd = &cobra.Command{
 
 func init() {
 	StartPodCmd.Flags().Float32Var(&bidPerGpu, "bid", 0, "bid per gpu for spot price")
+	StartPodCmd.Flags().IntVar(&gpuCount, "gpuCount", 1, "number of GPUs to request")
 }