optional_submodules.py

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)


if __name__ == "__main__":
    checkout_nccl()