cuda-dev

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',

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")
 

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