diff --git a/.github/workflows/nightly-release.yml b/.github/workflows/nightly-release.yml new file mode 100644 index 0000000000..fdece945f5 --- /dev/null +++ b/.github/workflows/nightly-release.yml @@ -0,0 +1,326 @@ +name: Nightly Release + +on: + schedule: + # Run at 00:00 UTC every day + - cron: '0 0 * * *' + workflow_dispatch: + inputs: + date_suffix: + description: 'Date suffix for dev version (YYYYMMDD, leave empty for today)' + required: false + type: string + pull_request: + # TODO: Remove this before merging - only for debugging this PR + +permissions: + contents: write + id-token: write + +jobs: + setup: + runs-on: ubuntu-latest + outputs: + dev_suffix: ${{ steps.set-suffix.outputs.dev_suffix }} + release_tag: ${{ steps.set-suffix.outputs.release_tag }} + version: ${{ steps.set-suffix.outputs.version }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set date suffix and release tag + id: set-suffix + run: | + # Read version from version.txt + VERSION=$(cat version.txt | tr -d '[:space:]') + + # Set date suffix + if [ -n "${{ inputs.date_suffix }}" ]; then + DEV_SUFFIX="${{ inputs.date_suffix }}" + else + DEV_SUFFIX=$(date -u +%Y%m%d) + fi + + # Create release tag with version + RELEASE_TAG="nightly-v${VERSION}-${DEV_SUFFIX}" + + echo "version=${VERSION}" >> $GITHUB_OUTPUT + echo "dev_suffix=${DEV_SUFFIX}" >> $GITHUB_OUTPUT + echo "release_tag=${RELEASE_TAG}" >> $GITHUB_OUTPUT + echo "Base version: ${VERSION}" + echo "Using dev suffix: ${DEV_SUFFIX}" + echo "Release tag: ${RELEASE_TAG}" + + build-flashinfer-python: + needs: setup + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: true + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build wheel + + - name: Build flashinfer-python sdist + env: + FLASHINFER_DEV_RELEASE_SUFFIX: ${{ needs.setup.outputs.dev_suffix }} + run: | + echo "Building flashinfer-python with dev suffix: ${FLASHINFER_DEV_RELEASE_SUFFIX}" + echo "Git commit: $(git rev-parse HEAD)" + python -m build --sdist + ls -lh dist/ + + - name: Upload flashinfer-python artifact + uses: actions/upload-artifact@v4 + with: + name: flashinfer-python-sdist + path: dist/*.tar.gz + retention-days: 7 + + build-flashinfer-cubin: + needs: setup + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: true + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build twine wheel + pip install setuptools>=61.0 requests filelock torch tqdm numpy apache-tvm-ffi==0.1.0b15 + + - name: Build flashinfer-cubin wheel + env: + FLASHINFER_DEV_RELEASE_SUFFIX: ${{ needs.setup.outputs.dev_suffix }} + run: | + echo "Building flashinfer-cubin with dev suffix: ${FLASHINFER_DEV_RELEASE_SUFFIX}" + echo "Git commit: $(git rev-parse HEAD)" + cd flashinfer-cubin + rm -rf dist build *.egg-info + python -m build --wheel + ls -lh dist/ + mkdir -p ../dist + cp dist/*.whl ../dist/ + + - name: Upload flashinfer-cubin artifact + uses: actions/upload-artifact@v4 + with: + name: flashinfer-cubin-wheel + path: dist/*.whl + retention-days: 7 + + build-flashinfer-jit-cache: + needs: setup + strategy: + fail-fast: false + matrix: + cuda: ["12.8", "12.9", "13.0"] + arch: ['x86_64', 'aarch64'] + + runs-on: [self-hosted, "${{ matrix.arch == 'aarch64' && 'arm64' || matrix.arch }}"] + + steps: + - name: Display Machine Information + run: | + echo "CPU: $(nproc) cores, $(lscpu | grep 'Model name' | cut -d':' -f2 | xargs)" + echo "RAM: $(free -h | awk '/^Mem:/ {print $7 " available out of " $2}')" + echo "Disk: $(df -h / | awk 'NR==2 {print $4 " available out of " $2}')" + echo "Architecture: $(uname -m)" + + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: true + + - name: Build wheel in container + env: + DOCKER_IMAGE: ${{ matrix.arch == 'aarch64' && format('pytorch/manylinuxaarch64-builder:cuda{0}', matrix.cuda) || format('pytorch/manylinux2_28-builder:cuda{0}', matrix.cuda) }} + FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda == '12.8' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0a' }} + FLASHINFER_DEV_RELEASE_SUFFIX: ${{ needs.setup.outputs.dev_suffix }} + run: | + # Extract CUDA major and minor versions + CUDA_MAJOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f1) + CUDA_MINOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f2) + export CUDA_MAJOR + export CUDA_MINOR + export CUDA_VERSION_SUFFIX="cu${CUDA_MAJOR}${CUDA_MINOR}" + + chown -R $(id -u):$(id -g) ${{ github.workspace }} + mkdir -p ${{ github.workspace }}/ci-cache + chown -R $(id -u):$(id -g) ${{ github.workspace }}/ci-cache + + # Run the build script inside the container with proper mounts + docker run --rm \ + -v ${{ github.workspace }}:/workspace \ + -v ${{ github.workspace }}/ci-cache:/ci-cache \ + -e FLASHINFER_CI_CACHE=/ci-cache \ + -e CUDA_VERSION="${{ matrix.cuda }}" \ + -e CUDA_MAJOR="$CUDA_MAJOR" \ + -e CUDA_MINOR="$CUDA_MINOR" \ + -e CUDA_VERSION_SUFFIX="$CUDA_VERSION_SUFFIX" \ + -e FLASHINFER_DEV_RELEASE_SUFFIX="${FLASHINFER_DEV_RELEASE_SUFFIX}" \ + -e ARCH="${{ matrix.arch }}" \ + -e FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST}" \ + --user $(id -u):$(id -g) \ + -w /workspace \ + ${{ env.DOCKER_IMAGE }} \ + bash /workspace/scripts/build_flashinfer_jit_cache_whl.sh + timeout-minutes: 180 + + - name: Display wheel size + run: du -h flashinfer-jit-cache/dist/* + + - name: Create artifact name + id: artifact-name + run: | + CUDA_NO_DOT=$(echo "${{ matrix.cuda }}" | tr -d '.') + echo "name=jit-cache-cu${CUDA_NO_DOT}-${{ matrix.arch }}" >> $GITHUB_OUTPUT + + - name: Upload flashinfer-jit-cache artifact + uses: actions/upload-artifact@v4 + with: + name: ${{ steps.artifact-name.outputs.name }} + path: flashinfer-jit-cache/dist/*.whl + retention-days: 7 + + create-release: + needs: [setup, build-flashinfer-python, build-flashinfer-cubin, build-flashinfer-jit-cache] + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Create GitHub Release (empty first) + env: + GH_TOKEN: ${{ github.token }} + run: | + TAG="${{ needs.setup.outputs.release_tag }}" + + # Delete existing release and tag if they exist + if gh release view "$TAG" &>/dev/null; then + echo "Deleting existing release: $TAG" + gh release delete "$TAG" --yes --cleanup-tag + fi + + # Create new release without assets first + gh release create "$TAG" \ + --title "Nightly Release v${{ needs.setup.outputs.version }}-${{ needs.setup.outputs.dev_suffix }}" \ + --notes "Automated nightly build for version ${{ needs.setup.outputs.version }} (dev${{ needs.setup.outputs.dev_suffix }})" \ + --prerelease + + - name: Download flashinfer-python artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-python-sdist + path: dist-python/ + + - name: Upload flashinfer-python to release + env: + GH_TOKEN: ${{ github.token }} + run: | + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-python/* --clobber + + - name: Download flashinfer-cubin artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-cubin-wheel + path: dist-cubin/ + + - name: Upload flashinfer-cubin to release + env: + GH_TOKEN: ${{ github.token }} + run: | + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-cubin/* --clobber + + - name: Upload flashinfer-jit-cache wheels to release (one at a time to avoid OOM) + env: + GH_TOKEN: ${{ github.token }} + run: | + # Upload jit-cache wheels one at a time to avoid OOM + # Each wheel can be several GB, so we download, upload, delete, repeat + mkdir -p dist-jit-cache + + for cuda in 128 129 130; do + for arch in x86_64 aarch64; do + ARTIFACT_NAME="jit-cache-cu${cuda}-${arch}" + echo "Processing ${ARTIFACT_NAME}..." + + # Download this specific artifact + gh run download ${{ github.run_id }} -n "${ARTIFACT_NAME}" -D dist-jit-cache/ || { + echo "Warning: Failed to download ${ARTIFACT_NAME}, skipping..." + continue + } + + # Upload to release + if [ -n "$(ls -A dist-jit-cache/)" ]; then + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-jit-cache/* --clobber + echo "✅ Uploaded ${ARTIFACT_NAME}" + fi + + # Clean up to save disk space before next iteration + rm -rf dist-jit-cache/* + done + done + + update-wheel-index: + needs: [setup, create-release] + runs-on: ubuntu-latest + steps: + - name: Checkout flashinfer repo + uses: actions/checkout@v4 + + - name: Download all artifacts + uses: actions/download-artifact@v4 + with: + path: artifacts/ + + - name: Collect wheels + run: | + mkdir -p dist + find artifacts/ -name "*.whl" -exec cp {} dist/ \; + ls -lh dist/ + + - name: Clone wheel index + run: git clone https://oauth2:${WHL_TOKEN}@github.com/flashinfer-ai/whl.git flashinfer-whl + env: + WHL_TOKEN: ${{ secrets.WHL_TOKEN }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Update wheel index + run: | + python3 scripts/update_whl_index.py \ + --dist-dir dist \ + --output-dir flashinfer-whl \ + --release-tag "${{ needs.setup.outputs.release_tag }}" + + - name: Push wheel index + run: | + cd flashinfer-whl + git config --local user.name "github-actions[bot]" + git config --local user.email "41898282+github-actions[bot]@users.noreply.github.com" + git add -A + git commit -m "update whl for nightly ${{ needs.setup.outputs.dev_suffix }}" + git push diff --git a/flashinfer-cubin/build_backend.py b/flashinfer-cubin/build_backend.py index 8815390ad5..89c531af99 100644 --- a/flashinfer-cubin/build_backend.py +++ b/flashinfer-cubin/build_backend.py @@ -11,15 +11,37 @@ # Add parent directory to path to import artifacts module sys.path.insert(0, str(Path(__file__).parent.parent)) + +def _get_git_version(): + """Get git commit hash.""" + import subprocess + + try: + git_version = ( + subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parent.parent, + stderr=subprocess.DEVNULL, + ) + .decode("ascii") + .strip() + ) + return git_version + except Exception: + return "unknown" + + # add flashinfer._build_meta, always override to ensure version is up-to-date build_meta_file = Path(__file__).parent.parent / "flashinfer" / "_build_meta.py" version_file = Path(__file__).parent.parent / "version.txt" if version_file.exists(): with open(version_file, "r") as f: version = f.read().strip() +git_version = _get_git_version() with open(build_meta_file, "w") as f: f.write('"""Build metadata for flashinfer package."""\n') f.write(f'__version__ = "{version}"\n') + f.write(f'__git_version__ = "{git_version}"\n') def _download_cubins(): @@ -60,6 +82,14 @@ def _create_build_metadata(): else: version = "0.0.0+unknown" + # Add dev suffix if specified + dev_suffix = os.environ.get("FLASHINFER_DEV_RELEASE_SUFFIX", "") + if dev_suffix: + version = f"{version}.dev{dev_suffix}" + + # Get git version + git_version = _get_git_version() + # Create build metadata in the source tree package_dir = Path(__file__).parent / "flashinfer_cubin" build_meta_file = package_dir / "_build_meta.py" @@ -67,6 +97,7 @@ def _create_build_metadata(): with open(build_meta_file, "w") as f: f.write('"""Build metadata for flashinfer-cubin package."""\n') f.write(f'__version__ = "{version}"\n') + f.write(f'__git_version__ = "{git_version}"\n') print(f"Created build metadata file with version {version}") return version diff --git a/flashinfer-cubin/download_cubins.py b/flashinfer-cubin/download_cubins.py index 2f2c847bbe..cbfd0d4f40 100644 --- a/flashinfer-cubin/download_cubins.py +++ b/flashinfer-cubin/download_cubins.py @@ -23,7 +23,7 @@ # Add parent directory to path to import flashinfer modules sys.path.insert(0, str(Path(__file__).parent.parent)) -from flashinfer.artifacts import download_artifacts +# from flashinfer.artifacts import download_artifacts from flashinfer.jit.cubin_loader import FLASHINFER_CUBINS_REPOSITORY @@ -69,7 +69,8 @@ def main(): # Use the existing download_artifacts function try: - download_artifacts() + # NOTE(Zihao): just for debugging, remove later + # download_artifacts() print("Download complete!") except Exception as e: print(f"Download failed: {e}") diff --git a/flashinfer-cubin/flashinfer_cubin/__init__.py b/flashinfer-cubin/flashinfer_cubin/__init__.py index b28fb09da5..abfe816282 100644 --- a/flashinfer-cubin/flashinfer_cubin/__init__.py +++ b/flashinfer-cubin/flashinfer_cubin/__init__.py @@ -63,5 +63,18 @@ def _get_version(): return "0.0.0" +def _get_git_version(): + # First try to read from build metadata (for wheel distributions) + try: + from . import _build_meta + + return _build_meta.__git_version__ + except (ImportError, AttributeError): + pass + + return "unknown" + + __version__ = _get_version() +__git_version__ = _get_git_version() __all__ = ["get_cubin_dir", "list_cubins", "get_cubin_path", "CUBIN_DIR"] diff --git a/flashinfer-jit-cache/build_backend.py b/flashinfer-jit-cache/build_backend.py index 7e0d9f4aaa..c8d736c602 100644 --- a/flashinfer-jit-cache/build_backend.py +++ b/flashinfer-jit-cache/build_backend.py @@ -22,15 +22,37 @@ # Add parent directory to path to import flashinfer modules sys.path.insert(0, str(Path(__file__).parent.parent)) + +def _get_git_version(): + """Get git commit hash.""" + import subprocess + + try: + git_version = ( + subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parent.parent, + stderr=subprocess.DEVNULL, + ) + .decode("ascii") + .strip() + ) + return git_version + except Exception: + return "unknown" + + # add flashinfer._build_meta, always override to ensure version is up-to-date build_meta_file = Path(__file__).parent.parent / "flashinfer" / "_build_meta.py" version_file = Path(__file__).parent.parent / "version.txt" if version_file.exists(): with open(version_file, "r") as f: version = f.read().strip() +git_version = _get_git_version() with open(build_meta_file, "w") as f: f.write('"""Build metadata for flashinfer package."""\n') f.write(f'__version__ = "{version}"\n') + f.write(f'__git_version__ = "{git_version}"\n') def get_version(): diff --git a/flashinfer-jit-cache/flashinfer_jit_cache/__init__.py b/flashinfer-jit-cache/flashinfer_jit_cache/__init__.py index 986232c63a..3be467e99a 100644 --- a/flashinfer-jit-cache/flashinfer_jit_cache/__init__.py +++ b/flashinfer-jit-cache/flashinfer_jit_cache/__init__.py @@ -29,8 +29,10 @@ def get_jit_cache_dir() -> str: try: from ._build_meta import __version__ as __version__ + from ._build_meta import __git_version__ as __git_version__ except ModuleNotFoundError: __version__ = "0.0.0+unknown" + __git_version__ = "unknown" __all__ = [ diff --git a/flashinfer-jit-cache/setup.py b/flashinfer-jit-cache/setup.py index e65b216cce..2c0a0c7dbe 100644 --- a/flashinfer-jit-cache/setup.py +++ b/flashinfer-jit-cache/setup.py @@ -16,6 +16,11 @@ def get_version(): else: version = "0.0.0" + # Add dev suffix if specified + dev_suffix = os.environ.get("FLASHINFER_DEV_RELEASE_SUFFIX", "") + if dev_suffix: + version = f"{version}.dev{dev_suffix}" + # Append CUDA version suffix if available cuda_suffix = os.environ.get("CUDA_VERSION_SUFFIX", "") if cuda_suffix: @@ -25,13 +30,34 @@ def get_version(): return version +def get_git_version(): + """Get git commit hash.""" + import subprocess + + try: + git_version = ( + subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parent.parent, + stderr=subprocess.DEVNULL, + ) + .decode("ascii") + .strip() + ) + return git_version + except Exception: + return "unknown" + + def generate_build_meta(): """Generate build metadata file.""" build_meta_file = Path(__file__).parent / "flashinfer_jit_cache" / "_build_meta.py" version = get_version() + git_version = get_git_version() with open(build_meta_file, "w") as f: f.write('"""Build metadata for flashinfer-jit-cache package."""\n') f.write(f'__version__ = "{version}"\n') + f.write(f'__git_version__ = "{git_version}"\n') class PlatformSpecificBdistWheel(bdist_wheel): diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index ccfae46ee7..ec08322649 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -18,8 +18,10 @@ try: from ._build_meta import __version__ as __version__ + from ._build_meta import __git_version__ as __git_version__ except ModuleNotFoundError: __version__ = "0.0.0+unknown" + __git_version__ = "unknown" from . import jit as jit diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 5062e7c576..81fc8e5eb8 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -525,6 +525,9 @@ def gen_all_modules( # Add cuDNN FMHA module jit_specs.append(gen_cudnn_fmha_module()) + # NOTE(Zihao): just for debugging, remove later + jit_specs = [gen_spdlog_module()] + # dedup names = set() ret: List[JitSpec] = [] diff --git a/scripts/build_flashinfer_jit_cache_whl.sh b/scripts/build_flashinfer_jit_cache_whl.sh index 30495ff377..fd1a313e14 100755 --- a/scripts/build_flashinfer_jit_cache_whl.sh +++ b/scripts/build_flashinfer_jit_cache_whl.sh @@ -27,8 +27,10 @@ echo "CUDA Major: ${CUDA_MAJOR}" echo "CUDA Minor: ${CUDA_MINOR}" echo "CUDA Version Suffix: ${CUDA_VERSION_SUFFIX}" echo "CUDA Architectures: ${FLASHINFER_CUDA_ARCH_LIST}" +echo "Dev Release Suffix: ${FLASHINFER_DEV_RELEASE_SUFFIX}" echo "MAX_JOBS: ${MAX_JOBS}" echo "Python Version: $(python3 --version)" +echo "Git commit: $(git rev-parse HEAD 2>/dev/null || echo 'unknown')" echo "Working directory: $(pwd)" echo "" @@ -62,6 +64,16 @@ echo "" echo "Built wheels:" ls -lh dist/ +# Verify version and git version +echo "" +echo "Verifying version and git version..." +pip install dist/*.whl +python -c " +import flashinfer_jit_cache +print(f'đŸ“Ļ Package version: {flashinfer_jit_cache.__version__}') +print(f'🔖 Git version: {flashinfer_jit_cache.__git_version__}') +" + # Copy wheels to output directory if specified if [ -n "${OUTPUT_DIR}" ]; then echo "" diff --git a/scripts/update_whl_index.py b/scripts/update_whl_index.py index 4902d0cc7f..01b6187996 100644 --- a/scripts/update_whl_index.py +++ b/scripts/update_whl_index.py @@ -1,17 +1,199 @@ +""" +Update wheel index for flashinfer packages. + +This script generates PEP 503 compatible simple repository index pages for: +- flashinfer-python (no CUDA suffix in version) +- flashinfer-cubin (no CUDA suffix in version) +- flashinfer-jit-cache (has CUDA suffix like +cu130) + +The index is organized by CUDA version for jit-cache, and flat for others. +""" + import hashlib import pathlib import re +import argparse +import sys +from typing import Optional + + +def get_cuda_version(wheel_name: str) -> Optional[str]: + """Extract CUDA version from wheel filename.""" + # Match patterns like +cu128, +cu129, +cu130 + match = re.search(r"\+cu(\d+)", wheel_name) + if match: + return match.group(1) + return None + + +def get_package_info(wheel_path: pathlib.Path) -> Optional[dict]: + """Extract package information from wheel filename.""" + wheel_name = wheel_path.name + + # Try flashinfer-python pattern + match = re.match(r"flashinfer_python-([0-9.]+(?:\.dev\d+)?)-", wheel_name) + if match: + version = match.group(1) + return { + "package": "flashinfer-python", + "version": version, + "cuda": None, + } + + # Try flashinfer-cubin pattern + match = re.match(r"flashinfer_cubin-([0-9.]+(?:\.dev\d+)?)-", wheel_name) + if match: + version = match.group(1) + return { + "package": "flashinfer-cubin", + "version": version, + "cuda": None, + } + + # Try flashinfer-jit-cache pattern (has CUDA suffix in version) + match = re.match(r"flashinfer_jit_cache-([0-9.]+(?:\.dev\d+)?\+cu\d+)-", wheel_name) + if match: + version = match.group(1) + cuda_ver = get_cuda_version(wheel_name) + return { + "package": "flashinfer-jit-cache", + "version": version, + "cuda": cuda_ver, + } + + return None + + +def compute_sha256(file_path: pathlib.Path) -> str: + """Compute SHA256 hash of a file.""" + with open(file_path, "rb") as f: + return hashlib.sha256(f.read()).hexdigest() + + +def update_index( + dist_dir: str = "dist", + output_dir: str = "whl", + base_url: str = "https://github.com/flashinfer-ai/flashinfer/releases/download", + release_tag: Optional[str] = None, +): + """ + Update wheel index from dist directory. + + Args: + dist_dir: Directory containing wheel files + output_dir: Output directory for index files + base_url: Base URL for wheel downloads + release_tag: GitHub release tag (e.g., 'nightly' or 'v0.3.1') + """ + dist_path = pathlib.Path(dist_dir) + if not dist_path.exists(): + print(f"Error: dist directory '{dist_dir}' does not exist") + sys.exit(1) + + wheels = sorted(dist_path.glob("*.whl")) + if not wheels: + print(f"No wheel files found in '{dist_dir}'") + sys.exit(1) + + print(f"Found {len(wheels)} wheel file(s)") + + for wheel_path in wheels: + print(f"\nProcessing: {wheel_path.name}") + + # Extract package information + info = get_package_info(wheel_path) + if not info: + print(" âš ī¸ Skipping: Could not parse wheel filename") + continue + + # Compute SHA256 + sha256 = compute_sha256(wheel_path) + + # Determine index directory + package = info["package"] + cuda = info["cuda"] + + if cuda: + # CUDA-specific index for jit-cache: whl/cu130/flashinfer-jit-cache/ + index_dir = pathlib.Path(output_dir) / f"cu{cuda}" / package + else: + # No CUDA version for python/cubin: whl/flashinfer-python/ + index_dir = pathlib.Path(output_dir) / package + + index_dir.mkdir(parents=True, exist_ok=True) + + # Construct download URL + tag = release_tag or f"v{info['version'].split('+')[0].split('.dev')[0]}" + download_url = f"{base_url}/{tag}/{wheel_path.name}#sha256={sha256}" + + # Update index.html + index_file = index_dir / "index.html" + + # Read existing content to avoid duplicates + existing_links = set() + if index_file.exists(): + with index_file.open("r") as f: + existing_links = set(f.readlines()) + else: + # Create new index file with HTML header + with index_file.open("w") as f: + f.write("\n") + f.write("\n") + f.write(f"Links for {package}\n") + f.write("\n") + f.write(f"

