Reduce default wk-libs dependencies (#1024)

* create an optional dependency group for wk-libs examples * remove nifti script to avoid sklearn dependency * make example deps optional * install dev deps in CI * extract cluster_tool extra dependencies * lazily load kubernetes * fix CI? * lazy import distributed for dask * updated readmes * installed python pacakges into virtual env in Docker * updated dockerfile * Update ci.yml * Update Dockerfile * fixes for kubernetes * dockerfile refactor * no venv for poetry * Update cluster_tools/Changelog.md Co-authored-by: Norman Rzepka <[email protected]> * Update cluster_tools/cluster_tools/executors/dask.py Co-authored-by: Norman Rzepka <[email protected]> * Update cluster_tools/cluster_tools/schedulers/kube.py Co-authored-by: Norman Rzepka <[email protected]> * Update webknossos/Changelog.md Co-authored-by: Norman Rzepka <[email protected]> --------- Co-authored-by: Norman Rzepka <[email protected]> Co-authored-by: Norman Rzepka <[email protected]>
scalableminds · Apr 18, 2024 · 096cd4b · 096cd4b
1 parent b62883e
commit 096cd4b
Show file tree

Hide file tree

Showing 19 changed files with 547 additions and 1,032 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -84,15 +84,22 @@ jobs:
           ./kind create cluster --config=tests/cluster-config.yaml
           ./kind export kubeconfig
 
+          cp ../requirements.txt .
           docker build -f tests/Dockerfile -t scalableminds/cluster-tools:latest .
           ./kind load docker-image scalableminds/cluster-tools:latest
 
       - name: Install dependencies (without docker)
-        if: ${{ matrix.executors != 'slurm' }}
+        if: ${{ matrix.executors == 'multiprocessing' }}
         run: |
           pip install -r ../requirements.txt
           poetry install
 
+      - name: Install dependencies (without docker)
+        if: ${{ matrix.executors == 'kubernetes' || matrix.executors == 'dask' }}
+        run: |
+          pip install -r ../requirements.txt
+          poetry install --all-extras
+
       - name: Check typing
         if: ${{ matrix.executors == 'multiprocessing' && matrix.python-version == '3.11' }}
         run: ./typecheck.sh
@@ -163,7 +170,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install -r ../requirements.txt
-          poetry install --extras all
+          poetry install --extras all --with examples --with dev
 
       - name: Check formatting
         if: ${{ matrix.group == 1 && matrix.python-version == '3.11' }}

diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
@@ -29,7 +29,7 @@ jobs:
     - name: Install dependencies
       run: |
         pip install poetry
-        poetry install --extras all
+        poetry install --extras all --with examples --with dev
 
     - name: Check if git is dirty
       run: |

diff --git a/cluster_tools/Changelog.md b/cluster_tools/Changelog.md
@@ -14,6 +14,7 @@ For upgrade instructions, please check the respective *Breaking Changes* section
 ### Added
 
 ### Changed
+- Moved the dependencies for Kubernetes and Dask support into optional extras that need to be installed separately. Either with `poetry install --all-extras or pip install -e "all"`. [#1024](https://github.com/scalableminds/webknossos-libs/pull/1024)
 
 ### Fixed
 

diff --git a/cluster_tools/README.md b/cluster_tools/README.md
@@ -4,7 +4,7 @@
 [![Code Style](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://docs.astral.sh/ruff/)
 
 
-This package provides python `Executor` classes for distributing tasks on a slurm cluster or via multi processing.
+This package provides python `Executor` classes for distributing tasks on a Slurm cluster, Kubernetes, Dask or via multi processing.
 
 ## Example
 
@@ -21,6 +21,23 @@ if __name__ == '__main__':
     assert result == [4, 9, 16]
 ```
 
+## Installation
+The `cluster_tools` package requires at least Python 3.8.
+
+You can install it from [pypi](https://pypi.org/project/cluster_tools/), e.g. via pip:
+
+```bash
+pip install cluster_tools
+```
+
+By default only the dependencies for running jobs on Slurm and via multiprocessing are installed. 
+For Kubernetes and Dask run:
+
+```bash
+pip install cluster_tools[kubernetes]
+pip install cluster_tools[dask]
+```
+
 ## Configuration
 
 ### Slurm
@@ -63,6 +80,8 @@ docker exec -it slurmctld bash
 docker exec -it c1 bash
 ```
 
+Make sure to install all extra dependencies, such as Kubernetes, with `poetry install --all-extras`.
+
 Tests can be executed with `cd tests && poetry run pytest -s tests.py` after entering the container.
 Linting can be run with `./lint.sh`.
 Code formatting (black) can be run with `./format.sh`.

diff --git a/cluster_tools/cluster_tools/executors/dask.py b/cluster_tools/cluster_tools/executors/dask.py
@@ -1,3 +1,4 @@
+import logging
 import os
 import re
 import signal
@@ -28,6 +29,9 @@
 if TYPE_CHECKING:
     from distributed import Client
 
+logger = logging.getLogger()
+
+
 _T = TypeVar("_T")
 _P = ParamSpec("_P")
 _S = TypeVar("_S")
@@ -110,6 +114,14 @@ class DaskExecutor(futures.Executor):
     def __init__(
         self, client: "Client", job_resources: Optional[Dict[str, Any]] = None
     ) -> None:
+        try:
+            import distributed  # noqa: F401 unused import
+        except ModuleNotFoundError:
+            logger.error(
+                'The distributed Python package for Dask is not installed. cluster_tools does not install this dependency be default. Run `pip install cluster_tools[dask]` or `poetry install --extras "dask"` to install Dask support.'
+            )
+            exit()
+
         self.client = client
         self.pending_futures = set()
         self.job_resources = job_resources

diff --git a/cluster_tools/cluster_tools/schedulers/kube.py b/cluster_tools/cluster_tools/schedulers/kube.py
@@ -1,5 +1,6 @@
 """Abstracts access to a Kubernetes cluster via its Python library."""
 
+import logging
 import os
 import re
 import sys
@@ -8,11 +9,10 @@
 from typing import Any, Dict, List, Literal, Optional, Tuple
 from uuid import uuid4
 
-import kubernetes
-import kubernetes.client.models as kubernetes_models
-
 from cluster_tools.schedulers.cluster_executor import ClusterExecutor
 
+logger = logging.getLogger()
+
 
 def _volume_name_from_path(path: Path) -> str:
     return f"{(hash(str(path)) & sys.maxsize):016x}"
@@ -30,6 +30,8 @@ def _deduplicate_mounts(mounts: List[Path]) -> List[Path]:
 
 class KubernetesClient:
     def __init__(self) -> None:
+        import kubernetes
+
         kubernetes.config.load_kube_config()
         self.core = kubernetes.client.api.core_v1_api.CoreV1Api()
         self.batch = kubernetes.client.api.batch_v1_api.BatchV1Api()
@@ -48,6 +50,14 @@ def __init__(
         additional_setup_lines: Optional[List[str]] = None,
         **kwargs: Any,
     ):
+        try:
+            import kubernetes  # noqa: F401 unused import
+        except ModuleNotFoundError:
+            logger.error(
+                'The Kubernetes Python package is not installed. cluster_tools does not install this dependency be default. Run `pip install cluster_tools[kubernetes]` or `poetry install --extras "kubernetes"` to install Kubernetes support.'
+            )
+            exit()
+
         super().__init__(
             debug=debug,
             keep_logs=keep_logs,
@@ -117,6 +127,9 @@ def inner_handle_kill(
         )
 
     def ensure_kubernetes_namespace(self) -> None:
+        import kubernetes
+        import kubernetes.client.models as kubernetes_models
+
         kubernetes_client = KubernetesClient()
         try:
             kubernetes_client.core.read_namespace(self.job_resources["namespace"])
@@ -145,6 +158,8 @@ def inner_submit(
     ) -> Tuple[List["Future[str]"], List[Tuple[int, int]]]:
         """Starts a Kubernetes pod that runs the specified shell command line."""
 
+        import kubernetes.client.models as kubernetes_models
+
         kubernetes_client = KubernetesClient()
         self.ensure_kubernetes_namespace()
         job_id = str(uuid4())

diff --git a/cluster_tools/dockered-slurm/docker-compose.yml b/cluster_tools/dockered-slurm/docker-compose.yml
@@ -22,8 +22,8 @@ services:
     volumes:
       - etc_munge:/etc/munge
       - etc_slurm:/etc/slurm
-      - ${PWD}/slurm.conf:/etc/slurm/slurm.conf
-      - ${PWD}/cgroup.conf:/etc/slurm/cgroup.conf
+      - ./slurm.conf:/etc/slurm/slurm.conf
+      - ./cgroup.conf:/etc/slurm/cgroup.conf
       - var_log_slurm:/var/log/slurm
     expose:
       - "6819"
@@ -40,8 +40,8 @@ services:
     volumes:
       - etc_munge:/etc/munge
       - etc_slurm:/etc/slurm
-      - ${PWD}/slurm.conf:/etc/slurm/slurm.conf
-      - ${PWD}/cgroup.conf:/etc/slurm/cgroup.conf
+      - ./slurm.conf:/etc/slurm/slurm.conf
+      - ./cgroup.conf:/etc/slurm/cgroup.conf
       - ./slurm_jobdir:/data
       - ..:/cluster_tools
       - var_log_slurm:/var/log/slurm
@@ -59,8 +59,8 @@ services:
     volumes:
       - etc_munge:/etc/munge
       - etc_slurm:/etc/slurm
-      - ${PWD}/slurm.conf:/etc/slurm/slurm.conf
-      - ${PWD}/cgroup.conf:/etc/slurm/cgroup.conf
+      - ./slurm.conf:/etc/slurm/slurm.conf
+      - ./cgroup.conf:/etc/slurm/cgroup.conf
       - ./slurm_jobdir:/data
       - ..:/cluster_tools
       - var_log_slurm:/var/log/slurm
@@ -77,8 +77,8 @@ services:
     volumes:
       - etc_munge:/etc/munge
       - etc_slurm:/etc/slurm
-      - ${PWD}/slurm.conf:/etc/slurm/slurm.conf
-      - ${PWD}/cgroup.conf:/etc/slurm/cgroup.conf
+      - ./slurm.conf:/etc/slurm/slurm.conf
+      - ./cgroup.conf:/etc/slurm/cgroup.conf
       - ./slurm_jobdir:/data
       - ..:/cluster_tools
       - var_log_slurm:/var/log/slurm