Add pip-index.py for PEP 503 simple-index style mirrors - #205
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds a standalone async mirroring script for PEP 503 / PEP 691 “simple” Python package indexes, including href rewriting and optional devpi channel expansion.
Changes:
- Introduces
pip-index.pycrawler/downloader built onaiohttpwith concurrency control. - Adds configurable href rewriting (
REWRITE_HOSTS,EXTRA_REWRITES) to point saved indexes back to local mirrors. - Adds optional devpi JSON-based discovery mode (
DEVPI_MODE) and optional PyTorch releases URL discovery.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| path = unquote(urlparse(url).path) | ||
| while path.startswith("/"): | ||
| path = path[1:] |
There was a problem hiding this comment.
Added safe_local_path() that joins the unquoted URL path under base and verifies the resolved path stays inside base (rejects '..', absolute path components, NUL). recursive_download now uses it for both index pages and downloaded files.
| if url.endswith("/") or url.endswith(".html"): | ||
| # index.html (current) or torch_stable.html (old) | ||
| async with sem: | ||
| logging.info(f"Getting {url}") | ||
| contents = await get_with_progress(client, url) | ||
| index_resp = contents.decode("utf-8") | ||
| if url.endswith("/"): | ||
| filename = "index.html" | ||
| else: | ||
| filename = url.split("/")[-1] | ||
| assert filename.endswith(".html"), f"Unexpected HTML file: {filename}" |
There was a problem hiding this comment.
Fixed. When the URL ends in .html the local path now drops the .html filename to obtain the parent directory, so we write index_resp to / instead of mkdir-ing on top of .
| if not dry_run: | ||
| index_resp = rewrite_index(index_resp) | ||
| os.makedirs(base / path, exist_ok=True) | ||
| with overwrite(base / path / filename, "w") as f: | ||
| f.write(index_resp) |
There was a problem hiding this comment.
Same as the line-281 reply.
| pass | ||
|
|
||
|
|
||
| async def get_with_progress(client: aiohttp.ClientSession, url: str) -> bytes: |
There was a problem hiding this comment.
Updated: I rolled the asyncio.to_thread(fh.write, chunk) wrap back. Looking at the rest of the repo, every other downloader (apt-sync.py, adoptium.py, github-release.py, homebrew-bottles.py, docker-ce.py, yum-sync.py, github-raw.py, pub-mirror.py, nixos-images.py, nix-channels.py) just calls f.write(chunk) synchronously. The aiohttp-side chunks are already small (iter_chunked(65536)) and disk writes here are well below the network rate, so handing each write off to a thread adds overhead without solving a real problem. Concurrency comes from the per-file JOBS semaphore on the network side. Kept stream_to_file for the .tmp -> rename atomicity. Reverted in de8d0d5.
| chunks = [] | ||
| try: | ||
| async for chunk in resp.content.iter_chunked(65536): | ||
| downloaded += len(chunk) | ||
| chunks.append(chunk) |
There was a problem hiding this comment.
Updated: I rolled the asyncio.to_thread(fh.write, chunk) wrap back. Looking at the rest of the repo, every other downloader (apt-sync.py, adoptium.py, github-release.py, homebrew-bottles.py, docker-ce.py, yum-sync.py, github-raw.py, pub-mirror.py, nixos-images.py, nix-channels.py) just calls f.write(chunk) synchronously. The aiohttp-side chunks are already small (iter_chunked(65536)) and disk writes here are well below the network rate, so handing each write off to a thread adds overhead without solving a real problem. Concurrency comes from the per-file JOBS semaphore on the network side. Kept stream_to_file for the .tmp -> rename atomicity. Reverted in de8d0d5.
| suburl = urljoin(upstream_base, suburl) | ||
| else: | ||
| suburl = urljoin(url, suburl) | ||
| tasks.append(asyncio.create_task(recursive_download(client, suburl))) |
There was a problem hiding this comment.
Same as the line-292 reply.
| if tasks: | ||
| await asyncio.gather(*tasks) |
There was a problem hiding this comment.
Same as the line-292 reply.
| if (base / path).exists(): | ||
| return |
There was a problem hiding this comment.
Same as the line-292 reply.
| return b"".join(chunks) | ||
| except Exception as e: | ||
| if attempt == 2: | ||
| raise e |
There was a problem hiding this comment.
Switched to bare 'raise' inside the except blocks (both in get_with_progress and recursive_download) so the original traceback is preserved.
| if e.status == 403: | ||
| logging.warning(f"Forbidden: {url}, skipping.") | ||
| else: | ||
| raise e |
There was a problem hiding this comment.
Same as the line-209 reply: bare 'raise' now.
Adapted from ustclug/ustcmirror-images pytorch/sync.py (originally
pytorch.py). Crawls PEP 503 / PEP 691 simple HTML indexes recursively
and rewrites href attributes so saved index pages point back to this
mirror. The same script is used for multiple jobs by parameterising
endpoint discovery and href rewriting via environment variables.
Additions on top of the upstream script:
- Multi-host href rewrite (REWRITE_HOSTS):
handles upstreams that emit absolute URLs across multiple origins
(e.g. download.pytorch.org + download-r2.pytorch.org).
- EXTRA_REWRITES:
redirect hrefs that point at sibling hosts (e.g. files.pythonhosted.org)
to a sibling local mirror prefix (e.g. /pypi/web), so transitive deps
stay on the same mirror without forcing users to set --extra-index-url.
- DEVPI_MODE:
query a devpi server's channel JSON API and crawl only that channel's
own projects via .../<channel>/+simple/<project>/ instead of walking
the full inherited PyPI namespace.
- aiohttp instead of httpx, so it can run inside the existing shared
tunathu/tunasync-scripts:latest image with no extra dependencies.
c472eb1 to
39aeb9a
Compare
Address review feedback: - Add safe_local_path() that resolves URL paths under base and rejects any candidate that escapes base (rejects '..', absolute paths, NUL). - Stop treating '.../torch_stable.html' as a directory: when the URL ends with .html, write into <parent>/<filename>, not <full path>/. - Add stream_to_file(): wheels are now written chunk-by-chunk through asyncio.to_thread(fh.write, chunk) into a sibling .tmp file, which is unlinked on failure and atomically renamed on success. The previous implementation buffered the whole response in memory before writing. - Maintain a process-wide visited URL set so cross-linked index pages do not re-enter recursive_download() for the same URL. - Replace 'raise e' in get_with_progress and recursive_download with a bare 'raise' so the original traceback is preserved.
Per Copilot review reflection: matching the rest of the repo (apt-sync.py, adoptium.py, github-release.py, etc.) which all do plain f.write(chunk) inside the download loop. With default JOBS=1 there is no concurrent coroutine to starve, and asyncio.to_thread per chunk adds ~50us thread scheduling overhead that outweighs any event-loop unblocking benefit.
…RITES hosts BREAKING: REWRITE_HOSTS is removed. Every absolute http(s) href is now rewritten to URLBASE by default. The crawler already downloads every linked file regardless of host, so the rewrite must cover every host too - otherwise pip silently bypasses the mirror whenever the upstream starts linking from a new host (seen with files.pythonhosted.org, pypi.nvidia.com, repo.amd.com). EXTRA_REWRITES rules are matched first and keep their own prefix; protocol-relative hrefs are handled too, and query/fragment (e.g. #sha256=...) are preserved. Files on EXTRA_REWRITES hosts are no longer downloaded; they are only rewritten, since they are expected to be served by the sibling local mirror at the configured prefix (equivalent to ustclug PR tuna#174's PYPI_URLBASE). This avoids storing duplicates of files already mirrored elsewhere (e.g. 115GB of pytorch wheels duplicated under pypi/web). EXTRA_REWRITES host matching is now case-insensitive.
Some index pages are blocked by upstream too (e.g. download.pytorch.org/whl/rocm7.14/cuda-bindings/ returns 403). Skip the whole subtree instead of aborting the run, and record it in forbidden_dirs so CLEANUP never treats its contents as stale. CLEANUP=1 deletes local files that no crawled URL maps to (stale/yanked content, leftover .tmp), then prunes emptied directories. Safety rails: refuses to delete anything when candidates exceed CLEANUP_MAX_DELETE (default 1000) so a partial crawl cannot wipe the tree; 403-skipped subtrees and (with NO_NIGHTLY=1) nightly paths are never candidates.
Summary
Add
pip-index.py, a generic mirror script for PEP 503 / PEP 691 "simple" index style upstreams.Adapted from
ustclug/ustcmirror-imagespytorch/sync.py(originallypytorch.py). Crawls a simple HTML index recursively and rewrites href attributes so saved index pages point back to this mirror. The same script can drive multiple tunasync jobs by parameterising endpoint discovery and href rewriting via environment variables.Additions over the upstream script
EXTRA_REWRITESfiles.pythonhosted.org) to a sibling local mirror prefix (e.g./pypi/web), so transitive deps stay on the same mirror without forcing users to set--extra-index-url. Files on these hosts are not downloaded — they are only rewritten, since the sibling mirror already serves them (equivalent to ustclug/ustcmirror-images#174'sPYPI_URLBASE).DEVPI_MODE.../<channel>/+simple/<project>/instead of walking the full inherited PyPI namespace.CLEANUP,CLEANUP_MAX_DELETE.tmp), then prune emptied directories. Refuses to delete anything when candidates exceedCLEANUP_MAX_DELETE(default 1000); 403-skipped subtrees and (withNO_NIGHTLY=1) nightly paths are never candidates.USE_PYTORCH_RELEASES,GET_ALL,NO_NIGHTLYURLBASE,JOBS,TIMEOUT,DRY_RUN/{TUNASYNC_MIRROR_NAME}/; JOBS and TIMEOUT have script-level defaults so most jobs need no env.aiohttpis used instead ofhttpx, so the script runs inside the existing sharedtunathu/tunasync-scripts:latestimage with no extra dependencies.The Issue #109-style fix (rewrite absolute
download-r2.pytorch.orgURLs) was originally proposed as ustclug/ustcmirror-images#154.REWRITE_HOSTSremovedEarlier revisions of this PR rewrote only hrefs whose host was listed in
REWRITE_HOSTS(defaulting to theTUNASYNC_UPSTREAM_URLhost). That allowlist repeatedly failed in production: every time the upstream started linking files from a new host, pip silently bypassed the mirror (seen withfiles.pythonhosted.org, thenpypi.nvidia.com, thenrepo.amd.com).Since the crawler already downloads every linked file regardless of host, the rewrite now covers every host too: all absolute and protocol-relative
http(s)hrefs are rewritten toURLBASEby default, withEXTRA_REWRITESrules matched first. Query strings and#sha256=...fragments are preserved. If you adopted an earlier revision of this script, dropREWRITE_HOSTSfrom your config.Robustness
download.pytorch.org/whl/rocm7.14/cuda-bindings/) are skipped with a warning instead of aborting the whole run, and are excluded from CLEANUP.Production validation (NJU mirror, tuning this script since 2026-05)
Drives two tunasync jobs in production:
pytorch(upstream download.pytorch.org, ~3.9TB) andjetson-pypi(devpi-backed, ~46GB). After the rewrite-all change: zero remaining external hrefs across all saved index pages; the first CLEANUP runs deleted 675 (pytorch) and 1931 (jetson-pypi) stale files, all verified gone upstream.