Make NVDRIVER configurable from env file

Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>

Author: ArangoGutierrez

pkg/provisioner/dependency.go

@@ -50,7 +50,7 @@ var (
 type ProvisionFunc func(tpl *bytes.Buffer, env v1alpha1.Environment) error
 
 func nvdriver(tpl *bytes.Buffer, env v1alpha1.Environment) error {
-	nvdriver := templates.NewNvDriver()
+	nvdriver := templates.NewNvDriver(env)
 	return nvdriver.Execute(tpl, env)
 }
 

pkg/provisioner/templates/nv-driver.go

@@ -34,21 +34,27 @@ wget https://developer.download.nvidia.com/compute/cuda/repos/$distribution/x86_
 sudo dpkg -i cuda-keyring_1.1-1_all.deb
 
 with_retry 3 10s sudo apt-get update
-install_packages_with_retry cuda-drivers
+install_packages_with_retry cuda-drivers{{if .Version}}={{.Version}}{{end}}
 
 nvidia-smi -L
 `
 
 type NvDriver struct {
+	Version string
 }
 
-func NewNvDriver() *NvDriver {
-	return &NvDriver{}
+func NewNvDriver(env v1alpha1.Environment) *NvDriver {
+	if env.Spec.NVIDIADriver.Version == "" {
+		return &NvDriver{}
+	}
+	return &NvDriver{
+		Version: env.Spec.NVIDIADriver.Version,
+	}
 }
 
 func (t *NvDriver) Execute(tpl *bytes.Buffer, env v1alpha1.Environment) error {
 	nvDriverTemplate := template.Must(template.New("nv-driver").Parse(NvDriverTemplate))
-	err := nvDriverTemplate.Execute(tpl, &NvDriver{})
+	err := nvDriverTemplate.Execute(tpl, t)
 	if err != nil {
 		return fmt.Errorf("failed to execute nv-driver template: %v", err)
 	}