diff --git a/src/spikeinterface/sorters/runsorter.py b/src/spikeinterface/sorters/runsorter.py index baec6aaac3..80608f8973 100644 --- a/src/spikeinterface/sorters/runsorter.py +++ b/src/spikeinterface/sorters/runsorter.py @@ -19,7 +19,16 @@ from ..core import BaseRecording, NumpySorting, load_extractor from ..core.core_tools import check_json, is_editable_mode from .sorterlist import sorter_dict -from .utils import SpikeSortingError, has_nvidia +from .utils import ( + SpikeSortingError, + has_nvidia, + has_docker, + has_docker_python, + has_singularity, + has_spython, + has_docker_nvidia_installed, + get_nvidia_docker_dependecies, +) from .container_tools import ( find_recording_folders, path_to_unix, @@ -169,6 +178,15 @@ def run_sorter( container_image = None else: container_image = docker_image + + if not has_docker(): + raise RuntimeError( + "Docker is not installed. Install docker on this machine to run sorting with docker." + ) + + if not has_docker_python(): + raise RuntimeError("The python `docker` package must be installed. Install with `pip install docker`") + else: mode = "singularity" assert not docker_image @@ -176,6 +194,19 @@ def run_sorter( container_image = None else: container_image = singularity_image + + if not has_singularity(): + raise RuntimeError( + "Singularity is not installed. Install singularity " + "on this machine to run sorting with singularity." + ) + + if not has_spython(): + raise RuntimeError( + "The python `spython` package must be installed to " + "run singularity. Install with `pip install spython`" + ) + return run_sorter_container( container_image=container_image, mode=mode, @@ -462,6 +493,15 @@ def run_sorter_container( if gpu_capability == "nvidia-required": assert has_nvidia(), "The container requires a NVIDIA GPU capability, but it is not available" extra_kwargs["container_requires_gpu"] = True + + if platform.system() == "Linux" and not has_docker_nvidia_installed(): + warn( + f"nvidia-required but none of \n{get_nvidia_docker_dependecies()}\n were found. " + f"This may result in an error being raised during sorting. Try " + "installing `nvidia-container-toolkit`, including setting the " + "configuration steps, if running into errors." + ) + elif gpu_capability == "nvidia-optional": if has_nvidia(): extra_kwargs["container_requires_gpu"] = True diff --git a/src/spikeinterface/sorters/tests/test_runsorter_dependency_checks.py b/src/spikeinterface/sorters/tests/test_runsorter_dependency_checks.py new file mode 100644 index 0000000000..83d6ec3161 --- /dev/null +++ b/src/spikeinterface/sorters/tests/test_runsorter_dependency_checks.py @@ -0,0 +1,170 @@ +import os +import pytest +from pathlib import Path +import shutil +import platform +from spikeinterface import generate_ground_truth_recording +from spikeinterface.sorters.utils import has_spython, has_docker_python, has_docker, has_singularity +from spikeinterface.sorters import run_sorter +import subprocess +import sys +import copy + + +def _monkeypatch_return_false(): + """ + A function to monkeypatch the `has_` functions, + ensuring the always return `False` at runtime. + """ + return False + + +def _monkeypatch_return_true(): + """ + Monkeypatch for some `has_` functions to + return `True` so functions that are later in the + `runsorter` code can be checked. + """ + return True + + +class TestRunersorterDependencyChecks: + """ + This class tests whether expected dependency checks prior to sorting are run. + The run_sorter function should raise an error if: + - singularity is not installed + - spython is not installed (python package) + - docker is not installed + - docker is not installed (python package) + when running singularity / docker respectively. + + Two separate checks should be run. First, that the + relevant `has_` function (indicating if the dependency + is installed) is working. Unfortunately it is not possible to + easily test this core singularity and docker installs, so this is not done. + `uninstall_python_dependency()` allows a test to check if the + `has_spython()` and `has_docker_dependency()` return `False` as expected + when these python modules are not installed. + + Second, the `run_sorters()` function should return the appropriate error + when these functions return that the dependency is not available. This is + easier to test as these `has_` reporting functions can be + monkeypatched to return False at runtime. This is done for these 4 + dependency checks, and tests check the expected error is raised. + + Notes + ---- + `has_nvidia()` and `has_docker_nvidia_installed()` are not tested + as these are complex GPU-related dependencies which are difficult to mock. + """ + + @pytest.fixture(scope="function") + def uninstall_python_dependency(self, request): + """ + This python fixture mocks python modules not being importable + by setting the relevant `sys.modules` dict entry to `None`. + It uses `yield` so that the function can tear-down the test + (even if it failed) and replace the patched `sys.module` entry. + + This function uses an `indirect` parameterization, meaning the + `request.param` is passed to the fixture at the start of the + test function. This is used to reuse code for nearly identical + `spython` and `docker` python dependency tests. + """ + dep_name = request.param + assert dep_name in ["spython", "docker"] + + try: + if dep_name == "spython": + import spython + else: + import docker + dependency_installed = True + except: + dependency_installed = False + + if dependency_installed: + copy_import = sys.modules[dep_name] + sys.modules[dep_name] = None + yield + if dependency_installed: + sys.modules[dep_name] = copy_import + + @pytest.fixture(scope="session") + def recording(self): + """ + Make a small recording to have something to pass to the sorter. + """ + recording, _ = generate_ground_truth_recording(durations=[10]) + return recording + + @pytest.mark.skipif(platform.system() != "Linux", reason="spython install only for Linux.") + @pytest.mark.parametrize("uninstall_python_dependency", ["spython"], indirect=True) + def test_has_spython(self, recording, uninstall_python_dependency): + """ + Test the `has_spython()` function, see class docstring and + `uninstall_python_dependency()` for details. + """ + assert has_spython() is False + + @pytest.mark.parametrize("uninstall_python_dependency", ["docker"], indirect=True) + def test_has_docker_python(self, recording, uninstall_python_dependency): + """ + Test the `has_docker_python()` function, see class docstring and + `uninstall_python_dependency()` for details. + """ + assert has_docker_python() is False + + def test_no_singularity_error_raised(self, recording, monkeypatch): + """ + When running a sorting, if singularity dependencies (singularity + itself or the `spython` package`) are not installed, an error is raised. + Beacause it is hard to actually uninstall these dependencies, the + `has_` functions that let `run_sorter` know if the dependency + are installed are monkeypatched. This is done so at runtime these always + return False. Then, test the expected error is raised when the dependency + is not found. + """ + monkeypatch.setattr(f"spikeinterface.sorters.runsorter.has_singularity", _monkeypatch_return_false) + + with pytest.raises(RuntimeError) as e: + run_sorter("kilosort2_5", recording, singularity_image=True) + + assert "Singularity is not installed." in str(e) + + def test_no_spython_error_raised(self, recording, monkeypatch): + """ + See `test_no_singularity_error_raised()`. + """ + # make sure singularity test returns true as that comes first + monkeypatch.setattr(f"spikeinterface.sorters.runsorter.has_singularity", _monkeypatch_return_true) + monkeypatch.setattr(f"spikeinterface.sorters.runsorter.has_spython", _monkeypatch_return_false) + + with pytest.raises(RuntimeError) as e: + run_sorter("kilosort2_5", recording, singularity_image=True) + + assert "The python `spython` package must be installed" in str(e) + + def test_no_docker_error_raised(self, recording, monkeypatch): + """ + See `test_no_singularity_error_raised()`. + """ + monkeypatch.setattr(f"spikeinterface.sorters.runsorter.has_docker", _monkeypatch_return_false) + + with pytest.raises(RuntimeError) as e: + run_sorter("kilosort2_5", recording, docker_image=True) + + assert "Docker is not installed." in str(e) + + def test_as_no_docker_python_error_raised(self, recording, monkeypatch): + """ + See `test_no_singularity_error_raised()`. + """ + # make sure docker test returns true as that comes first + monkeypatch.setattr(f"spikeinterface.sorters.runsorter.has_docker", _monkeypatch_return_true) + monkeypatch.setattr(f"spikeinterface.sorters.runsorter.has_docker_python", _monkeypatch_return_false) + + with pytest.raises(RuntimeError) as e: + run_sorter("kilosort2_5", recording, docker_image=True) + + assert "The python `docker` package must be installed" in str(e) diff --git a/src/spikeinterface/sorters/utils/__init__.py b/src/spikeinterface/sorters/utils/__init__.py index 6cad10b211..62317be6f2 100644 --- a/src/spikeinterface/sorters/utils/__init__.py +++ b/src/spikeinterface/sorters/utils/__init__.py @@ -1,2 +1,14 @@ from .shellscript import ShellScript -from .misc import SpikeSortingError, get_git_commit, has_nvidia, get_matlab_shell_name, get_bash_path +from .misc import ( + SpikeSortingError, + get_git_commit, + has_nvidia, + get_matlab_shell_name, + get_bash_path, + has_docker, + has_docker_python, + has_singularity, + has_spython, + has_docker_nvidia_installed, + get_nvidia_docker_dependecies, +) diff --git a/src/spikeinterface/sorters/utils/misc.py b/src/spikeinterface/sorters/utils/misc.py index 0a6b4a986c..9c8c3bba89 100644 --- a/src/spikeinterface/sorters/utils/misc.py +++ b/src/spikeinterface/sorters/utils/misc.py @@ -1,6 +1,7 @@ from __future__ import annotations from pathlib import Path +import subprocess # TODO: decide best format for this from subprocess import check_output, CalledProcessError from typing import List, Union @@ -80,3 +81,74 @@ def has_nvidia(): return device_count > 0 except RuntimeError: # Failed to dlopen libcuda.so return False + + +def _run_subprocess_silently(command): + """ + Run a subprocess command without outputting to stderr or stdout. + """ + output = subprocess.run(command, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + return output + + +def has_docker(): + return _run_subprocess_silently("docker --version").returncode == 0 + + +def has_singularity(): + return ( + _run_subprocess_silently("singularity --version").returncode == 0 + or _run_subprocess_silently("apptainer --version").returncode == 0 + ) + + +def has_docker_nvidia_installed(): + """ + On Linux, nvidia has a set of container dependencies + that are required for running GPU in docker. This is a little + complex and is described in more detail in the links below. + To summarise breifly, at least one of the `get_nvidia_docker_dependecies()` + is almost certainly required to run docker with GPU. + + https://github.com/NVIDIA/nvidia-docker/issues/1268 + https://www.howtogeek.com/devops/how-to-use-an-nvidia-gpu-with-docker-containers/ + + Returns + ------- + Whether at least one of the dependencies listed in + `get_nvidia_docker_dependecies()` is installed. + """ + all_dependencies = get_nvidia_docker_dependecies() + has_dep = [] + for dep in all_dependencies: + has_dep.append(_run_subprocess_silently(f"{dep} --version").returncode == 0) + return any(has_dep) + + +def get_nvidia_docker_dependecies(): + """ + See `has_docker_nvidia_installed()` + """ + return [ + "nvidia-docker", + "nvidia-docker2", + "nvidia-container-toolkit", + ] + + +def has_docker_python(): + try: + import docker + + return True + except ImportError: + return False + + +def has_spython(): + try: + import spython + + return True + except ImportError: + return False