Refactor kf-operator plugins configs and support setting different specs for different replica groups

[yubofredwang]

TL;DR

This PR refactors config fields for kubeflow operator plugins including: TFJob, MPIJob, PyTorchJob. This change is backward incompatible, so we will keep the original plugin configs untouched and use the new configs in v2 version of the plugins.

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

Code change includes:

1. Move TFJob, MPIJob, PyTorchJob from flyteidl/plugins/ to flyteidl/plugins/kubeflow/
2. Create flyteidl/plugins/kubeflow/common.proto that contains the definition RestartPolicy, RunPolicy and CleanPodPolicy
3. Add xxxReplicaSpec in each job type to allow settings of replicas, image, resources and restart_policy

Tracking Issue

fixes flyteorg/flyte#3308

Follow-up issue

[welcome]

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

• Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
• Sign off your commits (Reference: DCO Guide).

[codecov]

Codecov Report

Merging #386 (7ba6b07) into master (9189ff2) will increase coverage by 2.37%.
The diff coverage is n/a.

❗ Current head 7ba6b07 differs from pull request most recent head 96a45ea. Consider uploading reports for the commit 96a45ea to get more accurate results

@@            Coverage Diff             @@
##           master     #386      +/-   ##
==========================================
+ Coverage   76.11%   78.49%   +2.37%     
==========================================
  Files          18       18              
  Lines        1390     1195     -195     
==========================================
- Hits         1058      938     -120     
+ Misses        280      205      -75     
  Partials       52       52              

Flag	Coverage Δ
unittests	?

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

see 18 files with indirect coverage changes

[hamersaw]

Quick update per our offline discussion. The goal is to make these changes backwards compatible with the existing flytekit kubeflow operator APIs. So the proposed IDL spec will override the existing approach, so the kubeflow operator updates are something like:

@task(
    task_config=TfJob(
        num_workers=2, # @deprecated
        num_ps_replicas=1, # @deprecated
        num_chief_replicas=1 # @deprecated
        worker_spec=TfWorkflowSpec(
             replicas=3 # backend job will use 3 replicas, overriding the 2 provided above
        ),
        # similar specs for ps and chief replicas
    ),
)
def foo:
    # omitted
    
@task(
    task_config=TfJob(
        num_workers=2, # @deprecated - but used because worker_spec is not specified
        num_ps_replicas=1, # @deprecated - similarly used
        num_chief_replicas=1 # @deprecated  -similarly used
    ),
)
def foo:
    # omitted

IIUC this is the cleanest way to implement these changes. Very interested in flytekit maintainer thoughts - cc @wild-endeavor @eapolinario.

The largest open question is how this should be written in wire-format, there are two solutions:
(1) the proposed solution creates a new flyteidl proto message to support the new format. This will require flytekit to compile all new kubeflow jobs to use TaskType=1 (rather than the exist 0) in the TaskTemplate denoting the custom struct is using the old or new configuration definition. In this approach flytekit will be responsible for applying the override values and the deprecation messages will be on the flytekit kubeflow API fields.
(2) Merge the new and existing proto messages. In this scenario the idl message will directly reflect the kubeflow APIs. In the above example, the idl message will have both num_workers and workers_spec.replicas. In this approach flyteplugins will be responsible for applying the override values and the deprecation messages will be on both the flytekit kubeflow API fields and the flyteidl configuration proto.

@yubofredwang please feel free to correct this and / or clarify if there is more context here.

[hamersaw]

@yubofredwang we talked internally here, (hopefully) quick summary. The best approach is probably to use a new proto message and increment the TaskType version to 1. This means deprecating the top-level replica counts will be done in flytekit-side. Does this sound reasonable to you? We should be able to move forward with this PR then, but want to touch-base before pushing it through.

[yubofredwang]

@yubofredwang we talked internally here, (hopefully) quick summary. The best approach is probably to use a new proto message and increment the TaskType version to 1. This means deprecating the top-level replica counts will be done in flytekit-side. Does this sound reasonable to you? We should be able to move forward with this PR then, but want to touch-base before pushing it through.

Got it. My understanding is we make the following changes:

flyteidl: keep the new proto separated from the old protos of kf-operators
flyteplugins: support both TaskTypeVersion 0 and 1, the TaskTemplate will be in different format for the two versions
flytekit: bump the flytekit version and use the new proto only in the new flytekit version

My questions regarding to this approach is how the user upgrade process should look like. There are three related pieces here: user code, flytekit, flyteplugins

Case1: code unchanged, flytekit unchanged, flyteplugins updated
in this case, the existing workflow works without a problem since flyteplugins support both version

Case2: code unchanged, flytekit updated, flyteplugins unchanged
this use case, the existing workflow will fail due to flytekit does not accept old config. User should get error message instructing them to update code.

Case3: code updated, flytekit updated, flyteplugins unchanged
same as case1

Case4: code updated, flytekit unchanged, flyteplugins updated
Not likely to occur since users are not aware of the new flytekit interface

[ByronHsu]

Regardless of the compatibility issues, the idl looks good to me!

[hamersaw]

Looks great! My only nit is on proto location in respect to v0 of the kubeflow operators. Not sure of the cleanest way to handle this (maybe this is it!) but it seems a little odd to have v0 and v1 in different locations.

[yubofredwang]

Looks great! My only nit is on proto location in respect to v0 of the kubeflow operators. Not sure of the cleanest way to handle this (maybe this is it!) but it seems a little odd to have v0 and v1 in different locations.

The reason I put these into that specific location kubeflow/plugins was the following:

1. The three plugins proto share common parameters, which are specific to kubeflow plugins.
2. The new plugins now can keep the same name such as "DistributedMPITrainingTask". If I put it under same root directory, compiler will complain with conflicts since they are named the same.
3. I am kind of reluctant to rename the plugins to pytorch_v1.proto since later on after we deprecate v0, we might have to rename it again. Then at that time, we have to do another code change in both flytekit and flyteplugins repo

[welcome]

Congrats on merging your first pull request! 🎉