{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!

easybuild/easyconfigs/a/AlphaFold3/AlphaFold3-3.0.0.20250109_data_path.patch

@@ -0,0 +1,95 @@
+# Thomas Hoffmann, EMBL Heidelberg, structures-it@embl.de, 2025/01
+# - additionally search libciffpp components.cif in path definded by 
+#   environment variable LIBCIFPP_DATA_DIR
+# - allow to set default model dir and public databases dir by environment
+#   variables AF3_MODEL_DIR and DB_DIR, respectively. 
+diff -ru alphafold3-ea040346e10db1759170e723ef263316e64aa768/run_alphafold.py alphafold3-ea040346e10db1759170e723ef263316e64aa768_data_path/run_alphafold.py
+--- alphafold3-ea040346e10db1759170e723ef263316e64aa768/run_alphafold.py	2025-01-09 13:59:41.000000000 +0100
++++ alphafold3-ea040346e10db1759170e723ef263316e64aa768_data_path/run_alphafold.py	2025-01-15 18:48:19.060274236 +0100
+@@ -55,8 +55,14 @@
+
+
+ _HOME_DIR = pathlib.Path(os.environ.get('HOME'))
+-_DEFAULT_MODEL_DIR = _HOME_DIR / 'models'
+-_DEFAULT_DB_DIR = _HOME_DIR / 'public_databases'
++_DEFAULT_MODEL_DIR = os.environ.get('AF3_MODEL_DIR')
++if _DEFAULT_MODEL_DIR == None:
++  _DEFAULT_MODEL_DIR = _HOME_DIR / 'models'
++else: _DEFAULT_MODEL_DIR = pathlib.Path(_DEFAULT_MODEL_DIR)
++_DEFAULT_DB_DIR = os.environ.get('DB_DIR')
++if _DEFAULT_DB_DIR == None:
++  _DEFAULT_DB_DIR = _HOME_DIR / 'public_databases'
++else: _DEFAULT_DB_DIR = pathlib.Path(_DEFAULT_DB_DIR)
+
+
+ # Input and output paths.
+diff -ru alphafold3-ea040346e10db1759170e723ef263316e64aa768/src/alphafold3/build_data.py alphafold3-ea040346e10db1759170e723ef263316e64aa768_data_path/src/alphafold3/build_data.py
+--- alphafold3-ea040346e10db1759170e723ef263316e64aa768/src/alphafold3/build_data.py	2025-01-09 13:59:41.000000000 +0100
++++ alphafold3-ea040346e10db1759170e723ef263316e64aa768_data_path/src/alphafold3/build_data.py	2025-01-15 18:04:26.440308446 +0100
+@@ -17,16 +17,21 @@
+ import alphafold3.constants.converters
+ from alphafold3.constants.converters import ccd_pickle_gen
+ from alphafold3.constants.converters import chemical_component_sets_gen
+-
++from os import getenv
+
+ def build_data():
+   """Builds intermediate data."""
+-  for site_path in site.getsitepackages():
+-    path = pathlib.Path(site_path) / 'share/libcifpp/components.cif'
+-    if path.exists():
+-      cif_path = path
+-      break
++  cif_path = None
++  if getenv('LIBCIFPP_DATA_DIR'):
++    path = getenv('LIBCIFPP_DATA_DIR') + '/components.cif'
++    cif_path = path
+   else:
++    for site_path in site.getsitepackages():
++      path = pathlib.Path(site_path) / 'share/libcifpp/components.cif'
++      if path.exists():
++        cif_path = path
++        break
++  if (not cif_path):
+     raise ValueError('Could not find components.cif')
+
+   out_root = resources.files(alphafold3.constants.converters)
+diff -ru alphafold3-ea040346e10db1759170e723ef263316e64aa768/src/alphafold3/model/mkdssp_pybind.cc alphafold3-ea040346e10db1759170e723ef263316e64aa768_data_path/src/alphafold3/model/mkdssp_pybind.cc
+--- alphafold3-ea040346e10db1759170e723ef263316e64aa768/src/alphafold3/model/mkdssp_pybind.cc	2025-01-09 13:59:41.000000000 +0100
++++ alphafold3-ea040346e10db1759170e723ef263316e64aa768_data_path/src/alphafold3/model/mkdssp_pybind.cc	2025-01-15 14:34:35.267983665 +0100
+@@ -27,20 +27,23 @@
+ void RegisterModuleMkdssp(pybind11::module m) {
+   py::module site = py::module::import("site");
+   py::list paths = py::cast<py::list>(site.attr("getsitepackages")());
++  const char* libcifpp_data_dir_p = std::getenv("LIBCIFPP_DATA_DIR");
++  if (!libcifpp_data_dir_p) {
+   // Find the first path that contains the libcifpp components.cif file.
+-  bool found = false;
+-  for (const auto& py_path : paths) {
+-    auto path_str =
+-        std::filesystem::path(py::cast<absl::string_view>(py_path)) /
+-        "share/libcifpp/components.cif";
+-    if (std::filesystem::exists(path_str)) {
+-      setenv("LIBCIFPP_DATA_DIR", path_str.parent_path().c_str(), 0);
+-      found = true;
+-      break;
++     bool found = false;
++     for (const auto& py_path : paths) {
++       auto path_str =
++           std::filesystem::path(py::cast<absl::string_view>(py_path)) /
++           "share/libcifpp/components.cif";
++       if (std::filesystem::exists(path_str)) {
++       setenv("LIBCIFPP_DATA_DIR", path_str.parent_path().c_str(), 0);
++       found = true;
++       break;
++     }
++    }
++    if (!found) {
++      throw py::type_error("Could not find the libcifpp components.cif file.");
+     }
+-  }
+-  if (!found) {
+-    throw py::type_error("Could not find the libcifpp components.cif file.");
+   }
+   m.def(
+       "get_dssp",

easybuild/easyconfigs/a/AlphaFold3/AlphaFold3-3.0.0.20250109_relax_requirements.patch

@@ -0,0 +1,30 @@
+# Thomas Hoffmann, EMBL Heidelberg, structures-it@embl.de, 2025/01
+diff -ru alphafold3-ea040346e10db1759170e723ef263316e64aa768/pyproject.toml alphafold3-ea040346e10db1759170e723ef263316e64aa768_relax_requirements/pyproject.toml
+--- alphafold3-ea040346e10db1759170e723ef263316e64aa768/pyproject.toml	2025-01-09 13:59:41.000000000 +0100
++++ alphafold3-ea040346e10db1759170e723ef263316e64aa768_relax_requirements/pyproject.toml	2025-01-15 10:54:51.110182402 +0100
+@@ -17,17 +17,17 @@
+ dependencies = [
+     "absl-py",
+     "chex",
+-    "dm-haiku==0.0.13",
++    "dm-haiku>=0.0.13",
+     "dm-tree",
+-    "jax==0.4.34",
+-    "jax[cuda12]==0.4.34",
+-    "jax-triton==0.2.0",
+-    "jaxtyping==0.2.34",
++    "jax>=0.4.34",
++    #"jax[cuda12]==0.4.34",
++    "jax-triton>=0.2.0",
++    "jaxtyping>=0.2.34",
+     "numpy",
+-    "rdkit==2024.3.5",
+-    "triton==3.1.0",
++    "rdkit>=2024.3.5",
++    "triton>=3.1.0",
+     "tqdm",
+-    "typeguard==2.13.3",
++    "typeguard>=2.13.3",
+     "zstandard",
+ ]
+

easybuild/easyconfigs/a/AlphaFold3/AlphaFold3-3.0.0_disable_fetch.patch

@@ -0,0 +1,20 @@
+# Thomas Hoffmann, EMBL Heidelberg, structures-it@embl.de, 2025/01
+# disable FetchContent for pybind11, abseil-cpp, pybind11_abseil, cifpp and dssp
+# use find_package instead.
+diff -ru alphafold3-3.0.0/CMakeLists.txt alphafold3-3.0.0_disable_fetch/CMakeLists.txt
+--- alphafold3-3.0.0/CMakeLists.txt	2024-11-11 11:42:47.000000000 +0100
++++ alphafold3-3.0.0_disable_fetch/CMakeLists.txt	2025-01-14 17:14:41.925312870 +0100
+@@ -54,8 +54,11 @@
+   GIT_TAG 57560472b4260dc41f457706bc45fc6ef0bc0f10 # v4.4.7
+   EXCLUDE_FROM_ALL)
+
+-FetchContent_MakeAvailable(pybind11 abseil-cpp pybind11_abseil cifpp dssp)
+-
++#FetchContent_MakeAvailable(pybind11 abseil-cpp pybind11_abseil cifpp dssp)
++find_package(pybind11 CONFIG REQUIRED)
++find_package(absl CONFIG REQUIRED)
++find_package(pybind11_abseil CONFIG REQUIRED)
++find_package(dssp CONFIG REQUIRED)
+ find_package(
+   Python3
+   COMPONENTS Interpreter Development NumPy

easybuild/easyconfigs/a/AlphaFold3/AlphaFold3-3.0.1-foss-2024a-CUDA-12.6.0.eb

@@ -0,0 +1,119 @@
+# Thomas Hoffmann, EMBL Heidelberg, structures-it@embl.de, 2025/02
+#
+# DRAFT!
+easyblock = 'PythonBundle'
+
+name = 'AlphaFold3'
+version = '3.0.1'
+versionsuffix = '-CUDA-%(cudaver)s'
+
+homepage = 'https://deepmind.google/technologies/alphafold'
+description = """This package provides an implementation of the inference pipeline of AlphaFold3.
+See below for how to access the model parameters. You may only use AlphaFold3
+model parameters if received directly from Google. Use is subject to these
+terms of use:
+https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md"""
+
+toolchain = {'name': 'foss', 'version': '2024a'}
+
+github_account = 'google-deepmind'
+
+builddependencies = [
+    ('scikit-build', '0.17.6'),
+    ('poetry', '1.8.3'),
+    ('Python-bundle-PyPI', '2024.06'),
+    ('scikit-build-core', '0.10.6'),
+    ('pybind11', '2.13.6'),
+    ('pybind11_abseil', '202402.0'),
+]
+
+dependencies = [
+    ('Python', '3.12.3'),
+    ('CUDA', '12.6.0', '', SYSTEM),
+    ('tqdm', '4.66.5'),
+    ('jax', '0.4.35', versionsuffix),
+    ('dm-tree', '0.1.9'),
+    ('dm-haiku', '0.0.13', versionsuffix),
+    ('RDKit', '2024.09.5'),
+    ('jax-triton', '0.2.0', versionsuffix),
+    ('jaxtyping', '0.2.38', versionsuffix),
+    ('dssp', '4.4.10'),
+    ('Abseil', '20240722.0'),
+    ('HMMER', '3.4')
+]
+exts_list = [
+    ('zstandard', '0.23.0', {
+        'checksums': ['b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09'],
+    }),
+    ('typeguard', '2.13.3', {
+        'checksums': ['00edaa8da3a133674796cf5ea87d9f4b4c367d77476e185e80251cc13dfbb8c4'],
+    }),
+    (name, version, {
+        'patches': [
+            '%(name)s-3.0.0.20250109_relax_requirements.patch',
+            '%(name)s-3.0.0_disable_fetch.patch',
+            '%(name)s-3.0.0.20250109_data_path.patch',
+        ],
+        'postinstallcmds': [
+            'cp run_alphafold.py %(installdir)s/bin',
+            'chmod +x %(installdir)s/bin/run_alphafold.py',
+            '%(installdir)s/bin/build_data',
+        ],
+        'source_urls': [GITHUB_LOWER_SOURCE],
+        'sources': [{'download_filename': 'v%(version)s.tar.gz', 'filename': SOURCE_TAR_GZ}],
+        'checksums': [
+            {'AlphaFold3-3.0.1.tar.gz':
+             'cfe1ae5f6e5c0d7e64570255d296a440955beaf4a8b9f05aaf142f36478a4503'},
+            {'AlphaFold3-3.0.0.20250109_relax_requirements.patch':
+             '4c96c6d35542953d7dbae5cc43e0273c653430c472ffd4f5444d9096d83f00a2'},
+            {'AlphaFold3-3.0.0_disable_fetch.patch':
+             'e26da6d3a244ebf4185b7c8833326a5011a8af6f449eb6535c4147416336dd27'},
+            {'AlphaFold3-3.0.0.20250109_data_path.patch':
+             '00a71f24f430429bb8d432ee67253d29858c30fc7925e174fcb7e862285c8528'},
+        ],
+        'testinstall': True,
+        'runtest': 'python run_alphafold_test.py',
+    }),
+]
+
+fix_python_shebang_for = ['bin/run_alphafold.py']
+
+sanity_check_paths = {
+    'files': ['bin/run_alphafold.py'],
+    'dirs': ['lib'],
+}
+sanity_check_commands = [
+    ('run_alphafold.py --help|grep "AlphaFold 3 structure prediction script"')
+]
+
+modextravars = {
+    'DB_DIR': '/scratch/AlphaFold_DBs/%s/' % version,   # adapt
+    # acc. to:
+    # https://github.com/google-deepmind/alphafold3/
+    #  blob/main/docs/performance.md#compilation-time-workaround-with-xla-flags
+    'XLA_FLAGS': '--xla_gpu_enable_triton_gemm=false ',
+    # https://github.com/google-deepmind/alphafold3/blob/main/docs/performance.md#cuda-capability-7x-gpus:
+    # 'XLA_FLAGS': '"--xla_gpu_enable_triton_gemm=false --xla_disable_hlo_passes=custom-kernel-fusion-rewriter"',
+    # Unified memory:
+    'XLA_PYTHON_CLIENT_PREALLOCATE': 'false',
+    'TF_FORCE_UNIFIED_MEMORY': 'true',
+    'XLA_CLIENT_MEM_FRACTION': '3.2',
+}
+
+modluafooter = """
+setenv("AF3_MODEL_DIR", os.getenv("HOME") .. "/.alphafold3/models");
+eb_tmpdir = os.getenv("TMPDIR")
+if not eb_tmpdir then eb_tmpdir='/tmp' end
+setenv("TRITON_HOME", eb_tmpdir .. "/" .. os.getenv("USER") .. "/triton_home")
+"""
+
+modtclfooter = """
+setenv AF3_MODEL_DIR $::env(HOME)/.alphafold3/models
+if { [info exists ::env(TMPDIR)] } {
+  setenv AF3_MODEL_DIR $::env(TMPDIR)/$::env(USER)/triton_home
+} else {
+  setenv AF3_MODEL_DIR /tmp/$::env(USER)/triton_home
+}
+"""
+
+moduleclass = 'bio'

easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.9-gfbf-2024a.eb

@@ -0,0 +1,33 @@
+easyblock = 'PythonPackage'
+
+name = 'dm-tree'
+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': 'gfbf', 'version': '2024a'}
+
+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 = [
+    ('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'),
+]
+
+options = {'modulename': 'tree'}
+
+moduleclass = 'lib'