diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 88c7344683..b394b308e4 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -224,7 +224,8 @@ jobs: tests/sh/test_mac_intel_compat.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ tests/sh/test_tauri_install_exit_order.sh \ - tests/sh/test_torch_constraint.sh; do + tests/sh/test_torch_constraint.sh \ + tests/sh/test_torch_flavor.sh; do echo "::group::$s" bash "$s" echo "::endgroup::" diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index a772a6d102..a6f1401067 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -65,17 +65,20 @@ jobs: with: persist-credentials: false - # Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit - # test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke. - - name: setup.ps1 unit test (Resolve-CudaToolkit) + # Fast GPU-free gate: parse install.ps1 + setup.ps1 and run the PowerShell + # unit tests (CUDA-toolkit + torch-flavor helpers) before the heavy GGUF smoke. + - name: PowerShell installer unit tests shell: pwsh run: | - $errs = $null - [void][System.Management.Automation.Language.Parser]::ParseFile( - (Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs) - if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } - Write-Host "setup.ps1 parsed with no errors" + foreach ($f in @('install.ps1', 'studio/setup.ps1')) { + $errs = $null + [void][System.Management.Automation.Language.Parser]::ParseFile( + (Resolve-Path $f).Path, [ref]$null, [ref]$errs) + if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } + Write-Host "$f parsed with no errors" + } pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1 + pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1 - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: diff --git a/install.ps1 b/install.ps1 index cbb736b371..47851ed92e 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1903,6 +1903,64 @@ exit 0 substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow" return "$baseUrl/cu126" } + + # ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── + # torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu, + # matching setup.ps1's stale-venv parse. + function ConvertTo-TorchFlavorTag { + param([string]$TorchVersion) + if (-not $TorchVersion) { return $null } + if ($TorchVersion -match '\+(cu\d+)') { return $Matches[1] } + if ($TorchVersion -match '\+rocm') { return 'rocm' } + if ($TorchVersion -match '\+cpu') { return 'cpu' } + return 'cpu' + } + + # Expected tag from the index leaf: cuXXX / cpu / rocm ($ROCmIndexUrl or a + # gfx* leaf -> rocm). $null on an unknown leaf (odd mirror) so repair no-ops. + function Get-ExpectedTorchFlavorTag { + param([string]$TorchIndexUrl, [string]$ROCmIndexUrl) + if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' } + if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null } + $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + if ($leaf -match '^cu\d+$') { return $leaf } + if ($leaf -eq 'cpu') { return 'cpu' } + if ($leaf -match '^rocm') { return 'rocm' } + if ($leaf -match '^gfx') { return 'rocm' } + return $null + } + + # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses + # ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference. + function Get-InstalledTorchTag { + param([string]$PythonExe) + if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null } + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $PythonExe + $psi.Arguments = '-c "import torch; print(torch.__version__)"' + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.UseShellExecute = $false + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + # Drain BOTH streams async, then WaitForExit. A synchronous ReadToEnd() + # before the wait would block forever if a wedged "import torch" never + # closes stdout; leaving the redirected stderr undrained would deadlock a + # child that floods it past the pipe buffer. Async reads let a noisy-but- + # exiting probe finish, while a truly hung one still hits the 30s timeout + # and is killed -- bounded either way. + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + $finished = $proc.WaitForExit(30000) + if (-not $finished) { try { $proc.Kill() } catch {}; return $null } + $torchVer = $outTask.GetAwaiter().GetResult().Trim() + [void]$errTask.GetAwaiter().GetResult() + if ($proc.ExitCode -ne 0 -or -not $torchVer) { return $null } + return ConvertTo-TorchFlavorTag $torchVer + } catch { return $null } + } + $TorchIndexUrl = Get-TorchIndexUrl # ── GPU arch → newest compatible Windows ROCm wheel release ── @@ -2149,6 +2207,50 @@ exit 0 } } + # ── Enforce the installed torch flavor matches the detected GPU build ── + # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv + # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on + # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is + # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx* + # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install + # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops. + if (-not $SkipTorch) { + $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl + if ($expectedTorchTag -and $expectedTorchTag -ne 'cpu') { + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { + if ($expectedTorchTag -eq 'rocm' -and $ROCmIndexUrl) { + # AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path + # would have force-reinstalled. Repair from the same repo.amd.com index. + $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } + substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec torchvision torchaudio } + if ($torchFixExit -ne 0) { + Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) + } + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + } elseif ($expectedTorchTag -ne 'rocm') { + # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. + substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + if ($torchFixExit -ne 0) { + Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) + } + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + } + } + # Safety net (incl. AMD): GPU build expected but still CPU -> warn loudly. + if ($installedTorchTag -eq 'cpu') { + Write-Host "" + Write-Host " [WARN] PyTorch is CPU-only but a $expectedTorchTag GPU build was expected for this machine." -ForegroundColor Yellow + Write-Host " [WARN] Training and GPU inference will run on CPU until this is fixed." -ForegroundColor Yellow + Write-Host " [WARN] Re-run this installer, or reinstall the GPU build manually for your GPU." -ForegroundColor Yellow + } + } + } + # Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped # for --local: the editable install above already makes _PACKAGE_ROOT in # unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__). diff --git a/install.sh b/install.sh index f80d9b63b2..a0d8637d94 100755 --- a/install.sh +++ b/install.sh @@ -1990,6 +1990,45 @@ get_torch_index_url() { else echo "$_base/cpu"; fi } +# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── +# torch.__version__ ($1) -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu. +_torch_flavor_tag() { + case "$1" in + *+cu[0-9]*) printf '%s\n' "$1" | sed -n 's/.*+\(cu[0-9][0-9]*\).*/\1/p' ;; + *+rocm*) echo "rocm" ;; + *+cpu*) echo "cpu" ;; + "") echo "" ;; + *) echo "cpu" ;; + esac +} + +# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> +# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. +_expected_torch_flavor_tag() { + _u="${1%/}" + _leaf="${_u##*/}" + case "$_leaf" in + cu[0-9]*) echo "$_leaf" ;; + cpu) echo "cpu" ;; + rocm*|gfx*) echo "rocm" ;; + *) echo "" ;; + esac +} + +# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX / +# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv +# resolves (torch + every transitive dep) via --index-url -- the same URLs the +# fresh-install paths above already use -- so a stale wheel is auto-repairable. +# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. +_torch_index_repairable() { + _u="${1%/}" + _leaf="${_u##*/}" + case "$_leaf" in + cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;; + *) echo "no" ;; + esac +} + get_radeon_wheel_url() { # Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing # contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/, @@ -2513,7 +2552,7 @@ if [ "$_MIGRATED" = true ]; then # fresh reinstall. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" # Repair ROCm torch if overwritten during migrated install _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) @@ -2689,7 +2728,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # which is only useful once torch is present for training. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" ;; esac @@ -2735,7 +2774,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) if [ -z "$_has_hip" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." @@ -2764,6 +2803,41 @@ else fi fi +# ── Enforce the installed torch flavor matches the detected GPU build ── +# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv +# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on +# CPU. Reinstall the right wheel triplet when a GPU build is expected; if it +# can't be reinstalled, warn loudly. --no-torch / CPU-only / macOS: no-op. +if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then + _expected_torch_tag=$(_expected_torch_flavor_tag "$TORCH_INDEX_URL") + # Only act when a GPU build is expected (cuXXX / rocm); cpu and unknown skip. + if [ -n "$_expected_torch_tag" ] && [ "$_expected_torch_tag" != "cpu" ]; then + _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) + _installed_torch_tag="" + [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") + # Repair when flavor is wrong AND the index is plain --index-url reinstallable + # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only. + if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ + && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then + substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." + run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" \ + --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio + _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) + _installed_torch_tag="" + [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") + fi + # Safety net (incl. AMD/WSL): GPU build expected but still CPU -> warn loudly. + if [ "$_installed_torch_tag" = "cpu" ]; then + substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" + substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" + substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + fi + fi +fi + # ── Run studio setup ── tauri_log "STEP" "Running Studio setup" # When --local, use the repo's own setup.sh directly. diff --git a/tests/run_all.sh b/tests/run_all.sh index d84c930392..4311aa4276 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -9,6 +9,7 @@ sh "$TESTS_DIR/sh/test_get_torch_index_url.sh" sh "$TESTS_DIR/sh/test_mac_intel_compat.sh" sh "$TESTS_DIR/sh/test_torch_constraint.sh" sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh" +sh "$TESTS_DIR/sh/test_torch_flavor.sh" echo "" echo "=== Python tests ===" diff --git a/tests/sh/test_torch_flavor.sh b/tests/sh/test_torch_flavor.sh new file mode 100755 index 0000000000..1f25ad807c --- /dev/null +++ b/tests/sh/test_torch_flavor.sh @@ -0,0 +1,68 @@ +#!/bin/bash +# Unit tests for install.sh's torch-flavor helpers (_torch_flavor_tag, +# _expected_torch_flavor_tag, _torch_index_repairable) that drive the +# stale-CPU-PyTorch repair. Helpers are extracted from install.sh and sourced. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +PASS=0 +FAIL=0 + +# Extract the three helper functions from install.sh and source them. +_FUNC_FILE=$(mktemp) +{ + sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_expected_torch_flavor_tag()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_torch_index_repairable()/,/^}/p' "$INSTALL_SH" +} > "$_FUNC_FILE" +# shellcheck disable=SC1090 +. "$_FUNC_FILE" +rm -f "$_FUNC_FILE" + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1)) + fi +} + +echo "=== _torch_flavor_tag ===" +assert_eq "cu130 wheel" "cu130" "$(_torch_flavor_tag '2.10.0+cu130')" +assert_eq "cu128 wheel" "cu128" "$(_torch_flavor_tag '2.8.0+cu128')" +assert_eq "cu124 wheel" "cu124" "$(_torch_flavor_tag '2.5.1+cu124')" +assert_eq "cu118 wheel" "cu118" "$(_torch_flavor_tag '2.4.0+cu118')" +assert_eq "cpu wheel" "cpu" "$(_torch_flavor_tag '2.10.0+cpu')" +assert_eq "untagged -> cpu" "cpu" "$(_torch_flavor_tag '2.10.0')" +assert_eq "rocm wheel" "rocm" "$(_torch_flavor_tag '2.11.0+rocm7.1')" +assert_eq "nightly cu130" "cu130" "$(_torch_flavor_tag '2.10.0.dev20250601+cu130')" +assert_eq "cu130 with suffix" "cu130" "$(_torch_flavor_tag '2.10.0+cu130.post1')" +assert_eq "empty -> empty" "" "$(_torch_flavor_tag '')" +assert_eq "garbage -> cpu" "cpu" "$(_torch_flavor_tag 'not-a-version')" + +echo "=== _expected_torch_flavor_tag ===" +assert_eq "cu130 index" "cu130" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu130')" +assert_eq "cu130 trailing /" "cu130" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu130/')" +assert_eq "cu128 index" "cu128" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu128')" +assert_eq "cpu index" "cpu" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cpu')" +assert_eq "rocm index" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7.2')" +assert_eq "amd gfx index" "rocm" "$(_expected_torch_flavor_tag 'https://repo.amd.com/rocm/whl/gfx120X-all/')" +assert_eq "mirror cu130 leaf" "cu130" "$(_expected_torch_flavor_tag 'https://my.mirror/pytorch/whl/cu130')" +assert_eq "unrecognized leaf" "" "$(_expected_torch_flavor_tag 'https://my.mirror/whl/simple')" +assert_eq "empty url" "" "$(_expected_torch_flavor_tag '')" + +echo "=== _torch_index_repairable ===" +assert_eq "cu130 repairable" "yes" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cu130')" +assert_eq "rocm7.2 repairable" "yes" "$(_torch_index_repairable 'https://download.pytorch.org/whl/rocm7.2')" +assert_eq "gfx repairable" "yes" "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx120X-all/')" +assert_eq "gfx1151 repairable" "yes" "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx1151/')" +assert_eq "cpu NOT repairable" "no" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cpu')" +assert_eq "unknown NOT repair" "no" "$(_torch_index_repairable 'https://my.mirror/whl/simple')" + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] diff --git a/tests/studio/test_torch_flavor.ps1 b/tests/studio/test_torch_flavor.ps1 new file mode 100644 index 0000000000..86303af16c --- /dev/null +++ b/tests/studio/test_torch_flavor.ps1 @@ -0,0 +1,52 @@ +#!/usr/bin/env pwsh +# Unit test for install.ps1's torch-flavor helpers (ConvertTo-TorchFlavorTag, +# Get-ExpectedTorchFlavorTag) that drive the stale-CPU-PyTorch repair. Pure +# helpers, AST-extracted and run in-process -- no GPU/venv needed. +# Run: pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1 + +$ErrorActionPreference = "Stop" +$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1") +$installPath = (Resolve-Path $installPath).Path + +# --- Parse install.ps1 (also serves as a syntax gate) and extract the helpers --- +$tokens = $null; $errors = $null +$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors) +if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" } + +foreach ($name in @("ConvertTo-TorchFlavorTag", "Get-ExpectedTorchFlavorTag")) { + $fn = $ast.FindAll({ param($n) + $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name + }, $true) + if ($fn.Count -ne 1) { throw "expected exactly one $name in install.ps1, found $($fn.Count)" } + # Pure helpers (no exit / external calls) -- safe to define in this scope. + Invoke-Expression $fn[0].Extent.Text +} + +$failures = 0 +function Check($name, $cond) { + if ($cond) { Write-Host " PASS $name" } + else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ } +} + +Write-Host "ConvertTo-TorchFlavorTag" +Check "2.10.0+cu130 -> cu130" ((ConvertTo-TorchFlavorTag "2.10.0+cu130") -eq "cu130") +Check "2.8.0+cu128 -> cu128" ((ConvertTo-TorchFlavorTag "2.8.0+cu128") -eq "cu128") +Check "2.10.0+cpu -> cpu" ((ConvertTo-TorchFlavorTag "2.10.0+cpu") -eq "cpu") +Check "2.10.0 (untagged) -> cpu" ((ConvertTo-TorchFlavorTag "2.10.0") -eq "cpu") +Check "2.11.0+rocm7.1 -> rocm" ((ConvertTo-TorchFlavorTag "2.11.0+rocm7.1") -eq "rocm") +Check "empty -> null" ($null -eq (ConvertTo-TorchFlavorTag "")) + +Write-Host "Get-ExpectedTorchFlavorTag" +Check "cu130 index -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cu130") -eq "cu130") +Check "trailing slash -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cu130/") -eq "cu130") +Check "cpu index -> cpu" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cpu") -eq "cpu") +Check "ROCm url -> rocm" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cpu" -ROCmIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all/") -eq "rocm") +Check "gfx index leaf -> rocm" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all/") -eq "rocm") +Check "rocm7.2 leaf -> rocm" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/rocm7.2") -eq "rocm") +Check "mirror cu130 leaf -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://my.mirror/whl/cu130") -eq "cu130") +Check "unrecognized leaf -> null" ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://my.mirror/whl/simple")) +Check "empty url -> null" ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl "")) + +Write-Host "" +if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 } +Write-Host "All checks passed" -ForegroundColor Green