| 146 | |
| 147 | |
| 148 | def get_gpu_memory_usage(): |
| 149 | if sys.platform == 'win32': |
| 150 | # TODO: add windows support |
| 151 | return None |
| 152 | cuda_home = find_cuda() |
| 153 | if cuda_home is None: |
| 154 | return None |
| 155 | cuda_home = Path(cuda_home) |
| 156 | source = """ |
| 157 | #include <cuda_runtime.h> |
| 158 | #include <iostream> |
| 159 | int main(){ |
| 160 | int nDevices; |
| 161 | cudaGetDeviceCount(&nDevices); |
| 162 | size_t free_m, total_m; |
| 163 | // output json format. |
| 164 | std::cout << "["; |
| 165 | for (int i = 0; i < nDevices; i++) { |
| 166 | cudaSetDevice(i); |
| 167 | cudaMemGetInfo(&free_m, &total_m); |
| 168 | std::cout << "[" << free_m << "," << total_m << "]"; |
| 169 | if (i != nDevices - 1) |
| 170 | std::cout << "," << std::endl; |
| 171 | } |
| 172 | std::cout << "]" << std::endl; |
| 173 | return 0; |
| 174 | } |
| 175 | """ |
| 176 | with tempfile.NamedTemporaryFile('w', suffix='.cc') as f: |
| 177 | f_path = Path(f.name) |
| 178 | f.write(source) |
| 179 | f.flush() |
| 180 | try: |
| 181 | # TODO: add windows support |
| 182 | cmd = ( |
| 183 | f"g++ {f.name} -o {f_path.stem} -std=c++11" |
| 184 | f" -I{cuda_home / 'include'} -L{cuda_home / 'lib64'} -lcudart") |
| 185 | print(cmd) |
| 186 | subprocess.check_output(cmd, shell=True, cwd=f_path.parent) |
| 187 | cmd = f"./{f_path.stem}" |
| 188 | usages = subprocess.check_output( |
| 189 | cmd, shell=True, cwd=f_path.parent).decode() |
| 190 | usages = json.loads(usages) |
| 191 | return usages |
| 192 | except: |
| 193 | return None |
| 194 | return None |
| 195 | |
| 196 | |
| 197 | if __name__ == "__main__": |