Setup root logger once & on import time

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/523

To avoid setting it up multiple times, add run_once() decorator.

Additionally make sure logging is configured for datalodaing workers, which have a different entry point, by moving setting up logging to the import time. Right now when a dataloader worker is created using spawn method from multiprocessing module, a new Python interpreter is created, with all the modules imported anew and with the entry point set to the method specified. This means that the entry point of the training framework is skipped, together with the logging setup.

With this change, the logging is configured on the import time, which means that when a dataloading process is created, even though the training main is not invoked, the logging is still configured because even though train_net is not invoked as an entry point, it's still imported in the child process.

Reviewed By: miqueljubert

Differential Revision: D44641142

fbshipit-source-id: 06ea85363d965b31d7f9ade3c2615ed9db67470b

d2go/setup.py

@@ -38,6 +38,7 @@ from mobile_cv.common.misc.py import FolderLock, MultiprocessingPdb, post_mortem
 logger = logging.getLogger(__name__)
 
 
+@run_once()
 def setup_root_logger(logging_level: int = logging.INFO) -> None:
     """
     Sets up the D2Go root logger. When a new logger is created, it lies in a tree.

tools/train_net.py

@@ -38,6 +38,9 @@ from torch.distributed.elastic.multiprocessing.errors import (
 )
 
 logger = logging.getLogger("d2go.tools.train_net")
+# Make sure logging is set up centrally even for e.g. dataloading workers which
+# have entry points outside of D2Go.
+setup_root_logger()
 
 
 def main(
@@ -47,7 +50,6 @@ def main(
     eval_only: bool = False,
     resume: bool = True,  # NOTE: always enable resume when running on cluster
 ) -> Union[TrainNetOutput, TestNetOutput]:
-
     logger.info("Starting main")
     error_handler = get_error_handler()
     logger.debug(f">>>>>>> Error handler is: {type(error_handler)=}, {error_handler=}")
@@ -190,7 +192,6 @@ def build_cli_args(
 
 
 if __name__ == "__main__":
-    setup_root_logger()
     logger.info("Starting CLI application")
     try:
         cli()