code: def test_beam_search(
      n_peers=20,
      total_experts=128,
      batch_size=32,
      beam_size=4,
      parallel_rpc=4,
      grid_dims=(32, 32, 32),
  ):
      dht_instances = [hivemind.DHT(start=True)]
      initial_peers = dht_instances[0].get_visible_maddrs()
      dht_instances += [
          hivemind.DHT(initial_peers=initial_peers, start=True)
          for _ in range(n_peers - 1)
      ]
  
      real_experts = sorted(
          {
              "expert."
              + ".".join(
                  [str(random.randint(0, dim - 1)) for dim in grid_dims]
              )
              for _ in range(total_experts)
          }
      )
      for batch_start in range(0, len(real_experts), batch_size):
          dht = random.choice(dht_instances)
          declare_experts(
              dht,
              real_experts[batch_start : batch_start + batch_size],
              get_dht_time() + 30,
          )
  
      neighbors = sum(
          [
              peer.get_visible_maddrs()
              for peer in random.sample(
                  dht_instances, min(3, len(dht_instances))
              )
          ],
          [],
      )
      you = hivemind.DHT(
          start=True, initial_peers=neighbors, parallel_rpc=parallel_rpc
      )
      beam_search = MoEBeamSearcher(you, "expert.", grid_dims)
  
      for i in range(10):
          topk_experts = beam_search.find_best_experts(
              [np.random.randn(dim) for dim in grid_dims], beam_size
          )
          assert all(
              isinstance(e, hivemind.RemoteExpert) for e in topk_experts
          )
          assert len(topk_experts) == beam_size
  
      for i in range(10):
          batch_experts = beam_search.batch_find_best_experts(
              [np.random.randn(batch_size, dim) for dim in grid_dims],
              beam_size=beam_size,
          )
          assert (
              isinstance(batch_experts, list)
              and len(batch_experts) == batch_size
          )
          assert all(
              isinstance(e, hivemind.RemoteExpert)
              for experts in batch_experts
              for e in experts
          )
          assert all(
              len(experts) == beam_size for experts in batch_experts
          )
  
file_path: hivemind/tests/test_dht_experts.py
imports: import asyncio
  import random
  import time
  import numpy
  import pytest
  import hivemind
  from hivemind import get_dht_time
  from hivemind.dht.node import DHTNode
  from hivemind.moe.client.beam_search import MoEBeamSearcher
  from hivemind.moe.expert_uid import ExpertInfo, is_valid_prefix, is_valid_uid, split_uid
  from hivemind.moe.server.dht_handler import declare_experts, get_experts
label_desc: Use image node_id 2 for calling a global function with example usage: test_beam_search(n_peers, total_experts, batch_size, beam_size, parallel_rpc, grid_dims) without return types