Added more unit test

Signed-off-by: Yuvraj <[email protected]>
flyteorg · Aug 30, 2021 · 7492373 · 7492373
1 parent 67e330b
commit 7492373
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/go/tasks/plugins/k8s/kfoperators/common/common_operator.go b/go/tasks/plugins/k8s/kfoperators/common/common_operator.go
@@ -132,7 +132,7 @@ func GetLogs(taskType string, name string, namespace string,
 		taskLogs = append(taskLogs, workerLog.TaskLogs...)
 	}
 
-	if taskType != MPITaskType {
+	if taskType != MPITaskType && taskType != PytorchTaskType {
 		// get all parameter servers logs
 		for psReplicaIndex := int32(0); psReplicaIndex < psReplicasCount; psReplicaIndex++ {
 			psReplicaLog, err := logPlugin.GetTaskLogs(tasklog.Input{

diff --git a/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go b/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go
@@ -118,10 +118,16 @@ func TestGetLogs(t *testing.T) {
 
 	workers := int32(1)
 	launcher := int32(1)
-	slots := int32(1)
 
-	jobLogs, err := GetLogs(MPITaskType, "test", "mpi-namespace", workers, launcher, slots)
+	jobLogs, err := GetLogs(MPITaskType, "test", "mpi-namespace", workers, launcher, 0)
 	assert.NoError(t, err)
 	assert.Equal(t, 1, len(jobLogs))
 	assert.Equal(t, fmt.Sprintf("k8s.com/#!/log/%s/%s-worker-0/pod?namespace=mpi-namespace", "mpi-namespace", "test"), jobLogs[0].Uri)
+
+	jobLogs, err = GetLogs(PytorchTaskType, "test", "pytorch-namespace", workers, launcher, 0)
+	assert.NoError(t, err)
+	assert.Equal(t, 2, len(jobLogs))
+	assert.Equal(t, fmt.Sprintf("k8s.com/#!/log/%s/%s-master-0/pod?namespace=pytorch-namespace", "pytorch-namespace", "test"), jobLogs[0].Uri)
+	assert.Equal(t, fmt.Sprintf("k8s.com/#!/log/%s/%s-worker-0/pod?namespace=pytorch-namespace", "pytorch-namespace", "test"), jobLogs[1].Uri)
+
 }