diff --git a/.github/workflows/cache-dit-examples.yml b/.github/workflows/cache-dit-examples.yml new file mode 100644 index 000000000..8b12f3e6d --- /dev/null +++ b/.github/workflows/cache-dit-examples.yml @@ -0,0 +1,58 @@ +# cache-dit examples guard - project thin trigger. +# +# Calls the common engine .github/workflows/examples-template.yml; +# this file only declares what varies per project: schedule, +# concurrency, upstream repo, dispatch input forwarding, and matrix +# parallelism. The guard loop (release-only monitor / manifest-check / +# run-example matrix / validate-results / save-monitor-state) lives in +# the engine. See docs/examples-guard-engine.md for the design. + +name: cache-dit-examples + +concurrency: + # schedule runs serialize; each manual run gets its own group so a + # dispatch is never blocked behind a scheduled run. + group: ${{ github.event_name == 'schedule' && 'cache-dit-examples-schedule' || format('manual-{0}', github.run_id) }} + # false: a run cancelled mid-way loses its outcome writeback, and the + # outcome is what drives the next cycle's retry. + cancel-in-progress: false + +on: + schedule: + - cron: '35 0 * * *' + workflow_dispatch: + inputs: + target_ref: + description: Branch, tag, or SHA to test (empty = latest upstream release; falls back to main on API failure) + required: false + default: '' + type: string + # PR trigger: changes to the guard assets (examples-guards/cache-dit/**) + # or this workflow file re-run the guard against the latest upstream + # release so a regression is caught before merge. paths filter avoids + # burning the self-hosted NPU runner on unrelated PRs. `pull_request` + # (not `pull_request_target`): contents: read is enough, no write-token + # risk. Mirrors the peft-quick-start example. + pull_request: + branches: [main] + paths: + - 'examples-guards/cache-dit/**' + - '.github/workflows/cache-dit-examples.yml' + +# Engine jobs query the GitHub Job API (validate-results); called +# workflows run with the caller's GITHUB_TOKEN permissions. +permissions: + actions: read + contents: read + +jobs: + cache-dit-examples: + uses: ./.github/workflows/examples-template.yml + with: + # Namespaces cache keys (examples-monitor-state-cache-dit_*), + # artifact names and the scheduling paths under + # examples-guards/cache-dit. + project: cache-dit + upstream_repo: vipshop/cache-dit + target_ref: ${{ inputs.target_ref }} + max_parallel: 4 diff --git a/examples-guards/cache-dit/examples_manifest.yaml b/examples-guards/cache-dit/examples_manifest.yaml new file mode 100644 index 000000000..9b5391c23 --- /dev/null +++ b/examples-guards/cache-dit/examples_manifest.yaml @@ -0,0 +1,90 @@ +version: 1 +scan: + root: examples + include_extensions: ['.py'] + unit: files +supported: + # overlay_args 约定: + # - 只写与脚本默认值不同的项;每个参数的注释解释其作用或取值原因。 + # - ${FLUX_MODEL_PATH} 引用 setup_example.sh 写入 GITHUB_ENV 的模型 + # 快照路径变量(FLUX.1-dev 本地快照)。 + # examples/generate.py: thin wrapper over cache_dit.generate.entrypoint() + # (unified CLI, README "python3 -m cache_dit.generate"). Positional name + # "flux" selects the FLUX.1-dev example; --model-path 由 setup_example.sh + # 调用 scripts/hub_cache.py 解析:共享缓存盘(HF_HOME/hub)里有必需文件 + # 齐全的快照就直接用,零网络;未命中则先下到本 job 在同一块盘上的暂存 + # 目录,校验后逐个文件改名发布进快照,后续 run(同池任意 pod)直接命中。 + # 下载按 allow_patterns 裁剪到 FluxPipeline.from_pretrained 真正读取的 + # diffusers 组件树(57.89 GB → 33.74 GB,排除库根下重复的单文件 + # flux1-dev.safetensors 与 ae.safetensors)。FLUX.1-dev 在 HF 为 gated + # repo 且本仓无 HF_TOKEN,HF 一腿必然 401,实际由 ModelScope 同布局镜像 + # 投递; + # --attn _native_npu is the NPU attention backend; --cache exercises + # cache-dit's DBCache path. + - path: examples/generate.py + profile: flux + runner: linux-aarch64-a2-1 + npu_devices: '0' + image: swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-910b-ubuntu22.04-py3.12 + overlay_args: + # 位置参数:选择 FLUX.1-dev 示例(unified CLI 的示例名,非 --flag) + - flux + # 模型来源:setup_example.sh 导出的 FLUX.1-dev 本地快照路径 + - --model-path ${FLUX_MODEL_PATH} + # attention 后端选 NPU 原生实现(默认 CUDA 路径在 NPU 上不可用) + - --attn _native_npu + # 启用 cache-dit 的 DBCache 缓存路径(本项目被测核心能力) + - --cache + # pipeline 组件 CPU offload:FLUX.1-dev bf16 全量权重超 910B4 单卡 + # 可用显存,offload 后方可单卡运行 + - --cpu-offload + # 预热步数取最小,smoke 只验证链路 + - --warmup 1 + # 正式推理重复次数取最小,smoke 只验证链路 + - --repeat 1 + # 去噪采样步数;smoke 取可覆盖 DBCache 命中路径的最小配置 + - --steps 8 + # 冷缓存首次要把 33.74 GB 拉进共享盘(HF gated → ModelScope 镜像, + # 集群出口实测 1.5–3.1 MB/s),这一步必须装进预算;暖缓存命中后 + # setup 只剩环境安装。120 分钟曾被下载撞满(run 38520772429: + # setup 占满 119.7 分钟,run example 一步没跑)。 + timeout_minutes: 180 +unsupported: + # --- examples/api/*:上游坏脚本,manifest 层无法规避 --- + # PR #752(feat: add cache-dit-generate cli tool)删除了 examples/utils.py, + # 但未同步更新 examples/api/ 下的脚本;4 个脚本开头的 + # `from utils import get_args, strify, MemoryTracker` 在干净 checkout 下 + # 直接 ImportError(utils.py 仅存于 tests/api/utils.py,不在 examples 的 + # sys.path 上)。上游修复后移回。 + - examples/api/run_cache_refresh_flux.py + - examples/api/run_cache_refresh_wan_2.2.py + - examples/api/run_steps_mask.py + - examples/api/run_transformer_only.py + # --- examples/generate.py 分布式变体(torchrun 4 卡,ulysses/ring/usp, + # README "Distributed Inference"):两条模型路线在 NPU 上均不可行, + # 上游适配后移回: + # 1) FLUX.1-dev:序列并行不分片权重,每个 rank 须装载完整 pipeline, + # ~28.9GB bf16 超出 32GB 910B4 单卡可用显存(29.49GB;run + # 36657738934:4 rank 全部在 pipe.to(device) OOM)。加卡无效 + # (8 卡同样 OOM,权重不随卡数分摊);--cpu-offload 在 torchrun + # 多进程下上游未验证。 + # 2) Z-Image-Turbo(~19GB,显存适配,run 36660209993 已完成加载并 + # 进入 warmup):diffusers Z-Image 的 RoPE 频率表为 complex64 + # (torch.polar 预计算,按位置 id 索引),torch_npu 的 aclnnIndex + # 无 DT_COMPLEX64 kernel,warmup 第一步即崩(EZ1001, error + # 161002);attention processor 的复数乘法同样存在缺口。上游官方 + # NPU 适配:https://github.com/huggingface/diffusers/pull/12979 + # (draft 未合入;cache-dit docs/supported_matrix/ASCEND_NPU.md + # 亦注明需合入该 PR 后方可支持 Z-Image-Turbo)。PR 合入后恢复 + # 条目并补 setup 的模型下载即可(torchrun -m cache_dit.generate + + # zimage_turbo),无需本地补丁。 + - examples/generate.py + # --- examples/ptq/flux2_svdq_ptq.py:依赖 svdq CUDA kernel 扩展 + #(cache_dit.kernels.svdq_extension_is_available),NPU 环境不可用。 + - examples/ptq/flux2_svdq_ptq.py + # --- examples/ray/ray_wrapper_example.py:脚本 CUDA 硬编码,NPU 不可用 + #(与 runner 卡数无关):--use-init-fn 的 worker init_fn、主进程 pipeline、 + # transformer 转移共 3 处 .to("cuda");Ray GPU actor 生态面向 CUDA; + # --quantize(float8/torchao)、--use-compile 亦为 CUDA-only。上游适配 + # NPU 后再评估。 + - examples/ray/ray_wrapper_example.py diff --git a/examples-guards/cache-dit/scripts/hub_cache.py b/examples-guards/cache-dit/scripts/hub_cache.py new file mode 100644 index 000000000..abdd04e6b --- /dev/null +++ b/examples-guards/cache-dit/scripts/hub_cache.py @@ -0,0 +1,386 @@ +#!/usr/bin/env python3 +"""Resolve cache-dit's overlay model paths through the shared HF hub cache. + +setup_example.sh passes the env var names its profile needs, for example +`hub_cache.py FLUX_MODEL_PATH`. For each one this appends +VAR= to $GITHUB_ENV for the manifest's overlay_args: + +- a cached snapshot that holds every required file and no corrupt + safetensors is used as is, without network access; +- otherwise the asset is downloaded, HuggingFace first and ModelScope + as the fallback, and published into the cache, so the next run hits + the cache. + +FLUX.1-dev is an HF gated repo (gated: auto) and this repo carries no +HF_TOKEN, so the HuggingFace leg answers 401 and the ModelScope mirror +delivers it; both publish the same diffusers layout, so either source +lands a usable snapshot. The fetch is trimmed by allow_patterns to the +component tree FluxPipeline.from_pretrained reads: the 23.8 GB +single-file flux1-dev.safetensors and the 0.3 GB ae.safetensors +duplicate components that are loaded from transformer/ and vae/, so the +payload drops from 57.9 GB to 33.7 GB and the slowest file leaves the +queue. + +The cache root is HF_HOME (default ~/.cache/huggingface), which on the +self-hosted runners lives on the runner pool's persistent NFS volume, +shared by every concurrent matrix leg on every pod. The hub's own +.locks do not exclude across pods there: in run 37764199490 two legs +fetching SD v1.5 appended to one blob and left it larger than the whole +file. So nothing is downloaded into the cache in place. A job +downloads into its own staging dir on the same volume, checks it, and +renames the files into the snapshot, which readers then see whole or +not at all. A mkdir lock per asset lets one leg download while the +others wait for it; correctness does not depend on the lock. +""" +from __future__ import annotations + +import os +import shutil +import sys +import tempfile +import threading +import time +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path + +from huggingface_hub import HfApi, snapshot_download +from modelscope import snapshot_download as ms_snapshot_download +from safetensors import SafetensorError, safe_open + +HUB_ROOT = Path(os.environ.get( + 'HF_HOME', os.path.expanduser('~/.cache/huggingface'))) / 'hub' +# Siblings of HUB_ROOT on the same volume, so a rename from staging into +# the snapshot is atomic, kept out of the hub's own cache layout. +STAGING_DIR = 'guard-staging' +LOCKS_DIR = 'guard-locks' +# ModelScope's own snapshot bookkeeping, not repository content. +MODELSCOPE_METADATA = {'.mdl', '.msc', '.mv'} +HEARTBEAT_SECONDS = 30 +# A lock this long without a heartbeat belongs to a killed job. +LOCK_STALE_SECONDS = 300 +LOCK_POLL_SECONDS = 10 +# Past this, download alongside the holder rather than spend the job's +# timeout waiting. +LOCK_WAIT_SECONDS = 900 +# Longer than any job lives: the job that made it was killed. +STAGING_STALE_SECONDS = 24 * 3600 + + +@dataclass(frozen=True) +class Asset: + hf_id: str + ms_id: str + # Files to fetch; the single-file weights that duplicate components + # are deliberately absent. + patterns: tuple[str, ...] + # Globs that must each match a snapshot file for a cache hit, so an + # interrupted download is refetched instead of used. + required: tuple[str, ...] + + +# Only the diffusers component tree the pipeline loads: the component +# configs and shard indexes, both text encoders, the transformer shards, +# the VAE and both tokenizers. Excluded on purpose (both the HF repo and +# the ModelScope mirror carry them at the repository root, and +# from_pretrained never reads them): flux1-dev.safetensors 23.803 GB, +# ae.safetensors 0.335 GB, dev_grid.jpg, README.md, LICENSE.md. +FLUX_PATTERNS = ( + 'model_index.json', '*/*.json', '*/model.safetensors', + '*/model-*.safetensors', '*/diffusion_pytorch_model*.safetensors', + 'tokenizer/*', 'tokenizer_2/*') +# The shard indexes stand in for their whole shard set: a snapshot +# missing a shard is also missing the index that lists it. +FLUX_REQUIRED = ( + 'model_index.json', + 'transformer/diffusion_pytorch_model.safetensors.index.json', + 'text_encoder_2/model.safetensors.index.json', + 'text_encoder/model.safetensors', + 'vae/diffusion_pytorch_model.safetensors') + +ASSETS = { + 'FLUX_MODEL_PATH': Asset( + 'black-forest-labs/FLUX.1-dev', 'black-forest-labs/FLUX.1-dev', + FLUX_PATTERNS, FLUX_REQUIRED), +} + + +def repo_dir(asset: Asset) -> Path: + return HUB_ROOT / f"models--{asset.hf_id.replace('/', '--')}" + + +def safetensors_header_ok(path: Path) -> bool | None: + """True if the header parses, False if the file is readable but + malformed, None if it cannot be read at all. + + A read failure on the shared NFS volume is not evidence of + corruption, so None must not get a shard deleted. + """ + try: + with safe_open(str(path), framework='numpy') as f: + list(f.keys()) + except SafetensorError: + return False + except Exception: # noqa: BLE001 - OSError and the like: cannot judge + return None + return True + + +def corrupt_shards(root: Path) -> list[Path]: + """Safetensors under root whose header or size is inconsistent.""" + return [p for p in sorted(root.rglob('*.safetensors')) + if safetensors_header_ok(p) is False] + + +def cached_snapshot(asset: Asset) -> Path | None: + """The snapshot refs/main points at, if it holds every required file + and none of its safetensors is corrupt.""" + refs = repo_dir(asset) / 'refs' / 'main' + if not refs.is_file(): + return None + snap = repo_dir(asset) / 'snapshots' / refs.read_text().strip() + if not snap.is_dir(): + return None + if not all(any(p.is_file() for p in snap.glob(g)) for g in asset.required): + return None + corrupt = corrupt_shards(snap) + if corrupt: + print(f'[corrupt] {asset.hf_id}: ' + f'{[p.relative_to(snap).as_posix() for p in corrupt]}', flush=True) + return None + return snap + + +def purge_corrupt(asset: Asset) -> None: + """Delete corrupt shards and the blobs behind them. + + snapshot_download copies a blob the cache already holds into the + download target without rechecking it, so a damaged blob is handed + forward on every retry until it is gone. + """ + snapshots = repo_dir(asset) / 'snapshots' + if not snapshots.is_dir(): + return + for snap in sorted(snapshots.iterdir()): + for shard in corrupt_shards(snap): + blob = Path(os.path.realpath(shard)) if shard.is_symlink() else None + print(f'[purge] {asset.hf_id}: {shard.relative_to(snap).as_posix()}', + flush=True) + shard.unlink(missing_ok=True) + if blob is not None: + blob.unlink(missing_ok=True) + + +def revision(asset: Asset) -> str: + """The hub commit to file the snapshot under.""" + try: + return HfApi().model_info(asset.hf_id).sha + except Exception: # noqa: BLE001 - HuggingFace may be the reason we fall back + refs = repo_dir(asset) / 'refs' / 'main' + return refs.read_text().strip() if refs.is_file() else 'modelscope' + + +def from_huggingface(asset: Asset, dest: Path) -> Path: + snapshot_download( + repo_id=asset.hf_id, allow_patterns=list(asset.patterns), + local_dir=str(dest)) + return dest + + +def from_modelscope(asset: Asset, dest: Path) -> Path: + # snapshot_download does not create its cache root: it then fails + # mid-transfer. + dest.mkdir(parents=True, exist_ok=True) + return Path(ms_snapshot_download( + asset.ms_id, cache_dir=str(dest), + allow_patterns=list(asset.patterns))) + + +def checked(src: Path) -> Path: + corrupt = corrupt_shards(src) + if corrupt: + raise RuntimeError('corrupt safetensors downloaded: ' + f'{[p.relative_to(src).as_posix() for p in corrupt]}') + return src + + +def intact(dest: Path, staged: Path) -> bool: + """Whether dest already holds staged's content, by size and header.""" + try: + if dest.stat().st_size != staged.stat().st_size: + return False + except FileNotFoundError: + return False + return dest.suffix != '.safetensors' or safetensors_header_ok(dest) is not False + + +def publish(asset: Asset, src: Path, sha: str) -> Path: + """Rename src's files into snapshots/, then point refs/main at it. + + Intact files are left in place: another job may be reading them. + """ + snap = repo_dir(asset) / 'snapshots' / sha + for item in sorted(src.rglob('*')): + rel = item.relative_to(src) + # .cache is huggingface_hub's local_dir download bookkeeping. + if (not item.is_file() or rel.parts[0] == '.cache' + or rel.as_posix() in MODELSCOPE_METADATA): + continue + dest = snap / rel + if intact(dest, item): + continue + dest.parent.mkdir(parents=True, exist_ok=True) + os.replace(item, dest) + refs = repo_dir(asset) / 'refs' + refs.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + 'w', dir=refs, prefix='.main-', delete=False) as fh: + # No trailing newline: the hub compares this string with the + # snapshot folder name without stripping it. + fh.write(sha) + os.replace(fh.name, refs / 'main') + return snap + + +def fetch(asset: Asset) -> Path: + """Download into a staging dir of this job, check it, publish it.""" + sha = revision(asset) + staging_root = HUB_ROOT.parent / STAGING_DIR + staging_root.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp( + prefix=f'{repo_dir(asset).name}-', dir=staging_root)) + try: + print(f'[download] {asset.hf_id} <- HuggingFace', flush=True) + try: + src = checked(from_huggingface(asset, staging / 'hf')) + except Exception as exc: # noqa: BLE001 - falls back below + print(f'[fallback] {asset.hf_id}: HuggingFace failed ' + f'({type(exc).__name__}: {exc}); trying ModelScope {asset.ms_id}', + flush=True) + src = checked(from_modelscope(asset, staging / 'ms')) + snap = publish(asset, src, sha) + finally: + shutil.rmtree(staging, ignore_errors=True) + if cached_snapshot(asset) != snap: + raise RuntimeError( + f'{asset.hf_id}: downloaded snapshot {snap} lacks {asset.required}') + print(f'[cached] {asset.hf_id} -> {snap}', flush=True) + return snap + + +def heartbeat(lock: Path, stop: threading.Event) -> None: + while not stop.wait(HEARTBEAT_SECONDS): + try: + os.utime(lock) + except OSError: + return + + +def break_if_stale(lock: Path) -> None: + try: + idle = time.time() - lock.stat().st_mtime + except FileNotFoundError: + return + if idle < LOCK_STALE_SECONDS: + return + print(f'[lock] {lock.name}: no heartbeat for {idle:.0f}s, breaking it', + flush=True) + # Rename first: of several waiters breaking it, only one succeeds. + stale = lock.with_name(f'{lock.name}.stale-{uuid.uuid4().hex}') + try: + os.rename(lock, stale) + except FileNotFoundError: + return + shutil.rmtree(stale, ignore_errors=True) + + +@contextmanager +def download_lock(asset: Asset) -> Iterator[None]: + """Hold the asset's lock for the body; after LOCK_WAIT_SECONDS of + waiting, run the body without it.""" + locks = HUB_ROOT.parent / LOCKS_DIR + locks.mkdir(parents=True, exist_ok=True) + lock = locks / repo_dir(asset).name + deadline = time.monotonic() + LOCK_WAIT_SECONDS + waiting = False + while True: + try: + lock.mkdir() + break + except FileExistsError: + pass + if time.monotonic() >= deadline: + print(f'[lock] {asset.hf_id}: still held after {LOCK_WAIT_SECONDS}s, ' + 'downloading without it', flush=True) + yield + return + if not waiting: + print(f'[lock] {asset.hf_id}: another job is downloading it, waiting', + flush=True) + waiting = True + break_if_stale(lock) + time.sleep(LOCK_POLL_SECONDS) + stop = threading.Event() + beat = threading.Thread(target=heartbeat, args=(lock, stop), daemon=True) + beat.start() + try: + yield + finally: + stop.set() + beat.join() + shutil.rmtree(lock, ignore_errors=True) + + +def resolve(var: str) -> Path: + asset = ASSETS[var] + snap = cached_snapshot(asset) + if snap is None: + with download_lock(asset): + # The previous holder may have published it meanwhile. + snap = cached_snapshot(asset) + if snap is None: + purge_corrupt(asset) + return fetch(asset) + print(f'[cache] {asset.hf_id} -> {snap}', flush=True) + return snap + + +def sweep_staging() -> None: + """Remove staging dirs that killed jobs left behind.""" + staging_root = HUB_ROOT.parent / STAGING_DIR + if not staging_root.is_dir(): + return + cutoff = time.time() - STAGING_STALE_SECONDS + for entry in staging_root.iterdir(): + try: + if entry.stat().st_mtime < cutoff: + shutil.rmtree(entry, ignore_errors=True) + except OSError: + continue + + +def main(argv: list[str]) -> int: + unknown = [var for var in argv if var not in ASSETS] + if not argv or unknown: + print(f'usage: hub_cache.py VAR... (known: {sorted(ASSETS)}; ' + f'unknown: {unknown})', file=sys.stderr) + return 2 + sweep_staging() + for var in argv: + try: + snap = resolve(var) + except Exception as exc: # noqa: BLE001 - reported as guard noise + print(f'{ASSETS[var].hf_id}: not in the shared cache and neither ' + f'HuggingFace nor ModelScope delivered it ({exc}) - guard ' + 'noise, not an example failure', file=sys.stderr) + return 1 + with open(os.environ['GITHUB_ENV'], 'a', encoding='utf-8') as fh: + fh.write(f'{var}={snap}\n') + print(f'{var}={snap}', flush=True) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main(sys.argv[1:])) diff --git a/examples-guards/cache-dit/scripts/run_example.sh b/examples-guards/cache-dit/scripts/run_example.sh new file mode 100644 index 000000000..af376652b --- /dev/null +++ b/examples-guards/cache-dit/scripts/run_example.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +# Run one cache-dit example from a CI working copy of the target tree. +# Overlay CLI args come from OVERLAY_ARGS (JSON array, as emitted by +# toJSON(matrix.example.overlay_args)); cache-dit examples parse sys.argv +# themselves (argparse), so args are passed straight through. +# EXEC overrides the launcher command when set (e.g. torchrun). +set -euo pipefail + +if [[ $# -lt 1 ]]; then + echo "usage: $0 " >&2 + exit 2 +fi + +EXAMPLE_REL="$1" +TARGET_ROOT="${TARGET_ROOT:?TARGET_ROOT is required}" +CI_OUTPUT_DIR="${CI_OUTPUT_DIR:?CI_OUTPUT_DIR is required}" + +EXAMPLE_PATH="$TARGET_ROOT/$EXAMPLE_REL" + +if [[ ! -f "$EXAMPLE_PATH" ]]; then + echo "example not found: $EXAMPLE_PATH" >&2 + exit 1 +fi + +mkdir -p "$CI_OUTPUT_DIR" + +if command -v python3 >/dev/null 2>&1; then + PYTHON=python3 +else + PYTHON=python +fi + +expand_overlay() { + "$PYTHON" - <<'PY' +import json +import os +import shlex + +raw = os.environ.get('OVERLAY_ARGS', '').strip() +if not raw or raw in ('null', '""'): + raise SystemExit(0) +try: + items = json.loads(raw) +except json.JSONDecodeError as exc: + raise SystemExit(f'OVERLAY_ARGS is not valid JSON: {exc}') from exc +if items in (None, ''): + raise SystemExit(0) +if not isinstance(items, list): + raise SystemExit( + f'OVERLAY_ARGS must be a JSON array, got {type(items).__name__}') +tokens = [] +for item in items: + if not isinstance(item, str): + raise SystemExit( + f'OVERLAY_ARGS items must be strings, got {type(item).__name__}') + tokens.extend(shlex.split(os.path.expandvars(item), posix=True)) +print(' '.join(shlex.quote(token) for token in tokens)) +PY +} + +eval "EXTRA_ARGS=( $(expand_overlay) )" + +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0}" +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +if [ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]; then + source /usr/local/Ascend/ascend-toolkit/set_env.sh +fi + +# Run from the target repo root so relative paths (generated image +# outputs, example data) resolve consistently. +cd "$TARGET_ROOT" + +echo "Running example: $EXAMPLE_PATH" +echo "With overlay args:" +if ((${#EXTRA_ARGS[@]})); then + printf ' %q\n' "${EXTRA_ARGS[@]}" +fi +echo "Using NPU devices: ${ASCEND_RT_VISIBLE_DEVICES}" + +if [[ -n "${EXEC:-}" ]]; then + eval "${EXEC} ${EXTRA_ARGS[*]}" +else + "$PYTHON" "$EXAMPLE_PATH" "${EXTRA_ARGS[@]}" +fi diff --git a/examples-guards/cache-dit/scripts/setup_example.sh b/examples-guards/cache-dit/scripts/setup_example.sh new file mode 100644 index 000000000..90f6b6602 --- /dev/null +++ b/examples-guards/cache-dit/scripts/setup_example.sh @@ -0,0 +1,112 @@ +#!/bin/bash +# setup_example.sh - Prepare environment for cache-dit example based on profile +# Called from the examples engine's run-example job. +# Positional argument: $1 is the manifest profile (e.g., flux). + +set -euo pipefail + +if [[ $# -lt 1 ]]; then + echo "usage: $0 " >&2 + exit 2 +fi + +PROFILE="$1" + +# wan2.2 dropped: its upstream scripts moved to unsupported (broken +# `from utils import ...` imports, see the manifest); keeping the +# profile would silently run the flux setup for it. +supported_profiles="flux" + +if echo "$supported_profiles" | grep -qw "$PROFILE"; then + echo "Profile '$PROFILE' is supported" +else + echo "Unsupported profile: $PROFILE" + echo "Supported profiles: $supported_profiles" + exit 1 +fi + +TARGET_ROOT="${TARGET_ROOT:?TARGET_ROOT is required}" +GITHUB_WORKSPACE="${GITHUB_WORKSPACE:?GITHUB_WORKSPACE is required}" +GITHUB_ENV="${GITHUB_ENV:?GITHUB_ENV is required}" + +HERE=$(cd "$(dirname "$0")" && pwd) + +# Install cache-dit and dependencies +echo "Installing cache-dit and dependencies..." + +# CANN 9.1.0 base image ships without torch/torch_npu; install the +# official pairing (same torch line as the diffusers project). +ASCEND_PIP_INDEX=https://repo.huaweicloud.com/ascend/repos/pypi + +pip_ascend() { + python3 -m pip install --extra-index-url "$ASCEND_PIP_INDEX" "$@" +} + +ensure_torch_stack() { + if python3 -c " +import torch, torch_npu +raise SystemExit( + 0 if torch.__version__.startswith('2.9.0') + and torch_npu.__version__.startswith('2.9.0') else 1) +"; then + echo "reusing image torch stack ($(python3 -c 'import torch; print(torch.__version__)'))" + return + fi + echo "installing torch==2.9.0 torch_npu==2.9.0.post2" + pip_ascend torch==2.9.0 torch_npu==2.9.0.post2 +} + +ensure_torch_stack +pip3 install --no-deps torchvision==0.24.0 # torch 2.9.0 pairing + +pip3 install -U cache-dit +pip3 install einops sentencepiece accelerate + +# Install diffusers for parallel support +pip3 install -U diffusers # 要求 >= 0.36.0(PyPI latest,避免走 github 代理) + +# hub_cache.py's fallback source, installed here because hub_cache.py +# imports it at module load (after diffusers so the latter's +# huggingface_hub pin stands). FLUX.1-dev is an HF gated repo +# (gated: auto) and this repo carries no HF_TOKEN, so the HuggingFace +# leg answers 401 and the ModelScope mirror delivers the snapshot. +# Pinned to the version verified in this image (1.40.x fails to import +# next to huggingface_hub<1.0). +pip3 install "modelscope==1.37.0" + +# Keep the engine-declared NPU visibility (matrix npu_devices exported +# via GITHUB_ENV); default to 0 only when running outside the engine. +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0}" +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True + +# Source CANN environment if available +if [ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]; then + source /usr/local/Ascend/ascend-toolkit/set_env.sh +fi + +# Resolve the example's model into GITHUB_ENV for the manifest's +# overlay_args expansion (${FLUX_MODEL_PATH} in the run-example step). +# scripts/hub_cache.py owns the whole path: a snapshot in the shared +# runner cache that holds every required file is used as is, with zero +# network; a miss is downloaded (HuggingFace first, ModelScope +# fallback) into this job's private staging dir on the same persistent +# volume and published into the cache, so every later run - on any pod +# of the pool - hits it. The fetch is trimmed by allow_patterns to the +# diffusers component tree (57.89 GB -> 33.74 GB). Same helper shape as +# the peft project. +resolve_hub_paths() { + python3 "$HERE/hub_cache.py" "$@" +} + +setup_flux() { + echo "profile=flux: resolving FLUX.1-dev model" + # HF id is the gated black-forest-labs repo; hub_cache.py trims the + # download to the diffusers component tree FluxPipeline.from_pretrained + # reads and falls back to the ModelScope mirror (same layout) when the + # HuggingFace leg 401s for lack of an HF_TOKEN. + resolve_hub_paths FLUX_MODEL_PATH +} + +setup_flux + +echo "Environment setup complete for profile: $PROFILE"