wehub-resource-sync
e06fe8e8c6
Self-hosted runner (benchmark) / Benchmark (aws-g5-4xlarge-cache) (push) Waiting to run
New model PR merged notification / Notify new model (push) Waiting to run
Update Transformers metadata / build_and_package (push) Waiting to run
Secret Leaks / trufflehog (push) Failing after 1s
Build documentation / build (push) Failing after 1s
Build documentation / build_other_lang (push) Failing after 0s
CodeQL Security Analysis / CodeQL Analysis (push) Failing after 0s
PR CI / pr-ci (push) Failing after 1s
Slow tests on important models (on Push - A10) / Get all modified files (push) Failing after 1s
Slow tests on important models (on Push - A10) / Model CI (push) Has been skipped
27 lines
861 B
Python
27 lines
861 B
Python
"""A simple script to set flexibly CUDA_VISIBLE_DEVICES in GitHub Actions CI workflow files."""
|
|
|
|
import argparse
|
|
import os
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--test_folder",
|
|
type=str,
|
|
default=None,
|
|
help="The test folder name of the model being tested. For example, `models/cohere`.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# `test_eager_matches_sdpa_generate` for `cohere` needs a lot of GPU memory!
|
|
# This depends on the runners. At this moment we are targeting our AWS CI runners.
|
|
if args.test_folder == "models/cohere":
|
|
cuda_visible_devices = "0,1,2,3"
|
|
elif "CUDA_VISIBLE_DEVICES" in os.environ:
|
|
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
else:
|
|
cuda_visible_devices = "0"
|
|
|
|
print(cuda_visible_devices)
|