[BUG] set primary_container_name annotation for pod template configured pods w/ multiple containers

[pvditt]

Tracking issue

Why are the changes needed?

PrimaryContainerKey annotation is not set on pods when a pod configured by PodTemplate has multiple containers. This causes for a task to be stuck in a running state.

What changes were proposed in this pull request?

Remove some abstraction with how primary_container_name annotation is set and have it set during post_init

How was this patch tested?

• Added unit test
• Ran wf locally to confirm that annotations were getting set correctly

Setup process

import requests
from flytekit.core.pod_template import PodTemplate
from kubernetes.client.models import (
    V1Container,
    V1PodSpec,
    V1ResourceRequirements,
    V1ContainerPort,
    V1Probe,
    V1HTTPGetAction
)


@task(
    pod_template=PodTemplate(
        primary_container_name="primary",
        pod_spec=V1PodSpec(
            containers=[
                V1Container(
                    name="primary",
                    resources=V1ResourceRequirements(
                        requests={"cpu": "2", "memory": "1Gi"},
                        limits={"cpu": "2", "memory": "1Gi"},
                    ),
                    ports=[V1ContainerPort(container_port=80, name="http-primary")],
                    readiness_probe=V1Probe(
                        initial_delay_seconds=20,
                        period_seconds=5,
                        failure_threshold=5,
                        http_get=V1HTTPGetAction(path="/", port=80),
                    ),
                    liveness_probe=V1Probe(
                        initial_delay_seconds=20,
                        period_seconds=5,
                        failure_threshold=5,
                        http_get=V1HTTPGetAction(path="/", port=80),
                    ),
                ),
                V1Container(
                    name="secondary",
                    image="python",
                    command=["/bin/sh"],
                    args=[
                        "-c",
                        "python -m http.server 80",
                    ],
                    ports=[V1ContainerPort(container_port=80, name="http-api")],
                    readiness_probe=V1Probe(
                        initial_delay_seconds=20,
                        period_seconds=5,
                        failure_threshold=5,
                        http_get=V1HTTPGetAction(path="/", port=80),
                    ),
                    liveness_probe=V1Probe(
                        initial_delay_seconds=20,
                        period_seconds=5,
                        failure_threshold=5,
                        http_get=V1HTTPGetAction(path="/", port=80),
                    ),
                    resources=V1ResourceRequirements(
                        requests={"cpu": "2", "memory": "1Gi"},
                        limits={"cpu": "2", "memory": "1Gi"},
                    ),
                ),
            ],
        ),
    ),
    requests=Resources(
        mem="1G",
    ),
)
def multiple_containers_pod_task() -> int:
    url = "http://localhost:80/"
    response = requests.get(url)
    print(response)
    return response.status_code


@workflow
def multiple_containers_pod_workflow() -> int:
    code = multiple_containers_pod_task()
    return code

Screenshots

Annotations: cluster-autoscaler.kubernetes.io/safe-to-evict: false primary_container_name: primary

Check all the applicable boxes

• [x] I updated the documentation accordingly.
• [x] All new and existing tests passed.
• [x] All commits are signed-off.

Related PRs

Docs link

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 98.51%. Comparing base (da0485f) to head (1c5e563). Report is 5 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #2382       +/-   ##
===========================================
+ Coverage   52.39%   98.51%   +46.11%     
===========================================
  Files         201       14      -187     
  Lines       19250      471    -18779     
  Branches     3568        0     -3568     
===========================================
- Hits        10086      464     -9622     
+ Misses       8731        7     -8724     
+ Partials      433        0      -433     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[hamersaw]

I think this should be a flytepropeller thing rather than flytekit. Really it's just to ensure correct execution of the task, I know we set it here for the Pod plugin?