setup.ps1

[CmdletBinding()]
param(
    [string]$ExistingTorchSitePackages = $env:VLLM_TORCH_SITE_PACKAGES
)

$ErrorActionPreference = "Stop"
$ProjectRoot = (Resolve-Path $PSScriptRoot).Path
$VendorRoot = Join-Path $ProjectRoot "vendor"
$Venv = Join-Path $ProjectRoot ".venv"
$VllmSource = Join-Path $VendorRoot "vllm"
$PluginRepo = Join-Path $VendorRoot "windows-rocm"
$PluginSource = Join-Path $PluginRepo "windows_rocm_plugin"
$Requirements = Join-Path $ProjectRoot "requirements-vllm-rdna4-windows.txt"
$SelfContainedRocmRequirements = Join-Path $ProjectRoot "requirements-rocm712-windows.txt"
$VllmCommit = "b1388b1fbf5aaef47937fabe98931211684666a6"
$PluginCommit = "7c16fdd7244217e2e45e2c5118b49dd942f8487f"
$VllmUrl = "https://github.com/vllm-project/vllm.git"
$PluginUrl = "https://github.com/ThePie88/vLLM-ROCm-Windows.git"

if (-not (Test-Path -LiteralPath $Requirements)) {
    throw "Missing dependency file: $Requirements"
}
if (-not (Test-Path -LiteralPath $SelfContainedRocmRequirements)) {
    throw "Missing ROCm dependency file: $SelfContainedRocmRequirements"
}

function Invoke-CapturedProcess {
    param(
        [Parameter(Mandatory = $true)][string]$FilePath,
        [Parameter(Mandatory = $false)][string[]]$Arguments = @()
    )

    $psi = [System.Diagnostics.ProcessStartInfo]::new()
    $psi.FileName = $FilePath
    $psi.UseShellExecute = $false
    $psi.CreateNoWindow = $true
    $psi.RedirectStandardOutput = $true
    $psi.RedirectStandardError = $true
    $quotedArguments = foreach ($argument in $Arguments) {
        $text = [string]$argument
        if ($text -notmatch '[\s"]') {
            $text
        } else {
            $escaped = $text -replace '(\\*)"', '$1$1\"'
            $escaped = $escaped -replace '(\\+)$', '$1$1'
            '"' + $escaped + '"'
        }
    }
    $psi.Arguments = [string]::Join(' ', [string[]]$quotedArguments)
    $process = [System.Diagnostics.Process]::new()
    $process.StartInfo = $psi
    try {
        [void]$process.Start()
        $stdout = $process.StandardOutput.ReadToEnd()
        $stderr = $process.StandardError.ReadToEnd()
        $process.WaitForExit()
        return [pscustomobject]@{
            ExitCode = $process.ExitCode
            StdOut = @($stdout -split "`r?`n" | Where-Object { $_ -ne "" })
            StdErr = @($stderr -split "`r?`n" | Where-Object { $_ -ne "" })
        }
    } finally {
        $process.Dispose()
    }
}

function Find-BasePython {
    if ($env:VLLM_BASE_PYTHON -and (Test-Path -LiteralPath $env:VLLM_BASE_PYTHON)) {
        return (Resolve-Path $env:VLLM_BASE_PYTHON).Path
    }
    $pyLauncher = Get-Command py.exe -ErrorAction SilentlyContinue
    if ($pyLauncher) {
        $result = Invoke-CapturedProcess $pyLauncher.Source @("-3.12", "-c", "import sys; print(sys.executable)")
        $candidate = ($result.StdOut | Select-Object -First 1).Trim()
        if ($result.ExitCode -eq 0 -and $candidate -and (Test-Path -LiteralPath $candidate)) {
            return $candidate
        }
    }
    $python = Get-Command python.exe -ErrorAction SilentlyContinue
    if ($python -and $python.Source -notmatch "WindowsApps") {
        $result = Invoke-CapturedProcess $python.Source @("-c", "import sys; print(sys.version_info[:2])")
        if ($result.ExitCode -eq 0 -and (($result.StdOut -join " ") -match "3, 12")) {
            return $python.Source
        }
    }
    $common = Join-Path $env:LOCALAPPDATA "Programs\Python\Python312\python.exe"
    if (Test-Path -LiteralPath $common) { return $common }
    throw "Python 3.12 x64 was not found. Install it or set VLLM_BASE_PYTHON."
}

