Run in parallel via CLI (#273)

- Switch from `tqdm` to `rich` for progress tracking
- Set up `pytest-xdist` to distribute tests over multiple CPU cores
- Reworked CLI to parallelize morphing over specified number of workers when there is more than one job and more than one worker requested. Passing `0` will use all available cores.
- Updated docs and test suite for new additions
- Set matplotlib backend to Agg in actions to hopefully fix tkinter issue on Windows runners

Author: stefmolin

.github/workflows/ci.yml

@@ -46,6 +46,9 @@ jobs:
             os: [macos-latest, ubuntu-latest, windows-latest]
             python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
 
+    env:
+      MPLBACKEND: Agg  # non-interactive backend for matplotlib
+
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
 

.github/workflows/generate-morphs.yml

@@ -64,10 +64,10 @@ jobs:
           echo "Detected changes to dataset(s): $DATASET_ALL_CHANGED_FILES"
           DATASET_ARGS=$(python bin/ci.py $DATASET_ALL_CHANGED_FILES)
           echo "Generating morphs for the following datasets: $DATASET_ARGS"
-          parallel -j0 data-morph \
+          data-morph \
                --start-shape $DATASET_ARGS \
-               --target-shape {} \
-               ::: bullseye heart rectangle star slant_up
+               --target-shape bullseye heart rectangle star slant_up \
+               --workers 0
 
       # If shapes are added or modified in this PR
       - name: Generate morphs from new or changed shapes
@@ -78,20 +78,20 @@ jobs:
           echo "Detected changes to shape(s): $SHAPE_ALL_CHANGED_FILES"
           SHAPE_ARGS=$(python bin/ci.py $SHAPE_ALL_CHANGED_FILES)
           echo "Generating morphs for the following shapes: $SHAPE_ARGS"
-          parallel -j0 data-morph \
+          data-morph \
                --start-shape music \
-               --target-shape {} \
-               ::: $SHAPE_ARGS
+               --target-shape $SHAPE_ARGS \
+               --workers 0
 
       # For core code changes, we want to do a couple morphs to see if they still look ok
       # Only need to run if neither of the previous two morphs ran
       - name: Morph shapes with core code changes
         if: steps.changed-files-yaml.outputs.dataset_any_changed != 'true' && steps.changed-files-yaml.outputs.shape_any_changed != 'true'
         run: |
-          parallel -j0 data-morph \
+          data-morph \
                --start-shape music \
-               --target-shape {} \
-               ::: bullseye heart rectangle star slant_up
+               --target-shape bullseye heart star \
+               --workers 0
 
       - uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08  # v4.6.0
         with:

docs/api.rst

@@ -13,6 +13,7 @@ API
    data_morph.data
    data_morph.morpher
    data_morph.plotting
+   data_morph.progress
    data_morph.shapes
 
 ----

docs/cli.rst

@@ -19,11 +19,12 @@ Examples
 
     $ data-morph --start-shape panda --target-shape star
 
-2. Morph the panda shape into all available target shapes:
+2. Morph the panda shape into all available target shapes distributing the work
+   to as many worker processes as possible:
 
    .. code-block:: console
 
-    $ data-morph --start-shape panda --target-shape all
+    $ data-morph --start-shape panda --target-shape all --workers 0
 
 3. Morph the cat, dog, and panda shapes into the circle and slant_down shapes:
 

docs/conf.py

@@ -59,12 +59,13 @@
 # -- intersphinx -------------------------------------------------------------
 
 intersphinx_mapping = {
-    'matplotlib': ('https://matplotlib.org/stable/', None),
-    'numpy': ('https://numpy.org/doc/stable/', None),
-    'pandas': ('https://pandas.pydata.org/pandas-docs/stable/', None),
-    'Pillow': ('https://pillow.readthedocs.io/en/stable/', None),
-    'pytest': ('https://pytest.org/en/stable/', None),
-    'python': ('https://docs.python.org/3/', None),
+    'matplotlib': ('https://matplotlib.org/stable', None),
+    'numpy': ('https://numpy.org/doc/stable', None),
+    'pandas': ('https://pandas.pydata.org/pandas-docs/stable', None),
+    'Pillow': ('https://pillow.readthedocs.io/en/stable', None),
+    'pytest': ('https://pytest.org/en/stable', None),
+    'python': ('https://docs.python.org/3', None),
+    'rich': ('https://rich.readthedocs.io/en/stable', None),
 }
 
 

pyproject.toml

@@ -48,14 +48,15 @@ dependencies = [
   "matplotlib>=3.3",
   "numpy>=1.20",
   "pandas>=1.2",
-  "tqdm>=4.64.1",
+  "rich>=13.9.4",
 ]
 optional-dependencies.dev = [
   "pre-commit",
   "pytest",
   "pytest-cov",
   "pytest-mock",
   "pytest-randomly",
+  "pytest-xdist",
 ]
 optional-dependencies.docs = [
   "pydata-sphinx-theme>=0.15.3",
@@ -140,6 +141,7 @@ addopts = [
   "-ra",
   "-l",
   "-v",
+  "-n=auto",                   # use as many workers as possible with pytest-xdist
   "--tb=short",
   "--import-mode=importlib",
   "--strict-markers",

src/data_morph/cli.py

@@ -3,23 +3,30 @@
 from __future__ import annotations
 
 import argparse
-import sys
+import itertools
+import multiprocessing
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 from . import __version__
 from .data.loader import DataLoader
 from .morpher import DataMorpher
+from .progress import DataMorphProgress
 from .shapes.factory import ShapeFactory
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+    from rich.progress import TaskID
+
 ARG_DEFAULTS = {
     'output_dir': 'morphed_data',
     'decimals': 2,
     'min_shake': 0.3,
     'iterations': 100_000,
     'freeze': 0,
+    'workers': 2,
 }
 
 
@@ -50,6 +57,16 @@ def generate_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         '--version', action='version', version=f'%(prog)s {__version__}'
     )
+    parser.add_argument(
+        '-w',
+        '--workers',
+        type=int,
+        default=ARG_DEFAULTS['workers'],
+        help=(
+            f'The number of workers. Default {ARG_DEFAULTS["workers"]}. '
+            'Pass 0 to use as many as possible.'
+        ),
+    )
 
     shape_config_group = parser.add_argument_group(
         'Shape Configuration (required)',
@@ -226,33 +243,151 @@ def generate_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def main(argv: Sequence[str] | None = None) -> None:
+def _morph(
+    data: str,
+    shape: str,
+    args: argparse.Namespace,
+    progress: multiprocessing.DictProxy,
+    task_id: TaskID,
+) -> None:
     """
-    Run Data Morph as a script.
+    Run the morphing algorithm.
 
     Parameters
     ----------
-    argv : Sequence[str] | None, optional
-        Makes it possible to pass in options without running on
-        the command line.
+    data : str
+        The dataset to use. This can be the name of a built-in dataset or a path to a
+        CSV file containing the data.
+    shape : str
+        The name of the target shape.
+    args : argparse.Namespace
+        Command line arguments.
+    progress : multiprocessing.DictProxy
+        The state of all task progresses.
+    task_id : TaskID
+        The task ID assigned by the progress tracker.
+
+    Notes
+    -----
+    This should only be used with :func:`._parallelize`.
     """
+    progress[task_id] = {'progress': 0, 'total': args.iterations}
 
-    args = generate_parser().parse_args(argv)
+    dataset = DataLoader.load_dataset(data, scale=args.scale)
+    shape = ShapeFactory(dataset).generate_shape(shape)
 
-    target_shapes = (
-        ShapeFactory.AVAILABLE_SHAPES
-        if args.target_shape == 'all' or 'all' in args.target_shape
-        else set(args.target_shape).intersection(ShapeFactory.AVAILABLE_SHAPES)
+    morpher = DataMorpher(
+        decimals=args.decimals,
+        output_dir=args.output_dir,
+        write_data=args.write_data,
+        seed=args.seed,
+        keep_frames=args.keep_frames,
+        forward_only_animation=args.forward_only,
+        num_frames=100,
+        in_notebook=False,
     )
-    if not target_shapes:
-        raise ValueError(
-            'No valid target shapes were provided. Valid options are '
-            f"""'{"', '".join(ShapeFactory.AVAILABLE_SHAPES)}'."""
-        )
 
+    _ = morpher.morph(
+        start_shape=dataset,
+        target_shape=shape,
+        iterations=args.iterations,
+        min_shake=args.shake,
+        ease_in=args.ease_in or args.ease,
+        ease_out=args.ease_out or args.ease,
+        freeze_for=args.freeze,
+        progress=progress,
+        task_id=task_id,
+    )
+
+
+def _parallelize(
+    total_jobs: int,
+    workers: int,
+    args: argparse.Namespace,
+    target_shapes: Sequence[str],
+) -> None:
+    """
+    Run morphing algorithm in parallel.
+
+    Parameters
+    ----------
+    total_jobs : int
+        The total number of morphing jobs that need to be run.
+    workers : int
+        The number of worker processes to use.
+    args : argparse.Namespace
+        The command line arguments.
+    target_shapes : Sequence[str]
+        The target shapes for morphing (datasets are in ``args.start_shape``).
+    """
+    only_show_running = total_jobs > workers
+
+    with (
+        DataMorphProgress() as progress_tracker,
+        multiprocessing.Manager() as manager,
+    ):
+        task_progress = manager.dict()
+        overall_progress_task = progress_tracker.add_task('[green]Overall progress')
+
+        with ProcessPoolExecutor(max_workers=workers) as executor:
+            futures = [
+                executor.submit(
+                    _morph,
+                    dataset,
+                    shape,
+                    args,
+                    task_progress,
+                    progress_tracker.add_task(
+                        f'{Path(dataset).stem} to {shape}', visible=False, start=False
+                    ),
+                )
+                for dataset, shape in itertools.product(args.start_shape, target_shapes)
+            ]
+
+            while True:
+                finished_jobs = sum(future.done() for future in futures)
+                progress_tracker.update(
+                    overall_progress_task,
+                    completed=sum(task['progress'] for task in task_progress.values()),
+                    total=total_jobs * args.iterations,
+                )
+                for task_id, update_data in task_progress.items():
+                    latest = update_data['progress']
+                    total = update_data['total']
+
+                    if not latest:
+                        # hack to make the elapsed time accurate for ones that start later on
+                        # this is necessary because rich.progress.Progress is not pickleable
+                        progress_tracker.start_task(task_id)
+
+                    progress_tracker.update(
+                        task_id,
+                        completed=latest,
+                        total=total,
+                        visible=latest < total
+                        if only_show_running
+                        else latest <= total,
+                    )
+                if finished_jobs == total_jobs:
+                    break
+
+            for future in futures:
+                future.result()
+
+
+def _serialize(args: argparse.Namespace, target_shapes: Sequence[str]) -> None:
+    """
+    Run the morphing algorithm serially.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        The command line arguments.
+    target_shapes : Sequence[str]
+        The target shapes for morphing (datasets are in ``args.start_shape``).
+    """
     for start_shape in args.start_shape:
         dataset = DataLoader.load_dataset(start_shape, scale=args.scale)
-        print(f"Processing starter shape '{dataset.name}'", file=sys.stderr)
 
         shape_factory = ShapeFactory(dataset)
         morpher = DataMorpher(
@@ -266,10 +401,7 @@ def main(argv: Sequence[str] | None = None) -> None:
             in_notebook=False,
         )
 
-        total_shapes = len(target_shapes)
-        for i, target_shape in enumerate(target_shapes, start=1):
-            if total_shapes > 1:
-                print(f'Morphing shape {i} of {total_shapes}', file=sys.stderr)
+        for target_shape in target_shapes:
             _ = morpher.morph(
                 start_shape=dataset,
                 target_shape=shape_factory.generate_shape(target_shape),
@@ -279,3 +411,38 @@ def main(argv: Sequence[str] | None = None) -> None:
                 ease_out=args.ease_out or args.ease,
                 freeze_for=args.freeze,
             )
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    """
+    Run Data Morph as a script.
+
+    Parameters
+    ----------
+    argv : Sequence[str] | None, optional
+        Makes it possible to pass in options without running on
+        the command line.
+    """
+
+    args = generate_parser().parse_args(argv)
+
+    target_shapes = (
+        ShapeFactory.AVAILABLE_SHAPES
+        if args.target_shape == 'all' or 'all' in args.target_shape
+        else set(args.target_shape).intersection(ShapeFactory.AVAILABLE_SHAPES)
+    )
+    if not target_shapes:
+        raise ValueError(
+            'No valid target shapes were provided. Valid options are '
+            f"""'{"', '".join(ShapeFactory.AVAILABLE_SHAPES)}'."""
+        )
+
+    total_jobs = len(args.start_shape) * len(target_shapes)
+
+    max_workers = multiprocessing.cpu_count()
+    workers = max_workers if not args.workers else min(args.workers, max_workers)
+
+    if total_jobs > 1 and workers > 1:
+        _parallelize(total_jobs, workers, args, target_shapes)
+    else:
+        _serialize(args, target_shapes)