add global make commands, add cluster-tools typechecking (#683)

* add format/typecheck/lint/test commands to makefile

* add wkcuber typing (unclear why CI doesn't complain, got errors locally)

* load .env also for mp start method

* upgrade mypy in wkcuber

* add typechecking for cluster_tools

* update development docs

* undo dotenv loading

* remove cloudpickle from deps, format

Author: jstriebel

.github/workflows/ci.yml

@@ -85,6 +85,10 @@ jobs:
           pip install poetry
           poetry install
 
+      - name: Check typing
+        if: ${{ matrix.executors == 'multiprocessing' }}
+        run: ./typecheck.sh
+
       - name: Check formatting
         if: ${{ matrix.executors == 'multiprocessing' }}
         run: ./format.sh check
@@ -150,8 +154,7 @@ jobs:
       run: ./lint.sh
 
     - name: Check typing
-      run: |
-        ./typecheck.sh
+      run: ./typecheck.sh
 
     - name: Python tests
       env:

Makefile

@@ -1,11 +1,16 @@
 packages_by_priority := webknossos wkcuber cluster_tools docs
 packages_by_dependency := cluster_tools webknossos wkcuber docs
+code_packages := cluster_tools webknossos wkcuber
 
 define in_each_pkg_by_dependency
   for PKG in $(packages_by_dependency); do echo $$PKG; cd $$PKG; $1; cd ..; done
 endef
 
-.PHONY: list_packages_by_priority update update-internal install
+define in_each_code_pkg
+  for PKG in $(code_packages); do echo $$PKG; cd $$PKG; $1; cd ..; done
+endef
+
+.PHONY: list_packages_by_priority update update-internal install format lint typecheck flt test
 
 list_packages_by_priority:
 	@echo $(packages_by_priority)
@@ -18,3 +23,18 @@ update-internal:
 
 install:
 	$(call in_each_pkg_by_dependency, poetry install)
+
+format:
+	$(call in_each_code_pkg, ./format.sh)
+
+lint:
+	$(call in_each_code_pkg, ./lint.sh)
+
+typecheck:
+	$(call in_each_code_pkg, ./typecheck.sh)
+
+flt:
+	$(call in_each_code_pkg, ./format.sh && ./lint.sh && ./typecheck.sh)
+
+test:
+	$(call in_each_code_pkg, ./test.sh)

cluster_tools/cluster_tools/__init__.py

@@ -1,4 +1,3 @@
-import logging
 import multiprocessing
 import os
 import tempfile

cluster_tools/cluster_tools/pickling.py

@@ -1,16 +1,7 @@
 import os
+import pickle
 import sys
 
-use_cloudpickle = "USE_CLOUDPICKLE" in os.environ
-
-if use_cloudpickle:
-    import cloudpickle
-
-    pickle_strategy = cloudpickle
-else:
-    import pickle
-
-    pickle_strategy = pickle
 from .util import warn_after
 
 WARNING_TIMEOUT = 10 * 60  # seconds
@@ -45,27 +36,23 @@ def get_suitable_pickle_protocol():
 
 @warn_after("pickle.dumps", WARNING_TIMEOUT)
 def dumps(*args, **kwargs):
-    return pickle_strategy.dumps(
-        *args, protocol=get_suitable_pickle_protocol(), **kwargs
-    )
+    return pickle.dumps(*args, protocol=get_suitable_pickle_protocol(), **kwargs)
 
 
 @warn_after("pickle.dump", WARNING_TIMEOUT)
 def dump(*args, **kwargs):
-    return pickle_strategy.dump(
-        *args, protocol=get_suitable_pickle_protocol(), **kwargs
-    )
+    return pickle.dump(*args, protocol=get_suitable_pickle_protocol(), **kwargs)
 
 
 @warn_after("pickle.loads", WARNING_TIMEOUT)
 def loads(*args, **kwargs):
     assert (
         "custom_main_path" not in kwargs
     ), "loads does not implement support for the argument custom_main_path"
-    return pickle_strategy.loads(*args, **kwargs)
+    return pickle.loads(*args, **kwargs)
 
 
-class RenameUnpickler(pickle_strategy.Unpickler):
+class RenameUnpickler(pickle.Unpickler):
     def find_class(self, module, name):
         renamed_module = module
         if module == "__main__" and self.custom_main_path is not None:

cluster_tools/cluster_tools/schedulers/cluster_executor.py

@@ -7,7 +7,9 @@
 from abc import abstractmethod
 from concurrent import futures
 from functools import partial
-from typing import Union
+from typing import List, Optional, Tuple
+
+from typing_extensions import Literal
 
 from cluster_tools import pickling
 from cluster_tools.pickling import file_path_to_absolute_module
@@ -111,7 +113,7 @@ def handle_kill(self, _signum, _frame):
     @abstractmethod
     def check_for_crashed_job(
         self, job_id_with_index
-    ) -> Union["failed", "ignore", "completed"]:
+    ) -> Literal["failed", "ignore", "completed"]:
         pass
 
     def _start(self, workerid, job_count=None, job_name=None):
@@ -137,7 +139,13 @@ def _start(self, workerid, job_count=None, job_name=None):
         return jobids_futures, job_index_ranges
 
     @abstractmethod
-    def inner_submit(self, *args, **kwargs):
+    def inner_submit(
+        self,
+        cmdline: str,
+        job_name: Optional[str] = None,
+        additional_setup_lines: Optional[List[str]] = None,
+        job_count: Optional[int] = None,
+    ) -> Tuple[List["futures.Future[str]"], List[Tuple[int, int]]]:
         pass
 
     def _cleanup(self, jobid):

cluster_tools/cluster_tools/schedulers/kube.py

@@ -4,11 +4,12 @@
 import re
 import sys
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import List, Optional, Tuple
 from uuid import uuid4
 
 import kubernetes
 import kubernetes.client.models as kubernetes_models
+from typing_extensions import Literal
 
 from .cluster_executor import ClusterExecutor
 
@@ -19,9 +20,9 @@ def volume_name_from_path(path: Path) -> str:
 
 def deduplicate_mounts(mounts: List[Path]) -> List[Path]:
     output = []
-    mounts = set(mounts)
-    for mount in mounts:
-        if not any(m in mount.parents for m in mounts):
+    unique_mounts = set(mounts)
+    for mount in unique_mounts:
+        if not any(m in mount.parents for m in unique_mounts):
             output.append(mount)
     return output
 
@@ -73,7 +74,7 @@ def get_current_job_id() -> Optional[str]:
         return os.environ.get("JOB_ID", None)
 
     @classmethod
-    def get_job_id_string(cls) -> str:
+    def get_job_id_string(cls) -> Optional[str]:
         job_id = cls.get_current_job_id()
         job_index = cls.get_job_array_index()
         if job_index is None:
@@ -104,21 +105,21 @@ def inner_submit(
         self,
         cmdline: str,
         job_name: Optional[str] = None,
+        additional_setup_lines: Optional[List[str]] = None,
         job_count: Optional[int] = None,
-        **_,
-    ):
+    ) -> Tuple[List["concurrent.futures.Future[str]"], List[Tuple[int, int]]]:
         """Starts a Kubernetes pod that runs the specified shell command line."""
 
         kubernetes_client = KubernetesClient()
         self.ensure_kubernetes_namespace()
         job_id = str(uuid4())
 
