| import subprocess |
| import os |
| import errno |
| |
| def download_file(url, local_fname=None, force_write=False): |
| # requests is not default installed |
| import requests |
| if local_fname is None: |
| local_fname = url.split('/')[-1] |
| if not force_write and os.path.exists(local_fname): |
| return local_fname |
| |
| dir_name = os.path.dirname(local_fname) |
| |
| if dir_name != "": |
| if not os.path.exists(dir_name): |
| try: # try to create the directory if it doesn't exists |
| os.makedirs(dir_name) |
| except OSError as exc: |
| if exc.errno != errno.EEXIST: |
| raise |
| |
| r = requests.get(url, stream=True) |
| assert r.status_code == 200, "failed to open %s" % url |
| with open(local_fname, 'wb') as f: |
| for chunk in r.iter_content(chunk_size=1024): |
| if chunk: # filter out keep-alive new chunks |
| f.write(chunk) |
| return local_fname |
| |
| def get_gpus(): |
| """ |
| return a list of GPUs |
| """ |
| try: |
| re = subprocess.check_output(["nvidia-smi", "-L"], universal_newlines=True) |
| except OSError: |
| return [] |
| return range(len([i for i in re.split('\n') if 'GPU' in i])) |
