From 02024427089134e951fe4063100b93c28e7e01a8 Mon Sep 17 00:00:00 2001 From: erman-gurses Date: Mon, 16 Feb 2026 08:46:15 -0800 Subject: [PATCH] Add Apex to the Pytorch Manifest file --- .../build_portable_linux_pytorch_wheels.yml | 3 +- .../generate_pytorch_manifest.py | 29 ++++++++++++++++++- .../tests/generate_pytorch_manifest_test.py | 15 +++++++++- 3 files changed, 44 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build_portable_linux_pytorch_wheels.yml b/.github/workflows/build_portable_linux_pytorch_wheels.yml index 5eb742c1e83..66266c71882 100644 --- a/.github/workflows/build_portable_linux_pytorch_wheels.yml +++ b/.github/workflows/build_portable_linux_pytorch_wheels.yml @@ -210,7 +210,8 @@ jobs: --pytorch-dir "external-builds/pytorch/pytorch" \ --pytorch-audio-dir "external-builds/pytorch/pytorch_audio" \ --pytorch-vision-dir "external-builds/pytorch/pytorch_vision" \ - --triton-dir "external-builds/pytorch/triton" + --triton-dir "external-builds/pytorch/triton" \ + --apex-dir "external-builds/pytorch/apex" - name: Configure AWS Credentials if: always() diff --git a/build_tools/github_actions/generate_pytorch_manifest.py b/build_tools/github_actions/generate_pytorch_manifest.py index 3f0b3a45baa..db4753e401e 100644 --- a/build_tools/github_actions/generate_pytorch_manifest.py +++ b/build_tools/github_actions/generate_pytorch_manifest.py @@ -3,7 +3,7 @@ Generate a manifest for PyTorch external builds. Writes a JSON manifest containing: - - pytorch/pytorch_audio/pytorch_vision(/triton): git commit + origin repo (+ branch best-effort) + - pytorch/pytorch_audio/pytorch_vision(/triton)(/apex): git commit + origin repo (+ branch best-effort) - therock: repo + commit + branch from GitHub Actions env (best-effort) Filename format: @@ -149,10 +149,12 @@ def build_sources( pytorch_audio_dir: Path, pytorch_vision_dir: Path, triton_dir: Path | None, + apex_dir: Path | None, pytorch_git_ref: str, pytorch_audio_git_ref: str | None, pytorch_vision_git_ref: str | None, triton_git_ref: str | None, + apex_git_ref: str | None, ) -> dict[str, dict[str, str]]: pt = git_head(pytorch_dir, label="pytorch") aud = git_head(pytorch_audio_dir, label="pytorch_audio") @@ -193,6 +195,16 @@ def build_sources( commit=tri.commit, repo=tri.repo, branch=tri_branch ).to_dict() + if apex_dir is not None: + ax = git_head(apex_dir, label="apex") + ax_branch = resolve_branch( + inferred=git_branch_best_effort(apex_dir), + provided=apex_git_ref, + ) + sources["apex"] = GitSourceInfo( + commit=ax.commit, repo=ax.repo, branch=ax_branch + ).to_dict() + return sources @@ -220,10 +232,12 @@ def generate_manifest_dict( pytorch_audio_dir: Path, pytorch_vision_dir: Path, triton_dir: Path | None, + apex_dir: Path | None, pytorch_git_ref: str, pytorch_audio_git_ref: str | None, pytorch_vision_git_ref: str | None, triton_git_ref: str | None, + apex_git_ref: str | None, ) -> dict[str, object]: """Generate the manifest dictionary""" sources = build_sources( @@ -231,10 +245,12 @@ def generate_manifest_dict( pytorch_audio_dir=pytorch_audio_dir, pytorch_vision_dir=pytorch_vision_dir, triton_dir=triton_dir, + apex_dir=apex_dir, pytorch_git_ref=pytorch_git_ref, pytorch_audio_git_ref=pytorch_audio_git_ref, pytorch_vision_git_ref=pytorch_vision_git_ref, triton_git_ref=triton_git_ref, + apex_git_ref=apex_git_ref, ) server_url = os.environ.get("GITHUB_SERVER_URL") @@ -294,6 +310,10 @@ def parse_args(argv: list[str]) -> argparse.Namespace: "--triton-git-ref", help="Optional ref for triton branch field (used if detached).", ) + ap.add_argument( + "--apex-git-ref", + help="Optional ref for apex branch field (used if detached).", + ) ap.add_argument("--pytorch-dir", type=Path, required=True) ap.add_argument("--pytorch-audio-dir", type=Path, required=True) ap.add_argument("--pytorch-vision-dir", type=Path, required=True) @@ -302,6 +322,11 @@ def parse_args(argv: list[str]) -> argparse.Namespace: type=Path, help="Optional triton checkout (Linux only).", ) + ap.add_argument( + "--apex-dir", + type=Path, + help="Optional apex checkout (Linux only).", + ) return ap.parse_args(argv) @@ -322,10 +347,12 @@ def main(argv: list[str]) -> None: pytorch_audio_dir=args.pytorch_audio_dir, pytorch_vision_dir=args.pytorch_vision_dir, triton_dir=args.triton_dir, + apex_dir=args.apex_dir, pytorch_git_ref=args.pytorch_git_ref, pytorch_audio_git_ref=args.pytorch_audio_git_ref, pytorch_vision_git_ref=args.pytorch_vision_git_ref, triton_git_ref=args.triton_git_ref, + apex_git_ref=args.apex_git_ref, ) out_path.write_text( diff --git a/build_tools/github_actions/tests/generate_pytorch_manifest_test.py b/build_tools/github_actions/tests/generate_pytorch_manifest_test.py index dc53adbb7f5..0b0c30039c4 100755 --- a/build_tools/github_actions/tests/generate_pytorch_manifest_test.py +++ b/build_tools/github_actions/tests/generate_pytorch_manifest_test.py @@ -78,11 +78,13 @@ def test_sources_only_manifest(self) -> None: audio_repo = self.tmp_path / "src_audio" vision_repo = self.tmp_path / "src_vision" triton_repo = self.tmp_path / "src_triton" + apex_repo = self.tmp_path / "src_apex" pytorch_head = "1111111111111111111111111111111111111111" audio_head = "2222222222222222222222222222222222222222" vision_head = "3333333333333333333333333333333333333333" triton_head = "4444444444444444444444444444444444444444" + apex_head = "5555555555555555555555555555555555555555" def fake_git_head(dirpath: Path, *, label: str) -> m.GitSourceInfo: p = dirpath.resolve() @@ -102,6 +104,10 @@ def fake_git_head(dirpath: Path, *, label: str) -> m.GitSourceInfo: return m.GitSourceInfo( commit=triton_head, repo="https://github.com/ROCm/triton.git" ) + if p == apex_repo.resolve(): + return m.GitSourceInfo( + commit=apex_head, repo="https://github.com/ROCm/apex.git" + ) raise AssertionError(f"Unexpected repo path: {p}") with mock.patch.object( @@ -123,6 +129,8 @@ def fake_git_head(dirpath: Path, *, label: str) -> m.GitSourceInfo: str(vision_repo), "--triton-dir", str(triton_repo), + "--apex-dir", + str(apex_repo), ] ) @@ -133,7 +141,7 @@ def fake_git_head(dirpath: Path, *, label: str) -> m.GitSourceInfo: self.assertEqual( set(data.keys()), - {"pytorch", "pytorch_audio", "pytorch_vision", "triton", "therock"}, + {"pytorch", "pytorch_audio", "pytorch_vision", "triton", "apex", "therock"}, ) self.assertEqual(data["pytorch"]["commit"], pytorch_head) @@ -156,6 +164,10 @@ def fake_git_head(dirpath: Path, *, label: str) -> m.GitSourceInfo: self.assertEqual(data["triton"]["repo"], "https://github.com/ROCm/triton.git") self.assertNotIn("branch", data["triton"]) + self.assertEqual(data["apex"]["commit"], apex_head) + self.assertEqual(data["apex"]["repo"], "https://github.com/ROCm/apex.git") + self.assertNotIn("branch", data["apex"]) + self.assertEqual(data["therock"]["repo"], "https://github.com/ROCm/TheRock.git") self.assertEqual( data["therock"]["commit"], "b3eda956a19d0151cbb4699739eb71f62596c8bb" @@ -219,6 +231,7 @@ def fake_git_head(dirpath: Path, *, label: str) -> m.GitSourceInfo: self.assertIn("pytorch_audio", data) self.assertIn("pytorch_vision", data) self.assertNotIn("triton", data) + self.assertNotIn("apex", data) self.assertEqual(data["pytorch"]["branch"], "nightly") self.assertNotIn("branch", data["pytorch_audio"])