Add pod_template and pod_template_name arguments for PythonAutoContainerTask, its downstream tasks, and @task.

[ByronHsu]

TL;DR

Add pod_template and pod_template_name arguments for PythonAutoContainerTask, its downstream tasks, and @task. (Following this guide)

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

1. Users can pass in pod_template in PythonAutoContainerTask, its downstream tasks, and @task to customize pod_template for each task instance. It also allows developers to set the default pod_template for plugins in flytekit.
2. Users can pass in pod_template_name in PythonAutoContainerTask, its downstream tasks, and @task to select which PodTemplate resources to use on the server side.

Example

• instance

task_with_minimum_pod_template = DummyAutoContainerTask(
    name="x",
    task_config=None,
    task_type="t",
    container_image="repo/image:0.0.0",
    pod_template=PodTemplate(
        primary_container_name="primary",
        labels={"lKeyA": "lValA"},
        annotations={"aKeyA": "aValA"},
        pod_spec=V1PodSpec(
            containers=[
                V1Container(
                    name="primary",
                ),
            ]
        ),
    ),
    pod_template_name="A",
)

• decorator

    @task(
        container_image="repo/image:0.0.0",
        pod_template=PodTemplate(
            primary_container_name="primary",
            labels={"lKeyA": "lValA"},
            annotations={"aKeyA": "aValA"},
            pod_spec=V1PodSpec(
                containers=[
                    V1Container(
                        name="primary",
                    ),
                ]
            ),
        ),
        pod_template_name="A",
    )
    def func_with_pod_template(i: str):
        print(i + 3)

Tracking Issue

flyteorg/flyte#3123

[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).

[hamersaw]

Also, support in ShellTask would be great too! Happy to meet up and discuss any of this. Thanks so much, looks great so far!

[ByronHsu]

@hamersaw Shelltask should be working because it inherits PythonAutoContainerTask. @wild-endeavor Do we need to add a pod_template test in all PythonAutoContainerTask downstream tasks?

[ByronHsu]

How to safely update requirement.txt?

[ByronHsu]

I added "kubernetes>=12.0.1" in setup.py and ran make requirements.txt, but many deps got added. Is this the expected result?

[eapolinario]

I added "kubernetes>=12.0.1" in setup.py and ran make requirements.txt, but many deps got added. Is this the expected result?

Yes, you can see in the rendered requirements.txt file the source for these new deps.

That said, we have an issue right now where the dev-requirements.txt requirements doesn't match the new requirements.txt. I'm going to push a commit to your branch to fix this.

[eapolinario]

flyteorg/flyteidl#358 needs to be merged and incorporated here (i.e. by bumping version of flyteidl in setup.py)

[ByronHsu]

@esadler-hbo tysm!

[hamersaw]

Also, I noticed that right now we can not specify a PodSpec with an empty containers list. This means we need unecessary verbosity if a user wants to so something simple like add a toleration to the Pod. It looks like we could use local_vars_configuration internally to stop client-side validation of the PodTemplate. I think this is worth spending a little time looking into - because then we can do something like:

@task(
    pod_template=PodTemplate(
        spec=V1PodSpec{
            tolerations=[
                V1Toleration{
                    key= "num-gpus"
                    operator= "Equal"
                    value= 1
                    effect= "NoSchedule"
                },
            ]
        )
    )
)
def foo(bar: int):
    # omitted

Without having to add the bloat of definition a list with an empty container.

@task(
    pod_template=PodTemplate(
        spec=V1PodSpec{
            containers=[
                V1Container(
                    name="primary",
                ),
            ],
            tolerations=[
                V1Toleration{
                    key= "num-gpus"
                    operator= "Equal"
                    value= 1
                    effect= "NoSchedule"
                },
            ]
        )
    )
)
def foo(bar: int):
    # omitted

[hamersaw]

Should we briefly explore name changes here? pod_template and pod_template_name are not very descriptive and certainly don't capture a client-side PodTemplate application vs server-side k8s resource PodTemplate configuration. Thoughts?

[ByronHsu]

yeah it is very confusing now, but honestly I don't have a better idea on my mind. I also feel they should be one of. If users can set both, it will be more confusing.

[ByronHsu]

Still needs discussion!

[eapolinario]

We sync'd offline and decided to go with these names.

[ByronHsu]

I did the following changes:

1. Import PodTemplate in the root level
2. Return the primary container name in get_config in pythonAutoContainerTask and its descendants (awsbatch are papermili the two using get_config currently).
3. Allow empty V1PodSpec
4. Bypass required container assert (link)

Bullet 4 is a bit tricky. I overwrote V1PodSpec containers setter to not do the validation. I assume the current behavior is that if users import PodTemplate, this code will be run. I felt this overwrite things should be moved to some entrypoint in flytekit. I need some python experts' help.

[eapolinario]

Also, I noticed that right now we can not specify a PodSpec with an empty containers list. This means we need unecessary verbosity if a user wants to so something simple like add a toleration to the Pod. It looks like we could use local_vars_configuration internally to stop client-side validation of the PodTemplate. I think this is worth spending a little time looking into - because then we can do something like:

@task(
    pod_template=PodTemplate(
        spec=V1PodSpec{
            tolerations=[
                V1Toleration{
                    key= "num-gpus"
                    operator= "Equal"
                    value= 1
                    effect= "NoSchedule"
                },
            ]
        )
    )
)
def foo(bar: int):
    # omitted

