Test with GPU and m+t4 machines

iterative/resource_task.go

@@ -55,7 +55,7 @@ func resourceTask() *schema.Resource {
 				Type:     schema.TypeInt,
 				ForceNew: true,
 				Optional: true,
-				Default:  30,
+				Default:  -1,
 			},
 			"spot": {
 				Type:     schema.TypeFloat,

task/aws/resources/data_source_image.go

@@ -34,7 +34,7 @@ func (i *Image) Read(ctx context.Context) error {
 	image := i.Identifier
 	images := map[string]string{
 		"ubuntu": "ubuntu@099720109477:x86_64:*ubuntu/images/hvm-ssd/ubuntu-focal-20.04*",
-		"nvidia": "ubuntu@898082745236:x86_64:Deep Learning AMI GPU CUDA 11.2.1 (Ubuntu 20.04) 20220306",
+		"nvidia": "ubuntu@898082745236:x86_64:Deep Learning AMI GPU CUDA 11.3.1 (Ubuntu 20.04) 20220303",
 	}
 	if val, ok := images[image]; ok {
 		image = val

task/aws/resources/resource_launch_template.go

@@ -86,7 +86,6 @@ func (l *LaunchTemplate) Create(ctx context.Context) error {
 					Ebs: &types.LaunchTemplateEbsBlockDeviceRequest{
 						DeleteOnTermination: aws.Bool(true),
 						Encrypted:           aws.Bool(false),
-						VolumeSize:          aws.Int32(int32(l.Attributes.Size.Storage)),
 						VolumeType:          types.VolumeType("gp2"),
 					},
 				},
@@ -110,6 +109,10 @@ func (l *LaunchTemplate) Create(ctx context.Context) error {
 		},
 	}
 
+	if size := l.Attributes.Size.Storage; size > 0 {
+		input.LaunchTemplateData.BlockDeviceMappings[0].Ebs.VolumeSize = aws.Int32(int32(size))
+	}
+
 	if _, err := l.Client.Services.EC2.CreateLaunchTemplate(ctx, &input); err != nil {
 		var e smithy.APIError
 		if errors.As(err, &e) && e.ErrorCode() == "InvalidLaunchTemplateName.AlreadyExistsException" {

task/az/resources/resource_virtual_machine_scale_set.go

@@ -85,7 +85,7 @@ func (v *VirtualMachineScaleSet) Create(ctx context.Context) error {
 	image := v.Attributes.Environment.Image
 	images := map[string]string{
 		"ubuntu": "ubuntu@Canonical:0001-com-ubuntu-server-focal:20_04-lts:latest",
-		"nvidia": "ubuntu@nvidia:ngc_base_image_version_b:gen2_21-11-0:latest#plan",
+		"nvidia": "ubuntu@microsoft-dsvm:ubuntu-2004:2004-gen2:latest",
 	}
 	if val, ok := images[image]; ok {
 		image = val
@@ -145,7 +145,6 @@ func (v *VirtualMachineScaleSet) Create(ctx context.Context) error {
 					OsDisk: &compute.VirtualMachineScaleSetOSDisk{
 						Caching:      compute.CachingTypesReadWrite,
 						CreateOption: compute.DiskCreateOptionTypesFromImage,
-						DiskSizeGB:   to.Int32Ptr(int32(v.Attributes.Size.Storage)),
 						ManagedDisk: &compute.VirtualMachineScaleSetManagedDiskParameters{
 							StorageAccountType: compute.StorageAccountTypesStandardLRS,
 						},
@@ -192,6 +191,10 @@ func (v *VirtualMachineScaleSet) Create(ctx context.Context) error {
 		},
 	}
 
+	if size := v.Attributes.Size.Storage; size > 0 {
+		settings.VirtualMachineScaleSetProperties.VirtualMachineProfile.StorageProfile.OsDisk.DiskSizeGB = to.Int32Ptr(int32(size))
+	}
+
 	if plan == "#plan" {
 		settings.Plan = &compute.Plan{
 			Publisher: to.StringPtr(publisher),

task/common/machine/script.go

@@ -102,6 +102,8 @@ fi
 
 rclone copy "$RCLONE_REMOTE/data" /tmp/tpi-task
 
+yes | /etc/profile.d/install-driver-prompt.sh # for GCP GPU machines
+
 sudo systemctl daemon-reload
 sudo systemctl enable tpi-task.service --now
 sudo systemctl disable --now apt-daily.timer

task/gcp/resources/data_source_image.go

@@ -32,7 +32,7 @@ func (i *Image) Read(ctx context.Context) error {
 	image := i.Identifier
 	images := map[string]string{
 		"ubuntu": "ubuntu@ubuntu-os-cloud/ubuntu-2004-lts",
-		"nvidia": "ubuntu@nvidia-ngc-public/nvidia-gpu-cloud-image-20211105",
+		"nvidia": "ubuntu@deeplearning-platform-release/common-cu113-ubuntu-2004",
 	}
 	if val, ok := images[image]; ok {
 		image = val

task/gcp/resources/resource_instance_template.go

@@ -131,7 +131,6 @@ func (i *InstanceTemplate) Create(ctx context.Context) error {
 					Mode:       "READ_WRITE",
 					InitializeParams: &compute.AttachedDiskInitializeParams{
 						SourceImage: i.Dependencies.Image.Resource.SelfLink,
-						DiskSizeGb:  int64(i.Attributes.Size.Storage),
 						DiskType:    "pd-balanced",
 					},
 				},
@@ -171,6 +170,10 @@ func (i *InstanceTemplate) Create(ctx context.Context) error {
 		},
 	}
 
+	if size := i.Attributes.Size.Storage; size > 0 {
+		definition.Properties.Disks[0].InitializeParams.DiskSizeGb = int64(size)
+	}
+
 	insertOperation, err := i.Client.Services.Compute.InstanceTemplates.Insert(i.Client.Credentials.ProjectID, definition).Do()
 	if err != nil {
 		if strings.HasSuffix(err.Error(), "alreadyExists") {

task/k8s/resources/resource_job.go

@@ -62,13 +62,13 @@ func (j *Job) Create(ctx context.Context) error {
 		"m":       "8-32000",
 		"l":       "32-128000",
 		"xl":      "64-256000",
-		"m+t4":    "4-16000+nvidia-tesla-t4*1",
-		"m+k80":   "4-64000+nvidia-tesla-k80*1",
-		"l+k80":   "32-512000+nvidia-tesla-k80*8",
-		"xl+k80":  "64-768000+nvidia-tesla-k80*16",
-		"m+v100":  "8-64000+nvidia-tesla-v100*1",
-		"l+v100":  "32-256000+nvidia-tesla-v100*4",
-		"xl+v100": "64-512000+nvidia-tesla-v100*8",
+		"m+t4":    "4-16000+nvidia*1",
+		"m+k80":   "4-64000+nvidia*1",
+		"l+k80":   "32-512000+nvidia*8",
+		"xl+k80":  "64-768000+nvidia*16",
+		"m+v100":  "8-64000+nvidia*1",
+		"l+v100":  "32-256000+nvidia*4",
+		"xl+v100": "64-512000+nvidia*8",
 	}
 	if val, ok := sizes[size]; ok {
 		size = val
@@ -77,7 +77,7 @@ func (j *Job) Create(ctx context.Context) error {
 	image := j.Attributes.Task.Environment.Image
 	images := map[string]string{
 		"ubuntu": "ubuntu",
-		"nvidia": "nvidia/cuda",
+		"nvidia": "nvidia/cuda:11.3.0-cudnn8-runtime-ubuntu20.04",
 	}
 	if val, ok := images[image]; ok {
 		image = val
@@ -91,7 +91,7 @@ func (j *Job) Create(ctx context.Context) error {
 	// Define the accelerator settings (i.e. GPU type, model, ...)
 	jobNodeSelector := map[string]string{}
 	jobAccelerator := match[3]
-	jobGPUType := "kubernetes.io/gpu"
+	jobGPUType := "nvidia.com/gpu"
 	jobGPUCount := match[4]
 
 	// Define the dynamic resource allocation limits for the job pods.

task/k8s/resources/resource_persistent_volume_claim.go

@@ -14,7 +14,7 @@ import (
 	"terraform-provider-iterative/task/k8s/client"
 )
 
-func NewPersistentVolumeClaim(client *client.Client, identifier common.Identifier, storageClass string, size uint64, many bool) *PersistentVolumeClaim {
+func NewPersistentVolumeClaim(client *client.Client, identifier common.Identifier, storageClass string, size int, many bool) *PersistentVolumeClaim {
 	p := new(PersistentVolumeClaim)
 	p.Client = client
 	p.Identifier = identifier.Long()
@@ -29,7 +29,7 @@ type PersistentVolumeClaim struct {
 	Identifier string
 	Attributes struct {
 		StorageClass string
-		Size         uint64
+		Size         int
 		Many         bool
 	}
 	Dependencies struct{}
@@ -42,6 +42,11 @@ func (p *PersistentVolumeClaim) Create(ctx context.Context) error {
 		accessMode = kubernetes_core.ReadWriteMany
 	}
 
+	size := p.Attributes.Size
+	if size <= 0 {
+		size = 1 // Most StorageClasses disregard size anyway
+	}
+
 	persistentVolumeClaimInput := kubernetes_core.PersistentVolumeClaim{
 		ObjectMeta: kubernetes_meta.ObjectMeta{
 			Name:        p.Identifier,
@@ -53,7 +58,7 @@ func (p *PersistentVolumeClaim) Create(ctx context.Context) error {
 			AccessModes: []kubernetes_core.PersistentVolumeAccessMode{accessMode},
 			Resources: kubernetes_core.ResourceRequirements{
 				Requests: kubernetes_core.ResourceList{
-					kubernetes_core.ResourceStorage: kubernetes_resource.MustParse(strconv.Itoa(int(p.Attributes.Size)) + "G"),
+					kubernetes_core.ResourceStorage: kubernetes_resource.MustParse(strconv.Itoa(size) + "G"),
 				},
 			},
 		},

task/k8s/task.go

@@ -31,7 +31,7 @@ func New(ctx context.Context, cloud common.Cloud, identifier common.Identifier,
 	}
 
 	persistentVolumeClaimStorageClass := ""
-	persistentVolumeClaimSize := uint64(task.Size.Storage)
+	persistentVolumeClaimSize := task.Size.Storage
 	persistentVolumeDirectory := task.Environment.Directory
 
 	match := regexp.MustCompile(`^([^:]+):(?:(\d+):)?(.+)$`).FindStringSubmatch(task.Environment.Directory)
@@ -42,7 +42,7 @@ func New(ctx context.Context, cloud common.Cloud, identifier common.Identifier,
 			if err != nil {
 				return nil, err
 			}
-			persistentVolumeClaimSize = uint64(number)
+			persistentVolumeClaimSize = int(number)
 		}
 		persistentVolumeDirectory = match[3]
 	}

task/task_test.go

@@ -76,15 +76,14 @@ func TestTask(t *testing.T) {
 
 			task := common.Task{
 				Size: common.Size{
-					Machine: "m",
-					Storage: 30,
+					Machine: "m+t4",
 				},
 				Environment: common.Environment{
-					Image: "ubuntu",
-					Script: `#!/bin/bash
-					    mkdir cache
+					Image: "nvidia",
+					Script: `#!/bin/sh -e
+						nvidia-smi
+						mkdir --parents cache output
 						touch cache/file
-						mkdir output
 						echo "$ENVIRONMENT_VARIABLE_DATA" | tee --append output/file
 						sleep 60
 						cat output/file
@@ -132,6 +131,7 @@ func TestTask(t *testing.T) {
 			for assert.NoError(t, newTask.Read(ctx)) {
 				logs, err := newTask.Logs(ctx)
 				require.NoError(t, err)
+				t.Log(logs)
 
 				for _, log := range logs {
 					if strings.Contains(log, oldData) &&
@@ -140,6 +140,14 @@ func TestTask(t *testing.T) {
 					}
 				}
 
+				status, err := newTask.Status(ctx)
+				require.NoError(t, err)
+				t.Log(status)
+
+				if status[common.StatusCodeFailed] > 0 {
+					break
+				}
+
 				time.Sleep(10 * time.Second)
 			}