Change flytekit Pytorch, TFJob and MPI plugins to use new kubeflow co…

…nfig (#1627) * upgrade tensorflow plugin to v1 Signed-off-by: Yubo Wang <[email protected]> * minor fix Signed-off-by: Yubo Wang <[email protected]> * fix tests and lints Signed-off-by: Yubo Wang <[email protected]> * move models file into task make backward compatible Signed-off-by: Yubo Wang <[email protected]> Signed-off-by: Yubo Wang <[email protected]> * add code example in README Signed-off-by: Yubo Wang <[email protected]> * bump flyteidl Signed-off-by: Yubo Wang <[email protected]> * add pytorch Signed-off-by: Yubo Wang <[email protected]> * add mpi and fix requirements.txt Signed-off-by: Yubo Wang <[email protected]> * lint and fmt Signed-off-by: Yubo Wang <[email protected]> * Regenerate requirements files using python 3.8 Signed-off-by: eduardo apolinario <[email protected]> --------- Signed-off-by: Yubo Wang <[email protected]> Signed-off-by: eduardo apolinario <[email protected]> Co-authored-by: Yubo Wang <[email protected]> Co-authored-by: eduardo apolinario <[email protected]>
flyteorg · Jun 29, 2023 · 368602b · 368602b
1 parent da597e5
commit 368602b
Show file tree

Hide file tree

Showing 19 changed files with 1,672 additions and 368 deletions.
diff --git a/plugins/flytekit-kf-mpi/README.md b/plugins/flytekit-kf-mpi/README.md
@@ -8,4 +8,67 @@ To install the plugin, run the following command:
 pip install flytekitplugins-kfmpi
 ```
 
-_Example coming soon!_
+## Code Example
+MPI usage:
+```python
+    @task(
+        task_config=MPIJob(
+            launcher=Launcher(
+                replicas=1,
+            ),
+            worker=Worker(
+                replicas=5,
+                requests=Resources(cpu="2", mem="2Gi"),
+                limits=Resources(cpu="4", mem="2Gi"),
+            ),
+            slots=2,
+        ),
+        cache=True,
+        requests=Resources(cpu="1"),
+        cache_version="1",
+    )
+    def my_mpi_task(x: int, y: str) -> int:
+        return x
+```
+
+
+Horovod Usage:
+You can override the command of a replica group by:
+```python
+    @task(
+        task_config=HorovodJob(
+            launcher=Launcher(
+                replicas=1,
+                requests=Resources(cpu="1"),
+                limits=Resources(cpu="2"),
+            ),
+            worker=Worker(
+                replicas=1,
+                command=["/usr/sbin/sshd", "-De", "-f", "/home/jobuser/.sshd_config"],
+                restart_policy=RestartPolicy.NEVER,
+            ),
+            slots=2,
+            verbose=False,
+            log_level="INFO",
+        ),
+    )
+    def my_horovod_task():
+        ...
+```
+
+
+
+
+## Upgrade MPI Plugin from V0 to V1
+MPI plugin is now updated from v0 to v1 to enable more configuration options.
+To migrate from v0 to v1, change the following:
+1. Update flytepropeller to v1.6.0
+2. Update flytekit version to v1.6.2
+3. Update your code from:
+```
+    task_config=MPIJob(num_workers=10),
+```
+to
+```
+    task_config=MPIJob(worker=Worker(replicas=10)),
+```
diff --git a/plugins/flytekit-kf-mpi/flytekitplugins/kfmpi/__init__.py b/plugins/flytekit-kf-mpi/flytekitplugins/kfmpi/__init__.py
@@ -10,4 +10,4 @@
    MPIJob
 """
 
-from .task import HorovodJob, MPIJob
+from .task import CleanPodPolicy, HorovodJob, Launcher, MPIJob, RestartPolicy, RunPolicy, Worker
-Original file line number
+Diff line change
@@ Expand Up / @@ -10,4 +10,4 @@ @@
        MPIJob
     """
-    from .task import HorovodJob, MPIJob
+    from .task import CleanPodPolicy, HorovodJob, Launcher, MPIJob, RestartPolicy, RunPolicy, Worker