easyblock = 'PythonBundle'

name = 'TensorFlow'
version = '2.7.1'
versionsuffix = '-CUDA-%(cudaver)s'

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

toolchain = {'name': 'foss', 'version': '2022a'}
toolchainopts = {'pic': True}

builddependencies = [
    ('Bazel', '3.7.2'),
    ('protobuf', '3.17.3'),
    # git 2.x required, see also https://github.com/tensorflow/tensorflow/issues/29053
    ('git', '2.33.1', '-nodocs'),
    ('pybind11', '2.7.1'),
    ('UnZip', '6.0'),
    ('LLVM', '12.0.1'),  # for debugging with llvm-symbolizer, to be removed
]
dependencies = [
    ('CUDA', '11.4.1', '', SYSTEM),
    ('cuDNN', '8.2.2.26', versionsuffix, SYSTEM),
    ('NCCL', '2.10.3', versionsuffix),
    ('Python', '3.9.6'),
    ('h5py', '3.6.0'),
    ('cURL', '7.78.0'),
    ('double-conversion', '3.1.5'),
    ('flatbuffers', '2.0.0'),
    ('giflib', '5.2.1'),
    ('hwloc', '2.5.0'),
    ('ICU', '69.1'),
    ('JsonCpp', '1.9.4'),
    ('libjpeg-turbo', '2.0.6'),
    ('LMDB', '0.9.29'),
    ('NASM', '2.15.05'),
    ('nsync', '1.24.0'),
    ('SQLite', '3.36'),
    ('protobuf-python', '3.17.3'),
    ('flatbuffers-python', '2.0'),
    ('libpng', '1.6.37'),
    ('snappy', '1.1.9'),
    ('zlib', '1.2.11'),
    ('networkx', '2.6.3'),  # required for pythran
]

use_pip = True
sanity_pip_check = True