-        job_id_future = concurrent.futures.Future()
+        job_id_future: "concurrent.futures.Future[str]" = concurrent.futures.Future()
         job_id_future.set_result(job_id)
         job_id_futures = [job_id_future]
 
         is_array_job = job_count is not None
-        number_of_subjobs = job_count if is_array_job else 1
+        number_of_subjobs = job_count if job_count is not None else 1
         ranges = [(0, number_of_subjobs)]
 
         requested_resources = {
@@ -232,12 +233,12 @@ def inner_submit(
 
     def check_for_crashed_job(
         self, job_id_with_index: str
-    ) -> Union["failed", "ignore", "completed"]:
+    ) -> Literal["failed", "ignore", "completed"]:
         kubernetes_client = KubernetesClient()
         [job_id, job_index] = (
             job_id_with_index.split("_")
             if "_" in job_id_with_index
-            else [job_id_with_index, 0]
+            else [job_id_with_index, "0"]
         )
         resp = kubernetes_client.core.list_namespaced_pod(
             namespace=self.job_resources["namespace"],

cluster_tools/cluster_tools/schedulers/pbs.py

@@ -4,14 +4,16 @@
 import os
 import re
 from concurrent import futures
-from typing import Union
+from typing import Dict, List, Optional, Tuple
+
+from typing_extensions import Literal
 
 from cluster_tools.util import call, chcall, random_string
 
 from .cluster_executor import ClusterExecutor
 
 # qstat vs. checkjob
-PBS_STATES = {
+PBS_STATES: Dict[str, List[str]] = {
     "Failure": [],
     "Success": [
         "C",  # Completed
@@ -71,8 +73,12 @@ def submit_text(self, job):
         return int(jobid)
 
     def inner_submit(
-        self, cmdline, job_name=None, additional_setup_lines=None, job_count=None
-    ):
+        self,
+        cmdline: str,
+        job_name: Optional[str] = None,
+        additional_setup_lines: Optional[List[str]] = None,
+        job_count: Optional[int] = None,
+    ) -> Tuple[List["futures.Future[str]"], List[Tuple[int, int]]]:
         """Starts a PBS job that runs the specified shell command line."""
         if additional_setup_lines is None:
             additional_setup_lines = []
@@ -117,14 +123,14 @@ def inner_submit(
         ]
 
         job_id = self.submit_text("\n".join(script_lines))
-        job_id_future = futures.Future()
+        job_id_future: "futures.Future[str]" = futures.Future()
         job_id_future.set_result(job_id)
 
         return [job_id_future], [(0, job_count or 1)]
 
     def check_for_crashed_job(
         self, job_id_with_index
-    ) -> Union["failed", "ignore", "completed"]:
+    ) -> Literal["failed", "ignore", "completed"]:
         if len(str(job_id_with_index).split("_")) >= 2:
             a, b = job_id_with_index.split("_")
             job_id_with_index = f"{a}[{b}]"

cluster_tools/cluster_tools/schedulers/slurm.py

@@ -7,7 +7,9 @@
 import sys
 import threading
 from functools import lru_cache
-from typing import Union
+from typing import List, Optional, Tuple
+
+from typing_extensions import Literal
 
 from cluster_tools.util import call, chcall, random_string
 
@@ -190,8 +192,12 @@ def cleanup_submit_threads(self):
         ]
 
     def inner_submit(
-        self, cmdline, job_name=None, additional_setup_lines=None, job_count=None
-    ):
+        self,
+        cmdline: str,
+        job_name: Optional[str] = None,
+        additional_setup_lines: Optional[List[str]] = None,
+        job_count: Optional[int] = None,
+    ) -> Tuple[List["concurrent.futures.Future[str]"], List[Tuple[int, int]]]:
         """Starts a Slurm job that runs the specified shell command line."""
         if additional_setup_lines is None:
             additional_setup_lines = []
@@ -213,7 +219,7 @@ def inner_submit(
         batch_size = max(min(max_array_size, max_submit_jobs), 1)
 
         scripts = []
-        job_id_futures = []
+        job_id_futures: List["concurrent.futures.Future[str]"] = []
         ranges = []
         number_of_jobs = job_count if job_count is not None else 1
         for job_index_start in range(0, number_of_jobs, batch_size):
@@ -256,7 +262,7 @@ def inner_submit(
 
     def check_for_crashed_job(
         self, job_id_with_index
-    ) -> Union["failed", "ignore", "completed"]:
+    ) -> Literal["failed", "ignore", "completed"]:
 
         job_states = []
 

cluster_tools/cluster_tools/tailf.py

@@ -5,12 +5,15 @@
 import os
 import sys
 import time
+from typing import Any, Callable
 
 
 class Tail(object):
     """Represents a tail command."""
 
-    def __init__(self, tailed_file, callback=sys.stdout.write):
+    def __init__(
+        self, tailed_file: str, callback: Callable[[str], Any] = sys.stdout.write
+    ) -> None:
         """Initiate a Tail instance.
         Check for file validity, assigns callback function to standard out.
 
@@ -21,7 +24,7 @@ def __init__(self, tailed_file, callback=sys.stdout.write):
         self.callback = callback
         self.is_cancelled = False
 
-    def follow(self, seconds=1):
+    def follow(self, seconds: int = 1) -> None:
         """Do a tail follow. If a callback function is registered it is called with every new line.
         Else printed to standard out.
 
@@ -44,14 +47,14 @@ def follow(self, seconds=1):
                 else:
                     self.callback(line)
 
-    def cancel(self):
+    def cancel(self) -> None:
         self.is_cancelled = True
 
-    def register_callback(self, func):
+    def register_callback(self, func: Callable[[str], Any]) -> None:
         """Overrides default callback function to provided function."""
         self.callback = func
 
-    def check_file_validity(self, file_):
+    def check_file_validity(self, file_: str) -> None:
         """Check whether the a given file exists, readable and is a file"""
         if not os.access(file_, os.F_OK):
             raise TailError("File '%s' does not exist" % (file_))