{ai,bio}[foss/2024a] AlphaFold3 v3.0.1, jax-triton v0.2.0, jaxtyping v0.2.38, dm-tree v0.1.8 w/ CUDA 12.6.0 DRAFT!

[ThomasHoffmann77]

(created using eb --new-pr)

• {devel}[foss/2024a] Triton v3.1.0 w/ CUDA 12.6.0 #22064
• {lib}[foss/2024a] dm-haiku v0.0.13 w/ CUDA 12.6.0 #22028
• {ai}[gfbf/2024a] jax v0.4.35, ml_dtypes v0.5.0 w/ CUDA 12.6.0 WIP #21924
• {lib}[GCC/13.3.0] pybind11_abseil v202402.0 #22153
• {bio}[GCC/13.3.0] dssp v4.4.10, libcifpp v7.0.8, libmcfp v1.3.4, mrc v1.3.13 #22155
• {chem}[foss/2024a] RDKit v2025_03_3 #22166
• {bio}[gompi/2024a] HMMER v3.4 #22172

TODO:

• "alphafold3 3.0.0 requires rdkit, which is not installed." although RDKit is dependency -> add minimal rdkit-2024.03.5.dist-info
• disable fetchcontent
• dssp
• pybind11_abseil
• RDKit: check whether 2024_09_5 (latest) would also work
• Check, whether $TRITON_HOME is used at all. Otherwise remove from module(xxx)footer.
• teststep: run_alphafold_tests.py
• use jax 0.6.0 ?

[github-actions]

Updated software dm-tree-0.1.9-gfbf-2024a.eb

Diff against dm-tree-0.1.8-GCCcore-12.3.0.eb

easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.8-GCCcore-12.3.0.eb

diff --git a/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.8-GCCcore-12.3.0.eb b/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.9-gfbf-2024a.eb
index ed35d0496f..8737dcb334 100644
--- a/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.8-GCCcore-12.3.0.eb
+++ b/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.9-gfbf-2024a.eb
@@ -1,23 +1,32 @@
 easyblock = 'PythonPackage'
 
 name = 'dm-tree'
-version = '0.1.8'
+version = '0.1.9'
 
 homepage = 'https://github.com/deepmind/tree'
 description = """dm-tree provides tree, a library for working with nested data structures. In a way,
 tree generalizes the builtin map function which only supports flat sequences, and
 allows to apply a function to each "leaf" preserving the overall structure."""
 
-toolchain = {'name': 'GCCcore', 'version': '12.3.0'}
+toolchain = {'name': 'gfbf', 'version': '2024a'}
 
-sources = [SOURCELOWER_TAR_GZ]
-checksums = ['0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430']
+sources = ['dm_tree-%(version)s.tar.gz']
+patches = ['dm-tree-0.1.9_pybind11.patch']
+checksums = [
+    {'dm_tree-0.1.9.tar.gz': 'a4c7db3d3935a5a2d5e4b383fc26c6b0cd6f78c6d4605d3e7b518800ecd5342b'},
+    {'dm-tree-0.1.9_pybind11.patch': '39ae1e40e6a027a6c6e085b20d33fe8de0e8103963c2ae0da665597ed9c6b41b'},
+]
 
 builddependencies = [
-    ('binutils', '2.40'),
-    ('CMake', '3.26.3'),
+    ('pybind11', '2.12.0'),
+]
+dependencies = [
+    ('Python', '3.12.3'),
+    ('absl-py', '2.1.0'),
+    ('Abseil', '20240722.0'),
+    ('SciPy-bundle', '2024.05'),
+    ('wrapt', '1.16.0'),
 ]
-dependencies = [('Python', '3.11.3')]
 
 options = {'modulename': 'tree'}
 

Diff against dm-tree-0.1.8-GCCcore-11.3.0.eb

easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.8-GCCcore-11.3.0.eb

diff --git a/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.8-GCCcore-11.3.0.eb b/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.9-gfbf-2024a.eb
index 4831e9d08f..8737dcb334 100644
--- a/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.8-GCCcore-11.3.0.eb
+++ b/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.9-gfbf-2024a.eb
@@ -1,23 +1,32 @@
 easyblock = 'PythonPackage'
 
 name = 'dm-tree'
-version = '0.1.8'
+version = '0.1.9'
 
 homepage = 'https://github.com/deepmind/tree'
 description = """dm-tree provides tree, a library for working with nested data structures. In a way,
 tree generalizes the builtin map function which only supports flat sequences, and
 allows to apply a function to each "leaf" preserving the overall structure."""
 
-toolchain = {'name': 'GCCcore', 'version': '11.3.0'}
+toolchain = {'name': 'gfbf', 'version': '2024a'}
 
