diff --git a/torch/distributed/elastic/timer/api.py b/torch/distributed/elastic/timer/api.py index 7c856f078d89a..313f6fa14940c 100644 --- a/torch/distributed/elastic/timer/api.py +++ b/torch/distributed/elastic/timer/api.py @@ -152,8 +152,8 @@ def clear_timers(self, worker_ids: set[Any]) -> None: def get_expired_timers(self, deadline: float) -> dict[str, list[TimerRequest]]: """ Returns all expired timers for each worker_id. An expired timer - is a timer for which the expiration_time is less than or equal to - the provided deadline. + is a timer for which the expiration_time is less than the + provided deadline. """ @abc.abstractmethod @@ -187,10 +187,11 @@ def _watchdog_loop(self): logger.exception("Error running watchdog") def _run_watchdog(self): + from datetime import datetime batch_size = max(1, self._request_queue.size()) timer_requests = self._request_queue.get(batch_size, self._max_interval) self.register_timers(timer_requests) - now = time.time() + now = datetime.now().timestamp() reaped_worker_ids = set() for worker_id, expired_timers in self.get_expired_timers(now).items(): logger.info( @@ -228,7 +229,9 @@ def stop(self) -> None: self._stop_signaled = True if self._watchdog_thread: logger.info("Stopping watchdog thread...") - self._watchdog_thread.join(self._max_interval) + self._watchdog_thread.join(self._max_interval * 2) + if self._watchdog_thread.is_alive(): + logger.warning("Watchdog thread did not stop in time") self._watchdog_thread = None else: logger.info("No watchdog thread running, doing nothing") @@ -267,6 +270,7 @@ def expires( with expires(after=10): torch.distributed.all_reduce(...) """ + from datetime import datetime if client is None: if _timer_client is None: raise RuntimeError("Configure timer client before using countdown timers.") @@ -275,7 +279,7 @@ def expires( # grab the caller file + lineno caller = getframeinfo(stack()[1][0]) scope = f"{caller.filename}#{caller.lineno}" - expiration = time.time() + after + expiration = datetime.now().timestamp() + after client.acquire(scope, expiration) try: yield