Merge pull request #3024 from JoeZiminski/more_docker_dependency_checks

Add more container dependency checks in `run_sorter`

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,13 +178,35 @@ 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
             if isinstance(singularity_image, bool):
                 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

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_<dependency>` functions,
+    ensuring the always return `False` at runtime.
+    """
+    return False
+
+
+def _monkeypatch_return_true():
+    """
+    Monkeypatch for some `has_<dependency>` 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_<dependency>` 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_<dependency>` 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_<dependency>` 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)

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,
+)

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