code: def check_gpu_setup():
      from mpi4py import MPI
      import os
  
      mpi_init()
      comm = MPI.COMM_WORLD
      rank = comm.Get_rank()
      data = comm.gather(
          len(os.environ.get("CUDA_VISIBLE_DEVICES").split(","))
      )
      if rank == 0:
          assert data == [2, 2, 2, 2]
  
file_path: ray/python/ray/tests/test_mpi.py
imports: import pytest
  import ray
  import sys
  import os
  import numpy
  from ray.runtime_env import mpi_init
label_desc: Use image node_id 6 for calling a global function with example usage: check_gpu_setup() without return types