Create the function decorator to enable logging in torchrec (pytorch#3145)

Summary:
Pull Request resolved: pytorch#3145

This diff creates the logging decorator that will use the torchrec logger to record the function's input, output/error and other job identifying parameters. This is also where we will perform the JK check to see if torchrec logging is enabled.

Reviewed By: saumishr, kausv

Differential Revision: D76294270

fbshipit-source-id: 58c174787b76ea57324fa469cf2fa465d88c9673

Author: nipung90

torchrec/distributed/logger.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and 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.
+
+# mypy: allow-untyped-defs
+import functools
+import inspect
+from typing import Any, Callable, TypeVar
+
+import torchrec.distributed.torchrec_logger as torchrec_logger
+from torchrec.distributed.torchrec_logging_handlers import TORCHREC_LOGGER_NAME
+from typing_extensions import ParamSpec
+
+
+__all__: list[str] = []
+
+global _torchrec_logger
+_torchrec_logger = torchrec_logger._get_or_create_logger(TORCHREC_LOGGER_NAME)
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+
+def _torchrec_method_logger(
+    **wrapper_kwargs: Any,
+) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:  # pyre-ignore
+    """This method decorator logs the input, output, and exception of wrapped events."""
+
+    def decorator(func: Callable[_P, _T]):  # pyre-ignore
+        @functools.wraps(func)
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
+            msg_dict = torchrec_logger._get_msg_dict(func.__name__, **kwargs)
+            try:
+                ## Add function input to log message
+                msg_dict["input"] = _get_input_from_func(func, *args, **kwargs)
+                # exceptions
+                result = func(*args, **kwargs)
+            except BaseException as error:
+                msg_dict["error"] = f"{error}"
+                _torchrec_logger.error(msg_dict)
+                raise
+            msg_dict["output"] = str(result)
+            _torchrec_logger.debug(msg_dict)
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def _get_input_from_func(
+    func: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
+) -> str:
+    signature = inspect.signature(func)
+    bound_args = signature.bind_partial(*args, **kwargs)
+    bound_args.apply_defaults()
+    input_vars = {param.name: param.default for param in signature.parameters.values()}
+    for key, value in bound_args.arguments.items():
+        if isinstance(value, (int, float)):
+            input_vars[key] = value
+        else:
+            input_vars[key] = str(value)
+    return str(input_vars)

torchrec/distributed/tests/test_logger.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and 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.
+
+# pyre-strict
+
+import unittest
+from typing import Any
+from unittest import mock
+
+from torchrec.distributed.logger import _get_input_from_func, _torchrec_method_logger
+
+
+class TestLogger(unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+
+        # Mock torchrec_logger._get_msg_dict
+        self.get_msg_dict_patcher = mock.patch(
+            "torchrec.distributed.torchrec_logger._get_msg_dict"
+        )
+        self.mock_get_msg_dict = self.get_msg_dict_patcher.start()
+        self.mock_get_msg_dict.return_value = {}
+
+        # Mock _torchrec_logger
+        self.logger_patcher = mock.patch("torchrec.distributed.logger._torchrec_logger")
+        self.mock_logger = self.logger_patcher.start()
+
+    def tearDown(self) -> None:
+        self.get_msg_dict_patcher.stop()
+        self.logger_patcher.stop()
+        super().tearDown()
+
+    def test_get_input_from_func_no_args(self) -> None:
+        """Test _get_input_from_func with a function that has no arguments."""
+
+        def test_func() -> None:
+            pass
+
+        result = _get_input_from_func(test_func)
+        self.assertEqual(result, "{}")
+
+    def test_get_input_from_func_with_args(self) -> None:
+        """Test _get_input_from_func with a function that has positional arguments."""
+
+        def test_func(_a: int, _b: str) -> None:
+            pass
+
+        result = _get_input_from_func(test_func, 42, "hello")
+        self.assertEqual(result, "{'_a': 42, '_b': 'hello'}")
+
+    def test_get_input_from_func_with_kwargs(self) -> None:
+        """Test _get_input_from_func with a function that has keyword arguments."""
+
+        def test_func(_a: int = 0, _b: str = "default") -> None:
+            pass
+
+        result = _get_input_from_func(test_func, _b="world")
+        self.assertEqual(result, "{'_a': 0, '_b': 'world'}")
+
+    def test_get_input_from_func_with_args_and_kwargs(self) -> None:
+        """Test _get_input_from_func with a function that has both positional and keyword arguments."""
+
+        def test_func(
+            _a: int, _b: str = "default", *_args: Any, **_kwargs: Any
+        ) -> None:
+            pass
+
+        result = _get_input_from_func(test_func, 42, "hello", "extra", key="value")
+        self.assertEqual(
+            result,
+            "{'_a': 42, '_b': 'hello', '_args': \"('extra',)\", '_kwargs': \"{'key': 'value'}\"}",
+        )
+
+    def test_torchrec_method_logger_success(self) -> None:
+        """Test _torchrec_method_logger with a successful function execution when logging is enabled."""
+        # Create a mock function that returns a value
+        mock_func = mock.MagicMock(return_value="result")
+        mock_func.__name__ = "mock_func"
+
+        # Apply the decorator
+        decorated_func = _torchrec_method_logger()(mock_func)
+
+        # Call the decorated function
+        result = decorated_func(42, key="value")
+
+        # Verify the result
+        self.assertEqual(result, "result")
+
+        # Verify that _get_msg_dict was called with the correct arguments
+        self.mock_get_msg_dict.assert_called_once_with("mock_func", key="value")
+
+        # Verify that the logger was called with the correct message
+        self.mock_logger.debug.assert_called_once()
+        msg_dict = self.mock_logger.debug.call_args[0][0]
+        self.assertEqual(msg_dict["output"], "result")
+
+    def test_torchrec_method_logger_exception(self) -> None:
+        """Test _torchrec_method_logger with a function that raises an exception when logging is enabled."""
+        # Create a mock function that raises an exception
+        mock_func = mock.MagicMock(side_effect=ValueError("test error"))
+        mock_func.__name__ = "mock_func"
+
+        # Apply the decorator
+        decorated_func = _torchrec_method_logger()(mock_func)
+
+        # Call the decorated function and expect an exception
+        with self.assertRaises(ValueError):
+            decorated_func(42, key="value")
+
+        # Verify that _get_msg_dict was called with the correct arguments
+        self.mock_get_msg_dict.assert_called_once_with("mock_func", key="value")
+
+        # Verify that the logger was called with the correct message
+        self.mock_logger.error.assert_called_once()
+        msg_dict = self.mock_logger.error.call_args[0][0]
+        self.assertEqual(msg_dict["error"], "test error")
+
+    def test_torchrec_method_logger_with_wrapper_kwargs(self) -> None:
+        """Test _torchrec_method_logger with wrapper kwargs."""
+        # Create a mock function that returns a value
+        mock_func = mock.MagicMock(return_value="result")
+        mock_func.__name__ = "mock_func"
+
+        # Apply the decorator with wrapper kwargs
+        decorated_func = _torchrec_method_logger(custom_kwarg="value")(mock_func)
+
+        # Call the decorated function
+        result = decorated_func(42, key="value")
+
+        # Verify the result
+        self.assertEqual(result, "result")
+
+        # Verify that _get_msg_dict was called with the correct arguments
+        self.mock_get_msg_dict.assert_called_once_with("mock_func", key="value")
+
+        # Verify that the logger was called with the correct message
+        self.mock_logger.debug.assert_called_once()
+        msg_dict = self.mock_logger.debug.call_args[0][0]
+        self.assertEqual(msg_dict["output"], "result")
+
+
+if __name__ == "__main__":
+    unittest.main()

torchrec/distributed/tests/test_torchrec_logger.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and 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.
+
+# pyre-strict
+
+import logging
+import unittest
+from unittest import mock
+
+import torch.distributed as dist
+
+from torchrec.distributed.logging_handlers import _log_handlers
+from torchrec.distributed.torchrec_logger import (
+    _DEFAULT_DESTINATION,
+    _get_logging_handler,
+    _get_msg_dict,
+    _get_or_create_logger,
+)
+
+
+class TestTorchrecLogger(unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        # Save the original _log_handlers to restore it after tests
+        self.original_log_handlers = _log_handlers.copy()
+
+        # Create a mock logging handler
+        self.mock_handler = mock.MagicMock(spec=logging.Handler)
+        _log_handlers[_DEFAULT_DESTINATION] = self.mock_handler
+
+        # Mock print function
+        self.print_patcher = mock.patch("builtins.print")
+        self.mock_print = self.print_patcher.start()
+
+    def tearDown(self) -> None:
+        # Restore the original _log_handlers
+        _log_handlers.clear()
+        _log_handlers.update(self.original_log_handlers)
+
+        # Stop print patcher
+        self.print_patcher.stop()
+
+        super().tearDown()
+
+    def test_get_logging_handler(self) -> None:
+        """Test _get_logging_handler function."""
+        # Test with default destination
+        handler, name = _get_logging_handler()
+
+        self.assertEqual(handler, self.mock_handler)
+        self.assertEqual(
+            name, f"{type(self.mock_handler).__name__}-{_DEFAULT_DESTINATION}"
+        )
+
+        # Test with custom destination
+        custom_dest = "custom_dest"
+        custom_handler = mock.MagicMock(spec=logging.Handler)
+        _log_handlers[custom_dest] = custom_handler
+
+        handler, name = _get_logging_handler(custom_dest)
+
+        self.assertEqual(handler, custom_handler)
+        self.assertEqual(name, f"{type(custom_handler).__name__}-{custom_dest}")
+
+    @mock.patch("logging.getLogger")
+    def test_get_or_create_logger(self, mock_get_logger: mock.MagicMock) -> None:
+        """Test _get_or_create_logger function."""
+        mock_logger = mock.MagicMock(spec=logging.Logger)
+        mock_get_logger.return_value = mock_logger
+
+        # Test with default destination
+        _get_or_create_logger()
+
+        # Verify logger was created with the correct name
+        handler_name = f"{type(self.mock_handler).__name__}-{_DEFAULT_DESTINATION}"
+        mock_get_logger.assert_called_once_with(f"torchrec-{handler_name}")
+
+        # Verify logger was configured correctly
+        mock_logger.setLevel.assert_called_once_with(logging.DEBUG)
+        mock_logger.addHandler.assert_called_once_with(self.mock_handler)
+        self.assertFalse(mock_logger.propagate)
+
+        # Verify formatter was set on the handler
+        self.mock_handler.setFormatter.assert_called_once()
+        formatter = self.mock_handler.setFormatter.call_args[0][0]
+        self.assertIsInstance(formatter, logging.Formatter)
+
+        # Test with custom destination
+        mock_get_logger.reset_mock()
+        self.mock_handler.reset_mock()
+
+        custom_dest = "custom_dest"
+        custom_handler = mock.MagicMock(spec=logging.Handler)
+        _log_handlers[custom_dest] = custom_handler
+
+        _get_or_create_logger(custom_dest)
+
+        # Verify logger was created with the correct name
+        handler_name = f"{type(custom_handler).__name__}-{custom_dest}"
+        mock_get_logger.assert_called_once_with(f"torchrec-{handler_name}")
+
+        # Verify custom handler was used
+        mock_logger.addHandler.assert_called_once_with(custom_handler)
+
+    def test_get_msg_dict_without_dist(self) -> None:
+        """Test _get_msg_dict function without dist initialized."""
+        # Mock dist.is_initialized to return False
+        with mock.patch("torch.distributed.is_initialized", return_value=False):
+            msg_dict = _get_msg_dict("test_func", kwarg1="val1")
+
+            # Verify msg_dict contains only func_name
+            self.assertEqual(len(msg_dict), 1)
+            self.assertEqual(msg_dict["func_name"], "test_func")
+
+    def test_get_msg_dict_with_dist(self) -> None:
+        """Test _get_msg_dict function with dist initialized."""
+        # Mock dist functions
+        with mock.patch.multiple(
+            dist,
+            is_initialized=mock.MagicMock(return_value=True),
+            get_world_size=mock.MagicMock(return_value=4),
+            get_rank=mock.MagicMock(return_value=2),
+        ):
+            # Test with group in kwargs
+            mock_group = mock.MagicMock()
+            msg_dict = _get_msg_dict("test_func", group=mock_group)
+
+            # Verify msg_dict contains all expected keys
+            self.assertEqual(len(msg_dict), 4)
+            self.assertEqual(msg_dict["func_name"], "test_func")
+            self.assertEqual(msg_dict["group"], str(mock_group))
+            self.assertEqual(msg_dict["world_size"], "4")
+            self.assertEqual(msg_dict["rank"], "2")
+
+            # Verify get_world_size and get_rank were called with the group
+            dist.get_world_size.assert_called_once_with(mock_group)
+            dist.get_rank.assert_called_once_with(mock_group)
+
+            # Test with process_group in kwargs
+            dist.get_world_size.reset_mock()
+            dist.get_rank.reset_mock()
+
+            mock_process_group = mock.MagicMock()
+            msg_dict = _get_msg_dict("test_func", process_group=mock_process_group)
+
+            # Verify msg_dict contains all expected keys
+            self.assertEqual(msg_dict["group"], str(mock_process_group))
+
+            # Verify get_world_size and get_rank were called with the process_group
+            dist.get_world_size.assert_called_once_with(mock_process_group)
+            dist.get_rank.assert_called_once_with(mock_process_group)
+
+
+if __name__ == "__main__":
+    unittest.main()