using this code for multi node tpu setup , from vllm import LLM, SamplingParams
import ray

from huggingface_hub import snapshot_download
import pandas as pd
import os
from transformers import AutoTokenizer
# Define a remote function

from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

from vllm import initialize_ray_cluster
from vllm.config import ParallelConfig
from vllm.executor.ray_utils import _wait_until_pg_removed
from vllm.utils import get_ip


@ray.remote(num_cpus=240, resources={"TPU": 4})
def run_inference(input_data):
    # Initialize the LLM model on each TPU
    llm = LLM(model=ss, enforce_eager=False,  max_model_len=4096, max_num_batched_tokens=32768, max_num_seqs=256, tensor_parallel_size=2, dtype='bfloat16', gpu_memory_utilization=0.96)
    
    sampling_params = SamplingParams(temperature=0, top_p=1.0, n=1, max_tokens=2048)
    output = llm.generate(input_data, sampling_params)
    return output

# Distribute the inference tasks across TPUs
def divide_list(lst, num_sublists):
    avg = len(lst) // num_sublists
    remainder = len(lst) % num_sublists
    start = 0
    sublists = []
    
    for i in range(num_sublists):
        end = start + avg + (1 if i < remainder else 0)
        sublists.append(lst[start:end])
        start = end
    
    return sublists



ss = snapshot_download('sparsh35/qwenmulticot')

def getchatml(ll):
    
    probi = [[{'role':'user','content':i}] for i in ll]
    probx = [tt.apply_chat_template(i,tokenize=False,add_generation_prompt=True) for i in probi]
    return probx

de = pd.read_csv('mathtestgen.csv')
p1 = list(de['p1'])
p2 = list(de['p2'])
p3 = list(de['p3'])
p4 = list(de['p4'])

allp = p1+p2+p3+p4

print(len(allp))

tt = AutoTokenizer.from_pretrained(ss)

allp = getchatml(allp)

# Example usage
batched_data = divide_list(allp, 4)
outputs = ray.get([run_inference.remote(data) for data in batched_data])

alans = [outputs[i].outputs[0].text for i in range(len(allp))]
alar = [outputs[i].outputs[0].finish_reason for i in range(len(allp))]

la = pd.DataFrame()
la['problems'] = allp
la['answer'] = alans
la['finish_reason'] = alar

la.to_csv('generatedtest.csv')
getting error as 2024-09-19 02:45:41,624 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 127.0.0.1:8265 
(pid=389594) INFO 09-19 02:45:43 importing.py:10] Triton not installed; certain GPU-related functions will not be available.
(pid=389594) 2024-09-19 02:45:44.135444: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
(pid=389594) To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
(run_inference pid=389594) Connecting to existing Ray cluster at address: 10.130.0.23:6379...
(run_inference pid=389594) Calling ray.init() again after it has already been called.
Traceback (most recent call last):
  File "/nfs_share/da.py", line 67, in <module>
    outputs = ray.get([run_inference.remote(data) for data in batched_data])
  File "/home/spars/.local/lib/python3.10/site-packages/ray/_private/auto_init_hook.py", line 21, in auto_init_wrapper
    return fn(*args, **kwargs)
  File "/home/spars/.local/lib/python3.10/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
  File "/home/spars/.local/lib/python3.10/site-packages/ray/_private/worker.py", line 2664, in get
    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
  File "/home/spars/.local/lib/python3.10/site-packages/ray/_private/worker.py", line 871, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(ValueError): ray::run_inference() (pid=389594, ip=10.130.0.23)
  File "/nfs_share/da.py", line 21, in run_inference
    llm = LLM(model=ss, enforce_eager=False,  max_model_len=4096, max_num_batched_tokens=32768, max_num_seqs=256, tensor_parallel_size=2, dtype='bfloat16', gpu_memory_utilization=0.96)
  File "/nfs_share/vllm/vllm/entrypoints/llm.py", line 179, in __init__
    self.llm_engine = LLMEngine.from_engine_args(
  File "/nfs_share/vllm/vllm/engine/llm_engine.py", line 553, in from_engine_args
    executor_class = cls._get_executor_cls(engine_config)
  File "/nfs_share/vllm/vllm/engine/llm_engine.py", line 499, in _get_executor_cls
    initialize_ray_cluster(engine_config.parallel_config)
  File "/nfs_share/vllm/vllm/executor/ray_utils.py", line 280, in initialize_ray_cluster
    raise ValueError(
ValueError: Current node has no TPU available. current_node_resource={'TPU-v4-32-head': 1.0, 'object_store_memory': 127737862333.0, 'fspar': 1.0, 'accelerator_type:TPU-V4': 1.0, 'node:10.130.0.23': 1.0, 'node:__internal_head__': 1.0, 'memory': 288076655207.0}. vLLM engine cannot start without TPU. Make sure you have at least 1 TPU available in a node current_node_id='c78f245f5dd463c035aada6ad9a223fa30c7acf3fa45868aeb269f1c' current_ip='10.130.0.23'.
(pid=250848, ip=10.130.0.21) INFO 09-19 02:45:43 importing.py:10] Triton not installed; certain GPU-related functions will not be available. [repeated 3x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)
(pid=250848, ip=10.130.0.21) 2024-09-19 02:45:44.244316: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. [repeated 3x across cluster]
(pid=250848, ip=10.130.0.21) To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. [repeated 3x across cluster]
(run_inference pid=250848, ip=10.130.0.21) Connecting to existing Ray cluster at address: 10.130.0.23:6379... [repeated 3x across cluster]
(run_inference pid=250848, ip=10.130.0.21) Calling ray.init() again after it has already been called. [repeated 3x across cluster]