class_docstr: Handles input pipeline for SimdMeshImpl.
class_docstr_tok: ["Handles","input","pipeline","for","SimdMeshImpl","."]
docstr: Finds which host should read which sub-batch.
docstr_tok: ["Finds","which","host","should","read","which","sub-batch","."]
code: def _get_hosts_to_hold_ds(self, pnum_map):
      """Finds which host should read which sub-batch."""
      assert _NONE_PNUM not in pnum_map
  
      # This records how many datasets (ds) are already stored on each host.
      num_dss_per_host = [0] * self._p_dev.num_hosts
  
      # A list of host_ids that holds datasets (ds).
      hosts_to_hold_ds = []
  
      def _get_num_pnums_per_host(sub_batch_pnum_map):
          num_pnums_per_host = [0] * self._p_dev.num_hosts
          for pnum in sub_batch_pnum_map.flatten():
              num_pnums_per_host[
                  self._p_dev.ordered_host_ids[pnum]
              ] += 1
          return num_pnums_per_host
  
      def _find_host_id_with_most_pnums_and_least_ds(
          num_pnums_per_host, num_dss_per_host
      ):
          host_metics = [
              (
                  host_id,
                  num_pnums_per_host[host_id],
                  num_dss_per_host[host_id],
              )
              for host_id in range(self._p_dev.num_hosts)
          ]
          # Major max key: num_pnums
          # Minor max key: -num_dss. We need to find a relatively spare host.
          host_id, _, _ = max(
              host_metics, key=lambda keys: (keys[1], -keys[2])
          )
          return host_id
  
      for sub_batch_pnum_map in pnum_map:
          num_pnums_per_host = _get_num_pnums_per_host(
              sub_batch_pnum_map
          )
          host_id = _find_host_id_with_most_pnums_and_least_ds(
              num_pnums_per_host, num_dss_per_host
          )
          num_dss_per_host[host_id] += 1
          hosts_to_hold_ds.append(host_id)
  
      return hosts_to_hold_ds
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import threading
  import time
  import numpy
  from six.moves import range
  from six.moves import zip
  import tensorflow.compat.v1
  from tensorflow.python.framework import ops
  from tensorflow.python.tpu import tpu_feed
label_desc: Use image node_id 6 for calling the SimdMeshImplInputReader obj's underlying member method code with example usage: obj._get_hosts_to_hold_ds(pnum_map) and returns: hosts_to_hold_ds, num_pnums_per_host, host_id