add nvidia lib path to env before running

Author: SomeBottle

tests/unit/test_nvidia.py

@@ -9,6 +9,7 @@
 from unittest.mock import patch, Mock
 from udocker.config import Config
 from udocker.engine.nvidia import NvidiaMode
+from udocker.utils.uenv import Uenv
 
 collections.Callable = collections.abc.Callable
 
@@ -350,6 +351,29 @@ def test_13__copy_single_file(self, mock_exists, mock_statmod, mock_copy2,
         self.assertTrue(mock_copy2.called) # should copy file
         self.assertTrue(mock_chmod.called)
 
+    @patch('udocker.engine.nvidia.os.path.isdir')
+    def test_14_merge_path_env(self, mock_isdir):
+        """Test14 NvidiaMode.merge_path_env()."""
+        nvmode = NvidiaMode(self.local, self.cont_id)
+        env_path = "/opt/python/bin"
+        nvi_lib_dir = nvmode._find_cont_dir()
+        mock_isdir.return_value = True
+
+        # Case 1: LD_LIBRARY_PATH not in env
+        env = Uenv(["PATH=%s" % env_path])
+        nvmode.merge_path_env(env)
+        self.assertEqual(nvi_lib_dir, env.getenv("LD_LIBRARY_PATH"))
+        self.assertEqual(env_path, env.getenv("PATH"))
+        self.assertTrue(mock_isdir.called)
+
+        # Case 2: LD_LIBRARY_PATH in env
+        original_ld_library_path = "/usr/local/nvidia/lib:/usr/local/nvidia/lib64"
+        env = Uenv(["PATH=%s" % env_path, "LD_LIBRARY_PATH=%s" % original_ld_library_path])
+        nvmode.merge_path_env(env)
+        self.assertTrue(nvi_lib_dir in env.getenv("LD_LIBRARY_PATH"))
+        self.assertTrue(original_ld_library_path in env.getenv("LD_LIBRARY_PATH"))
+        self.assertEqual(env_path, env.getenv("PATH"))
+
 
 if __name__ == '__main__':
     main()

udocker/engine/base.py

@@ -9,6 +9,7 @@
 from udocker.genstr import is_genstr
 from udocker.msg import Msg
 from udocker.config import Config
+from udocker.engine.nvidia import NvidiaMode
 from udocker.helper.nixauth import NixAuthentication
 from udocker.helper.hostinfo import HostInfo
 from udocker.helper.osinfo import OSInfo
@@ -566,6 +567,10 @@ def _run_env_set(self):
         self.opt["env"].append("container_execmode=" +
                                self.exec_mode.get_mode())
         cont_name = self.container_names
+        # Add nvidia library path to LD_LIBRARY_PATH
+        nvidia_mode = NvidiaMode(self.localrepo, self.container_id)
+        if nvidia_mode.get_mode():
+            nvidia_mode.merge_path_env(self.opt["env"])
         # if Python 3
         if sys.version_info[0] >= 3:
             names = str(cont_name).translate(str.maketrans('', '', " '\"[]"))

udocker/engine/nvidia.py

@@ -149,7 +149,7 @@ def _find_host_dir(self):
 
     def _find_cont_dir(self):
         """Find the location of the host target directory for libraries"""
-        # Ensure these directories exist in the container
+        # Support different linux distributions
         for dst_dir in ("/usr/lib/x86_64-linux-gnu", "/usr/lib64"):
             if os.path.isdir(self.container_root + '/' + dst_dir):
                 Msg().out("Debug: Cont. location nvidia", dst_dir, l=Msg.DBG)
@@ -164,8 +164,9 @@ def _installation_exists(self, nvi_host_dir_list, nvi_cont_dir):
                 self._files_exist(nvi_cont_dir, lib_list)
             self._files_exist("/etc", Config.conf['nvi_etc_list'])
             self._files_exist("/usr/bin", Config.conf['nvi_bin_list'])
-            Msg().out("Info: Cont, has files from previous nvidia install")
+            Msg().out("Info: Cont, nvidia not set up")
         except OSError:
+            Msg().out("Info: Cont, has files from previous nvidia install")
             return True
         return False
 
@@ -209,3 +210,19 @@ def get_devices(self):
                 dev_list.append(expanded_devs)
         Msg().out("Debug: nvidia device list", dev_list, l=Msg.DBG)
         return dev_list
+
+    def merge_path_env(self, env):
+        """Add nvidia library path into LD_LIBRARY_PATH of env(Uenv)"""
+        nvi_cont_dir = self._find_cont_dir()
+        if not nvi_cont_dir:
+            return
+        ld_library_path = env.getenv("LD_LIBRARY_PATH")
+        if not ld_library_path:
+            ld_library_list = []
+        else:
+            ld_library_list = [p.strip() for p in ld_library_path.split(':')]
+        if nvi_cont_dir in ld_library_list:
+            return
+        ld_library_list.append(nvi_cont_dir)
+        env.setenv("LD_LIBRARY_PATH", ':'.join(ld_library_list))
+        Msg().out("Debug: add nvidia library path", nvi_cont_dir, "to LD_LIBRARY_PATH", l=Msg.DBG)