REF: Reorganize how multi-device parallelism is implemented

Purpose

Reduce communications overhead for multi-device reconstructions and hopefully speed up large reconstructions.

Separate multi-device parallelism from reconstruction logic (ptycho solver implementation).

Approach

Instead of synchronizing every update to the probe/object across all devices, use an approach that is more like distributed consensus where devices only communicate with their neighbors for object updates and synchronization only happens once per epoch.

The effect of this approach on performance is minor reduction in time per epoch for all numbers of GPUs, but increase in the number of epochs required for convergence for multi-gpu reconstructions because information is exchanged between GPUs less frequently.

Screenshot from 2024-07-09 12-28-25

Pre-Merge Checklists

Submitter

[ ] Write a helpfully descriptive pull request title.
[ ] Organize changes into logically grouped commits with descriptive commit messages.
[ ] Document all new functions.
[ ] Click 'details' on the readthedocs check to view the updated docs.
[ ] Write tests for new functions or explain why they are not needed.
[ ] Address any complaints from pep8speaks.

Reviewer

[ ] Actually read all of the code.
[ ] Run the new code yourself; the included tests should make this easy.
[ ] Write a summary of the changes as you understand them.
[ ] Thank the submitter.

Hello @carterbox! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

In the file src/tike/communicators/pool.py:

Line 459:81: E501 line too long (81 > 80 characters)

In the file src/tike/ptycho/probe.py:

Line 158:81: E501 line too long (81 > 80 characters) Line 180:22: E225 missing whitespace around operator Line 218:22: E225 missing whitespace around operator Line 247:22: E225 missing whitespace around operator Line 398:29: E203 whitespace before ':' Line 398:40: E203 whitespace before ':' Line 410:48: E203 whitespace before ':' Line 410:55: E203 whitespace before ':' Line 411:27: E203 whitespace before ':' Line 411:38: E203 whitespace before ':' Line 419:25: E203 whitespace before ':' Line 419:32: E203 whitespace before ':' Line 428:25: E203 whitespace before ':' Line 428:32: E203 whitespace before ':' Line 429:29: E203 whitespace before ':' Line 429:36: E203 whitespace before ':' Line 436:41: E203 whitespace before ':' Line 436:48: E203 whitespace before ':' Line 438:29: E203 whitespace before ':' Line 453:21: E203 whitespace before ':' Line 453:32: E203 whitespace before ':'

In the file src/tike/ptycho/ptycho.py:

Line 395:81: E501 line too long (101 > 80 characters) Line 396:81: E501 line too long (105 > 80 characters) Line 397:81: E501 line too long (103 > 80 characters) Line 407:81: E501 line too long (101 > 80 characters) Line 408:81: E501 line too long (105 > 80 characters) Line 409:81: E501 line too long (103 > 80 characters) Line 419:81: E501 line too long (101 > 80 characters) Line 420:81: E501 line too long (105 > 80 characters) Line 421:81: E501 line too long (103 > 80 characters) Line 500:81: E501 line too long (87 > 80 characters) Line 710:1: E302 expected 2 blank lines, found 1 Line 721:81: E501 line too long (84 > 80 characters) Line 749:81: E501 line too long (88 > 80 characters) Line 868:81: E501 line too long (81 > 80 characters) Line 871:81: E501 line too long (86 > 80 characters)

In the file src/tike/ptycho/solvers/lstsq.py:

Line 83:81: E501 line too long (85 > 80 characters) Line 97:81: E501 line too long (85 > 80 characters) Line 290:81: E501 line too long (85 > 80 characters) Line 320:81: E501 line too long (85 > 80 characters)

In the file src/tike/ptycho/solvers/options.py:

Line 280:81: E501 line too long (82 > 80 characters)

In the file src/tike/ptycho/solvers/rpie.py:

Line 90:81: E501 line too long (85 > 80 characters) Line 275:81: E501 line too long (82 > 80 characters) Line 345:81: E501 line too long (81 > 80 characters)

In the file tests/communicators/test_pool.py:

Line 146:81: E501 line too long (81 > 80 characters) Line 162:81: E501 line too long (82 > 80 characters) Line 164:36: E203 whitespace before ':' Line 166:40: E203 whitespace before ':'

AdvancedPhotonSource / tike