Repository: jax-ml/jax
Pull request #24861 (2024-11-12 18:38:06)

[JAX/version] Support adding custom version suffix.
This PR attempts to add a custom version suffix for JAX dev builds.

The versioning is important when people want to hash the compilation result of the JIT program. The users can use whatever local version of XLA repo and pass a version suffix constructed on their own side as part of the hash key.

This will greatly enhance the cache key especially when developers are doing dev builds based on different commits or local change.

yliu120: > Can you please squash your commits? Thanks!

Done.

Patch:
diff --git a/jax/version.py b/jax/version.py
index c27caf979..3e8a8291e 100644
--- a/jax/version.py
+++ b/jax/version.py
@@ -60,7 +60,11 @@ def _version_from_git_tree(base_version: str) -> str | None:
   except:
     return None
   else:
-    return f"{base_version}.dev{datestring}+{commit_hash}"
+    version = f"{base_version}.dev{datestring}+{commit_hash}"
+    suffix = os.environ.get("JAX_CUSTOM_VERSION_SUFFIX", None)
+    if suffix:
+      return version + "." + suffix
+    return version
 
 
 def _get_version_for_build() -> str:
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index af2fab8ed..e7ae7fe71 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -21,8 +21,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "9ab7d704d7fe7e73fc3976adc2ccec070bc9a2ea"
-XLA_SHA256 = "6944ceaa425cacd30a54cca3cd6c4cb88b79f219d421fb97fa87ffbf06007143"
+XLA_COMMIT = "195f45b7082930033f6533a160b0f8f7f1cbfb40"
+XLA_SHA256 = "75e77091bae789175f3de24efee9debf8835b167770490db75571bf65c27b727"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/tests/version_test.py b/tests/version_test.py
index 7ce98c858..51297a971 100644
--- a/tests/version_test.py
+++ b/tests/version_test.py
@@ -26,11 +26,11 @@ from jax._src import test_util as jtu
 
 # This is a subset of the full PEP440 pattern; for example we skip pre & 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'
-  (?:\+(?P<local>[a-zA-Z0-9_]+))?     # optional local version; like '+g6643af3c3'
-  $                                   # end of string
+  ^                                    # 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'
+  (?:\+(?P<local>[a-zA-Z0-9_.]+))?     # optional local version; like '+g6643af3c3'
+  $                                    # end of string
 """, re.VERBOSE)
 
 
@@ -61,11 +61,12 @@ def assert_no_subprocess_call():
 
 
 @contextlib.contextmanager
-def assert_subprocess_call():
+def assert_subprocess_call(stdout: bytes | None = None):
   """Run code, asserting that subprocess.Popen *is* called at least once."""
   with mock.patch("subprocess.Popen") as mock_Popen:
+    mock_Popen.return_value.communicate.return_value = (stdout, b"")
     yield
-  mock_Popen.assert_called()
+  mock_Popen.return_value.communicate.assert_called()
 
 
 class JaxVersionTest(unittest.TestCase):
@@ -126,7 +127,7 @@ class JaxVersionTest(unittest.TestCase):
       self.assertValidVersion(version)
 
     with jtu.set_env(JAX_RELEASE=None, JAXLIB_RELEASE=None,
-                      JAX_NIGHTLY=None, JAXLIB_NIGHTLY="1"):
+                     JAX_NIGHTLY=None, JAXLIB_NIGHTLY="1"):
       with assert_no_subprocess_call():
         version = jax.version._get_version_for_build()
       datestring = datetime.date.today().strftime("%Y%m%d")
@@ -134,19 +135,28 @@ class JaxVersionTest(unittest.TestCase):
       self.assertValidVersion(version)
 
     with jtu.set_env(JAX_RELEASE="1", JAXLIB_RELEASE=None,
-                      JAX_NIGHTLY=None, JAXLIB_NIGHTLY=None):
+                     JAX_NIGHTLY=None, JAXLIB_NIGHTLY=None):
       with assert_no_subprocess_call():
         version = jax.version._get_version_for_build()
       self.assertEqual(version, base_version)
       self.assertValidVersion(version)
 
     with jtu.set_env(JAX_RELEASE=None, JAXLIB_RELEASE="1",
-                      JAX_NIGHTLY=None, JAXLIB_NIGHTLY=None):
+                     JAX_NIGHTLY=None, JAXLIB_NIGHTLY=None):
       with assert_no_subprocess_call():
         version = jax.version._get_version_for_build()
       self.assertEqual(version, base_version)
       self.assertValidVersion(version)
 
+    with jtu.set_env(JAX_RELEASE=None, JAXLIB_RELEASE=None,
+                     JAX_NIGHTLY=None, JAXLIB_NIGHTLY=None,
+                     JAX_CUSTOM_VERSION_SUFFIX="test"):
+      with assert_subprocess_call(stdout=b"1731433958-1c0f1076e"):
+        version = jax.version._get_version_for_build()
+      self.assertTrue(version.startswith(f"{base_version}.dev"))
+      self.assertTrue(version.endswith("test"))
+      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")