-sources = [SOURCELOWER_TAR_GZ]
-checksums = ['0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430']
+sources = ['dm_tree-%(version)s.tar.gz']
+patches = ['dm-tree-0.1.9_pybind11.patch']
+checksums = [
+    {'dm_tree-0.1.9.tar.gz': 'a4c7db3d3935a5a2d5e4b383fc26c6b0cd6f78c6d4605d3e7b518800ecd5342b'},
+    {'dm-tree-0.1.9_pybind11.patch': '39ae1e40e6a027a6c6e085b20d33fe8de0e8103963c2ae0da665597ed9c6b41b'},
+]
 
 builddependencies = [
-    ('binutils', '2.38'),
-    ('CMake', '3.23.1'),
+    ('pybind11', '2.12.0'),
+]
+dependencies = [
+    ('Python', '3.12.3'),
+    ('absl-py', '2.1.0'),
+    ('Abseil', '20240722.0'),
+    ('SciPy-bundle', '2024.05'),
+    ('wrapt', '1.16.0'),
 ]
-dependencies = [('Python', '3.10.4')]
 
 options = {'modulename': 'tree'}
 

Diff against dm-tree-0.1.6-GCCcore-10.3.0.eb

easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.6-GCCcore-10.3.0.eb

diff --git a/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.6-GCCcore-10.3.0.eb b/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.9-gfbf-2024a.eb
index efd0aee27a..8737dcb334 100644
--- a/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.6-GCCcore-10.3.0.eb
+++ b/easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.9-gfbf-2024a.eb
@@ -1,23 +1,32 @@
 easyblock = 'PythonPackage'
 
 name = 'dm-tree'
-version = '0.1.6'
+version = '0.1.9'
 
 homepage = 'https://github.com/deepmind/tree'
 description = """dm-tree provides tree, a library for working with nested data structures. In a way,
 tree generalizes the builtin map function which only supports flat sequences, and
 allows to apply a function to each "leaf" preserving the overall structure."""
 
-toolchain = {'name': 'GCCcore', 'version': '10.3.0'}
+toolchain = {'name': 'gfbf', 'version': '2024a'}
 
-sources = [SOURCELOWER_TAR_GZ]
-checksums = ['6776404b23b4522c01012ffb314632aba092c9541577004ab153321e87da439a']
+sources = ['dm_tree-%(version)s.tar.gz']
+patches = ['dm-tree-0.1.9_pybind11.patch']
+checksums = [
+    {'dm_tree-0.1.9.tar.gz': 'a4c7db3d3935a5a2d5e4b383fc26c6b0cd6f78c6d4605d3e7b518800ecd5342b'},
+    {'dm-tree-0.1.9_pybind11.patch': '39ae1e40e6a027a6c6e085b20d33fe8de0e8103963c2ae0da665597ed9c6b41b'},
+]
 
 builddependencies = [
-    ('binutils', '2.36.1'),
-    ('Bazel', '3.7.2'),
+    ('pybind11', '2.12.0'),
+]
+dependencies = [
+    ('Python', '3.12.3'),
+    ('absl-py', '2.1.0'),
+    ('Abseil', '20240722.0'),
+    ('SciPy-bundle', '2024.05'),
+    ('wrapt', '1.16.0'),
 ]
-dependencies = [('Python', '3.9.5')]
 
 options = {'modulename': 'tree'}
 

[ThomasHoffmann77]

dm-tree pybind11 patch taken from debian: https://udd.debian.org/patches.cgi?src=dm-tree&version=0.1.8-3

[ThomasHoffmann77]

