mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
0cb85c188f
To expand it to optional eigen checkout later Pull Request resolved: https://github.com/pytorch/pytorch/pull/155947 Approved by: https://github.com/Skylion007
45 lines
1.1 KiB
Python
45 lines
1.1 KiB
Python
import os
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
|
|
repo_root = Path(__file__).absolute().parent.parent
|
|
third_party_path = os.path.join(repo_root, "third_party")
|
|
|
|
|
|
def read_nccl_pin() -> str:
|
|
nccl_file = "nccl-cu12.txt"
|
|
if os.getenv("DESIRED_CUDA", "").startswith("11") or os.getenv(
|
|
"CUDA_VERSION", ""
|
|
).startswith("11"):
|
|
nccl_file = "nccl-cu11.txt"
|
|
nccl_pin_path = os.path.join(
|
|
repo_root, ".ci", "docker", "ci_commit_pins", nccl_file
|
|
)
|
|
with open(nccl_pin_path) as f:
|
|
return f.read().strip()
|
|
|
|
|
|
def checkout_nccl() -> None:
|
|
release_tag = read_nccl_pin()
|
|
print(f"-- Checkout nccl release tag: {release_tag}")
|
|
nccl_basedir = os.path.join(third_party_path, "nccl")
|
|
if not os.path.exists(nccl_basedir):
|
|
subprocess.check_call(
|
|
[
|
|
"git",
|
|
"clone",
|
|
"--depth",
|
|
"1",
|
|
"--branch",
|
|
release_tag,
|
|
"https://github.com/NVIDIA/nccl.git",
|
|
"nccl",
|
|
],
|
|
cwd=third_party_path,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
checkout_nccl()
|