code: def collect_testmodules():
      all_test_files = []
      return_code, stderr, stdout = run_shell_command(
          ["python3", "-m", "pytest", "--collect-only", "tests"]
      )
      if return_code != 0:
          print(stdout)
          print(stderr)
          print("Test module discovery failed.")
          exit(return_code)
      for line in stdout.split("\n"):
          match = re.match("<Module (.*)>", line)
          if match:
              test_file = match.group(1)
              all_test_files.append(test_file)
      print("---------- collected test modules ----------")
      print("Found %d test modules." % (len(all_test_files)))
      print("\n".join(all_test_files))
      print("--------------------------------------------")
      return all_test_files
  
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 2 for calling a global function with example usage: collect_testmodules() and returns: all_test_files