Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: Flax #3123

Merged
merged 25 commits into from
Feb 23, 2023
Merged

feat: Flax #3123

merged 25 commits into from
Feb 23, 2023

Conversation

aarnphm
Copy link
Contributor

@aarnphm aarnphm commented Oct 20, 2022 •
edited
Loading

This PR adds support for Flax.

  • add data container for Jax.

I will add everything here and then cherry-pick out commits later.

Signed-off-by: Aaron Pham [email protected]

@codecov
Copy link

codecov bot commented Oct 20, 2022 •
edited
Loading

Codecov Report

Merging #3123 (4e46247) into main (346e185) will decrease coverage by 0.43%.
The diff coverage is 0.00%.

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #3123      +/-   ##
==========================================
- Coverage   30.92%   30.49%   -0.43%     
==========================================
  Files         140      143       +3     
  Lines       11671    11834     +163     
  Branches     1929     1948      +19     
==========================================
  Hits         3609     3609              
- Misses       7818     7981     +163     
  Partials      244      244
Impacted Files Coverage Δ
src/bentoml/_internal/frameworks/common/jax.py 0.00% <0.00%> (ø)
src/bentoml/_internal/frameworks/flax.py 0.00% <0.00%> (ø)
src/bentoml/_internal/runner/container.py 0.00% <0.00%> (ø)
src/bentoml/flax.py 0.00% <0.00%> (ø)

@aarnphm aarnphm linked an issue Oct 24, 2022 that may be closed by this pull request
feature: Jax ecosystem #1763
Closed
@aarnphm aarnphm removed a link to an issue Oct 24, 2022
feature: Jax ecosystem #1763
Closed
@aarnphm aarnphm linked an issue Oct 24, 2022 that may be closed by this pull request
feature: Jax ecosystem #1763
Closed
@aarnphm aarnphm mentioned this pull request Oct 24, 2022
Closed
@aarnphm aarnphm marked this pull request as ready for review October 25, 2022 18:33
@aarnphm aarnphm requested a review from a team as a code owner October 25, 2022 18:33
@aarnphm aarnphm requested review from ssheng and removed request for a team October 25, 2022 18:33
@aarnphm aarnphm changed the title feat(WIP): Jax ecosystem feat: Jax ecosystem Oct 25, 2022
@aarnphm aarnphm requested a review from larme November 2, 2022 10:09
@aarnphm aarnphm force-pushed the feat/flax-ecosystem branch 5 times, most recently from 938694c to 3b406aa Compare November 7, 2022 02:54
larme
larme requested changes Nov 15, 2022
View reviewed changes
examples/flax/MNIST/README.md Show resolved Hide resolved
examples/flax/MNIST/train.py Show resolved Hide resolved
src/bentoml/_internal/frameworks/common/jax.py Outdated Show resolved Hide resolved
src/bentoml/_internal/frameworks/flax.py Show resolved Hide resolved
tests/integration/frameworks/models/flax.py Outdated Show resolved Hide resolved
src/bentoml/_internal/frameworks/flax.py Outdated Show resolved Hide resolved
src/bentoml/_internal/frameworks/flax.py Show resolved Hide resolved
src/bentoml/_internal/frameworks/flax.py Outdated Show resolved Hide resolved
src/bentoml/_internal/frameworks/flax.py Outdated Show resolved Hide resolved
@aarnphm aarnphm requested a review from larme November 15, 2022 02:49
@aarnphm
Copy link
Contributor Author

aarnphm commented Feb 18, 2023

Hi @larme, PTAL.

I'm intending to use this to build some of the example projects using Flax.

@aarnphm aarnphm force-pushed the feat/flax-ecosystem branch 6 times, most recently from ee90559 to 0286f49 Compare February 18, 2023 20:25
@aarnphm
chore: export imports
fadb5e1 
Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
chore: address comments and convert payload to np.array
7651748 
Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
chore: enable plugins on CI only
5457a07 
Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
chore: add instruction to run standalone
a5eaed3 
Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
feat: adding hermetic options for training
8dac69f 
(default to False)

Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
fix: requirements
535d726 
Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
chore: add instruction for containerize
0396a27 
Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
chore: update readme
c681bb5 
Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
fix: tests
21863d7 
Signed-off-by: Aaron Pham <[email protected]>
@aarnphm
chore: add gpu requirements
934f1d1 
Signed-off-by: aarnphm-ec2-dev <[email protected]>
@aarnphm
fix: ignore user_defined_callback if it is a lambda function
a947056 
Signed-off-by: aarnphm-ec2-dev <[email protected]>
@aarnphm
tests: run async_predict inside test
b346a26 
Signed-off-by: aarnphm-ec2-dev <[email protected]>
@aarnphm
fix: typing
5672420 
Signed-off-by: Aaron <[email protected]>
@aarnphm
chore: update tests
e3c1021 
Signed-off-by: Aaron <[email protected]>
@aarnphm aarnphm requested a review from larme February 21, 2023 22:59
@aarnphm
Copy link
Contributor Author

aarnphm commented Feb 21, 2023 •
edited
Loading

I will follow up with a example PR for flax and frameworks docs afterward.

@aarnphm aarnphm requested a review from sauyon February 22, 2023 05:23
@sauyon
Copy link
Contributor

sauyon commented Feb 22, 2023

Remind me to take a look at this tomorrow if Larme doesn't get a chance to.

larme
larme approved these changes Feb 23, 2023
View reviewed changes
Copy link
Member

@larme larme left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍LGTM

Let's add an equality test in a later PR

@aarnphm
Copy link
Contributor Author

aarnphm commented Feb 23, 2023

👍LGTM

Let's add an equality test in a later PR

I have already added that.

@aarnphm aarnphm changed the title feat: Jax ecosystem feat: Flax Feb 23, 2023
@aarnphm aarnphm merged commit e5af0c3 into bentoml:main Feb 23, 2023
@aarnphm aarnphm deleted the feat/flax-ecosystem branch February 23, 2023 05:35
sauyon
sauyon reviewed Mar 2, 2023
View reviewed changes
pyproject.toml
@@ -236,7 +236,7 @@ exclude = '''
'''

[tool.pytest.ini_options]
addopts = ["-rfEX", "-pbentoml.testing.pytest.plugin"]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@aarnphm why was this in this PR anyway?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh because the test in the examples will always running train everytime the test is running.

THe plugins should be for our internal test only.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can revert this back imo.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why should it be for our internal tests only?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given we have implemented as part of our library, it seems to me like it should be general-purpose?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i will put up a PR to revert this.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@sauyon sauyon sauyon left review comments

@larme larme larme approved these changes

@ssheng ssheng Awaiting requested review from ssheng ssheng is a code owner automatically assigned from bentoml/bentoml-reviewers

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

feature: Jax ecosystem
3 participants
@aarnphm @sauyon @larme