Modify version test to consider "rc" versions as well

I was testing the RC promotion workflow and found that the version test failed as it does not consider pre-releases. Therefore, this commit modifies the `VERSION_PATTERN` to also consider "rc" wheels.

Fixes https://github.com/jax-ml/jax/actions/runs/13705984545/job/38331236497

PiperOrigin-RevId: 735444828
This commit is contained in:
Nitin Srinivasan 2025-03-10 11:09:40 -07:00 committed by jax authors
parent 5cb29949d4
commit d41e96835b

View File

@ -24,11 +24,15 @@ import jax
from jax._src.lib import check_jaxlib_version
from jax._src import test_util as jtu
# This is a subset of the full PEP440 pattern; for example we skip pre & post releases
# This is a subset of the full PEP440 pattern; for example we skip post releases
VERSION_PATTERN = re.compile(r"""
^ # start of string
(?P<version>[0-9]+\.[0-9]+\.[0-9]+) # main version; like '0.4.16'
(?:\.dev(?P<dev>[0-9]+))? # optional dev version; like '.dev20230908'
(?:
(?:rc(?P<rc>[0-9]+))? # optional rc version; like 'rc1'
| # or
(?:\.dev(?P<dev>[0-9]+))? # optional dev version; like '.dev20230908'
)?
(?:\+(?P<local>[a-zA-Z0-9_.]+))? # optional local version; like '+g6643af3c3'
$ # end of string
""", re.VERBOSE)
@ -170,6 +174,18 @@ class JaxVersionTest(unittest.TestCase):
self.assertEqual(version, f"{base_version}.dev20250101+1c0f1076erc1")
self.assertValidVersion(version)
with jtu.set_env(
JAX_RELEASE="1",
JAXLIB_RELEASE=None,
JAX_NIGHTLY=None,
JAXLIB_NIGHTLY=None,
WHEEL_VERSION_SUFFIX="rc0",
):
with assert_no_subprocess_call():
version = jax.version._get_version_for_build()
self.assertEqual(version, f"{base_version}rc0")
self.assertValidVersion(version)
def testVersions(self):
check_jaxlib_version(jax_version="1.2.3", jaxlib_version="1.2.3",
minimum_jaxlib_version="1.2.3")