| 108 | |
| 109 | |
| 110 | def check_gpu(): |
| 111 | if not (sys.platform == "linux" and shutil.which("lshw")): |
| 112 | # Can't check GPU on non-Linux platforms |
| 113 | return |
| 114 | |
| 115 | # See if we can check if a GPU is present in case of later failures, |
| 116 | # but don't block on execution since this isn't critical |
| 117 | try: |
| 118 | proc = cmd( |
| 119 | ["lshw", "-json", "-C", "display"], |
| 120 | stdout=subprocess.PIPE, |
| 121 | stderr=subprocess.PIPE, |
| 122 | encoding="utf-8", |
| 123 | ) |
| 124 | stdout = proc.stdout.strip().strip(",") |
| 125 | stdout = json.loads(stdout) |
| 126 | except (subprocess.CalledProcessError, json.decoder.JSONDecodeError): |
| 127 | # Do nothing if any step failed |
| 128 | return |
| 129 | |
| 130 | if isinstance(stdout, dict): |
| 131 | # Sometimes lshw outputs a single item as a dict instead of a list of |
| 132 | # dicts, so wrap it up if necessary |
| 133 | stdout = [stdout] |
| 134 | if not isinstance(stdout, list): |
| 135 | return |
| 136 | |
| 137 | vendors = [s.get("vendor", "").lower() for s in stdout] |
| 138 | if not any("nvidia" in vendor for vendor in vendors): |
| 139 | warnings.append( |
| 140 | "nvidia GPU not found in 'lshw', maybe use --cpu flag when running 'docs' command?" |
| 141 | ) |
| 142 | |
| 143 | |
| 144 | def gen_name(s: str) -> str: |