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.

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

Add JAX-defined augmentation examples #5426

Merged
merged 8 commits into from
Apr 16, 2024

Conversation

stiepan
Copy link
Member

@stiepan stiepan commented Apr 12, 2024

Category:

Other (e.g. Documentation, Tests, Configuration)

Description:

Adds a basic jax_function usage example and a more involved example presenting multi process shard-aware processing.

Additional information:

Affected modules and functionalities:

Key points relevant for the review:

Tests:

  • Existing tests apply
  • New tests added
    • Python tests
    • GTests
    • Benchmark
    • Other
  • N/A

Checklist

Documentation

  • Existing documentation applies
  • Documentation updated
    • Docstring
    • Doxygen
    • RST
    • Jupyter
    • Other
  • N/A

DALI team only

Requirements

  • Implements new requirements
  • Affects existing requirements
  • N/A

REQ IDs: N/A

JIRA TASK: DALI-3925

stiepan added 2 commits April 12, 2024 16:20
Signed-off-by: Kamil Tokarski <[email protected]>
Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@stiepan
Copy link
Member Author

stiepan commented Apr 12, 2024

!build

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14196087]: BUILD STARTED

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14196087]: BUILD PASSED

@awolant awolant self-assigned this Apr 15, 2024
Copy link

review-notebook-app bot commented Apr 15, 2024

View / edit / reply to this conversation on ReviewNB

awolant commented on 2024-04-15T10:50:50Z
----------------------------------------------------------------

Line #3.    os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # the other process passes 1

Maybe you could add some explanation that in a real world example this would have been set up by a launching script or something like that.


@klecki
Copy link
Contributor

klecki commented Apr 15, 2024

I think you should consider adding those files under:

  • qa/TL1_jupyter_multigpu/test.sh
  • qa/TL0_jupyter/test_jax.sh

Copy link

review-notebook-app bot commented Apr 15, 2024

View / edit / reply to this conversation on ReviewNB

klecki commented on 2024-04-15T15:00:10Z
----------------------------------------------------------------

that access samples -> accesses?


Copy link

review-notebook-app bot commented Apr 15, 2024

View / edit / reply to this conversation on ReviewNB

klecki commented on 2024-04-15T15:00:10Z
----------------------------------------------------------------

, the other with the -> will process

or while the other ...


@klecki
Copy link
Contributor

klecki commented Apr 15, 2024

I would also consider adding a runnable script containing the other process to the multipgpu example, starting it alongside the jupyter and testing that it actually works.

stiepan added 2 commits April 15, 2024 17:37
Signed-off-by: Kamil Tokarski <[email protected]>
@stiepan
Copy link
Member Author

stiepan commented Apr 15, 2024

!build

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14254191]: BUILD STARTED

Signed-off-by: Kamil Tokarski <[email protected]>
@stiepan
Copy link
Member Author

stiepan commented Apr 15, 2024

!build

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14257313]: BUILD STARTED

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14257313]: BUILD FAILED

Signed-off-by: Kamil Tokarski <[email protected]>
@stiepan
Copy link
Member Author

stiepan commented Apr 16, 2024

!build

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14273394]: BUILD STARTED

Signed-off-by: Kamil Tokarski <[email protected]>
@stiepan
Copy link
Member Author

stiepan commented Apr 16, 2024

!build

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14274963]: BUILD STARTED

@stiepan
Copy link
Member Author

stiepan commented Apr 16, 2024

!build

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14277138]: BUILD STARTED

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [14277138]: BUILD PASSED

@stiepan stiepan merged commit f3c920a into NVIDIA:main Apr 16, 2024
7 checks passed
stiepan added a commit that referenced this pull request Apr 16, 2024
* Add JAX operator examples: introducion and multi-gpu
* Fix typo in error message
---------

Signed-off-by: Kamil Tokarski <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants