Merge pull request #2273 from samuelgarcia/docker_si_install_options

run_sorter in docker/singlurity : add option for installing spikeinterface

.github/workflows/test_containers_docker.yml

@@ -5,8 +5,6 @@ on: workflow_dispatch
 
 jobs:
   test-images:
-    env:
-      SPIKEINTERFACE_DEV_PATH: ${{ github.workspace }}
     name: Test on (${{ matrix.os }})
     runs-on: ${{ matrix.os }}
     strategy:
@@ -30,10 +28,6 @@ jobs:
           pip install pytest
           pip install -e .[full]
           pip install docker
-      - name: Test that containers install the local CI version of spikeinterface
-        run: |
-          echo $SPIKEINTERFACE_DEV_PATH
-          python -c "import os; assert os.getenv('SPIKEINTERFACE_DEV_PATH') is not None"
       - name: Run test docker containers
         run: |
           pytest -vv --capture=tee-sys -rA src/spikeinterface/sorters/external/tests/test_docker_containers.py

.github/workflows/test_containers_singularity.yml

@@ -5,8 +5,6 @@ on: workflow_dispatch
 
 jobs:
   test-images:
-    env:
-      SPIKEINTERFACE_DEV_PATH: ${{ github.workspace }}
     name: Test on (${{ matrix.os }})
     runs-on: ${{ matrix.os }}
     strategy:
@@ -34,11 +32,6 @@ jobs:
           pip install pytest
           pip install -e .[full]
           pip install spython
-      - name: Test that containers install the local CI version of spikeinterface
-        run: |
-          echo $SPIKEINTERFACE_DEV_PATH
-          python -c "import os; assert os.getenv('SPIKEINTERFACE_DEV_PATH') is not None"
-          ls -l
       - name: Run test singularity containers
         run: |
           pytest -vv --capture=tee-sys -rA src/spikeinterface/sorters/external/tests/test_singularity_containers.py

.github/workflows/test_containers_singularity_gpu.yml

@@ -46,6 +46,5 @@ jobs:
       - name: Run test singularity containers with GPU
         env:
           REPO_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
-          SPIKEINTERFACE_DEV_PATH: ${{ github.workspace }}
         run: |
           pytest -vv --capture=tee-sys -rA src/spikeinterface/sorters/external/tests/test_singularity_containers_gpu.py

src/spikeinterface/__init__.py