function Ensure-Checkout {
    param(
        [Parameter(Mandatory = $true)][string]$Path,
        [Parameter(Mandatory = $true)][string]$Url,
        [Parameter(Mandatory = $true)][string]$Commit
    )

    $git = Get-Command git.exe -ErrorAction SilentlyContinue
    if (-not $git) { throw "Git is required to download the pinned source checkouts." }

    if (-not (Test-Path -LiteralPath $Path)) {
        New-Item -ItemType Directory -Force -Path (Split-Path $Path -Parent) | Out-Null
        Write-Output "CLONE_START url=$Url path=$Path"
        & $git.Source clone --filter=blob:none --no-checkout $Url $Path
        if ($LASTEXITCODE -ne 0) { throw "Failed to clone $Url" }
    } elseif (-not (Test-Path -LiteralPath (Join-Path $Path ".git"))) {
        throw "Checkout path exists but is not a Git repository: $Path"
    }

    $currentResult = Invoke-CapturedProcess $git.Source @("-C", $Path, "rev-parse", "HEAD")
    $current = if ($currentResult.StdOut.Count -gt 0) { $currentResult.StdOut[-1].Trim() } else { "" }
    if ($current -ne $Commit) {
        if ($current) {
            $status = Invoke-CapturedProcess $git.Source @("-C", $Path, "status", "--porcelain")
            if ($status.StdOut.Count -gt 0) {
                throw "Cannot switch $Path to $Commit because it contains local changes."
            }
        }
        & $git.Source -C $Path fetch --depth 1 origin $Commit
        if ($LASTEXITCODE -ne 0) { throw "Failed to fetch commit $Commit" }
        & $git.Source -C $Path checkout --detach $Commit
        if ($LASTEXITCODE -ne 0) { throw "Failed to checkout commit $Commit" }
    }
    Write-Output "CHECKOUT_READY path=$Path commit=$Commit"
}

Ensure-Checkout -Path $VllmSource -Url $VllmUrl -Commit $VllmCommit
Ensure-Checkout -Path $PluginRepo -Url $PluginUrl -Commit $PluginCommit

if (-not (Test-Path -LiteralPath (Join-Path $VllmSource "vllm\__init__.py"))) {
    throw "The pinned vLLM checkout is incomplete: $VllmSource"
}
if (-not (Test-Path -LiteralPath (Join-Path $PluginSource "pyproject.toml"))) {
    throw "The pinned Windows ROCm plugin checkout is incomplete: $PluginSource"
}

function Apply-PatchOnce {
    param(
        [Parameter(Mandatory = $true)][string]$Repository,
        [Parameter(Mandatory = $true)][string]$Patch,
        [Parameter(Mandatory = $true)][string]$MarkerFile,
        [Parameter(Mandatory = $true)][string]$Marker
    )

    if (Select-String -LiteralPath $MarkerFile -Pattern $Marker -Quiet) {
        return
    }
    $git = (Get-Command git.exe -ErrorAction Stop).Source
    Write-Output "PATCH_APPLY patch=$Patch"
    & $git -C $Repository apply --whitespace=nowarn $Patch
    if ($LASTEXITCODE -ne 0) { throw "Failed to apply patch: $Patch" }
}

Apply-PatchOnce -Repository $VllmSource `
    -Patch (Join-Path $ProjectRoot "patches\0001-windows-rocm-fallbacks.patch") `
    -MarkerFile (Join-Path $VllmSource "vllm\__init__.py") `
    -Marker "vllm_windows_rocm.bootstrap"
Apply-PatchOnce -Repository $VllmSource `
    -Patch (Join-Path $ProjectRoot "patches\0004-windows-hf-uri-separator.patch") `
    -MarkerFile (Join-Path $VllmSource "vllm\model_executor\model_loader\weight_utils.py") `
    -Marker "hf_model_path"
