add GPU tests on Bowie cluster

[Delcior]

No description provided.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87.47%. Comparing base (ca9db44) to head (844d7eb).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1939   +/-   ##
=======================================
  Coverage   87.47%   87.47%           
=======================================
  Files         430      430           
  Lines       11159    11159           
=======================================
  Hits         9761     9761           
  Misses       1398     1398           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[slayoo]

the new JAX backend also support GPU
anyhow, wouldn't it be simpler to just run all tests?

[Delcior]

There's a lot of failing tests and output on UI is absolutely unreadable. I thought it'd be nice to have an option to run particular subset of tests or even a signle test to make it easier to debug (this feature works only when you trigger GPU Tests via workflow_dispatch. For PRs it always runs all the tests).

Two questions:

1. Do you want to keep 'run custom tests' feature for workflow_dispatch?
2. By "just run all the tests" you mean tests/unit_tests or tests/?

[slayoo]

I meant the -k option

[Delcior]

With -k 'GPU or ThrustRTC' we have:

• pipeline duration: ~7-8 minutes
• tests collected: ~400

without -k 'GPU or ThrustRTC' we have:

• pipeline duration: > 1h (I don't know exactly how long, because I hit my timeout limit set in workflow. Link)
• tests collected: ~2000
• duplicate work (as these tests already run in other workflows)

I see 3 options:

1. Leave -k GPU or ThrustRTC'. AI told me that this option currently omits 4 GPU tests.
2. Leave -k GPU or ThrustRTC', but start an effort to mark GPU tests using pytest markers so that we can select the GPU tests via -m option (much more reliable than selecting by name as we currently attempt to do)
3. Increase the timeout, remove -k GPU or ThrustRTC' and run all the tests with one run taking > 1h.

I have very limited knowledge about PySDM codebase, but I think the safest and most profitable option is 2.
What do you think?

Side note: pytest runs the tests on a single core, so I installed pytest-xdist to use all 4 cores from the Jetson node, but it didn't help much as it also hit the 1h timeout link. (not saying that other options/optimizations wouldn't help, just highlighting our hardware limitations).

[slayoo]

Regarding multi-threading, I'd expect that pytest parallelisation will not help much as most of the PySDM logic uses multi-threading internally.

I agree that there is no point in running all tests on Bowie, since we can test on ARM64 on GitHub. The point is how to pick GPU tests. These are effectively already marked with GPU via the backend_class and backend_instance fixtures which resolve to GPU, hence "-k GPU" is kind of OK. However, the direction we're embracing is to phase out the ThrustRTC backend at all, in favour of the new JAX backend (already in main), and the -k GPU skips JAX (even if we flag with GPU label). So, for now, let's perhaps do -k GPU or ThrustRTC or JAX?

@Bodzio-2, how to make JAX tests actually use the GPU?

[Bodzio-2]

JAX has a config option to set the default device. That could be fit into the backend initialization to select based on user preference. Also, from what I've read JAX doesn't automatically do any multi-threading on CPU.

*GPU on JAX works only on Linux or dubiously via WSL

[slayoo]

relying on a "default" setting could be tricky as we do want to be able to run JAX-CPU and JAX-GPU simulations in one notebook, or at least warn the user if such setting is not possible

[Bodzio-2]

You can run code on a specific device wrapping it with a: with jax.default_device(jax.devices("gpu")[0]):.

For a more user-friendly experience I'd imagine an argument to the JAX backend constructor to attempt using gpu/cpu (with a warning/exception if a device wasn't found).

[slayoo]

would something like this be OK: #1941
(haven't tested, just a draft)

[Delcior]

TLDR: GPU tests with JAX will not work at all on our current Bowie Cluster due to CUDA == 10.2.

PySDM uses JAX >= 0.8.0, but according to the documentation this package alone gives support for CPU-only backend. We'd need to install jax[cuda*] to get the GPU one.

JAX supports CUDA >= 12 (docs). Support for CUDA 10.x was dropped ~5 years ago (link) and our worker nodes are equipped with CUDA==10.2 which we cannot upgrade. That means, in the current setup, we will not be able to run these JAX GPU tests on the Bowie Cluster.