docstr: Test jax_visible_devices works in distributed settings.
docstr_tok: ["Test","jax_visible_devices","works","in","distributed","settings","."]
code: def test_distributed_jax_visible_devices(self):
      """Test jax_visible_devices works in distributed settings."""
      if not jtu.test_device_matches(["gpu"]):
          raise unittest.SkipTest("Tests only for GPU.")
  
      port = portpicker.pick_unused_port()
      num_gpus = 4
      num_gpus_per_task = 1
      num_tasks = num_gpus // num_gpus_per_task
  
      with contextlib.ExitStack() as exit_stack:
          subprocesses = []
          for task in range(num_tasks):
              env = os.environ.copy()
              env["JAX_PORT"] = str(port)
              env["NUM_TASKS"] = str(num_tasks)
              env["TASK"] = str(task)
              visible_devices = ",".join(
                  str((task * num_gpus_per_task) + i)
                  for i in range(num_gpus_per_task)
              )
  
              if jtu.is_device_rocm():
                  program = (
                      "import jax, os; "
                      f'jax.config.update("jax_rocm_visible_devices", "{visible_devices}"); '
                      "jax.distributed.initialize("
                      "f'localhost:{os.environ[\"JAX_PORT\"]}', "
                      'int(os.environ["NUM_TASKS"]), int(os.environ["TASK"])); '
                      's = jax.pmap(lambda x: jax.lax.psum(x, "i"), axis_name="i")(jax.numpy.ones(jax.local_device_count())); '
                      "print(f'{jax.local_device_count()},{jax.device_count()},{s}', end=\"\"); "
                  )
              else:
                  program = (
                      "import jax, os; "
                      f'jax.config.update("jax_cuda_visible_devices", "{visible_devices}"); '
                      "jax.distributed.initialize("
                      "f'localhost:{os.environ[\"JAX_PORT\"]}', "
                      'int(os.environ["NUM_TASKS"]), int(os.environ["TASK"])); '
                      's = jax.pmap(lambda x: jax.lax.psum(x, "i"), axis_name="i")(jax.numpy.ones(jax.local_device_count())); '
                      "print(f'{jax.local_device_count()},{jax.device_count()},{s}', end=\"\"); "
                  )
              args = [sys.executable, "-c", program]
              proc = subprocess.Popen(
                  args,
                  env=env,
                  stdout=subprocess.PIPE,
                  stderr=subprocess.PIPE,
                  universal_newlines=True,
              )
              subprocesses.append(exit_stack.enter_context(proc))
  
          try:
              for proc in subprocesses:
                  out, err = proc.communicate()
                  self.assertEqual(
                      proc.returncode,
                      0,
                      msg=f"Process failed:\n\n{err}",
                  )
                  self.assertRegex(
                      out,
                      f"{num_gpus_per_task},{num_gpus},\\[{num_gpus}.\\]$",
                  )
          finally:
              for proc in subprocesses:
                  proc.kill()
  
imports: import contextlib
  import os
  import shutil
  import subprocess
  import sys
  import threading
  import unittest
  import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import jax
  from jax import config
  from jax._src import core
  from jax._src import distributed
  from jax._src import maps
  from jax._src import test_util
  from jax._src import util
  from jax.experimental import pjit
  import jax.numpy
label_desc: Use image node_id 2 for calling the MultiProcessGpuTest obj's underlying member method code with example usage: obj.test_distributed_jax_visible_devices() without return types