Introduce multiprocess logger (#337)

Author: muellerzr

examples/by_feature/tracking.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
-import logging
 import os
 
 import torch
@@ -30,9 +29,6 @@
 )
 
 
-logger = logging.getLogger(__name__)
-
-
 ########################################################################
 # This is a fully working simple example to use Accelerate,
 # specifically showcasing the experiment tracking capability,

src/accelerate/accelerator.py

@@ -26,6 +26,7 @@
 from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state
 from .data_loader import prepare_data_loader
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
+from .logging import get_logger
 from .optimizer import AcceleratedOptimizer
 from .scheduler import AcceleratedScheduler
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
@@ -52,10 +53,7 @@
 
     from .deepspeed_utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
 
-import logging
-
-
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class Accelerator:

src/accelerate/checkpointing.py

@@ -28,10 +28,10 @@
 if is_tpu_available():
     import torch_xla.core.xla_model as xm
 
-import logging
+from .logging import get_logger
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 def save_accelerator_state(

src/accelerate/logging.py

@@ -0,0 +1,63 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from .state import AcceleratorState
+
+
+class MultiProcessAdapter(logging.LoggerAdapter):
+    """
+    An adapter to assist with logging in multiprocess.
+
+    `log` takes in an additional `main_process_only` kwarg, which dictates whether it should be called on all processes
+    or only the main executed one. Default is `main_process_only=True`.
+    """
+
+    @staticmethod
+    def _should_log(main_process_only):
+        "Check if log should be performed"
+        return not main_process_only or (main_process_only and AcceleratorState().local_process_index == 0)
+
+    def log(self, level, msg, *args, **kwargs):
+        """
+        Delegates logger call after checking if we should log.
+
+        Accepts a new kwarg of `main_process_only`, which will dictate whether it will be logged across all processes
+        or only the main executed one. Default is `True` if not passed
+        """
+        main_process_only = kwargs.pop("main_process_only", True)
+        if self.isEnabledFor(level) and self._should_log(main_process_only):
+            msg, kwargs = self.process(msg, kwargs)
+            self.logger.log(level, msg, *args, **kwargs)
+
+
+def get_logger(name: str):
+    """
+    Returns a `logging.Logger` for `name` that can handle multiprocessing.
+
+    If a log should be called on all processes, pass `main_process_only=False`
+
+    E.g.
+    ```python
+    logger.info("My log", main_process_only=False)
+    logger.debug("My log", main_process_only=False)
+    ```
+
+    Args:
+        name (`str`):
+            The name for the logger, such as `__file__`
+    """
+    logger = logging.getLogger(name)
+    return MultiProcessAdapter(logger, {})

src/accelerate/tracking.py

@@ -15,11 +15,11 @@
 # Expectation:
 # Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
 
-import logging
 import os
 from abc import ABCMeta, abstractmethod, abstractproperty
 from typing import List, Optional, Union
 
+from .logging import get_logger
 from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
 
 
@@ -41,7 +41,7 @@
     _available_trackers.append(LoggerType.COMETML)
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 def get_available_trackers():