Add Node Selector attribute to KF plugins

[alykhantejani]

Why are the changes needed?

Currently for flyteplugins-kftensorflow (similarly for pytorch and mpi) there are use cases where one would want to use preemptible/interuptible nodes for workers and chief but not for parameter servers. This is currently not possible as we cannot specify node selectors for each replica type. This PR, along with a flyteidl PR (https://github.com/flyteorg/flyte/pull/4873) introduces this functionality.

How was this patch tested?

Example run in flyte-sandbox

import os
import time

from flytekit import task, workflow, ImageSpec, dynamic, Resources

from flytekitplugins.kftensorflow import TfJob, Chief, Worker, PS

kubeflow_plugin = "git+https://github.com/alykhantejani/flytekit.git@1a0302028868355dd3e9a148f260a58a947a3730#subdirectory=plugins/flytekit-kf-tensorflow"
envd_plugin = "git+https://github.com/alykhantejani/flytekit.git@1a0302028868355dd3e9a148f260a58a947a3730#subdirectory=plugins/flytekit-envd"
kubeflow_idl = "git+https://github.com/alykhantejani/flyte.git@9dcc23813925ccc7ebbdfe9f31eae704141937af#subdirectory=flyteidl"
flytekit = 'git+https://github.com/alykhantejani/flytekit.git@1a0302028868355dd3e9a148f260a58a947a3730'

custom_image = ImageSpec(
    packages=[envd_plugin, flytekit, kubeflow_idl, kubeflow_plugin],
    apt_packages=["git"],
    name="kftensorflow-flyte-plugin",
    registry="localhost:30000",
)


@task(
task_config=TfJob(worker=Worker(replicas=1), ps=PS(replicas=1), chief=Chief(replicas=1, node_selectors={"something": "value"})),
container_image=custom_image
)
def launch(name: str) -> None:
    print(os.environ)



@workflow
def tf_job_test() -> None:
    launch(name="hello world")

Check all the applicable boxes

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

Related PRs

flyteidl: https://github.com/flyteorg/flyte/pull/4873

[welcome[bot]]

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