File size: 861 Bytes
cb2428f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
"""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)
|