From df967f6a60f51c8220d722da551d996128588c8d Mon Sep 17 00:00:00 2001 From: Greg Macnamara Date: Mon, 10 Nov 2025 20:04:24 -0800 Subject: [PATCH] Add hash and equality functions to planner types (#3522) Summary: This diff adds explicit `__hash__` and `__eq__` methods to custom objects in the torchrec planner that previously relied on Python's default object identity-based hashing. 1. **DeviceHardware** (fbcode/torchrec/distributed/planner/types.py): - Hash based on rank, storage, and perf fields - Equality checks all three fields 2. **BasicCommsBandwidths** (fbcode/torchrec/distributed/planner/types.py): - Hash based on inter_host_bw and intra_host_bw - Equality checks both bandwidth values Reviewed By: iamzainhuda Differential Revision: D85731424 --- torchrec/distributed/planner/types.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/torchrec/distributed/planner/types.py b/torchrec/distributed/planner/types.py index db185d47c..32a750290 100644 --- a/torchrec/distributed/planner/types.py +++ b/torchrec/distributed/planner/types.py @@ -144,6 +144,18 @@ class DeviceHardware: storage: Storage perf: Perf + def __hash__(self) -> int: + return hash((self.rank, self.storage, self.perf)) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, DeviceHardware): + return False + return ( + self.rank == other.rank + and self.storage == other.storage + and self.perf == other.perf + ) + class CustomTopologyData: """ @@ -247,6 +259,17 @@ def get_bw( else: return self.inter_host_bw + def __hash__(self) -> int: + return hash((self._inter_host_bw, self._intra_host_bw)) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, BasicCommsBandwidths): + return False + return ( + self._inter_host_bw == other._inter_host_bw + and self._intra_host_bw == other._intra_host_bw + ) + class Topology: """