Apply default pod template to PytorchJob pods #297

fg91 · 2022-11-24T14:52:48Z

TL;DR

Currently, the default pod template is applied to Python tasks but not to the kf-operators tasks like the Pytorch task.

This is a problem since e.g. for the pytorch dataloaders, one often has to increase the shared memory, wich on K8s is done by mounting an emptyDir volume. Currently there exists no mechanism to do that for Pytorch tasks.

In this PR, the v1.PodSpec of the default pod template is applied to the PodSpec of the PytorchJob.

Type

Bug Fix
Feature
Plugin

Are all requirements met?

Complete description

Tracking Issue

Remove the 'fixes' keyword if there will be multiple PRs to fix the linked issue

fixes https://github.com/flyteorg/flyte/issues/

Follow-up issue

NA
OR
https://github.com/flyteorg/flyte/issues/

codecov · 2022-11-24T15:01:44Z

Codecov Report

Merging #297 (b5ed78a) into master (85c512e) will increase coverage by 0.20%.
The diff coverage is 70.83%.

@@            Coverage Diff             @@
##           master     #297      +/-   ##
==========================================
+ Coverage   62.09%   62.29%   +0.20%     
==========================================
  Files         145      145              
  Lines       11497    11564      +67     
==========================================
+ Hits         7139     7204      +65     
  Misses       3815     3815              
- Partials      543      545       +2

