Added unit test for Heuristic Storage Reservation (#3511)

Summary:
Pull Request resolved: #3511

**Context:**
Heuristic Storage reservation is a common component for all planner that checks if the given module along with the constraints can be sharded across the topology.

**In this diff:**

We added a UT to validate the error for storage use in the storage reservation process. If the given module is larger than the provided topology. We need to OOM the process asap with appropriate error to notify the PG

Reviewed By: kausv, mserturk

Differential Revision: D85892579

fbshipit-source-id: 03eb679e6cabf8c030092c46d93a7b030e4c3814

Author: Raahul Kalyaan Jakka

torchrec/distributed/planner/tests/test_storage_reservations.py

@@ -19,7 +19,7 @@
     _get_module_size,
     HeuristicalStorageReservation,
 )
-from torchrec.distributed.planner.types import Topology
+from torchrec.distributed.planner.types import PlannerError, PlannerErrorType, Topology
 
 from torchrec.distributed.test_utils.test_model import TestTowerInteraction
 from torchrec.distributed.types import ModuleSharder
@@ -36,6 +36,36 @@ def __init__(self, shardable_sparse: nn.Module) -> None:
 
 
 class TestHeuristicalStorageReservation(unittest.TestCase):
+
+    def test_validate_storage_reservations_errors(self) -> None:
+        tables = [
+            EmbeddingBagConfig(
+                num_embeddings=1_000_000,
+                embedding_dim=1024,
+                name="table_0",
+                feature_names=["feature_0"],
+            ),
+        ]
+
+        ebc = EmbeddingBagCollection(tables)
+        model = TestModel(shardable_sparse=ebc)
+
+        # Reserving 100% of HBM to make sure the heuristic storage reservation fails
+        heuristical_storage_reservation = HeuristicalStorageReservation(percentage=1)
+        with self.assertRaises(PlannerError) as context:
+            heuristical_storage_reservation.reserve(
+                topology=Topology(world_size=1, compute_device="cuda"),
+                batch_size=1024,
+                module=model,
+                sharders=cast(
+                    List[ModuleSharder[nn.Module]], [EmbeddingBagCollectionSharder()]
+                ),
+            )
+
+        self.assertEqual(
+            context.exception.error_type, PlannerErrorType.INSUFFICIENT_STORAGE
+        )
+
     def test_storage_reservations_ebc(self) -> None:
         tables = [
             EmbeddingBagConfig(