Make core function start_dask_run a utils.

Author: domfournier

simpeg_drivers/driver.py

@@ -14,10 +14,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod, ABC
-import cProfile
-import pstats
 
-import contextlib
 from copy import deepcopy
 import sys
 from datetime import datetime, timedelta
@@ -26,7 +23,7 @@
 from time import time
 
 import numpy as np
-from dask.distributed import get_client, Client, LocalCluster, performance_report
+from dask.distributed import get_client, Client
 
 from geoapps_utils.base import Driver, Options
 from geoapps_utils.run import load_ui_json_as_dict
@@ -74,7 +71,7 @@
 from simpeg_drivers.joint.options import BaseJointOptions
 from simpeg_drivers.utils.nested import tile_locations
 from simpeg_drivers.utils.regularization import cell_neighbors, set_rotated_operators
-from simpeg_drivers.utils.utils import validate_out_group
+from simpeg_drivers.utils.utils import validate_out_group, start_dask_run
 
 mlogger = logging.getLogger("distributed")
 mlogger.setLevel(logging.WARNING)
@@ -500,47 +497,7 @@ def start_dask_run(
         :param n_workers: Number of workers to use.
         :param n_threads: Number of threads to use.
         """
-        ui_json = load_ui_json_as_dict(json_path)
-
-        n_workers = ui_json.get("n_workers", n_workers)
-        n_threads = ui_json.get("n_threads", n_threads)
-        save_report = ui_json.get("performance_report", False)
-
-        if (n_workers is not None and n_workers > 1) or n_threads is not None:
-            cluster = LocalCluster(
-                processes=True,
-                n_workers=n_workers,
-                threads_per_worker=n_threads,
-            )
-        else:
-            cluster = None
-
-        profiler = cProfile.Profile()
-        profiler.enable()
-
-        with (
-            cluster.get_client()
-            if cluster is not None
-            else contextlib.nullcontext() as context_client
-        ):
-            # Full run
-            with (
-                performance_report(filename=json_path.parent / "dask_profile.html")
-                if (save_report and isinstance(context_client, Client))
-                else contextlib.nullcontext()
-            ):
-                cls.start(json_path)
-                sys.stdout.close()
-
-        profiler.disable()
-
-        if save_report:
-            with open(
-                json_path.parent / "runtime_profile.txt", encoding="utf-8", mode="w"
-            ) as s:
-                ps = pstats.Stats(profiler, stream=s)
-                ps.sort_stats("cumulative")
-                ps.print_stats()
+        start_dask_run(cls, json_path, n_workers=n_workers, n_threads=n_threads)
 
     @property
     def workers(self) -> list[tuple[str]]:

simpeg_drivers/plate_simulation/driver.py

@@ -33,7 +33,7 @@
 from simpeg_drivers.plate_simulation.models.parametric import Plate
 from simpeg_drivers.plate_simulation.models.series import DikeSwarm, Geology
 from simpeg_drivers.plate_simulation.options import PlateSimulationOptions
-from simpeg_drivers.utils.utils import validate_out_group
+from simpeg_drivers.utils.utils import start_dask_run, validate_out_group
 
 
 logger = get_logger(__name__, propagate=False)
@@ -278,8 +278,19 @@ def replicate(
             plates.append(new)
         return plates
 
+    @classmethod
+    def start_dask_run(
+        cls, json_path: Path, n_workers: int | None = None, n_threads: int | None = None
+    ):
+        """
+        Sets Dask config settings.
+
+        :param json_path: Path to input file (.ui.json) for the application.
+        :param n_workers: Number of workers to use.
+        :param n_threads: Number of threads to use.
+        """
+        start_dask_run(cls, json_path, n_workers=n_workers, n_threads=n_threads)
 
-PlateSimulationDriver.start_dask_run = InversionDriver.start_dask_run
 
 if __name__ == "__main__":
     file = Path(sys.argv[1])

simpeg_drivers/plate_simulation/match/driver.py

@@ -38,6 +38,7 @@
 from simpeg_drivers.plate_simulation.options import ModelOptions, PlateSimulationOptions
 from simpeg_drivers.utils.utils import (
     get_default_parallelization_params,
+    start_dask_run,
     validate_out_group,
 )
 
@@ -382,8 +383,18 @@ def run_scores(self, spatial_projection, data) -> tuple[np.ndarray, np.ndarray]:
 
         return scores, centers
 
+    @classmethod
+    def start_dask_run(
+        cls, json_path: Path, n_workers: int | None = None, n_threads: int | None = None
+    ):
+        """
+        Sets Dask config settings.
 
-PlateMatchDriver.start_dask_run = BaseDriver.start_dask_run
+        :param json_path: Path to input file (.ui.json) for the application.
+        :param n_workers: Number of workers to use.
+        :param n_threads: Number of threads to use.
+        """
+        start_dask_run(cls, json_path, n_workers=n_workers, n_threads=n_threads)
 
 
 def is_up_dip(data: np.ndarray) -> bool:
@@ -514,7 +525,6 @@ def batch_files_score(
 
 if __name__ == "__main__":
     file = Path(sys.argv[1]).resolve()
+    n_w, n_t = get_default_parallelization_params(file)
 
-    n_workers, n_threads = get_default_parallelization_params(file)
-
-    PlateMatchDriver.start_dask_run(file, n_workers=n_workers, n_threads=n_threads)
+    PlateMatchDriver.start_dask_run(file, n_workers=n_w, n_threads=n_t)

simpeg_drivers/plate_simulation/sweep/driver.py

@@ -37,7 +37,7 @@
 from simpeg_drivers.plate_simulation.options import PlateSimulationOptions
 from simpeg_drivers.plate_simulation.sweep.options import SweepOptions
 from simpeg_drivers.plate_simulation.sweep.uijson import PlateSweepUIJson
-from simpeg_drivers.utils.utils import validate_out_group
+from simpeg_drivers.utils.utils import start_dask_run, validate_out_group
 
 
 logger = get_logger(name=__name__, level_name=False, propagate=False, add_name=False)
@@ -192,8 +192,18 @@ def run_trial(
         del plate_sim
         return None
 
+    @classmethod
+    def start_dask_run(
+        cls, json_path: Path, n_workers: int | None = None, n_threads: int | None = None
+    ):
+        """
+        Sets Dask config settings.
 
-PlateSweepDriver.start_dask_run = BaseDriver.start_dask_run
+        :param json_path: Path to input file (.ui.json) for the application.
+        :param n_workers: Number of workers to use.
+        :param n_threads: Number of threads to use.
+        """
+        start_dask_run(cls, json_path, n_workers=n_workers, n_threads=n_threads)
 
 
 def forms_to_values(data: dict) -> dict:

simpeg_drivers/utils/utils.py

@@ -11,13 +11,18 @@
 
 from __future__ import annotations
 
+import contextlib
+import cProfile
 import multiprocessing
+import pstats
+import sys
 from collections.abc import Sequence
 from copy import deepcopy
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import numpy as np
+from dask.distributed import Client, LocalCluster, performance_report
 from discretize import TensorMesh, TreeMesh
 from discretize.utils import mesh_utils
 from geoapps_utils.base import Options
@@ -724,3 +729,59 @@ def validate_out_group(options: Options) -> SimPEGGroup:
         out_group.metadata = None
 
     return out_group
+
+
+def start_dask_run(
+    class_type,
+    json_path: Path,
+    n_workers: int | None = None,
+    n_threads: int | None = None,
+):
+    """
+    Sets Dask config settings.
+
+    :param json_path: Path to input file (.ui.json) for the application.
+    :param n_workers: Number of workers to use.
+    :param n_threads: Number of threads to use.
+    """
+    ui_json = load_ui_json_as_dict(json_path)
+
+    n_workers = ui_json.get("n_workers", n_workers)
+    n_threads = ui_json.get("n_threads", n_threads)
+    save_report = ui_json.get("performance_report", False)
+
+    if (n_workers is not None and n_workers > 1) or n_threads is not None:
+        cluster = LocalCluster(
+            processes=True,
+            n_workers=n_workers,
+            threads_per_worker=n_threads,
+        )
+    else:
+        cluster = None
+
+    profiler = cProfile.Profile()
+    profiler.enable()
+
+    with (
+        cluster.get_client()
+        if cluster is not None
+        else contextlib.nullcontext() as context_client
+    ):
+        # Full run
+        with (
+            performance_report(filename=json_path.parent / "dask_profile.html")
+            if (save_report and isinstance(context_client, Client))
+            else contextlib.nullcontext()
+        ):
+            class_type.start(json_path)
+            sys.stdout.close()
+
+    profiler.disable()
+
+    if save_report:
+        with open(
+            json_path.parent / "runtime_profile.txt", encoding="utf-8", mode="w"
+        ) as s:
+            ps = pstats.Stats(profiler, stream=s)
+            ps.sort_stats("cumulative")
+            ps.print_stats()