code: def benchmark_throughput(
      num_experts=16,
      num_handlers=None,
      num_clients=128,
      num_batches_per_client=16,
      expert_cls="ffn",
      hid_dim=1024,
      batch_size=2048,
      max_batch_size=None,
      backprop=True,
      device=None,
  ):
      assert (
          not hasattr(torch.cuda, "is_initialized")
          or not torch.cuda.is_initialized()
          or torch.device(device) == torch.device("cpu")
      )
      assert expert_cls in name_to_block
      max_batch_size = max_batch_size or batch_size * 4
      num_handlers = max(1, num_handlers or num_clients // 2)
      benchmarking_failed = mp.Event()
      can_start = mp.Event()
      timestamps = dict(started=time.perf_counter())
  
      try:
          server_dht = DHT(start=True)
          clients = [
              mp.Process(
                  target=client_process,
                  name=f"client_process-{i}",
                  args=(
                      can_start,
                      benchmarking_failed,
                      server_dht.get_visible_maddrs(),
                      server_dht.peer_id,
                      num_experts,
                      batch_size,
                      hid_dim,
                      num_batches_per_client,
                      backprop,
                  ),
                  daemon=True,
              )
              for i in range(num_clients)
          ]
  
          for client in clients:
              client.start()
  
          timestamps["launched_clients"] = timestamps[
              "began_launching_server"
          ] = time.perf_counter()
  
          device = device or (
              "cuda" if torch.cuda.is_available() else "cpu"
          )
          module_backends = {}
          for i in range(num_experts):
              expert = torch.jit.script(
                  name_to_block[expert_cls](hid_dim)
              )
              module_backends[f"expert.{i}"] = ModuleBackend(
                  name=f"expert.{i}",
                  module=expert,
                  optimizer=torch.optim.Adam(expert.parameters()),
                  args_schema=(BatchTensorDescriptor(hid_dim),),
                  outputs_schema=BatchTensorDescriptor(hid_dim),
                  max_batch_size=max_batch_size,
              )
          timestamps["created_experts"] = time.perf_counter()
  
          server = Server(
              dht=server_dht,
              module_backends=module_backends,
              num_connection_handlers=num_handlers,
              device=device,
          )
          server.start()
          server.ready.wait()
  
          timestamps["server_ready"] = time.perf_counter()
          can_start.set()
  
          for client in clients:
              client.join()
  
          timestamps["clients_finished"] = time.perf_counter()
  
      except BaseException as e:
          benchmarking_failed.set()
          raise e
      finally:
          for client in clients:
              if client.is_alive():
                  client.terminate()
          server.shutdown()
          timestamps["server_shutdown_finished"] = time.perf_counter()
          server.join()
  
      sys.stdout.flush()
      sys.stderr.flush()
      time_between = (
          lambda key1, key2: abs(timestamps[key2] - timestamps[key1])
          if (key1 in timestamps and key2 in timestamps)
          else float("nan")
      )
      total_examples = batch_size * num_clients * num_batches_per_client
  
      logger.info(
          "Benchmark finished, status:"
          + ["Success", "Failure"][benchmarking_failed.is_set()]
      )
      logger.info(
          f"Server parameters: num_experts={num_experts}, num_handlers={num_handlers}, "
          f"max_batch_size={max_batch_size}, expert_cls={expert_cls}, hid_dim={hid_dim}, device={device}"
      )
      logger.info(
          f"Client parameters: num_clients={num_clients}, num_batches_per_client={num_batches_per_client}, "
          f"batch_size={batch_size}, backprop={backprop}"
      )
  
      logger.info("Results: ")
      logger.info(
          f"\tServer startup took {time_between('began_launching_server', 'server_ready') :.3f} s. "
          f"({time_between('began_launching_server', 'created_experts') :.3f} s. experts + "
          f"{time_between('created_experts', 'server_ready') :.3f} s. networking)"
      )
      logger.info(
          f"\tProcessed {total_examples} examples in {time_between('server_ready', 'clients_finished') :.3f}"
      )
      logger.info(
          f"\tThroughput for {'forward + backward' if backprop else 'forward'} passes: "
          f"{total_examples / time_between('server_ready', 'clients_finished') :.3f} samples / s."
      )
      logger.info(
          f"\tBenchmarking took {time_between('started', 'server_shutdown_finished') :.3f} s."
      )
      if benchmarking_failed.is_set():
          logger.info(
              "Note: benchmark code failed, timing/memory results only indicate time till failure!"
          )
      print_device_info(device)
      sys.stdout.flush()
      sys.stderr.flush()
  
      assert not benchmarking_failed.is_set()
  
file_path: hivemind/benchmarks/benchmark_throughput.py
imports: import argparse
  import multiprocessing
  import random
  import sys
  import time
  import torch
  from hivemind.dht import DHT
  from hivemind.moe.client.expert import RemoteExpert
  from hivemind.moe.client.remote_expert_worker import RemoteExpertWorker
  from hivemind.moe.expert_uid import ExpertInfo
  from hivemind.moe.server import ModuleBackend, Server
  from hivemind.moe.server.layers import name_to_block
  from hivemind.p2p import P2P
  from hivemind.utils.limits import increase_file_limit
  from hivemind.utils.logging import get_logger, use_hivemind_log_handler
  from hivemind.utils.tensor_descr import BatchTensorDescriptor
label_desc: Use image node_id 3 for calling a global function with example usage: benchmark_throughput(num_experts, num_handlers, num_clients, num_batches_per_client, expert_cls, hid_dim, batch_size, max_batch_size, backprop, device) without return types