diff --git a/.github/workflows/self-push.yml b/.github/workflows/self-push.yml index e21520d0e8..0bb5dbe87c 100644 --- a/.github/workflows/self-push.yml +++ b/.github/workflows/self-push.yml @@ -40,7 +40,8 @@ jobs: - name: Run all non-slow tests on GPU env: - TF_FORCE_GPU_ALLOW_GROWTH: yes + TF_FORCE_GPU_ALLOW_GROWTH: "true" + # TF_GPU_MEMORY_LIMIT: 4096 OMP_NUM_THREADS: 1 USE_CUDA: yes run: | diff --git a/.github/workflows/self-scheduled.yml b/.github/workflows/self-scheduled.yml index d963ab5920..0473949058 100644 --- a/.github/workflows/self-scheduled.yml +++ b/.github/workflows/self-scheduled.yml @@ -41,7 +41,7 @@ jobs: - name: Run all tests on GPU env: - TF_FORCE_GPU_ALLOW_GROWTH: yes + TF_FORCE_GPU_ALLOW_GROWTH: "true" OMP_NUM_THREADS: 1 RUN_SLOW: yes USE_CUDA: yes diff --git a/tests/test_modeling_tf_common.py b/tests/test_modeling_tf_common.py index bcfb6bfe5d..e6f70d6bfa 100644 --- a/tests/test_modeling_tf_common.py +++ b/tests/test_modeling_tf_common.py @@ -21,14 +21,26 @@ import tempfile from transformers import is_tf_available, is_torch_available -from .utils import require_tf +from .utils import _tf_gpu_memory_limit, require_tf if is_tf_available(): import tensorflow as tf import numpy as np - # from transformers.modeling_bert import BertModel, BertConfig, BERT_PRETRAINED_MODEL_ARCHIVE_MAP + if _tf_gpu_memory_limit is not None: + gpus = tf.config.list_physical_devices("GPU") + for gpu in gpus: + # Restrict TensorFlow to only allocate x GB of memory on the GPUs + try: + tf.config.experimental.set_virtual_device_configuration( + gpu, [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)] + ) + logical_gpus = tf.config.experimental.list_logical_devices("GPU") + print("Logical GPUs", logical_gpus) + except RuntimeError as e: + # Virtual devices must be set before GPUs have been initialized + print(e) def _config_zero_init(config): diff --git a/tests/utils.py b/tests/utils.py index 163628d3a7..6036cefd9d 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -29,8 +29,22 @@ def parse_flag_from_env(key, default=False): return _value +def parse_int_from_env(key, default=None): + try: + value = os.environ[key] + except KeyError: + _value = default + else: + try: + _value = int(value) + except ValueError: + raise ValueError("If set, {} must be a int.".format(key)) + return _value + + _run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False) _run_custom_tokenizers = parse_flag_from_env("RUN_CUSTOM_TOKENIZERS", default=False) +_tf_gpu_memory_limit = parse_int_from_env("TF_GPU_MEMORY_LIMIT", default=None) def slow(test_case): diff --git a/utils/link_tester.py b/utils/link_tester.py index 0ef165c401..ff53eb7049 100644 --- a/utils/link_tester.py +++ b/utils/link_tester.py @@ -14,6 +14,9 @@ import requests REGEXP_FIND_S3_LINKS = r"""([\"'])(https:\/\/s3)(.*)?\1""" +S3_BUCKET_PREFIX = "https://s3.amazonaws.com/models.huggingface.co/bert" + + def list_python_files_in_repository(): """ List all python files in the repository. @@ -36,7 +39,7 @@ def find_all_links(file_paths): for path in file_paths: links += scan_code_for_links(path) - return links + return [link for link in links if link != S3_BUCKET_PREFIX] def scan_code_for_links(source):