@@ -25,3 +25,9 @@
 import spikeinterface.full as si
 
 """
+
+# This flag must be set to False for release
+# This avoids using versioning that contains ".dev0" (and this is a better choice)
+# This is mainly useful when using run_sorter in a container and spikeinterface install
+DEV_MODE = True
+# DEV_MODE = False

src/spikeinterface/core/core_tools.py

@@ -443,3 +443,13 @@ def convert_bytes_to_str(byte_value: int) -> str:
         byte_value /= 1024
         i += 1
     return f"{byte_value:.2f} {suffixes[i]}"
+
+
+def is_editable_mode() -> bool:
+    """
+    Check if spikeinterface is installed in editable mode
+    pip install -e .
+    """
+    import spikeinterface
+
+    return (Path(spikeinterface.__file__).parents[2] / "README.md").exists()

src/spikeinterface/sorters/__init__.py

@@ -1,4 +1,5 @@
 from .basesorter import BaseSorter
 from .sorterlist import *
-from .runsorter import *
+from .container_tools import ContainerClient, install_package_in_container
+from .runsorter import run_sorter, run_sorter_local, run_sorter_container, read_sorter_folder
 from .launcher import run_sorter_jobs, run_sorter_by_property

src/spikeinterface/sorters/container_tools.py

@@ -0,0 +1,275 @@
+from pathlib import Path
+import platform
+import os
+import random
+import string
+
+# TODO move this inside functions
+try:
+    HAS_DOCKER = True
+    import docker
+except ModuleNotFoundError:
+    HAS_DOCKER = False
+
+
+from spikeinterface.core.core_tools import recursive_path_modifier
+
+
+def find_recording_folders(d):
+    """Finds all recording folders 'paths' in a dict"""
+    folders_to_mount = []
+
+    def append_parent_folder(p):
+        p = Path(p)
+        folders_to_mount.append(p.resolve().absolute().parent)
+        return p
+
+    _ = recursive_path_modifier(d, append_parent_folder, target="path", copy=True)
+
+    try:  # this will fail if on different drives (Windows)
+        base_folders_to_mount = [Path(os.path.commonpath(folders_to_mount))]
+    except ValueError:
+        base_folders_to_mount = folders_to_mount
+
+    # let's not mount root if dries are /home/..., /mnt1/...
+    if len(base_folders_to_mount) == 1:
+        if len(str(base_folders_to_mount[0])) == 1:
+            base_folders_to_mount = folders_to_mount
+
+    return base_folders_to_mount
+
+
+def path_to_unix(path):
+    """Convert a Windows path to unix format"""
+    path = Path(path)
+    if platform.system() == "Windows":
+        path = Path(str(path)[str(path).find(":") + 1 :])
+    return path.as_posix()
+
+
+def windows_extractor_dict_to_unix(d):
+    d = recursive_path_modifier(d, path_to_unix, target="path", copy=True)
+    return d
+
+
+class ContainerClient:
+    """
+    Small abstraction class to run commands in:
+      * docker with "docker" python package
+      * singularity with  "spython" python package
+    """
+
+    def __init__(self, mode, container_image, volumes, py_user_base, extra_kwargs):
+        """
+        Parameters
+        ----------
+        mode: "docker" | "singularity"
+            The container mode
+        container_image: str
+            container image name and tag
+        volumes: dict
+            dict of volumes to bind
+        py_user_base: str
+            Python user base folder to set as PYTHONUSERBASE env var in Singularity mode
+            Prevents from overwriting user's packages when running pip install
+        extra_kwargs: dict
+            Extra kwargs to start container
+        """
+        assert mode in ("docker", "singularity")
+        self.mode = mode
+        self.py_user_base = py_user_base
+        container_requires_gpu = extra_kwargs.get("container_requires_gpu", None)
+
+        if mode == "docker":
+            if not HAS_DOCKER:
+                raise ModuleNotFoundError("No module named 'docker'")
+            client = docker.from_env()
+            if container_requires_gpu is not None:
+                extra_kwargs.pop("container_requires_gpu")
+                extra_kwargs["device_requests"] = [docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])]
+
+            if self._get_docker_image(container_image) is None:
+                print(f"Docker: pulling image {container_image}")
+                client.images.pull(container_image)
+
+            self.docker_container = client.containers.create(container_image, tty=True, volumes=volumes, **extra_kwargs)
+
+        elif mode == "singularity":
+            assert self.py_user_base, "py_user_base folder must be set in singularity mode"
+            from spython.main import Client
+
+            # load local image file if it exists, otherwise search dockerhub
+            sif_file = Client._get_filename(container_image)
+            singularity_image = None
+            if Path(container_image).exists():
+                singularity_image = container_image
+            elif Path(sif_file).exists():
+                singularity_image = sif_file
+            else:
+                if HAS_DOCKER:
+                    docker_image = self._get_docker_image(container_image)
+                    if docker_image and len(docker_image.tags) > 0:
+                        tag = docker_image.tags[0]
+                        print(f"Building singularity image from local docker image: {tag}")
+                        singularity_image = Client.build(f"docker-daemon://{tag}", sif_file, sudo=False)
+                if not singularity_image:
+                    print(f"Singularity: pulling image {container_image}")
+                    singularity_image = Client.pull(f"docker://{container_image}")
+
+            if not Path(singularity_image).exists():
+                raise FileNotFoundError(f"Unable to locate container image {container_image}")
+
+            # bin options
+            singularity_bind = ",".join([f'{volume_src}:{volume["bind"]}' for volume_src, volume in volumes.items()])
+            options = ["--bind", singularity_bind]
+
+            # gpu options
+            if container_requires_gpu:
+                # only nvidia at the moment
+                options += ["--nv"]
+
+            self.client_instance = Client.instance(singularity_image, start=False, options=options)
+
+    @staticmethod
+    def _get_docker_image(container_image):
+        docker_client = docker.from_env(timeout=300)
+        try:
+            docker_image = docker_client.images.get(container_image)
+        except docker.errors.ImageNotFound:
+            docker_image = None
+        return docker_image
+
+    def start(self):
+        if self.mode == "docker":
+            self.docker_container.start()
+        elif self.mode == "singularity":
+            self.client_instance.start()
+
+    def stop(self):
+        if self.mode == "docker":
+            self.docker_container.stop()
+            self.docker_container.remove(force=True)
+        elif self.mode == "singularity":
+            self.client_instance.stop()
+
+    def run_command(self, command):
+        if self.mode == "docker":
+            res = self.docker_container.exec_run(command)
+            return res.output.decode(encoding="utf-8", errors="ignore")
+        elif self.mode == "singularity":
+            from spython.main import Client
+
+            options = ["--cleanenv", "--env", f"PYTHONUSERBASE={self.py_user_base}"]
+            res = Client.execute(self.client_instance, command, options=options)
+            if isinstance(res, dict):
+                res = res["message"]
+            return res
+
+
+def install_package_in_container(
+    container_client,
+    package_name,
+    installation_mode="pypi",
+    extra=None,
+    version=None,
+    tag=None,
+    github_url=None,
+    container_folder_source=None,
+    verbose=False,
+):
+    """
+    Install a package in a container with different modes:
+
+    * pypi: pip install package_name
+    * github: pip install {github_url}/archive/{tag/version}.tar.gz#egg=package_name
+    * folder: pip install folder
+
+    Parameters
+    ----------
+    container_client: ContainerClient
+        The container client
+    package_name: str
+        The package name
+    installation_mode: str
+        The installation mode
+    extra: str
+        Extra pip install arguments, e.g. [full]
+    version: str
+        The package version to install
+    tag: str
+        The github tag to install
+    github_url: str
+        The github url to install (needed for github mode)
+    container_folder_source: str
+        The container folder source (needed for folder mode)
+    verbose: bool
+        If True, print output of pip install command
+
+    Returns
+    -------
+    res_output: str
+        The output of the pip install command
+    """
+    assert installation_mode in ("pypi", "github", "folder")
+
+    if "[" in package_name:
+        raise ValueError("Extra pip install should not be in package_name but like this extra='[full]'")
+
+    if extra is not None:
+        assert extra[0] == "[" and extra[-1] == "]", "extra should be like this: '[full]'"
+
+    if verbose:
+        print(f"Installing {package_name} with {installation_mode} in container")
+
+    if installation_mode == "pypi":
+        cmd = f"pip install --user --upgrade --no-input --no-build-isolation {package_name}"
+
+        if extra is not None:
+            cmd += f"{extra}"
+
+        if version is not None:
+            cmd += f"=={version}"
+        res_output = container_client.run_command(cmd)
+
+    elif installation_mode == "github":
+        if version is None and tag is None:
+            tag_or_version = "main"
+        elif tag is not None:
+            tag_or_version = tag
+        elif version is not None:
+            tag_or_version = version
+
+        if github_url is None:
+            github_url = "https://github.com/SpikeInterface/spikeinterface"
+
+        cmd = f"pip install --user --upgrade --no-input {github_url}/archive/{tag_or_version}.tar.gz#egg={package_name}"
+        if extra is not None:
+            cmd += f"{extra}"
+        res_output = container_client.run_command(cmd)
+
+    elif installation_mode == "folder":
+        assert tag is None
+
+        if container_client.mode == "singularity":
+            folder_copy = f"{container_client.py_user_base}/sources/"
+        else:
+            folder_copy = "/sources/"
+
+        # create a folder for source copy
+        rand_str = "".join(random.choices(string.ascii_uppercase + string.digits, k=8))
+        folder_copy += rand_str
+        cmd = f"mkdir -p {folder_copy}"
+        res_output = container_client.run_command(cmd)
+
+        cmd = f"cp -r {container_folder_source} {folder_copy}/{package_name}"
+        res_output = container_client.run_command(cmd)
+
+        cmd = f"pip install --user --no-input {folder_copy}/{package_name}"
+        if extra is not None:
+            cmd += f"{extra}"
+        res_output = container_client.run_command(cmd)
+
+    else:
+        raise ValueError(f"install_package_incontainer, wrong installation_mode={installation_mode}")
+
+    return res_output

src/spikeinterface/sorters/external/tests/test_docker_containers.py

@@ -4,6 +4,7 @@
 import pytest
 from pathlib import Path
 
+from spikeinterface.core.core_tools import is_editable_mode
 import spikeinterface.extractors as se
 import spikeinterface.sorters as ss
 
@@ -18,8 +19,7 @@
 
 def check_gh_settings():
     if ON_GITHUB:
-        si_dev_path = os.getenv("SPIKEINTERFACE_DEV_PATH")
-        assert si_dev_path is not None, "Tests on GITHUB CI must run with the SPIKEINTERFACE_DEV_PATH"
+        assert is_editable_mode(), "Tests on GITHUB CI must run with SpikeInterface in editable mode"
 
 
 def generate_run_kwargs():
@@ -29,6 +29,7 @@ def generate_run_kwargs():
     test_recording.set_channel_offsets(0)
     run_kwargs = dict(recording=test_recording, verbose=True)
     run_kwargs["docker_image"] = True
+    run_kwargs["installation_mode"] = "dev"
     return run_kwargs