class_docstr: Greedy Partitioner
  
  Args:
      sort_by (SortBy): Sort sharding options by storage or perf in
          descending order (i.e., large tables will be placed first).
      balance_modules (bool): Whether to sort by modules first, where
          smaller modules will be sorted first. In effect, this will place
          tables in each module in a balanced way.
docstr: Places sharding options on topology based on each sharding option's
  `partition_by` attribute.
  The topology, storage, and perfs are updated at the end of the placement.
  
  Args:
      proposal (List[ShardingOption]): list of populated sharding options.
      storage_constraint (Topology): device topology.
  
  Returns:
      List[ShardingOption]: list of sharding options for selected plan.
  
  Example::
  
      sharding_options = [
              ShardingOption(partition_by="uniform",
                      shards=[
                          Shards(storage=1, perf=1),
                          Shards(storage=1, perf=1),
                      ]),
              ShardingOption(partition_by="uniform",
                      shards=[
                          Shards(storage=2, perf=2),
                          Shards(storage=2, perf=2),
                      ]),
              ShardingOption(partition_by="device",
                      shards=[
                          Shards(storage=3, perf=3),
                          Shards(storage=3, perf=3),
                      ])
              ShardingOption(partition_by="device",
                      shards=[
                          Shards(storage=4, perf=4),
                          Shards(storage=4, perf=4),
                      ]),
          ]
      topology = Topology(world_size=2)
  
      # First [sharding_options[0] and sharding_options[1]] will be placed on the
      # topology with the uniform strategy, resulting in
  
      topology.devices[0].perf.total = (1,2)
      topology.devices[1].perf.total = (1,2)
  
      # Finally sharding_options[2] and sharding_options[3]] will be placed on the
      # topology with the device strategy (see docstring of `partition_by_device` for
      # more details).
  
      topology.devices[0].perf.total = (1,2) + (3,4)
      topology.devices[1].perf.total = (1,2) + (3,4)
  
      # The topology updates are done after the end of all the placements (the other
      # in the example is just for clarity).
