Metadata-Version: 2.4
Name: aimirror
Version: 0.2.4
Summary: AI时代的下载镜像加速器 - 支持Docker/PyPI/HuggingFace的多线程加速下载代理
Author-email: huanglin <huanglingg@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/livehl/aimirror
Project-URL: Repository, https://github.com/livehl/aimirror
Project-URL: Issues, https://github.com/livehl/aimirror/issues
Keywords: proxy,download,accelerator,docker,pypi,huggingface,cache
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: System Administrators
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Internet :: Proxy Servers
Classifier: Topic :: System :: Networking
Classifier: Topic :: Utilities
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: fastapi>=0.100.0
Requires-Dist: uvicorn[standard]>=0.23.0
Requires-Dist: httpx[http2]>=0.25.0
Requires-Dist: aiohttp>=3.9.0
Requires-Dist: aiofiles>=23.0.0
Requires-Dist: pyyaml>=6.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Dynamic: license-file

# 🚀 aimirror

[![Python](https://img.shields.io/badge/Python-3.11%2B-blue)](https://www.python.org/)
[![FastAPI](https://img.shields.io/badge/FastAPI-0.100%2B-009688)](https://fastapi.tiangolo.com/)
[![License](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)
[![PyPI](https://img.shields.io/badge/PyPI-aimirror-blue)](https://pypi.org/project/aimirror/)

> AI 时代的下载镜像加速器 —— 被慢速网络逼疯的工程师的自救工具
> 
> **一个服务 = PyPI + Docker Hub + CRAN + HuggingFace 全加速，还能任意扩展更多源**

## 💡 项目背景

作为一名 AI 工程师，每天的工作离不开：
- `pip install torch` —— 几百 MB 的 wheel 包下载到地老天荒
- `docker pull nvidia/cuda` —— 几个 GB 的镜像层反复下载
- `huggingface-cli download` —— 模型文件从 HuggingFace 蜗牛般爬过来

公司内网有代理，但单线程下载大文件依然慢得让人崩溃。重复下载相同的包？不存在的缓存。忍无可忍，于是写了这个工具。

**aimirror** = 智能路由 + 并行分片下载 + 本地缓存，让下载速度飞起来。

## ✨ 功能特性

- **⚡ 并行下载** —— HTTP Range 分片，多线程并发，榨干带宽
- **💾 智能缓存** —— 基于文件 digest 去重，LRU 自动淘汰
- **🎯 动态路由** —— 小文件直接代理，大文件自动并行
- **🔗 多源支持** —— Docker Hub、PyPI、CRAN、HuggingFace 开箱即用
- **🔌 任意扩展** —— 只要是 HTTP 下载，配置一条规则即可几十倍加速

## 🔥 性能实测

### PyPI 包安装加速对比

使用 `uv pip install` 安装 148 个依赖包（含 torch、transformers 等大包）：

| 模式 | 解析依赖 | 准备包 | 总耗时 | 加速比 |
|------|---------|--------|--------|--------|
| ❌ 仅代理 (900KB/s) | 17m 04s | 14m 20s | **~31 分钟** | 1x |
| ✅ aimirror (170MB/s) | 34.78s | 45.81s | **~80 秒** | **23x** |

> 💡 **实测环境**: 公司内网代理，带宽瓶颈明显。使用 aimirror 后从 900KB/s 飙升至 170MB/s，**提速近 200 倍**！

### 缓存效果

| 场景 | 耗时 | 说明 |
|------|------|------|
| 首次下载 | 80s | 并行下载 + 写入缓存 |
| 缓存命中 | <1s | 本地直接返回，秒开 |

### 多源同时加速

**一个 aimirror 服务，同时加速多种包管理器：**

| 包管理器 | 配置方式 | 加速效果 |
|---------|---------|---------|
| **pip/uv** | `export HTTPS_PROXY=http://localhost:8081` 或 `pip install -i http://localhost:8081/simple` | PyPI 包 23x 加速 |
| **docker** | `/etc/docker/daemon.json` 中设置 `registry-mirrors` | 镜像拉取并行分片 |
| **R CRAN** | `options(repos = c(CRAN = "http://localhost:8081"))` | R 包下载加速 |
| **huggingface-cli** | `export HF_ENDPOINT=http://localhost:8081` | 模型文件秒下 |
| **conda** | `.condarc` 中配置 `channels` | 同理可扩展 |
| **npm/maven** | 配置 registry 指向 aimirror | 任意 HTTP 源均可 |

> 🔌 **扩展能力**: 只要是 HTTP 下载，在 `config.yaml` 中添加一条规则即可接入加速，无需启动多个服务。

## 🏗️ 架构

```mermaid
flowchart LR
    subgraph Client["客户端"]
        PIP[pip install]
        DOCKER[docker pull]
        HF[huggingface-cli]
    end

    subgraph aimirror["aimirror 服务"]
        ROUTER[路由匹配器<br/>router.py]
        PROXY[直接代理<br/>小文件]
        DOWNLOADER[并行下载器<br/>downloader.py]
        CACHE[缓存管理器<br/>cache.py]
    end

    subgraph UpstreamProxy["上游代理 (可选)"]
        COMPANY_PROXY[公司代理<br/>http://proxy.company.com:8080]
    end

    subgraph Upstream["上游服务"]
        PYPI[PyPI]
        DOCKER_HUB[Docker Hub]
        HF_HUB[HuggingFace]
    end

    PIP --> ROUTER
    DOCKER --> ROUTER
    HF --> ROUTER

    ROUTER -->|小文件| PROXY
    ROUTER -->|大文件| DOWNLOADER

    PROXY --> UpstreamProxy
    DOWNLOADER --> CACHE
    CACHE -->|未命中| UpstreamProxy
    CACHE -->|命中| Client

    UpstreamProxy -->|可选| Upstream
    UpstreamProxy -.->|直连| Upstream

    DOCKER_HUB -->|返回文件| CACHE
    PYPI -->|返回文件| CACHE
    HF_HUB -->|返回文件| CACHE
```

## 🚀 快速开始

### 方式一：pip 安装（推荐）

```bash
# 安装
pip install aimirror

# 启动
aimirror

# 使用
curl http://localhost:8081/health
```

### 方式二：源码安装

```bash
# 克隆仓库
git clone https://github.com/livehl/aimirror.git
cd aimirror

# 安装依赖
pip install -r requirements.txt

# 启动
python main.py

# 使用
curl http://localhost:8081/health
```

## 🔧 客户端配置

**pip**
```bash
pip install torch --index-url http://localhost:8081/simple --trusted-host localhost:8081
```

**Docker**
```json
{
  "registry-mirrors": ["http://localhost:8081"]
}
```

**HuggingFace (huggingface-cli)**
```bash
# 设置环境变量
export HF_ENDPOINT=http://localhost:8081

# 下载模型（支持所有文件类型：.gguf, .bin, .safetensors, .json 等）
huggingface-cli download TheBloke/Llama-2-7B-GGUF llama-2-7b.Q4_K_M.gguf

# 下载整个仓库
huggingface-cli download meta-llama/Llama-2-7b-hf --local-dir ./models
```

或使用 Python:
```python
import os
os.environ["HF_ENDPOINT"] = "http://localhost:8081"

from huggingface_hub import hf_hub_download, snapshot_download

# 下载单个文件
hf_hub_download(repo_id="TheBloke/Llama-2-7B-GGUF", filename="llama-2-7b.Q4_K_M.gguf")

# 下载整个仓库
snapshot_download(repo_id="meta-llama/Llama-2-7b-hf", local_dir="./models")
```

## 📖 API

| 路径 | 说明 |
|------|------|
| `/*` | 代理到对应上游 (Docker/PyPI/CRAN/HuggingFace) |
| `/health` | 健康检查 |
| `/stats` | 缓存统计 |

## 🐳 Docker 部署

```bash
# 使用 PyPI 安装（推荐）
pip install aimirror
aimirror

# 或使用 GitHub Container Registry
docker pull ghcr.io/livehl/aimirror:latest

# 运行（基础版）
docker run -d -p 8081:8081 -v $(pwd)/cache:/data/fast_proxy/cache ghcr.io/livehl/aimirror:latest

# 运行（带自定义配置）
docker run -d -p 8081:8081 \
  -v $(pwd)/config.yaml:/app/config.yaml \
  -v $(pwd)/cache:/data/fast_proxy/cache \
  ghcr.io/livehl/aimirror:latest
```

## ⚙️ 配置示例

```yaml
server:
  host: "0.0.0.0"
  port: 8081
  upstream_proxy: "http://proxy.company.com:8080"  # 公司代理（可选）

cache:
  dir: "./cache"
  max_size_gb: 100

rules:
  - name: docker-blob
    pattern: "/v2/.*/blobs/sha256:[a-f0-9]+"
    upstream: "https://registry-1.docker.io"
    strategy: parallel
    min_size: 1048576
    concurrency: 20
    chunk_size: 10485760

  - name: pip-wheel
    pattern: "/packages/.+\.whl$"
    upstream: "https://pypi.org"
    strategy: parallel
    min_size: 1048576
    concurrency: 20
    chunk_size: 5242880

  # HuggingFace 文件下载（支持所有模型文件，临时签名 URL 缓存优化）
  - name: huggingface-files
    pattern: '/.*/(blob|resolve)/main/.+'
    upstream: "https://huggingface.co"
    strategy: parallel
    min_size: 1048576
    concurrency: 20
    chunk_size: 10485760
    cache_key_source: original  # 使用原始 URL 作为缓存 key
    path_rewrite:
      - search: "/blob/main/"
        replace: "/resolve/main/"

  # 示例：扩展任意 HTTP 下载站点
  # - name: my-custom-repo
  #   pattern: '/downloads/.+\.(tar\.gz|zip|bin)$'
  #   upstream: "https://downloads.example.com"
  #   strategy: parallel
  #   min_size: 10485760    # 10MB 以上启用并行
  #   concurrency: 16
  #   chunk_size: 20971520  # 20MB 分片

  - name: default
    pattern: ".*"
    upstream: "https://pypi.org"
    strategy: proxy
```

### 配置说明

| 字段 | 说明 |
|------|------|
| `server.upstream_proxy` | 可选的公司代理，用于连接外网 |
| `rules[].upstream` | 上游源 base URL |
| `rules[].strategy` | `proxy` 直接代理 / `parallel` 并行下载 |
| `rules[].path_rewrite` | 路径重写规则（如 HuggingFace blob→resolve） |
| `rules[].cache_key_source` | `original` 使用原始URL作为缓存key（解决临时签名问题） |

## 🧪 测试

### 运行测试

```bash
# 运行简单测试（无需 pytest）
python test_simple.py

# 运行完整测试套件（需要 pytest）
pytest test_proxy.py -v
```

### 手动验证

**测试 PyPI 代理**
```bash
curl -o /dev/null "http://localhost:8081/packages/fb/d7/71b982339efc4fff3c622c6fefecddfd3e0b35b60c5f822872d5b806bb71/torch-1.0.0-cp27-cp27m-manylinux1_x86_64.whl" \
  -w "HTTP: %{http_code}, Size: %{size_download}, Time: %{time_total}s\n"
```

**测试 HuggingFace 代理**
```bash
export HF_ENDPOINT=http://localhost:8081

# 测试下载 GGUF 模型文件
huggingface-cli download TheBloke/Llama-2-7B-GGUF llama-2-7b.Q4_K_M.gguf

# 测试下载 safetensors 格式模型
huggingface-cli download meta-llama/Llama-2-7b-hf model-00001-of-00002.safetensors

# 测试下载整个仓库
huggingface-cli download sentence-transformers/all-MiniLM-L6-v2 --local-dir ./test-model
```

**测试 Docker Registry 代理**
```bash
# 获取 token
TOKEN=$(curl -s "https://auth.docker.io/token?service=registry.docker.io&scope=repository:library/nginx:pull" \
  | grep -o '"token":"[^"]*"' | cut -d'"' -f4)

# 下载 blob
curl -o /dev/null "http://localhost:8081/v2/library/nginx/blobs/sha256:abc123" \
  -H "Authorization: Bearer $TOKEN" \
  -w "HTTP: %{http_code}, Size: %{size_download}, Time: %{time_total}s\n"
```

**验证缓存命中**
```bash
# 第一次下载（并行下载）
time curl -o /tmp/test1.gguf "http://localhost:8081/unsloth/model/resolve/main/file.gguf"

# 第二次下载（缓存命中，应该快很多）
time curl -o /tmp/test2.gguf "http://localhost:8081/unsloth/model/resolve/main/file.gguf"

# 查看缓存统计
curl http://localhost:8081/stats | jq
```

## 📄 License

MIT
