Skip to content

Commit 8ea1956

Browse files
committed
Add automated PyPI index with separate CUDA/ROCm backends
- Add script to auto-generate PyPI-compatible index from GitHub releases - Create separate indexes for CUDA and ROCm wheels based on release tags - Add GitHub Actions workflow to deploy indexes to GitHub Pages - Generate landing page with installation instructions - Ignore auto-generated pypi-index directory Indexes will be available at: - ROCm: https://embeddedllm.github.io/fastsafetensors-rocm/rocm/simple/ - CUDA: https://embeddedllm.github.io/fastsafetensors-rocm/cuda/simple/ Signed-off-by: tjtanaa <[email protected]>
1 parent c543e16 commit 8ea1956

File tree

3 files changed

+285
-1
lines changed

3 files changed

+285
-1
lines changed
Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,228 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Generate PyPI-compatible simple index from GitHub releases.
4+
This script fetches all releases and creates separate indexes for CUDA and ROCm wheels.
5+
"""
6+
7+
import json
8+
import os
9+
import sys
10+
from pathlib import Path
11+
from urllib.request import urlopen, Request
12+
13+
def fetch_releases(repo_owner, repo_name):
14+
"""Fetch all releases from GitHub API."""
15+
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/releases"
16+
headers = {
17+
"Accept": "application/vnd.github.v3+json",
18+
"User-Agent": "PyPI-Index-Generator"
19+
}
20+
21+
# Add GitHub token if available (for higher rate limits)
22+
token = os.environ.get("GITHUB_TOKEN")
23+
if token:
24+
headers["Authorization"] = f"token {token}"
25+
26+
request = Request(url, headers=headers)
27+
28+
try:
29+
with urlopen(request) as response:
30+
return json.loads(response.read().decode())
31+
except Exception as e:
32+
print(f"Error fetching releases: {e}", file=sys.stderr)
33+
sys.exit(1)
34+
35+
def categorize_backend(release_tag):
36+
"""Determine backend (cuda/rocm) from release tag."""
37+
tag_lower = release_tag.lower()
38+
39+
if "rocm" in tag_lower:
40+
return "rocm"
41+
elif "cuda" in tag_lower:
42+
return "cuda"
43+
else:
44+
# Default to cuda for untagged releases
45+
return "cuda"
46+
47+
def extract_wheels_by_backend(releases):
48+
"""Extract wheel files from releases, categorized by backend."""
49+
wheels_by_backend = {
50+
"cuda": [],
51+
"rocm": []
52+
}
53+
54+
for release in releases:
55+
backend = categorize_backend(release.get("tag_name", ""))
56+
57+
for asset in release.get("assets", []):
58+
name = asset.get("name", "")
59+
if name.endswith(".whl"):
60+
wheels_by_backend[backend].append({
61+
"name": name,
62+
"url": asset.get("browser_download_url"),
63+
"version": release.get("tag_name"),
64+
})
65+
66+
return wheels_by_backend
67+
68+
def generate_root_index(output_dir, packages):
69+
"""Generate the root simple index."""
70+
html = """<!DOCTYPE html>
71+
<html>
72+
<head>
73+
<meta charset="UTF-8">
74+
<title>Simple Index</title>
75+
</head>
76+
<body>
77+
<h1>Simple Index</h1>
78+
"""
79+
80+
for package in sorted(packages):
81+
html += f' <a href="{package}/">{package}</a><br/>\n'
82+
83+
html += """</body>
84+
</html>
85+
"""
86+
87+
output_path = output_dir / "index.html"
88+
output_path.write_text(html)
89+
print(f"Generated: {output_path}")
90+
91+
def generate_package_index(output_dir, package_name, wheels):
92+
"""Generate package-specific index with all wheels."""
93+
html = f"""<!DOCTYPE html>
94+
<html>
95+
<head>
96+
<meta charset="UTF-8">
97+
<title>Links for {package_name}</title>
98+
</head>
99+
<body>
100+
<h1>Links for {package_name}</h1>
101+
"""
102+
103+
# Sort wheels by version and Python version
104+
sorted_wheels = sorted(wheels, key=lambda w: (w["name"], w["version"]), reverse=True)
105+
106+
for wheel in sorted_wheels:
107+
# Extract package name from wheel filename to ensure consistency
108+
wheel_name = wheel["name"]
109+
url = wheel["url"]
110+
html += f' <a href="{url}#sha256=">{wheel_name}</a><br/>\n'
111+
112+
html += """</body>
113+
</html>
114+
"""
115+
116+
package_dir = output_dir / package_name
117+
package_dir.mkdir(parents=True, exist_ok=True)
118+
119+
output_path = package_dir / "index.html"
120+
output_path.write_text(html)
121+
print(f"Generated: {output_path}")
122+
123+
def generate_landing_page(base_dir, repo_name):
124+
"""Generate a landing page for the PyPI index."""
125+
html = f"""<!DOCTYPE html>
126+
<html>
127+
<head>
128+
<meta charset="UTF-8">
129+
<title>{repo_name} - PyPI Index</title>
130+
<style>
131+
body {{ font-family: system-ui, -apple-system, sans-serif; max-width: 800px; margin: 50px auto; padding: 20px; }}
132+
h1 {{ color: #2c3e50; }}
133+
.backend {{ margin: 30px 0; padding: 20px; background: #f8f9fa; border-radius: 8px; }}
134+
code {{ background: #e9ecef; padding: 2px 6px; border-radius: 4px; font-size: 0.9em; }}
135+
pre {{ background: #2c3e50; color: #ecf0f1; padding: 15px; border-radius: 6px; overflow-x: auto; }}
136+
</style>
137+
</head>
138+
<body>
139+
<h1>{repo_name} - PyPI Index</h1>
140+
<p>Choose the appropriate index URL based on your GPU backend:</p>
141+
142+
<div class="backend">
143+
<h2>🔥 ROCm (AMD GPUs)</h2>
144+
<p>For AMD GPUs using ROCm:</p>
145+
<pre>pip install fastsafetensors --index-url https://embeddedllm.github.io/{repo_name}/rocm/simple/</pre>
146+
</div>
147+
148+
<div class="backend">
149+
<h2>💚 CUDA (NVIDIA GPUs)</h2>
150+
<p>For NVIDIA GPUs using CUDA:</p>
151+
<pre>pip install fastsafetensors --index-url https://embeddedllm.github.io/{repo_name}/cuda/simple/</pre>
152+
</div>
153+
154+
<h3>Version Specific Installation</h3>
155+
<pre>pip install fastsafetensors==0.1.15 --index-url https://embeddedllm.github.io/{repo_name}/rocm/simple/</pre>
156+
157+
<h3>In requirements.txt</h3>
158+
<pre>--index-url https://embeddedllm.github.io/{repo_name}/rocm/simple/
159+
fastsafetensors>=0.1.15</pre>
160+
161+
<hr>
162+
<p><small>Direct access: <a href="rocm/simple/">ROCm Index</a> | <a href="cuda/simple/">CUDA Index</a></small></p>
163+
</body>
164+
</html>
165+
"""
166+
167+
output_path = base_dir / "index.html"
168+
output_path.write_text(html)
169+
print(f"Generated landing page: {output_path}")
170+
171+
def main():
172+
# Configuration
173+
repo_owner = os.environ.get("GITHUB_REPOSITORY_OWNER", "EmbeddedLLM")
174+
repo_full = os.environ.get("GITHUB_REPOSITORY", "EmbeddedLLM/fastsafetensors-rocm")
175+
repo_name = repo_full.split("/")[-1]
176+
177+
print(f"Fetching releases from {repo_owner}/{repo_name}...")
178+
releases = fetch_releases(repo_owner, repo_name)
179+
print(f"Found {len(releases)} releases")
180+
181+
# Extract wheels categorized by backend
182+
wheels_by_backend = extract_wheels_by_backend(releases)
183+
184+
total_wheels = sum(len(wheels) for wheels in wheels_by_backend.values())
185+
print(f"Found {total_wheels} total wheel files")
186+
print(f" CUDA: {len(wheels_by_backend['cuda'])} wheels")
187+
print(f" ROCm: {len(wheels_by_backend['rocm'])} wheels")
188+
189+
if total_wheels == 0:
190+
print("Warning: No wheel files found in any release", file=sys.stderr)
191+
return
192+
193+
# Generate indexes for each backend
194+
for backend, wheels in wheels_by_backend.items():
195+
if not wheels:
196+
print(f"Skipping {backend} index (no wheels found)")
197+
continue
198+
199+
print(f"\nGenerating {backend.upper()} index...")
200+
output_dir = Path(f"pypi-index/{backend}/simple")
201+
output_dir.mkdir(parents=True, exist_ok=True)
202+
203+
# Group wheels by package name
204+
packages = {}
205+
for wheel in wheels:
206+
# Extract package name from wheel filename (before first dash)
207+
package_name = wheel["name"].split("-")[0]
208+
if package_name not in packages:
209+
packages[package_name] = []
210+
packages[package_name].append(wheel)
211+
212+
# Generate indexes
213+
generate_root_index(output_dir, packages.keys())
214+
215+
for package_name, package_wheels in packages.items():
216+
generate_package_index(output_dir, package_name, package_wheels)
217+
218+
print(f" Generated {backend.upper()} index with {len(packages)} package(s)")
219+
220+
# Generate landing page
221+
base_dir = Path("pypi-index")
222+
generate_landing_page(base_dir, repo_name)
223+
224+
print(f"\n✓ Successfully generated indexes for all backends")
225+
print(f" Total wheels: {total_wheels}")
226+
227+
if __name__ == "__main__":
228+
main()
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
name: Deploy PyPI Index to GitHub Pages
2+
3+
on:
4+
release:
5+
types: [published]
6+
workflow_dispatch:
7+
8+
permissions:
9+
contents: read
10+
pages: write
11+
id-token: write
12+
13+
concurrency:
14+
group: "pages"
15+
cancel-in-progress: false
16+
17+
jobs:
18+
build:
19+
runs-on: ubuntu-latest
20+
steps:
21+
- name: Checkout
22+
uses: actions/checkout@v4
23+
24+
- name: Set up Python
25+
uses: actions/setup-python@v5
26+
with:
27+
python-version: '3.11'
28+
29+
- name: Generate PyPI index from releases
30+
env:
31+
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
32+
run: |
33+
chmod +x .github/scripts/generate_pypi_index.py
34+
python .github/scripts/generate_pypi_index.py
35+
36+
- name: Setup Pages
37+
uses: actions/configure-pages@v4
38+
39+
- name: Upload artifact
40+
uses: actions/upload-pages-artifact@v3
41+
with:
42+
path: ./pypi-index
43+
44+
deploy:
45+
environment:
46+
name: github-pages
47+
url: ${{ steps.deployment.outputs.page_url }}
48+
runs-on: ubuntu-latest
49+
needs: build
50+
steps:
51+
- name: Deploy to GitHub Pages
52+
id: deployment
53+
uses: actions/deploy-pages@v4

.gitignore

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,4 +15,7 @@ examples/paddle_case/log
1515
# Auto-generated hipified files and directories (created during ROCm build)
1616
fastsafetensors/cpp/hip/
1717
fastsafetensors/cpp/*.hip.*
18-
fastsafetensors/cpp/hip_compat.h
18+
fastsafetensors/cpp/hip_compat.h
19+
20+
# Auto-generated PyPI index (generated by GitHub Actions)
21+
pypi-index/

0 commit comments

Comments
 (0)