Support for Multiple Devices and Parallelism via MPI

[YigitElma]

Support for multi-device optimization.

• Updates set_device and desc.backend for initialization of multiple devices. These are the cause of slightly low coverage (around 90% for patch)
• Adds _device_id and _device to _Objective class (defaults to 0) for making parallelization work with other objectives
• Uses mpi4py for compute and jvp methods of parallel objective functions (there is no separate class but if user passes objectives with different devices and an mpi communicator, parallel strategy will be used.) To make this work a user has to use context manager properly
• Adds _rank to _Objective that specifies which objective should run on which rank. Default is None, and if all sub-objectives have None, ObjectiveFunction defaults to np.arange(len(objs)). There are some limitations for the user. For example, 2 different rank cannot use the same device (maybe it can, but it shouldn't in practice). One tricky part about this is, for multiple node case, same device_id can mean different devices. The reason we have rank is to be able to distinguish that (MPI.size is equal to total number of devices in all nodes, whereas max(device_id) can be only the number of GPUs connected to a single node)

Things to consider in future PRs,

• Maybe implement a new optimizer that uses distributed matrix operations instead of QR and SVD.
• Find good distributed linear algebra routines for SVD and QR (necessary for speed and low memory usage since the current implementation will use CPU if the Jacobian is big)
  • Something like TSQR implementation via sharding? This should reduce QR of (CM)xN matrix into C MxN QRs plus a single (CN)xN QR. By some careful implementation (without storing Q, like in our custom qr_multiply), this can be faster.
• Use CUDA-aware MPI (possibly without mpi4jax). I added desc_config["mpi-cuda"] for related parts, it should be easy to implement in the future
• Incorporate sharding if possible?

Note for reviewers:

• Actual code changes live in objective_funs.py, _constraint_wrappers.py and optimizer.py
• All the changes in _Objective class are adding _device_id and _rank
• There are around 1100 lines of code just for documentation and ~750 lines for testing

Resolves #1071 (but with mpi4py)
Resolves #1601

[codecov]

Codecov Report

❌ Patch coverage is 91.20521% with 27 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.48%. Comparing base (b08eb42) to head (4f04d45).

Files with missing lines	Patch %	Lines
desc/backend.py	51.85%	13 Missing ⚠️
desc/objectives/objective_funs.py	95.65%	9 Missing ⚠️
desc/optimize/_constraint_wrappers.py	92.85%	3 Missing ⚠️
desc/optimize/optimizer.py	93.33%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1495      +/-   ##
==========================================
+ Coverage   94.37%   94.48%   +0.11%     
==========================================
  Files         101      100       -1     
  Lines       28953    29153     +200     
==========================================
+ Hits        27323    27546     +223     
+ Misses       1630     1607      -23     

Files with missing lines	Coverage Δ
desc/objectives/_bootstrap.py	95.18% <ø> (ø)
desc/objectives/_coils.py	99.34% <ø> (ø)
desc/objectives/_equilibrium.py	96.63% <ø> (+0.42%)	⬆️
desc/objectives/_fast_ion.py	95.06% <ø> (ø)
desc/objectives/_free_boundary.py	96.78% <ø> (ø)
desc/objectives/_generic.py	99.58% <ø> (ø)
desc/objectives/_geometry.py	96.96% <ø> (ø)
desc/objectives/_neoclassical.py	94.87% <ø> (ø)
desc/objectives/_omnigenity.py	97.04% <ø> (ø)
desc/objectives/_power_balance.py	91.75% <ø> (ø)
... and 7 more

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +6.00 +/- 3.80     | +4.87e-02 +/- 3.09e-02 |  8.61e-01 +/- 2.7e-02  |  8.12e-01 +/- 1.5e-02  |
 test_equilibrium_init_lowres            |     +6.94 +/- 7.78     | +4.41e-01 +/- 4.94e-01 |  6.79e+00 +/- 4.6e-01  |  6.35e+00 +/- 1.7e-01  |
 test_objective_compile_atf              |     +2.30 +/- 5.15     | +1.33e-01 +/- 2.99e-01 |  5.93e+00 +/- 2.4e-01  |  5.80e+00 +/- 1.8e-01  |
 test_objective_compute_atf              |     -0.83 +/- 14.46    | -1.75e-05 +/- 3.04e-04 |  2.09e-03 +/- 1.3e-04  |  2.11e-03 +/- 2.7e-04  |
 test_objective_jac_atf                  |     +0.85 +/- 3.04     | +1.42e-02 +/- 5.04e-02 |  1.68e+00 +/- 3.2e-02  |  1.66e+00 +/- 3.9e-02  |
 test_perturb_1                          |     +0.18 +/- 1.53     | +2.03e-02 +/- 1.74e-01 |  1.14e+01 +/- 1.1e-01  |  1.14e+01 +/- 1.3e-01  |
 test_proximal_jac_atf                   |     +0.08 +/- 2.01     | +4.14e-03 +/- 1.09e-01 |  5.41e+00 +/- 7.2e-02  |  5.40e+00 +/- 8.1e-02  |
 test_proximal_freeb_compute             |     +0.85 +/- 3.32     | +1.15e-03 +/- 4.53e-03 |  1.38e-01 +/- 3.0e-03  |  1.36e-01 +/- 3.4e-03  |
 test_solve_fixed_iter                   |     -0.27 +/- 3.28     | -6.41e-02 +/- 7.88e-01 |  2.39e+01 +/- 6.3e-01  |  2.40e+01 +/- 4.7e-01  |
 test_LinearConstraintProjection_build   |     -0.69 +/- 4.55     | -4.65e-02 +/- 3.08e-01 |  6.72e+00 +/- 2.3e-01  |  6.77e+00 +/- 2.1e-01  |
 test_objective_compute_ripple           |     +0.86 +/- 4.32     | +1.83e-03 +/- 9.15e-03 |  2.14e-01 +/- 7.0e-03  |  2.12e-01 +/- 5.9e-03  |
 test_objective_grad_ripple              |     -0.96 +/- 3.17     | -8.75e-03 +/- 2.90e-02 |  9.06e-01 +/- 1.9e-02  |  9.15e-01 +/- 2.2e-02  |
 test_objective_quadratic_flux_compute   |     -1.04 +/- 18.56    | -5.69e-04 +/- 1.02e-02 |  5.42e-02 +/- 8.4e-03  |  5.48e-02 +/- 5.8e-03  |
 test_build_transform_fft_midres         |     +1.93 +/- 4.96     | +1.76e-02 +/- 4.53e-02 |  9.32e-01 +/- 2.8e-02  |  9.14e-01 +/- 3.5e-02  |
 test_build_transform_fft_highres        |     +2.15 +/- 2.72     | +2.59e-02 +/- 3.28e-02 |  1.23e+00 +/- 3.0e-02  |  1.20e+00 +/- 1.2e-02  |
 test_equilibrium_init_medres            |     +3.56 +/- 3.08     | +2.57e-01 +/- 2.22e-01 |  7.47e+00 +/- 1.8e-01  |  7.21e+00 +/- 1.3e-01  |
 test_objective_compile_dshape_current   |     +0.01 +/- 2.36     | +5.41e-04 +/- 1.01e-01 |  4.26e+00 +/- 5.0e-02  |  4.26e+00 +/- 8.7e-02  |
 test_objective_compute_dshape_current   |     -4.01 +/- 11.93    | -3.07e-05 +/- 9.12e-05 |  7.34e-04 +/- 7.1e-05  |  7.65e-04 +/- 5.7e-05  |
 test_objective_jac_dshape_current       |     +1.17 +/- 25.82    | +2.95e-04 +/- 6.51e-03 |  2.55e-02 +/- 4.8e-03  |  2.52e-02 +/- 4.4e-03  |
 test_perturb_2                          |     +0.38 +/- 1.48     | +6.14e-02 +/- 2.38e-01 |  1.61e+01 +/- 1.8e-01  |  1.61e+01 +/- 1.5e-01  |
 test_proximal_jac_atf_with_eq_update    |     -0.81 +/- 3.47     | -1.01e-01 +/- 4.32e-01 |  1.23e+01 +/- 6.3e-02  |  1.25e+01 +/- 4.3e-01  |
 test_proximal_freeb_jac                 |     -0.52 +/- 3.76     | -2.54e-02 +/- 1.85e-01 |  4.91e+00 +/- 4.4e-02  |  4.93e+00 +/- 1.8e-01  |
 test_solve_fixed_iter_compiled          |     +0.47 +/- 3.05     | +2.94e-02 +/- 1.89e-01 |  6.24e+00 +/- 1.1e-01  |  6.21e+00 +/- 1.5e-01  |
 test_objective_compute_ripple_bounce1d  |     -1.03 +/- 6.64     | -3.18e-03 +/- 2.04e-02 |  3.05e-01 +/- 1.8e-02  |  3.08e-01 +/- 9.7e-03  |
 test_objective_grad_ripple_bounce1d     |     +0.78 +/- 3.00     | +7.56e-03 +/- 2.92e-02 |  9.80e-01 +/- 1.9e-02  |  9.72e-01 +/- 2.2e-02  |
 test_objective_quadratic_flux_jac       |     +1.49 +/- 0.76     | +1.26e-01 +/- 6.46e-02 |  8.60e+00 +/- 2.2e-02  |  8.47e+00 +/- 6.1e-02  |

Github CI performance can be noisy. When evaluating the benchmarks, developers should take this into account.

[dpanici]

#763
check for overlap with this one

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[YigitElma]

If you have been using this branch, I made a recent change. Now, instead of giving rank_per_objective to ObjectiveFunction, you should give it to each sub-objective. Default behavior of not having any rank is the same.

[YigitElma]

I can finalize this PR for parallel constraints in Proximal too after #2239