update

Lightning-AI · May 13, 2024 · 264f7e1 · 264f7e1
1 parent 1b3f4e7
commit 264f7e1
Showing 1 changed file with 12 additions and 0 deletions.
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