Skip to content
This repository was archived by the owner on Oct 9, 2023. It is now read-only.

Commit 9a2bbba

Browse files
yubofredwangYubo Wang
yubofredwang
and
Yubo Wang
authored
Change kubeflow plugins to allow settings specs for different replica (#345)
* change pytorch plugin to accept new pytorch task idl Signed-off-by: Yubo Wang <yubwang@linkedin.com> * merge elastic config in Signed-off-by: Yubo Wang <yubwang@linkedin.com> * add unit tests for pytorch Signed-off-by: Yubo Wang <yubwang@linkedin.com> * add tfjob Signed-off-by: Yubo Wang <yubwang@linkedin.com> * add mpi job Signed-off-by: Yubo Wang <yubwang@linkedin.com> * add test to commone operator Signed-off-by: Yubo Wang <yubwang@linkedin.com> * update flyteidl Signed-off-by: Yubo Wang <yubwang@linkedin.com> * add function header comments Signed-off-by: Yubo Wang <yubwang@linkedin.com> * fix lint Signed-off-by: Yubo Wang <yubwang@linkedin.com> --------- Signed-off-by: Yubo Wang <yubwang@linkedin.com> Co-authored-by: Yubo Wang <yubwang@linkedin.com>
1 parent 76a80ec commit 9a2bbba

10 files changed

Lines changed: 975 additions & 121 deletions

File tree

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ require (
1212
github.com/aws/aws-sdk-go-v2/service/athena v1.0.0
1313
github.com/bstadlbauer/dask-k8s-operator-go-client v0.1.0
1414
github.com/coocood/freecache v1.1.1
15-
github.com/flyteorg/flyteidl v1.3.19
15+
github.com/flyteorg/flyteidl v1.5.2
1616
github.com/flyteorg/flytestdlib v1.0.15
1717
github.com/go-test/deep v1.0.7
1818
github.com/golang/protobuf v1.5.2

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -232,8 +232,8 @@ github.com/evanphx/json-patch v4.12.0+incompatible/go.mod h1:50XU6AFN0ol/bzJsmQL
232232
github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
233233
github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
234234
github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
235-
github.com/flyteorg/flyteidl v1.3.19 h1:i79Dh7UoP8Z4LEJ2ox6jlfZVJtFZ+r4g84CJj1gh22Y=
236-
github.com/flyteorg/flyteidl v1.3.19/go.mod h1:Pkt2skI1LiHs/2ZoekBnyPhuGOFMiuul6HHcKGZBsbM=
235+
github.com/flyteorg/flyteidl v1.5.2 h1:DZPzYkTg92qA4e17fd0ZW1M+gh1gJKh/VOK+F4bYgM8=
236+
github.com/flyteorg/flyteidl v1.5.2/go.mod h1:ckLjB51moX4L0oQml+WTCrPK50zrJf6IZJ6LPC0RB4I=
237237
github.com/flyteorg/flytestdlib v1.0.15 h1:kv9jDQmytbE84caY+pkZN8trJU2ouSAmESzpTEhfTt0=
238238
github.com/flyteorg/flytestdlib v1.0.15/go.mod h1:ghw/cjY0sEWIIbyCtcJnL/Gt7ZS7gf9SUi0CCPhbz3s=
239239
github.com/flyteorg/stow v0.3.6 h1:jt50ciM14qhKBaIrB+ppXXY+SXB59FNREFgTJqCyqIk=

‎go/tasks/plugins/k8s/kfoperators/common/common_operator.go‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,11 @@ import (
55
"sort"
66
"time"
77

8+
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/flytek8s"
89
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/tasklog"
910

1011
"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
12+
kfplugins "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/plugins/kubeflow"
1113
flyteerr "github.com/flyteorg/flyteplugins/go/tasks/errors"
1214
"github.com/flyteorg/flyteplugins/go/tasks/logs"
1315
pluginsCore "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core"
@@ -21,6 +23,12 @@ const (
2123
PytorchTaskType = "pytorch"
2224
)
2325

26+
type ReplicaEntry struct {
27+
PodSpec *v1.PodSpec
28+
ReplicaNum int32
29+
RestartPolicy commonOp.RestartPolicy
30+
}
31+
2432
// ExtractMPICurrentCondition will return the first job condition for MPI
2533
func ExtractMPICurrentCondition(jobConditions []commonOp.JobCondition) (commonOp.JobCondition, error) {
2634
if jobConditions != nil {
@@ -180,3 +188,72 @@ func OverridePrimaryContainerName(podSpec *v1.PodSpec, primaryContainerName stri
180188
}
181189
}
182190
}
191+
192+
// ParseRunPolicy converts a kubeflow plugin RunPolicy object to a k8s RunPolicy object.
193+
func ParseRunPolicy(flyteRunPolicy kfplugins.RunPolicy) commonOp.RunPolicy {
194+
runPolicy := commonOp.RunPolicy{}
195+
if flyteRunPolicy.GetBackoffLimit() != 0 {
196+
var backoffLimit = flyteRunPolicy.GetBackoffLimit()
197+
runPolicy.BackoffLimit = &backoffLimit
198+
}
199+
var cleanPodPolicy = ParseCleanPodPolicy(flyteRunPolicy.GetCleanPodPolicy())
200+
runPolicy.CleanPodPolicy = &cleanPodPolicy
201+
if flyteRunPolicy.GetActiveDeadlineSeconds() != 0 {
202+
var ddlSeconds = int64(flyteRunPolicy.GetActiveDeadlineSeconds())
203+
runPolicy.ActiveDeadlineSeconds = &ddlSeconds
204+
}
205+
if flyteRunPolicy.GetTtlSecondsAfterFinished() != 0 {
206+
var ttl = flyteRunPolicy.GetTtlSecondsAfterFinished()
207+
runPolicy.TTLSecondsAfterFinished = &ttl
208+
}
209+
210+
return runPolicy
211+
}
212+
213+
// Get k8s clean pod policy from flyte kubeflow plugins clean pod policy.
214+
func ParseCleanPodPolicy(flyteCleanPodPolicy kfplugins.CleanPodPolicy) commonOp.CleanPodPolicy {
215+
cleanPodPolicyMap := map[kfplugins.CleanPodPolicy]commonOp.CleanPodPolicy{
216+
kfplugins.CleanPodPolicy_CLEANPOD_POLICY_NONE: commonOp.CleanPodPolicyNone,
217+
kfplugins.CleanPodPolicy_CLEANPOD_POLICY_ALL: commonOp.CleanPodPolicyAll,
218+
kfplugins.CleanPodPolicy_CLEANPOD_POLICY_RUNNING: commonOp.CleanPodPolicyRunning,
219+
}
220+
return cleanPodPolicyMap[flyteCleanPodPolicy]
221+
}
222+
223+
// Get k8s restart policy from flyte kubeflow plugins restart policy.
224+
func ParseRestartPolicy(flyteRestartPolicy kfplugins.RestartPolicy) commonOp.RestartPolicy {
225+
restartPolicyMap := map[kfplugins.RestartPolicy]commonOp.RestartPolicy{
226+
kfplugins.RestartPolicy_RESTART_POLICY_NEVER: commonOp.RestartPolicyNever,
227+
kfplugins.RestartPolicy_RESTART_POLICY_ON_FAILURE: commonOp.RestartPolicyOnFailure,
228+
kfplugins.RestartPolicy_RESTART_POLICY_ALWAYS: commonOp.RestartPolicyAlways,
229+
}
230+
return restartPolicyMap[flyteRestartPolicy]
231+
}
232+
233+
// OverrideContainerSpec overrides the specified container's properties in the given podSpec. The function
234+
// updates the image, resources and command arguments of the container that matches the given containerName.
235+
func OverrideContainerSpec(podSpec *v1.PodSpec, containerName string, image string, resources *core.Resources, args []string) error {
236+
for idx, c := range podSpec.Containers {
237+
if c.Name == containerName {
238+
if image != "" {
239+
podSpec.Containers[idx].Image = image
240+
}
241+
if resources != nil {
242+
// if resources requests and limits both not set, we will not override the resources
243+
if len(resources.Requests) >= 1 || len(resources.Limits) >= 1 {
244+
resources, err := flytek8s.ToK8sResourceRequirements(resources)
245+
if err != nil {
246+
return flyteerr.Errorf(flyteerr.BadTaskSpecification, "invalid TaskSpecificat ion on Resources [%v], Err: [%v]", resources, err.Error())
247+
}
248+
podSpec.Containers[idx].Resources = *resources
249+
}
250+
} else {
251+
podSpec.Containers[idx].Resources = v1.ResourceRequirements{}
252+
}
253+
if len(args) != 0 {
254+
podSpec.Containers[idx].Args = args
255+
}
256+
}
257+
}
258+
return nil
259+
}

‎go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@ import (
55
"testing"
66
"time"
77

8+
"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
89
"github.com/flyteorg/flyteplugins/go/tasks/logs"
910

1011
pluginsCore "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core"
1112
commonOp "github.com/kubeflow/common/pkg/apis/common/v1"
1213
"github.com/stretchr/testify/assert"
1314
corev1 "k8s.io/api/core/v1"
15+
v1 "k8s.io/api/core/v1"
16+
"k8s.io/apimachinery/pkg/api/resource"
1417
)
1518

1619
func TestExtractMPICurrentCondition(t *testing.T) {
@@ -183,3 +186,101 @@ func TestGetLogs(t *testing.T) {
183186
assert.Equal(t, fmt.Sprintf("k8s.com/#!/log/%s/%s-chiefReplica-0/pod?namespace=tensorflow-namespace", "tensorflow-namespace", "test"), jobLogs[2].Uri)
184187

185188
}
189+
190+
func dummyPodSpec() v1.PodSpec {
191+
return v1.PodSpec{
192+
Containers: []v1.Container{
193+
{
194+
Name: "primary container",
195+
Args: []string{"pyflyte-execute", "--task-module", "tests.flytekit.unit.sdk.tasks.test_sidecar_tasks", "--task-name", "simple_sidecar_task", "--inputs", "{{.input}}", "--output-prefix", "{{.outputPrefix}}"},
196+
Resources: v1.ResourceRequirements{
197+
Limits: v1.ResourceList{
198+
"cpu": resource.MustParse("2"),
199+
"memory": resource.MustParse("200Mi"),
200+
"gpu": resource.MustParse("1"),
201+
},
202+
Requests: v1.ResourceList{
203+
"cpu": resource.MustParse("1"),
204+
"memory": resource.MustParse("100Mi"),
205+
"gpu": resource.MustParse("1"),
206+
},
207+
},
208+
VolumeMounts: []v1.VolumeMount{
209+
{
210+
Name: "volume mount",
211+
},
212+
},
213+
},
214+
{
215+
Name: "secondary container",
216+
Resources: v1.ResourceRequirements{
217+
Limits: v1.ResourceList{
218+
"gpu": resource.MustParse("2"),
219+
},
220+
Requests: v1.ResourceList{
221+
"gpu": resource.MustParse("2"),
222+
},
223+
},
224+
},
225+
},
226+
Volumes: []v1.Volume{
227+
{
228+
Name: "dshm",
229+
},
230+
},
231+
Tolerations: []v1.Toleration{
232+
{
233+
Key: "my toleration key",
234+
Value: "my toleration value",
235+
},
236+
},
237+
}
238+
}
239+
240+
func TestOverrideContainerSpec(t *testing.T) {
241+
podSpec := dummyPodSpec()
242+
err := OverrideContainerSpec(
243+
&podSpec, "primary container", "testing-image",
244+
&core.Resources{
245+
Requests: []*core.Resources_ResourceEntry{
246+
{Name: core.Resources_CPU, Value: "250m"},
247+
},
248+
Limits: []*core.Resources_ResourceEntry{
249+
{Name: core.Resources_CPU, Value: "500m"},
250+
},
251+
},
252+
[]string{"python", "-m", "run.py"},
253+
)
254+
assert.NoError(t, err)
255+
assert.Equal(t, 2, len(podSpec.Containers))
256+
assert.Equal(t, "testing-image", podSpec.Containers[0].Image)
257+
assert.NotNil(t, podSpec.Containers[0].Resources.Limits)
258+
assert.NotNil(t, podSpec.Containers[0].Resources.Requests)
259+
// verify resources not overridden if empty resources
260+
assert.True(t, podSpec.Containers[0].Resources.Requests.Cpu().Equal(resource.MustParse("250m")))
261+
assert.True(t, podSpec.Containers[0].Resources.Limits.Cpu().Equal(resource.MustParse("500m")))
262+
assert.Equal(t, []string{"python", "-m", "run.py"}, podSpec.Containers[0].Args)
263+
}
264+
265+
func TestOverrideContainerSpecEmptyFields(t *testing.T) {
266+
podSpec := dummyPodSpec()
267+
err := OverrideContainerSpec(&podSpec, "primary container", "", &core.Resources{}, []string{})
268+
assert.NoError(t, err)
269+
assert.Equal(t, 2, len(podSpec.Containers))
270+
assert.NotNil(t, podSpec.Containers[0].Resources.Limits)
271+
assert.NotNil(t, podSpec.Containers[0].Resources.Requests)
272+
// verify resources not overridden if empty resources
273+
assert.True(t, podSpec.Containers[0].Resources.Requests.Cpu().Equal(resource.MustParse("1")))
274+
assert.True(t, podSpec.Containers[0].Resources.Requests.Memory().Equal(resource.MustParse("100Mi")))
275+
assert.True(t, podSpec.Containers[0].Resources.Limits.Cpu().Equal(resource.MustParse("2")))
276+
assert.True(t, podSpec.Containers[0].Resources.Limits.Memory().Equal(resource.MustParse("200Mi")))
277+
}
278+
279+
func TestOverrideContainerNilResources(t *testing.T) {
280+
podSpec := dummyPodSpec()
281+
err := OverrideContainerSpec(&podSpec, "primary container", "", nil, []string{})
282+
assert.NoError(t, err)
283+
assert.Equal(t, 2, len(podSpec.Containers))
284+
assert.Nil(t, podSpec.Containers[0].Resources.Limits)
285+
assert.Nil(t, podSpec.Containers[0].Resources.Requests)
286+
}

0 commit comments

Comments
 (0)