easyblock = 'PythonBundle'

name = 'TensorFlow'
version = '2.10.0'

homepage = 'https://www.tensorflow.org/'
description = "An open-source software library for Machine Intelligence"

toolchain = {'name': 'fosscuda', 'version': '2020b'}
toolchainopts = {'pic': True}

builddependencies = [
    ('Bazel', '3.7.2'),
    ('protobuf', '3.14.0'),
    # git 2.x required, see also https://github.com/tensorflow/tensorflow/issues/29053
    ('git', '2.28.0', '-nodocs'),
    ('pybind11', '2.6.0'),
    ('pkgconfig', '1.5.1', '-python'),  # For h5py
    ('UnZip', '6.0'),
]
dependencies = [
    ('cuDNN', '8.0.4.30', '-CUDA-%(cudaver)s', True),
    ('NCCL', '2.8.3', '-CUDA-%(cudaver)s'),
    ('Python', '3.8.6'),
    ('h5py', '3.1.0'),
    ('cURL', '7.72.0'),
    ('double-conversion', '3.1.5'),
    ('flatbuffers', '1.12.0'),
    ('giflib', '5.2.1'),
    ('hwloc', '2.2.0'),
    ('ICU', '67.1'),
    ('JsonCpp', '1.9.4'),
    ('libjpeg-turbo', '2.0.5'),
    ('LMDB', '0.9.24'),
    ('NASM', '2.15.05'),
    ('nsync', '1.24.0'),
    ('SQLite', '3.33.0'),
    ('PCRE', '8.44'),
    ('protobuf-python', '3.14.0'),
    ('flatbuffers-python', '1.12'),
    ('typing-extensions', '3.7.4.3'),
    ('libpng', '1.6.37'),
    ('snappy', '1.1.8'),
    ('zlib', '1.2.11'),
]

