python.tool.gpu_check to load_nvidia_uvm if not already loaded

Author: drewgilliam

python/vsi/test/test_gpu_check.py

@@ -0,0 +1,165 @@
+import os
+import unittest
+from unittest.mock import patch, MagicMock
+
+import vsi.tools.gpu_check
+
+
+# dummy class to represent the loaded libcudart.so library
+class MockLibrary:
+  def __init__(self, file):
+    self._name = file
+    self.cudaGetDeviceCount = MagicMock(return_value="cudaGetDeviceCount")
+
+
+# unit tests
+class GpuCheckTest(unittest.TestCase):
+
+  def _check_logs(self, logs, logs_expected):
+    '''Check for expected logs'''
+
+    # check length
+    self.assertEqual(len(logs.output), len(logs_expected),
+                     f'Logs are of different length\n{logs.output=}')
+
+    # check that expected substring is in each log item
+    for idx, (item, expected) in enumerate(zip(logs.output, logs_expected)):
+      self.assertIn(expected, item,
+                    f'Log mismatch line {idx}\n{logs.output=}')
+
+
+  @patch('vsi.tools.gpu_check.pathlib.Path.rglob')
+  @patch('vsi.tools.gpu_check.ctypes.cdll.LoadLibrary', side_effect=MockLibrary)
+  def test_find_cudart(self, mock_LoadLibrary, mock_rglob):
+    '''Test find_cudart'''
+
+    # rglob return for 3 search_dirs
+    search_dirs = ['/foo', '/bar', '/path/to']
+    mock_rglob.side_effect = [list(), list(), ['/path/to/liubcudart.so']]
+
+    # call function & check return
+    result = vsi.tools.gpu_check.find_cudart(search_dirs)
+    self.assertEqual(result._name, '/path/to/liubcudart.so')
+
+
+  @patch('vsi.tools.gpu_check.pathlib.Path.rglob')
+  def test_find_cudart_error(self, mock_rglob):
+    '''Test find_cudart without discovering a file'''
+
+    # rglob return for 1 search_dirs
+    search_dirs = ['/foo']
+    mock_rglob.side_effect = [list()]
+
+    # raise OSError
+    with self.assertRaises(OSError):
+      _ = vsi.tools.gpu_check.find_cudart(search_dirs)
+
+
+  @patch('vsi.tools.gpu_check.ctypes.cdll.LoadLibrary', side_effect=MockLibrary)
+  def test_load_cudart_file(self, mock_LoadLibrary):
+    '''Test load_cudart from file'''
+
+    # call function & check return
+    result = vsi.tools.gpu_check.load_cudart('/path/to/file')
+    self.assertEqual(result._name, '/path/to/file')
+
+
+  @patch.dict(os.environ, {"LD_LIBRARY_PATH": ""})
+  @patch('vsi.tools.gpu_check.find_cudart')
+  @patch('vsi.tools.gpu_check.pathlib.Path.glob')
+  def test_load_cudart_dirs(self, mock_glob, mock_find_cudart):
+    '''Test load_cudart from ``/usr/local/cuda``'''
+
+    # glob result for ['/usr/local', '/usr']
+    mock_glob.side_effect = [['/usr/local/cuda'], ['/usr/cuda']]
+
+    # find_cudart result
+    def _find_cudart(search_dirs):
+      for search_dir in search_dirs:
+        if '/usr/local/cuda' in str(search_dir):
+          return MockLibrary('/usr/local/cuda/libcudart.so')
+      raise None
+
+    mock_find_cudart.side_effect = _find_cudart
+
+    # call function & check return
+    result = vsi.tools.gpu_check.load_cudart()
+    self.assertEqual(result._name, '/usr/local/cuda/libcudart.so')
+
+
+  @patch.dict(os.environ, {"LD_LIBRARY_PATH": "/foo:/bar:/path/to"})
+  @patch('vsi.tools.gpu_check.find_cudart')
+  @patch('vsi.tools.gpu_check.pathlib.Path.glob')
+  def test_load_cudart_ld(self, mock_glob, mock_find_cudart):
+    '''Test load_cudart from ``/usr/local/cuda``'''
+
+    # glob result for ['/usr/local', '/usr']
+    mock_glob.side_effect = [['/usr/local/cuda'], ['/usr/cuda']]
+
+    # find_cudart result
+    def _find_cudart(search_dirs):
+      for search_dir in search_dirs:
+        if '/path/to' in str(search_dir):
+          return MockLibrary('/path/to/libcudart.so')
+      raise None
+
+    mock_find_cudart.side_effect = _find_cudart
+
+    # call function & check return
+    result = vsi.tools.gpu_check.load_cudart()
+    self.assertEqual(result._name, '/path/to/libcudart.so')
+
+
+  @patch('vsi.tools.gpu_check.load_cudart', return_value=MockLibrary('libcudart.so'))
+  def test_load_nvidia_uvm(self, mock_load_cudart):
+    '''Test load_nvidia_uvm'''
+
+    # call function & check logs
+    with self.assertLogs(level='DEBUG') as logs:
+      vsi.tools.gpu_check.load_nvidia_uvm()
+
+    logs_expected = [
+      "found libcudart.so : libcudart.so",
+      "cudaGetDeviceCount : exit_code='cudaGetDeviceCount', gpu_count=-1",
+    ]
+    self._check_logs(logs, logs_expected)
+
+
+  @patch('vsi.tools.gpu_check.glob.glob', return_value=None)
+  def test_gpu_check_no_gpu(self, mock_glob):
+    '''Test gpu_check with nothing in /dev/nvidia[0-9]'''
+
+    with self.assertLogs(level='DEBUG') as logs:
+      vsi.tools.gpu_check.gpu_check()
+
+    logs_expected = ["Skip gpu_check : /dev/nvidia[0-9] missing"]
+    self._check_logs(logs, logs_expected)
+
+
+  @patch('vsi.tools.gpu_check.glob.glob', return_value=True)
+  @patch('vsi.tools.gpu_check.os.path.exists', return_value=True)
+  def test_gpu_check_uvm_exists(self, mock_exists, mock_glob):
+    '''Test gpu_check with existing /dev/nvidia-uvm'''
+
+    with self.assertLogs(level='DEBUG') as logs:
+      vsi.tools.gpu_check.gpu_check()
+
+    logs_expected = ["Skip gpu_check : /dev/nvidia-uvm already loaded"]
+    self._check_logs(logs, logs_expected)
+
+
+  @patch('vsi.tools.gpu_check.glob.glob', return_value=True)
+  @patch('vsi.tools.gpu_check.os.path.exists')
+  @patch('vsi.tools.gpu_check.load_nvidia_uvm', return_value=MockLibrary('libcudart.so'))
+  def test_gpu_check_uvm_success(self, mock_load_nvidia_uvm, mock_exists, mock_glob):
+    '''Test gpu_check with load_nvidia_uvm success'''
+
+    # /dev/nvidia-uvm check fails then succeeds
+    mock_exists.side_effect = [False, True]
+
+    # run test
+    with self.assertLogs(level='DEBUG') as logs:
+      vsi.tools.gpu_check.gpu_check()
+    logs_expected = ["/dev/nvidia-uvm has been successfully loaded"]
+    self._check_logs(logs, logs_expected)
+

