Merge pull request #101 from LSSTDESC/dask-fixes

Fixes for more recent dask versions

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9, "3.10"]
+        python-version: [3.8, 3.9, "3.10", "3.11", "3.12"]
 
     steps:
     - name: Checkout repository

ceci/stage.py

@@ -415,6 +415,8 @@ def describe_configuration(cls):
     @classmethod
     def _describe_configuration_text(cls):
         s = []
+        if cls.config_options is None:
+            return "<This class has no configuration options>"
         for name, val in cls.config_options.items():
             if isinstance(val, StageParameter):
                 if val.required:
@@ -831,14 +833,12 @@ def start_dask(self):
 
         return is_client
 
-    @staticmethod
-    def stop_dask():
+    def stop_dask(self):
         """
         End the dask event loop
         """
-        from dask_mpi import send_close_signal
-
-        send_close_signal()
+        self.dask_client.retire_workers()
+        self.dask_client.shutdown()
 
     def split_tasks_by_rank(self, tasks):
         """Iterate through a list of items, yielding ones this process is responsible for/

pyproject.toml

@@ -66,8 +66,8 @@ viz = [
 ]
 
 dask = [
-    "dask",
-    "dask_mpi @ git+https://github.com/joezuntz/dask-mpi",
+    "dask[distributed] >= 2023.5.0",
+    "dask_mpi >= 2022.4.0",
 ]
 
 test = [
@@ -90,6 +90,6 @@ all = [
     "pytest-mock",
     "mockmpi",
     "h5py",
-    "dask[distributed]",
-    "dask_mpi @ git+https://github.com/joezuntz/dask-mpi",
+    "dask[distributed] >= 2023.5.0",
+    "dask_mpi >= 2022.4.0",
 ]

tests/test_main.py

@@ -47,7 +47,7 @@ def test_flow_chart_dot():
 def test_run_parsl():
     run1("launcher.name=parsl", "launcher.max_threads=3")
 
-
+@pytest.mark.skip(reason="CWL currently broken")
 def test_run_cwl():
     run1("launcher.name=cwl", "launcher.dir=tests/cwl") == 0
 

tests/test_stage.py

@@ -153,6 +153,12 @@ def run(self):
     assert stage_1_cmd.config.a == 6.
     assert stage_1_cmd.config.b == 'puffins are not extinct?'
 
+    class AbstractStage(PipelineStage):
+        name="AbstractStage"
+        config_options = None
+
+    AbstractStage.describe_configuration()
+
 
     # This one should not work
     class TestStage_2(PipelineStage):