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

Update PJRT plugin API version to 0.57 #19241

Merged
merged 3 commits into from
Nov 25, 2024

Conversation

PragmaTwice
Copy link
Member

@PragmaTwice PragmaTwice commented Nov 21, 2024

It closes #19223.

integrations/pjrt/third_party/pjrt_c_api/xla/pjrt/c/pjrt_c_api.h is updated to the latest (with API version from 0.38 to 0.57), fetching from https://github.com/openxla/xla/blob/a454e14ab0b10e35fb8ad73bd6db7d93782114f6/xla/pjrt/c/pjrt_c_api.h.

A blank implementation of PJRT_Plugin_Attributes is now provided since an unimplemented PJRT_Plugin_Attributes will lead to initialization failure of PJRT plugin (and thus crashes) in recent versions of PJRT clients.

Also the JAX version in the CI workflow is updated from 0.4.20 to 0.4.35 and subsequently more tests can be enabled.

ci-exactly: build_packages, test_pjrt

Signed-off-by: PragmaTwice <twice@apache.org>
@PragmaTwice PragmaTwice force-pushed the update-pjrt-api branch 2 times, most recently from 8e96573 to cd195e2 Compare November 23, 2024 09:16
@PragmaTwice PragmaTwice marked this pull request as ready for review November 23, 2024 09:19
Signed-off-by: PragmaTwice <twice@apache.org>
Signed-off-by: PragmaTwice <twice@apache.org>
Copy link
Member

@ScottTodd ScottTodd left a comment

Choose a reason for hiding this comment

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

Thanks!

Comment on lines +2167 to +2170
// PJRT_Plugin_Attributes should be implemented since it will always be
// called from the PJRT client in the initial phase.
// here we provide a blank implementation to avoid crash due to unimplemented.
api->PJRT_Plugin_Attributes =
Copy link
Member

Choose a reason for hiding this comment

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

Ah, classic. An internal bug link for API documentation.
openxla/xla@619ca7b

// Returns an array of plugin attributes which are key-value pairs. One example
// attribute is the minimum supported StableHLO version.
// TODO(b/280349977): standardize the list of attributes.
typedef PJRT_Error* PJRT_Plugin_Attributes(PJRT_Plugin_Attributes_Args* args);

Fine to implement as you have it here for now.

@ScottTodd ScottTodd added the integrations/pjrt OpenXLA PJRT Integration Work label Nov 25, 2024
@ScottTodd ScottTodd merged commit ee6749d into iree-org:main Nov 25, 2024
25 checks passed
Groverkss pushed a commit to Groverkss/iree that referenced this pull request Dec 1, 2024
It closes iree-org#19223.

`integrations/pjrt/third_party/pjrt_c_api/xla/pjrt/c/pjrt_c_api.h` is
updated to the latest (with API version from 0.38 to 0.57), fetching
from
https://github.com/openxla/xla/blob/a454e14ab0b10e35fb8ad73bd6db7d93782114f6/xla/pjrt/c/pjrt_c_api.h.

A blank implementation of `PJRT_Plugin_Attributes` is now provided since
an unimplemented `PJRT_Plugin_Attributes` will lead to initialization
failure of PJRT plugin (and thus crashes) in recent versions of PJRT
clients.

Also the JAX version in the CI workflow is updated from 0.4.20 to 0.4.35
and subsequently more tests can be enabled.

ci-exactly: build_packages, test_pjrt

---------

Signed-off-by: PragmaTwice <twice@apache.org>
giacs-epic pushed a commit to giacs-epic/iree that referenced this pull request Dec 4, 2024
It closes iree-org#19223.

`integrations/pjrt/third_party/pjrt_c_api/xla/pjrt/c/pjrt_c_api.h` is
updated to the latest (with API version from 0.38 to 0.57), fetching
from
https://github.com/openxla/xla/blob/a454e14ab0b10e35fb8ad73bd6db7d93782114f6/xla/pjrt/c/pjrt_c_api.h.

A blank implementation of `PJRT_Plugin_Attributes` is now provided since
an unimplemented `PJRT_Plugin_Attributes` will lead to initialization
failure of PJRT plugin (and thus crashes) in recent versions of PJRT
clients.

Also the JAX version in the CI workflow is updated from 0.4.20 to 0.4.35
and subsequently more tests can be enabled.

ci-exactly: build_packages, test_pjrt

---------

Signed-off-by: PragmaTwice <twice@apache.org>
Signed-off-by: Giacomo Serafini <179146510+giacs-epic@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
integrations/pjrt OpenXLA PJRT Integration Work
Projects
None yet
Development

Successfully merging this pull request may close these issues.

The API version of PJRT plugin is lower than jax>0.4.20
2 participants