use_pip = True
sanity_pip_check = True
# Dependencies created and updated using findPythonDeps.sh:
# https://gist.github.com/Flamefire/49426e502cd8983757bd01a08a10ae0d
exts_list = [
    ('Markdown', '3.3.4', {
        'checksums': ['31b5b491868dcc87d6c24b7e3d19a0d730d59d3e46f4eea6430a321bed387a49'],
    }),
    ('pyasn1-modules', '0.2.8', {
        'checksums': ['905f84c712230b2c592c19470d3ca8d552de726050d1d1716282a1f6146be65e'],
    }),
    ('rsa', '4.7.2', {
        'checksums': ['9d689e6ca1b3038bc82bf8d23e944b6b6037bc02301a574935b2dd946e0353b9'],
    }),
    ('cachetools', '4.2.2', {
        'checksums': ['61b5ed1e22a0924aed1d23b478f37e8d52549ff8a961de2909c69bf950020cff'],
    }),
    ('google-auth', '1.30.0', {
        'modulename': 'google.auth',
        'checksums': ['9ad25fba07f46a628ad4d0ca09f38dcb262830df2ac95b217f9b0129c9e42206'],
    }),
    ('oauthlib', '3.1.0', {
        'checksums': ['bee41cc35fcca6e988463cacc3bcb8a96224f470ca547e697b604cc697b2f889'],
    }),
    ('requests-oauthlib', '1.3.0', {
        'checksums': ['b4261601a71fd721a8bd6d7aa1cc1d6a8a93b4a9f5e96626f8e4d91e8beeaa6a'],
    }),
    ('google-auth-oauthlib', '0.4.4', {
        'checksums': ['09832c6e75032f93818edf1affe4746121d640c625a5bef9b5c96af676e98eee'],
    }),
    ('Werkzeug', '2.0.0', {
        'checksums': ['3389bbfe6d40c6dd25e6d3f974155163c8b3de5bbda6a89342d4ab93fae80ba0'],
    }),
    ('absl-py', '0.12.0', {
        'modulename': 'absl',
        'checksums': ['b44f68984a5ceb2607d135a615999b93924c771238a63920d17d3387b0d229d5'],
    }),
    ('astunparse', '1.6.3', {
        'checksums': ['5ad93a8456f0d084c3456d059fd9a92cce667963232cbf763eac3bc5b7940872'],
    }),
    ('grpcio', '1.34.1', {
        'modulename': 'grpc',
        'checksums': ['1c746a3cd8a830d8d916a9d0476a786aaa98c5cc2a096344af2be955e439f8ac'],
        'preinstallopts': "export GRPC_PYTHON_BUILD_EXT_COMPILER_JOBS=%(parallel)s && ",
    }),
    ('gviz-api', '1.9.0', {
        'source_tmpl': 'gviz_api-%(version)s.tar.gz',
        'checksums': ['43d13ccc21834d0501b33a291ef3265e933dbb4bbdca3d34b1ed0a048c0ef640'],
    }),
    ('tensorboard_data_server', '0.6.1', {
        'source_tmpl': SOURCE_PY3_WHL,  # Requires Rust to build, take the dummy for now
        'checksums': ['809fe9887682d35c1f7d1f54f0f40f98bb1f771b14265b453ca051e2ce58fca7'],
    }),
    ('tensorboard', version, {
        'source_tmpl': SOURCE_PY3_WHL,
        'checksums': ['e167460085b6528956b33bab1c970c989cdce47a6616273880733f5e7bde452e'],
    }),
    ('tensorboard_plugin_wit', '1.8.0', {
        'source_tmpl': SOURCE_PY3_WHL,
        'checksums': ['2a80d1c551d741e99b2f197bb915d8a133e24adb8da1732b840041860f91183a'],
    }),
    ('tensorboard_plugin_profile', '2.4.0', {
        'checksums': ['dfbf254ee960440e3b2518324f876a6d6704c60b936887d99214fa36988a206a'],
    }),
    ('google-pasta', '0.2.0', {
        'modulename': 'pasta',
        'checksums': ['c9f2c8dfc8f96d0d5808299920721be30c9eec37f2389f28904f454565c8a16e'],
    }),
    ('termcolor', '1.1.0', {
        'checksums': ['1d6d69ce66211143803fbc56652b41d73b4a400a2891d7bf7a1cdf4c02de613b'],
    }),
    ('tensorflow_estimator', version, {
        'source_tmpl': SOURCE_WHL,
        'checksums': ['d1fe76dee8b1dcab865d807a0246da0a9c4a635b1eba6e9545bf216c3aad6955'],
    }),
    ('astor', '0.8.1', {
        'checksums': ['6a6effda93f4e1ce9f618779b2dd1d9d84f1e32812c23a29b3fff6fd7f63fa5e'],
    }),
    ('gast', '0.4.0', {
        'checksums': ['40feb7b8b8434785585ab224d1568b857edb18297e5a3047f1ba012bc83b42c1'],
    }),
    ('opt_einsum', '3.3.0', {
        'checksums': ['59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549'],
    }),
    ('wrapt', '1.12.1', {
        'checksums': ['b62ffa81fb85f4332a4f609cab4ac40709470da05643a082ec1eb88e6d9b97d7'],
    }),
    ('Keras_Preprocessing', '1.1.2', {
        'checksums': ['add82567c50c8bc648c14195bf544a5ce7c1f76761536956c3d2978970179ef3'],
    }),
    ('dill', '0.3.4', {
        'source_tmpl': '%(name)s-%(version)s.zip',
        'checksums': ['9f9734205146b2b353ab3fec9af0070237b6ddae78452af83d2fca84d739e675'],
    }),
    ('tblib', '1.7.0', {
        'checksums': ['059bd77306ea7b419d4f76016aef6d7027cc8a0785579b5aad198803435f882c'],
    }),
    ('portpicker', '1.3.1', {
        'checksums': ['d2cdc776873635ed421315c4d22e63280042456bbfa07397817e687b142b9667'],
    }),
    ('keras_nightly', '2.5.0.dev2021032900', {
        'modulename': 'keras',
        'source_tmpl': SOURCE_WHL,
        'checksums': ['6ba70f738f4008222de7e7fdd5b2b18c48c49b897a9fca54c844854e25964011'],
    }),
    (name, version, {
        'source_tmpl': 'v%(version)s.tar.gz',
        'source_urls': ['https://github.com/tensorflow/tensorflow/archive/'],
        'patches': [
            'TensorFlow-2.1.0_fix-cuda-build.patch',
            'TensorFlow-2.4.0_add-ldl.patch',
            'TensorFlow-2.4.0_dont-use-var-lock.patch',
            'TensorFlow-2.4.1_fix-min-priority-test.patch',
            'TensorFlow-2.5.0_add-default-shell-env.patch',
            'TensorFlow-2.5.0_add-protobuf-deps.patch',
            'TensorFlow-2.5.0_add-support-for-large-core-systems.patch',
            'TensorFlow-2.5.0_disable-avx512-extensions.patch',
            'TensorFlow-2.5.0-fix-alias-violation-in-absl.patch',
            'TensorFlow-2.5.0_fix-alignment-in-matmul-test.patch',
            'TensorFlow-2.5.0_fix-arm-vector-intrinsics.patch',
            'TensorFlow-2.5.0_fix-crash-on-shutdown.patch',
            'TensorFlow-2.5.0_remove-duplicate-gpu-tests.patch',
        ],
        'checksums': [
            '233875ea27fc357f6b714b2a0de5f6ff124b50c1ee9b3b41f9e726e9e677b86c',  # v2.5.0.tar.gz
            '78c20aeaa7784b8ceb46238a81e8c2461137d28e0b576deeba8357d23fbe1f5a',  # TensorFlow-2.1.0_fix-cuda-build.patch
            '917ee7282e782e48673596d8917c3207e60e0851bb9acf230a2a439b067af2e3',  # TensorFlow-2.4.0_add-ldl.patch
            # TensorFlow-2.4.0_dont-use-var-lock.patch
            'b14f2493fd2edf79abd1c4f2dde6c98a3e7d5cb9c25ab9386df874d5f072d6b5',
            # TensorFlow-2.4.1_fix-min-priority-test.patch
            '389febce9a0612fd457daf4cb18c02f77fc7311bacae0963602a3198d9f2737f',
            # TensorFlow-2.5.0_add-default-shell-env.patch
            '09b0c5d4ff04f56a9657875471ed78001d4201cac795aeff62019d582115b468',
            # TensorFlow-2.5.0_add-protobuf-deps.patch
            '2aa79b89cff13e81f83e385761917d5d6dbdffd6b0366d90580761b958f14363',
            # TensorFlow-2.5.0_add-support-for-large-core-systems.patch
            '915f3477d6407fafd48269fe1e684a05ce361d9b9b85e58686682df87760f636',
            # TensorFlow-2.5.0_disable-avx512-extensions.patch
            '3655ce24c97569ac9738c07cac85347ba6f5c815ada95b19b606ffa46d4dda03',
            # TensorFlow-2.5.0-fix-alias-violation-in-absl.patch
            '12454fda3330fb45cd380377e283f04488b40e0b8ae7378e786ddf731a581f75',
            # TensorFlow-2.5.0_fix-alignment-in-matmul-test.patch
            '6a4d6cbf45a622b8a2c3ea0b1c0171f01f595684d9c57d415bb39b1b27e1180f',
            # TensorFlow-2.5.0_fix-arm-vector-intrinsics.patch
            '6abfadc0f67ff3b510d70430843201cb46d7bd65db045ec9b482af70e0c8c0c8',
            # TensorFlow-2.5.0_fix-crash-on-shutdown.patch
            '578c7493221ebd3dc25ca43d63a72cbb28fdf4112b1e2baa7390f25781bd78fd',
            # TensorFlow-2.5.0_remove-duplicate-gpu-tests.patch
            'b940d438e036faac24453bff2cf1834c5e1359e87e84d1f1999fa7a30b278fec',
        ],
        'test_script': 'TensorFlow-2.x_mnist-test.py',
        'test_tag_filters_cpu': '-gpu,-tpu,-no_cuda_on_cpu_tap,-no_pip,-no_oss,-oss_serial,-benchmark-test,-v1only',
        'test_tag_filters_gpu': ('gpu,-no_gpu,-nogpu,-gpu_cupti,-no_cuda11,-no_pip,-no_oss,-oss_serial,'
                                 '-benchmark-test,-v1only'),
        'test_targets': [
            '//tensorflow/core/...',
            '-//tensorflow/core:example_java_proto',
            '-//tensorflow/core/example:example_protos_closure',
            '//tensorflow/cc/...',
            '//tensorflow/c/...',
            '//tensorflow/python/...',
            # Fails on some nodes but C API isn't installed anyway
            '-//tensorflow/c/eager:c_api_test_gpu',
            '-//tensorflow/c/eager:c_api_distributed_test',
            '-//tensorflow/c/eager:c_api_distributed_test_gpu',
            # Race condition with port picker: https://github.com/tensorflow/tensorflow/issues/46602
            '-//tensorflow/c/eager:c_api_cluster_test_gpu',
            '-//tensorflow/c/eager:c_api_remote_function_test_gpu',
            '-//tensorflow/c/eager:c_api_remote_test_gpu',
            # Fails to open its own test.xml(?)
            '-//tensorflow/core/common_runtime:collective_param_resolver_local_test',
            # Fails on non-AVX-512 systems: https://github.com/tensorflow/tensorflow/issues/46532
            '-//tensorflow/core/common_runtime:mkl_layout_pass_test',
            '-//tensorflow/core/kernels/mkl:mkl_fused_ops_test',
            # Fails on AMD EPYC systems: https://github.com/tensorflow/tensorflow/issues/52151
            '-//tensorflow/core/kernels/mkl:mkl_fused_batch_norm_op_test',
        ],
        'testopts': "--test_timeout=3600 --test_size_filters=small",
        'testopts_gpu': "--test_timeout=3600 --test_size_filters=small " +
                        "--run_under=//tensorflow/tools/ci_build/gpu_build:parallel_gpu_execute",
        'with_xla': True,
    }),
]

cuda_compute_capabilities = ['8.0']

moduleclass = 'lib'