Without having to add the bloat of definition a list with an empty container.

@task(
    pod_template=PodTemplate(
        spec=V1PodSpec{
            containers=[
                V1Container(
                    name="primary",
                ),
            ],
            tolerations=[
                V1Toleration{
                    key= "num-gpus"
                    operator= "Equal"
                    value= 1
                    effect= "NoSchedule"
                },
            ]
        )
    )
)
def foo(bar: int):
    # omitted

Unfortunately, there's no way to turn off validation through this method you mentioned, Dan. Take a look at this code:

In [35]: from kubernetes.client import Configuration

In [36]: conf = Configuration()

In [37]: conf.client_side_validation = True  # the default

In [38]: V1PodSpec(local_vars_configuration=conf)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[38], line 1
----> 1 V1PodSpec(local_vars_configuration=conf)

File /tmp/tmp.m5bYaiOZ4k/venv/lib/python3.10/site-packages/kubernetes/client/models/v1_pod_spec.py:166, in V1PodSpec.__init__(self, active_deadline_seconds, affinity, automount_service_account_token, containers, dns_config, dns_policy, enable_service_links, ephemeral_containers, host_aliases, host_ipc, host_network, host_pid, host_users, hostname, image_pull_secrets, init_containers, node_name, node_selector, os, overhead, preemption_policy, priority, priority_class_name, readiness_gates, restart_policy, runtime_class_name, scheduler_name, security_context, service_account, service_account_name, set_hostname_as_fqdn, share_process_namespace, subdomain, termination_grace_period_seconds, tolerations, topology_spread_constraints, volumes, local_vars_configuration)
    164 if automount_service_account_token is not None:
    165     self.automount_service_account_token = automount_service_account_token
--> 166 self.containers = containers
    167 if dns_config is not None:
    168     self.dns_config = dns_config

File /tmp/tmp.m5bYaiOZ4k/venv/lib/python3.10/site-packages/kubernetes/client/models/v1_pod_spec.py:322, in V1PodSpec.containers(self, containers)
    314 """Sets the containers of this V1PodSpec.
    315
    316 List of containers belonging to the pod. Containers cannot currently be added or removed. There must be at least one container in a Pod. Cannot be updated.  # noqa: E501
   (...)
    319 :type: list[V1Container]
    320 """
    321 if self.local_vars_configuration.client_side_validation and containers is None:  # noqa: E501
--> 322     raise ValueError("Invalid value for `containers`, must not be `None`")  # noqa: E501
    324 self._containers = containers

ValueError: Invalid value for `containers`, must not be `None`

In [39]: conf.client_side_validation = False

In [40]: V1PodSpec(local_vars_configuration=conf)
Out[40]:
... # this works

That said, we don't necessarily have to define a container, the only restriction is that the list of containers is not None, i.e. this works:

In [41]: V1PodSpec(containers=[])
Out[41]:
{'active_deadline_seconds': None,
 'affinity': None,
 'automount_service_account_token': None,
 'containers': [],
 'dns_config': None,
 'dns_policy': None,
...

I'm conflicted. On one hand, we could wrap k8s V1PodSpec with our own V1PodSpec (so that we could disable client-side validation), OTOH, leaning on this behavior of the current k8s object where we define an empty list is not terrible (and we have a short feedback loop in case we break it).

One thing that feels off for me with this idea of replacing the validation of the k8s client validation for the containers specifically is that if there's code that relies on that validation (for whatever reason) in a task, that code will no longer work.

wdyt?

[hamersaw]

wdyt?

I think this sounds great. Lets leave it as is, users have to define an empty containers list is the pod spec is defined. Going this route also means if it's a problem we can address it in the future - much easier then going the opposite direction.

[codecov]

Codecov Report

Merging #1425 (01b10c0) into master (ea39054) will increase coverage by 0.07%.
The diff coverage is 80.98%.

@@            Coverage Diff             @@
##           master    #1425      +/-   ##
==========================================
+ Coverage   69.17%   69.25%   +0.07%     
==========================================
  Files         300      301       +1     
  Lines       28150    28309     +159     
  Branches     2649     2663      +14     
==========================================
+ Hits        19472    19604     +132     
- Misses       8166     8193      +27     
  Partials      512      512              

Impacted Files	Coverage Δ
flytekit/core/base_task.py	46.06% <0.00%> (-0.19%)	⬇️
flytekit/core/container_task.py	27.45% <ø> (ø)
flytekit/core/pod_template.py	0.00% <0.00%> (ø)
flytekit/core/task.py	34.09% <0.00%> (-0.80%)	⬇️
setup.py	0.00% <ø> (ø)
tests/flytekit/common/parameterizers.py	100.00% <ø> (ø)
...ests/flytekit/unit/models/test_workflow_closure.py	100.00% <ø> (ø)
flytekit/models/task.py	36.31% <50.00%> (+0.15%)	⬆️
flytekit/core/python_auto_container.py	59.11% <76.92%> (+9.57%)	⬆️
...ts/flytekit/unit/core/test_python_function_task.py	92.38% <96.15%> (+1.24%)	⬆️
... and 3 more

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

[eapolinario]

Just a few minor comments and we should be good to go.

[eapolinario]

We sync'd offline and decided to go with these names.

[eapolinario]

This is awesome!

[welcome]

Congrats on merging your first pull request! 🎉