Skip to content

FIX: Load HuggingFace models without blocking the event loop #2211

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
romanlutz merged 3 commits into from
Jul 17, 2026
Merged

FIX: Load HuggingFace models without blocking the event loop #2211

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
14e35b9
FIX: Load HuggingFace models without blocking the event loop
Copilot Jul 16, 2026
18cf36e
Merge remote-tracking branch 'origin/main' into romanlutz-validate-c-…
Copilot Jul 17, 2026
aabb189
MAINT: Remove unused custom HF download helpers
Copilot Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 1 addition & 8 deletions doc/code/targets/11_message_normalizer.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,13 +320,6 @@
"No HuggingFace token provided. Gated models may fail to load without authentication.\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
Expand Down Expand Up @@ -535,7 +528,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.12"
"version": "3.14.4"
}
},
"nbformat": 4,
Expand Down
74 changes: 61 additions & 13 deletions doc/code/targets/use_huggingface_chat_target.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,14 @@
"id": "1",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"./git/copilot-worktrees/PyRIT/romanlutz-cautious-meme/.venv/Lib/site-packages/confusables/__init__.py:46: SyntaxWarning: \"\\*\" is an invalid escape sequence. Such sequences will not work in the future. Did you mean \"\\\\*\"? A raw string is also an option.\n",
" space_regex = \"[\\*_~|`\\-\\.]*\" if include_character_padding else ''\n"
]
},
{
"name": "stdout",
"output_type": "stream",
Expand All @@ -54,14 +62,48 @@
"name": "stdout",
"output_type": "stream",
"text": [
"No new upgrade operations detected.\n",
"[pyrit:alembic] No new upgrade operations detected.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running model: Qwen/Qwen2-0.5B-Instruct\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "f637f0a22e814d5390fcfab383ecff0b",
"model_id": "2cb1ed2b5d6c4fa98456d1c217564010",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Downloading (incomplete total...): 0.00B [00:00, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "76523dac8fe04b87921ad444853385c3",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Fetching 10 files: 0%| | 0/10 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "650a316e795b4e84abdc46b06bc79cc1",
"version_major": 2,
"version_minor": 0
},
Expand All @@ -76,7 +118,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"Average response time for Qwen/Qwen2-0.5B-Instruct: 32.08 seconds\n",
"Average response time for Qwen/Qwen2-0.5B-Instruct: 7.43 seconds\n",
"\n",
"\n",
"\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n",
Expand All @@ -88,11 +130,11 @@
"\u001b[1m 📋 Basic Information\u001b[0m\n",
"\u001b[36m • Objective: What is 3*3? Give me the solution.\u001b[0m\n",
"\u001b[36m • Attack Type: PromptSendingAttack\u001b[0m\n",
"\u001b[36m • Conversation ID: 434c7b92-3756-4f64-8c9d-f61a181e1ed8\u001b[0m\n",
"\u001b[36m • Conversation ID: 3e01b421-e41b-4c47-a51b-ea192f7cdb24\u001b[0m\n",
"\n",
"\u001b[1m ⚡ Execution Metrics\u001b[0m\n",
"\u001b[32m • Turns Executed: 1\u001b[0m\n",
"\u001b[32m • Execution Time: 1m 2s\u001b[0m\n",
"\u001b[32m • Execution Time: 12.20s\u001b[0m\n",
"\n",
"\u001b[1m 🎯 Outcome\u001b[0m\n",
"\u001b[33m • Status: ❓ UNDETERMINED\u001b[0m\n",
Expand All @@ -109,12 +151,14 @@
"\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
"\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n",
"\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
"\u001b[33m The solution, after expanding \\(3 \\cdot 3\\), is that it equals 9.\u001b[0m\n",
"\u001b[33m The answer to \"What is 3 * 3?\" is 9.\u001b[0m\n",
"\u001b[33m \u001b[0m\n",
"\u001b[33m This is because the product of two numbers, where the first number (\u001b[0m\n",
"\n",
"\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
"\n",
"\u001b[2m\u001b[37m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
"\u001b[2m\u001b[37m Report generated at: 2026-05-21 23:13:40 UTC \u001b[0m\n",
"\u001b[2m\u001b[37m Report generated at: 2026-07-16 05:18:39 UTC \u001b[0m\n",
"\n",
"\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n",
"\u001b[1m\u001b[33m ❓ ATTACK RESULT: UNDETERMINED ❓ \u001b[0m\n",
Expand All @@ -125,11 +169,11 @@
"\u001b[1m 📋 Basic Information\u001b[0m\n",
"\u001b[36m • Objective: What is 4*4? Give me the solution.\u001b[0m\n",
"\u001b[36m • Attack Type: PromptSendingAttack\u001b[0m\n",
"\u001b[36m • Conversation ID: 3b782d88-adb6-47b1-bb3f-699ec338d66b\u001b[0m\n",
"\u001b[36m • Conversation ID: 76b845b6-0e72-4a80-a793-09652c4d7405\u001b[0m\n",
"\n",
"\u001b[1m ⚡ Execution Metrics\u001b[0m\n",
"\u001b[32m • Turns Executed: 1\u001b[0m\n",
"\u001b[32m • Execution Time: 1.40s\u001b[0m\n",
"\u001b[32m • Execution Time: 2.63s\u001b[0m\n",
"\n",
"\u001b[1m 🎯 Outcome\u001b[0m\n",
"\u001b[33m • Status: ❓ UNDETERMINED\u001b[0m\n",
Expand All @@ -146,13 +190,17 @@
"\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
"\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n",
"\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
"\u001b[33m The answer to \"4 * 4\" is 16.\u001b[0m\n",
"\u001b[33m The result of multiplying 4 by itself four times is:\u001b[0m\n",
"\u001b[33m 256.\u001b[0m\n",
"\u001b[33m \u001b[0m\n",
"\u001b[33m Here's why:\u001b[0m\n",
"\u001b[33m First, we multiply 4 and\u001b[0m\n",
"\n",
"\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
"\n",
"\u001b[2m\u001b[37m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
"\u001b[2m\u001b[37m Report generated at: 2026-05-21 23:13:40 UTC \u001b[0m\n",
"Qwen/Qwen2-0.5B-Instruct: 32.08 seconds\n"
"\u001b[2m\u001b[37m Report generated at: 2026-07-16 05:18:39 UTC \u001b[0m\n",
"Qwen/Qwen2-0.5B-Instruct: 7.43 seconds\n"
]
}
],
Expand Down Expand Up @@ -236,7 +284,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.12"
"version": "3.14.4"
}
},
"nbformat": 4,
Expand Down
121 changes: 10 additions & 111 deletions pyrit/common/download_hf_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,125 +2,24 @@
# Licensed under the MIT license.

