code: def run_test(testmodule, gpu_tokens):
      global LAST_CODE
      with GPU_LOCK:
          if LAST_CODE != 0:
              return
          target_gpu = gpu_tokens.pop()
      env_vars = {
          "HIP_VISIBLE_DEVICES": str(target_gpu),
          "XLA_PYTHON_CLIENT_ALLOCATOR": "default",
      }
      cmd = [
          "python3",
          "-m",
          "pytest",
          "--reruns",
          "3",
          "-x",
          testmodule,
      ]
      return_code, stderr, stdout = run_shell_command(
          cmd, env_vars=env_vars
      )
      with GPU_LOCK:
          gpu_tokens.append(target_gpu)
          if LAST_CODE == 0:
              print(
                  "Running tests in module %s on GPU %d:"
                  % (testmodule, target_gpu)
              )
              print(stdout)
              print(stderr)
              LAST_CODE = return_code
      return
  
file_path: jax/build/rocm/run_single_gpu.py
imports: import argparse
  import os
  import re
  import subprocess
  import threading
  from concurrent.futures import ThreadPoolExecutor
label_desc: Use image node_id 3 for calling a global function with example usage: run_test(testmodule, gpu_tokens) without return types