Skip to content

Commit eeca23c

Browse files
committed
Add canonicalization of python versions
The ci_build parsing will now drop the revision portion of the requested python version targets to build for, since users were passing in old buggy versions. Using only major.minor will result in the latest of the minor series being used, and won't affect any build outputs like bytecode or ABI bindings. Those are tied to major.minor versions of CPython. (cherry picked from commit 5bf0875)
1 parent 37fccbb commit eeca23c

File tree

2 files changed

+48
-6
lines changed

2 files changed

+48
-6
lines changed

build/rocm/ci_build

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,9 @@ import sys
2727
from typing import List
2828

2929

30-
DEFAULT_GPU_DEVICE_TARGETS = "gfx906,gfx908,gfx90a,gfx942,gfx1030,gfx1100,gfx1101,gfx1200,gfx1201"
30+
DEFAULT_GPU_DEVICE_TARGETS = (
31+
"gfx906,gfx908,gfx90a,gfx942,gfx1030,gfx1100,gfx1101,gfx1200,gfx1201"
32+
)
3133

3234

3335
def image_by_name(name):
@@ -44,7 +46,7 @@ def dist_wheels(
4446
rocm_build_job="",
4547
rocm_build_num="",
4648
compiler="gcc",
47-
gpu_device_targets : List[str] = None,
49+
gpu_device_targets: List[str] = None,
4850
):
4951
if not gpu_device_targets:
5052
gpu_device_targets = DEFAULT_GPU_DEVICE_TARGETS.split(",")
@@ -170,7 +172,7 @@ def dist_docker(
170172
tag="rocm/jax-dev",
171173
dockerfile=None,
172174
keep_image=True,
173-
gpu_device_targets : List[str] = None,
175+
gpu_device_targets: List[str] = None,
174176
):
175177
if not dockerfile:
176178
dockerfile = "build/rocm/Dockerfile.ms"
@@ -255,6 +257,24 @@ def test(image_name):
255257
subprocess.check_call(cmd)
256258

257259

260+
def canonicalize_python_versions(versions: List[str]):
261+
if isinstance(versions, str):
262+
raise ValueError("'versions' must be a list of strings: versions=%r" % versions)
263+
264+
cleaned = []
265+
for v in versions:
266+
tup = v.split(".")
267+
major = tup[0]
268+
minor = tup[1]
269+
rev = None
270+
if len(tup) > 2 and tup[2]:
271+
rev = tup[2]
272+
273+
cleaned.append("%s.%s" % (major, minor))
274+
275+
return cleaned
276+
277+
258278
def parse_gpu_targets(targets_string):
259279
# catch case where targets_string was empty.
260280
# None should already be caught by argparse, but
@@ -353,11 +373,12 @@ def parse_args():
353373
def main():
354374
args = parse_args()
355375
gpu_device_targets = parse_gpu_targets(args.gpu_device_targets)
376+
python_versions = canonicalize_python_versions(args.python_versions)
356377

357378
if args.action == "dist_wheels":
358379
dist_wheels(
359380
args.rocm_version,
360-
args.python_versions,
381+
python_versions,
361382
args.xla_source_dir,
362383
args.rocm_build_job,
363384
args.rocm_build_num,
@@ -371,7 +392,7 @@ def main():
371392
elif args.action == "dist_docker":
372393
dist_wheels(
373394
args.rocm_version,
374-
args.python_versions,
395+
python_versions,
375396
args.xla_source_dir,
376397
args.rocm_build_job,
377398
args.rocm_build_num,
@@ -381,7 +402,7 @@ def main():
381402
dist_docker(
382403
args.rocm_version,
383404
args.base_docker,
384-
args.python_versions,
405+
python_versions,
385406
args.xla_source_dir,
386407
rocm_build_job=args.rocm_build_job,
387408
rocm_build_num=args.rocm_build_num,

build/rocm/test_ci_build.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,27 @@ def test_parse_gpu_targets_invalid_arch(self):
5555
targets = ["gfx908", "gfx940", "--oops", "/jax"]
5656
self.assertRaises(ValueError, ci_build.parse_gpu_targets, " ".join(targets))
5757

58+
def test_canonicalize_python_versions(self):
59+
versions = ["3.10.0", "3.11.0", "3.12.0"]
60+
exp = ["3.10", "3.11", "3.12"]
61+
res = ci_build.canonicalize_python_versions(versions)
62+
self.assertEqual(res, exp)
63+
64+
def test_canonicalize_python_versions_scalar(self):
65+
versions = ["3.10.0"]
66+
exp = ["3.10"]
67+
res = ci_build.canonicalize_python_versions(versions)
68+
self.assertEqual(res, exp)
69+
70+
def test_canonicalize_python_versions_no_revision_part(self):
71+
versions = ["3.10", "3.11"]
72+
res = ci_build.canonicalize_python_versions(versions)
73+
self.assertEqual(res, versions)
74+
75+
def test_canonicalize_python_versions_string(self):
76+
versions = "3.10.0"
77+
self.assertRaises(ValueError, ci_build.canonicalize_python_versions, versions)
78+
5879

5980
if __name__ == "__main__":
6081
unittest.main()

0 commit comments

Comments
 (0)