Feat: Add torch elastic task type #1583

fg91 · 2023-04-05T20:13:42Z

TL;DR

This PR brings torch elastic training (torchrun) to the pytorch plugin:

from flytekitplugins.kfpytorch import Elastic

@task(
    task_config=Elastic(
        replicas=4,
        nproc_per_node=4,
        ...
    ),
    ...
)
def train(...):
    ...

Type

Bug Fix
Feature
Plugin

Are all requirements met?

Complete description

See flyteorg/flyte#3614 for motivation and a more detailed description.

Tracking Issue

flyteorg/flyte#3614

Follow-up issue

codecov · 2023-04-05T20:28:13Z

Codecov Report

Merging #1583 (d5ab6ac) into master (0063b29) will increase coverage by 0.18%.
The diff coverage is 79.06%.

@@            Coverage Diff             @@
##           master    #1583      +/-   ##
==========================================
+ Coverage   69.93%   70.12%   +0.18%     
==========================================
  Files         319      328       +9     
  Lines       29569    29973     +404     
  Branches     5332     5408      +76     
==========================================
+ Hits        20680    21018     +338     
- Misses       8370     8419      +49     
- Partials      519      536      +17

Impacted Files	Coverage Δ
flytekit/core/task.py	`33.33% <0.00%> (-0.76%)`	⬇️
flytekit/image_spec/__init__.py	`0.00% <0.00%> (ø)`
flytekit/loggers.py	`0.00% <0.00%> (ø)`
flytekit/tools/translator.py	`75.17% <ø> (+0.78%)`	⬆️
plugins/setup.py	`0.00% <ø> (ø)`
setup.py	`0.00% <ø> (ø)`
tests/flytekit/unit/cli/pyflyte/test_backfill.py	`100.00% <ø> (ø)`
tests/flytekit/unit/core/test_type_engine.py	`98.62% <ø> (ø)`
flytekit/remote/remote.py	`39.93% <20.00%> (-0.14%)`	⬇️
flytekit/tools/repo.py	`81.66% <33.33%> (-1.24%)`	⬇️
... and 41 more

... and 7 files with indirect coverage changes

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

fg91 · 2023-04-08T20:27:23Z

plugins/flytekit-torch-elastic/flytekitplugins/torchelastic/task.py

@@ -13,8 +13,7 @@

 @dataclass
 class Elastic(object):
-    min_replicas: int = 1
-    max_replicas: int = 1
+    nnodes: typing.Union[int, str] = 1


Let's use the way torchrun exposes this to the user, e.g. --nnodes 2:3

plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py

Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: Fabio Grätz <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: Fabio Graetz <[email protected]>

* Unify sqlalchemy Dockerfiles Signed-off-by: eduardo apolinario <[email protected]> * Simplifyte sqlalchemy Dockerfile Signed-off-by: eduardo apolinario <[email protected]> * Set pythonpath to /root Signed-off-by: eduardo apolinario <[email protected]> --------- Signed-off-by: eduardo apolinario <[email protected]> Co-authored-by: eduardo apolinario <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: Fabio Graetz <[email protected]>

* wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * pyflyte build Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * Support serialize and package Signed-off-by: Kevin Su <[email protected]> * more tests Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * move to plugin Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * fixed tested Signed-off-by: Kevin Su <[email protected]> * fixed tested Signed-off-by: Kevin Su <[email protected]> * more tests Signed-off-by: Kevin Su <[email protected]> * Add support passing yaml in pyflyte run Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

