Search for profiling tools executable paths when they are required

nvcc4jupyter/path_utils.py

@@ -0,0 +1,61 @@
+"""
+Helper functions relating to file paths.
+"""
+
+import os
+from glob import glob
+from typing import List, Optional
+
+CUDA_SEARCH_PATHS: List[str] = [
+    "/opt/nvidia/nsight-compute",
+    "/usr/local/cuda",
+    "/opt",
+    "/usr",
+]
+
+
+def is_executable(fpath: str) -> bool:
+    """Check if file exists and is executable"""
+    return os.path.isfile(fpath) and os.access(fpath, os.X_OK)
+
+
+def which(name: str) -> Optional[str]:
+    """Find an executable by name by searching the PATH directories"""
+    for path_dir in os.environ.get("PATH", "").split(os.pathsep):
+        exec_path = os.path.join(path_dir, name)
+        if is_executable(exec_path):
+            return exec_path
+    return None
+
+
+def find_executable(
+    name: str, search_paths: Optional[List[str]] = None
+) -> Optional[str]:
+    """
+    Find an executable, either by searching in the directories of the PATH
+    environment variable or, if that did not work, by searching recursively
+    in directories a list given as parameter.
+
+    Args:
+        name: The name of the executable to be found.
+        search_paths: If None, only executables that are available from PATH
+            will be found. Otherwise, will recursively search these
+            directories. Defaults to None.
+
+    Returns:
+        The path to the executable if it is found, and None otherwise.
+    """
+    if search_paths is None:
+        search_paths = []
+
+    which_path = which(name)
+    if which_path is not None:
+        return which_path
+
+    for search_path in search_paths:
+        search_path = os.path.abspath(search_path)
+        search_path = os.path.join(search_path, f"**/{name}")
+        for exec_path in glob(search_path, recursive=True):
+            return exec_path
+
+    return None

nvcc4jupyter/plugin.py

@@ -9,13 +9,20 @@ import shutil
 import subprocess
 import tempfile
 import uuid
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 # pylint: disable=import-error
 from IPython.core.interactiveshell import InteractiveShell
 from IPython.core.magic import Magics, cell_magic, line_magic, magics_class
 
-from . import parsers
+from .parsers import (
+    Profiler,
+    get_parser_cuda,
+    get_parser_cuda_group_delete,
+    get_parser_cuda_group_run,
+    get_parser_cuda_group_save,
+)
+from .path_utils import CUDA_SEARCH_PATHS, find_executable
 
 DEFAULT_EXEC_FNAME = "cuda_exec.out"
 SHARED_GROUP_NAME = "shared"
@@ -37,14 +44,19 @@ class NVCCPlugin(Magics):
         super().__init__(shell)
         self.shell: InteractiveShell  # type hint not provided by parent class
 
-        self.parser_cuda = parsers.get_parser_cuda()
-        self.parser_cuda_group_save = parsers.get_parser_cuda_group_save()
-        self.parser_cuda_group_delete = parsers.get_parser_cuda_group_delete()
-        self.parser_cuda_group_run = parsers.get_parser_cuda_group_run()
+        self.parser_cuda = get_parser_cuda()
+        self.parser_cuda_group_save = get_parser_cuda_group_save()
+        self.parser_cuda_group_delete = get_parser_cuda_group_delete()
+        self.parser_cuda_group_run = get_parser_cuda_group_run()
 
         self.workdir = tempfile.mkdtemp()
         print(f'Source files will be saved in "{self.workdir}".')
 
+        self.profiler_paths: Dict[Profiler, Optional[str]] = {
+            Profiler.NCU: None,
+            Profiler.NSYS: None,
+        }
+
     def _save_source(
         self, source_name: str, source_code: str, group_name: str
     ) -> None:
@@ -135,12 +147,42 @@ class NVCCPlugin(Magics):
 
         return executable_fpath
 
+    def _get_profiler_path(self, profiler: Profiler) -> str:
+        """
+        Get the path of the executable of a given profiling tool. Searches
+        the directories of the PATH environment variable and some extra
+        directories where CUDA is usually installed.
+
+        Args:
+            profiler: The profiler whose executable should be found.
+
+        Raises:
+            RuntimeError: If the profiler executable could not be found.
+
+        Returns:
+            The file path of the executable.
+        """
+        profiler_path = self.profiler_paths[profiler]
+        if profiler_path is not None:
+            return profiler_path
+
+        profiler_path = find_executable(profiler.value, CUDA_SEARCH_PATHS)
+        if profiler_path is None:
+            raise RuntimeError(
+                f'Could not find the "{profiler.value}" profiling tool.'
+                " Consider searching for where it is installed and adding its"
+                " directory to the PATH environment variable."
+            )
+
+        self.profiler_paths[profiler] = profiler_path
+        return profiler_path
+
     def _run(  # pylint: disable=too-many-arguments
         self,
         exec_fpath: str,
         timeit: bool = False,
         profile: bool = False,
-        profiler: parsers.Profiler = parsers.Profiler.NCU,
+        profiler: Profiler = Profiler.NCU,
         profiler_args: str = "",
     ) -> str:
         """
@@ -175,7 +217,8 @@ class NVCCPlugin(Magics):
         else:
             run_args = []
             if profile:
-                run_args.extend([profiler.value] + profiler_args.split())
+                profiler_path = self._get_profiler_path(profiler)
+                run_args.extend([profiler_path] + profiler_args.split())
             run_args.append(exec_fpath)
             output = subprocess.check_output(
                 run_args, stderr=subprocess.STDOUT

tests/fixtures/scripts/searchforme

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+echo "This is just used to test the path_utils.find_executable function"

tests/test_path_utils.py

@@ -0,0 +1,16 @@
+import os
+
+from nvcc4jupyter.path_utils import find_executable
+
+
+def test_which():
+    assert find_executable("ls") == "/usr/bin/ls"
+
+
+def test_find_executable(fixtures_path: str):
+    exec_path = find_executable("searchforme", [fixtures_path])
+    assert exec_path is not None
+
+    exec_dir, exec_fname = os.path.split(exec_path)
+    assert exec_fname == "searchforme"
+    assert os.path.basename(exec_dir) == "scripts"