add env var port live reco

add env var port live reco

Author: samsja

README.md

@@ -103,3 +103,4 @@ HF_HUB_ETAG_TIMEOUT=500
 | `ZERO_BAND_GLOBAL_STORE_POLLING_INTERVAL_SECONDS` | Number of seconds between polls to the store when waiting for values | `0.1` |
 | `ZERO_BAND_EDM_HEARTBEAT_INTERVAL_SECONDS` | Interval in seconds between heartbeats | `2` |
 | `ZERO_BAND_EDM_HEARTBEAT_TIMEOUT_SECONDS` | Time in seconds after which a node is considered dead if no heartbeat is received | `10` |
+| `ZERO_BAND_LIVE_RECO_PORT` | Port number for the live recovery server | `8000` |  

src/zeroband/checkpoint.py

@@ -34,6 +34,8 @@
 
 SHM_PATH = "/dev/shm/zeroband"
 
+ZERO_BAND_LIVE_RECO_PORT = int(os.environ.get("ZERO_BAND_LIVE_RECO_PORT", "8000"))
+
 
 @dataclass
 class TrainingProgress(Stateful):
@@ -135,7 +137,9 @@ def __init__(
         self.async_save_process: list[multiprocessing.Process] = []
 
         if live_ckpt_server:
-            self.live_server = CkptLiveServer(port=8000 + self.world_info.global_rank, ckpt_path=SHM_PATH)
+            self.live_server = CkptLiveServer(
+                port=ZERO_BAND_LIVE_RECO_PORT + self.world_info.global_rank, ckpt_path=SHM_PATH
+            )
 
     def _init_state(self):
         # states can only be stateful object, hence we need to wrap Model and Optimizer
@@ -288,11 +292,11 @@ def download_and_load_ckpt_from_peers(self):
         if self.world_info.local_rank == 0:
             # only local rank download the ckpt
             wget(
-                source=f"http://localhost:{8000+dest_rank}/latest/diloco_{dest_rank}",
+                source=f"http://localhost:{ZERO_BAND_LIVE_RECO_PORT+dest_rank}/latest/diloco_{dest_rank}",
                 destination=path,
             )
             wget(
-                source=f"http://localhost:{8000+dest_rank}/latest/diloco_{dest_rank}/.metadata",
+                source=f"http://localhost:{ZERO_BAND_LIVE_RECO_PORT+dest_rank}/latest/diloco_{dest_rank}/.metadata",
                 destination=path,
             )
         dist.barrier()