feat: update cse script package changes for aks windows gpu feature (#3794)

parts/windows/windowscsehelper.ps1

@@ -54,6 +54,16 @@ $global:WINDOWS_CSE_ERROR_NO_CUSTOM_DATA_BIN=49 # Return this error code in csec
 $global:WINDOWS_CSE_ERROR_NO_CSE_RESULT_LOG=50 # Return this error code in csecmd.ps1 when C:\AzureData\CSEResult.log does not exist
 $global:WINDOWS_CSE_ERROR_COPY_LOG_COLLECTION_SCRIPTS=51
 $global:WINDOWS_CSE_ERROR_RESIZE_OS_DRIVE=52
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_FAILED=53
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_TIMEOUT=54
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_VM_SIZE_NOT_SUPPORTED=55
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_URL_NOT_SET=56
+$global:WINDOWS_CSE_ERROR_GPU_SKU_INFO_NOT_FOUND=57
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_DOWNLOAD_FAILURE=58
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INVALID_SIGNATURE=59
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_EXCEPTION=60
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_URL_NOT_EXE=61
+
 
 # NOTE: KubernetesVersion does not contain "v"
 $global:MinimalKubernetesVersionWithLatestContainerd = "1.28.0" # Will change it to the correct version when we support new Windows containerd version
@@ -143,6 +153,18 @@ function Set-ExitCode
     exit $ExitCode
 }
 
