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

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

@@ -0,0 +1,69 @@
+# Thomas Hoffmann, EMBL Heidelberg, [email protected], 2024/12
+#
+# DRAFT!
+easyblock = 'PythonBundle'
+
+name = 'AlphaFold3'
+version = '3.0.0'
+versionsuffix = '-CUDA-%(cudaver)s'
+homepage = 'https://deepmind.google/technologies/alphafold'
+description = """This package provides an implementation of the inference pipeline of AlphaFold  
+3. See below for how to access the model parameters. You may only use AlphaFold 
+3 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'
+use_pip = True
+
+builddependencies = [
+    ('scikit-build', '0.17.6'),
+    ('poetry', '1.8.3'),
+    ('Python-bundle-PyPI', '2024.06'),
+    ('scikit-build-core', '0.10.6'),
+]
+
+dependencies = [
+    ('Python', '3.12.3'),
+    ('CUDA', '12.6.0', '', SYSTEM),
+    ('tqdm', '4.67.0'),
+    ('jax', '0.4.34', versionsuffix),
+    ('dm-tree', '0.1.8'),
+    ('dm-haiku', '0.0.13', versionsuffix),
+    ('RDKit', '2024.03.5'),
+    ('jax-triton', '0.2.0', versionsuffix),
+    ('jaxtyping', '0.2.36', versionsuffix),
+    # TODO
+]
+
+exts_list = [
+    ('zstandard', '0.23.0', {
+        'checksums': ['b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09'],
+    }),
+    ('chex', '0.1.87', {
+        'checksums': ['0096d89cc8d898bb521ef4bfbf5c24549022b0e5b301f529ab57238896fe6c5d'],
+    }),
+    (name, version, {
+        'patches': ['%(name)s-%(version)s_relax_requirements.patch'],
+        'source_urls': ['https://github.com/%(github_account)s/%(name)s/archive'],
+        'sources': ['v%(version)s.tar.gz'],
+        'checksums': [
+            {'v3.0.0.tar.gz': 'f4fde43daf4cdf25e4da712f440b8aa8bd8798e2be936a254a562f6cb6eb32df'},
+            {'AlphaFold3-3.0.0_relax_requirements.patch':
+             'ea2c11638dfc60e51753301412a1867d17b4f772bdfe6a2c834f8bfa8310646d'},
+        ],
+    }),
+]
+
+sanity_pip_check = True
+
+sanity_check_paths = {
+    'files': [],
+    'dirs': ['lib'],
+}
+
+# sanity_check_cmds = []  # TODO
+
+moduleclass = 'bio'

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

@@ -0,0 +1,19 @@
+diff -ru alphafold3-3.0.0/pyproject.toml alphafold3-3.0.0_relax_requirements/pyproject.toml
+--- alphafold3-3.0.0/pyproject.toml	2024-11-11 11:42:47.000000000 +0100
++++ alphafold3-3.0.0_relax_requirements/pyproject.toml	2024-12-19 17:09:46.784161491 +0100
+@@ -19,12 +19,12 @@
+     "chex",
+     "dm-haiku==0.0.13",
+     "dm-tree",
+-    "jax==0.4.34",
+-    "jax[cuda12]==0.4.34",
++    "jax>=0.4.34",
++    #"jax[cuda12]==0.4.34",
+     "jax-triton==0.2.0",
+     "jaxtyping",
+     "numpy",
+-    "rdkit==2024.3.5",
++    "rdkit",
+     "triton==3.1.0",
+     "tqdm",
+     "zstandard",

easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.8-GCC-13.3.0.eb

@@ -0,0 +1,35 @@
+easyblock = 'PythonPackage'
+
+name = 'dm-tree'
+version = '0.1.8'
+
+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': 'GCC', 'version': '13.3.0'}
+
+sources = [SOURCELOWER_TAR_GZ]
+patches = ['dm-tree-0.1.8_pybind11.patch']
+checksums = [
+    {'dm-tree-0.1.8.tar.gz': '0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430'},
+    {'dm-tree-0.1.8_pybind11.patch': '1ed433eda86d46333d735dfea9eecf9d007ba0ca8b1bc610c6f02cd478b8d131'},
+]
+
+builddependencies = [
+    ('pybind11', '2.12.0'),
+]
+dependencies = [
+    ('Python', '3.12.3'),
+    ('Abseil', '20240722.0'),
+]
+
+download_dep_fail = True
+
+use_pip = True
+sanity_pip_check = True
+
+options = {'modulename': 'tree'}
+
+moduleclass = 'lib'

easybuild/easyconfigs/d/dm-tree/dm-tree-0.1.8_pybind11.patch

@@ -0,0 +1,140 @@
+# Taken from: https://data.gpo.zugaina.org/benzene-overlay/dev-python/dm-tree/files/dm-tree-0.1.8-Simplify-setup.py-by-using-pybind11.setup_helpers.Py.patch
+From: Michael Fladischer <[email protected]>
+Date: Mon, 31 Jul 2023 21:13:42 +0000
+Subject: Simplify setup.py by using pybind11.setup_helpers.Pybind11Extension.
+
+---
+ setup.py | 103 ++++++++++++---------------------------------------------------
+ 1 file changed, 19 insertions(+), 84 deletions(-)
+
+diff --git a/setup.py b/setup.py
+index d1862a1..8381bff 100644
+--- a/setup.py
++++ b/setup.py
+@@ -15,14 +15,10 @@
+ """Setup for pip package."""
+
+ import os
+-import platform
+-import shutil
+-import subprocess
+-import sys
+-import sysconfig
+
+-import setuptools
+-from setuptools.command import build_ext
++from setuptools import find_packages, setup
++from pybind11.setup_helpers import Pybind11Extension, build_ext
++
+
+ here = os.path.dirname(os.path.abspath(__file__))
+
+@@ -48,80 +44,7 @@ def _parse_requirements(path):
+     ]
+
+
+-class CMakeExtension(setuptools.Extension):
+-  """An extension with no sources.
+-
+-  We do not want distutils to handle any of the compilation (instead we rely
+-  on CMake), so we always pass an empty list to the constructor.
+-  """
+-
+-  def __init__(self, name, source_dir=''):
+-    super().__init__(name, sources=[])
+-    self.source_dir = os.path.abspath(source_dir)
+-
+-
+-class BuildCMakeExtension(build_ext.build_ext):
+-  """Our custom build_ext command.
+-
+-  Uses CMake to build extensions instead of a bare compiler (e.g. gcc, clang).
+-  """
+-
+-  def run(self):
+-    self._check_build_environment()
+-    for ext in self.extensions:
+-      self.build_extension(ext)
+-
+-  def _check_build_environment(self):
+-    """Check for required build tools: CMake, C++ compiler, and python dev."""
+-    try:
+-      subprocess.check_call(['cmake', '--version'])
+-    except OSError as e:
+-      ext_names = ', '.join(e.name for e in self.extensions)
+-      raise RuntimeError(
+-          f'CMake must be installed to build the following extensions: {ext_names}'
+-      ) from e
+-    print('Found CMake')
+-
+-  def build_extension(self, ext):
+-    extension_dir = os.path.abspath(
+-        os.path.dirname(self.get_ext_fullpath(ext.name)))
+-    build_cfg = 'Debug' if self.debug else 'Release'
+-    cmake_args = [
+-        f'-DPython3_ROOT_DIR={sys.prefix}',
+-        f'-DPython3_EXECUTABLE={sys.executable}',
+-        f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extension_dir}',
+-        f'-DCMAKE_BUILD_TYPE={build_cfg}'
+-    ]
+-    if platform.system() != 'Windows':
+-      cmake_args.extend([
+-          f'-DPython3_LIBRARY={sysconfig.get_paths()["stdlib"]}',
+-          f'-DPython3_INCLUDE_DIR={sysconfig.get_paths()["include"]}',
+-      ])
+-    if platform.system() == 'Darwin' and os.environ.get('ARCHFLAGS'):
+-      osx_archs = []
+-      if '-arch x86_64' in os.environ['ARCHFLAGS']:
+-        osx_archs.append('x86_64')
+-      if '-arch arm64' in os.environ['ARCHFLAGS']:
+-        osx_archs.append('arm64')
+-      cmake_args.append(f'-DCMAKE_OSX_ARCHITECTURES={";".join(osx_archs)}')
+-    os.makedirs(self.build_temp, exist_ok=True)
+-    subprocess.check_call(
+-        ['cmake', ext.source_dir] + cmake_args, cwd=self.build_temp)
+-    subprocess.check_call(
+-        ['cmake', '--build', '.', f'-j{os.cpu_count()}', '--config', build_cfg],
+-        cwd=self.build_temp)
+-
+-    # Force output to <extension_dir>/. Amends CMake multigenerator output paths
+-    # on Windows and avoids Debug/ and Release/ subdirs, which is CMake default.
+-    tree_dir = os.path.join(extension_dir, 'tree')  # pylint:disable=unreachable
+-    for cfg in ('Release', 'Debug'):
+-      cfg_dir = os.path.join(extension_dir, cfg)
+-      if os.path.isdir(cfg_dir):
+-        for f in os.listdir(cfg_dir):
+-          shutil.move(os.path.join(cfg_dir, f), tree_dir)
+-
+-
+-setuptools.setup(
++setup(
+     name='dm-tree',
+     version=_get_tree_version(),
+     url='https://github.com/deepmind/tree',
+@@ -131,11 +54,23 @@ setuptools.setup(
+     long_description=open(os.path.join(here, 'README.md')).read(),
+     long_description_content_type='text/markdown',
+     # Contained modules and scripts.
+-    packages=setuptools.find_packages(),
++    packages=find_packages(),
+     tests_require=_parse_requirements('requirements-test.txt'),
+     test_suite='tree',
+-    cmdclass=dict(build_ext=BuildCMakeExtension),
+-    ext_modules=[CMakeExtension('_tree', source_dir='tree')],
++    cmdclass={"build_ext": build_ext},
++    ext_modules=[
++        Pybind11Extension(
++            'tree._tree',
++            sources=['tree/tree.cc'],
++            libraries=[
++                'absl_int128',
++                'absl_raw_hash_set',
++                'absl_raw_logging_internal',
++                'absl_strings',
++                'absl_throw_delegate',
++            ]
++        )
++    ],
+     zip_safe=False,
+     # PyPI package information.
+     classifiers=[

easybuild/easyconfigs/j/jax-triton/jax-triton-0.2.0-foss-2024a-CUDA-12.6.0.eb

@@ -0,0 +1,36 @@
+# Thomas Hoffmann, EMBL Heidelberg, [email protected], 2024/12
+
+easyblock = 'PythonBundle'
+
+name = 'jax-triton'
+version = '0.2.0'
+versionsuffix = '-CUDA-%(cudaver)s'
+
+homepage = 'https://pypi.org/project/jax-triton'
+description = """The jax-triton repository contains integrations between JAX and Triton."""
+
+toolchain = {'name': 'foss', 'version': '2024a'}
+
+dependencies = [
+    ('Python', '3.12.3'),
+    ('CUDA', '12.6.0', '', SYSTEM),
+    ('jax', '0.4.34', versionsuffix),
+    ('Triton', '3.1.0', versionsuffix),
+]
+
+use_pip = True
+
+exts_list = [
+    ('jax_triton', version, {
+        'checksums': ['ed564a5ffb9e404557dc8d296e7eb30e501da5a4d7b03408acd0837c1c618c21'],
+    }),
+]
+
+sanity_pip_check = True
+
+sanity_check_paths = {
+    'files': [],
+    'dirs': ['lib'],
+}
+
+moduleclass = 'ai'

easybuild/easyconfigs/j/jaxtyping/jaxtyping-0.2.36-foss-2024a-CUDA-12.6.0.eb

@@ -0,0 +1,40 @@
+# Thomas Hoffmann, EMBL Heidelberg, [email protected], 2024/12
+
+easyblock = 'PythonBundle'
+
+name = 'jaxtyping'
+version = '0.2.36'
+versionsuffix = '-CUDA-%(cudaver)s'
+
+homepage = 'https://github.com/patrick-kidger/jaxtyping'
+description = """ Type annotations and runtime type-checking for: 
+1. shape and dtype of JAX arrays;  
+   (Now also supports PyTorch, NumPy, and TensorFlow!)
+2. PyTrees.  """
+toolchain = {'name': 'foss', 'version': '2024a'}
+
+builddependencies = [
+    ('poetry', '1.8.3')
+]
+dependencies = [
+    ('Python', '3.12.3'),
+    ('CUDA', '12.6.0', '', SYSTEM),
+    ('jax', '0.4.34', versionsuffix),
+]
+
+use_pip = True
+
+exts_list = [
+    (name, version, {
+        'checksums': ['781ac44a3cf8982063d7ee48b5008ccfad7b13793bf878eb3058d5319aa08f0f'],
+    }),
+]
+
+sanity_pip_check = True
+
+sanity_check_paths = {
+    'files': [],
+    'dirs': ['lib'],
+}
+
+moduleclass = 'ai'