Links for {package}

\n") + print(f" 📝 Created new index file: {index_dir}/index.html") + + # Create new link + new_link = f'{wheel_path.name}
\n' + + if new_link in existing_links: + print(f" â„šī¸ Already in index: {index_dir}/index.html") + else: + with index_file.open("a") as f: + f.write(new_link) + print(f" ✅ Added to index: {index_dir}/index.html") + + print(f" đŸ“Ļ Package: {package}") + print(f" 🔖 Version: {info['version']}") + if cuda: + print(f" 🎮 CUDA: cu{cuda}") + print(f" 📍 URL: {download_url}") + + +def main(): + parser = argparse.ArgumentParser( + description="Update wheel index for flashinfer packages" + ) + parser.add_argument( + "--dist-dir", + default="dist", + help="Directory containing wheel files (default: dist)", + ) + parser.add_argument( + "--output-dir", + default="whl", + help="Output directory for index files (default: whl)", + ) + parser.add_argument( + "--base-url", + default="https://github.com/flashinfer-ai/flashinfer/releases/download", + help="Base URL for wheel downloads", + ) + parser.add_argument( + "--release-tag", + help="GitHub release tag (e.g., 'nightly' or 'v0.3.1'). If not specified, will be derived from version.", + ) + + args = parser.parse_args() + + update_index( + dist_dir=args.dist_dir, + output_dir=args.output_dir, + base_url=args.base_url, + release_tag=args.release_tag, + ) + -for path in sorted(pathlib.Path("dist").glob("*.whl")): - with open(path, "rb") as f: - sha256 = hashlib.sha256(f.read()).hexdigest() - ver, cu, torch = re.findall( - r"flashinfer_python-([0-9.]+(?:\.post[0-9]+)?)\+cu(\d+)torch([0-9.]+)-", - path.name, - )[0] - index_dir = pathlib.Path(f"flashinfer-whl/cu{cu}/torch{torch}/flashinfer-python") - index_dir.mkdir(exist_ok=True) - base_url = "https://github.com/flashinfer-ai/flashinfer/releases/download" - full_url = f"{base_url}/v{ver}/{path.name}#sha256={sha256}" - with (index_dir / "index.html").open("a") as f: - f.write(f'{path.name}
\n') +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 6c663c1cb6..8fdc0a76a5 100755 --- a/setup.py +++ b/setup.py @@ -30,12 +30,35 @@ def write_if_different(path: Path, content: str) -> None: def get_version(): + import os + package_version = (root / "version.txt").read_text().strip() - return f"{package_version}" + dev_suffix = os.environ.get("FLASHINFER_DEV_RELEASE_SUFFIX", "") + if dev_suffix: + package_version = f"{package_version}.dev{dev_suffix}" + return package_version + + +def get_git_version(): + """Get git commit hash.""" + import subprocess + + try: + git_version = ( + subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=root, stderr=subprocess.DEVNULL + ) + .decode("ascii") + .strip() + ) + return git_version + except Exception: + return "unknown" def generate_build_meta() -> None: build_meta_str = f"__version__ = {get_version()!r}\n" + build_meta_str += f"__git_version__ = {get_git_version()!r}\n" write_if_different(root / "flashinfer" / "_build_meta.py", build_meta_str)