[BE][1/N]Add sharding spec logger for ShardedTensor (#99748)

wz337 · pytorchmergebot · commit ca8625f45662 · 2023-04-22T04:05:21.000Z
Set up a nullHandler() on the OSS side. Next step is to set up the counterpart in internal. This is part of the effort for ShardedTensor deprecation. We want to log internal use cases for different sharding spec. Pull Request resolved: #99748 Approved by: https://github.com/H-Huang, https://github.com/fegin
diff --git a/test/distributed/_shard/sharded_tensor/test_logger.py b/test/distributed/_shard/sharded_tensor/test_logger.py
@@ -0,0 +1,21 @@
+# Owner(s): ["oncall: distributed"]
+
+import logging
+
+from torch.distributed._shard.sharded_tensor.logger import _get_or_create_logger
+from torch.testing._internal.common_utils import (
+    TestCase,
+    run_tests,
+)
+
+
+class ShardingSpecLoggerTest(TestCase):
+    def test_get_or_create_logger(self):
+        logger = _get_or_create_logger()
+        self.assertIsNotNone(logger)
+        self.assertEqual(1, len(logger.handlers))
+        self.assertIsInstance(logger.handlers[0], logging.NullHandler)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/torch/distributed/_shard/sharded_tensor/logger.py b/torch/distributed/_shard/sharded_tensor/logger.py
@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from typing import List, Tuple
+
+from torch.distributed._shard.sharded_tensor.logging_handlers import (
+    _log_handlers,
+)
+
+__all__: List[str] = []
+
+
+def _get_or_create_logger() -> logging.Logger:
+    logging_handler, log_handler_name = _get_logging_handler()
+    logger = logging.getLogger(f"sharding-spec-{log_handler_name}")
+    logger.setLevel(logging.DEBUG)
+    formatter = logging.Formatter(
+        "%(asctime)s %(filename)s:%(lineno)s %(levelname)s p:%(processName)s t:%(threadName)s: %(message)s"
+    )
+    logging_handler.setFormatter(formatter)
+    logger.propagate = False
+    logger.addHandler(logging_handler)
+    return logger
+
+
+def _get_logging_handler(
+    destination: str = "default",
+) -> Tuple[logging.Handler, str]:
+    log_handler = _log_handlers[destination]
+    log_handler_name = type(log_handler).__name__
+    return (log_handler, log_handler_name)
diff --git a/torch/distributed/_shard/sharded_tensor/logging_handlers.py b/torch/distributed/_shard/sharded_tensor/logging_handlers.py
@@ -0,0 +1,16 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from typing import Dict, List
+
+__all__: List[str] = []
+
+_log_handlers: Dict[str, logging.Handler] = {
+    "default": logging.NullHandler(),
+}
