Merge pull request #28 from AllenCellModeling/slurm_dask_update

Slurm dask update

Author: jcass11

mti_nma/bin/all.py

@@ -107,23 +107,31 @@ def run(
                 # Create or get log dir
                 # Do not include ms
                 log_dir_name = datetime.now().isoformat().split(".")[0]
-                log_dir = Path(f"~/.dask_logs/mti_nma/{log_dir_name}").expanduser()
+                log_dir = Path(f".dask_logs/{log_dir_name}").expanduser()
                 # Log dir settings
-                log_dir.mkdir(parents=True)
+                log_dir.mkdir(parents=True, exist_ok=True)
+
+                # Configure dask config
+                dask.config.set(
+                    {
+                        "scheduler.work-stealing": False,
+                        "logging.distributed.worker": "info",
+                    }
+                )
 
                 # Create cluster
                 log.info("Creating SLURMCluster")
                 cluster = SLURMCluster(
-                    cores=1,
-                    memory="8GB",
+                    cores=4,
+                    memory="20GB",
                     queue="aics_cpu_general",
                     walltime="10:00:00",
                     local_directory=str(log_dir),
-                    log_directory=str(log_dir)
+                    log_directory=str(log_dir),
                 )
                 log.info("Created SLURMCluster")
 
-                # Scale workers
+                # Scale cluster
                 cluster.scale(60)
 
                 # Use the port from the created connector to set executor address

mti_nma/steps/dask_utils.py

@@ -10,7 +10,7 @@
 from vtk.util import numpy_support
 
 from datastep import Step, log_run_params
-from .. import dask_utils
+import aics_dask_utils
 from .nma_utils import get_vtk_verts_faces, run_nma, draw_whist, get_eigvec_mags
 from .nma_utils import color_vertices_by_magnitude
 
@@ -205,7 +205,7 @@ def run(
         path_input_mesh = avg_df["AvgShapeFilePathStl"].iloc[0]
 
         # Distribute blender mode heatmap generation
-        with dask_utils.DistributedHandler(distributed_executor_address) as handler:
+        with aics_dask_utils.DistributedHandler(distributed_executor_address) as handler:
             futures = handler.client.map(
                 color_vertices_by_magnitude,
                 [path_blender for i in range(len(mode_list))],

mti_nma/steps/nma/nma.py

@@ -11,7 +11,7 @@
 
 from datastep import Step, log_run_params
 
-from .. import dask_utils
+import aics_dask_utils
 
 ###############################################################################
 
@@ -157,7 +157,7 @@ def run(
         self.manifest = pd.DataFrame([])
 
         # Process each cell in the dataframe
-        with dask_utils.DistributedHandler(distributed_executor_address) as handler:
+        with aics_dask_utils.DistributedHandler(distributed_executor_address) as handler:
             futures = handler.client.map(
                 self._process_cell,
                 # Convert dataframe iterrows into two lists of items to iterate over

mti_nma/steps/shparam/shparam.py

@@ -12,7 +12,7 @@
 from aicsimageio import AICSImage, writers
 from datastep import Step, log_run_params
 
-from .. import dask_utils
+import aics_dask_utils
 from .singlecell_utils import crop_object, query_data_from_labkey
 
 ###############################################################################
@@ -181,7 +181,7 @@ def run(
             df = df.sample(n=nsamples)
 
             # Process each FOV in dataframe
-            with dask_utils.DistributedHandler(distributed_executor_address) as handler:
+            with aics_dask_utils.DistributedHandler(distributed_executor_address) as handler:
                 futures = handler.client.map(
                     self._process_fov,
                     # Convert dataframe iterrows into two lists of items to iterate over

mti_nma/steps/singlecell/singlecell.py

@@ -45,6 +45,7 @@
 ]
 
 requirements = [
+    "aics_dask_utils==0.1.0",
     "aicsimageio>=3.1.2",
     "aicsshparam>=0.0.7",
     "bokeh",