add logging (#265)

Author: tushar00jain

torchft/futures.py

@@ -1,4 +1,5 @@
 import asyncio
+import os
 import queue
 import sys
 import threading
@@ -15,6 +16,8 @@
 
 T = TypeVar("T")
 
+WATCHDOG_TIMEOUT_SEC = "TORCHFT_WATCHDOG_TIMEOUT_SEC"
+
 
 class _TimerHandle:
     def __init__(self) -> None:
@@ -61,7 +64,9 @@ def __init__(self) -> None:
 
         # Give this much time the the `_event_loop_thread` to confirm that
         # it is not stuck
-        self._watchdog_interval = timedelta(seconds=30)
+        self._watchdog_interval = timedelta(
+            seconds=int(os.environ.get(WATCHDOG_TIMEOUT_SEC, "30"))
+        )
 
         # This queue is used to delete events on the main thread as cudaEventDestroy
         # can block if the CUDA queue is full.

torchft/local_sgd.py

@@ -762,6 +762,7 @@ def _step_post_hook(
             # can be overrepresented.
             self._manager.start_quorum()
             fragment = self._current_fragment()
+            logger.info(f"Preparing fragment={fragment} step={self._local_step}")
             self._fragments[fragment].prepare_sync()
 
         if self._local_step < self._sync_every:
@@ -770,6 +771,9 @@ def _step_post_hook(
         if self._local_step == self._sync_every:
             # Time to sync a fragment
             fragment = self._current_fragment()
+            logger.info(
+                f"Syncing fragment={fragment} step={self._local_step} manager_step={self._manager.current_step()}"
+            )
             self._fragments[fragment].perform_sync()
 
             # If the allreduce truly failed, we'll keep retrying this fragment.