From f253f43c9d39b943c7e31c1ebb7a1376c362c4db Mon Sep 17 00:00:00 2001 From: Prozac614 Date: Tue, 30 Dec 2025 11:39:27 +0800 Subject: [PATCH] [diffusion] CI: fix LoRA downloading issues and respect offline flag (#15813) --- python/pyproject_other.toml | 2 +- .../runtime/utils/hf_diffusers_utils.py | 154 +++++++++++++++--- .../sglang/multimodal_gen/test/run_suite.py | 7 +- scripts/ci/ci_install_dependency.sh | 1 + 4 files changed, 136 insertions(+), 28 deletions(-) diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 94bee0042..70780e714 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -105,7 +105,7 @@ test = [ ] diffusion = [ "diffusers @ git+https://github.com/huggingface/diffusers.git@6290fdfda40610ce7b99920146853614ba529c6e", - "opencv-python==4.10.0.84", + "opencv-python-headless==4.10.0.84", "imageio==2.36.0", "imageio-ffmpeg==0.5.1", "PyYAML==6.0.1", diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index f572cc54e..63e6d6e73 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -21,6 +21,7 @@ import contextlib import json import os +import time from functools import reduce from pathlib import Path from typing import Any, Optional, cast @@ -29,6 +30,13 @@ from diffusers.loaders.lora_base import ( _best_guess_weight_name, # watch out for potetential removal from diffusers ) from huggingface_hub import snapshot_download +from huggingface_hub.errors import ( + LocalEntryNotFoundError, + RepositoryNotFoundError, + RevisionNotFoundError, +) +from requests.exceptions import ConnectionError as RequestsConnectionError +from requests.exceptions import RequestException from transformers import AutoConfig, PretrainedConfig from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES @@ -208,15 +216,20 @@ def maybe_download_lora( Returns: Local path to the model """ + allow_patterns = ["*.json", "*.safetensors", "*.bin"] - local_path = maybe_download_model(model_name_or_path, local_dir, download) + local_path = maybe_download_model( + model_name_or_path, + local_dir, + download, + is_lora=True, + allow_patterns=allow_patterns, + ) # return directly if local_path is a file if os.path.isfile(local_path): return local_path - weight_name = _best_guess_weight_name( - model_name_or_path, file_extension=".safetensors" - ) + weight_name = _best_guess_weight_name(local_path, file_extension=".safetensors") return os.path.join(local_path, weight_name) @@ -349,7 +362,11 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: def maybe_download_model( - model_name_or_path: str, local_dir: str | None = None, download: bool = True + model_name_or_path: str, + local_dir: str | None = None, + download: bool = True, + is_lora: bool = False, + allow_patterns: list[str] | None = None, ) -> str: """ Check if the model path is a Hugging Face Hub model ID and download it if needed. @@ -358,6 +375,7 @@ def maybe_download_model( model_name_or_path: Local path or Hugging Face Hub model ID local_dir: Local directory to save the model download: Whether to download the model from Hugging Face Hub + is_lora: If True, skip model completeness verification (LoRA models don't have transformer/vae directories) Returns: Local path to the model @@ -374,9 +392,9 @@ def maybe_download_model( and os.path.exists(vae_dir) ) - # If the path exists locally, verify it's complete + # 1. Local path check: if path exists locally, verify it's complete (skip for LoRA) if os.path.exists(model_name_or_path): - if _verify_model_complete(model_name_or_path): + if is_lora or _verify_model_complete(model_name_or_path): logger.info("Model already exists locally and is complete") return model_name_or_path else: @@ -386,34 +404,118 @@ def maybe_download_model( model_name_or_path, ) - # Otherwise, assume it's a HF Hub model ID and try to download it + # 2. Cache-first strategy (Fast Path) + # Try to read from HF cache without network access try: logger.info( - "Downloading model snapshot from HF Hub for %s...", model_name_or_path + "Checking for cached model in HF Hub cache for %s...", model_name_or_path ) - with get_lock(model_name_or_path).acquire(poll_interval=2): - local_path = snapshot_download( - repo_id=model_name_or_path, - ignore_patterns=["*.onnx", "*.msgpack"], - local_dir=local_dir, + local_path = snapshot_download( + repo_id=model_name_or_path, + ignore_patterns=["*.onnx", "*.msgpack"], + local_dir=local_dir, + local_files_only=True, + resume_download=True, + max_workers=8, + etag_timeout=60, + ) + if is_lora or _verify_model_complete(local_path): + logger.info("Found complete model in cache at %s", local_path) + return str(local_path) + else: + # Model found in cache but incomplete + if not download: + raise ValueError( + f"Model {model_name_or_path} found in cache but is incomplete and download=False." + ) + logger.info( + "Model found in cache but incomplete, will download from HF Hub" ) - # Verify downloaded model is complete - if not _verify_model_complete(local_path): - logger.warning( - "Downloaded model at %s is incomplete, retrying with force_download=True", - local_path, + except LocalEntryNotFoundError: + if not download: + raise ValueError( + f"Model {model_name_or_path} not found in local cache and download=False." + ) + logger.info("Model not found in cache, will download from HF Hub") + except Exception as e: + logger.warning( + "Unexpected error while checking cache for %s: %s, will attempt download", + model_name_or_path, + e, + ) + if not download: + raise ValueError( + f"Error checking cache for {model_name_or_path} and download=False: {e}" + ) from e + + # 3. Download strategy (with retry mechanism) + MAX_RETRIES = 5 + for attempt in range(MAX_RETRIES): + try: + logger.info( + "Downloading model snapshot from HF Hub for %s (attempt %d/%d)...", + model_name_or_path, + attempt + 1, + MAX_RETRIES, ) with get_lock(model_name_or_path).acquire(poll_interval=2): local_path = snapshot_download( repo_id=model_name_or_path, ignore_patterns=["*.onnx", "*.msgpack"], + allow_patterns=allow_patterns, local_dir=local_dir, - force_download=True, + resume_download=True, + max_workers=8, + etag_timeout=120, ) - logger.info("Downloaded model to %s", local_path) - return str(local_path) - except Exception as e: - raise ValueError( - f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}" - ) from e + # Verify downloaded model is complete (skip for LoRA) + if not is_lora and not _verify_model_complete(local_path): + logger.warning( + "Downloaded model at %s is incomplete, retrying with force_download=True", + local_path, + ) + with get_lock(model_name_or_path).acquire(poll_interval=2): + local_path = snapshot_download( + repo_id=model_name_or_path, + ignore_patterns=["*.onnx", "*.msgpack"], + local_dir=local_dir, + resume_download=True, + max_workers=8, + etag_timeout=60, + force_download=True, + ) + if not _verify_model_complete(local_path): + raise ValueError( + f"Downloaded model at {local_path} is still incomplete after forced re-download. " + "The model repository may be missing required components (model_index.json, transformer/, or vae/)." + ) + + logger.info("Downloaded model to %s", local_path) + return str(local_path) + + except (RepositoryNotFoundError, RevisionNotFoundError) as e: + raise ValueError( + f"Model or revision not found at {model_name_or_path}. " + f"Please check the model ID or ensure you have access to the repository. Error: {e}" + ) from e + except (RequestException, RequestsConnectionError) as e: + if attempt == MAX_RETRIES - 1: + raise ValueError( + f"Could not find model at {model_name_or_path} and failed to download from HF Hub " + f"after {MAX_RETRIES} attempts due to network error: {e}" + ) from e + wait_time = 2**attempt + logger.warning( + "Download failed (attempt %d/%d) due to network error: %s. " + "Retrying in %d seconds...", + attempt + 1, + MAX_RETRIES, + e, + wait_time, + ) + time.sleep(wait_time) + except Exception as e: + raise ValueError( + f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}" + ) from e diff --git a/python/sglang/multimodal_gen/test/run_suite.py b/python/sglang/multimodal_gen/test/run_suite.py index 9f47799cd..1b990329a 100644 --- a/python/sglang/multimodal_gen/test/run_suite.py +++ b/python/sglang/multimodal_gen/test/run_suite.py @@ -198,7 +198,12 @@ def run_pytest(files, filter_expr=None): is_flaky_ci_assertion = "SafetensorError" in full_output - if not (is_perf_assertion or is_flaky_ci_assertion): + is_oom_error = ( + "out of memory" in full_output.lower() + or "oom killer" in full_output.lower() + ) + + if not (is_perf_assertion or is_flaky_ci_assertion or is_oom_error): return returncode logger.info(f"Max retry exceeded") diff --git a/scripts/ci/ci_install_dependency.sh b/scripts/ci/ci_install_dependency.sh index 0307cc2f0..78e1ab65f 100755 --- a/scripts/ci/ci_install_dependency.sh +++ b/scripts/ci/ci_install_dependency.sh @@ -96,6 +96,7 @@ fi # Clean up existing installations $PIP_UNINSTALL_CMD sgl-kernel sglang $PIP_UNINSTALL_SUFFIX || true $PIP_UNINSTALL_CMD flashinfer-python flashinfer-cubin flashinfer-jit-cache $PIP_UNINSTALL_SUFFIX || true +$PIP_UNINSTALL_CMD opencv-python opencv-python-headless $PIP_UNINSTALL_SUFFIX || true # Install the main package EXTRAS="dev"