import asyncio
import logging
from pathlib import Path

import aiofiles
import httpx
from huggingface_hub import HfApi

logger = logging.getLogger(__name__)


def get_available_files(model_id: str, token: str) -> list[str]:
"""
Fetch available files for a model from the Hugging Face repository.

Returns:
List of available file names.

Raises:
ValueError: If no files are found for the model.
"""
api = HfApi()
try:
model_info = api.model_info(model_id, token=token)
available_files = [file.rfilename for file in (model_info.siblings or [])]

# Perform simple validation: raise a ValueError if no files are available
if not len(available_files):
raise ValueError(f"No available files found for the model: {model_id}")

return available_files
except Exception as e:
logger.info(f"Error fetching model files for {model_id}: {e}")
return []
from huggingface_hub import snapshot_download


async def download_specific_files_async(
model_id: str, file_patterns: list[str] | None, token: str, cache_dir: Path
) -> None:
"""
Download specific files from a Hugging Face model repository.
Download a Hugging Face model snapshot without blocking the event loop.

If file_patterns is None, downloads all files.
"""
cache_dir.mkdir(parents=True, exist_ok=True)

available_files = get_available_files(model_id, token)
# If no file patterns are provided, download all available files
if file_patterns is None:
files_to_download = available_files
logger.info(f"Downloading all files for model {model_id}.")
else:
# Filter files based on the patterns provided
files_to_download = [file for file in available_files if any(pattern in file for pattern in file_patterns)]
if not files_to_download:
logger.info(f"No files matched the patterns provided for model {model_id}.")
return

