more changes

Signed-off-by: Yuvraj <[email protected]>

.github/workflows/pythonbuild.yml

@@ -61,9 +61,9 @@ jobs:
           - flytekit-greatexpectations
           - flytekit-hive
           - flytekit-k8s-pod
+          - flytekit-kf-mpi
           - flytekit-kf-pytorch
           - flytekit-kf-tensorflow
-          - flytekit-kf-mpi
           - flytekit-papermill
           - flytekit-spark
           - flytekit-sqlalchemy

plugins/README.md

@@ -10,7 +10,7 @@ All Flytekit plugins maintained by the core team are added here. It is not neces
 | Hive Queries                 | ```bash pip install flytekitplugins-hive ```         | Installs SDK to author Hive Queries that can be executed on a configured hive backend using Flyte backend plugin            | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-spark.svg)](https://pypi.python.org/pypi/flytekitplugins-hive/)         | Backend       |
 | K8s distributed PyTorch Jobs | ```bash pip install flytekitplugins-kfpytorch ```    | Installs SDK to author Distributed pyTorch Jobs in python using Kubeflow PyTorch Operator                                   | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-spark.svg)](https://pypi.python.org/pypi/flytekitplugins-kfpytorch/)    | Backend       |
 | K8s native tensorflow Jobs   | ```bash pip install flytekitplugins-kftensorflow ``` | Installs SDK to author Distributed tensorflow Jobs in python using Kubeflow Tensorflow Operator                             | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-spark.svg)](https://pypi.python.org/pypi/flytekitplugins-kftensorflow/) | Backend       |
-| K8s native mpi Jobs          | ```bash pip install flytekitplugins-kfmpi ```        | Installs SDK to author Distributed mpi Jobs in python using Kubeflow mpi Operator                             | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-spark.svg)](https://pypi.python.org/pypi/flytekitplugins-kftensorflow/) | Backend       |
+| K8s native MPI Jobs          | ```bash pip install flytekitplugins-kfmpi ```        | Installs SDK to author Distributed MPI Jobs in python using Kubeflow MPI Operator                             | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-spark.svg)](https://pypi.python.org/pypi/flytekitplugins-kftensorflow/) | Backend       |
 | Papermill based Tasks        | ```bash pip install flytekitplugins-papermill ```    | Execute entire notebooks as Flyte Tasks and pass inputs and outputs between them and python tasks                           | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-spark.svg)](https://pypi.python.org/pypi/flytekitplugins-papermill/)    | Flytekit-only |
 | Pod Tasks                    | ```bash pip install flytekitplugins-pod ```          | Installs SDK to author Pods in python. These pods can have multiple containers, use volumes and have non exiting side-cars  | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-spark.svg)](https://pypi.python.org/pypi/flytekitplugins-pod/)          | Flytekit-only |
 | spark                        | ```bash pip install flytekitplugins-spark ```        | Installs SDK to author Spark jobs that can be executed natively on Kubernetes with a supported backend Flyte plugin         | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-spark.svg)](https://pypi.python.org/pypi/flytekitplugins-spark/)        | Backend       |

plugins/flytekit-kf-mpi/flytekitplugins/kfmpi/task.py

@@ -1,6 +1,6 @@
 """
-This Plugin adds the capability of running distributed tensorflow training to Flyte using backend plugins, natively on
-Kubernetes. It leverages `TF Job <https://github.com/kubeflow/mpi-operator>`_ Plugin from kubeflow.
+This Plugin adds the capability of running distributed MPI training to Flyte using backend plugins, natively on
+Kubernetes. It leverages `MPI Job <https://github.com/kubeflow/mpi-operator>`_ Plugin from kubeflow.
 """
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List
@@ -58,11 +58,19 @@ class MPIJob(object):
         num_launcher_replicas: Number of launcher server replicas to use
 
         slots: Number of slots per worker used in hostfile
+
+        per_replica_requests: [optional] lower-bound resources for each replica spawned for this job
+        (i.e. both for (main)master and workers).  Default is set by platform-level configuration.
+
+        per_replica_limits: [optional] upper-bound resources for each replica spawned for this job. If not specified
+        the scheduled resource may not have all the resources
     """
 
     slots: int
     num_launcher_replicas: int = 1
     num_workers: int = 1
+    per_replica_requests: Optional[Resources] = None
+    per_replica_limits: Optional[Resources] = None
 
 
 class MPIFunctionTask(PythonFunctionTask[MPIJob]):
@@ -98,7 +106,7 @@ def __init__(self, task_config: MPIJob, task_function: Callable, **kwargs):
             task_config=task_config,
             task_function=task_function,
             task_type=self._MPI_JOB_TASK_TYPE,
-            **kwargs,
+            **{**kwargs, "requests": task_config.per_replica_requests, "limits": task_config.per_replica_limits}
         )
 
     def get_command(self, settings: SerializationSettings) -> List[str]:

plugins/setup.py

@@ -14,9 +14,9 @@
     "flytekitplugins-great_expectations": "flytekit-greatexpectations",
     "flytekitplugins-hive": "flytekit-hive",
     "flytekitplugins-pod": "flytekit-k8s-pod",
+    "flytekitplugins-kfmpi": "flytekit-kf-mpi",
     "flytekitplugins-kfpytorch": "flytekit-kf-pytorch",
     "flytekitplugins-kftensorflow": "flytekit-kf-tensorflow",
-    "flytekitplugins-kfmpi": "flytekit-kf-mpi",
     "flytekitplugins-pandera": "flytekit-pandera",
     "flytekitplugins-papermill": "flytekit-papermill",
     "flytekitplugins-snowflake": "flytekit-snowflake",