Windows AMI - Cuda 12.6 (#5880)

tinglvv · web-flow · commit fbebf8226935 · 2024-11-08T23:08:36.000-05:00
related to pytorch/pytorch#138440 cc @atalman
diff --git a/aws/ami/windows/scripts/Installers/Install-CUDA-Tools.ps1 b/aws/ami/windows/scripts/Installers/Install-CUDA-Tools.ps1
@@ -34,6 +34,11 @@ Switch ($cudaVersion) {
     $cudnn_subfolder = "cudnn-windows-x86_64-9.1.0.70_cuda12-archive"
     $installerArgs += " cuda_profiler_api_$cudaVersion nvjitlink_$cudaVersion"
   }
+  "12.6" {
+    $toolkitInstaller = "cuda_12.6.2_560.94_windows.exe"
+    $cudnn_subfolder = "cudnn-windows-x86_64-9.5.0.50_cuda12-archive"
+    $installerArgs += " cuda_profiler_api_$cudaVersion nvjitlink_$cudaVersion"
+  }
 }
 
 
diff --git a/aws/ami/windows/windows.pkr.hcl b/aws/ami/windows/windows.pkr.hcl
@@ -120,6 +120,13 @@ build {
     ]
   }
 
+  provisioner "powershell" {
+    environment_vars = ["CUDA_VERSION=12.6"]
+    scripts = [
+      "${path.root}/scripts/Installers/Install-CUDA-Tools.ps1",
+    ]
+  }
+
   # Uninstall Windows Defender, it brings more trouble than it's worth. Do this
   # last as it screws up the installation of other services like sshd somehow
   provisioner "powershell" {

Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,11 @@ Switch ($cudaVersion) {`
`34`	`34`	`$cudnn_subfolder = "cudnn-windows-x86_64-9.1.0.70_cuda12-archive"`
`35`	`35`	`$installerArgs += " cuda_profiler_api_$cudaVersion nvjitlink_$cudaVersion"`
`36`	`36`	`}`
	`37`	`+ "12.6" {`
	`38`	`+ $toolkitInstaller = "cuda_12.6.2_560.94_windows.exe"`
	`39`	`+ $cudnn_subfolder = "cudnn-windows-x86_64-9.5.0.50_cuda12-archive"`
	`40`	`+ $installerArgs += " cuda_profiler_api_$cudaVersion nvjitlink_$cudaVersion"`
	`41`	`+ }`
`37`	`42`	`}`
`38`	`43`
`39`	`44`
Original file line number	Diff line number	Diff line change
`@@ -120,6 +120,13 @@ build {`
`120`	`120`	`]`
`121`	`121`	`}`
`122`	`122`
	`123`	`+ provisioner "powershell" {`
	`124`	`+ environment_vars = ["CUDA_VERSION=12.6"]`
	`125`	`+ scripts = [`
	`126`	`+ "${path.root}/scripts/Installers/Install-CUDA-Tools.ps1",`
	`127`	`+ ]`
	`128`	`+ }`
	`129`	`+`
`123`	`130`	`# Uninstall Windows Defender, it brings more trouble than it's worth. Do this`
`124`	`131`	`# last as it screws up the installation of other services like sshd somehow`
`125`	`132`	`provisioner "powershell" {`