# Dependencies created and updated using findPythonDeps.sh:
# https://gist.github.com/Flamefire/49426e502cd8983757bd01a08a10ae0d
exts_list = [
    # gast deps beniget and pythran require different version from Scipy-bundle
    ('beniget', '0.3.0', {
        'checksums': ['062c893be9cdf87c3144fb15041cce4d81c67107c1591952cd45fdce789a0ff1'],
    }),
    ('pythran', '0.9.11', {
        'checksums': ['a317f91e2aade9f6550dc3bf40b5caeb45b7e012daf27e2b3e4ad928edb01667'],
    }),
    ('wrapt', '1.13.3', {
        'checksums': ['1fea9cd438686e6682271d36f3481a9f3636195578bab9ca3382e2f5f01fc185'],
    }),
    ('termcolor', '1.1.0', {
        'checksums': ['1d6d69ce66211143803fbc56652b41d73b4a400a2891d7bf7a1cdf4c02de613b'],
    }),
    ('tensorflow_estimator', '2.7.0', {
        'source_tmpl': '%(name)s-%(version)s-py2.py3-none-any.whl',
        'checksums': ['325b5a224864379242b7b76c6987ca544239be82579d33e68ec7c2bda57abc9d'],
    }),
    ('Werkzeug', '2.0.3', {
        'checksums': ['b863f8ff057c522164b6067c9e28b041161b4be5ba4d0daceeaa50a163822d3c'],
    }),
    ('tensorboard_plugin_wit', '1.8.1', {
        'source_tmpl': '%(name)s-%(version)s-py3-none-any.whl',
        'checksums': ['ff26bdd583d155aa951ee3b152b3d0cffae8005dc697f72b44a8e8c2a77a8cbe'],
    }),
    ('tensorboard_data_server', '0.6.1', {
        'source_tmpl': '%(name)s-%(version)s-py3-none-any.whl',
        'checksums': ['809fe9887682d35c1f7d1f54f0f40f98bb1f771b14265b453ca051e2ce58fca7'],
    }),
    ('Markdown', '3.3.6', {
        'checksums': ['76df8ae32294ec39dcf89340382882dfa12975f87f45c3ed1ecdb1e8cefc7006'],
    }),
    ('grpcio', '1.43.0', {
        'modulename': 'grpc',
        'preinstallopts': "export GRPC_PYTHON_BUILD_EXT_COMPILER_JOBS=%(parallel)s && ",
        'checksums': ['735d9a437c262ab039d02defddcb9f8f545d7009ae61c0114e19dda3843febe5'],
    }),
    ('oauthlib', '3.2.0', {
        'checksums': ['23a8208d75b902797ea29fd31fa80a15ed9dc2c6c16fe73f5d346f83f6fa27a2'],
    }),
    ('requests-oauthlib', '1.3.1', {
        'checksums': ['75beac4a47881eeb94d5ea5d6ad31ef88856affe2332b9aafb52c6452ccf0d7a'],
    }),
    ('rsa', '4.8', {
        'checksums': ['5c6bd9dc7a543b7fe4304a631f8a8a3b674e2bbfc49c2ae96200cdbe55df6b17'],
    }),
    ('pyasn1-modules', '0.2.8', {
        'checksums': ['905f84c712230b2c592c19470d3ca8d552de726050d1d1716282a1f6146be65e'],
    }),
    ('cachetools', '5.0.0', {
        'checksums': ['486471dfa8799eb7ec503a8059e263db000cdda20075ce5e48903087f79d5fd6'],
    }),
    ('google-auth', '2.6.0', {
        'modulename': 'google.auth',
        'checksums': ['ad160fc1ea8f19e331a16a14a79f3d643d813a69534ba9611d2c80dc10439dad'],
    }),
    ('google-auth-oauthlib', '0.4.6', {
        'checksums': ['a90a072f6993f2c327067bf65270046384cda5a8ecb20b94ea9a687f1f233a7a'],
    }),
    ('tensorboard', '2.8.0', {
        'source_tmpl': '%(name)s-%(version)s-py3-none-any.whl',
        'checksums': ['65a338e4424e9079f2604923bdbe301792adce2ace1be68da6b3ddf005170def'],
    }),
    ('opt_einsum', '3.3.0', {
        'checksums': ['59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549'],
    }),
    ('Keras_Preprocessing', '1.1.2', {
        'checksums': ['add82567c50c8bc648c14195bf544a5ce7c1f76761536956c3d2978970179ef3'],
    }),
    ('keras', '2.7.0', {
        'source_tmpl': '%(name)s-%(version)s-py2.py3-none-any.whl',
        'checksums': ['0c33ae1f728064ca0d35dfba999e9c316f03623bf5688c82fb83cc74a80ea248'],
    }),
    ('google-pasta', '0.2.0', {
        'modulename': 'pasta',
        'checksums': ['c9f2c8dfc8f96d0d5808299920721be30c9eec37f2389f28904f454565c8a16e'],
    }),
    ('gast', '0.4.0', {
        'checksums': ['40feb7b8b8434785585ab224d1568b857edb18297e5a3047f1ba012bc83b42c1'],
    }),
    ('astunparse', '1.6.3', {
        'checksums': ['5ad93a8456f0d084c3456d059fd9a92cce667963232cbf763eac3bc5b7940872'],
    }),
    ('gviz-api', '1.10.0', {
        'source_tmpl': 'gviz_api-%(version)s.tar.gz',
        'checksums': ['846692dd8cc73224fc31b18e41589bd934e1cc05090c6576af4b4b26c2e71b90'],
    }),
    ('tensorboard_plugin_profile', '2.5.0', {
        'checksums': ['f832698d87a773b9a017fc4dd5cf598a1ff942ccfbf3392c83fe12c949ab9f52'],
    }),
    ('astor', '0.8.1', {
        'checksums': ['6a6effda93f4e1ce9f618779b2dd1d9d84f1e32812c23a29b3fff6fd7f63fa5e'],
    }),
    ('dill', '0.3.4', {
        'source_tmpl': '%(name)s-%(version)s.zip',
        'checksums': ['9f9734205146b2b353ab3fec9af0070237b6ddae78452af83d2fca84d739e675'],
    }),
    ('tblib', '1.7.0', {
        'checksums': ['059bd77306ea7b419d4f76016aef6d7027cc8a0785579b5aad198803435f882c'],
    }),
    ('portpicker', '1.5.0', {
        'checksums': ['e13b148008adeb2793cf8b55bcd20fdcec4f763f2d3bf3c45f5e5e5d1df7d228'],
    }),
    ('absl-py', '0.13.0', {
        'modulename': 'absl',
        'checksums': ['6953272383486044699fd0e9f00aad167a27e08ce19aae66c6c4b10e7e767793'],
    }),
    (name, version, {
        '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.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-arm-vector-intrinsics.patch',
            'TensorFlow-2.5.0_fix-crash-on-shutdown.patch',
            'TensorFlow-2.7.1_fix_protobuf_error_message.patch',
            'TensorFlow-2.7.1_remove-duplicate-gpu-tests.patch',
            'TensorFlow-2.7.1_remove-io-gcs-dependency.patch',
            'TensorFlow-2.7.1_remove-libclang-dependency.patch',
            'TensorFlow-2.7.1_fix_cpu_count.patch',
        ],
        'source_tmpl': 'v%(version)s.tar.gz',
        'source_urls': ['https://github.com/tensorflow/tensorflow/archive/'],
        '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',
            # All tests in this directory fail with segfault (TensorFlow Graph IR)
            '-//tensorflow/core/ir/importexport/tests/roundtrip/...',
            # Fails on e.g. AMD EPYC and known to be flaky on Mac: https://github.com/tensorflow/tensorflow/issues/56717
            '-//tensorflow/python/data/experimental/kernel_tests/service:fault_tolerance_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,
        'checksums': [
            'abebe2cf5ca379e18071693ca5f45b88ade941b16258a21cc1f12d77d5387a21',  # v2.7.1.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.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-arm-vector-intrinsics.patch
            '6abfadc0f67ff3b510d70430843201cb46d7bd65db045ec9b482af70e0c8c0c8',
            # TensorFlow-2.5.0_fix-crash-on-shutdown.patch
            '578c7493221ebd3dc25ca43d63a72cbb28fdf4112b1e2baa7390f25781bd78fd',
            # TensorFlow-2.7.1_fix_protobuf_error_message.patch
            '301ce21845987dc7868624880c0d85e60aa5491c5a0bae0f6ebc8db5727960a5',
            # TensorFlow-2.7.1_remove-duplicate-gpu-tests.patch
            'f78526a34d85f4dda59b160b576aa3a3126db7073a58c4e3b1424923f6a21483',
            # TensorFlow-2.7.1_remove-io-gcs-dependency.patch
            '01c439749938f9fc23fe8de17592460aa46e201575a3f9465bfb6b39ca50ef20',
            # TensorFlow-2.7.1_remove-libclang-dependency.patc
            'd202058c0e8a7cdef3157fe12327e23c7291ce2724a15653df0aae3b40c4dadc',
            # TensorFlow-2.7.1_fix_cpu_count.patch
            '5427a4cff0afc2fe5b24776ae9ca3616c56a79c1fde0025b37bec24837bb0698',
        ],
    }),
]

moduleclass = 'lib'