Add support for VGPU in Azure

data/cloud/azure.yaml

@@ -0,0 +1,9 @@
+profile::gpu::install::vgpu::installer: bin
+profile::gpu::install::vgpu::bin::source: https://go.microsoft.com/fwlink/?linkid=874272
+profile::gpu::install::vgpu::bin::gridd_content: |
+  IgnoreSP=FALSE
+  EnableUI=FALSE
+profile::gpu::install::vgpu::grid_vgpu_types:
+  - "^Standard_NV(6|12|18|36|72)ad[m]*s_A10_v5$"
+  - "^Standard_NV(12|24|48)s_v3$"
+  - "^Standard_NC(4|8|16|64)as_T4_v3$"

hiera.yaml

@@ -41,8 +41,10 @@ hierarchy:
       hostname: "%{facts.networking.hostname}"
   - name: "Software stack"
     path: "software_stack/%{facts.software_stack}.yaml"
-  - name: "Cloud provider"
+  - name: "Cloud provider region"
     path: "cloud/%{facts.cloud.provider}/%{facts.cloud.region}.yaml"
+  - name: "Cloud provider"
+    path: "cloud/%{facts.cloud.provider}.yaml"
   - name: "OS version"
     path: "os/%{facts.os.family}/%{facts.os.release.major}.yaml"
   - name: "Other YAML hierarchy levels"

site/profile/functions/is_grid_vgpu.pp

@@ -0,0 +1,9 @@
+function profile::is_grid_vgpu() >> Bool {
+  if $facts['nvidia_grid_vgpu'] {
+    true
+  } else {
+    $grid_vgpu_types = lookup('profile::gpu::install::vgpu::grid_vgpu_types', undef, undef, [])
+    $type = lookup('terraform.self.specs.type')
+    $grid_vgpu_types.any|$regex| { $type =~ Regexp($regex) }
+  }
+}

site/profile/manifests/gpu.pp

@@ -1,7 +1,7 @@
 class profile::gpu {
   if $facts['nvidia_gpu_count'] > 0 {
     require profile::gpu::install
-    if ! $facts['nvidia_grid_vgpu'] {
+    if ! profile::is_grid_vgpu() {
       service { 'nvidia-persistenced':
         ensure => 'running',
         enable => true,
@@ -32,7 +32,7 @@
     source_pp => 'puppet:///modules/profile/gpu/nvidia-gpu.pp',
   }
 
-  if ! $facts['nvidia_grid_vgpu'] {
+  if ! profile::is_grid_vgpu() {
     include profile::gpu::install::passthrough
     Class['profile::gpu::install::passthrough'] -> Exec['dkms_nvidia']
   } else {
@@ -43,7 +43,7 @@
   # Binary installer do not build drivers with DKMS
   $installer = lookup('profile::gpu::install::vgpu::installer', undef, undef, '')
   $nvidia_kmod = ['nvidia', 'nvidia_drm', 'nvidia_modeset', 'nvidia_uvm']
-  if ! $facts['nvidia_grid_vgpu'] or $installer != 'bin' {
+  if ! profile::is_grid_vgpu() or $installer != 'bin' {
     exec { 'dkms_nvidia':
       command => "dkms autoinstall -m nvidia -k ${facts['kernelrelease']}",
       path    => ['/usr/bin', '/usr/sbin'],
@@ -218,6 +218,7 @@
 class profile::gpu::install::vgpu (
   Enum['rpm', 'bin', 'none'] $installer = 'none',
   String $nvidia_ml_py_version = '11.515.75',
+  Array[String] $grid_vgpu_types = [],
 ) {
   if $installer == 'rpm' {
     include profile::gpu::install::vgpu::rpm
@@ -274,7 +275,8 @@
 
 class profile::gpu::install::vgpu::bin (
   String $source,
-  String $gridd_source,
+  String $gridd_content = undef,
+  String $gridd_source = undef,
 ) {
   exec { 'vgpu-driver-install-bin':
     command => "curl -L ${source} -o /tmp/NVIDIA-driver.run && sh /tmp/NVIDIA-driver.run --ui=none --no-questions --disable-nouveau && rm /tmp/NVIDIA-driver.run", # lint:ignore:140chars
@@ -290,11 +292,19 @@
     ],
   }
 
+  if $gridd_content {
+    $gridd_definition = { 'content' => $gridd_content }
+  } elsif $gridd_source {
+    $gridd_definition = { 'source' => $gridd_source }
+  } else {
+    $gridd_definition = {}
+  }
+
   file { '/etc/nvidia/gridd.conf':
     ensure => file,
     mode   => '0644',
     owner  => 'root',
     group  => 'root',
-    source => $gridd_source,
+    *      => $gridd_definition,
   }
 }