Flag	Coverage Δ
unittests	`62.29% <70.83%> (+0.20%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
...o/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go	`67.96% <53.33%> (-3.15%)`	⬇️
...s/plugins/k8s/kfoperators/tensorflow/tensorflow.go	`71.55% <53.33%> (-3.21%)`	⬇️
go/tasks/plugins/k8s/kfoperators/mpi/mpi.go	`70.96% <58.82%> (-2.43%)`	⬇️
go/tasks/pluginmachinery/flytek8s/pod_helper.go	`84.06% <85.71%> (+1.39%)`	⬆️
go/tasks/plugins/array/awsbatch/executor.go	`36.95% <0.00%> (+0.52%)`	⬆️
go/tasks/plugins/array/awsbatch/transformer.go	`79.00% <0.00%> (+6.07%)`	⬆️
go/tasks/plugins/array/awsbatch/monitor.go	`67.42% <0.00%> (+7.07%)`	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

fg91 · 2022-11-24T18:59:54Z

@hamersaw I compiled an image and ran it in our cluster, it works. However, how should I add tests for this?

Ideally I would want to duplicate this test but apply a pod template. However, for this I would have to mock/patch flytek8s.DefaultPodTemplateStore.LoadOrDefault which is called here.

This function which is the equivalent for normal pods is not tested in the case of an existing pod template either from what I can see.

Methods for mocking like mentioned here wouldn't work without refactoring and unfortunately I don't understand how you employ mockery. Do you see a way to mock flytek8s.DefaultPodTemplateStore.LoadOrDefault with mockery or would you suggest to only add a unit test for the new function MergePodSpecs?

hamersaw · 2022-11-28T19:00:03Z

All great questions. TL;DR I think testing the MergePodSpecs function is enough here.

Ideally I would want to duplicate this test but apply a pod template. However, for this I would have to mock/patch flytek8s.DefaultPodTemplateStore.LoadOrDefault which is called here.

You could manually inject a PodTemplate into the DefaultPodTemplateStore using the Store function. To get an idea of what this looks like you can see here. I would be careful about how this might effect other tests.

Methods for mocking like mentioned here wouldn't work without refactoring and unfortunately I don't understand how you employ mockery. Do you see a way to mock flytek8s.DefaultPodTemplateStore.LoadOrDefault with mockery or would you suggest to only add a unit test for the new function MergePodSpecs?

So mockery works by generating stubs for golang interfaces. Currently the PodTemplateStore is just defined as a struct. We could easily add an interface abstraction and use mockery to generate mocks. I'm not sure that this is necessary as long as we are sure the merging logic is sound.

go/tasks/pluginmachinery/flytek8s/pod_helper.go

hamersaw

I think this looks great. A few nits to maybe look into.

Also I see the PodTemplateSpec here and here have an ObjectMeta field. In the BuildPodWithSpec we maintain the metadata to bring things like labels and annotations over. Do you know if we set ObjectMeta on the PodTemplateSpec will be be maintained in the Pods launched from the operator?

fg91 · 2022-11-30T10:08:47Z

Also I see the PodTemplateSpec here and here have an ObjectMeta field. In the BuildPodWithSpec we maintain the metadata to bring things like labels and annotations over. Do you know if we set ObjectMeta on the PodTemplateSpec will be be maintained in the Pods launched from the operator?

No, the operator doesn't carry them over. I created a PytorchJob with the following ObjectMeta:

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: pytorch-simple
  namespace: test
  labels:
    foo: bar
  annotations:
    afoo: abar

It results in Pods that look like this:

apiVersion: v1
kind: Pod
metadata:
  creationTimestamp: "2022-11-30T10:07:42Z"
  labels:
    group-name: kubeflow.org
    job-name: pytorch-simple
    job-role: master
    replica-index: "0"
    replica-type: master
    training.kubeflow.org/job-name: pytorch-simple
    training.kubeflow.org/job-role: master
    training.kubeflow.org/operator-name: pytorchjob-controller
    training.kubeflow.org/replica-index: "0"
    training.kubeflow.org/replica-type: master
  name: pytorch-simple-master-0
  namespace: test
  ownerReferences:
  - apiVersion: kubeflow.org/v1
    blockOwnerDeletion: true
    controller: true
    kind: PyTorchJob
    name: pytorch-simple
    uid: 84b4e76d-de63-46a5-a829-c1595b7fd4bd
  resourceVersion: "30970195"
  uid: 7e5232d1-0deb-4013-bbba-cf350232c1db

fg91 · 2022-11-30T19:24:11Z

@hamersaw can you please take a look at the new changes? I tried to make sensible commits that make the review easy, maybe look at them one by one.

About unit tests:

The existing test TestBuildPodWithSpec tests the podSpec merging behaviour but it currently doesn't test that the object metadata is copied over.

I copied the merging test logic from TestBuildPodWithSpec into the new test TestMergePodSpecs
I created a new test TestBuildPodWithSpec2 that checks that podSpec merging is taking place in BuildPodWithSpec without testing any details of it. In addition, I now test that the object metadata is copied over.
I left the TestBuildPodWithSpec test untouched even though the entire logic was moved into TestBuildPodWithSpec in order to give confidence that the refactoring didn't break anything that the old test would have caught. Maybe we should remove this test now though after this has been shown?

Having added these lines ...

if podTemplate != nil {
    mergedPodSpec, err := flytek8s.MergePodSpecs(

... reduced the test coverage though. Do you have a suggestion how we can make the build green without mocking anything here? And do we have to?

hamersaw · 2022-12-07T16:37:30Z

No, the operator doesn't carry them over. I created a PytorchJob with the following ObjectMeta:

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: pytorch-simple
  namespace: test
  labels:
    foo: bar
  annotations:
    afoo: abar

Oh sorry, I don't think I explained in-depth enough. I was thinking the ObjectMeta should be set on the PodTemplateSpec like here. I haven't been able to see this carried over in the kf pytorch operator, but it looks like the MPI controller just copies the PodSpec, which should carry over ObjectMeta fields. Does this sound right?

hamersaw · 2022-12-07T16:40:54Z

IIUC the updated default container names means that we need to create a container in the default PodTemplate called pytorch (because that's the pytorch default container name) which will be used as the base for all pytorch containers? So the default container will not be used? I'm assuming this is because you want to create a single PodTemplate for the entire deployment and use separate default container configuration for regular k8s tasks and pytorch jobs, can you say a little more about this? I kind of like the strict 'default' applies to everything approach, but just want to make sure I understand the change.

fg91 · 2022-12-08T12:03:49Z

IIUC the updated default container names means that we need to create a container in the default PodTemplate called pytorch (because that's the pytorch default container name) which will be used as the base for all pytorch containers? So the default container will not be used? I'm assuming this is because you want to create a single PodTemplate for the entire deployment and use separate default container configuration for regular k8s tasks and pytorch jobs, can you say a little more about this? I kind of like the strict 'default' applies to everything approach, but just want to make sure I understand the change.

No, I'm currently using this pod template with a single container named default:

apiVersion: v1
kind: PodTemplate
metadata:
  name: flyte-template
  namespace: flyte
template:
  spec:
    containers:
    - image: foo
      imagePullPolicy: Always
      name: default
      volumeMounts:
      - mountPath: /dev/shm
        name: dshm
    volumes:
    - emptyDir:
        medium: Memory
      name: dshm

With the current implementation this translates into this pod:

apiVersion: v1
kind: Pod
...
spec:
  containers:
  - name: pytorch
    ...
    volumeMounts:
    - mountPath: /dev/shm
      name: dshm
  volumes:
  - emptyDir:
      medium: Memory
    name: dshm
...

hamersaw · 2022-12-08T16:48:43Z

No, I'm currently using this pod template with a single container named default:

Ah OK, small nit then. If we are not using the defaultContainerName parameter on the MergePodSpecs function can we just remove it? After that, looks great - lets merge!

fg91 · 2022-12-08T17:10:29Z

Ah OK, small nit then. If we are not using the defaultContainerName parameter on the MergePodSpecs function can we just remove it? After that, looks great - lets merge!

I just realised that in MergePodSpecs I accidentally use the global vars defaultContainerTemplateName and primaryContainerTemplateName instead of the function args primaryContainerName and defaultContainerName:

// merge template Containers
	var mergedContainers []v1.Container
	var defaultContainerTemplate, primaryContainerTemplate *v1.Container
	for i := 0; i < len(podTemplatePodSpecCopy.Containers); i++ {
		if podTemplatePodSpecCopy.Containers[i].Name == defaultContainerTemplateName {
			defaultContainerTemplate = &podTemplatePodSpecCopy.Containers[i]
		} else if podTemplatePodSpecCopy.Containers[i].Name == primaryContainerTemplateName {
			primaryContainerTemplate = &podTemplatePodSpecCopy.Containers[i]
		}
	}

That is why defaultContainerName isn't used.

Let me take another close look at this and also at setting the ObjectMeta on the PodTemplateSpec before we merge please.

fg91 · 2022-12-10T13:02:55Z

I tested the Pytorch Job plugin with the following template:

apiVersion: v1
kind: PodTemplate
metadata:
  name: flyte-template
  namespace: flyte
template:
  metadata:
    labels:
      foo: bar                              # <- new
  spec:
    containers:
    - image: foo
      imagePullPolicy: Always
      name: default
      volumeMounts:
      - mountPath: /dev/shm
        name: dshm
    volumes:
    - emptyDir:
        medium: Memory
      name: dshm

The template leads to the following pod:

apiVersion: v1
kind: Pod
metadata:
  labels:
    foo: bar
    group-name: kubeflow.org
    ....
  name: avjfhs8nsrq4jqm46hzr-feihpdka-0-master-0
  namespace: development
  ...
spec:
  affinity: {}
  containers:
  - args:
    - pyflyte-fast-execute
    ...
    name: pytorch
    ...
    volumeMounts:
    - mountPath: /dev/shm
      name: dshm
    ...
  volumes:
  - emptyDir:
      medium: Memory
    name: dshm
  ...

go/tasks/pluginmachinery/flytek8s/pod_helper_test.go

fg91 · 2022-12-10T13:15:49Z

go/tasks/plugins/k8s/kfoperators/mpi/mpi.go

@@ -67,6 +67,21 @@ func (mpiOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx plu
 		return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create pod spec: [%v]", err.Error())
 	}

+	common.OverrideDefaultContainerName(taskCtx, podSpec, kubeflowv1.MPIJobDefaultContainerName)


@hamersaw Do you think adding this line is incorrect?
As opposed to the pytorch and tfjob plugins, the container name is currently not overridden with mpi even though this is the default container name (see here and here).

This is interesting. I'm not seeing the GetDefaultContainerName function being used anywhere. I'm wonder if either, it doesn't care about the container name or if the launcher pod automatically updates container names to reflect this.

cc @bimtauer is this something you know anything about? or can review this update?

@hamersaw I will try and have a look by the end of the week!

go/tasks/plugins/k8s/kfoperators/mpi/mpi.go

go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go

Signed-off-by: Fabio Grätz <[email protected]>

hamersaw

Looks great! Clean up tests and get another pair of eyes on it and lets merge!

go/tasks/pluginmachinery/flytek8s/pod_helper_test.go

go/tasks/plugins/k8s/kfoperators/mpi/mpi.go

hamersaw · 2022-12-12T16:42:08Z

go/tasks/plugins/k8s/kfoperators/mpi/mpi.go

@@ -67,6 +67,21 @@ func (mpiOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx plu
 		return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create pod spec: [%v]", err.Error())
 	}

+	common.OverrideDefaultContainerName(taskCtx, podSpec, kubeflowv1.MPIJobDefaultContainerName)


This is interesting. I'm not seeing the GetDefaultContainerName function being used anywhere. I'm wonder if either, it doesn't care about the container name or if the launcher pod automatically updates container names to reflect this.

cc @bimtauer is this something you know anything about? or can review this update?

hamersaw · 2022-12-12T16:45:36Z

cc @bimtauer - this PR adds the default PodTemplate to the MPI task type. IIUC this should address this issue on IAM roles with annotations for the MPI tasks. Could you take a look through this to sanity check MPI updates? It would be VERY helpful!

Signed-off-by: Fabio Grätz <[email protected]>

fg91 · 2022-12-12T18:19:37Z

Looks great! Clean up tests and get another pair of eyes on it and lets merge!

Renamed my test and deleted the old one 👍

* Merge pod template spec with pod spec in separate func Signed-off-by: Fabio Grätz <[email protected]> * Merge pod template spec with pod spec in separate func Signed-off-by: Fabio Grätz <[email protected]> * Apply pod template to pytorch job pod spec Signed-off-by: Fabio Grätz <[email protected]> * Handle nil podspecs before merging Signed-off-by: Fabio Grätz <[email protected]> * Pass both default and primare container name to MergePodSpecs Signed-off-by: Fabio Grätz <[email protected]> * Move podSpec.DeepCopy into MergePodSpecs Signed-off-by: Fabio Grätz <[email protected]> * Add tests Signed-off-by: Fabio Grätz <[email protected]> * Merge pod template into tfjob and mpijob Signed-off-by: Fabio Grätz <[email protected]> * Lint Signed-off-by: Fabio Grätz <[email protected]> * Correct usage of default and primate container (template) name Signed-off-by: Fabio Grätz <[email protected]> * Override mpi default container name Signed-off-by: Fabio Grätz <[email protected]> * Carry over ObjectMeta from pod template Signed-off-by: Fabio Grätz <[email protected]> * Remove old `TestBuildPodWithSpec` test Signed-off-by: Fabio Grätz <[email protected]> Signed-off-by: Fabio Grätz <[email protected]> Co-authored-by: Fabio Grätz <[email protected]>

fg91 marked this pull request as draft November 24, 2022 14:53

fg91 force-pushed the fabio-pod-template-kfoperators branch 4 times, most recently from d25aade to 2ec7da8 Compare November 24, 2022 18:44

fg91 changed the title ~~Fabio pod template kfoperators~~ Apply default pod template to PytorchJob pods Nov 24, 2022

fg91 changed the title ~~Apply default pod template to PytorchJob pods~~ Apply default pod template to PytorchJob pods # minor Nov 24, 2022

fg91 changed the title ~~Apply default pod template to PytorchJob pods # minor~~ Apply default pod template to PytorchJob pods # patch Nov 24, 2022

fg91 requested a review from hamersaw November 24, 2022 19:00

fg91 marked this pull request as ready for review November 24, 2022 19:03

fg91 assigned fg91 and hamersaw Nov 24, 2022

hamersaw reviewed Nov 28, 2022

View reviewed changes

go/tasks/pluginmachinery/flytek8s/pod_helper.go Show resolved Hide resolved

hamersaw reviewed Nov 28, 2022

View reviewed changes

fg91 requested a review from hamersaw December 6, 2022 16:30

fg91 commented Dec 10, 2022

View reviewed changes

Fabio Grätz added 3 commits December 10, 2022 14:33

Merge pod template spec with pod spec in separate func

1dd1495

Signed-off-by: Fabio Grätz <[email protected]>

Merge pod template spec with pod spec in separate func

01b3885

Signed-off-by: Fabio Grätz <[email protected]>

Apply pod template to pytorch job pod spec

22f00bc

Signed-off-by: Fabio Grätz <[email protected]>

Fabio Grätz added 9 commits December 10, 2022 14:33

Handle nil podspecs before merging

03afd51

Signed-off-by: Fabio Grätz <[email protected]>

Pass both default and primare container name to MergePodSpecs

4d92d8a

Signed-off-by: Fabio Grätz <[email protected]>

Move podSpec.DeepCopy into MergePodSpecs

2f6a24b

Signed-off-by: Fabio Grätz <[email protected]>

Add tests

5efe69b

Signed-off-by: Fabio Grätz <[email protected]>

Merge pod template into tfjob and mpijob

1980e8a

Signed-off-by: Fabio Grätz <[email protected]>

Lint

d6f3a74

Signed-off-by: Fabio Grätz <[email protected]>

Correct usage of default and primate container (template) name

f71fbc8

Signed-off-by: Fabio Grätz <[email protected]>

Override mpi default container name

1c8f4de

Signed-off-by: Fabio Grätz <[email protected]>

Carry over ObjectMeta from pod template

ffbe158

Signed-off-by: Fabio Grätz <[email protected]>

fg91 force-pushed the fabio-pod-template-kfoperators branch from af9d869 to ffbe158 Compare December 10, 2022 13:33

hamersaw reviewed Dec 12, 2022

View reviewed changes

Remove old TestBuildPodWithSpec test

b5ed78a

Signed-off-by: Fabio Grätz <[email protected]>

fg91 force-pushed the fabio-pod-template-kfoperators branch from 2db21b3 to b5ed78a Compare December 12, 2022 18:17

fg91 mentioned this pull request Dec 14, 2022

[Core feature] Add PodTemplate configuration to flytekit task definitions flyteorg/flyte#3123

Closed

2 tasks

hamersaw approved these changes Dec 14, 2022

View reviewed changes

hamersaw changed the title ~~Apply default pod template to PytorchJob pods # patch~~ Apply default pod template to PytorchJob pods Dec 14, 2022

hamersaw merged commit 428b7a4 into flyteorg:master Dec 16, 2022

hamersaw mentioned this pull request Dec 16, 2022

[BUG] IAM roles are not passed on to MPI Launcher and worker pods flyteorg/flyte#3043

Closed

2 tasks

fg91 deleted the fabio-pod-template-kfoperators branch December 16, 2022 16:49

fg91 mentioned this pull request Jul 3, 2024

Add myself to code owners of flytekit-kf-pytorch flyteorg/flytekit#2556

Merged

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Apply default pod template to PytorchJob pods #297

Apply default pod template to PytorchJob pods #297

fg91 commented Nov 24, 2022 •

edited

Loading

codecov bot commented Nov 24, 2022 •

edited

Loading

fg91 commented Nov 24, 2022 •

edited

Loading

hamersaw commented Nov 28, 2022

hamersaw left a comment

fg91 commented Nov 30, 2022

fg91 commented Nov 30, 2022 •

edited

Loading

hamersaw commented Dec 7, 2022

hamersaw commented Dec 7, 2022

fg91 commented Dec 8, 2022

hamersaw commented Dec 8, 2022 •

edited

Loading

fg91 commented Dec 8, 2022

fg91 commented Dec 10, 2022

fg91 Dec 10, 2022

hamersaw Dec 12, 2022

bimtauer Dec 15, 2022

hamersaw left a comment

hamersaw Dec 12, 2022

hamersaw commented Dec 12, 2022

fg91 commented Dec 12, 2022

Apply default pod template to PytorchJob pods #297

Apply default pod template to PytorchJob pods #297

Conversation

fg91 commented Nov 24, 2022 • edited Loading

TL;DR

Type

Are all requirements met?

Complete description

Tracking Issue

Follow-up issue

codecov bot commented Nov 24, 2022 • edited Loading

Codecov Report

fg91 commented Nov 24, 2022 • edited Loading

hamersaw commented Nov 28, 2022

hamersaw left a comment

Choose a reason for hiding this comment

fg91 commented Nov 30, 2022

fg91 commented Nov 30, 2022 • edited Loading

hamersaw commented Dec 7, 2022

hamersaw commented Dec 7, 2022

fg91 commented Dec 8, 2022

hamersaw commented Dec 8, 2022 • edited Loading

fg91 commented Dec 8, 2022

fg91 commented Dec 10, 2022

fg91 Dec 10, 2022

Choose a reason for hiding this comment

hamersaw Dec 12, 2022

Choose a reason for hiding this comment

bimtauer Dec 15, 2022

Choose a reason for hiding this comment

hamersaw left a comment

Choose a reason for hiding this comment

hamersaw Dec 12, 2022

Choose a reason for hiding this comment

hamersaw commented Dec 12, 2022

fg91 commented Dec 12, 2022

fg91 commented Nov 24, 2022 •

edited

Loading

codecov bot commented Nov 24, 2022 •

edited

Loading

fg91 commented Nov 24, 2022 •

edited

Loading

fg91 commented Nov 30, 2022 •

edited

Loading

hamersaw commented Dec 8, 2022 •

edited

Loading