分布式训练training-operator和pytorch-distributed RANK变量不统一解决

#!/usr/bin/env bash 
CONFIG=$1 
GPUS=$2 
NNODES=${NNODES:-1} 
NODE_RANK=${NODE_RANK:-0} # 如果NODE_RANK没有被设置为系统变量，则使用默认值0 
PORT=${PORT:-29500} 
MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} 
PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 
python -m torch.distributed.launch \ 
    --nnodes=$NNODES \ 
    --node_rank=$NODE_RANK \ # 作为torch.distributed.launch参数的一部分 
    --master_addr=$MASTER_ADDR \ 
    --nproc_per_node=$GPUS \ 
    --master_port=$PORT \ 
    $(dirname "$0")/train.py \ 
    $CONFIG \ 
    --seed 0 \ 
    --launcher pytorch ${@:3}

func setPodEnv(obj interface{}, podTemplateSpec *corev1.PodTemplateSpec, rtype, index string) error {
   pytorchjob, ok := obj.(*kubeflowv1.PyTorchJob)
   if !ok {
      return fmt.Errorf("%+v is not a type of PyTorchJob", obj)
   }
   for i := range podTemplateSpec.Spec.Containers {
      // Initialize the environment variables.
      if len(podTemplateSpec.Spec.Containers[i].Env) == 0 {
         podTemplateSpec.Spec.Containers[i].Env = make([]corev1.EnvVar, 0)
      }
      // Set PYTHONUNBUFFERED to true, to disable output buffering.
      // Ref https://stackoverflow.com/questions/59812009/what-is-the-use-of-pythonunbuffered-in-docker-file.
      podTemplateSpec.Spec.Containers[i].Env = append(
         podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
            Name:  "PYTHONUNBUFFERED",
            Value: "0",
         })
      // If the master is not null, then we need to set the MASTER_ADDR and RANK.
      if pytorchjob.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeMaster] != nil {
         envVars, err := GetMasterEnvVarGenerator().Generate(pytorchjob)
         if err != nil {
            return err
         }
         // Set master related environment variables.
         podTemplateSpec.Spec.Containers[i].Env = append(
            podTemplateSpec.Spec.Containers[i].Env, envVars...)
         // Set world size and rank.
         rank, err := strconv.Atoi(index)
         if err != nil {
            return err
         }
         if rtype == strings.ToLower(string(kubeflowv1.PyTorchJobReplicaTypeWorker)) {
            rank = rank + 1
         }
         totalReplicas := getTotalReplicas(pytorchjob)
         podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
            Name:  "WORLD_SIZE",
            Value: strconv.Itoa(int(totalReplicas)),
         })
         podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
            Name:  "RANK",
            Value: strconv.Itoa(rank),
         })
         // 新增一个名为NODE_RANK的环境变量
         podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
            Name:  "NODE_RANK",
            Value: strconv.Itoa(rank),
         })
      }
      // Set the elastic environment variables if the elasticPolicy is not null.
      if pytorchjob.Spec.ElasticPolicy != nil {
         envVars, err := GetElasticEnvVarGenerator().Generate(pytorchjob)
         if err != nil {
            return err
         }
         // Set elastic related environment variables.
         podTemplateSpec.Spec.Containers[i].Env = append(
            podTemplateSpec.Spec.Containers[i].Env, envVars...)
      }
   }
   return nil
}

 # Build manager binary.
go build -o bin/manager cmd/training-operator.v1/main.go
 # Build docker image with the manager.
docker build -t ${IMG} -f build/images/training-operator/Dockerfile .
 # Push docker image with the manager.
docker push ${IMG}

kubectl apply -k .