diff --git a/src/rocker/nvidia_extension.py b/src/rocker/nvidia_extension.py index f6b13d4..226e860 100644 --- a/src/rocker/nvidia_extension.py +++ b/src/rocker/nvidia_extension.py @@ -28,12 +28,35 @@ from .core import RockerExtension from .em import empy_expand +GLVND_VERSION_POLICY_LATEST_LTS='latest_lts' + +NVIDIA_GLVND_VALID_VERSIONS=['16.04', '18.04','20.04', '22.04', '24.04'] def get_docker_version(): docker_version_raw = get_docker_client().version()['Version'] # Fix for version 17.09.0-ce return Version(docker_version_raw.split('-')[0]) +def glvnd_version_from_policy(image_version, policy): + # Default policy GLVND_VERSION_POLICY_LATEST_LTS + if not policy: + policy = GLVND_VERSION_POLICY_LATEST_LTS + + if policy == GLVND_VERSION_POLICY_LATEST_LTS: + if image_version in ['16.04', '16.10', '17.04', '17.10']: + return '16.04' + if image_version in ['18.04', '18.10', '19.04', '19.10']: + return '18.04' + if image_version in ['20.04', '20.10', '21.04', '21.10']: + return '20.04' + if image_version in ['22.04', '22.10', '23.04', '23.10']: + return '22.04' + # 24.04 is not available yet + # if image_version in ['24.04', '24.10', '25.04', '25.10']: + # return '24.04' + return '22.04' + return None + class X11(RockerExtension): @staticmethod def get_name(): @@ -112,6 +135,10 @@ def get_environment_subs(self, cliargs={}): 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 + nvidia_glvnd_version = cliargs.get('nvidia_glvnd_version', None) + if not nvidia_glvnd_version: + nvidia_glvnd_version = glvnd_version_from_policy(ver, cliargs.get('nvidia_glvnd_policy', None) ) + self._env_subs['nvidia_glvnd_version'] = nvidia_glvnd_version return self._env_subs @@ -141,6 +168,18 @@ def register_arguments(parser, defaults={}): const='auto', default=defaults.get(Nvidia.get_name(), None), help="Enable nvidia. Default behavior is to pick flag based on docker version.") + parser.add_argument('--nvidia-glvnd-version', + choices=NVIDIA_GLVND_VALID_VERSIONS, + # nargs='1', + # const='auto', + default=defaults.get('nvidia-glvnd-version', None), + help="Explicitly select an nvidia glvnd version") + parser.add_argument('--nvidia-glvnd-policy', + choices=[GLVND_VERSION_POLICY_LATEST_LTS], + # nargs='1', + # const='auto', + default=defaults.get('nvidia-glvnd-policy', GLVND_VERSION_POLICY_LATEST_LTS), + help="Set an nvidia glvnd version policy if version is unset") class Cuda(RockerExtension): @staticmethod diff --git a/src/rocker/templates/nvidia_preamble.Dockerfile.em b/src/rocker/templates/nvidia_preamble.Dockerfile.em index 1bab533..5b9681c 100644 --- a/src/rocker/templates/nvidia_preamble.Dockerfile.em +++ b/src/rocker/templates/nvidia_preamble.Dockerfile.em @@ -1,3 +1 @@ -# Ubuntu 16.04 with nvidia-docker2 beta opengl support -@{suffix = '16.04' if image_distro_version == '16.04' else '18.04'}@ -FROM nvidia/opengl:1.0-glvnd-devel-ubuntu@(suffix) as glvnd +FROM nvidia/opengl:1.0-glvnd-devel-ubuntu@(nvidia_glvnd_version) as glvnd diff --git a/test/test_nvidia.py b/test/test_nvidia.py index d63d979..80203f6 100644 --- a/test/test_nvidia.py +++ b/test/test_nvidia.py @@ -172,7 +172,15 @@ def test_nvidia_extension_basic(self): preamble = p.get_preamble(mock_cliargs) - self.assertIn('FROM nvidia/opengl:1.0-glvnd-devel-', preamble) + self.assertIn('FROM nvidia/opengl:1.0-glvnd-devel-ubuntu18.04', preamble) + + mock_cliargs = {'base_image': 'ubuntu:jammy'} + preamble = p.get_preamble(mock_cliargs) + self.assertIn('FROM nvidia/opengl:1.0-glvnd-devel-ubuntu22.04', preamble) + + mock_cliargs = {'base_image': 'ubuntu:jammy', 'nvidia_glvnd_version': '20.04'} + preamble = p.get_preamble(mock_cliargs) + self.assertIn('FROM nvidia/opengl:1.0-glvnd-devel-ubuntu20.04', preamble) docker_args = p.get_docker_args(mock_cliargs) #TODO(tfoote) restore with #37 self.assertIn(' -e DISPLAY -e TERM', docker_args)