# Generate download URLs directly
base_url = f"https://huggingface.co/{model_id}/resolve/main/"
urls = [base_url + file for file in files_to_download]

# Download the files
await download_files_async(urls, token, cache_dir)
Comment thread
romanlutz marked this conversation as resolved.


async def download_chunk_async(
url: str, headers: dict[str, str], start: int, end: int, client: httpx.AsyncClient
) -> bytes:
"""
Download a chunk of the file with a specified byte range.

Returns:
The content of the downloaded chunk.
"""
range_header = {"Range": f"bytes={start}-{end}", **headers}
response = await client.get(url, headers=range_header)
response.raise_for_status()
return response.content


async def download_file_async(url: str, token: str, download_dir: Path, num_splits: int) -> None:
"""Download a file in multiple segments (splits) using byte-range requests."""
headers = {"Authorization": f"Bearer {token}"}
async with httpx.AsyncClient(follow_redirects=True) as client:
# Get the file size to determine chunk size
response = await client.head(url, headers=headers)
response.raise_for_status()
file_size = int(response.headers["Content-Length"])
chunk_size = file_size // num_splits

# Prepare tasks for each chunk
tasks = []
file_name = url.split("/")[-1]
file_path = Path(download_dir, file_name)

for i in range(num_splits):
start = i * chunk_size
end = start + chunk_size - 1 if i < num_splits - 1 else file_size - 1
tasks.append(download_chunk_async(url, headers, start, end, client))

# Download all chunks concurrently
chunks = await asyncio.gather(*tasks)

# Write chunks to the file in order
async with aiofiles.open(file_path, "wb") as f:
for chunk in chunks:
await f.write(chunk)
logger.info(f"Downloaded {file_name} to {file_path}")


async def download_files_async(
urls: list[str], token: str, download_dir: Path, num_splits: int = 3, parallel_downloads: int = 4
) -> None:
"""Download multiple files with parallel downloads and segmented downloading."""
# Limit the number of parallel downloads
semaphore = asyncio.Semaphore(parallel_downloads)

async def download_with_limit_async(url: str) -> None:
async with semaphore:
await download_file_async(url, token, download_dir, num_splits)

# Run downloads concurrently, but limit to parallel_downloads at a time
await asyncio.gather(*(download_with_limit_async(url) for url in urls))
await asyncio.to_thread(
snapshot_download,
repo_id=model_id,
allow_patterns=file_patterns,
token=token,
local_dir=cache_dir,
)
14 changes: 3 additions & 11 deletions pyrit/prompt_target/hugging_face/hugging_face_chat_target.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,20 +295,12 @@ async def load_model_and_tokenizer_async(self) -> None:
Path(cache_dir),
)

# Load the tokenizer and model from the specified directory
# Load the tokenizer and model from the downloaded local snapshot.
logger.info(f"Loading model {self.model_id} from cache path: {cache_dir}...")
self.tokenizer = AutoTokenizer.from_pretrained(
self.model_id or "", cache_dir=cache_dir, trust_remote_code=self.trust_remote_code
)
self.model = AutoModelForCausalLM.from_pretrained(
self.model_id or "",
cache_dir=cache_dir,
trust_remote_code=self.trust_remote_code,
**optional_model_kwargs,
)
self._load_from_path(str(cache_dir), **optional_model_kwargs)

# Move the model to the correct device
self.model = cast("Any", self.model).to(self.device)
self.model = self.model.to(self.device)

# Debug prints to check types
logger.info(f"Model loaded: {type(self.model)}")
Expand Down
Loading