code: def _pick_bucket(
      self, buckets: List[Bucket], maybe_old_bucket: Optional[Bucket]
  ) -> Bucket:
      # We return a bucket (the lexicographically smallest one) among those
      # that minimize the I/O cost of loading them (from scratch) or switching
      # to them (if another bucket was already loaded). The cost of loading a
      # bucket is the cost of loading all the embedding it needs, and is
      # proportional to the number of entities it needs. The cost of switching
      # buckets is the cost of storing the embeddings that were needed before
      # but are not needed anymore and, conversely, loading the ones that
      # weren't needed but now are: embeddings that were needed and still are
      # come for free! Thus this function will tend to keep at least one side
      # of the bucket in place and, depending on the "bipartitedness" of the
      # graph (ranging from each entity appearing on only one side to all
      # entities appearing on both sides), it may optimize further, for
      # example by swapping the two sides. When loading from scratch, it will
      # pick by preference diagonal buckets (i.e., of the form (i, i)), unless
      # the graph is fully bipartite.
      # TODO If init_tree is enabled, during the first pass we might want to
      # _not_ choose diagonal buckets as they tend to keep the number of
      # initialized embeddings small and thus delay the time at which other
      # trainers can start working.
      old_entities_parts: Set[Tuple[EntityName, Partition]] = set()
      if maybe_old_bucket is not None:
          old_entities_parts.update(
              (entity, maybe_old_bucket.lhs)
              for entity in self.entities_lhs
          )
          old_entities_parts.update(
              (entity, maybe_old_bucket.rhs)
              for entity in self.entities_rhs
          )
  
      buckets_by_cost: Dict[int, List[Bucket]] = defaultdict(list)
      for bucket in buckets:
          new_entities_parts: Set[Tuple[EntityName, Partition]] = set()
          new_entities_parts.update(
              (entity, bucket.lhs) for entity in self.entities_lhs
          )
          new_entities_parts.update(
              (entity, bucket.rhs) for entity in self.entities_rhs
          )
  
          cost = sum(
              self.average_entity_counts[entity]
              for entity, _ in new_entities_parts.symmetric_difference(
                  old_entities_parts
              )
          )
  
          buckets_by_cost[cost].append(bucket)
  
      min_cost = min(buckets_by_cost.keys())
      cheapest_buckets = buckets_by_cost[min_cost]
  
      # TODO It may be interesting to get a random bucket among the acquirable
      # ones (after filtering by highest affinity, if possible), rather than
      # the lexicographically smallest one, to better mix the edges, but we
      # should first empirically verify that this doesn't degrade accuracy.
      # return random.choice(cheapest_buckets)
  
      return cheapest_buckets[0]
  
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 5 for calling the LockServer obj's underlying member method code with example usage: obj._pick_bucket(buckets, maybe_old_bucket) and returns: cheapest_buckets