AlphaFold 3

[laraPPr]

• link to support ticket: 2024120960000091
• website: https://github.com/google-deepmind/alphafold3
• installation docs: https://github.com/google-deepmind/alphafold3/blob/main/docs/installation.md
• toolchain: foss/2024a?
• easyblock to use: PythonBundle
• required dependencies:
  • https://github.com/google-deepmind/alphafold3/blob/main/requirements.txt
• optional dependencies:
  • ...
• notes:
  • ...
• effort: (TBD)
• other install methods
  • conda: yes (link?) / no
  • container image: yes (link?) / no
  • pre-built binaries (RHEL8 Linux x86_64): yes (link?) / no
  • easyconfig outside EasyBuild: yes / no

[lexming]

Duplicate of #454?

[laraPPr]

We can collapse them if we need to. But would also like to set our inuits people on it as well.

[laraPPr]

Someone already started on an EasyConfig easybuilders/easybuild-easyconfigs#21864

[boegel]

@pavelToman Can you look into testing easybuilders/easybuild-easyconfigs#22028 (submit test report), and figure out what needs to happen to get easybuilders/easybuild-easyconfigs#21924 across the finish line?

And then maybe also talk to @ThomasHoffmann77 and figure out what else needs to happen to get AlphaFold 3 fully supported in EasyBuild...

[ThomasHoffmann77]

@pavelToman Can you look into testing easybuilders/easybuild-easyconfigs#22028 (submit test report), and figure out what needs to happen to get easybuilders/easybuild-easyconfigs#21924 across the finish line?

And then maybe also talk to @ThomasHoffmann77 and figure out what else needs to happen to get AlphaFold 3 fully supported in EasyBuild...

• I moved the checksums in {lib}[foss/2024a] dm-haiku v0.0.13, tensorstore v0.1.69 w/ CUDA 12.6.0 easybuilders/easybuild-easyconfigs#22028 tensorstore
• current progress in AF3 easyconfig: see also {ai,bio}[foss/2024a] AlphaFold3 v3.0.0.20250109, jax-triton v0.2.0, jaxtyping v0.2.36, dm-tree v0.1.8 w/ CUDA 12.6.0 DRAFT! easybuilders/easybuild-easyconfigs#22067
• I am currently struggling with Clang sanity checks for Triton 3.1.0 ({devel}[foss/2024a] Triton v3.1.0 w/ CUDA 12.6.0 easybuilders/easybuild-easyconfigs#22064). For test purpose I omitted the sanity check step for Clang 19.0.0 at our site in order to get Triton built.

@pavelToman @boegel Update: I switched to the latest AlphaFold3 commit to fix some problems. I got the module built at our site. However, at a testrun* it turned out, that easybuilders/easybuild-easyconfigs#21924 is missing jax_cuda12_plugin and jax_cuda12_pjrt. I am working on it.
Furthermore, the jax module is still missing the test-step and bazel downloads need to be switched off.

*) AF3 testrun of the data pipeline ran through. Inference step failed.

[boegel]

@pavelToman Can you look into the easyconfig PRs that @ThomasHoffmann77 mentioned? (quick review + test)

[ThomasHoffmann77]

@pavelToman Can you look into the easyconfig PRs that @ThomasHoffmann77 mentioned? (quick review + test)

@boegel @pavelToman regarding jax: the problem with the missing jax_cuda_{plugin,pjrt} has been fixed. I also ran all tests of which 5 were failing.
Recently I was working on an update of this PR for offline build pr-downloading the complete set of bazel artifact into a dist directory. I am currently struggling with replacing hermetic python by EB python. Furthermore, I did not find a way to avoid some downloads of python packages, yet (-> patch rules_python).
I did not upload my latest work, yet.

The jax(lib) easyblock might require an update to get things less complicated. Maybe @Flamefire could help here?

Maybe it would also make sense to switch to the latest jax (0.5.0) (?)

[Flamefire]

I don't have much experience with jax and haven't built TensorFlow for quite a while now. But I can see what I can do based on what you got already. I think we have quite a bit of experience bottled in our existing ECs using Bazel.
I'm a bit busy right now so might not be able to do much soon