Bump version, add runtest and a nitpick on the description. (Note, test step doesn't currently pass with these changes).

Thanks @VRehnberg

[VRehnberg]

Loading the a built module (--ignore-test-failure) and running tests with PYTHONPATH="$(dirname $(which run_alphafold.py)):$PYTHONPATH" python, then run_alphafold_data_test.py passes everything. However, run_alphafold_data_test.py fails. With this Jax patch suggested in #21924 (comment) I have a single remaining failing test which probably is just slightly too strict tolerances:

======================================================================
FAIL: test_inference_default_bucket (__main__.InferenceTest)
InferenceTest.test_inference_default_bucket
test_inference_default_bucket(bucket=None, exp_ranking_scores=[0.69, 0.69, 0.72, 0.75, 0.7])
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/apps/Test/software/absl-py/2.1.0-GCCcore-13.3.0/lib/python3.12/site-packages/absl/testing/parameterized.py", line 319, in bound_param_test
    return test_method(self, **testcase_params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/cephyr/users/vikren/Alvis/af3_test/run_alphafold_test.py", line 335, in test_inference
    np.testing.assert_array_almost_equal(
  File "/apps/Arch/software/Python/3.12.3-GCCcore-13.3.0/lib/python3.12/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/apps/Arch/software/SciPy-bundle/2024.05-gfbf-2024a/lib/python3.12/site-packages/numpy/testing/_private/utils.py", line 1034, in assert_array_almost_equal
    assert_array_compare(compare, x, y, err_msg=err_msg, verbose=verbose,
  File "/apps/Arch/software/Python/3.12.3-GCCcore-13.3.0/lib/python3.12/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/apps/Arch/software/SciPy-bundle/2024.05-gfbf-2024a/lib/python3.12/site-packages/numpy/testing/_private/utils.py", line 797, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Arrays are not almost equal to 2 decimals

Mismatched elements: 4 / 5 (80%)
Max absolute difference: 0.02214217
Max relative difference: 0.03209009
 x: array([0.71, 0.71, 0.74, 0.74, 0.72])
 y: array([0.69, 0.69, 0.72, 0.75, 0.7 ])

----------------------------------------------------------------------
Ran 4 tests in 390.760s

FAILED (failures=1)

[VRehnberg]

Update: I've been providing this in a test tree for some users at our site for some time now and was recently in contact with two of the ones that have tried it and it seems to be working fine for them.

[boegel]

Update: I've been providing this in a test tree for some users at our site for some time now and was recently in contact with two of the ones that have tried it and it seems to be working fine for them.

Did you just install all dependencies from the open PRs and then this on top, no changes?
Just trying to get a better view on this, I'm quite keen on getting this PR merged soon actually...

[VRehnberg]

Did you just install all dependencies from the open PRs and then this on top, no changes?

I added one change to jax #21924 (comment) and maybe changed some of the environment variables as suggested in comments. No other changes as far as I recall.

[Flamefire]

As I'm working on PyTorch 2.7: Is it possible to use Triton 3.3.1 for this from this PR? That is the one required by PyTorch and I'm not sure it would work with the older one

Check, whether $TRITON_HOME is used at all. Otherwise remove from module(xxx)footer.

If it is it might be better to set it in the triton module.

[pavelToman]

Hello @ThomasHoffmann77, can I help with AlphaFold3? We have another researchers asking for this.
I can help with the jax I hope - I made updates to the jaxlib easyblock here: easybuilders/easybuild-easyblocks#3852, but it is only for non-CUDA versions of new jax. So I can look either in the CUDA versions, if this helps.

[ThomasHoffmann77]

Hello @ThomasHoffmann77, can I help with AlphaFold3? We have another researchers asking for this. I can help with the jax I hope - I made updates to the jaxlib easyblock here: easybuilders/easybuild-easyblocks#3852, but it is only for non-CUDA versions of new jax. So I can look either in the CUDA versions, if this helps.

Hi @pavelToman! Help would be very much appreciated as I am currently quite busy with other projects. My concerns about the current state of #21924 is that it still downloads plenty of bazel artifacts and I see some problems with using hermetic python. As far as I know more recent Bazel versions, as used for jax 0.6/7 have some options to better control offline builds.
For AF3 a cuda capable build of jax is required.

[pavelToman]

Hello @ThomasHoffmann77, can I help with AlphaFold3? We have another researchers asking for this. I can help with the jax I hope - I made updates to the jaxlib easyblock here: easybuilders/easybuild-easyblocks#3852, but it is only for non-CUDA versions of new jax. So I can look either in the CUDA versions, if this helps.

Hi @pavelToman! Help would be very much appreciated as I am currently quite busy with other projects. My concerns about the current state of #21924 is that it still downloads plenty of bazel artifacts and I see some problems with using hermetic python. As far as I know more recent Bazel versions, as used for jax 0.6/7 have some options to better control offline builds. For AF3 a cuda capable build of jax is required.

Ok, working on jax-0.6.2-gfbf-2024a-CUDA-12.6.0 and easyblock jaxlib.py update!

[pavelToman]

The jax-0.6.2 with CUDA-12.6.0 is ready!

• {ai}[gfbf/2024a] jax v0.6.2 w/ CUDA 12.6.0 #24141
• Update jaxlib easyblock for jax 0.6.2 with CUDA-12.6.0 easybuild-easyblocks#3951

[pavelToman]

Hello @ThomasHoffmann77, the jax-0.6.2 with CUDA seems ready, what else could I do to let AF3 works? What about dm-haiku-0.0.13, this one should now works with the new jax?

[ThomasHoffmann77]

Hello @ThomasHoffmann77, the jax-0.6.2 with CUDA seems ready, what else could I do to let AF3 works? What about dm-haiku-0.0.13, this one should now works with the new jax?

Hi @pavelToman! Great news! Thanks a lot. Some of AF3's dependencies also depend on jax and need to be adapted: jaxtyping, jax-triton, dm-haiku, and Optax

[pavelToman]

I've just open Draft PR for now (could be merged here latter):

• {ai}[foss/2024a] AlphaFold3 v3.0.1-20250908, dm-haiku v0.0.14, Flax v0.10.2, ... w/ CUDA 12.6.0 #24177

It is the copy of this easyconfig and patches for AF3.
I just update jax to v0.6.2-CUDA-12.6.0 and all the AF3 deps with jax as dep: Optax, Flax, dm-haiku, jaxtyping, jax-triton
Also Triton is updated to v3.3.1.

Since I still do not have googles models to fully test the installation, could you with models test it please?