class_docstr: Stats for a sharding planner execution.
class_docstr_tok: ["Stats","for","a","sharding","planner","execution","."]
docstr: Logs stats for a given sharding plan.
  
  Provides a tabular view of stats for the given sharding plan with per device
  storage usage (HBM and DDR), perf, input, output, and number/type of shards.
  
  Args:
      sharding_plan (ShardingPlan): sharding plan chosen by the planner.
      topology (Topology): device topology.
      batch_size (int): batch size.
      storage_reservation (StorageReservation): reserves storage for unsharded
          parts of the model
      num_proposals (int): number of proposals evaluated.
      num_plans (int): number of proposals successfully partitioned.
      run_time (float): time taken to find plan (in seconds).
      best_plan (List[ShardingOption]): plan with expected performance.
      constraints (Optional[Dict[str, ParameterConstraints]]): dict of parameter
          names to provided ParameterConstraints.
      debug (bool): whether to enable debug mode.
docstr_tok: ["Logs","stats","for","a","given","sharding","plan",".","Provides","a","tabular","view","of","stats","for","the","given","sharding","plan","with","per","device","storage","usage","(","HBM","and","DDR",")",",","perf",",","input",",","output",",","and","number\/type","of","shards",".","Args",":","sharding_plan","(","ShardingPlan",")",":","sharding","plan","chosen","by","the","planner",".","topology","(","Topology",")",":","device","topology",".","batch_size","(","int",")",":","batch","size",".","storage_reservation","(","StorageReservation",")",":","reserves","storage","for","unsharded","parts","of","the","model","num_proposals","(","int",")",":","number","of","proposals","evaluated",".","num_plans","(","int",")",":","number","of","proposals","successfully","partitioned",".","run_time","(","float",")",":","time","taken","to","find","plan","(","in","seconds",")",".","best_plan","(","List","[","ShardingOption","]",")",":","plan","with","expected","performance",".","constraints","(","Optional","[","Dict","[","str",",","ParameterConstraints","]","]",")",":","dict","of","parameter","names","to","provided","ParameterConstraints",".","debug","(","bool",")",":","whether","to","enable","debug","mode","."]
code: def log(
      self,
      sharding_plan: ShardingPlan,
      topology: Topology,
      batch_size: int,
      storage_reservation: StorageReservation,
      num_proposals: int,
      num_plans: int,
      run_time: float,
      best_plan: List[ShardingOption],
      constraints: Optional[Dict[str, ParameterConstraints]] = None,
      sharders: Optional[List[ModuleSharder[nn.Module]]] = None,
      debug: bool = True,
  ) -> None:
      """
      Logs stats for a given sharding plan.
  
      Provides a tabular view of stats for the given sharding plan with per device
      storage usage (HBM and DDR), perf, input, output, and number/type of shards.
  
      Args:
          sharding_plan (ShardingPlan): sharding plan chosen by the planner.
          topology (Topology): device topology.
          batch_size (int): batch size.
          storage_reservation (StorageReservation): reserves storage for unsharded
              parts of the model
          num_proposals (int): number of proposals evaluated.
          num_plans (int): number of proposals successfully partitioned.
          run_time (float): time taken to find plan (in seconds).
          best_plan (List[ShardingOption]): plan with expected performance.
          constraints (Optional[Dict[str, ParameterConstraints]]): dict of parameter
              names to provided ParameterConstraints.
          debug (bool): whether to enable debug mode.
      """
  
      shard_by_fqn = {
          module_name + "." + param_name: value
          for module_name, param_dict in sharding_plan.plan.items()
          # pyre-ignore - this is a EmbeddingShardingPlan below
          for param_name, value in param_dict.items()
      }
      stats: Dict[int, Dict[str, Any]] = {
          rank: {"type": {}, "input_sizes": 0.0, "output_sizes": 0.0}
          for rank in range(topology.world_size)
      }
  
      used_sharding_types = set()
      compute_kernels_to_count = defaultdict(int)
  
      reserved_hbm_percent = (
          storage_reservation._percentage
          if isinstance(
              storage_reservation,
              (
                  FixedPercentageStorageReservation,
                  HeuristicalStorageReservation,
                  InferenceStorageReservation,
              ),
          )
          else 0.0
      )
      dense_storage = (
          storage_reservation._dense_storage
          if isinstance(
              storage_reservation,
              (
                  HeuristicalStorageReservation,
                  InferenceStorageReservation,
              ),
          )
          and storage_reservation._dense_storage is not None
          else Storage(0, 0)
      )
      assert dense_storage
      kjt_storage = (
          storage_reservation._kjt_storage
          if isinstance(
              storage_reservation,
              (
                  HeuristicalStorageReservation,
                  InferenceStorageReservation,
              ),
          )
          and storage_reservation._kjt_storage
          else Storage(0, 0)
      )
      assert kjt_storage
  
      for sharding_option in best_plan:
          fqn = sharding_option.fqn
  
          compute_kernels_to_count[sharding_option.compute_kernel] += 1
  
          if shard_by_fqn.get(fqn) is None:
              continue
          shard: ParameterSharding = shard_by_fqn[fqn]
  
          ranks, input_sizes, output_sizes = self._get_shard_stats(
              shard=shard,
              sharding_option=sharding_option,
              world_size=topology.world_size,
              local_world_size=topology.local_world_size,
              constraints=constraints,
          )
          sharding_type_abbr = _get_sharding_type_abbr(
              shard.sharding_type
          )
          used_sharding_types.add(sharding_type_abbr)
  
          for i, rank in enumerate(ranks):
              count = stats[rank]["type"].get(sharding_type_abbr, 0)
              stats[rank]["type"][sharding_type_abbr] = count + 1
              stats[rank]["input_sizes"] += input_sizes[i]
              stats[rank]["output_sizes"] += output_sizes[i]
  
      used_hbm = [0] * topology.world_size
      used_ddr = [0] * topology.world_size
      perf = [
          Perf(fwd_compute=0, fwd_comms=0, bwd_compute=0, bwd_comms=0)
          for _ in range(topology.world_size)
      ]
      for sharding_option in best_plan:
          for shard in sharding_option.shards:
              shard_storage = cast(Storage, shard.storage)
              rank = cast(int, shard.rank)
              used_hbm[rank] += shard_storage.hbm
              used_ddr[rank] += shard_storage.ddr
              perf[rank] += cast(Perf, shard.perf)
  
      used_hbm = [
          hbm + dense_storage.hbm + kjt_storage.hbm for hbm in used_hbm
      ]
      used_ddr = [
          ddr + dense_storage.ddr + kjt_storage.ddr for ddr in used_ddr
      ]
  
      table: List[List[Union[str, int]]] = [
          [
              "Rank",
              "HBM (GB)",
              "DDR (GB)",
              "Perf (ms)",
              "Input (MB)",
              "Output (MB)",
              "Shards",
          ],
          [
              "------",
              "----------",
              "----------",
              "-----------",
              "------------",
              "-------------",
              "--------",
          ],
      ]
  
      for rank, device in enumerate(topology.devices):
          used_hbm_gb = bytes_to_gb(used_hbm[rank])
          used_hbm_ratio = (
              used_hbm[rank]
              / ((1 - reserved_hbm_percent) * device.storage.hbm)
              if topology.compute_device == "cuda"
              and ((1 - reserved_hbm_percent) * device.storage.hbm) != 0
              else 0
          )
          used_ddr_gb = bytes_to_gb(used_ddr[rank])
          used_ddr_ratio = (
              used_ddr[rank] / device.storage.ddr
              if device.storage.ddr > 0
              else 0
          )
          for sharding_type in used_sharding_types:
              if sharding_type not in stats[rank]["type"]:
                  stats[rank]["type"][sharding_type] = 0
  
          rank_hbm = f"{round(used_hbm_gb, 1)} ({used_hbm_ratio:.0%})"
          rank_ddr = f"{round(used_ddr_gb, 1)} ({used_ddr_ratio:.0%})"
          rank_perf = _format_perf_breakdown(perf[rank])
          rank_input = f"{round(stats[rank]['input_sizes'], 2)}"
          rank_output = f"{round(stats[rank]['output_sizes'], 2)}"
          rank_shards = " ".join(
              f"{sharding_type}: {num_tables}"
              for sharding_type, num_tables in sorted(
                  stats[rank]["type"].items()
              )
          )
          table.append(
              [
                  rank,
                  rank_hbm,
                  rank_ddr,
                  rank_perf,
                  rank_input,
                  rank_output,
                  rank_shards,
              ]
          )
      formatted_table = _format_table(table)
      self._width = max(self._width, len(formatted_table[0]) + 8)
  
      if debug:
          param_table: List[List[Union[str, int]]] = [
              [
                  "FQN",
                  "Sharding",
                  "Compute Kernel",
                  "Perf (ms)",
                  "Storage (HBM, DDR)",
                  "Cache Load Factor",
                  "Pooling Factor",
                  "Num Poolings",
                  "Output",
                  "Weighing",
                  "Sharder",
                  "Features",
                  "Emb Dim (CW Dim)",
                  "Hash Size",
                  "Ranks",
              ],
              [
                  "-----",
                  "----------",
                  "----------------",
                  "-----------",
                  "--------------------",
                  "-------------------",
                  "----------------",
                  "--------------",
                  "--------",
                  "----------",
                  "---------",
                  "----------",
                  "------------------",
                  "-----------",
                  "-------",
              ],
          ]
          feat_batch_sizes = [
              constraints[so.name].batch_sizes
              if constraints and constraints.get(so.name)
              else None
              for so in best_plan
          ]
  
          sharder_map: Dict[str, ModuleSharder[nn.Module]] = {
              get_sharder_name(sharder.module_type): sharder
              # pyre-ignore - this is a ModuleSharder below
              for sharder in sharders
              if sharders
          }
  
          if include_batch_sizes := any(feat_batch_sizes):
              param_table[0].append("Batch Sizes")
              param_table[1].append("-------------")
          for i, so in enumerate(best_plan):
              ranks = sorted(
                  [cast(int, shard.rank) for shard in so.shards]
              )
              ranks = _collapse_consecutive_ranks(ranks)
  
              so_perf = Perf(
                  fwd_compute=0, fwd_comms=0, bwd_compute=0, bwd_comms=0
              )
              for shard in so.shards:
                  so_perf += cast(Perf, shard.perf)
  
              shard_perfs = _format_perf_breakdown(so_perf)
  
              so_storage = Storage(hbm=0, ddr=0)
              for shard in so.shards:
                  so_storage += cast(Storage, shard.storage)
  
              shard_storages = _format_storage_breakdown(so_storage)
  
              pooling_factor = str(round(sum(so.input_lengths), 3))
              num_poolings = (
                  cast(List[float], constraints[so.name].num_poolings)
                  if constraints
                  and constraints.get(so.name)
                  and constraints[so.name].num_poolings
                  else [NUM_POOLINGS] * len(so.input_lengths)
              )
              num_poolings = str(round(sum(num_poolings), 3))
              output = "pooled" if so.is_pooled else "sequence"
              weighing = "weighted" if so.is_weighted else "unweighted"
              sharder = sharder_map.get(
                  get_sharder_name(type(so.module[1])), None
              )
              sharder_name = type(sharder).__name__
              num_features = len(so.input_lengths)
              embedding_dim = (
                  f"{so.tensor.shape[1]} ({so.shards[0].size[1]})"
                  if so.sharding_type == ShardingType.COLUMN_WISE.value
                  or so.sharding_type
                  == ShardingType.TABLE_COLUMN_WISE.value
                  else f"{so.tensor.shape[1]}"
              )
              cache_load_factor = _get_cache_load_factor(
                  sharding_option=so
              )
              cache_load_factor = (
                  str(cache_load_factor)
                  if cache_load_factor is not None
                  else "None"
              )
              hash_size = so.tensor.shape[0]
              param_table.append(
                  [
                      so.fqn,
                      _get_sharding_type_abbr(so.sharding_type),
                      so.compute_kernel,
                      shard_perfs,
                      shard_storages,
                      cache_load_factor,
                      pooling_factor,
                      num_poolings,
                      output,
                      weighing,
                      sharder_name,
                      num_features,
                      embedding_dim,
                      hash_size,
                      ",".join(ranks) if sharding_plan.plan else "None",
                  ]
              )
              if include_batch_sizes:
                  bs = feat_batch_sizes[i]
                  param_table[-1].append(
                      _reduce_int_list(bs) if bs else "n/a"
                  )
          formatted_param_table = _format_table(param_table)
          self._width = max(
              self._width, len(formatted_param_table[0]) + 6
          )
  
      self._stats_table.clear()
      self._stats_table.append("#" * self._width)
      header_text = "--- Planner Statistics ---"
      self._stats_table.append(f"#{header_text: ^{self._width-2}}#")
  
      iter_text = (
          f"--- Evaluated {num_proposals} proposal(s), "
          f"found {num_plans} possible plan(s), "
          f"ran for {run_time:.2f}s ---"
      )
      self._stats_table.append(f"#{iter_text: ^{self._width-2}}#")
  
      divider = "-" * (self._width - 4)
      self._stats_table.append(f"#{divider: ^{self._width-2}}#")
  
      if sharding_plan.plan:
          for row in formatted_table:
              self._stats_table.append(f"# {row: <{self._width-3}}#")
  
          perf_breakdown = "Perf: Total perf (Forward compute, Forward comms, Backward compute, Backward comms, Prefetch compute)"
          legend = "Input: MB/iteration, Output: MB/iteration, Shards: number of tables"
          hbm_info = "HBM: estimated peak memory usage for shards, dense tensors, and features (KJT)"
          self._stats_table.append(f"#{'' : ^{self._width-2}}#")
          self._stats_table.append(
              f"# {perf_breakdown: <{self._width-3}}#"
          )
          self._stats_table.append(f"# {legend: <{self._width-3}}#")
          self._stats_table.append(f"# {hbm_info: <{self._width-3}}#")
  
      if debug:
          self._stats_table.append(f"#{'' : ^{self._width-2}}#")
          self._stats_table.append(
              f"# {'Parameter Info:' : <{self._width-3}}#"
          )
          for row in formatted_param_table:
              self._stats_table.append(f"# {row: <{self._width-3}}#")
  
      batch_size_text = f"Batch Size: {batch_size}"
      self._stats_table.append(f"#{'' : ^{self._width-2}}#")
      self._stats_table.append(
          f"# {batch_size_text : <{self._width-3}}#"
      )
  
      if not sharding_plan.plan:
          rank_size_text = f"World Size: {topology.world_size}"
          self._stats_table.append(f"#{'' : ^{self._width-2}}#")
          self._stats_table.append(
              f"# {rank_size_text : <{self._width-3}}#"
          )
  
      self._log_compute_kernel_stats(compute_kernels_to_count)
  
      if debug:
          if sharding_plan.plan:
              self._log_max_perf_and_max_hbm(perf, used_hbm)
          self._log_storage_reservation_stats(
              storage_reservation,
              topology,
              reserved_hbm_percent,
              dense_storage,
              kjt_storage,
          )
          if sharding_plan.plan:
              self._log_imbalance_tables(best_plan)
  
      self._stats_table.append("#" * self._width)
  
      for row in self._stats_table:
          logger.info(row)
  
imports: import logging
  from collections import defaultdict
  from typing import Any, cast, Dict, List, Optional, Tuple, Union
  from torch import nn
  from torchrec.distributed.planner.constants import BIGINT_DTYPE, NUM_POOLINGS
  from torchrec.distributed.planner.shard_estimators import _calculate_shard_io_sizes
  from torchrec.distributed.planner.storage_reservations import FixedPercentageStorageReservation, HeuristicalStorageReservation, InferenceStorageReservation
  from torchrec.distributed.planner.types import ParameterConstraints, Perf, ShardingOption, Stats, Storage, StorageReservation, Topology
  from torchrec.distributed.planner.utils import _find_imbalance_tables, bytes_to_gb, bytes_to_mb, sharder_name
  from torchrec.distributed.types import ModuleSharder, ParameterSharding, ShardingPlan, ShardingType
label_desc: Use image node_id 2 for calling the EmbeddingStats obj's underlying member method code with example usage: obj.log(sharding_plan, topology, batch_size, storage_reservation, num_proposals, num_plans, run_time, best_plan, constraints, sharders, debug) without return types