Merge pull request #201 from hexinw-nvidia/rendezvous_v2

Barrier Based Rendezvous

Author: hexinw-nvidia

pyproject.toml

@@ -33,6 +33,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry.dependencies]
 nv-one-logger-core = ">=2.1.0"
+nv-one-logger-training-telemetry = ">=2.3.0"
 torch = ">=2.3.0"
 packaging = "*"
 python = ">=3.10"

src/nvidia_resiliency_ext/fault_tolerance/_ft_rendezvous.py

@@ -1210,6 +1210,7 @@ class FtRendezvousHandler(RendezvousHandler):
     _op_executor: _RendezvousOpExecutor
     _heartbeat_lock: threading.Lock
     _keep_alive_timer: Optional[_PeriodicTimer]
+    _worker_group: Optional[Any] = None  # Store reference to worker group
 
     @classmethod
     def from_backend(
@@ -1326,6 +1327,10 @@ def _record(
             rank=rank,
         )
 
+    def set_worker_group(self, worker_group: Any) -> None:
+        """Set the worker group reference for this handler."""
+        self._worker_group = worker_group
+
     @property
     def settings(self) -> RendezvousSettings:
         """Get the settings of the rendezvous."""