diff --git a/.github/workflows/test_pytorch_wheels.yml b/.github/workflows/test_pytorch_wheels.yml index 9eddb93d835..d5c64ab729d 100644 --- a/.github/workflows/test_pytorch_wheels.yml +++ b/.github/workflows/test_pytorch_wheels.yml @@ -214,6 +214,10 @@ jobs: - name: Install test requirements run: | python -m pip install -r external-builds/pytorch/requirements-test.txt + # numpy is pinned to what upstream validates for the checked-out torch + # ref, derived from its .ci/docker/requirements-ci.txt (avoids drift). + python external-builds/pytorch/derive_test_requirements.py \ + | python -m pip install -r - pip freeze - name: Run rocm-sdk sanity tests diff --git a/external-builds/pytorch/derive_test_requirements.py b/external-builds/pytorch/derive_test_requirements.py new file mode 100644 index 00000000000..92ce64deac4 --- /dev/null +++ b/external-builds/pytorch/derive_test_requirements.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +# Copyright Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Print pinned test requirements sourced from the vendored PyTorch checkout. + +Extracts the pins for the requested packages (default: ``numpy``) from the +checked-out PyTorch's ``.ci/docker/requirements-ci.txt`` and prints them, so +TheRock installs the exact versions upstream validates for the torch ref under +test instead of a hand-copied pin that silently drifts. Environment markers +(e.g. ``; python_version ...``) are preserved so pip selects the right line. + +Usage: + python derive_test_requirements.py + python derive_test_requirements.py --package numpy --package scipy + python derive_test_requirements.py | python -m pip install -r - +""" + +import argparse +import re +import sys +from pathlib import Path + +THIS_DIR = Path(__file__).resolve().parent +DEFAULT_REQUIREMENTS_CI = ( + THIS_DIR / "pytorch" / ".ci" / "docker" / "requirements-ci.txt" +) + + +def extract_pins(requirements_ci: Path, packages: list[str]) -> list[str]: + """Return the requirement lines for each package, in the given order.""" + lines = requirements_ci.read_text().splitlines() + pins: list[str] = [] + for package in packages: + pattern = re.compile(rf"^\s*{re.escape(package)}\s*(==|>=|<=|<|>|~=|;|$)") + matches = [line.rstrip() for line in lines if pattern.match(line)] + if not matches: + raise ValueError(f"no pin for '{package}' found in {requirements_ci}") + pins.extend(matches) + return pins + + +def main(argv: list[str]) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--package", + action="append", + dest="packages", + help="Package to extract (repeatable). Defaults to numpy.", + ) + parser.add_argument( + "--requirements-ci", + type=Path, + default=DEFAULT_REQUIREMENTS_CI, + help=f"PyTorch requirements-ci.txt to read. Default: {DEFAULT_REQUIREMENTS_CI}", + ) + args = parser.parse_args(argv) + + if not args.requirements_ci.is_file(): + parser.error(f"requirements-ci.txt not found: {args.requirements_ci}") + + try: + pins = extract_pins(args.requirements_ci, args.packages or ["numpy"]) + except ValueError as e: + parser.error(str(e)) + + print("\n".join(pins)) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/external-builds/pytorch/requirements-test.txt b/external-builds/pytorch/requirements-test.txt index afaa73b3e48..82dd8c2dc00 100644 --- a/external-builds/pytorch/requirements-test.txt +++ b/external-builds/pytorch/requirements-test.txt @@ -6,5 +6,7 @@ pytest-xdist>=3.5.0 expecttest>=0.3.0 hypothesis ninja -numpy +# numpy is intentionally omitted here: it is derived per torch-ref from the +# vendored PyTorch checkout at test time (see derive_test_requirements.py), +# mirroring upstream's .ci/docker/requirements-ci.txt without manual drift. psutil