From 876a65d9e54cc0c105c10bca3fd38a5811b513da Mon Sep 17 00:00:00 2001 From: gkroiz Date: Mon, 8 Apr 2024 23:11:59 +0000 Subject: [PATCH] Add dist env detection via env vars --- src/litdata/utilities/env.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/litdata/utilities/env.py b/src/litdata/utilities/env.py index 7276a87e..874e597b 100644 --- a/src/litdata/utilities/env.py +++ b/src/litdata/utilities/env.py @@ -57,7 +57,16 @@ def detect(cls) -> "_DistributedEnv": global_rank = 0 num_nodes = 1 - if world_size is None or world_size == -1: + if os.environ["WORLD_SIZE"] is not None: + world_size = os.environ["WORLD_SIZE"] + + if os.environ["GLOBAL_RANK"] is not None: + global_rank = os.environ["GLOBAL_RANK"] + + if os.environ["NNODES"] is not None: + num_nodes = os.environ["NNODES"] + + if world_size is in [None, -1, 0]: world_size = 1 return cls(world_size=world_size, global_rank=global_rank, num_nodes=num_nodes)