fix : 修复存储层 rglob 死锁、DDP LOCAL_RANK 绑定

astrai/parallel/executor.py

@@ -2,6 +2,7 @@
 
 import contextlib
 import logging
+import os
 from contextlib import contextmanager
 from typing import Optional, Tuple
 
@@ -181,7 +182,7 @@ class DDPExecutor(BaseExecutor):
         if not self.use_distributed:
             logger.warning("DDP backend selected but world_size=1, model not wrapped")
             return model
-        local_rank = get_rank()
+        local_rank = int(os.environ.get("LOCAL_RANK", get_rank()))
         model = DDP(
             model,
             device_ids=[local_rank],

astrai/parallel/setup.py

@@ -44,11 +44,12 @@ def setup_parallel(
         yield None
         return
 
-    device_id = torch.device(device_type, rank)
+    local_rank = int(os.environ["LOCAL_RANK"]) if "LOCAL_RANK" in os.environ else rank
+    device_id = torch.device(device_type, local_rank)
 
     os.environ["MASTER_ADDR"] = master_addr
     os.environ["MASTER_PORT"] = master_port
-    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(local_rank)
     os.environ["WORLD_SIZE"] = str(world_size)
     os.environ["LOCAL_DEVICE"] = str(device_id)
 
@@ -126,7 +127,23 @@ def spawn_parallel_fn(
     start_method: str = "spawn",
     **kwargs,
 ):
-    # clear environment variables
+    # Multi-node support: if RANK env var is set, init process group
+    # and run function directly (no local spawn).
+    if "RANK" in os.environ:
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ["WORLD_SIZE"])
+        with setup_parallel(
+            rank=rank,
+            world_size=world_size,
+            backend=backend,
+            master_addr=os.environ.get("MASTER_ADDR", master_addr),
+            master_port=os.environ.get("MASTER_PORT", master_port),
+            device_type=device_type,
+        ):
+            func(**kwargs)
+        return
+
+    # clear environment variables (single-node path)
     for key in [
         "MASTER_ADDR",
         "MASTER_PORT",