jax
jax copied to clipboard
google
Add multi host pjit tests
Added multi host pjit tests to the existing test setup of 2 nodes and 1 process per GPU
I don't think there's any need to repeat the tests here. Why not just subclass the original pjit tests and run them in Slurm?
Is that straightforward (basically, are the upstream single host pjit tests extensible enough that they could be run on multi-host multi-GPU scenario)? I'm asking this because these tests needed to be adapted with some significant effort from TPU -> GPU case.
I agree that in the longer term the existing test cases should be generic enough to be run on any setup/backend, but for now I think we should get these tests in so that they are testing at least something other than just psum and pmap in the GPU CI.
(Also, @skye mentioned that she was perhaps looking at some of those upstream pjit tests and making them generic?)
After #5498 is complete, We can support most of extended query mode now, including:
- extended query mode message flow
- all text format param
- all text format result
- all kind of prepare statement with param
There are still some remained type format not be supported, including:
- [x] binary format param
- [ ] struct
- [ ] list
- [x] interval
- [x] timestampz
- [x] binary format result
- [ ] struct
- [ ] list
- [x] interval
- [x] timestampz
If the client uses text format, extended query mode can be used. If the client uses binary format, except the type still not be supported, the extended query mode can be used too.