|
| 1 | +function Get-LatestPythonPatchVersionFromPyEnvWin { |
| 2 | + <# |
| 3 | + .SYNOPSIS |
| 4 | + Resolves the latest patch version for a given Python major.minor (e.g. |
| 5 | + '3.12') by parsing `pyenv install --list` output on Windows (pyenv-win). |
| 6 | + .PARAMETER Version |
| 7 | + A string in the form 'M.m' (e.g., '3.10', '3.11', '3.12'). |
| 8 | + #> |
| 9 | + [CmdletBinding()] |
| 10 | + param( |
| 11 | + [Parameter(Mandatory, Position = 0)] |
| 12 | + [ValidatePattern('^\d+\.\d+$')] |
| 13 | + [string]$Version |
| 14 | + ) |
| 15 | + |
| 16 | + # Verify pyenv exists. |
| 17 | + if (-not (Get-Command pyenv -ErrorAction SilentlyContinue)) { |
| 18 | + throw [System.InvalidOperationException]::new( |
| 19 | + 'pyenv-win ("pyenv") not found on PATH.' |
| 20 | + ) |
| 21 | + } |
| 22 | + |
| 23 | + $listOutput = & pyenv install --list 2>&1 |
| 24 | + if ($LASTEXITCODE -ne 0 -or -not $listOutput) { |
| 25 | + $joined = $listOutput -join "`n" |
| 26 | + throw [System.InvalidOperationException]::new( |
| 27 | + "Failed to run 'pyenv install --list'. Output:`n$joined" |
| 28 | + ) |
| 29 | + } |
| 30 | + |
| 31 | + # Build a list of patch numbers that match the requested minor version. |
| 32 | + $versionPrefix = "$Version." |
| 33 | + $patchNumbers = @() |
| 34 | + foreach ($line in $listOutput) { |
| 35 | + $candidate = $line.Trim() |
| 36 | + if (-not $candidate) { continue } |
| 37 | + |
| 38 | + # Accept any major version; the StartsWith check guarantees we only |
| 39 | + # keep the wanted minor. |
| 40 | + if (-not $candidate.StartsWith($versionPrefix)) { continue } |
| 41 | + if ($candidate -notmatch '^\d+\.\d+\.\d+$') { continue } |
| 42 | + |
| 43 | + $patchNumbers += [int]($candidate.Split('.')[2]) |
| 44 | + } |
| 45 | + |
| 46 | + if ($patchNumbers.Count -eq 0) { |
| 47 | + throw [System.InvalidOperationException]::new( |
| 48 | + "No installable CPython versions found for prefix " + |
| 49 | + "'$Version' in pyenv-win list." |
| 50 | + ) |
| 51 | + } |
| 52 | + |
| 53 | + $latestPatch = ($patchNumbers | Sort-Object -Descending)[0] |
| 54 | + return "$Version.$latestPatch" |
| 55 | +} |
| 56 | + |
| 57 | +function Install-PythonViaPyEnvWin { |
| 58 | + <# |
| 59 | + .SYNOPSIS |
| 60 | + Ensures a Python version for the given major.minor exists via |
| 61 | + pyenv-win, activates it for the current shell, and returns the |
| 62 | + path to python.exe. |
| 63 | + .PARAMETER Version |
| 64 | + A string in the form 'M.m' (e.g., '3.12'). |
| 65 | + #> |
| 66 | + param( |
| 67 | + [Parameter(Mandatory, Position = 0)] |
| 68 | + [ValidatePattern('^\d+\.\d+$')] |
| 69 | + [string]$Version |
| 70 | + ) |
| 71 | + |
| 72 | + $fullVersion = Get-LatestPythonPatchVersionFromPyEnvWin ` |
| 73 | + -Version $Version |
| 74 | + |
| 75 | + Write-Host "Installing Python $fullVersion via pyenv..." |
| 76 | + Write-Host "pyenv install $fullVersion" |
| 77 | + ($null = & pyenv install $fullVersion | Out-Host) |
| 78 | + if ($LASTEXITCODE -ne 0) { |
| 79 | + throw [System.InvalidOperationException]::new( |
| 80 | + "Failed to install Python $fullVersion via pyenv." |
| 81 | + ) |
| 82 | + } |
| 83 | + Write-Host "Successfully installed Python $fullVersion via pyenv." |
| 84 | + |
| 85 | + ($null = & pyenv local $fullVersion | Out-Host) |
| 86 | + if ($LASTEXITCODE -ne 0) { |
| 87 | + throw [System.InvalidOperationException]::new( |
| 88 | + "Failed to set Python $fullVersion as local via pyenv." |
| 89 | + ) |
| 90 | + } |
| 91 | + Write-Host "Successfully set Python $fullVersion as local via pyenv." |
| 92 | + |
| 93 | + # Avoid the shim (i.e. shims/python.bat) because it will attempt to set |
| 94 | + # a codepage via `chcp` that we probably won't have installed on our |
| 95 | + # Server Core-based image. |
| 96 | + $exe = (Resolve-Path -LiteralPath $(pyenv which python)).Path |
| 97 | + Write-Host "python.exe path: $exe" |
| 98 | + |
| 99 | + # Add the root and Scripts directory to $Env:PATH. |
| 100 | + $rootDir = $exe.Replace("\python.exe", "") |
| 101 | + $scriptsDir = $exe.Replace("python.exe", "Scripts") |
| 102 | + $pathPrefix = $rootDir + ";" + $scriptsDir + ";" |
| 103 | + $Env:PATH = $pathPrefix + $Env:PATH |
| 104 | + |
| 105 | + # Upgrade pip using the found exe. This is necessary because some older |
| 106 | + # versions of pip (e.g. 23.10) don't support arguments like `--wheeldir`. |
| 107 | + ($null = & $exe -m pip install --upgrade pip --no-cache-dir | Out-Host) |
| 108 | + if ($LASTEXITCODE -ne 0) { |
| 109 | + throw [System.InvalidOperationException]::new("pip upgrade failed") |
| 110 | + } |
| 111 | + |
| 112 | + Write-Host "pip successfully upgraded, running pyenv rehash..." |
| 113 | + ($null = & pyenv rehash | Out-Host) |
| 114 | + if ($LASTEXITCODE -ne 0) { |
| 115 | + throw [System.InvalidOperationException]::new("pyenv rehash failed") |
| 116 | + } |
| 117 | + Write-Host "Successfully ran pyenv rehash." |
| 118 | + |
| 119 | + return $exe |
| 120 | +} |
| 121 | + |
| 122 | +function Get-Python { |
| 123 | + <# |
| 124 | + .SYNOPSIS |
| 125 | + Returns the path of the Python interpreter satisfying the supplied |
| 126 | + version, potentially installing it via pyenv-win if it's not already |
| 127 | + installed. |
| 128 | + #> |
| 129 | + [CmdletBinding()] |
| 130 | + param( |
| 131 | + [Parameter(Mandatory, Position = 0)] |
| 132 | + [ValidatePattern('^\d+\.\d+$')] |
| 133 | + [string]$Version |
| 134 | + ) |
| 135 | + |
| 136 | + # Look for a plain 'python.exe' already on the path. |
| 137 | + try { |
| 138 | + $candidate = (Get-Command python -ErrorAction Stop).Source |
| 139 | + $foundVer = & $candidate -c "import sys; print(f'{sys.version_info[0]}.{sys.version_info[1]}')" 2>$null |
| 140 | + if ($foundVer -eq $Version) { |
| 141 | + Write-Host "Found matching Python $foundVer at $candidate." |
| 142 | + return $candidate.Trim() |
| 143 | + } |
| 144 | + else { |
| 145 | + Write-Host "Found python.exe but version $foundVer != requested version $Version." |
| 146 | + } |
| 147 | + } |
| 148 | + catch { |
| 149 | + Write-Host "Unable to query existing 'python' on PATH: $_" |
| 150 | + } |
| 151 | + |
| 152 | + # If we reach here, we'll need to install the requested version via pyenv. |
| 153 | + try { |
| 154 | + $exe = Install-PythonViaPyEnvWin -Version $Version |
| 155 | + return $exe.Trim() |
| 156 | + } |
| 157 | + catch { |
| 158 | + throw [System.InvalidOperationException]::new( |
| 159 | + "Requested Python $Version not found and installation " + |
| 160 | + "via pyenv-win failed: $($_.Exception.Message)" |
| 161 | + ) |
| 162 | + } |
| 163 | +} |
| 164 | + |
| 165 | +function Get-CudaMajor { |
| 166 | + <# |
| 167 | + .SYNOPSIS |
| 168 | + Gets the CUDA major version for this container instance (e.g. '12' or |
| 169 | + '13'). Defaults to '13' if no match can be found. |
| 170 | + #> |
| 171 | + if ($env:CUDA_PATH) { |
| 172 | + $nvcc = Join-Path $env:CUDA_PATH "bin/nvcc.exe" |
| 173 | + if (Test-Path $nvcc) { |
| 174 | + $out = & $nvcc --version 2>&1 |
| 175 | + $text = ($out -join "`n") |
| 176 | + if ($text -match 'release\s+(\d+)\.') { return $Matches[1] } |
| 177 | + } |
| 178 | + # Fallback: parse major from CUDA_PATH like ...\v13.0 or ...\CUDA\13 |
| 179 | + $pathMatch = [regex]::Match($env:CUDA_PATH, 'v?(\d+)(?:\.\d+)?') |
| 180 | + if ($pathMatch.Success) { return $pathMatch.Groups[1].Value } |
| 181 | + } |
| 182 | + return '13' |
| 183 | +} |
| 184 | + |
| 185 | +function Convert-ToUnixPath { |
| 186 | + Param([Parameter(Mandatory = $true)][string]$p) |
| 187 | + return ($p -replace "\\", "/") |
| 188 | +} |
| 189 | + |
| 190 | +function Get-CudaCcclWheel { |
| 191 | + <# |
| 192 | + .SYNOPSIS |
| 193 | + Returns the path of the cuda-cccl wheel artifact to use in the context |
| 194 | + of a GitHub Actions CI test script. |
| 195 | + #> |
| 196 | + Param() |
| 197 | + |
| 198 | + $repoRoot = Get-RepoRoot |
| 199 | + if ($env:GITHUB_ACTIONS) { |
| 200 | + Push-Location $repoRoot |
| 201 | + try { |
| 202 | + $wheelArtifactName = (& bash -lc "ci/util/workflow/get_wheel_artifact_name.sh").Trim() |
| 203 | + if (-not $wheelArtifactName) { throw 'Failed to resolve wheel artifact name' } |
| 204 | + $repoRootPosix = Convert-ToUnixPath $repoRoot |
| 205 | + # Ensure output from downloader goes to console, not function return pipeline |
| 206 | + $null = (& bash -lc "ci/util/artifacts/download.sh $wheelArtifactName $repoRootPosix" 2>&1 | Out-Host) |
| 207 | + if ($LASTEXITCODE -ne 0) { throw "Failed to download wheel artifact '$wheelArtifactName'" } |
| 208 | + } |
| 209 | + finally { Pop-Location } |
| 210 | + } |
| 211 | + |
| 212 | + $wheelhouse = Join-Path $repoRoot 'wheelhouse' |
| 213 | + $wheelPath = Get-OnePathMatch -Path $wheelhouse -Pattern '^cuda_cccl-.*\.whl' -File |
| 214 | + return $wheelPath |
| 215 | +} |
| 216 | + |
| 217 | +function Get-OnePathMatch { |
| 218 | + <# |
| 219 | + .SYNOPSIS |
| 220 | + Returns a single path (file or directory) match for a given pattern, |
| 221 | + throwing an error if there were no matches or more than one match. |
| 222 | + #> |
| 223 | + [CmdletBinding(DefaultParameterSetName = 'FileSet')] |
| 224 | + param( |
| 225 | + [Parameter(Mandatory)] |
| 226 | + [string] $Path, |
| 227 | + |
| 228 | + [Parameter(Mandatory)] |
| 229 | + [string] $Pattern, |
| 230 | + |
| 231 | + [Parameter(Mandatory, ParameterSetName = 'FileSet')] |
| 232 | + [switch] $File, |
| 233 | + |
| 234 | + [Parameter(Mandatory, ParameterSetName = 'DirSet')] |
| 235 | + [switch] $Directory, |
| 236 | + |
| 237 | + [switch] $Recurse |
| 238 | + ) |
| 239 | + |
| 240 | + if (-not (Test-Path -LiteralPath $Path -PathType Container)) { |
| 241 | + throw "Path not found or not a directory: $Path" |
| 242 | + } |
| 243 | + |
| 244 | + $gciArgs = @{ |
| 245 | + LiteralPath = $Path |
| 246 | + ErrorAction = 'SilentlyContinue' |
| 247 | + } |
| 248 | + |
| 249 | + if ($Recurse) { $gciArgs['Recurse'] = $true } |
| 250 | + if ($PSCmdlet.ParameterSetName -eq 'FileSet') { |
| 251 | + $gciArgs['File'] = $true |
| 252 | + } |
| 253 | + else { |
| 254 | + $gciArgs['Directory'] = $true |
| 255 | + } |
| 256 | + |
| 257 | + $pathMatches = @( |
| 258 | + Get-ChildItem @gciArgs | |
| 259 | + Where-Object { $_.Name -match $Pattern } | |
| 260 | + Select-Object -ExpandProperty FullName |
| 261 | + ) |
| 262 | + |
| 263 | + if ($pathMatches.Count -ne 1) { |
| 264 | + $kind = if ($PSCmdlet.ParameterSetName -eq 'FileSet') { 'file' } |
| 265 | + else { 'directory' } |
| 266 | + $indented = ($pathMatches | ForEach-Object { " $_" }) -join "`n" |
| 267 | + |
| 268 | + $msg = @" |
| 269 | +Expected exactly one $kind name matching regex: |
| 270 | + $Pattern |
| 271 | +under: |
| 272 | + $Path |
| 273 | +Found: |
| 274 | + $($pathMatches.Count) |
| 275 | +
|
| 276 | +$indented |
| 277 | +"@ |
| 278 | + throw $msg |
| 279 | + } |
| 280 | + |
| 281 | + return $pathMatches[0] |
| 282 | +} |
| 283 | + |
| 284 | +Export-ModuleMember -Function Get-Python, Get-CudaMajor, Convert-ToUnixPath, Get-RepoRoot, Get-CudaCcclWheel, Get-OnePathMatch |
0 commit comments