Pass worker_spec_command to mpi plugin to support horovod (flyteorg…

…#341) * wip Signed-off-by: byhsu <[email protected]> * add comment Signed-off-by: byhsu <[email protected]> * more comment Signed-off-by: byhsu <[email protected]> * fix style Signed-off-by: byhsu <[email protected]> --------- Signed-off-by: byhsu <[email protected]> Co-authored-by: byhsu <[email protected]>
eapolinario · Apr 19, 2023 · 46796b6 · 46796b6
1 parent f5f4182
commit 46796b6
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go b/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go
@@ -3,6 +3,7 @@ package mpi
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/plugins"
@@ -21,6 +22,8 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
+const workerSpecCommandKey = "worker_spec_command"
+
 type mpiOperatorResourceHandler struct {
 }
 
@@ -45,6 +48,7 @@ func (mpiOperatorResourceHandler) BuildIdentityResource(ctx context.Context, tas
 // Defines a func to create the full resource object that will be posted to k8s.
 func (mpiOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext) (client.Object, error) {
 	taskTemplate, err := taskCtx.TaskReader().Read(ctx)
+	taskTemplateConfig := taskTemplate.GetConfig()
 
 	if err != nil {
 		return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "unable to fetch task specification [%v]", err.Error())
@@ -69,11 +73,19 @@ func (mpiOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx plu
 	common.OverridePrimaryContainerName(podSpec, primaryContainerName, kubeflowv1.MPIJobDefaultContainerName)
 
 	// workersPodSpec is deepCopy of podSpec submitted by flyte
-	// WorkerPodSpec doesn't need any Argument & command. It will be trigger from launcher pod
 	workersPodSpec := podSpec.DeepCopy()
 
+	// If users don't specify "worker_spec_command" in the task config, the command/args are empty.
+	// However, in some cases, the workers need command/args.
+	// For example, in horovod tasks, each worker runs a command launching ssh daemon.
+
+	workerSpecCommand := []string{}
+	if val, ok := taskTemplateConfig[workerSpecCommandKey]; ok {
+		workerSpecCommand = strings.Split(val, " ")
+	}
+
 	for k := range workersPodSpec.Containers {
-		workersPodSpec.Containers[k].Args = []string{}
+		workersPodSpec.Containers[k].Args = workerSpecCommand
 		workersPodSpec.Containers[k].Command = []string{}
 	}