Enable logging for the plan() function, ShardEstimators and TrainingPipeline class constructors (meta-pytorch#3576)

Summary: Pull Request resolved: meta-pytorch#3576

Differential Revision: D87910772

Author: nipung90

torchrec/distributed/embeddingbag.py

@@ -120,6 +120,15 @@
 except OSError:
     pass
 
+try:
+    # This is a safety measure against torch package issues for when
+    # Torchrec is included in the inference side model code. We should
+    # remove this once we are sure all model side packages have the required
+    # dependencies
+    from torchrec.distributed.logger import _torchrec_method_logger
+except Exception:
+    from torchrec.distributed.empty_logger import _torchrec_method_logger
+
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -466,6 +475,7 @@ class ShardedEmbeddingBagCollection(
     This is part of the public API to allow for manual data dist pipelining.
     """
 
+    @_torchrec_method_logger()
     def __init__(
         self,
         module: EmbeddingBagCollectionInterface,
@@ -2021,6 +2031,7 @@ class ShardedEmbeddingBag(
     This is part of the public API to allow for manual data dist pipelining.
     """
 
+    @_torchrec_method_logger()
     def __init__(
         self,
         module: nn.EmbeddingBag,

torchrec/distributed/empty_logger.py

@@ -0,0 +1,28 @@
+#!/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.
+
+import functools
+from typing import Any, Callable, ParamSpec, TypeVar
+
+
+_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:
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

torchrec/distributed/planner/planners.py

@@ -70,6 +70,15 @@
 )
 from torchrec.distributed.utils import get_device_type, none_throws
 
+try:
+    # This is a safety measure against torch package issues for when
+    # Torchrec is included in the inference side model code. We should
+    # remove this once we are sure all model side packages have the required
+    # dependencies
+    from torchrec.distributed.logger import _torchrec_method_logger
+except Exception:
+    from torchrec.distributed.empty_logger import _torchrec_method_logger
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -498,6 +507,7 @@ def collective_plan(
             sharders,
         )
 
+    @_torchrec_method_logger()
     def plan(
         self,
         module: nn.Module,

torchrec/distributed/planner/shard_estimators.py

@@ -53,6 +53,15 @@
 
 from torchrec.modules.embedding_modules import EmbeddingBagCollectionInterface
 
+try:
+    # This is a safety measure against torch package issues for when
+    # Torchrec is included in the inference side model code. We should
+    # remove this once we are sure all model side packages have the required
+    # dependencies
+    from torchrec.distributed.logger import _torchrec_method_logger
+except Exception:
+    from torchrec.distributed.empty_logger import _torchrec_method_logger
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -955,6 +964,7 @@ class EmbeddingStorageEstimator(ShardEstimator):
         is_inference (bool): If the model is inference model. Default to False.
     """
 
+    @_torchrec_method_logger()
     def __init__(
         self,
         topology: Topology,

torchrec/distributed/shard.py

@@ -30,6 +30,15 @@
 from torchrec.distributed.utils import init_parameters
 from torchrec.modules.utils import reset_module_states_post_sharding
 
+try:
+    # This is a safety measure against torch package issues for when
+    # Torchrec is included in the inference side model code. We should
+    # remove this once we are sure all model side packages have the required
+    # dependencies
+    from torchrec.distributed.logger import _torchrec_method_logger
+except Exception:
+    from torchrec.distributed.empty_logger import _torchrec_method_logger
+
 
 def _join_module_path(path: str, name: str) -> str:
     return (path + "." + name) if path else name
@@ -146,6 +155,7 @@ def _shard(
 
 # pyre-ignore
 @contract()
+@_torchrec_method_logger()
 def shard_modules(
     module: nn.Module,
     env: Optional[ShardingEnv] = None,
@@ -194,6 +204,7 @@ def init_weights(m):
     return _shard_modules(module, env, device, plan, sharders, init_params)
 
 
+@_torchrec_method_logger()
 def _shard_modules(  # noqa: C901
     module: nn.Module,
     # TODO: Consolidate to using Dict[str, ShardingEnv]

torchrec/distributed/test_utils/test_model_parallel_base.py

@@ -248,12 +248,26 @@ def test_sharding_ebc_input_validation_enabled(self, mock_jk: Mock) -> None:
             values=torch.tensor([1, 2, 3, 4, 5]),
             lengths=torch.tensor([1, 2, 0, 2]),
             offsets=torch.tensor([0, 1, 3, 3, 5]),
-        )
+        ).to(self.device)
         mock_jk.return_value = True
 
         with self.assertRaisesRegex(ValueError, "keys must be unique"):
             model(kjt)
-        mock_jk.assert_called_once_with("pytorch/torchrec:enable_kjt_validation")
+
+        # Count only calls with the input "pytorch/torchrec:enable_kjt_validation"
+        # This ignores any other calls to justknobs_check() with other inputs
+        # and protects the test from breaking when new JK checks are added.
+        validation_calls = [
+            call
+            for call in mock_jk.call_args_list
+            if len(call[0]) > 0
+            and call[0][0] == "pytorch/torchrec:enable_kjt_validation"
+        ]
+        self.assertEqual(
+            1,
+            len(validation_calls),
+            "There should be exactly one call to JK with pytorch/torchrec:enable_kjt_validation",
+        )
 
     @patch("torch._utils_internal.justknobs_check")
     def test_sharding_ebc_validate_input_only_once(self, mock_jk: Mock) -> None:
@@ -271,7 +285,20 @@ def test_sharding_ebc_validate_input_only_once(self, mock_jk: Mock) -> None:
             model(kjt)
             model(kjt)
 
-        mock_jk.assert_called_once_with("pytorch/torchrec:enable_kjt_validation")
+        # Count only calls with the input "pytorch/torchrec:enable_kjt_validation"
+        # This ignores any other calls to justknobs_check() with other inputs
+        # and protects the test from breaking when new JK checks are added.
+        validation_calls = [
+            call
+            for call in mock_jk.call_args_list
+            if len(call[0]) > 0
+            and call[0][0] == "pytorch/torchrec:enable_kjt_validation"
+        ]
+        self.assertEqual(
+            1,
+            len(validation_calls),
+            "There should be exactly one call to JK with pytorch/torchrec:enable_kjt_validation",
+        )
         matched_logs = list(
             filter(lambda s: "Validating input features..." in s, logs.output)
         )
@@ -294,7 +321,20 @@ def test_sharding_ebc_input_validation_disabled(self, mock_jk: Mock) -> None:
         except ValueError:
             self.fail("Input validation should not be enabled.")
 
-        mock_jk.assert_called_once_with("pytorch/torchrec:enable_kjt_validation")
+        # Count only calls with the input "pytorch/torchrec:enable_kjt_validation"
+        # This ignores any other calls to justknobs_check() with other inputs
+        # and protects the test from breaking when new JK checks are added.
+        validation_calls = [
+            call
+            for call in mock_jk.call_args_list
+            if len(call[0]) > 0
+            and call[0][0] == "pytorch/torchrec:enable_kjt_validation"
+        ]
+        self.assertEqual(
+            1,
+            len(validation_calls),
+            "There should be exactly one call to JK with pytorch/torchrec:enable_kjt_validation",
+        )
 
     def _create_sharded_model(
         self, embedding_dim: int = 128, num_embeddings: int = 256

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -73,6 +73,14 @@
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 from torchrec.streamable import Pipelineable
 
+try:
+    # This is a safety measure against torch package issues for when
+    # Torchrec is included in the inference side model code. We should
+    # remove this once we are sure all model side packages have the required
+    # dependencies
+    from torchrec.distributed.logger import _torchrec_method_logger
+except Exception:
+    from torchrec.distributed.empty_logger import _torchrec_method_logger
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -106,6 +114,8 @@ class TrainPipeline(abc.ABC, Generic[In, Out]):
     def progress(self, dataloader_iter: Iterator[In]) -> Out:
         pass
 
+    # pyre-ignore [56]
+    @_torchrec_method_logger()
     def __init__(self) -> None:
         # pipeline state such as in foward, in backward etc, used in training recover scenarios
         self._state: PipelineState = PipelineState.IDLE

torchrec/modules/mc_embedding_modules.py

@@ -20,6 +20,15 @@
 from torchrec.modules.mc_modules import ManagedCollisionCollection
 from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor, KeyedTensor
 
+try:
+    # This is a safety measure against torch package issues for when
+    # Torchrec is included in the inference side model code. We should
+    # remove this once we are sure all model side packages have the required
+    # dependencies
+    from torchrec.distributed.logger import _torchrec_method_logger
+except Exception:
+    from torchrec.distributed.empty_logger import _torchrec_method_logger
+
 
 def evict(
     evictions: Dict[str, Optional[torch.Tensor]],
@@ -125,6 +134,7 @@ class ManagedCollisionEmbeddingCollection(BaseManagedCollisionEmbeddingCollectio
 
     """
 
+    @_torchrec_method_logger()
     def __init__(
         self,
         embedding_collection: EmbeddingCollection,
@@ -164,6 +174,7 @@ class ManagedCollisionEmbeddingBagCollection(BaseManagedCollisionEmbeddingCollec
 
     """
 
+    @_torchrec_method_logger()
     def __init__(
         self,
         embedding_bag_collection: EmbeddingBagCollection,