Add more exception handling and support for class constructors (meta-pytorch#3517)

Summary:
Pull Request resolved: meta-pytorch#3517

This diff adds a few features to the static logger

1) Better exception handling to make sure any failures are silent and don't affect the actual training job.
2) Support for handling class constructors. This includes changing the function name to include the name of the class for constructors and moving the collection of inputs to after the function call succeeds.
3) Moved the static logger enablement stricts to a configerator file instead of using justknobs metadata. The justknobs will only be used as a kill switch now

Reviewed By: kausv, saumishr

Differential Revision: D76294357

fbshipit-source-id: 69e5db646e8557d7864647901773f52bbffb0e12

Author: nipung90

torchrec/distributed/logger.py

@@ -8,7 +8,8 @@
 # mypy: allow-untyped-defs
 import functools
 import inspect
-from typing import Any, Callable, TypeVar
+import logging
+from typing import Any, Callable, Dict, TypeVar
 
 import torchrec.distributed.torchrec_logger as torchrec_logger
 from torchrec.distributed.torchrec_logging_handlers import TORCHREC_LOGGER_NAME
@@ -34,16 +35,25 @@ def decorator(func: Callable[_P, _T]):  # pyre-ignore
         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}"
+                ## Add function input to log message
+                msg_dict["input"] = _get_input_from_func(
+                    func, msg_dict, *args, **kwargs
+                )
                 _torchrec_logger.error(msg_dict)
                 raise
-            msg_dict["output"] = str(result)
-            _torchrec_logger.debug(msg_dict)
+            ## Add function input to log message
+            try:
+                msg_dict["input"] = _get_input_from_func(
+                    func, msg_dict, *args, **kwargs
+                )
+                msg_dict["output"] = str(result)
+                _torchrec_logger.debug(msg_dict)
+            except Exception as error:
+                logging.info(f"Torchrec logger: Failed in static logger: {error}")
             return result
 
         return wrapper
@@ -52,15 +62,29 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
 
 
 def _get_input_from_func(
-    func: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
+    func: Callable[_P, _T],
+    msg_dict: Dict[str, Any],
+    *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)
+    try:
+        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 key == "self" and func.__name__ == "__init__":
+                # Add class name to function name if the function is a constructor
+                msg_dict["func_name"] = (
+                    f"{value.__class__.__name__}.{msg_dict['func_name']}"
+                )
+            if isinstance(value, (int, float)):
+                input_vars[key] = value
+            else:
+                input_vars[key] = str(value)
+        return str(input_vars)
+    except Exception as error:
+        logging.error(f"Torchrec Logger: Error in _get_input_from_func: {error}")
+        return "Error in _get_input_from_func: " + str(error)

torchrec/distributed/tests/test_logger.py

@@ -23,7 +23,12 @@ def setUp(self) -> None:
             "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 = {}
+
+        # Return a dictionary with func_name that can be modified by _get_input_from_func
+        def mock_get_msg_dict_impl(func_name: str, **kwargs: Any) -> dict[str, Any]:
+            return {"func_name": func_name}
+
+        self.mock_get_msg_dict.side_effect = mock_get_msg_dict_impl
 
         # Mock _torchrec_logger
         self.logger_patcher = mock.patch("torchrec.distributed.logger._torchrec_logger")
@@ -40,7 +45,8 @@ def test_get_input_from_func_no_args(self) -> None:
         def test_func() -> None:
             pass
 
-        result = _get_input_from_func(test_func)
+        msg_dict = {"func_name": "test_func"}
+        result = _get_input_from_func(test_func, msg_dict)
         self.assertEqual(result, "{}")
 
     def test_get_input_from_func_with_args(self) -> None:
@@ -49,7 +55,8 @@ def test_get_input_from_func_with_args(self) -> None:
         def test_func(_a: int, _b: str) -> None:
             pass
 
-        result = _get_input_from_func(test_func, 42, "hello")
+        msg_dict = {"func_name": "test_func"}
+        result = _get_input_from_func(test_func, msg_dict, 42, "hello")
         self.assertEqual(result, "{'_a': 42, '_b': 'hello'}")
 
     def test_get_input_from_func_with_kwargs(self) -> None:
@@ -58,7 +65,8 @@ def test_get_input_from_func_with_kwargs(self) -> None:
         def test_func(_a: int = 0, _b: str = "default") -> None:
             pass
 
-        result = _get_input_from_func(test_func, _b="world")
+        msg_dict = {"func_name": "test_func"}
+        result = _get_input_from_func(test_func, msg_dict, _b="world")
         self.assertEqual(result, "{'_a': 0, '_b': 'world'}")
 
     def test_get_input_from_func_with_args_and_kwargs(self) -> None:
