Implement glvnd version option and default policy (#281)

* Implement glvnd version option and default policy

* remove unused arg options

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,14 @@ 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,
+            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],
+            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

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

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)