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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
|
import os
from pathlib import Path
from subprocess import check_call
repo_root = Path(__file__).absolute().parent.parent
third_party_path = repo_root / "third_party"
def _read_file(path: Path) -> str:
with path.open(encoding="utf-8") as f:
return f.read().strip()
def _checkout_by_tag(repo: str, tag: str) -> None:
check_call(
[
"git",
"clone",
"--depth",
"1",
"--branch",
tag,
repo,
],
cwd=third_party_path,
)
def read_nccl_pin() -> str:
nccl_file = "nccl-cu12.txt"
if os.getenv("DESIRED_CUDA", os.getenv("CUDA_VERSION", "")).startswith("11"):
nccl_file = "nccl-cu11.txt"
nccl_pin_path = repo_root / ".ci" / "docker" / "ci_commit_pins" / nccl_file
return _read_file(nccl_pin_path)
def checkout_nccl() -> None:
release_tag = read_nccl_pin()
print(f"-- Checkout nccl release tag: {release_tag}")
nccl_basedir = third_party_path / "nccl"
if not nccl_basedir.exists():
_checkout_by_tag("https://github.com/NVIDIA/nccl", release_tag)
def checkout_eigen() -> None:
eigen_tag = _read_file(third_party_path / "eigen_pin.txt")
print(f"-- Checkout Eigen release tag: {eigen_tag}")
eigen_basedir = third_party_path / "eigen"
if not eigen_basedir.exists():
_checkout_by_tag("https://gitlab.com/libeigen/eigen", eigen_tag)
if __name__ == "__main__":
import sys
if len(sys.argv) == 1:
# If no arguments are given checkout all optional dependency
checkout_nccl()
checkout_eigen()
else:
# Otherwise just call top-level function of choice
globals()[sys.argv[1]]()
|