…ror messages (#1582) * improve input type conversion error Signed-off-by: Niels Bantilan <[email protected]> * fix lint Signed-off-by: Niels Bantilan <[email protected]> * fix Signed-off-by: Niels Bantilan <[email protected]> * add tests Signed-off-by: Niels Bantilan <[email protected]> * add tests Signed-off-by: Niels Bantilan <[email protected]> * add rich Signed-off-by: Niels Bantilan <[email protected]> * fix lint Signed-off-by: Niels Bantilan <[email protected]> * remove prototyping script, update loggers Signed-off-by: Niels Bantilan <[email protected]> * update __init__.py Signed-off-by: Niels Bantilan <[email protected]> * update logger Signed-off-by: Niels Bantilan <[email protected]> * update logger Signed-off-by: Niels Bantilan <[email protected]> * fix GE and pandera tests Signed-off-by: Niels Bantilan <[email protected]> * fix lint Signed-off-by: Niels Bantilan <[email protected]> --------- Signed-off-by: Niels Bantilan <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* Add support nested FlyteFile in pyflyte run Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* Set cache_regions=True in the case of s3fs * Linting Signed-off-by: eduardo apolinario <[email protected]> * Fix tests Signed-off-by: eduardo apolinario <[email protected]> * More linting Signed-off-by: eduardo apolinario <[email protected]> --------- Signed-off-by: eduardo apolinario <[email protected]> Co-authored-by: eduardo apolinario <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* set default ComSpec when running on Windows Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* add activate-launchplan command to pyflyte Signed-off-by: Artem Petrov <[email protected]> * fix linting Signed-off-by: Artem Petrov <[email protected]> * rework launchplan command Signed-off-by: Artem Petrov <[email protected]> * fix linting Signed-off-by: Artem Petrov <[email protected]> --------- Signed-off-by: Artem Petrov <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: Yee Hing Tong <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* checkpoint Signed-off-by: Mike Zhong <[email protected]> * Experimental implementation works well. Instead of messing with the class, we utilize the interpolation and context to dynamically generate a wrapper around the desired script. In the wrapper, we cd to the ctx.working_directory, export the env variables (handled by an additional method to convert dict to str), and then pass the arguments to the script Signed-off-by: Mike Zhong <[email protected]> * fix spacing Signed-off-by: Mike Zhong <[email protected]> * remove breakpoint Signed-off-by: Mike Zhong <[email protected]> * Output ctx.working_directory as single output Signed-off-by: Mike Zhong <[email protected]> * more doc strings Signed-off-by: Mike Zhong <[email protected]> * more comments Signed-off-by: Mike Zhong <[email protected]> * Added comments Signed-off-by: Mike Zhong <[email protected]> * Add tests and test files from other branch Signed-off-by: Mike Zhong <[email protected]> * minor fix, have function return the instance rather than create. It seems to get registered when flyte packages your project Signed-off-by: Mike Zhong <[email protected]> * fix tests Signed-off-by: Mike Zhong <[email protected]> * remove set flags not supported by sh Signed-off-by: Mike Zhong <[email protected]> * fix spellcheck lint errors Signed-off-by: Mike Zhong <[email protected]> * don't have a windows equivalent test script so bypassing those tests for now Signed-off-by: Mike Zhong <[email protected]> * Add typing to make_export_string_from_env_dict params Signed-off-by: Mike Zhong <[email protected]> * fixup doc string Signed-off-by: Mike Zhong <[email protected]> * Refactored the new behavior out into a new class. Did not change implementation details at all Signed-off-by: Mike Zhong <[email protected]> * Address linter issues and up test cov Signed-off-by: Mike Zhong <[email protected]> * Skip test on windows, no equivalent script Signed-off-by: Mike Zhong <[email protected]> * Run black and isort, address SC2236 Signed-off-by: Mike Zhong <[email protected]> * Refactored class name to remove _, make utility function require name so multiple uses in a workflow don't break Signed-off-by: Mike Zhong <[email protected]> * fix utility function Signed-off-by: Mike Zhong <[email protected]> * fix utility function call in tests Signed-off-by: Mike Zhong <[email protected]> * Fix isort on test_shell.py Signed-off-by: Mike Zhong <[email protected]> * adding logging settings to papermill plugin Signed-off-by: Calvin Leather <[email protected]> * set level to info logging * improved comments/docs --------- Signed-off-by: Mike Zhong <[email protected]> Signed-off-by: Calvin Leather <[email protected]> Co-authored-by: Mike Zhong <[email protected]> Co-authored-by: Mike Zhong <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* set lower bound for s3fs Signed-off-by: Kevin Su <[email protected]> * update doc-requirements.txt Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* Add horovod task to mpi plugin Signed-off-by: byhsu <[email protected]> * Remove unused Signed-off-by: byhsu <[email protected]> * add test Signed-off-by: byhsu <[email protected]> * inherit from mpifunctiontask Signed-off-by: byhsu <[email protected]> * fix format Signed-off-by: byhsu <[email protected]> * fix format Signed-off-by: byhsu <[email protected]> --------- Signed-off-by: byhsu <[email protected]> Co-authored-by: byhsu <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: esad <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

) Bumps [tensorflow](https://github.com/tensorflow/tensorflow) from 2.10.0 to 2.11.1. - [Release notes](https://github.com/tensorflow/tensorflow/releases) - [Changelog](https://github.com/tensorflow/tensorflow/blob/master/RELEASE.md) - [Commits](tensorflow/tensorflow@v2.10.0...v2.11.1) --- updated-dependencies: - dependency-name: tensorflow dependency-type: direct:development ... Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: eduardo apolinario <[email protected]> Co-authored-by: eduardo apolinario <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* pyflyte run imperative workflow Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* Backfill fix - Backfill was using incorrect arguments - backfill should use the argument that user provides or none at all Signed-off-by: Ketan Umare <[email protected]> * Updated code Signed-off-by: Ketan Umare <[email protected]> * fixed unit test Signed-off-by: Ketan Umare <[email protected]> --------- Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

* Pyflyte prettified Signed-off-by: Ketan Umare <[email protected]> * added dependency Signed-off-by: Ketan Umare <[email protected]> --------- Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: Fabio Graetz <[email protected]>

…'elastic'] Signed-off-by: Fabio Graetz <[email protected]>

Signed-off-by: Fabio Graetz <[email protected]>

fg91 · 2023-04-23T11:01:54Z

Tests are failing because flyteidl needs to be released and updated here.

fg91 mentioned this pull request Apr 5, 2023

Feat: Add torchrun plugin #1576

Closed

8 tasks

fg91 commented Apr 8, 2023

View reviewed changes

fg91 commented Apr 11, 2023

