diff --git a/.github/workflows/update-quick-start-module.yml b/.github/workflows/update-quick-start-module.yml index 7d070eb7ff8a..8c6a452e8dd0 100644 --- a/.github/workflows/update-quick-start-module.yml +++ b/.github/workflows/update-quick-start-module.yml @@ -32,6 +32,12 @@ jobs: package-type: all os: windows channel: "nightly" + windows-arm64-nightly-matrix: + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: all + os: windows-arm64 + channel: "nightly" macos-arm64-nightly-matrix: uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main with: @@ -52,6 +58,13 @@ jobs: package-type: all os: windows channel: "release" + windows-arm64-release-matrix: + needs: [windows-arm64-nightly-matrix] + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: all + os: windows + channel: "release" macos-arm64-release-matrix: needs: [macos-arm64-nightly-matrix] uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main @@ -61,8 +74,8 @@ jobs: channel: "release" update-quick-start: - needs: [linux-nightly-matrix, windows-nightly-matrix, macos-arm64-nightly-matrix, - linux-release-matrix, windows-release-matrix, macos-arm64-release-matrix] + needs: [linux-nightly-matrix, windows-nightly-matrix, windows-arm64-nightly-matrix, macos-arm64-nightly-matrix, + linux-release-matrix, windows-release-matrix, windows-arm64-release-matrix, macos-arm64-release-matrix] runs-on: "ubuntu-latest" environment: pytorchbot-env steps: @@ -78,17 +91,21 @@ jobs: env: LINUX_NIGHTLY_MATRIX: ${{ needs.linux-nightly-matrix.outputs.matrix }} WINDOWS_NIGHTLY_MATRIX: ${{ needs.windows-nightly-matrix.outputs.matrix }} + WINDOWS_ARM64_NIGHTLY_MATRIX: ${{ needs.windows-arm64-nightly-matrix.outputs.matrix }} MACOS_NIGHTLY_MATRIX: ${{ needs.macos-arm64-nightly-matrix.outputs.matrix }} LINUX_RELEASE_MATRIX: ${{ needs.linux-release-matrix.outputs.matrix }} WINDOWS_RELEASE_MATRIX: ${{ needs.windows-release-matrix.outputs.matrix }} + WINDOWS_ARM64_RELEASE_MATRIX: ${{ needs.windows-arm64-release-matrix.outputs.matrix }} MACOS_RELEASE_MATRIX: ${{ needs.macos-arm64-release-matrix.outputs.matrix }} run: | set -ex printf '%s\n' "$LINUX_NIGHTLY_MATRIX" > linux_nightly_matrix.json printf '%s\n' "$WINDOWS_NIGHTLY_MATRIX" > windows_nightly_matrix.json + printf '%s\n' "$WINDOWS_ARM64_NIGHTLY_MATRIX" > windows-arm64_nightly_matrix.json printf '%s\n' "$MACOS_NIGHTLY_MATRIX" > macos_nightly_matrix.json printf '%s\n' "$LINUX_RELEASE_MATRIX" > linux_release_matrix.json printf '%s\n' "$WINDOWS_RELEASE_MATRIX" > windows_release_matrix.json + printf '%s\n' "$WINDOWS_ARM64_RELEASE_MATRIX" > windows-arm64_release_matrix.json printf '%s\n' "$MACOS_RELEASE_MATRIX" > macos_release_matrix.json python3 ./scripts/gen_quick_start_module.py --autogenerate > assets/quick-start-module.js rm *_matrix.json diff --git a/scripts/gen_quick_start_module.py b/scripts/gen_quick_start_module.py index 5fd20d79949f..24a670b06fbf 100755 --- a/scripts/gen_quick_start_module.py +++ b/scripts/gen_quick_start_module.py @@ -22,6 +22,7 @@ class OperatingSystem(Enum): LINUX: str = "linux" WINDOWS: str = "windows" MACOS: str = "macos" + WINDOWS_ARM64: str = "windows-arm64" PRE_CXX11_ABI = "pre-cxx11" @@ -130,7 +131,7 @@ def update_versions(versions, release_matrix, release_version): ) versions["latest_stable"] = version - # Perform update of the json file from release matrix + # Perform update of the JSON file from the release matrix for os_key, os_vers in versions["versions"][version].items(): for pkg_key, pkg_vers in os_vers.items(): for acc_key, instr in pkg_vers.items(): @@ -145,7 +146,6 @@ def update_versions(versions, release_matrix, release_version): if (x["package_type"], x["gpu_arch_type"], x["gpu_arch_version"]) == (package_type, gpu_arch_type, gpu_arch_version) ] - if pkg_arch_matrix: if package_type != "libtorch": instr["command"] = pkg_arch_matrix[0]["installation"] @@ -200,17 +200,42 @@ def gen_install_matrix(versions) -> Dict[str, str]: key = f"{ver},{pkg_key},{os_key},{acc_key},{extra_key}" note = instr["note"] lines = [note] if note is not None else [] - if pkg_key == "libtorch": - ivers = instr["versions"] - if ivers is not None: - lines += [ - f"{lab}
{val}" - for (lab, val) in ivers.items() - ] + if os_key == "windows": + if pkg_key == "libtorch": + ivers = instr["versions"] + if ivers is not None: + # Flatten x64/arm64 links into separate lines with arch in label + for lab, val in ivers.items(): + if isinstance(val, dict): + for arch, url in val.items(): + if url: + lines.append(f"{lab[:-1]} {arch}:
{url}") + else: + lines.append(f"{lab}
{val}") + elif pkg_key == "pip": + command = instr.get("command") + if isinstance(command, dict): + for arch, cmd in command.items(): + lines.append(f"{arch}: {cmd}") + elif command is not None: + lines.append(command) + else: + command = instr.get("command") + if command is not None: + lines.append(command) else: - command = instr["command"] - if command is not None: - lines.append(command) + if pkg_key == "libtorch": + ivers = instr["versions"] + if ivers is not None: + lines += [ + f"{lab}
{val}" + for (lab, val) in ivers.items() + ] + else: + command = instr.get("command") + if command is not None: + lines.append(command) + result[key] = "
".join(lines) return result @@ -235,6 +260,44 @@ def gen_ver_list(chan, gpu_arch_type): acc_arch_ver_map[chan][label] = ("cuda", cuda_ver) +def merge_windows_arch_entries(entries): + """ + Merge x64 and arm64 entries for Windows + """ + from collections import defaultdict + + def entry_key(entry): + # Exclude validation_runner and installation from the key + return tuple( + (k, v) + for k, v in sorted(entry.items()) + if k not in ("validation_runner", "installation", "upload_to_base_bucket") + ) + + grouped = defaultdict(dict) + for entry in entries: + key = entry_key(entry) + arch = "arm64" if "arm64" in str(entry.get("validation_runner", "")).lower() else "x64" + grouped[key][arch] = entry + + merged = [] + for key, arch_dict in grouped.items(): + if "x64" in arch_dict and "arm64" in arch_dict: + base = {k: v for k, v in arch_dict["x64"].items() if k not in ("validation_runner", "installation")} + base["validation_runner"] = { + "x64": arch_dict["x64"]["validation_runner"], + "arm64": arch_dict["arm64"]["validation_runner"], + } + base["installation"] = { + "x64": arch_dict["x64"]["installation"], + "arm64": arch_dict["arm64"]["installation"], + } + merged.append(base) + else: + merged.extend(arch_dict.values()) + return merged + + def main(): parser = argparse.ArgumentParser() parser.add_argument("--autogenerate", dest="autogenerate", action="store_true") @@ -248,7 +311,13 @@ def main(): for val in ("nightly", "release"): release_matrix[val] = {} for osys in OperatingSystem: - release_matrix[val][osys.value] = read_matrix_for_os(osys, val) + if osys == OperatingSystem.WINDOWS_ARM64: + winarm64_matrix = read_matrix_for_os(osys, val) + windowsx64_matrix = release_matrix[val][OperatingSystem.WINDOWS.value] + merged = merge_windows_arch_entries(windowsx64_matrix + winarm64_matrix) + release_matrix[val][OperatingSystem.WINDOWS.value] = merged + else: + release_matrix[val][osys.value] = read_matrix_for_os(osys, val) write_releases_file(release_matrix)