Add dist env detection via env vars (#95)

* Add dist env detection via env vars * minor typo fix * Use os.environ.get * fix mypy issues * update --------- Co-authored-by: thomas chaton <[email protected]>
Lightning-AI · May 13, 2024 · 58b158a · 58b158a
1 parent dc29634
commit 58b158a
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/src/litdata/utilities/env.py b/src/litdata/utilities/env.py
@@ -60,6 +60,10 @@ def detect(cls) -> "_DistributedEnv":
         if world_size is None or world_size == -1:
             world_size = 1
 
+        world_size = int(os.environ.get("WORLD_SIZE", world_size))
+        global_rank = int(os.environ.get("GLOBAL_RANK", global_rank))
+        num_nodes = int(os.environ.get("NNODES", num_nodes))
+
         return cls(world_size=world_size, global_rank=global_rank, num_nodes=num_nodes)
 
     @classmethod

diff --git a/tests/utilities/test_env.py b/tests/utilities/test_env.py
@@ -0,0 +1,12 @@
+from litdata.utilities.env import _DistributedEnv
+
+
+def test_distributed_env_from_env(monkeypatch):
+    monkeypatch.setenv("WORLD_SIZE", 2)
+    monkeypatch.setenv("GLOBAL_RANK", 1)
+    monkeypatch.setenv("NNODES", 2)
+
+    dist_env = _DistributedEnv.detect()
+    assert dist_env.world_size == 2
+    assert dist_env.global_rank == 1
+    assert dist_env.num_nodes == 2