Skip to content

Add Apex to the Pytorch Manifest file #3431

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
erman-gurses merged 1 commit into from
Feb 20, 2026
Merged

Add Apex to the Pytorch Manifest file #3431

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .github/workflows/build_portable_linux_pytorch_wheels.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
29 changes: 28 additions & 1 deletion build_tools/github_actions/generate_pytorch_manifest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -220,21 +232,25 @@ 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(
pytorch_dir=pytorch_dir,
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")
Expand Down Expand Up @@ -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)
Expand All @@ -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)


Expand All @@ -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(
Expand Down
15 changes: 14 additions & 1 deletion build_tools/github_actions/tests/generate_pytorch_manifest_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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(
Expand All @@ -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),
]
)

Expand All @@ -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)
Expand All @@ -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"
Expand Down Expand Up @@ -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"])
Expand Down
Loading