Update env var RANK to NODE_RANK.

Author: qiuosier

ads/jobs/builders/runtimes/pytorch_runtime.py

@@ -205,7 +205,8 @@ def run(self, dsc_job, **kwargs):
  if not envs:
  envs = {}
  # Huggingface accelerate requires machine rank
- envs["RANK"] = str(i)
+ # Here we use NODE_RANK to store the machine rank
+ envs["NODE_RANK"] = str(i)
  envs["WORLD_SIZE"] = str(replicas)
  if main_run:
  envs["MAIN_JOB_RUN_OCID"] = main_run.id

ads/jobs/templates/driver_pytorch.py

@@ -694,7 +694,7 @@ def __init__(self, code_dir: str = driver_utils.DEFAULT_CODE_DIR) -> None:
  # --multi_gpu will be set automatically if there is more than 1 GPU
  # self.multi_gpu = bool(self.node_count > 1 or self.gpu_count > 1)
  self.num_machines = self.node_count
- self.machine_rank = os.environ["RANK"]
+ self.machine_rank = os.environ["NODE_RANK"]
  # Total number of processes across all nodes
  # Here we assume all nodes are having the same shape
  self.num_processes = (self.gpu_count if self.gpu_count else 1) * self.node_count