diff --git a/setup.py b/setup.py index 12d6e145..dab3c455 100644 --- a/setup.py +++ b/setup.py @@ -50,6 +50,7 @@ 'name = rocker.extensions:Name', 'network = rocker.extensions:Network', 'nvidia = rocker.nvidia_extension:Nvidia', + 'cuda_dev = rocker.nvidia_extension:CudaDev', 'privileged = rocker.extensions:Privileged', 'pulse = rocker.extensions:PulseAudio', 'ssh = rocker.ssh_extension:Ssh', diff --git a/src/rocker/nvidia_extension.py b/src/rocker/nvidia_extension.py index 327ade7f..4e20b7cc 100644 --- a/src/rocker/nvidia_extension.py +++ b/src/rocker/nvidia_extension.py @@ -93,7 +93,7 @@ def get_environment_subs(self, cliargs={}): self._env_subs = {} self._env_subs['user_id'] = os.getuid() self._env_subs['username'] = getpass.getuser() - + # non static elements test every time detected_os = detect_os(cliargs['base_image'], print, nocache=cliargs.get('nocache', False)) if detected_os is None: @@ -133,4 +133,59 @@ def register_arguments(parser, defaults={}): default=defaults.get(Nvidia.get_name(), None), help="Enable nvidia") +class CudaDev(RockerExtension): + @staticmethod + def get_name(): + return 'cuda_dev' + + def __init__(self): + self._env_subs = None + self.name = CudaDev.get_name() + self.supported_distros = ['Ubuntu', 'Debian GNU/Linux'] + self.supported_versions = ['20.04', '10'] + + def get_environment_subs(self, cliargs={}): + if not self._env_subs: + self._env_subs = {} + self._env_subs['user_id'] = os.getuid() + self._env_subs['username'] = getpass.getuser() + + # non static elements test every time + detected_os = detect_os(cliargs['base_image'], print, nocache=cliargs.get('nocache', False)) + if detected_os is None: + print("WARNING unable to detect os for base image '%s', maybe the base image does not exist" % cliargs['base_image']) + sys.exit(1) + dist, ver, codename = detected_os + + self._env_subs['image_distro_id'] = dist + if self._env_subs['image_distro_id'] not in self.supported_distros: + print("WARNING distro id %s not supported by Cuda supported " % self._env_subs['image_distro_id'], self.supported_distros) + sys.exit(1) + self._env_subs['image_distro_version'] = ver + if self._env_subs['image_distro_version'] not in self.supported_versions: + print("WARNING distro %s version %s not in supported list by Nvidia supported versions" % (dist, ver), self.supported_versions) + sys.exit(1) + # TODO(tfoote) add a standard mechanism for checking preconditions and disabling plugins + + return self._env_subs + + def get_preamble(self, cliargs): + return '' + # preamble = pkgutil.get_data('rocker', 'templates/%s_preamble.Dockerfile.em' % self.name).decode('utf-8') + # return em.expand(preamble, self.get_environment_subs(cliargs)) + + def get_snippet(self, cliargs): + snippet = pkgutil.get_data('rocker', 'templates/%s_snippet.Dockerfile.em' % self.name).decode('utf-8') + return em.expand(snippet, self.get_environment_subs(cliargs)) + + def get_docker_args(self, cliargs): + return "" + # Runtime requires --nvidia option too + + @staticmethod + def register_arguments(parser, defaults={}): + parser.add_argument(name_to_argument(CudaDev.get_name()), + action='store_true', + default=defaults.get('cuda_dev', None), + help="Enable nvidia cuda-dev support") diff --git a/src/rocker/templates/cuda_dev_snippet.Dockerfile.em b/src/rocker/templates/cuda_dev_snippet.Dockerfile.em new file mode 100644 index 00000000..c6069eb6 --- /dev/null +++ b/src/rocker/templates/cuda_dev_snippet.Dockerfile.em @@ -0,0 +1,23 @@ +RUN apt-get update && apt-get install -y --no-install-recommends \ + gnupg2 curl ca-certificates && \ + apt-key del 7fa2af80 && \ + curl -L -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb && \ + dpkg -i cuda-keyring_1.0-1_all.deb && \ + sudo sed -i '/developer\.download\.nvidia\.com\/compute\/cuda\/repos/d' /etc/apt/sources.list && \ + apt-get purge --autoremove -y curl \ + && rm -rf /var/lib/apt/lists/* + +RUN apt-get update && apt-get install -y --no-install-recommends \ + cuda nvidia-cuda-dev \ + && rm -rf /var/lib/apt/lists/* + +# File conflict problem with libnvidia-ml.so.1 and libcuda.so.1 +# https://github.com/NVIDIA/nvidia-docker/issues/1551 +RUN rm -rf /usr/lib/x86_64-linux-gnu/libnv* +RUN rm -rf /usr/lib/x86_64-linux-gnu/libcuda* + +ENV PATH /usr/local/cuda/bin${PATH:+:${PATH}} +ENV LD_LIBRARY_PATH /usr/local/cuda/lib64/stubs:/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}} + +ENV NVIDIA_VISIBLE_DEVICES ${NVIDIA_VISIBLE_DEVICES:-all} +ENV NVIDIA_DRIVER_CAPABILITIES ${NVIDIA_DRIVER_CAPABILITIES:+$NVIDIA_DRIVER_CAPABILITIES,}compute,utility