Skip to content

Commit 7924c6f

Browse files
gnahzgfacebook-github-bot
authored andcommitted
add __str__ func to ShardingOption for esay debugging purpose (#1981)
Summary: Pull Request resolved: #1981 As titled Reviewed By: henrylhtsang Differential Revision: D57215579 fbshipit-source-id: b9a72aea12ce294f2325f1de45489e1149c153e0
1 parent 9beeb6c commit 7924c6f

File tree

1 file changed

+14
-0
lines changed

1 file changed

+14
-0
lines changed

torchrec/distributed/planner/types.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,9 @@ def __hash__(self) -> int:
315315
)
316316
)
317317

318+
def __str__(self) -> str:
319+
return f"Shard size: {tuple(self.size)}, offset: {tuple(self.offset)}, storage: {str(self.storage)}, perf: {str(self.perf)}, rank: {self.rank}"
320+
318321

319322
class ShardingOption:
320323
"""
@@ -491,6 +494,17 @@ def __deepcopy__(
491494
setattr(result, k, deepcopy(v, memo))
492495
return result
493496

497+
def __str__(self) -> str:
498+
str_obj: str = ""
499+
str_obj += f"name: {self.name}"
500+
str_obj += f"\nsharding type: {self.sharding_type}"
501+
str_obj += f"\ncompute kernel: {self.compute_kernel}"
502+
str_obj += f"\nnum shards: {len(self.shards)}"
503+
for shard in self.shards:
504+
str_obj += f"\n\t{str(shard)}"
505+
506+
return str_obj
507+
494508

495509
class PartitionByType(Enum):
496510
"""

0 commit comments

Comments
 (0)