@@ -69,11 +77,19 @@ def test_func(
         ) -> 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'}\"}",
+        msg_dict = {"func_name": "test_func"}
+        result = _get_input_from_func(
+            test_func, msg_dict, 42, "hello", "extra", key="value"
         )
+        self.assertIn("_a", result)
+        self.assertIn("42", result)
+        self.assertIn("_b", result)
+        self.assertIn("hello", result)
+        self.assertIn("_args", result)
+        self.assertIn("extra", result)
+        self.assertIn("_kwargs", result)
+        self.assertIn("key", result)
+        self.assertIn("value", result)
 
     def test_torchrec_method_logger_success(self) -> None:
         """Test _torchrec_method_logger with a successful function execution when logging is enabled."""
@@ -142,6 +158,57 @@ def test_torchrec_method_logger_with_wrapper_kwargs(self) -> None:
         msg_dict = self.mock_logger.debug.call_args[0][0]
         self.assertEqual(msg_dict["output"], "result")
 
+    def test_torchrec_method_logger_constructor_with_args(self) -> None:
+        """Test _torchrec_method_logger with a class constructor that has positional arguments."""
+
+        class TestClass:
+            @_torchrec_method_logger()
+            def __init__(self, _a: int, _b: str) -> None:
+                pass
+
+        # Create an instance which will call __init__
+        _ = TestClass(42, "hello")
+
+        # Verify that the logger was called
+        self.mock_logger.debug.assert_called_once()
+        msg_dict = self.mock_logger.debug.call_args[0][0]
+        # Verify that class name was prepended to function name
+        self.assertEqual(msg_dict["func_name"], "TestClass.__init__")
+        # Verify the input contains the arguments
+        self.assertIn("_a", msg_dict["input"])
+        self.assertIn("42", msg_dict["input"])
+        self.assertIn("_b", msg_dict["input"])
+        self.assertIn("hello", msg_dict["input"])
+
+    def test_torchrec_method_logger_constructor_with_args_and_kwargs(self) -> None:
+        """Test _torchrec_method_logger with a class constructor that has both positional and keyword arguments."""
+
+        class TestClass:
+            @_torchrec_method_logger()
+            def __init__(
+                self, _a: int, _b: str = "default", *_args: Any, **_kwargs: Any
+            ) -> None:
+                pass
+
+        # Create an instance which will call __init__
+        _ = TestClass(42, "hello", "extra", key="value")
+
+        # Verify that the logger was called
+        self.mock_logger.debug.assert_called_once()
+        msg_dict = self.mock_logger.debug.call_args[0][0]
+        # Verify that class name was prepended to function name
+        self.assertEqual(msg_dict["func_name"], "TestClass.__init__")
+        # Verify the input contains the arguments
+        self.assertIn("_a", msg_dict["input"])
+        self.assertIn("42", msg_dict["input"])
+        self.assertIn("_b", msg_dict["input"])
+        self.assertIn("hello", msg_dict["input"])
+        self.assertIn("_args", msg_dict["input"])
+        self.assertIn("extra", msg_dict["input"])
+        self.assertIn("_kwargs", msg_dict["input"])
+        self.assertIn("key", msg_dict["input"])
+        self.assertIn("value", msg_dict["input"])
+
 
 if __name__ == "__main__":
     unittest.main()

torchrec/distributed/torchrec_logger.py

@@ -42,12 +42,16 @@ def _get_logging_handler(
 
 
 def _get_msg_dict(func_name: str, **kwargs: Any) -> dict[str, Any]:
-    msg_dict = {
-        "func_name": f"{func_name}",
-    }
-    if dist.is_initialized():
-        group = kwargs.get("group") or kwargs.get("process_group")
-        msg_dict["group"] = f"{group}"
-        msg_dict["world_size"] = f"{dist.get_world_size(group)}"
-        msg_dict["rank"] = f"{dist.get_rank(group)}"
-    return msg_dict
+    try:
+        msg_dict = {
+            "func_name": f"{func_name}",
+        }
+        if dist.is_initialized():
+            group = kwargs.get("group") or kwargs.get("process_group")
+            msg_dict["group"] = f"{group}"
+            msg_dict["world_size"] = f"{dist.get_world_size(group)}"
+            msg_dict["rank"] = f"{dist.get_rank(group)}"
+        return msg_dict
+    except Exception as error:
+        logging.error(f"Torchrec Logger: Error in _get_msg_dict: {error}")
+        return {"_get_msg_dict_error": str(error)}