Apply-PatchOnce -Repository $PluginRepo `
    -Patch (Join-Path $ProjectRoot "patches\0002-windows-selector-event-loop.patch") `
    -MarkerFile (Join-Path $PluginSource "vllm_windows_rocm\torchdist_shim.py") `
    -Marker "WindowsSelectorEventLoopPolicy"
Apply-PatchOnce -Repository $PluginRepo `
    -Patch (Join-Path $ProjectRoot "patches\0003-native-c-ext-windows-rocm.patch") `
    -MarkerFile (Join-Path $PluginRepo "experiments\vllm_c_ext\build_c_ext.py") `
    -Marker "fms-runtime-lib=dll"

if (-not (Test-Path -LiteralPath (Join-Path $Venv "Scripts\python.exe"))) {
    $basePython = Find-BasePython
    Write-Output "VENV_CREATE python=$basePython path=$Venv"
    & $basePython -m venv $Venv
    if ($LASTEXITCODE -ne 0) { throw "Failed to create the Python virtual environment." }
}
$Python = Join-Path $Venv "Scripts\python.exe"
$siteDir = Join-Path $Venv "Lib\site-packages"
$bridge = Join-Path $siteDir "_existing_rocm_torch.pth"

$localTorch = Invoke-CapturedProcess $Python @("-c", "import torch; print(torch.__file__)")
if ($localTorch.ExitCode -eq 0 -and $localTorch.StdOut.Count -gt 0) {
    if (Test-Path -LiteralPath $bridge) {
        Remove-Item -LiteralPath $bridge -Force
        Write-Output "REMOVED_EXTERNAL_TORCH_BRIDGE path=$bridge"
    }
} elseif ($ExistingTorchSitePackages -and (Test-Path -LiteralPath $ExistingTorchSitePackages)) {
    New-Item -ItemType Directory -Force -Path $siteDir | Out-Null
    Set-Content -LiteralPath $bridge -Value $ExistingTorchSitePackages -Encoding ASCII
} else {
    if (-not (Test-Path -LiteralPath $SelfContainedRocmRequirements)) {
        throw "ROCm requirements file not found: $SelfContainedRocmRequirements"
    }
    Write-Output "TORCH_INSTALL_START requirements=$SelfContainedRocmRequirements"
    & $Python -m pip install --no-cache-dir -r $SelfContainedRocmRequirements
    if ($LASTEXITCODE -ne 0) { throw "Native ROCm/PyTorch installation failed." }
}

$torchCheck = Invoke-CapturedProcess $Python @("-c", "import torch; print(torch.__file__)")
if ($torchCheck.ExitCode -ne 0) {
    throw "The project environment cannot import ROCm PyTorch."
}

$env:VLLM_TARGET_DEVICE = "empty"
& $Python -m pip install --upgrade pip "setuptools==80.10.2" wheel
if ($LASTEXITCODE -ne 0) { throw "pip bootstrap failed." }

$commonRequirements = Join-Path $VllmSource "requirements\common.txt"
& $Python -m pip install -r $commonRequirements
if ($LASTEXITCODE -ne 0) { throw "vLLM common dependency installation failed." }
& $Python -m pip install -r $Requirements
if ($LASTEXITCODE -ne 0) { throw "Windows ROCm dependency installation failed." }
& $Python -m pip install "llguidance==1.3.0" "xgrammar==0.1.34" "apache-tvm-ffi==0.1.13.post3"
if ($LASTEXITCODE -ne 0) { throw "Structured-output dependency installation failed." }

& $Python -m pip install --no-build-isolation --no-deps -e $VllmSource
if ($LASTEXITCODE -ne 0) { throw "vLLM editable installation failed." }
& $Python -m pip install --no-deps -e $PluginSource
if ($LASTEXITCODE -ne 0) { throw "Windows ROCm plugin installation failed." }

& (Join-Path $ProjectRoot "prepare_runtime.ps1")
if ($LASTEXITCODE -ne 0) { throw "Runtime preparation failed." }
Write-Output "SETUP_OK venv=$Venv"