diff --git a/nvcc4jupyter/path_utils.py b/nvcc4jupyter/path_utils.py new file mode 100644 index 0000000..b6cb27a --- /dev/null +++ b/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 diff --git a/nvcc4jupyter/plugin.py b/nvcc4jupyter/plugin.py index 56f32b5..612d321 100644 --- a/nvcc4jupyter/plugin.py +++ b/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 diff --git a/tests/fixtures/scripts/searchforme b/tests/fixtures/scripts/searchforme new file mode 100755 index 0000000..d698fec --- /dev/null +++ b/tests/fixtures/scripts/searchforme @@ -0,0 +1,3 @@ +#!/bin/bash + +echo "This is just used to test the path_utils.find_executable function" diff --git a/tests/test_path_utils.py b/tests/test_path_utils.py new file mode 100644 index 0000000..4969d8f --- /dev/null +++ b/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"