add simple test for workerspec arg + expected resources

Signed-off-by: ljstrnadiii <[email protected]>

go/tasks/plugins/k8s/dask/dask.go

@@ -189,9 +189,8 @@ func createWorkerSpec(cluster plugins.DaskWorkerGroup, defaults defaults) (*dask
 		}
 		// If limits includes gpu, assume dask cuda worker cli startup
 		// https://docs.rapids.ai/api/dask-cuda/nightly/quickstart.html#dask-cuda-worker
-		// TODO: is this how gpu resources are called?
-		if limits.Gpu() != nil {
-			workerArgs[0] = "dask cuda worker"
+		if !limits.Name(flytek8s.ResourceNvidiaGPU, "0").IsZero() {
+			workerArgs[0] = "dask-cuda-worker"
 		}
 	}
 

go/tasks/plugins/k8s/dask/dask_test.go

@@ -340,6 +340,34 @@ func TestBuildResourceDaskDefaultResoureRequirements(t *testing.T) {
 	assert.Contains(t, workerSpec.Containers[0].Args, "2G")
 }
 
+func TestBuildResourceGPUCudaWorkerArgs(t *testing.T) {
+	protobufResources := core.Resources{
+		Limits: []*core.Resources_ResourceEntry{
+			{
+				Name:  core.Resources_GPU,
+				Value: "1",
+			},
+		},
+	}
+	expectedResources, _ := flytek8s.ToK8sResourceRequirements(&protobufResources)
+
+	flyteWorkflowResources := v1.ResourceRequirements{}
+
+	daskResourceHandler := daskResourceHandler{}
+	taskTemplate := dummyDaskTaskTemplate("", &protobufResources)
+	taskContext := dummyDaskTaskContext(taskTemplate, &flyteWorkflowResources, false)
+	resource, err := daskResourceHandler.BuildResource(context.TODO(), taskContext)
+	assert.Nil(t, err)
+	assert.NotNil(t, resource)
+	daskJob, ok := resource.(*daskAPI.DaskJob)
+	assert.True(t, ok)
+
+	// Default Workers
+	workerSpec := daskJob.Spec.Cluster.Spec.Worker.Spec
+	assert.Equal(t, *expectedResources, workerSpec.Containers[0].Resources)
+	assert.Contains(t, workerSpec.Containers[0].Args, "dask-cuda-worker")
+}
+
 func TestBuildResourcesDaskCustomResoureRequirements(t *testing.T) {
 	protobufResources := core.Resources{
 		Requests: []*core.Resources_ResourceEntry{