feat(sdk): Add in ability to override the template handler for a BaseOp.

[alexlatchford]

Description of your changes:

Argo has some great functionality the SDK doesn't expose by default. In particular, we're interested in using the PodSecurityContext and podSpecPatch functionality but needing to get in an upstream PR to leverage the functionality seems a little slow so wondering if this is an approach is something maintainers would consider.

No cherry-picking required.

---

Example usage:

from typing import Dict
from kfp.compiler._op_to_template import _op_to_template
from kfp.dsl._container_op import BaseOp, ContainerOp

def my_template_handler(op: BaseOp) -> Dict:
    template = _op_to_template(op)
    # See https://argoproj.github.io/argo/fields/#podsecuritycontext
    template['securityContext'] = {"fsGroup": 2000}
    return template

...

op = ContainerOp(...)
op.set_template_handler(my_template_handler)

Going to guess this may not pass review given a desire to hide away the underlying scheduler but open to improving this as needed.

[k8s-ci-robot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by: alexlatchford
To complete the pull request process, please assign ark-kun after the PR has been reviewed.
You can assign the PR to them by writing /assign @ark-kun in a comment when ready.

The full list of commands accepted by this bot can be found here.

Needs approval from an approver in each of these files:

• sdk/OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[k8s-ci-robot]

@alexlatchford: The following test failed, say /retest to rerun all failed tests:

Test name	Commit	Details	Rerun command
kubeflow-pipelines-tfx-python36	576f10c	link	/test kubeflow-pipelines-tfx-python36

Full PR test history. Your PR dashboard. Please help us cut down on flakes by linking to an open issue when you hit one in your PR.

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository. I understand the commands that are listed here.

[Ark-kun]

We have plans to implement ResourceOp on top of ContainerOp. It's also the way that Argo implements this internally - just a container with the executor which calls kubectl.
This might give you more flexibility and portability.

We have some reservations regarding exposing Argo-specific structures and types (e.g. Template) in the public SDK surface.

What do you think?

[alexlatchford]

We have plans to implement ResourceOp on top of ContainerOp. It's also the way that Argo implements this internally - just a container with the executor which calls kubectl.
This might give you more flexibility and portability.

Maybe I'm lacking a little knowledge here about ResourceOp but my understanding was that functionality was for creating additional resources whereas podSpecPatch and podSecurityContext would modify existing steps?

I guess taking a step back from an end user perspective what we're trying to achieve is:

1. Skirt this bug which means AWS IAM role tokens are mounted as root, we want to leverage podSecurityContext to force volume mounts to have a specific fsGroup to force them to be readable as non-root users. Currently all of our pipeline steps which interact with AWS resources need to run as root 😅
2. We have some larger pipelines wanting to be run atop KFP and dynamically set the resources of containers at runtime to fit data partitions we don't know sizes until then will save us some significant compute costs, we're thinking podSpecPatch is the desired approach from Argo to achieve that.

I'd be happy with a ResourceOp solution if it can hit these goals 😄

We have some reservations regarding exposing Argo-specific structures and types (e.g. Template) in the public SDK surface.

Makes sense, looking at the API I suspected this to be a blocker 😄

[Bobgy]

/assign @chensun @neuromage

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[rimolive]

Closing this PR. No activity for more than a year.

/close

[google-oss-prow]

@rimolive: Closed this PR.

In response to this:

Closing this PR. No activity for more than a year.

/close

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.