diff --git a/scripts/gen_quick_start_module.py b/scripts/gen_quick_start_module.py index 1f074da6031c..d3b4a3c42e32 100755 --- a/scripts/gen_quick_start_module.py +++ b/scripts/gen_quick_start_module.py @@ -107,18 +107,51 @@ def get_package_type(pkg_key: str, os_key: OperatingSystem) -> str: def get_gpu_info(acc_key, instr, acc_arch_map): + if DEFAULT in instr: + return acc_arch_map["accnone"] if acc_key not in acc_arch_map: return (None, None) - gpu_arch_type, gpu_arch_version = acc_arch_map[acc_key] - if DEFAULT in instr: - gpu_arch_type, gpu_arch_version = acc_arch_map["accnone"] - return (gpu_arch_type, gpu_arch_version) + return acc_arch_map[acc_key] # This method is used for generating new published_versions.json file # It will modify versions json object with installation instructions # Provided by generate install matrix Github Workflow, stored in release_matrix # json object. +CUDA_LABELS = ["cuda.x", "cuda.y", "cuda.z"] +UNAVAILABLE_NOTE = "NOTE: Not available for this build" + + +def sync_cuda_entries(version_entry, acc_arch_map): + """Match the CUDA entries of a version to the channel's CUDA labels. + + A channel with fewer CUDA versions (e.g. two) drops the unused labels, so + the entry carries no stale commands for boxes the page hides. A new stable + is copied from preview, so any label the channel needs but the entry lacks + is re-created from a sibling entry. Only the entry being updated (preview + or the latest stable) is touched; older versions keep their labels. + """ + for os_vers in version_entry.values(): + for pkg_key, pkg_vers in os_vers.items(): + siblings = [pkg_vers[k] for k in CUDA_LABELS if k in pkg_vers] + if not siblings: + continue + rebuilt = {} + for acc_key, instr in pkg_vers.items(): + if acc_key in CUDA_LABELS and acc_key not in acc_arch_map: + continue + rebuilt[acc_key] = instr + if acc_key in CUDA_LABELS: + # Re-add missing labels right after their predecessor so + # cuda.x/y/z stay in order. + for label in CUDA_LABELS[CUDA_LABELS.index(acc_key) + 1 :]: + if label in pkg_vers: + break + if label in acc_arch_map: + rebuilt[label] = copy.deepcopy(siblings[-1]) + os_vers[pkg_key] = rebuilt + + def update_versions(versions, release_matrix, release_version): version = "preview" template = "preview" @@ -132,6 +165,8 @@ def update_versions(versions, release_matrix, release_version): ) versions["latest_stable"] = version + sync_cuda_entries(versions["versions"][version], acc_arch_map) + # Perform update of the json file from release matrix for os_key, os_vers in versions["versions"][version].items(): for pkg_key, pkg_vers in os_vers.items(): @@ -149,15 +184,26 @@ def update_versions(versions, release_matrix, release_version): ] if gpu_arch_type is None: + # The channel ships no build for this accelerator (e.g. no + # ROCm version in the matrix): don't keep a stale command. + if not instr.get("note"): + instr["note"] = UNAVAILABLE_NOTE + instr["command" if package_type != "libtorch" else "versions"] = None continue + if pkg_arch_matrix: + # A new stable is copied from preview; drop the note an + # accelerator got there while the channel lacked it. + if instr.get("note") == UNAVAILABLE_NOTE: + instr["note"] = None + if package_type == "libtorch" and instr.get("versions") is None: + instr["versions"] = {} + # A CUDA version the channel ships but this OS does not (e.g. no # Windows build): say so instead of keeping a stale command. if gpu_arch_type == "cuda": if pkg_arch_matrix: instr["note"] = None - if package_type == "libtorch" and instr.get("versions") is None: - instr["versions"] = {} else: instr["note"] = ( f"NOTE: CUDA {gpu_arch_version} is not available " @@ -250,17 +296,32 @@ def gen_ver_list(chan, gpu_arch_type): if x["gpu_arch_type"] == gpu_arch_type } + # Compare numerically: as strings, "7.14" sorts before "7.2". + def version_key(ver): + return tuple(int(part) for part in ver.split(".")) + for chan in ("nightly", "release"): cuda_ver_list = gen_ver_list(chan, "cuda") rocm_ver_list = gen_ver_list(chan, "rocm") - cuda_list = sorted(cuda_ver_list.values()) - acc_arch_ver_map[chan]["rocm5.x"] = ("rocm", max(rocm_ver_list.values())) - cuda_labels = ["cuda.x", "cuda.y", "cuda.z"] - for cuda_ver, label in zip(cuda_list, cuda_labels): + cuda_list = sorted(set(cuda_ver_list.values()), key=version_key) + if len(cuda_list) > len(CUDA_LABELS): + raise ValueError( + f"{chan} getting-started matrix has {len(cuda_list)} CUDA versions " + f"({', '.join(cuda_list)}), but the page has only " + f"{len(CUDA_LABELS)} CUDA boxes" + ) + if rocm_ver_list: + acc_arch_ver_map[chan]["rocm5.x"] = ( + "rocm", + max(rocm_ver_list.values(), key=version_key), + ) + else: + acc_arch_ver_map[chan].pop("rocm5.x", None) + for cuda_ver, label in zip(cuda_list, CUDA_LABELS): acc_arch_ver_map[chan][label] = ("cuda", cuda_ver) - # Drop labels the channel has no CUDA version for, so the page strikes - # them out instead of showing the hard-coded defaults. - for label in cuda_labels[len(cuda_list):]: + # Drop labels the channel has no CUDA version for, so the page hides + # them instead of showing the hard-coded defaults. + for label in CUDA_LABELS[len(cuda_list):]: acc_arch_ver_map[chan].pop(label, None)