View reviewed changes

plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py Outdated Show resolved Hide resolved

fg91 force-pushed the fabio/feat/torch-elastic-plugin branch from bcb423f to e647383 Compare April 16, 2023 13:48

fg91 mentioned this pull request Apr 22, 2023

[Core feature] Support torch elastic training/torchrun flyteorg/flyte#3614

Closed

2 tasks

fg91 commented Apr 22, 2023

View reviewed changes

plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py Outdated Show resolved Hide resolved

fg91 self-assigned this Apr 23, 2023

fg91 added the enhancement New feature or request label Apr 23, 2023

Fabio Grätz and others added 21 commits April 23, 2023 10:44

Add torch elastic task type

c3de752

Signed-off-by: Fabio Graetz <[email protected]>

updated

fd0ae59

Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

updated

dd2a334

Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Don't pass min_nodes, max_nodes but nnodes only

855bf84

Signed-off-by: Fabio Grätz <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Add docstrings

4245bf6

Signed-off-by: Fabio Grätz <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Cleanup test

a6ac14d

Signed-off-by: Fabio Grätz <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Replace wrong occurences of mpi in torch-elastic plugin

35930b9

Signed-off-by: Fabio Grätz <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Extend kfpytorch plugin README

4245d39

Signed-off-by: Fabio Graetz <[email protected]>

Move elastic task into existing kf-pytorch plugin

40c83dd

Signed-off-by: Fabio Graetz <[email protected]>

Add Elastic config to plugin's __init__

440ea47

Signed-off-by: Fabio Graetz <[email protected]>

Set elastic config in pytorchjob proto

2d916af

Signed-off-by: Fabio Graetz <[email protected]>

removed unnecessary model files and simplified codebase

cc960b6

Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Configure rdzv endpoint

7f31a5f

Signed-off-by: Fabio Graetz <[email protected]>

Configure worker count also for elastic training

bf97ece

Signed-off-by: Fabio Graetz <[email protected]>

Fix exception scope and handle non-rank-0 outputs

fcc54f4

Signed-off-by: Fabio Graetz <[email protected]>

Remove todo item about handling exception scope, now done

e5f7c58

Signed-off-by: Fabio Graetz <[email protected]>

Add note about c10d backend

59ba4a5

Signed-off-by: Fabio Graetz <[email protected]>

Add classifier for Python 3.11 (#1589)

7367757

Signed-off-by: Fabio Graetz <[email protected]>

pingsutw and others added 18 commits April 23, 2023 10:44

nit (#1596)

25b00b1

Signed-off-by: Yee Hing Tong <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Set lower bound for s3fs (#1598)

e701f96

* set lower bound for s3fs Signed-off-by: Kevin Su <[email protected]> * update doc-requirements.txt Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

add ImageRender class for Flyte Decks (#1599)

c7d69fa

Signed-off-by: esad <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Handle batchable lists of length 0 (#1600)

6a85430

Signed-off-by: eduardo apolinario <[email protected]> Co-authored-by: eduardo apolinario <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Pyflyte prettified (#1602)

c46e2e6

* Pyflyte prettified Signed-off-by: Ketan Umare <[email protected]> * added dependency Signed-off-by: Ketan Umare <[email protected]> --------- Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Fabio Graetz <[email protected]>

Let user set min and max replicas explicitly and not via nnodes

63db8b4

Signed-off-by: Fabio Graetz <[email protected]>

Catch torch import error and configure for flytekitplugins-kfpytorch[…

5d24b29

…'elastic'] Signed-off-by: Fabio Graetz <[email protected]>

Add more tests

c58eaa2

Signed-off-by: Fabio Graetz <[email protected]>

Update plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py

379a68b

Signed-off-by: Fabio Graetz <[email protected]>

fg91 force-pushed the fabio/feat/torch-elastic-plugin branch from 223c82d to 379a68b Compare April 23, 2023 10:45

Lint

d5ab6ac

Signed-off-by: Fabio Graetz <[email protected]>

fg91 marked this pull request as ready for review April 23, 2023 11:01

fg91 requested review from wild-endeavor, kumare3, eapolinario, pingsutw and cosmicBboy as code owners April 23, 2023 11:01

fg91 closed this Apr 23, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feat: Add torch elastic task type #1583

Feat: Add torch elastic task type #1583

fg91 commented Apr 5, 2023 •

edited

Loading

codecov bot commented Apr 5, 2023 •

edited

Loading

fg91 Apr 8, 2023

fg91 commented Apr 23, 2023

Feat: Add torch elastic task type #1583

Feat: Add torch elastic task type #1583

Conversation

fg91 commented Apr 5, 2023 • edited Loading

TL;DR

Type

Are all requirements met?

Complete description

Tracking Issue

Follow-up issue

codecov bot commented Apr 5, 2023 • edited Loading

Codecov Report

fg91 Apr 8, 2023

Choose a reason for hiding this comment

fg91 commented Apr 23, 2023

fg91 commented Apr 5, 2023 •

edited

Loading

codecov bot commented Apr 5, 2023 •

edited

Loading