add node_rank to torchrun cmd if rdzv_backend is 'static' (pytorch#761)

* add node_rank to torchrun cmd if rdzv_backend is 'static' * add unit test for static backend option
KPostOffice · Sep 7, 2023 · 9882f12 · 9882f12
1 parent b2e3075
commit 9882f12
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/torchx/components/dist.py b/torchx/components/dist.py
@@ -267,6 +267,10 @@ def ddp(
         "--role",
         "",
     ]
+    # TODO 'node_rank' is made optional as it currently does not work with the AWS Batch scheduler.
+    # node_rank is only used when rdzv_backend is 'static'
+    if rdzv_backend == "static":
+        cmd += ["--node_rank", f"{macros.replica_id}"]
     if script is not None:
         cmd += [script]
     elif m is not None:

diff --git a/torchx/components/test/dist_test.py b/torchx/components/test/dist_test.py
@@ -38,6 +38,12 @@ def test_ddp_debug(self) -> None:
         for k, v in _TORCH_DEBUG_FLAGS.items():
             self.assertEqual(env[k], v)
 
+    def test_ddp_rdzv_backend_static(self) -> None:
+        app = ddp(script="foo.py", rdzv_backend="static")
+        cmd = app.roles[0].args[1]
+        self.assertTrue("--rdzv_backend static" in cmd)
+        self.assertTrue("--node_rank" in cmd)
+
 
 class SpmdTest(ComponentTestCase):
     def test_validate_spmd(self) -> None: