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,7 @@ 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=+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