Addressed comments and made some cleanup in tests.

Signed-off-by: meetkuma <[email protected]>

Author: quic-meetkuma

QEfficient/finetune/experimental/core/logger.py

@@ -12,7 +12,21 @@
 from typing import Optional
 from transformers.utils.logging import get_logger as hf_get_logger
 
-from .utils.dist_utils import get_rank
+from .utils.dist_utils import get_local_rank
+
+
+# -----------------------------------------------------------------------------
+# Logger usage:
+# Initialize logger:
+#   logger = Logger("my_logger", log_file="logs/output.log", level=logging.DEBUG)
+# Log messages:
+#   logger.info("This is an info message")
+#   logger.error("This is an error message")
+#   logger.log_rank_zero("This message is logged only on rank 0")
+#   logger.log_exception("An error occurred", exception, raise_exception=False)
+# Attach file handler later if needed:
+#   logger.prepare_for_logs(output_dir="logs", log_level="DEBUG")
+# -----------------------------------------------------------------------------
 
 
 class Logger:
@@ -86,7 +100,7 @@ def log_rank_zero(self, message: str, level: int = logging.INFO) -> None:
             message: Message to log
             level: Logging level
         """
-        if get_rank() == 0:
+        if get_local_rank() == 0:
             self.logger.log(level, message)
 
     def log_exception(self, message: str, exception: Exception, raise_exception: bool = True) -> None:
@@ -104,13 +118,13 @@ def log_exception(self, message: str, exception: Exception, raise_exception: boo
         if raise_exception:
             raise exception
 
-    def prepare_for_logs(self, output_dir: str, save_metrics: bool = True, log_level: str = "INFO") -> None:
+    def prepare_for_logs(self, output_dir: Optional[str] = None, log_level: str = "INFO") -> None:
         """
-        Prepare logger for training logs.
+        Prepare existing logger to log to both console and file with specified
+        output directory and log level.
 
         Args:
             output_dir: Output directory for logs
-            save_metrics: Whether to save metrics to file
             log_level: Logging level as string
         """
         # Convert string log level to logging constant
@@ -122,7 +136,7 @@ def prepare_for_logs(self, output_dir: str, save_metrics: bool = True, log_level
             handler.setLevel(level)
 
         # Add file handler if saving metrics
-        if save_metrics:
+        if output_dir:
             log_file = Path(output_dir) / "training.log"
             log_file.parent.mkdir(parents=True, exist_ok=True)
 

QEfficient/finetune/experimental/core/utils/dist_utils.py

@@ -14,12 +14,19 @@ def is_dist_available_and_initialized() -> bool:
 
 
 def get_rank() -> int:
-    """Get the rank of the current process in distributed training."""
+    """Return the global rank of the current process, else 0."""
     if not is_dist_available_and_initialized():
         return 0
     return dist.get_rank()
 
 
+def get_local_rank() -> int:
+    """Return the local rank of the current process on its node, else 0."""
+    if not is_dist_available_and_initialized():
+        return 0
+    return dist.get_node_local_rank()
+
+
 def get_world_size() -> int:
     """Get the total number of processes in distributed training."""
     if not is_dist_available_and_initialized():

QEfficient/finetune/experimental/tests/test_logger.py

@@ -63,21 +63,21 @@ def test_log_levels(self, caplog):
             assert "Error message" in caplog.text
             assert "Critical message" in caplog.text
 
-    @patch("QEfficient.finetune.experimental.core.logger.get_rank")
-    def test_log_rank_zero(self, mock_get_rank, caplog):
+    @patch("QEfficient.finetune.experimental.core.logger.get_local_rank")
+    def test_log_rank_zero_positive_case(self, mock_get_local_rank, caplog):
         """Test rank zero logging functionality"""
-        mock_get_rank.return_value = 0
+        mock_get_local_rank.return_value = 0
         logger = Logger("rank_test_logger")
 
         with caplog.at_level(logging.INFO):
             logger.log_rank_zero("Rank zero message")
 
             assert "Rank zero message" in caplog.text
 
-    @patch("QEfficient.finetune.experimental.core.logger.get_rank")
-    def test_log_rank_zero_not_zero(self, mock_get_rank, caplog):
-        """Test that non-rank zero messages are not logged"""
-        mock_get_rank.return_value = 1
+    @patch("QEfficient.finetune.experimental.core.logger.get_local_rank")
+    def test_log_rank_zero_negative_case(self, mock_get_local_rank, caplog):
+        """Test to verify that only rank‑zero messages are logged"""
+        mock_get_local_rank.return_value = 1
         logger = Logger("rank_test_logger")
 
         with caplog.at_level(logging.INFO):
@@ -112,7 +112,7 @@ def test_prepare_for_logs(self, tmp_path):
         logger = Logger("prepare_test_logger")
 
         # Prepare for logs
-        logger.prepare_for_logs(str(output_dir), save_metrics=True, log_level="DEBUG")
+        logger.prepare_for_logs(str(output_dir), log_level="DEBUG")
 
         # Check file handler was added
         file_handlers = [h for h in logger.logger.handlers if isinstance(h, logging.FileHandler)]
@@ -125,13 +125,12 @@ def test_prepare_for_logs(self, tmp_path):
         # Check log level was updated
         assert logger.logger.level == logging.DEBUG
 
-    def test_prepare_for_logs_no_metrics(self, tmp_path):
-        """Test preparing logger without saving metrics"""
-        output_dir = tmp_path / "output"
+    def test_prepare_for_logs_no_file_handler(self):
+        """Test preparing logger without saving to file"""
         logger = Logger("prepare_test_logger")
 
         # Prepare for logs without saving metrics
-        logger.prepare_for_logs(str(output_dir), save_metrics=False, log_level="INFO")
+        logger.prepare_for_logs(log_level="INFO")
 
         # Check no file handler was added
         file_handlers = [h for h in logger.logger.handlers if isinstance(h, logging.FileHandler)]
@@ -149,7 +148,7 @@ def test_prepare_for_logs_already_has_file_handler(self, tmp_path):
         logger.logger.addHandler(file_handler)
 
         # Prepare for logs again
-        logger.prepare_for_logs(str(output_dir), save_metrics=True, log_level="INFO")
+        logger.prepare_for_logs(str(output_dir), log_level="INFO")
 
         # Should still have only one file handler
         file_handlers = [h for h in logger.logger.handlers if isinstance(h, logging.FileHandler)]
@@ -204,7 +203,7 @@ def test_complete_workflow(self, tmp_path, caplog):
             logger.log_exception("Caught exception", e, raise_exception=False)
 
         # Test rank zero logging
-        with patch("QEfficient.finetune.experimental.core.logger.get_rank") as mock_rank:
+        with patch("QEfficient.finetune.experimental.core.logger.get_local_rank") as mock_rank:
             mock_rank.return_value = 0
             logger.log_rank_zero("Rank zero test")