Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 71.45%. Comparing base (0e70789) to head (55e8469).
Report is 4 commits behind head on master.

❗ Current head 55e8469 differs from pull request most recent head f5bb285

Please upload reports for the commit f5bb285 to get more accurate results.

@@            Coverage Diff             @@
##           master     #236      +/-   ##
==========================================
+ Coverage   70.21%   71.45%   +1.24%     
==========================================
  Files          34       34              
  Lines        2078     2172      +94     
  Branches      157      164       +7     
==========================================
+ Hits         1459     1552      +93     
+ Misses        559      558       -1     
- Partials       60       62       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Installing this branch with

pip install --no-build-isolation --verbose git+https://github.com/mpi4jax/mpi4jax.git@pv/test-install

works and correctly picks up the Nvidia packages if they are installed.

However one then needs to set the LD_LIBRARY_PATH because the path is not hardcoded... with something like

export LD_LIBRARY_PATH=[...]/ENV_NAME/lib/python3.11/site-packages/nvidia/cuda_runtime/lib:$LD_LIBRARY_PATH

Nevertheless, this still gives raise to some annoying ptas errors..

Though it works for standard code

But this is nice. It allows to avoid requiring CUDA installed....

This now sets the rpath so it bakes the compile-time libraries if it detects it's obtained from pip...

@dionhaefner what can we do to package this properly?

This looks generally fine to me.

Setting rpath means that mpi4jax may break violently if the user ever updates the CUDA libs but I guess it's the best we can do, barring using dlopen at runtime like XLA does.

I think the most important task now is to test all modes of installation on CI. That should give us enough confidence to ship this new feature pretty much as-is.

@dionhaefner I added a test for the build infrastructure using PyPi distributed cuda.

I also added the relevant instructions and a mention that [cuda12] is not yet supported.

Anything else?

Is it fine for you if I merge and tag a new release?

Fine with me, I could pick some nits with the logic in setup.py and the wording in the docs, but feels like this is up for a refactor soon anyway once we have to support jax[cuda12] (when jax[cuda12_pip] goes away).

yes, that's what I was thinking. It seems that the issues with supporting jax[cuda12] are not up to us, but because the present interface based on the pjrt runtime does not support defining custom calls from python.

They are fixing it on jax those days. I guess I'll try again in a week or two.