Skip to main content

A lightweight scheduler reading nvidia-smi and updating torch environment variables to run on the recommended GPU.

Project description

pytorch_run_on_recommended_gpu

A lightweight script that interactively updates CUDA_VISIBLE_DEVICES for pytorch

Install

pip install pytorch_run_on_recommended_cuda

Usage from CLI

Perform a dry run

pytorch_run_on_recommended_cuda

Run a script and select a GPU manually

pytorch_run_on_recommended_cuda <path_to_script>

Run a script from the next available GPU

pytorch_run_on_recommended_cuda --select * <path_to_script>

Run a script from the next two available GPUs

pytorch_run_on_recommended_cuda --select ** <path_to_script>

Run a script from GPU ids 6 and 7

pytorch_run_on_recommended_cuda --select 6 7 <path_to_script>

Usage from .py file

import os
from pytorch_run_on_recommended_gpu.run_on_recommended_gpu import get_cuda_environ_vars as get_vars

os.environ.update(get_vars('*')
print(get_vars('*')))
import torch # Import torch after you have updated the vars.

How it looks like

### Recommended gpus on this machine (descending order) ###
  ID  Card name                Util    Mem free  Cuda              User(s)
----  --------------------  -------  ----------  ----------------  ---------
   4  Tesla V100-SXM3-32GB      0 %   31889 MiB  11.8 (470.82.01)
   3  Tesla V100-SXM3-32GB      0 %   31887 MiB  11.8 (470.82.01)
   5  Tesla V100-SXM3-32GB      0 %   31737 MiB  11.8 (470.82.01)
   2  Tesla V100-SXM3-32GB      0 %   31341 MiB  11.8 (470.82.01)
   0  Tesla V100-SXM3-32GB      0 %   31263 MiB  11.8 (470.82.01)
   1  Tesla V100-SXM3-32GB      0 %   31038 MiB  11.8 (470.82.01)
  11  Tesla V100-SXM3-32GB      0 %   23012 MiB  11.8 (470.82.01)
   7  Tesla V100-SXM3-32GB      0 %   15481 MiB  11.8 (470.82.01)
  10  Tesla V100-SXM3-32GB     21 %    1025 MiB  11.8 (470.82.01)
   6  Tesla V100-SXM3-32GB     50 %   29296 MiB  11.8 (470.82.01)
   8  Tesla V100-SXM3-32GB     50 %   28988 MiB  11.8 (470.82.01)
   9  Tesla V100-SXM3-32GB     51 %   28988 MiB  11.8 (470.82.01)
  15  Tesla V100-SXM3-32GB   ! 99 %   22636 MiB  11.8 (470.82.01)
  13  Tesla V100-SXM3-32GB   ! 99 %   21441 MiB  11.8 (470.82.01)
  14  Tesla V100-SXM3-32GB  ! 100 %   22610 MiB  11.8 (470.82.01)
  12  Tesla V100-SXM3-32GB  ! 100 %   22141 MiB  11.8 (470.82.01)

Which GPUs shall be used? Give stars or ids. Input=* [ENTER]

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pytorch_run_on_recommended_gpu-1.0.5.tar.gz (19.8 kB view details)

Uploaded Source

Built Distribution

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page