Create a parameter for nvidia-ml-py version

ComputeCanada · Jun 26, 2023 · 5ad52db · 5ad52db
1 parent b0ce9ad
commit 5ad52db
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/site/profile/manifests/gpu.pp b/site/profile/manifests/gpu.pp
@@ -132,6 +132,7 @@
 
 class profile::gpu::install::vgpu (
   Enum['rpm', 'bin', 'none'] $installer = 'none',
+  String $nvidia_ml_py_version = '11.515.75',
 ) {
   if $installer == 'rpm' {
     include profile::gpu::install::vgpu::rpm
@@ -146,7 +147,7 @@
   $py3_version = lookup('os::redhat::python3::version')
 
   exec { 'pip install nvidia-ml-py':
-    command => "/usr/bin/pip${py3_version} install --force-reinstall nvidia-ml-py==11.515.75",
+    command => "/usr/bin/pip${py3_version} install --force-reinstall nvidia-ml-py==${nvidia_ml_py_version}",
     creates => "/usr/local/lib/python${py3_version}/site-packages/pynvml.py",
     before  => Service['slurm-job-exporter'],
     require => Package['python3'],