add GROUP_RANK as node rank

Author: pmcurtin

src/torchrunx/agent.py

@@ -83,6 +83,7 @@ def main(launcher_agent_group: LauncherAgentGroup, logger_hostname: str, logger_
                     backend=launcher_payload.backend,
                     rank=worker_global_ranks[i],
                     local_rank=i,
+                    node_rank=agent_rank,
                     local_world_size=num_workers,
                     world_size=worker_world_size,
                     hostname=launcher_payload.hostnames[agent_rank],

src/torchrunx/worker.py

@@ -32,6 +32,7 @@ class WorkerArgs:
     backend: Literal["nccl", "gloo", "mpi", "ucc", "auto"] | None
     rank: int
     local_rank: int
+    node_rank: int
     local_world_size: int
     world_size: int
     hostname: str
@@ -79,6 +80,7 @@ def worker_entrypoint(serialized_worker_args: SerializedWorkerArgs) -> Any | Exc
 
     os.environ["RANK"] = str(worker_args.rank)
     os.environ["LOCAL_RANK"] = str(worker_args.local_rank)
+    os.environ["GROUP_RANK"] = str(worker_args.node_rank)
     os.environ["LOCAL_WORLD_SIZE"] = str(worker_args.local_world_size)
     os.environ["WORLD_SIZE"] = str(worker_args.world_size)
     os.environ["MASTER_ADDR"] = worker_args.main_agent_hostname