# IT4Innovations
# LK 2023

easyblock = 'PythonBundle'

name = 'TensorFlow'
version = '2.4.1'

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

toolchain = {'name': 'foss', '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 = [
    ('Python', '3.8.6'),
    # Next 2 are for h5py
    ('SciPy-bundle', '2020.11'),
    ('HDF5', '1.10.7'),
    ('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'),
]

sanity_pip_check = True
use_pip = True


# Dependencies created and updated using findPythonDeps.sh:
# https://gist.github.com/Flamefire/49426e502cd8983757bd01a08a10ae0d
exts_list = [
    # TF 2.4 needs h5py 2.10, later versions allow 3.x
    ('h5py', '2.10.0', {
        'patches': ['h5py-2.10.0_avoid-mpi-init.patch'],
        'checksums': [
            '84412798925dc870ffd7107f045d7659e60f5d46d1c70c700375248bf6bf512d',  # h5py-2.10.0.tar.gz
            '6bacb71f5d9fbd7bd9a01018d7fe21b067a2317f33c4a7c21fde9cd404c1603f',  # h5py-2.10.0_avoid-mpi-init.patch
        ],
        'preinstallopts': 'HDF5_MPI=ON HDF5_DIR="$EBROOTHDF5" ',
    }),
    ('Markdown', '3.3.3', {
        'checksums': ['5d9f2b5ca24bc4c7a390d22323ca4bad200368612b5aaa7796babf971d2b2f18'],
    }),
    ('pyasn1-modules', '0.2.8', {
        'checksums': ['905f84c712230b2c592c19470d3ca8d552de726050d1d1716282a1f6146be65e'],
    }),
    ('rsa', '4.7', {
        'checksums': ['69805d6b69f56eb05b62daea3a7dbd7aa44324ad1306445e05da8060232d00f4'],
    }),
    ('cachetools', '4.2.0', {
        'checksums': ['3796e1de094f0eaca982441c92ce96c68c89cced4cd97721ab297ea4b16db90e'],
    }),
    ('google-auth', '1.24.0', {
        'modulename': 'google.auth',
        'checksums': ['0b0e026b412a0ad096e753907559e4bdb180d9ba9f68dd9036164db4fdc4ad2e'],
    }),
    ('oauthlib', '3.1.0', {
        'checksums': ['bee41cc35fcca6e988463cacc3bcb8a96224f470ca547e697b604cc697b2f889'],
    }),
    ('requests-oauthlib', '1.3.0', {
        'checksums': ['b4261601a71fd721a8bd6d7aa1cc1d6a8a93b4a9f5e96626f8e4d91e8beeaa6a'],
    }),
    ('google-auth-oauthlib', '0.4.2', {
        'checksums': ['65b65bc39ad8cab15039b35e5898455d3d66296d0584d96fe0e79d67d04c51d9'],
    }),
    ('Werkzeug', '1.0.1', {
        'checksums': ['6c80b1e5ad3665290ea39320b91e1be1e0d5f60652b964a3070216de83d2e47c'],
    }),
    ('absl-py', '0.10.0', {
        'modulename': 'absl',
        'checksums': ['b20f504a7871a580be5268a18fbad48af4203df5d33dbc9272426cb806245a45'],
    }),
    ('astunparse', '1.6.3', {
        'checksums': ['5ad93a8456f0d084c3456d059fd9a92cce667963232cbf763eac3bc5b7940872'],
    }),
    ('grpcio', '1.32.0', {
        'modulename': 'grpc',
        'checksums': ['01d3046fe980be25796d368f8fc5ff34b7cf5e1444f3789a017a7fe794465639'],
        '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', version, {
        'source_tmpl': 'tensorboard-%(version)s-py3-none-any.whl',
        'unpack_sources': False,
        'checksums': ['7b8c53c396069b618f6f276ec94fc45d17e3282d668979216e5d30be472115e4'],
    }),
    ('tensorboard-plugin-wit', '1.8.0', {
        'source_tmpl': 'tensorboard_plugin_wit-%(version)s-py3-none-any.whl',
        'unpack_sources': False,
        'checksums': ['2a80d1c551d741e99b2f197bb915d8a133e24adb8da1732b840041860f91183a'],
    }),
    ('tensorboard-plugin-profile', '2.4.0', {
        'source_tmpl': 'tensorboard_plugin_profile-%(version)s.tar.gz',
        'checksums': ['dfbf254ee960440e3b2518324f876a6d6704c60b936887d99214fa36988a206a'],
    }),
    ('google-pasta', '0.2.0', {
        'modulename': 'pasta',
        'checksums': ['c9f2c8dfc8f96d0d5808299920721be30c9eec37f2389f28904f454565c8a16e'],
    }),
    ('termcolor', '1.1.0', {
        'checksums': ['1d6d69ce66211143803fbc56652b41d73b4a400a2891d7bf7a1cdf4c02de613b'],
    }),
    ('tensorflow-estimator', '2.4.0', {
        'source_tmpl': 'tensorflow_estimator-%(version)s-py2.py3-none-any.whl',
        'unpack_sources': False,
        'checksums': ['5b7b7bf2debe19a8794adacc43e8ba6459daa4efaf54d3302623994a359b17f0'],
    }),
    ('astor', '0.8.1', {
        'checksums': ['6a6effda93f4e1ce9f618779b2dd1d9d84f1e32812c23a29b3fff6fd7f63fa5e'],
    }),
    ('gast', '0.3.3', {
        'checksums': ['b881ef288a49aa81440d2c5eb8aeefd4c2bb8993d5f50edae7413a85bfdb3b57'],
    }),
    ('opt-einsum', '3.3.0', {
        'source_tmpl': 'opt_einsum-%(version)s.tar.gz',
        'checksums': ['59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549'],
    }),
    ('wrapt', '1.12.1', {
        'checksums': ['b62ffa81fb85f4332a4f609cab4ac40709470da05643a082ec1eb88e6d9b97d7'],
    }),
    ('Keras-Preprocessing', '1.1.2', {
        'source_tmpl': 'Keras_Preprocessing-%(version)s.tar.gz',
        'checksums': ['add82567c50c8bc648c14195bf544a5ce7c1f76761536956c3d2978970179ef3'],
    }),
    ('dill', '0.3.4', {
        'source_tmpl': '%(name)s-%(version)s.zip',
        'checksums': ['4ad2e60f025526219df9ac151b19760a'],
    }),
    ('tblib', '1.7.0', {
        'checksums': ['059bd77306ea7b419d4f76016aef6d7027cc8a0785579b5aad198803435f882c'],
    }),
    ('portpicker', '1.3.1', {
        'checksums': ['d2cdc776873635ed421315c4d22e63280042456bbfa07397817e687b142b9667'],
    }),
    (name, version, {
        'patches': [
            'TensorFlow-2.1.0_fix-cuda-build.patch',
            'TensorFlow-2.3.1_disable-float128-on-power.patch',
            'TensorFlow-2.4.0_add-default-shell-env.patch',
            'TensorFlow-2.4.0_fix-eigen-on-power.patch',
            'TensorFlow-2.4.0_downgrade-required-versions.patch',
            'TensorFlow-2.4.0_add-protobuf-deps.patch',
            'TensorFlow-2.4.0_add-missing-llvm-deps.patch',
            'TensorFlow-2.4.0_fix-absl-on-power.patch',
            'TensorFlow-2.4.0_dont-use-var-lock.patch',
            'TensorFlow-2.4.1_add-ldl.patch',
            'TensorFlow-2.4.1_remove-duplicate-gpu-tests.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_targets': [
            '//tensorflow/core/...',
            '-//tensorflow/core:example_java_proto',
            '-//tensorflow/core/example:example_protos_closure',
            '//tensorflow/cc/...',
            '//tensorflow/c/...',
            '//tensorflow/python/...',
            # Requires access to CUPTI (privileged!)
            '-//tensorflow/core/profiler/internal/gpu:device_tracer_test',
            # 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',
            '-//tensorflow/c/eager:c_api_remote_function_test_gpu',
            '-//tensorflow/c/eager:c_api_remote_test_gpu',
            # Invalid alignment: https://github.com/tensorflow/tensorflow/issues/46460
            '-//tensorflow/core/kernels:sparse_matmul_op_test',
            '-//tensorflow/core/kernels:sparse_matmul_op_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',
        ],
        'testopts': "--test_timeout=3600 --test_size_filters=small",
        'with_xla': True,
        'checksums': [
            'f681331f8fc0800883761c7709d13cda11942d4ad5ff9f44ad855e9dc78387e0',  # v2.4.1.tar.gz
            '78c20aeaa7784b8ceb46238a81e8c2461137d28e0b576deeba8357d23fbe1f5a',  # TensorFlow-2.1.0_fix-cuda-build.patch
            # TensorFlow-2.3.1_disable-float128-on-power.patch
            'a8369e92559cce4566ff831e615420fd98b83106c73ee3af7fbd5335c5fcd569',
            # TensorFlow-2.4.0_add-default-shell-env.patch
            '033a870b458406f4cabab878c4183cdf54cd7627648f24c8cc9f1964ff24d292',
            # TensorFlow-2.4.0_fix-eigen-on-power.patch
            '761303a9d5a0e78b833610891194e497a45f1218cfcb0ac4854a0d31d3156c2b',
            # TensorFlow-2.4.0_downgrade-required-versions.patch
            'ac8f53662ad0bb660089c17bb433faf1ce3c2238234a2c3be7d92c8e02c1a8d6',
            # TensorFlow-2.4.0_add-protobuf-deps.patch
            '0ae2f4f1f1dc5b85f61edbe8a4a91dcffd5cf7588fa46c6a143455e4bc41a64e',
            # TensorFlow-2.4.0_add-missing-llvm-deps.patch
            'ce8c41e75f9854e547d03de99ae72cf41afa967b20439e864bc4af3dd4f565c7',
            # TensorFlow-2.4.0_fix-absl-on-power.patch
            'cd0e1181cfc56c259cb57b7f9590816db5f5cde3fab39b3386c001359ca6e9fa',
            # TensorFlow-2.4.0_dont-use-var-lock.patch
            'b14f2493fd2edf79abd1c4f2dde6c98a3e7d5cb9c25ab9386df874d5f072d6b5',
            'ab57637e3e64a5ff8345573effed2b229c33579b9ac3c8ceb7e554d236890d4a',  # TensorFlow-2.4.1_add-ldl.patch
            # TensorFlow-2.4.1_remove-duplicate-gpu-tests.patch
            'f384cd97e004673ae530b6e5cc826f8fddb46646991f1a228fc07b2aac87ee80',
        ],
    }),
]

moduleclass = 'lib'