code: def create_buckets_ordered_by_affinity(
      nparts_lhs: int, nparts_rhs: int, *, generator: random.Random
  ) -> List[Bucket]:
      """Try having consecutive buckets share as many partitions as possible.
  
      Start from a random bucket. Until there are buckets left, try to choose the
      next one so that it has as many partitions in common as possible with the
      previous one. When multiple options are available, pick one randomly.
  
      """
      if nparts_lhs <= 0 or nparts_rhs <= 0:
          return []
  
      # TODO Change this function to use the same cost model as the LockServer
      # when computing affinity (based on the number of entities to save and load)
      # rather than just the number of partitions in common. Pay attention to keep
      # the complexity of this algorithm linear in the number of buckets. This
      # comment is too short to give a full description, but the idea is that only
      # a few transitions are possible between a bucket and the next: the one that
      # preserves all (ent, part) pairs, the one that preserves only the lhs ones,
      # only the rhs ones, only the intersection of the two, or none at all. So we
      # can keep a dict from sets of (ent, part) to lists of buckets, and insert
      # each bucket into four of those lists, namely the ones for all its (ent,
      # part), its lhs ones, its rhs ones and the intersection of its lhs and rhs
      # ones. Then, when looking for the next bucket, we figure out the transition
      # that is cheapest (among the options defined above), determine the set of
      # (ent, part) we need to move to in order to achieve that transition type
      # and we look up in the dict to find a bucket containing those (ent, part).
  
      # This is our "source of truth" on what buckets we haven't outputted yet. It
      # can be queried in constant time.
      remaining: Set[Bucket] = set()
      # These are our random orders: we shuffle them once and then pop from the
      # end. Each bucket appears in several of them. They are updated lazily,
      # which means they may contain buckets that have already been outputted.
      all_buckets: List[Bucket] = []
      buckets_per_partition: List[List[Bucket]] = [
          [] for _ in range(max(nparts_lhs, nparts_rhs))
      ]
  
      for lhs in range(nparts_lhs):
          for rhs in range(nparts_rhs):
              b = Bucket(lhs, rhs)
              remaining.add(b)
              all_buckets.append(b)
              buckets_per_partition[lhs].append(b)
              buckets_per_partition[rhs].append(b)
  
      generator.shuffle(all_buckets)
      for buckets in buckets_per_partition:
          generator.shuffle(buckets)
  
      b = all_buckets.pop()
      remaining.remove(b)
      order = [b]
  
      while remaining:
          transposed_b = Bucket(b.rhs, b.lhs)
          if transposed_b in remaining:
              remaining.remove(transposed_b)
              order.append(transposed_b)
              if not remaining:
                  break
  
          same_as_lhs = buckets_per_partition[b.lhs]
          same_as_rhs = buckets_per_partition[b.rhs]
          while len(same_as_lhs) > 0 or len(same_as_rhs) > 0:
              (chosen,) = generator.choices(
                  [same_as_lhs, same_as_rhs],
                  weights=[len(same_as_lhs), len(same_as_rhs)],
              )
              next_b = chosen.pop()
              if next_b in remaining:
                  break
          else:
              while True:
                  next_b = all_buckets.pop()
                  if next_b in remaining:
                      break
          remaining.remove(next_b)
          order.append(next_b)
          b = next_b
  
      return order
  
imports: import logging
  import random
  from abc import ABC, abstractmethod
  from collections import defaultdict
  from statistics import mean
  from typing import Dict, List, NamedTuple, Optional, Set, Tuple
  from torchbiggraph.config import BucketOrder
  from torchbiggraph.distributed import Startable
  from torchbiggraph.rpc import Client, Server
  from torchbiggraph.stats import Stats, StatsHandler
  from torchbiggraph.types import Bucket, EntityName, Partition, Rank, Side
label_desc: Use image node_id 4 for calling a global function with example usage: create_buckets_ordered_by_affinity(nparts_lhs, nparts_rhs) and returns: order, list