Providing containers with newer Python versions(?)

[olupton]

Context
JAX's policy is to support Python versions up to 45 months old. At the moment, wheels are published for Python 3.{9,10,11,12} (e.g. here). Python 3.8 is not yet EOL (https://devguide.python.org/versions), but has fallen outside the 45 month window and was dropped from JAX 0.4.14.

The older supported Python versions are several years old, and various new features have been introduced that may help improve the JAX user experience. A recent example is jax-ml/jax#18794, which should improve the profiling experience by improving the metadata linking HLO operations back to Python source code. These improvements are only available to Python 3.11+ users.

Containers
At the moment, we publish containers (https://github.com/NVIDIA/JAX-Toolbox/pkgs/container/jax) based on Ubuntu 22.04 LTS, which ships with Python 3.10. These are built on top of CUDA containers, which are not available for newer versions of Ubuntu.

The next LTS version of Ubuntu will be 24.04, which is due in April 2024 (https://wiki.ubuntu.com/Releases) and seems likely to be used on Python 3.12. It will clearly take some time from the Ubuntu version being released to CUDA containers based on it becoming available.

Non-system Python versions
It's possible to install Python 3.11 using the system package manager in the existing Ubuntu 22.04-based containers, but one gets a broken rc version (jax-ml/jax#18625 (comment)).

The deadsnakes PPA can be used (https://launchpad.net/~deadsnakes/+archive/ubuntu/ppa) to install up-to-date patch versions of various minor Python 3 versions, but it is a third-party repository (albeit a very popular one). If we went down this route, we would be committed to installing all dependencies via pip as opposed to the system package manager (numpy etc.), but that is existing practice anyway.

Discussion
We should use, and encourage others to use, these features rather than sitting too far behind the cutting edge.

#34 also explored the the idea that we should test more Python versions in our CI.

In a local experiment, it wasn't very difficult to build + install + run with a deadsnakes build of Python 3.11; some minor fixes were needed because scripts assumed that python and pip refer to the same Python version, which wasn't trivially the case.

cc: @nouiz @yhtang

[yhtang]

Note that we already have scripts internally to build Python 3.X from source. The question remains is what strategy we should take in terms of shipping multiple Python versions, and what the ensuing testing/QA implications are.

[nouiz]

For now, we will only support the current python from the container. Closing.