Find best n idle devices and return a string of device ids using the `nvidia-smi` command.
(n=2)
| 837 | |
| 838 | |
| 839 | def query_memory(n=2): |
| 840 | """ |
| 841 | Find best n idle devices and return a string of device ids using the `nvidia-smi` command. |
| 842 | """ |
| 843 | bash_string = "nvidia-smi --query-gpu=power.draw,temperature.gpu,memory.used --format=csv,noheader,nounits" |
| 844 | |
| 845 | try: |
| 846 | print(f"query memory with n={n}") |
| 847 | p1 = Popen(bash_string.split(), stdout=PIPE) |
| 848 | output, error = p1.communicate() |
| 849 | free_memory = [x.split(",") for x in output.decode("utf-8").split("\n")[:-1]] |
| 850 | free_memory = np.asarray(free_memory, dtype=float).T |
| 851 | free_memory[1] += free_memory[0] # combine 0/1 column measures |
| 852 | ids = np.lexsort(free_memory)[:n] |
| 853 | except (TypeError, ValueError, IndexError, OSError): |
| 854 | ids = range(n) if isinstance(n, int) else [] |
| 855 | return ",".join(f"{int(x)}" for x in ids) |
| 856 | |
| 857 | |
| 858 | def test_local_inversion(invertible_xform, to_invert, im, dict_key=None): |
no test coverage detected
searching dependent graphs…