code: def __init__(
      self, mesh_shape, node_count, model: Callable, version=1
  ):
      assert ray.is_initialized()  # needs a valid ray cluster to start
      self.nodes = []
      self.node_count = node_count
      self.dp, self.mp = mesh_shape
      self.version = version
  
      start = time.time()
  
      for i in range(node_count):
          self.nodes.append(
              NetworkRunner.options(max_concurrency=2).remote(
                  mesh_shape, model
              )
          )
  
      for n in self.nodes:
          n.run.remote()
  
      params = []
      for n in self.nodes:
          params.append(n.get_params.remote())
  
      self.param_count = ray.get(params)[0]
      print(f"Ray actors created in {time.time() - start:.06}s")
  
file_path: mesh-transformer-jax/mesh_transformer/TPU_cluster.py
imports: import itertools
  import json
  import time
  import ray
  from typing import Callable
  import numpy
  from mesh_transformer.train_actor import NetworkRunner
  from google.cloud import storage
  from smart_open import open
  from func_timeout import func_set_timeout
label_desc: Use image node_id 1 to create a new TPUCluster object  with example: obj = TPUCluster(mesh_shape, node_count, model, version)