docstr: Query the information of all the GPUs on local machine
docstr_tok: ["Query","the","information","of","all","the","GPUs","on","local","machine"]
code: def new_query(debug=False, id=None) -> "GPUStatCollection":
      """Query the information of all the GPUs on local machine"""
  
      nvml.ensure_initialized()
      log = util.DebugHelper()
  
      def _decode(b: Union[str, bytes]) -> str:
          if isinstance(b, bytes):
              return b.decode("utf-8")
          assert isinstance(b, str)
          return b
  
      def get_gpu_info(handle: NVMLHandle) -> NvidiaGPUInfo:
          """Get one GPU information specified by nvml handle"""
  
          def safepcall(fn: Callable[[], Any], error_value: Any):
              # Ignore the exception from psutil when the process is gone
              # at the moment of querying. See #144.
              return util.safecall(
                  fn,
                  error_value=error_value,
                  exc_types=(
                      psutil.AccessDenied,
                      psutil.NoSuchProcess,
                      FileNotFoundError,
                  ),
              )
  
          def get_process_info(nv_process) -> ProcessInfo:
              """Get the process information of specific pid"""
              process = {}
              if (
                  nv_process.pid
                  not in GPUStatCollection.global_processes
              ):
                  GPUStatCollection.global_processes[
                      nv_process.pid
                  ] = psutil.Process(pid=nv_process.pid)
              ps_process: psutil.Process = (
                  GPUStatCollection.global_processes[nv_process.pid]
              )
  
              # TODO: ps_process is being cached, but the dict below is not.
              process["username"] = safepcall(ps_process.username, "?")
              # cmdline returns full path;
              # as in `ps -o comm`, get short cmdnames.
              _cmdline = safepcall(ps_process.cmdline, [])
              if not _cmdline:
                  # sometimes, zombie or unknown (e.g. [kworker/8:2H])
                  process["command"] = "?"
                  process["full_command"] = ["?"]
              else:
                  process["command"] = os.path.basename(_cmdline[0])
                  process["full_command"] = _cmdline
              # Bytes to MBytes
              # if drivers are not TTC this will be None.
              usedmem = (
                  nv_process.usedGpuMemory // MB
                  if nv_process.usedGpuMemory
                  else None
              )
              process["gpu_memory_usage"] = usedmem
  
              process["cpu_percent"] = safepcall(
                  ps_process.cpu_percent, 0.0
              )
              process["cpu_memory_usage"] = safepcall(
                  lambda: round(
                      (ps_process.memory_percent() / 100.0)
                      * psutil.virtual_memory().total
                  ),
                  0.0,
              )
  
              process["pid"] = nv_process.pid
              return process
  
          def safenvml(fn):
              @functools.wraps(fn)
              def _wrapped(*args, **kwargs):
                  try:
                      return fn(*args, **kwargs)
                  except N.NVMLError as e:
                      log.add_exception(fn.__name__, e)
                      return None  # Not supported
  
              return _wrapped
  
          gpu_info = NvidiaGPUInfo()
          gpu_info["index"] = N.nvmlDeviceGetIndex(handle)
  
          gpu_info["name"] = _decode(N.nvmlDeviceGetName(handle))
          gpu_info["uuid"] = _decode(N.nvmlDeviceGetUUID(handle))
  
          gpu_info["temperature.gpu"] = safenvml(
              N.nvmlDeviceGetTemperature
          )(handle, N.NVML_TEMPERATURE_GPU)
  
          gpu_info["fan.speed"] = safenvml(N.nvmlDeviceGetFanSpeed)(
              handle
          )
  
          # memory: in Bytes
          # Note that this is a compat-patched API (see gpustat.nvml)
          memory = N.nvmlDeviceGetMemoryInfo(handle)
          gpu_info["memory.used"] = int(memory.used) // MB
          gpu_info["memory.total"] = int(memory.total) // MB
  
          # GPU utilization
          utilization = safenvml(N.nvmlDeviceGetUtilizationRates)(
              handle
          )
          gpu_info["utilization.gpu"] = (
              int(utilization.gpu) if utilization is not None else None
          )
  
          utilization = safenvml(N.nvmlDeviceGetEncoderUtilization)(
              handle
          )
          gpu_info["utilization.enc"] = (
              utilization[0] if utilization is not None else None
          )
  
          utilization = safenvml(N.nvmlDeviceGetDecoderUtilization)(
              handle
          )
          gpu_info["utilization.dec"] = (
              utilization[0] if utilization is not None else None
          )
  
          # Power
          power = safenvml(N.nvmlDeviceGetPowerUsage)(handle)
          gpu_info["power.draw"] = (
              power // 1000 if power is not None else None
          )
  
          power_limit = safenvml(N.nvmlDeviceGetEnforcedPowerLimit)(
              handle
          )
          gpu_info["enforced.power.limit"] = (
              power_limit // 1000 if power_limit is not None else None
          )
  
          # Processes
          nv_comp_processes = safenvml(
              N.nvmlDeviceGetComputeRunningProcesses
          )(handle)
          nv_graphics_processes = safenvml(
              N.nvmlDeviceGetGraphicsRunningProcesses
          )(handle)
  
          if (
              nv_comp_processes is None
              and nv_graphics_processes is None
          ):
              processes = None
          else:
              processes = []
              nv_comp_processes = nv_comp_processes or []
              nv_graphics_processes = nv_graphics_processes or []
              # A single process might run in both of graphics and compute mode,
              # However we will display the process only once
              seen_pids = set()
              for nv_process in (
                  nv_comp_processes + nv_graphics_processes
              ):
                  if nv_process.pid in seen_pids:
                      continue
                  seen_pids.add(nv_process.pid)
                  try:
                      process = get_process_info(nv_process)
                      processes.append(process)
                  except psutil.NoSuchProcess:
                      # TODO: add some reminder for NVML broken context
                      # e.g. nvidia-smi reset  or  reboot the system
                      pass
                  except psutil.AccessDenied:
                      pass
                  except FileNotFoundError:
                      # Ignore the exception which probably has occured
                      # from psutil, due to a non-existent PID (see #95).
                      # The exception should have been translated, but
                      # there appears to be a bug of psutil. It is unlikely
                      # FileNotFoundError is thrown in different situations.
                      pass
  
              # TODO: Do not block if full process info is not requested
              time.sleep(0.1)
              for process in processes:
                  pid = process["pid"]
                  cache_process: psutil.Process = (
                      GPUStatCollection.global_processes[pid]
                  )
                  process["cpu_percent"] = safepcall(
                      cache_process.cpu_percent, 0
                  )
          gpu_info["processes"] = processes
  
          GPUStatCollection.clean_processes()
          return gpu_info
  
      # 1. get the list of gpu and status
      gpu_list = []
      device_count = N.nvmlDeviceGetCount()
  
      if id is None:
          gpus_to_query = range(device_count)
      elif isinstance(id, str):
          gpus_to_query = [int(i) for i in id.split(",")]
      elif isinstance(id, Sequence):
          gpus_to_query = [int(i) for i in id]
      else:
          raise TypeError(f"Unknown id: {id}")
  
      for index in gpus_to_query:
          try:
              handle: NVMLHandle = N.nvmlDeviceGetHandleByIndex(index)
              gpu_info = get_gpu_info(handle)
              gpu_stat = GPUStat(gpu_info)
          except N.NVMLError_Unknown as e:
              gpu_stat = InvalidGPU(index, "((Unknown Error))", e)
          except N.NVMLError_GpuIsLost as e:
              gpu_stat = InvalidGPU(index, "((GPU is lost))", e)
  
          if isinstance(gpu_stat, InvalidGPU):
              log.add_exception("GPU %d" % index, gpu_stat.exception)
          gpu_list.append(gpu_stat)
  
      # 2. additional info (driver version, etc).
      # TODO: check this only once, no need to call multiple times
      try:
          driver_version = _decode(N.nvmlSystemGetDriverVersion())
          check_driver_nvml_version(driver_version)
      except N.NVMLError as e:
          log.add_exception("driver_version", e)
          driver_version = None  # N/A
  
      if debug:
          log.report_summary()
  
      return GPUStatCollection(gpu_list, driver_version=driver_version)
  
imports: import functools
  from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Sequence, Union, cast
  import json
  import locale
  import os.path
  import platform
  import sys
  import time
  from datetime import datetime
  from io import StringIO
  import psutil
  from blessed import Terminal
  from gpustat import util
  from gpustat import nvml
  from gpustat.nvml import pynvml
  from gpustat.nvml import check_driver_nvml_version
label_desc: Use image node_id 3 for calling the GPUStatCollection obj's underlying member method code with example usage: obj.new_query(debug, id) and returns: GPUStatCollection, b, gpu_info, b, util, process, _wrapped, fn, None