update CI scripts #1552

Workflow file for this run

	name: CI

	on:
	push:
	branches:
	- main
	- "pull-request/[0-9]+"
	tags:
	- '**'
	# `merge_group` fires when a PR enters GitHub's merge queue. GitHub creates
	# a temporary `gh-readonly-queue/main/pr-<n>-<sha>` ref containing the PR
	# rebased on top of `main` + any earlier queued PRs and runs the required
	# checks against it. The job names below must match the required-check
	# names configured in branch protection.
	merge_group:
	branches: [main]

	jobs:
	cpu:
	runs-on: linux-amd64-cpu8
	container:
	image: nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04
	env:
	UV_PROJECT_ENVIRONMENT: /tmp/flashdreams-venv
	UV_LINK_MODE: copy
	# Exercise the dependency branch used by current Python releases.
	UV_PYTHON: "3.12"
	steps:
	- name: Get PR info
	if: startsWith(github.ref_name, 'pull-request/')
	id: get-pr-info
	uses: nv-gha-runners/get-pr-info@main

	- name: Install lint system dependencies
	run: \|
	apt-get update -qq
	DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \
	python3 python3-dev python3-venv \
	git curl ca-certificates

	- name: Checkout
	uses: actions/checkout@v4

	- name: Setup proxy cache
	uses: nv-gha-runners/setup-proxy-cache@main

	- name: Setup uv
	uses: astral-sh/setup-uv@v6
	with:
	enable-cache: true
	cache-suffix: "cpu"

	- name: Install lint dependencies
	run: \|
	# ty needs third-party packages to resolve imports. Skip packages
	# that cannot be installed on CPU-only runners:
	# - transformer-engine-torch: source-only, requires GPU arch to compile
	uv venv --clear
	uv sync --locked --extra dev --group lint \
	--no-install-package transformer-engine-torch

	- name: Run linter checks
	run: \|
	# Container runs as root but the workspace is owned by the runner
	# uid -- git refuses to operate without safe.directory.
	git config --global --add safe.directory "$(pwd)"
	uv run --group lint pre-commit run -a

	- name: Install remaining system dependencies
	run: \|
	DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \
	ffmpeg \
	gcc g++ ninja-build \
	libnccl-dev \
	unzip
	rm -rf /var/lib/apt/lists/*

	- name: Build flashdreams wheel
	run: uv build --wheel --package flashdreams

	- name: Run CPU unit tests
	run: uv run --no-sync pytest -m ci_cpu -v

	# Documentation is built and deployed by `.github/workflows/doc.yml`.

	gpu:
	runs-on: linux-amd64-gpu-rtxpro6000-latest-2
	timeout-minutes: 120
	container:
	image: nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04
	options: --gpus all
	env:
	# Place the venv outside the workspace so uv does not interfere with
	# the checkout directory owned by the runner.
	UV_PROJECT_ENVIRONMENT: /tmp/flashdreams-venv
	UV_LINK_MODE: copy
	# imgui-bundle publishes Linux ARM64 wheels for Python 3.11+, but not
	# 3.10. Keep GB300 CI on a wheel-supported interpreter and fail instead
	# of attempting its large, platform-sensitive native source build.
	UV_PYTHON: "3.11"
	UV_NO_BUILD_PACKAGE: imgui-bundle
	# Limit parallel CUDA compilation jobs to avoid OOM (each nvcc job
	# uses ~9GB peak memory).
	MAX_JOBS: 8
	OMNIDREAMS_QUALITY_REFERENCE_REPO: ${{ vars.OMNIDREAMS_QUALITY_REFERENCE_REPO \|\| 'nvidia/omni-dreams-samples' }}
	OMNIDREAMS_QUALITY_REFERENCE_PATH: ${{ vars.OMNIDREAMS_QUALITY_REFERENCE_PATH \|\| 'data/quality_references/single_view/omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae/239560dc-33d1-11ef-9720-00044bcbccac/reference_compare_region.mp4' }}
	OMNIDREAMS_QUALITY_REFERENCE_CLIP: /tmp/omnidreams_quality_reference/reference_compare_region.mp4
	OMNIDREAMS_QUALITY_ARTIFACT_DIR: /tmp/omnidreams_quality_artifacts
	steps:
	- name: Get PR info
	if: startsWith(github.ref_name, 'pull-request/')
	id: get-pr-info
	uses: nv-gha-runners/get-pr-info@main

	- name: Detect GPU architecture
	id: gpu-arch
	run: \|
	nvidia-smi
	# Get compute capability (e.g. "12.0") and strip the dot for nvcc
	# arch format (e.g. "120").
	compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \| head -1 \| tr -d '[:space:]')
	arch=$(echo "${compute_cap}" \| tr -d '.')
	echo "arch=${arch}" >> "$GITHUB_OUTPUT"
	echo "Detected GPU compute capability: ${compute_cap} -> sm_${arch}"

	- name: Checkout
	uses: actions/checkout@v4

	- name: Install system dependencies
	run: \|
	apt-get update -qq
	DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \
	python3 python3-dev python3-venv \
	ffmpeg \
	gcc g++ ninja-build \
	libgl1-mesa-dev \
	libx11-dev libxcursor-dev libxext-dev libxi-dev libxinerama-dev libxrandr-dev \
	libnccl-dev \
	curl git ca-certificates unzip
	rm -rf /var/lib/apt/lists/*

	- name: Setup proxy cache
	uses: nv-gha-runners/setup-proxy-cache@main

	- name: Setup uv
	uses: astral-sh/setup-uv@v6
	with:
	enable-cache: true
	cache-suffix: "gpu-sm${{ steps.gpu-arch.outputs.arch }}"
	# Disable default prune --ci which removes pre-built wheels but also
	# invalidates git-sourced builds (timestamps change on cache restore).
	# Instead we selectively trim the cache in a post step below.
	prune-cache: false

	- name: Install dependencies
	env:
	NVTE_CUDA_ARCHS: ${{ steps.gpu-arch.outputs.arch }}
	run: \|
	uv venv --clear
	uv sync --locked --extra dev

	- name: Verify GPU availability
	run: \|
	echo "Verifying GPU access..."
	nvidia-smi

	- name: Download Omnidreams quality reference
	env:
	HF_TOKEN: ${{ secrets.HF_TOKEN }}
	run: \|
	mkdir -p "$(dirname "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}")"
	if [ -n "${HF_TOKEN:-}" ]; then
	uv run --no-sync hf download \
	"${OMNIDREAMS_QUALITY_REFERENCE_REPO}" \
	"${OMNIDREAMS_QUALITY_REFERENCE_PATH}" \
	--repo-type dataset \
	--local-dir "$(dirname "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}")" \
	--token "${HF_TOKEN}"
	else
	uv run --no-sync hf download \
	"${OMNIDREAMS_QUALITY_REFERENCE_REPO}" \
	"${OMNIDREAMS_QUALITY_REFERENCE_PATH}" \
	--repo-type dataset \
	--local-dir "$(dirname "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}")"
	fi
	cp \
	"$(dirname "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}")/${OMNIDREAMS_QUALITY_REFERENCE_PATH}" \
	"${OMNIDREAMS_QUALITY_REFERENCE_CLIP}"
	test -s "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}"

	- name: Run GPU unit tests
	env:
	HF_TOKEN: ${{ secrets.HF_TOKEN }}
	FLASHDREAMS_OMNIDREAMS_QUALITY_REFERENCE_CLIP: ${{ env.OMNIDREAMS_QUALITY_REFERENCE_CLIP }}
	FLASHDREAMS_OMNIDREAMS_QUALITY_EXAMPLE_DATA: "1"
	FLASHDREAMS_OMNIDREAMS_QUALITY_TOTAL_BLOCKS: "4"
	FLASHDREAMS_OMNIDREAMS_QUALITY_ARTIFACT_DIR: ${{ env.OMNIDREAMS_QUALITY_ARTIFACT_DIR }}
	run: uv run --no-sync pytest -m ci_gpu -v

	- name: Upload Omnidreams quality artifacts
	if: always()
	uses: actions/upload-artifact@v4
	with:
	name: omnidreams-quality-artifacts
	path: ${{ env.OMNIDREAMS_QUALITY_ARTIFACT_DIR }}
	if-no-files-found: ignore

	- name: Trim uv cache for upload
	if: always()
	run: \|
	cache_dir="${UV_CACHE_DIR:-/github/home/.cache/uv}"
	echo "=== Cache size before trim ==="
	du -sh "${cache_dir}" \|\| true
	du -sh "${cache_dir}"/*/ 2>/dev/null \|\| true

	# Remove pre-built (downloaded) wheels -- these are re-downloaded
	# quickly from the proxy cache on each run.
	rm -rf "${cache_dir}/wheels-v6"

	# Remove unzipped wheel archives -- uv can re-extract from .whl
	# files in sdists-v9/ on next run.
	rm -rf "${cache_dir}/archive-v0"

	# Remove build artifacts from git checkouts (build/, *.egg-info,
	# __pycache__) that bloat the cache but aren't needed for reuse.
	find "${cache_dir}/git-v0/checkouts" \
	$ -name "build" -o -name "*.egg-info" -o -name "__pycache__" $ \
	-type d -exec rm -rf {} + 2>/dev/null \|\| true

	# Remove editable installs -- these are in-repo workspace members
	# that get rebuilt instantly from source on every sync.
	rm -rf "${cache_dir}/sdists-v9/editable"

	echo ""
	echo "=== Cache size after trim ==="
	du -sh "${cache_dir}" \|\| true
	du -sh "${cache_dir}"/*/ 2>/dev/null \|\| true
	echo ""
	echo "=== Cached built wheels (sdists-v9) ==="
	find "${cache_dir}/sdists-v9" -name "*.whl" -exec ls -lh {} \; 2>/dev/null \|\| true

	# ===========================================================================
	# Publish to PyPI -- runs only on main after cpu + gpu jobs pass.
	# ===========================================================================
	publish-pypi:
	if: github.ref == 'refs/heads/main' && github.event_name == 'push'
	needs: [cpu, gpu]
	runs-on: linux-amd64-cpu8
	steps:
	- name: Checkout
	uses: actions/checkout@v4

	- name: Setup uv
	uses: astral-sh/setup-uv@v6
	with:
	enable-cache: true

	- name: Build wheel
	run: uv build --wheel --package flashdreams

	- name: Previously published check
	id: previously_published
	run: \|
	version=$(sed -n 's/^__version__ = "$.*$"/\1/p' flashdreams/flashdreams/_version.py)
	echo "version=$version"
	if curl -sSf "https://pypi.org/simple/flashdreams/" \
	\| grep -q "flashdreams-${version}-"; then
	echo "Version $version already exists on PyPI -- skipping publish"
	echo "skip=true" >> "$GITHUB_OUTPUT"
	else
	echo "Version $version not found on PyPI -- will publish"
	echo "skip=false" >> "$GITHUB_OUTPUT"
	fi

	- name: Publish to PyPI
	if: steps.previously_published.outputs.skip != 'true'
	run: >
	uv publish
	--publish-url https://upload.pypi.org/legacy/
	--check-url https://pypi.org/simple/
	--token ${{ secrets.PYPI_API_TOKEN }}
	dist/*

Navigation Menu

Uh oh!

update CI scripts #1552

Workflow file

update CI scripts #1552

Uh oh!

Workflow file for this run