Add MultiProcessMock for cross-process testing

Summary:
**WHY:**
Distributed testing in TorchRec requires the ability to mock functions and methods across multiple processes. Without this capability, tests that spawn multiple processes cannot easily mock expensive operations, external dependencies, or specific behaviors, making it difficult to write comprehensive unit tests for distributed PyTorch code. This leads to either incomplete test coverage or tests that require actual expensive resources.

This change is also for T245908194

**WHAT:**
1. Created `MultiProcessMock` class in `/data/sandcastle/boxes/fbsource/fbcode/torchrec/distributed/test_utils/multi_process.py`:
   - Added `mocks` list to maintain mock configurations
   - Added `add_mock(target, return_value, side_effect, **kwargs)` to register new mocks
   - Added `apply_mocks()` to apply all registered mocks in child processes
   - Added `clear_mocks()` to clear all registered mocks

2. Integrated `MultiProcessMock` with `MultiProcessTestBase`:
   - Added `_mock_manager` instance variable initialized in `__init__`
   - Added public `add_mock(...)` method for users to register mocks
   - Created `_callable_wrapper_with_mocks()` static method to apply mocks before calling test functions
   - Modified `_run_multi_process_test()` to pass mock manager to child processes and use the wrapper
   - Modified `_run_multi_process_test_per_rank()` to pass mock manager to child processes and use the wrapper

3. Added comprehensive test suite in `/data/sandcastle/boxes/fbsource/fbcode/torchrec/distributed/tests/test_multi_process_mock.py`:
   - 6 unit tests for `MultiProcessMock` class functionality
   - 6 integration tests including baseline tests to ensure backward compatibility
   - Added BUCK target for the new test file

**TEST:**
All 12 tests pass successfully:
- Unit tests verify mock storage, multiple mocks, side effects, clearing, and application
- Integration tests verify cross-process mocking with return values, multiple mocks, and side effects
- Baseline tests confirm backward compatibility - existing tests without mocks continue to work exactly as before
- Fixed pickling issue by using module-level functions for side effects (required for multiprocessing)

Ran: `buck2 test torchrec/distributed/tests:test_multi_process_mock`
Result: Tests finished: Pass 12. Fail 0. Fatal 0. Skip 0.

Differential Revision: D88696058

Author: tcauth

torchrec/distributed/test_utils/multi_process.py

@@ -14,6 +14,7 @@
 import os
 import unittest
 from typing import Any, Callable, Dict, List, Optional
+from unittest.mock import patch
 
 import torch
 import torch.distributed as dist
@@ -25,6 +26,65 @@
 )
 
 
+class MultiProcessMock:
+    """
+    Manages cross-process mocks for multi-process testing.
+
+    This class maintains a collection of mocks that can be applied across
+    different processes in distributed testing scenarios.
+    """
+
+    def __init__(self) -> None:
+        self.mocks: List[Dict[str, Any]] = []
+
+    def add_mock(
+        self,
+        target: str,
+        return_value: Any = None,
+        side_effect: Any = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Add a new cross-process mock.
+
+        Args:
+            target: The target to mock (e.g., 'module.function')
+            return_value: The return value for the mock
+            side_effect: The side effect for the mock
+            **kwargs: Additional arguments to pass to the mock
+        """
+        mock_config = {
+            "target": target,
+            "return_value": return_value,
+            "side_effect": side_effect,
+            **kwargs,
+        }
+        self.mocks.append(mock_config)
+
+    def apply_mocks(self) -> List[Any]:
+        """
+        Apply all registered mocks and return context managers.
+
+        Returns:
+            List of active mock context managers
+        """
+        active_patches = []
+        for mock_config in self.mocks:
+            target = mock_config["target"]
+            return_value = mock_config.get("return_value")
+            side_effect = mock_config.get("side_effect")
+
+            patcher = patch(target, return_value=return_value, side_effect=side_effect)
+            active_patch = patcher.__enter__()
+            active_patches.append((patcher, active_patch))
+
+        return active_patches
+
+    def clear_mocks(self) -> None:
+        """Clear all registered mocks."""
+        self.mocks.clear()
+
+
 class MultiProcessContext:
     def __init__(
         self,
@@ -111,6 +171,32 @@ def __init__(
             self._mp_init_mode: str = mp_init_mode
         logging.info(f"Using {self._mp_init_mode} for multiprocessing")
 
+        # Initialize MultiProcessMock
+        self._mock_manager = MultiProcessMock()
+
+    def add_mock(
+        self,
+        target: str,
+        return_value: Any = None,
+        side_effect: Any = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Add a new cross-process mock that will be applied during test execution.
+
+        Args:
+            target: The target to mock (e.g., 'module.function')
+            return_value: The return value for the mock
+            side_effect: The side effect for the mock
+            **kwargs: Additional arguments to pass to the mock
+        """
+        self._mock_manager.add_mock(
+            target=target,
+            return_value=return_value,
+            side_effect=side_effect,
+            **kwargs,
+        )
+
     @seed_and_log
     def setUp(self) -> None:
         os.environ["MASTER_ADDR"] = str("localhost")
@@ -149,8 +235,10 @@ def _run_multi_process_test(
         for rank in range(world_size):
             kwargs["rank"] = rank
             kwargs["world_size"] = world_size
+            kwargs["_mock_manager"] = self._mock_manager
             p = ctx.Process(
-                target=callable,
+                target=self._callable_wrapper_with_mocks,
+                args=(callable,),
                 kwargs=kwargs,
             )
             p.start()
@@ -176,9 +264,11 @@ def _run_multi_process_test_per_rank(
             kwargs = {}
             kwargs["rank"] = rank
             kwargs["world_size"] = world_size
+            kwargs["_mock_manager"] = self._mock_manager
             kwargs.update(kwargs_per_rank[rank])
             p = ctx.Process(
-                target=callable,
+                target=self._callable_wrapper_with_mocks,
+                args=(callable,),
                 kwargs=kwargs,
             )
             p.start()
@@ -188,6 +278,36 @@ def _run_multi_process_test_per_rank(
             p.join()
             self.assertEqual(0, p.exitcode)
 
+    @staticmethod
+    def _callable_wrapper_with_mocks(
+        callable: Callable[..., None],
+        _mock_manager: Optional[MultiProcessMock] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Wrapper that applies mocks before calling the target callable.
+
+        Args:
+            callable: The function to call
+            _mock_manager: Optional mock manager containing mocks to apply
+            **kwargs: Additional keyword arguments to pass to the callable
+        """
+        active_patches = []
+        try:
+            # Apply mocks if a mock manager is provided
+            if _mock_manager is not None:
+                active_patches = _mock_manager.apply_mocks()
+
+            # Remove _mock_manager from kwargs before calling the target
+            kwargs.pop("_mock_manager", None)
+
+            # Call the actual test callable
+            callable(**kwargs)
+        finally:
+            # Clean up all patches
+            for patcher, _ in active_patches:
+                patcher.__exit__(None, None, None)
+
 
 def _wrapper_func_for_multiprocessing(args):  # pyre-ignore[2, 3]
     """Wrapper function that unpacks arguments and calls the original func"""