├── .gitignore
├── README.md
├── setGPU.py
└── setup.py


/.gitignore:
--------------------------------------------------------------------------------
1 | __pycache__
2 | *.egg-info
3 | build
4 | dist


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # setGPU
 2 | 
 3 | A small Python library that automatically sets `CUDA_VISIBLE_DEVICES`
 4 | to the least-loaded GPU on multi-GPU systems and can be used by:
 5 | 
 6 | 1. Putting `import setGPU` before any import
 7 |    that will use a GPU like Torch, TensorFlow, or JAX.
 8 | 2. Defining an alias such as
 9 |    `alias setGPU='eval $(python3 -m setGPU)'`
10 |    and calling that to set the GPU in the shell before running
11 |    another program that uses the GPU.
12 | 
13 | # Installation
14 | 
15 | ```
16 | pip install git+https://github.com/bamos/setGPU.git
17 | ```
18 | 
19 | # Dependencies
20 | 
21 | + [Jongwook Choi's](https://wook.kr) [gpustat](https://github.com/wookayin/gpustat) library (`pip install gpustat`)
22 | 
23 | # Licensing
24 | 
25 | This code is in the public domain.
26 | 


--------------------------------------------------------------------------------
/setGPU.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import os
 4 | import sys
 5 | import gpustat
 6 | import random
 7 | 
 8 | stats = gpustat.GPUStatCollection.new_query()
 9 | ids = map(lambda gpu: int(gpu.entry['index']), stats)
10 | ratios = map(lambda gpu: float(gpu.entry['memory.used'])/float(gpu.entry['memory.total']), stats)
11 | pairs = list(zip(ids, ratios))
12 | random.shuffle(pairs)
13 | bestGPU = min(pairs, key=lambda x: x[1])[0]
14 | 
15 | print(f'Setting GPU to {bestGPU}', file=sys.stderr)
16 | os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
17 | os.environ['CUDA_VISIBLE_DEVICES'] = str(bestGPU)
18 | print(f"export CUDA_VISIBLE_DEVICES={bestGPU}")
19 | 


--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
 1 | from setuptools import find_packages, setup
 2 | 
 3 | setup(
 4 |     name='setGPU',
 5 |     version='0.0.8',
 6 |     description="A small Python library that automatically sets CUDA_VISIBLE_DEVICES to the least-loaded GPU on multi-GPU systems.",
 7 |     author='Brandon Amos',
 8 |     author_email='bamos@cs.cmu.edu',
 9 |     platforms=['any'],
10 |     license="Public Domain",
11 |     url='https://github.com/bamos/setGPU',
12 |     py_modules=['setGPU'],
13 |     install_requires=[
14 |         'gpustat'
15 |     ]
16 | )
17 | 


--------------------------------------------------------------------------------