code: def big_batch_search(
      index,
      xq,
      k,
      method="knn_function",
      pairwise_distances=faiss.pairwise_distances,
      knn=faiss.knn,
      verbose=0,
      threaded=0,
      use_float16=False,
      prefetch_threads=1,
      computation_threads=1,
      q_assign=None,
      checkpoint=None,
      checkpoint_freq=7200,
      start_list=0,
      end_list=None,
      crash_at=-1,
  ):
      """
      Search queries xq in the IVF index, with a search function that collects
      batches of query vectors per inverted list. This can be faster than the
      regular search indexes.
      Supports IVFFlat, IVFPQ and IVFScalarQuantizer.
  
      Supports three computation methods:
      method = "index":
          build a flat index and populate it separately for each index
      method = "pairwise_distances":
          decompress codes and compute all pairwise distances for the queries
          and index and add result to heap
      method = "knn_function":
          decompress codes and compute knn results for the queries
  
      threaded=0: sequential execution
      threaded=1: prefetch next bucket while computing the current one
      threaded=2: prefetch prefetch_threads buckets at a time.
  
      compute_threads>1: the knn function will get an additional thread_no that
          tells which worker should handle this.
  
      In threaded mode, the computation is tiled with the bucket perparation and
      the writeback of results (useful to maximize GPU utilization).
  
      use_float16: convert all matrices to float16 (faster for GPU gemm)
  
      q_assign: override coarse assignment, should be a matrix of size nq * nprobe
  
      checkpointing (only for threaded > 1):
      checkpoint: file where the checkpoints are stored
      checkpoint_freq: when to perform checkpoinging. Should be a multiple of threaded
  
      start_list, end_list: process only a subset of invlists
      """
      nprobe = index.nprobe
  
      assert method in ("index", "pairwise_distances", "knn_function")
  
      mem_queries = xq.nbytes
      mem_assign = len(xq) * nprobe * np.dtype("int32").itemsize
      mem_res = (
          len(xq)
          * k
          * (np.dtype("int64").itemsize + np.dtype("float32").itemsize)
      )
      mem_tot = mem_queries + mem_assign + mem_res
      if verbose > 0:
          logging.info(
              f"memory: queries {mem_queries} assign {mem_assign} "
              f"result {mem_res} total {mem_tot} = {mem_tot / (1<<30):.3f} GiB"
          )
  
      bbs = BigBatchSearcher(
          index, xq, k, verbose=verbose, use_float16=use_float16
      )
  
      comp = BlockComputer(
          index,
          method=method,
          pairwise_distances=pairwise_distances,
          knn=knn,
      )
  
      bbs.decode_func = comp.decode_func
  
      bbs.by_residual = comp.by_residual
      if q_assign is None:
          bbs.coarse_quantization()
      else:
          bbs.q_assign = q_assign
      bbs.reorder_assign()
  
      if end_list is None:
          end_list = index.nlist
  
      completed = set()
      if checkpoint is not None:
          assert (start_list, end_list) == (0, index.nlist)
          if os.path.exists(checkpoint):
              logging.info(f"recovering checkpoint: {checkpoint}")
              completed = bbs.read_checkpoint(checkpoint)
              logging.info(f"   already completed: {len(completed)}")
          else:
              logging.info("no checkpoint: starting from scratch")
  
      if threaded == 0:
          # simple sequential version
  
          for l in range(start_list, end_list):
              bbs.report(l)
              q_subset, xq_l, list_ids, xb_l = bbs.prepare_bucket(l)
              t0i = time.time()
              D, I = comp.block_search(xq_l, xb_l, list_ids, k)
              bbs.t_accu[2] += time.time() - t0i
              bbs.add_results_to_heap(q_subset, D, list_ids, I)
  
      elif threaded == 1:
          # parallel version with granularity 1
  
          def add_results_and_prefetch(to_add, l):
              """perform the addition for the previous bucket and
              prefetch the next (if applicable)"""
              if to_add is not None:
                  bbs.add_results_to_heap(*to_add)
              if l < index.nlist:
                  return bbs.prepare_bucket(l)
  
          prefetched_bucket = bbs.prepare_bucket(start_list)
          to_add = None
          pool = ThreadPool(1)
  
          for l in range(start_list, end_list):
              bbs.report(l)
              prefetched_bucket_a = pool.apply_async(
                  add_results_and_prefetch, (to_add, l + 1)
              )
              q_subset, xq_l, list_ids, xb_l = prefetched_bucket
              bbs.start_t_accu()
              D, I = comp.block_search(xq_l, xb_l, list_ids, k)
              bbs.stop_t_accu(2)
              to_add = q_subset, D, list_ids, I
              bbs.start_t_accu()
              prefetched_bucket = prefetched_bucket_a.get()
              bbs.stop_t_accu(4)
  
          bbs.add_results_to_heap(*to_add)
          pool.close()
      else:
  
          def task_manager_thread(
              task,
              pool_size,
              start_task,
              end_task,
              completed,
              output_queue,
              input_queue,
          ):
              try:
                  with ThreadPool(pool_size) as pool:
                      res = [
                          pool.apply_async(
                              task, args=(i, output_queue, input_queue)
                          )
                          for i in range(start_task, end_task)
                          if i not in completed
                      ]
                      for r in res:
                          r.get()
                      pool.close()
                      pool.join()
                  output_queue.put(None)
              except:
                  traceback.print_exc()
                  _thread.interrupt_main()
                  raise
  
          def task_manager(*args):
              task_manager = threading.Thread(
                  target=task_manager_thread,
                  args=args,
              )
              task_manager.daemon = True
              task_manager.start()
              return task_manager
  
          def prepare_task(task_id, output_queue, input_queue=None):
              try:
                  logging.info(f"Prepare start: {task_id}")
                  q_subset, xq_l, list_ids, xb_l = bbs.prepare_bucket(
                      task_id
                  )
                  output_queue.put(
                      (task_id, q_subset, xq_l, list_ids, xb_l)
                  )
                  logging.info(f"Prepare end: {task_id}")
              except:
                  traceback.print_exc()
                  _thread.interrupt_main()
                  raise
  
          def compute_task(task_id, output_queue, input_queue):
              try:
                  logging.info(f"Compute start: {task_id}")
                  t_wait_out = 0
                  while True:
                      t0 = time.time()
                      logging.info(f"Compute input: task {task_id}")
                      input_value = input_queue.get()
                      t_wait_in = time.time() - t0
                      if input_value is None:
                          # signal for other compute tasks
                          input_queue.put(None)
                          break
                      (
                          centroid,
                          q_subset,
                          xq_l,
                          list_ids,
                          xb_l,
                      ) = input_value
                      logging.info(
                          f"Compute work: task {task_id}, centroid {centroid}"
                      )
                      t0 = time.time()
                      if computation_threads > 1:
                          D, I = comp.block_search(
                              xq_l, xb_l, list_ids, k, thread_id=task_id
                          )
                      else:
                          D, I = comp.block_search(
                              xq_l, xb_l, list_ids, k
                          )
                      t_compute = time.time() - t0
                      logging.info(
                          f"Compute output: task {task_id}, centroid {centroid}"
                      )
                      t0 = time.time()
                      output_queue.put(
                          (
                              centroid,
                              t_wait_in,
                              t_wait_out,
                              t_compute,
                              q_subset,
                              D,
                              list_ids,
                              I,
                          )
                      )
                      t_wait_out = time.time() - t0
                  logging.info(f"Compute end: {task_id}")
              except:
                  traceback.print_exc()
                  _thread.interrupt_main()
                  raise
  
          prepare_to_compute_queue = Queue(2)
          compute_to_main_queue = Queue(2)
          compute_task_manager = task_manager(
              compute_task,
              computation_threads,
              0,
              computation_threads,
              set(),
              compute_to_main_queue,
              prepare_to_compute_queue,
          )
          prepare_task_manager = task_manager(
              prepare_task,
              prefetch_threads,
              start_list,
              end_list,
              completed,
              prepare_to_compute_queue,
              None,
          )
  
          t_checkpoint = time.time()
          while True:
              logging.info("Waiting for result")
              value = compute_to_main_queue.get()
              if not value:
                  break
              (
                  centroid,
                  t_wait_in,
                  t_wait_out,
                  t_compute,
                  q_subset,
                  D,
                  list_ids,
                  I,
              ) = value
              # to test checkpointing
              if centroid == crash_at:
                  1 / 0
              bbs.t_accu[2] += t_compute
              bbs.t_accu[4] += t_wait_in
              bbs.t_accu[5] += t_wait_out
              logging.info(f"Adding to heap start: centroid {centroid}")
              bbs.add_results_to_heap(q_subset, D, list_ids, I)
              logging.info(f"Adding to heap end: centroid {centroid}")
              completed.add(centroid)
              bbs.report(centroid)
              if checkpoint is not None:
                  if time.time() - t_checkpoint > checkpoint_freq:
                      logging.info("writing checkpoint")
                      bbs.write_checkpoint(checkpoint, completed)
                      t_checkpoint = time.time()
  
          prepare_task_manager.join()
          compute_task_manager.join()
  
      bbs.tic("finalize heap")
      bbs.rh.finalize()
      bbs.toc()
  
      return bbs.rh.D, bbs.rh.I
  
imports: import time
  import pickle
  import os
  import logging
  from multiprocessing.pool import ThreadPool
  import threading
  import _thread
  from queue import Queue
  import traceback
  import datetime
  import numpy
  import faiss
  from faiss.contrib.inspect_tools import get_invlist
label_desc: Use image node_id 1 for calling a global function with example usage: big_batch_search(index, xq, k, method, pairwise_distances, knn, verbose, threaded, use_float16, prefetch_threads, computation_threads, q_assign, checkpoint, checkpoint_freq, start_list, end_list, crash_at) and returns: bbs,  bbs, task_manager, bbs