more changes

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

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

@@ -11,10 +11,21 @@
 from flytekit import PythonFunctionTask
 from flytekit.extend import SerializationSettings, TaskPlugins
 from flytekit.models import common as _common
-from flytekit.models import model as _model
+from flytekit.models import task as _task_model
 
 
 class MPIJobModel(_common.FlyteIdlEntity):
+    """It will define the model for MPI plugin
+
+    Args:
+        num_workers: integer determining the number of worker replicas spawned in the cluster for this job
+        (in addition to 1 master).
+
+        num_launcher_replicas: Number of launcher server replicas to use
+
+        slots: Number of slots per worker used in hostfile
+
+    """
     def __init__(self, num_workers, num_launcher_replicas, slots):
         self._num_workers = num_workers
         self._num_launcher_replicas = num_launcher_replicas
@@ -70,8 +81,8 @@ class MPIJob(object):
     slots: int
     num_launcher_replicas: int = 1
     num_workers: int = 1
-    per_replica_requests: Optional[_model.Resources] = None
-    per_replica_limits: Optional[_model.Resources] = None
+    per_replica_requests: Optional[_task_model.Resources] = None
+    per_replica_limits: Optional[_task_model.Resources] = None
 
 
 class MPIFunctionTask(PythonFunctionTask[MPIJob]):