Merge pull request #177 from NVIDIA/rankmonitorserver_inherit_env

Configure RankMonitorServer to inherit env vars from launcher

Author: namitdhameja

src/nvidia_resiliency_ext/fault_tolerance/launcher.py

@@ -478,6 +478,7 @@ def setup_rank_monitors(self, envs: Dict[int, Dict[str, str]]) -> None:
                     ipc_socket_path=rmon_ipc_socket,
                     is_restarter_logger=is_restarter_logger,
                     mp_ctx=fork_mp_ctx,
+                    env=worker_env,
                 )
 
     def shutdown_rank_monitors(self):

src/nvidia_resiliency_ext/fault_tolerance/rank_monitor_server.py

@@ -126,7 +126,8 @@ def __init__(
         cfg: FaultToleranceConfig,
         ipc_socket_path: str,
         rank_monitor_ready_event,
-        logger: RankMonitorLogger,
+        logger: logging.Logger,
+        is_restarter_logger: bool,
     ):
         """
         Initializes the RankMonitorServer object.
@@ -151,7 +152,10 @@ def __init__(
         self.connection_lock = asyncio.Lock()
         self.rank_monitor_ready_event = rank_monitor_ready_event
         self.logger = logger
-        self.state_machine = RankMonitorStateMachine(logger)
+        self.rmlogger = RankMonitorLogger(
+            level=logger.level, is_restarter_logger=is_restarter_logger
+        )
+        self.state_machine = RankMonitorStateMachine(self.rmlogger)
         self._periodic_restart_task = None
         self.health_checker = GPUHealthCheck(
             interval=self.cfg.node_health_check_interval, on_failure=self._handle_unhealthy_node
@@ -264,7 +268,7 @@ async def _handle_init_msg(self, msg, writer):
         # Update NIC health checker on the rank to monitor.
         if self.nic_health_checker is not None:
             self.nic_health_checker.set_nic_device(local_rank=self.rank_info.local_rank)
-        self.logger.set_connected_rank(msg.rank_info.global_rank)
+        self.rmlogger.set_connected_rank(msg.rank_info.global_rank)
         await write_obj_to_ipc_stream(OkMsg(cfg=self.cfg), writer)
 
     async def _handle_heartbeat_msg(self, msg, writer):
@@ -318,7 +322,7 @@ def _handle_ipc_connection_lost(self):
                 f"Section(s) {open_section_names} were still open. you can use`.end_all_sections` to avoid this warning"
             )
             self.open_sections.clear()
-        self.logger.set_connected_rank(None)
+        self.rmlogger.set_connected_rank(None)
         if self.connection_lock.locked():
             self.connection_lock.release()
 
@@ -546,7 +550,8 @@ def run(
                 cfg,
                 ipc_socket_path,
                 rank_monitor_ready_event,
-                rmlogger,
+                logger,
+                is_restarter_logger,
             )
             asyncio.run(inst._rank_monitor_loop())
             logger.debug("Leaving RankMonitorServer process")

src/nvidia_resiliency_ext/inprocess/wrap.py

@@ -219,6 +219,7 @@ def __init__(
         self.finalize = finalize
         self.health_check = health_check
 
+        setup_logger(node_local_tmp_prefix="wrapper")
         # Construct internal restart_health_check by chaining user's health_check with GPU and NVL checks
         self._construct_restart_health_check()
 

src/nvidia_resiliency_ext/shared_utils/health_check.py

@@ -215,6 +215,7 @@ def __init__(
             on_failure (Optional[Callable]): Callback function to handle health check failures.
         """
         super().__init__()
+        logger = logging.getLogger(LogConfig.name)
         self.device_index = device_index
         self.interval = interval
         self.on_failure = on_failure

src/nvidia_resiliency_ext/shared_utils/log_node_local_tmp.py

@@ -72,7 +72,7 @@ def _get_backup_files(self):
     def _log_file_namer(self):
         backup_files = self._get_backup_files()
         if self.fname is None and backup_files:
-            return backup_files[0]
+            return backup_files[-1]
         rank_str = str(self.rank_id) if self.rank_id is not None else "unknown"
         file_prefix = f"rank_{rank_str}_{self.proc_name}.msg."
         return f"{file_prefix}{int(time.time()*1000)}"