fix bug: allocated gpus is 0 when using arena top node (#628)

pkg/apis/utils/pods.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"sort"
-	"strconv"
 	"strings"
 	"time"
 
@@ -28,21 +27,21 @@ func CPUCountInPod(pod *v1.Pod) float64 {
 }
 
 func GPUCountInPod(pod *v1.Pod) int {
-	total := 0
+	total := int64(0)
 	for _, count := range ResourceInContainers(pod, types.NvidiaGPUResourceName) {
-		c, _ := parseInt(count)
+		c := count.(int64)
 		total += c
 	}
-	return total
+	return int(total)
 }
 
 func AliyunGPUCountInPod(pod *v1.Pod) int {
-	total := 0
+	total := int64(0)
 	for _, count := range ResourceInContainers(pod, types.AliyunGPUResourceName) {
-		c, _ := parseInt(count)
+		c := count.(int64)
 		total += c
 	}
-	return total
+	return int(total)
 }
 
 func ResourceInContainers(pod *v1.Pod, resourceName string) map[int]interface{} {
@@ -143,12 +142,12 @@ func DefinePodPhaseStatus(pod v1.Pod) (string, int, int, int) {
 }
 
 func GPUMemoryCountInPod(pod *v1.Pod) int {
-	total := 0
+	total := int64(0)
 	for _, count := range ResourceInContainers(pod, types.GPUShareResourceName) {
-		c, _ := parseInt(count)
+		c := count.(int64)
 		total += c
 	}
-	return total
+	return int(total)
 }
 
 func GetContainerAllocation(pod *v1.Pod) map[int]map[string]int {
@@ -316,9 +315,9 @@ func GetDurationOfPod(pod *v1.Pod) time.Duration {
 }
 
 func parseInt(i interface{}) (int, error) {
-	s, ok := i.(string)
+	s, ok := i.(int64)
 	if !ok {
 		return 0, errors.New("invalid value")
 	}
-	return strconv.Atoi(s)
+	return int(s), nil
 }

pkg/podlogs/logger.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"strings"
 
 	"github.com/kubeflow/arena/pkg/apis/config"
 	"github.com/kubeflow/arena/pkg/apis/types"
@@ -96,6 +97,10 @@ func (p *PodLogger) ensureContainerStarted() error {
 		log.Debugf("pod:%s,pod phase: %v\n", p.InstanceName, pod.Status.Phase)
 		log.Debugf("pod print status: %s\n", status)
 		switch podPhase := pod.Status.Phase; {
+		case podPhase == v1.PodPending && strings.Index(status, "Init:") == 0:
+			return nil
+		case podPhase == v1.PodPending && strings.Index(status, "PodInitializing") == 0:
+			return nil
 		case podPhase == v1.PodRunning && status != "Unknown":
 			return nil
 		case podPhase == v1.PodFailed || podPhase == v1.PodSucceeded:

pkg/training/logs.go

@@ -50,21 +50,19 @@ func AcceptJobLog(jobName string, trainingType types.TrainingJobType, args *type
 		podStatuses[pod.Name] = status
 	}
 	// 4.if the instance name is invalid,return error
-	_, ok := podStatuses[args.InstanceName]
+	status, ok := podStatuses[args.InstanceName]
 	if !ok {
 		return fmt.Errorf("invalid instance name %v in job %v,please use 'arena get %v' to make sure instance name.",
 			args.InstanceName,
 			jobName,
 			jobName,
 		)
 	}
-	// 5.if the instance status is not running,return error
-	//if status != "Running" {
-	//	return fmt.Errorf("failed to get logs of instance %v,because it is not running,please use 'arena get %v' to make sure instance status",
-	//		args.InstanceName,
-	//		jobName,
-	//	)
-	//}
+	if strings.Index(status, "Init:") == 0 || strings.Index(status, "PodInitializing") == 0 {
+		if args.ContainerName == "" {
+			args.ContainerName = "init-code"
+		}
+	}
 	logger := podlogs.NewPodLogger(args)
 	_, err = logger.AcceptLogs()
 	return err