docstr_tok: ["Places","sharding","options","on","topology","based","on","each","sharding","option's","`","partition_by","`","attribute",".","The","topology",",","storage",",","and","perfs","are","updated","at","the","end","of","the","placement",".","Args",":","proposal","(","List","[","ShardingOption","]",")",":","list","of","populated","sharding","options",".","storage_constraint","(","Topology",")",":","device","topology",".","Returns",":","List","[","ShardingOption","]",":","list","of","sharding","options","for","selected","plan",".","Example",":",":","sharding_options","=","[","ShardingOption","(","partition_by=","''","uniform","''",",","shards=","[","Shards","(","storage=1",",","perf=1",")",",","Shards","(","storage=1",",","perf=1",")",",","]",")",",","ShardingOption","(","partition_by=","''","uniform","''",",","shards=","[","Shards","(","storage=2",",","perf=2",")",",","Shards","(","storage=2",",","perf=2",")",",","]",")",",","ShardingOption","(","partition_by=","''","device","''",",","shards=","[","Shards","(","storage=3",",","perf=3",")",",","Shards","(","storage=3",",","perf=3",")",",","]",")","ShardingOption","(","partition_by=","''","device","''",",","shards=","[","Shards","(","storage=4",",","perf=4",")",",","Shards","(","storage=4",",","perf=4",")",",","]",")",",","]","topology","=","Topology","(","world_size=2",")","#","First","[","sharding_options","[","0","]","and","sharding_options","[","1","]","]","will","be","placed","on","the","#","topology","with","the","uniform","strategy",",","resulting","in","topology.devices","[","0","]",".perf.total","=","(","1,2",")","topology.devices","[","1","]",".perf.total","=","(","1,2",")","#","Finally","sharding_options","[","2","]","and","sharding_options","[","3","]","]","will","be","placed","on","the","#","topology","with","the","device","strategy","(","see","docstring","of","`","partition_by_device","`","for","#","more","details",")",".","topology.devices","[","0","]",".perf.total","=","(","1,2",")","+","(","3,4",")","topology.devices","[","1","]",".perf.total","=","(","1,2",")","+","(","3,4",")","#","The","topology","updates","are","done","after","the","end","of","all","the","placements","(","the","other","#","in","the","example","is","just","for","clarity",")","."]
code: def partition(
      self,
      proposal: List[ShardingOption],
      storage_constraint: Topology,
  ) -> List[ShardingOption]:
      """
      Places sharding options on topology based on each sharding option's
      `partition_by` attribute.
      The topology, storage, and perfs are updated at the end of the placement.
  
      Args:
          proposal (List[ShardingOption]): list of populated sharding options.
          storage_constraint (Topology): device topology.
  
      Returns:
          List[ShardingOption]: list of sharding options for selected plan.
  
      Example::
  
          sharding_options = [
                  ShardingOption(partition_by="uniform",
                          shards=[
                              Shards(storage=1, perf=1),
                              Shards(storage=1, perf=1),
                          ]),
                  ShardingOption(partition_by="uniform",
                          shards=[
                              Shards(storage=2, perf=2),
                              Shards(storage=2, perf=2),
                          ]),
                  ShardingOption(partition_by="device",
                          shards=[
                              Shards(storage=3, perf=3),
                              Shards(storage=3, perf=3),
                          ])
                  ShardingOption(partition_by="device",
                          shards=[
                              Shards(storage=4, perf=4),
                              Shards(storage=4, perf=4),
                          ]),
              ]
          topology = Topology(world_size=2)
  
          # First [sharding_options[0] and sharding_options[1]] will be placed on the
          # topology with the uniform strategy, resulting in
  
          topology.devices[0].perf.total = (1,2)
          topology.devices[1].perf.total = (1,2)
  
          # Finally sharding_options[2] and sharding_options[3]] will be placed on the
          # topology with the device strategy (see docstring of `partition_by_device` for
          # more details).
  
          topology.devices[0].perf.total = (1,2) + (3,4)
          topology.devices[1].perf.total = (1,2) + (3,4)
  
          # The topology updates are done after the end of all the placements (the other
          # in the example is just for clarity).
      """
  
      _topology: Topology = copy.deepcopy(storage_constraint)
      minheap_devices: Optional[List[OrderedDeviceHardware]] = None
      _host_level_devices = (
          GreedyPerfPartitioner._get_host_level_devices(_topology)
      )
  
      # first partition the uniform sharding options (RW & DP)
      uniform_sharding_options = _get_uniform_sharding_options(proposal)
      GreedyPerfPartitioner._uniform_partition(
          uniform_sharding_options, _topology.devices
      )
  
      # group the rest sharding options by colocation type (co-host, co-device, none)
      # and sort the groups by storage in reverse order
      sharding_option_groups = (
          _group_and_sort_non_uniform_sharding_options(
              proposal,
              sort_by=self._sort_by,
              balance_modules=self._balance_modules,
          )
      )
  
      for sharding_option_group in sharding_option_groups:
          if (
              sharding_option_group.sharding_options[0].partition_by
              == PartitionByType.HOST.value
          ):
              GreedyPerfPartitioner._cohost_partition(
                  sharding_option_group, _host_level_devices
              )
              # _cohost_partition invalidates minheap_devices, force rebuild before using
              minheap_devices = None
          elif (
              sharding_option_group.sharding_options[0].partition_by
              == PartitionByType.DEVICE.value
          ):
              if minheap_devices is None:
                  minheap_devices = (
                      GreedyPerfPartitioner._establish_minheap(
                          _topology.devices, _topology.local_world_size
                      )
                  )
              assert (
                  len(sharding_option_group.sharding_options) == 1
              ), f"Unexpected length for sharding options: {len(sharding_option_group.sharding_options)}"
              GreedyPerfPartitioner._device_partition(
                  sharding_option_group.sharding_options[0],
                  minheap_devices,
              )
          else:
              raise RuntimeError(
                  f"Unexpected sharding option group {sharding_option_group}"
              )
      # pyre-ignore [16]: `GreedyPerfPartitioner` has no attribute `_topology`.
      self._topology: Topology = _topology
      return proposal
  
imports: import copy
  import heapq
  import logging
  from dataclasses import dataclass
  from enum import Enum
  from typing import cast, Dict, List, Optional
  from torchrec.distributed.planner.perf_models import NoopPerfModel
  from torchrec.distributed.planner.types import DeviceHardware, PartitionByType, Partitioner, Perf, PerfModel, PlannerError, PlannerErrorType, ShardingOption, Storage, Topology
  from torchrec.distributed.planner.utils import bytes_to_gb, reset_shard_rank
  from torchrec.distributed.types import ShardingType
label_desc: Use image node_id 2 for calling the GreedyPerfPartitioner obj's underlying member method code with example usage: obj.partition(proposal, storage_constraint) and returns: proposal