+function Postpone-RestartComputer 
+{
+    Write-Log "Creating an one-time task to restart the VM"
+    $action = New-ScheduledTaskAction -Execute "powershell.exe" -Argument " -Command `"Restart-Computer -Force`""
+    $principal = New-ScheduledTaskPrincipal -UserId SYSTEM -LogonType ServiceAccount -RunLevel Highest
+    # trigger this task once
+    $trigger = New-JobTrigger -At  (Get-Date).AddSeconds(15).DateTime -Once
+    $definition = New-ScheduledTask -Action $action -Principal $principal -Trigger $trigger -Description "Restart computer after provisioning the VM"
+    Register-ScheduledTask -TaskName "restart-computer" -InputObject $definition
+    Write-Log "Created an one-time task to restart the VM"
+}
+
 function Create-Directory
 {
     Param(
@@ -286,4 +308,4 @@ function Install-Containerd-Based-On-Kubernetes-Version {
     $ContainerdUrl = $ContainerdUrl + $containerdPackage
   }
   Install-Containerd -ContainerdUrl $ContainerdUrl -CNIBinDir $CNIBinDir -CNIConfDir $CNIConfDir -KubeDir $KubeDir
-}
+}

pkg/templates/templates_generated.go

@@ -8204,6 +8204,16 @@ $global:WINDOWS_CSE_ERROR_NO_CUSTOM_DATA_BIN=49 # Return this error code in csec
 $global:WINDOWS_CSE_ERROR_NO_CSE_RESULT_LOG=50 # Return this error code in csecmd.ps1 when C:\AzureData\CSEResult.log does not exist
 $global:WINDOWS_CSE_ERROR_COPY_LOG_COLLECTION_SCRIPTS=51
 $global:WINDOWS_CSE_ERROR_RESIZE_OS_DRIVE=52
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_FAILED=53
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_TIMEOUT=54
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_VM_SIZE_NOT_SUPPORTED=55
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_URL_NOT_SET=56
+$global:WINDOWS_CSE_ERROR_GPU_SKU_INFO_NOT_FOUND=57
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_DOWNLOAD_FAILURE=58
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INVALID_SIGNATURE=59
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_EXCEPTION=60
+$global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_URL_NOT_EXE=61
+
 
 # NOTE: KubernetesVersion does not contain "v"
 $global:MinimalKubernetesVersionWithLatestContainerd = "1.28.0" # Will change it to the correct version when we support new Windows containerd version
@@ -8293,6 +8303,18 @@ function Set-ExitCode
     exit $ExitCode
 }
 
+function Postpone-RestartComputer 
+{
+    Write-Log "Creating an one-time task to restart the VM"
+    $action = New-ScheduledTaskAction -Execute "powershell.exe" -Argument " -Command `+"`"+`"Restart-Computer -Force`+"`"+`""
+    $principal = New-ScheduledTaskPrincipal -UserId SYSTEM -LogonType ServiceAccount -RunLevel Highest
+    # trigger this task once
+    $trigger = New-JobTrigger -At  (Get-Date).AddSeconds(15).DateTime -Once
+    $definition = New-ScheduledTask -Action $action -Principal $principal -Trigger $trigger -Description "Restart computer after provisioning the VM"
+    Register-ScheduledTask -TaskName "restart-computer" -InputObject $definition
+    Write-Log "Created an one-time task to restart the VM"
+}
+
 function Create-Directory
 {
     Param(
@@ -8436,8 +8458,7 @@ function Install-Containerd-Based-On-Kubernetes-Version {
     $ContainerdUrl = $ContainerdUrl + $containerdPackage
   }
   Install-Containerd -ContainerdUrl $ContainerdUrl -CNIBinDir $CNIBinDir -CNIConfDir $CNIConfDir -KubeDir $KubeDir
-}
-`)
+}`)
 
 func windowsWindowscsehelperPs1Bytes() ([]byte, error) {
 	return _windowsWindowscsehelperPs1, nil

staging/cse/windows/debug/collect-windows-logs.ps1

@@ -309,6 +309,19 @@ if ($tempFile -ne "") {
   $paths += $tempFile
 }
 
+$gpuTemp = Join-Path ([System.IO.Path]::GetTempPath()) ([System.IO.Path]::GetRandomFileName())
+New-Item -Type Directory $gpuTemp
+
+$nvidiaInstallLogFolder="C:\AzureData\NvidiaInstallLog"
+$logFiles = Get-ChildItem (Join-Path $nvidiaInstallLogFolder *.log)
+$logFiles | Foreach-Object {
+  Write-Host "Copying $_ to temp"
+  $tempFile = Copy-Item $_ $gpuTemp -Passthru -ErrorAction Ignore
+  if ($tempFile) {
+    $paths += $tempFile
+  }
+}
+
 Write-Host "Compressing all logs to $zipName"
 $paths | Format-Table FullName, Length -AutoSize
 Compress-Archive -LiteralPath $paths -DestinationPath $zipName

staging/cse/windows/nvidiagpudriverfunc.ps1

@@ -0,0 +1,145 @@
+function Start-InstallGPUDriver {
+    param(
+        [Parameter(Mandatory = $true)]
+        [bool]$EnableInstall,
+        # when the vm size does not have gpu, this value is an empty string.
+        [Parameter(Mandatory = $false)]
+        [string]$GpuDriverURL
+    )
+  
+    if (-not $EnableInstall) {
+        Write-Log "ConfigGPUDriverIfNeeded is false. GPU driver installation skipped as per configuration."
+        return
+    }
+  
+    Write-Log "ConfigGPUDriverIfNeeded is true. GPU driver installation started as per configuration."
+  
+    if ([string]::IsNullOrEmpty($GpuDriverURL)) {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_URL_NOT_SET -ErrorMessage "DriverURL is not properly specified."
+    }
+
+    $fileName = [IO.Path]::GetFileName($GpuDriverURL)
+    if (-not $fileName.EndsWith(".exe")) {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_URL_NOT_EXE -ErrorMessage "DriverURL does not point to a exe file"
+    }
+
+    $Target = "C:\AzureData\$fileName"
+    $LogFolder = "C:\AzureData"
+
+    Write-Log "Attempting to install Nvidia driver..."
+
+    Write-Log "Downloading from $GpuDriverURL to $Target"
+    DownloadFileOverHttp -Url $GpuDriverURL -DestinationPath $Target -ExitCode $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_DOWNLOAD_FAILURE
+    Write-Log "Installer download complete"
+    
+    VerifySignature $Target 
+
+    Write-Log "Installing $Target ..."
+    try {
+        $InstallLogFolder = "$LogFolder\NvidiaInstallLog"
+        $Arguments = "-s -n -log:$InstallLogFolder -loglevel:6"
+    
+        $p = Start-Process -FilePath $Target -ArgumentList $Arguments -PassThru
+        
+        $Timeout = 10 * 60 # in seconds. 10 minutes for timeout of the installation
+
+        # This is for testability. Start-Process mock returns a hashtable.
+        if (-not ($p -is [hashtable])) {
+            Wait-Process -InputObject $p -Timeout $Timeout -ErrorAction Stop
+        }
+
+        Handle-InstallationResult -ErrorCode $p.ExitCode
+    }
+    catch {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_EXCEPTION -ErrorMessage "Exception: $($_.ToString())"
+    }
+}
+
+function Handle-InstallationResult {
+    param(
+        [Parameter(Mandatory = $true)]
+        [int]$ErrorCode
+    )
+
+    if ($ErrorCode -ne 0 -and $ErrorCode -ne 1) {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_FAILED -ErrorMessage "GPU Driver Installation Failed! Code: $ErrorCode"
+    }
+
+    if ($ErrorCode -eq 0) {
+        Write-Log "GPU Driver Installation Success. Code: $ErrorCode"
+        # check if vm size is nv series. if so, set RebootNeeded to be true
+        try {
+            $Compute = Get-VmData
+            $vmSize = $Compute.vmSize
+        }
+        catch {
+            Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_SKU_INFO_NOT_FOUND -ErrorMessage "Failed to query the SKU information."
+        }
+
+        $IsNVSeries = $vmSize -ne $null -and $vmSize -match "standard_nv" #case insensitive
+
+        if ($IsNVSeries -eq $true) {
+            Write-Log "Reboot is needed for this GPU Driver..."
+            $global:RebootNeeded = $true
+        }
+    }
+    elseif ($ErrorCode -eq 1) {
+        Write-Log "GPU Driver Installation Success. Code: $ErrorCode. Reboot is needed for this installation..."
+        $global:RebootNeeded = $true
+    }
+
+    Remove-InstallerFile -InstallerPath $Target
+}
+  
+function Get-VmData {
+    [OutputType([hashtable])]
+
+    $arguments = @{
+        Headers = @{"Metadata" = "true" }
+        URI     = "http://169.254.169.254/metadata/instance/compute?api-version=2017-08-01"
+        Method  = "get"
+    }
+
+    try {
+        $Compute = Retry-Command -Command Invoke-RestMethod -Args $arguments -Retries 5 -RetryDelaySeconds 10
+    }
+    catch {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_SKU_INFO_NOT_FOUND -ErrorMessage "Failed to query the SKU information."
+    }
+
+    return $Compute
+}
+
+function VerifySignature([string] $targetFile) {
+    Write-Log "VerifySignature - Start"
+    Write-Log "Verifying signature for $targetFile"
+    $fileCertificate = Get-AuthenticodeSignature $targetFile
+
+    if ($fileCertificate.Status -ne "Valid") {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INVALID_SIGNATURE -ErrorMessage "Signature embedded in $($Target) is not valid."
+    }
+
+    if ($fileCertificate.SignerCertificate.Subject -eq $fileCertificate.SignerCertificate.Issuer) {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INVALID_SIGNATURE -ErrorMessage "Signer certificate's Subject in $($Target) matches the Issuer: The certificate is self-signed"
+    }
+
+    Write-Log "Signature for $targetFile is valid and is not self-signed"
+    Write-Log "VerifySignature - End"
+}
+
+function Remove-InstallerFile {
+    param(
+        [Parameter(Mandatory = $true)]
+        [string]$InstallerPath
+    )
+  
+    Write-Log "Attempting to remove installer file at $InstallerPath..."
+  
+    try {
+        Remove-Item -Path $InstallerPath -Force
+        Write-Log "Installer file removed successfully."
+    }
+    catch {
+        Write-Log "Failed to remove installer file. Error: $($_.Exception.Message)"
+    }
+}

staging/cse/windows/nvidiagpudriverfunc.tests.ps1

@@ -0,0 +1,278 @@
+BeforeAll {
+    . $PSScriptRoot\..\..\..\parts\windows\windowscsehelper.ps1
+    . $PSCommandPath.Replace('.tests.ps1', '.ps1')
+}
+
+Describe 'Start-InstallGPUDriver' {
+
+    BeforeEach {
+        $global:RebootNeeded = $false
+
+        Mock Set-ExitCode -MockWith {
+            Param(
+                [Parameter(Mandatory = $true)][int]
+                $ExitCode,
+                [Parameter(Mandatory = $true)][string]
+                $ErrorMessage
+            )
+            Write-Host "Set ExitCode to $ExitCode and exit. Error: $ErrorMessage"
+            throw $ExitCode
+        } -Verifiable
+
+        Mock DownloadFileOverHttp -MockWith {
+            param (
+                [string]$Url,
+                [string]$DestinationPath,
+                [int]$ExitCode
+            )
+            $DestinationPath | Should -Be $TargetPath
+        } -Verifiable
+
+        Mock VerifySignature -MockWith {
+            param (
+                [string]$targetFile
+            )
+            $targetFile | Should -Be $TargetPath
+        } -Verifiable
+
+        Mock Start-Process -MockWith {
+            param (
+                [string]$FilePath,
+                [string]$ArgumentList
+            )
+            $FilePath | Should -Be $TargetPath
+            $process = New-Object System.Diagnostics.Process
+            return @($process)
+        } -Verifiable
+
+        Mock Wait-Process -MockWith {
+            param (
+                [System.Diagnostics.Process[]]$InputObject,
+                [int]$Timeout
+            )
+            return
+        } -Verifiable
+
+        Mock Get-VmData -MockWith {
+            return @{ vmSize = "Standard_NV" }
+        } -Verifiable
+
+        Mock Remove-InstallerFile
+    }
+
+    Context 'When EnableInstall is false' {
+        It "Should skip installation" {
+            try {
+                Start-InstallGPUDriver -EnableInstall $false
+            }
+            catch {
+                Throw "Unexpected exception during UT: $_"
+            }
+            Assert-MockCalled -CommandName 'Set-ExitCode' -Exactly -Times 0
+        }
+    }
+
+    Context 'When EnableInstall is true' {
+        It "Should skip installation if GpuDriverURL is empty" {
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL ''
+            }
+            catch {
+                Write-Host "Expected exception: $_"
+            }
+            Assert-MockCalled -CommandName 'Set-ExitCode' -Exactly -Times 1 -ParameterFilter { $ExitCode -eq $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_URL_NOT_SET }
+        }
+
+        It "Should skip installation if GpuDriverURL does not point to an exe file" {
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL 'https://example.com/somefile.zip'
+            }
+            catch {
+                Write-Host "Expected exception: $_"
+            }
+            Assert-MockCalled -CommandName 'Set-ExitCode' -Exactly -Times 1 -ParameterFilter { $ExitCode -eq $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_URL_NOT_EXE }
+        }
+
+        It "Should exit when the signature is not valid" {
+            $GpuDriverURL = 'https://example.com/gpudriver.exe'
+            $TargetPath = "C:\AzureData\gpudriver.exe"
+
+            Mock VerifySignature -MockWith {
+                param (
+                    [string]$targetFile
+                )
+                $targetFile | Should -Be $TargetPath
+                Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INVALID_SIGNATURE -ErrorMessage "Signature embedded in $($Target) is not valid."
+            } -Verifiable
+            
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL $GpuDriverURL
+            }
+            catch {
+                Write-Host "Expected exception: $_"
+            }
+            Assert-MockCalled -CommandName 'Set-ExitCode' -Exactly -Times 1 -ParameterFilter { $ExitCode -eq $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INVALID_SIGNATURE }
+        }
+
+        It "Should exit when Start-Process does not return a valid output for Wait-Process to use as a valid argument." {
+            $GpuDriverURL = 'https://example.com/gpudriver.exe'
+            $TargetPath = "C:\AzureData\gpudriver.exe"
+
+            Mock Start-Process -MockWith {
+                param (
+                    [string]$FilePath,
+                    [string]$ArgumentList
+                )
+                return
+            } -Verifiable
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL $GpuDriverURL
+            }
+            catch {
+                Write-Host "Expected exception: $_"
+            }
+
+            Assert-MockCalled -CommandName 'DownloadFileOverHttp' -Exactly -Times 1 -ParameterFilter { $Url -eq $GpuDriverURL }
+            Assert-MockCalled -CommandName 'VerifySignature' -Exactly -Times 1 -ParameterFilter { $targetFile -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Start-Process' -Exactly -Times 1 -ParameterFilter { $FilePath -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Set-ExitCode' -Exactly -Times 1 -ParameterFilter { $ExitCode -eq $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_EXCEPTION }
+        }
+
+        It "Should run Wait-Process when Start-Process does not return a hashtable mock output" {
+            $GpuDriverURL = 'https://example.com/gpudriver.exe'
+            $TargetPath = "C:\AzureData\gpudriver.exe"
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL $GpuDriverURL
+            }
+            catch {
+                Throw "Unexpected exception during UT: $_"
+            }
+
+            Assert-MockCalled -CommandName 'DownloadFileOverHttp' -Exactly -Times 1 -ParameterFilter { $Url -eq $GpuDriverURL }
+            Assert-MockCalled -CommandName 'VerifySignature' -Exactly -Times 1 -ParameterFilter { $targetFile -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Start-Process' -Exactly -Times 1 -ParameterFilter { $FilePath -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Wait-Process' -Exactly -Times 1
+        }
+
+        It "Should exit when installation code is not 0 or 1" {
+            $GpuDriverURL = 'https://example.com/gpudriver.exe'
+            $TargetPath = "C:\AzureData\gpudriver.exe"
+
+            # The ExitCode of System.Diagnostics.Process is readonly so we have to create a custom hashtable.
+            # However, Wait-Process does not accept hastable as a parameter so we'll need to skip it in the code.
+            Mock Start-Process -MockWith {
+                param (
+                    [string]$FilePath,
+                    [string]$ArgumentList
+                )
+                $FilePath | Should -Be $TargetPath
+                $process = @{ ExitCode = 9 }
+                return $process
+            } -Verifiable
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL $GpuDriverURL
+            }
+            catch {
+                Write-Host "Expected exception: $_"
+            }
+
+            Assert-MockCalled -CommandName 'DownloadFileOverHttp' -Exactly -Times 1 -ParameterFilter { $Url -eq $GpuDriverURL }
+            Assert-MockCalled -CommandName 'VerifySignature' -Exactly -Times 1 -ParameterFilter { $targetFile -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Start-Process' -Exactly -Times 1 -ParameterFilter { $FilePath -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Wait-Process' -Exactly -Times 0
+            Assert-MockCalled -CommandName 'Set-ExitCode' -Exactly -Times 1 -ParameterFilter { $ExitCode -eq $global:WINDOWS_CSE_ERROR_GPU_DRIVER_INSTALLATION_FAILED }
+        }
+
+        It "Should set RebootNeeded to be true when installation code is 1" {
+            $GpuDriverURL = 'https://example.com/gpudriver.exe'
+            $TargetPath = "C:\AzureData\gpudriver.exe"
+
+            Mock Start-Process -MockWith {
+                param (
+                    [string]$FilePath,
+                    [string]$ArgumentList
+                )
+                $FilePath | Should -Be $TargetPath
+                $process = @{ ExitCode = 1 }
+                return $process
+            } -Verifiable
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL $GpuDriverURL
+            }
+            catch {
+                Write-Host "Expected exception: $_"
+            }
+
+            Assert-MockCalled -CommandName 'DownloadFileOverHttp' -Exactly -Times 1 -ParameterFilter { $Url -eq $GpuDriverURL }
+            Assert-MockCalled -CommandName 'VerifySignature' -Exactly -Times 1 -ParameterFilter { $targetFile -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Start-Process' -Exactly -Times 1 -ParameterFilter { $FilePath -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Wait-Process' -Exactly -Times 0
+            Assert-MockCalled -CommandName 'Remove-InstallerFile' -Exactly -Times 1
+            $global:RebootNeeded | Should -Be $true
+        }
+
+        It "Should run set RebootNeeded to be true when installation code is 0 and vm size is nv series" {
+            $GpuDriverURL = 'https://example.com/gpudriver.exe'
+            $TargetPath = "C:\AzureData\gpudriver.exe"
+
+            Mock Start-Process -MockWith {
+                param (
+                    [string]$FilePath,
+                    [string]$ArgumentList
+                )
+                $FilePath | Should -Be $TargetPath
+                $process = @{ ExitCode = 0 }
+                return $process
+            } -Verifiable
+
+            Mock Get-VmData -MockWith {
+                return @{ vmSize = "Standard_NV" }
+            } -Verifiable
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL $GpuDriverURL
+            }
+            catch {
+                Throw "Unexpected exception during UT: $_"
+            }
+
+            Assert-MockCalled -CommandName 'DownloadFileOverHttp' -Exactly -Times 1 -ParameterFilter { $Url -eq $GpuDriverURL }
+            Assert-MockCalled -CommandName 'VerifySignature' -Exactly -Times 1 -ParameterFilter { $targetFile -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Start-Process' -Exactly -Times 1 -ParameterFilter { $FilePath -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Wait-Process' -Exactly -Times 0
+            Assert-MockCalled -CommandName 'Remove-InstallerFile' -Exactly -Times 1
+            $global:RebootNeeded | Should -Be $true
+        }
+
+        It "Should run set RebootNeeded to be false when installation code is 0 and vm size is nc series" {
+            $GpuDriverURL = 'https://example.com/gpudriver.exe'
+            $TargetPath = "C:\AzureData\gpudriver.exe"
+
+            Mock Start-Process -MockWith {
+                param (
+                    [string]$FilePath,
+                    [string]$ArgumentList
+                )
+                $FilePath | Should -Be $TargetPath
+                $process = @{ ExitCode = 0 }
+                return $process
+            } -Verifiable
+
+            Mock Get-VmData -MockWith {
+                return @{ vmSize = "Standard_NC" }
+            } -Verifiable
+            try {
+                Start-InstallGPUDriver -EnableInstall $true -GpuDriverURL $GpuDriverURL
+            }
+            catch {
+                Throw "Unexpected exception during UT: $_"
+            }
+
+            Assert-MockCalled -CommandName 'DownloadFileOverHttp' -Exactly -Times 1 -ParameterFilter { $Url -eq $GpuDriverURL }
+            Assert-MockCalled -CommandName 'VerifySignature' -Exactly -Times 1 -ParameterFilter { $targetFile -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Start-Process' -Exactly -Times 1 -ParameterFilter { $FilePath -eq $TargetPath }
+            Assert-MockCalled -CommandName 'Wait-Process' -Exactly -Times 0
+            Assert-MockCalled -CommandName 'Remove-InstallerFile' -Exactly -Times 1
+            $global:RebootNeeded | Should -Be $false
+        }
+    }
+}

staging/cse/windows/provisioningscripts/loggenerator.ps1

@@ -87,6 +87,15 @@ $kLogFiles | Foreach-Object {
     Create-SymbolLinkFile -SrcFile (Join-Path "C:\k\" $_) -DestFile (Join-Path $aksLogFolder $_)
 }
 
+$nvidiaInstallLogFolder="C:\AzureData\NvidiaInstallLog"
+if (Test-Path $nvidiaInstallLogFolder) {
+    $logFiles=Get-ChildItem (Join-Path $nvidiaInstallLogFolder *.log)
+    $logFiles | Foreach-Object {
+        $fileName = [IO.Path]::GetFileName($_)
+        Create-SymbolLinkFile -SrcFile (Join-Path $nvidiaInstallLogFolder $fileName) -DestFile (Join-Path $aksLogFolder $fileName)
+    }
+}
+
 $calicoLogFolder="C:\CalicoWindows\logs\"
 if (Test-Path $calicoLogFolder) {
     $calicoLogFiles = @(