python/vsi/tools/gpu_check.py

@@ -0,0 +1,163 @@
+import argparse
+import ctypes
+import glob
+import logging
+import os
+import pathlib
+
+logger = logging.getLogger(__name__)
+
+
+def find_cudart(search_dirs):
+  '''
+  Find, load, and return the ``libcudart.so`` library.
+
+  Paramters
+  ---------
+  search_dirs : :obj:`list`
+    List of directories to search
+  '''
+  for search_dir in search_dirs:
+    for file in pathlib.Path(search_dir).rglob('libcudart.so*'):
+      try:
+        return ctypes.cdll.LoadLibrary(file)
+      except OSError:
+        continue
+
+  raise OSError("Failed to find & load libcudart.so")
+
+
+def load_cudart(file=None):
+  '''
+  Load ``libcudart.so`` library.  If not provided as an input, search
+  standard locations for the library.
+
+  Paramters
+  ---------
+  file : :obj:`str`
+    Location of ``libcudart.so``
+  '''
+
+  # load directly
+  if file:
+    return ctypes.cdll.LoadLibrary(file)
+
+  # list of directories to search (via recursive glob)
+  # - /usr/local/cuda*
+  # - /usr/cuda*
+  # - ${LD_LIBRARY_PATH}
+  # - /usr
+  search_dirs = list()
+
+  search_dirs.extend(pathlib.Path('/usr/local').glob('cuda*'))
+  search_dirs.extend(pathlib.Path('/usr').glob('cuda*'))
+
+  LD_LIBRARY_PATH = os.getenv('LD_LIBRARY_PATH')
+  if LD_LIBRARY_PATH:
+    search_dirs.extend(LD_LIBRARY_PATH.split(os.pathsep))
+
+  search_dirs.append('/usr')
+
+  # ensure all search directories are pathlib
+  search_dirs = [pathlib.Path(d) for d in search_dirs]
+
+  # search for cudart
+  return find_cudart(search_dirs)
+
+
+def load_nvidia_uvm(file=None):
+  '''
+  A function to load the nvidia uvm device
+
+  Some (older) Linux Operating systems do not load ``/dev/nvidia-uvm`` on boot
+  to runlevel 3 (headless). This results in the ``nvidia-uvm`` module not being
+  loaded.
+
+  Unfortunately, a simple modprobe does not fix the issue, but a CUDA call on
+  the host (not in a container) will.
+
+  This scripts attempts to locate a ``libcudart.so`` library and calls the
+  ``cudaGetDeviceCount`` function, which loads the ``/dev/nvidia-uvm`` driver.
+  If it cannot locate the cuda runtime, you can give it the location as an
+  argument.
+
+  The CUDA Runtime is required on the host.
+  '''
+
+  # check file exists
+  if file:
+    file = pathlib.Path(file)
+    if not file.is_file():
+      raise OSError(f"File does not exist {file}")
+
+  # load libcudart.so
+  try:
+    cudart = load_cudart(file)
+  except OSError:
+    if file:
+      raise OSError(f"Failed to load cuda runtime from {file}")
+    else:
+      raise OSError("Failed to find & load cuda runtime. Try passing the "
+                    "full path of cuda runtime as an argument.")
+
+  # report
+  logger.debug(f"found libcudart.so : {cudart._name}")
+
+  # run cudaGetDeviceCount
+  cudart.cudaGetDeviceCount.argtypes = (ctypes.POINTER(ctypes.c_int), )
+  gpu_count = ctypes.c_int(-1)
+  exit_code = cudart.cudaGetDeviceCount(ctypes.pointer(gpu_count))
+
+  # report
+  logger.debug(f"cudaGetDeviceCount : {exit_code=}, gpu_count={gpu_count.value}")
+
+
+def gpu_check(file=None):
+  '''Try to load nvidia-uvm if not already loaded'''
+
+  # Only bother checking if there are any nvidia cards present
+  if not glob.glob('/dev/nvidia[0-9]'):
+    logger.debug('Skip gpu_check : /dev/nvidia[0-9] missing')
+    return
+
+  # is nvidia-uvm already loaded
+  if os.path.exists('/dev/nvidia-uvm'):
+    logger.debug('Skip gpu_check : /dev/nvidia-uvm already loaded')
+    return
+
+  # call load_nvidia_uvm
+  try:
+    load_nvidia_uvm(file)
+  except OSError as err:
+    logger.warning(f'load_nvidia_uvm failure : {err}')
+
+  # nvidia-uvm report
+  if os.path.exists('/dev/nvidia-uvm'):
+    logger.debug("/dev/nvidia-uvm has been successfully loaded")
+  else:
+    logger.critical("load_nvidia_uvm ran but /dev/nvidia-uvm is still not loaded")
+
+
+def main():
+  '''
+  Command line interface to :func:`gpu_check`
+  '''
+
+  # argument parser
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--file', type=pathlib.Path,
+                      default=None, required=False,
+                      help="Path to libcudart.so")
+
+  args = parser.parse_args()
+
+  # basic logging when called from command line
+  log_format = "[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s : %(message)s"
+  logging.basicConfig(level='DEBUG', format=log_format)
+
+  # run gpu_check
+  gpu_check(args.file)
+
+
+if __name__ == '__main__':
+  main()