From 14e35b957ff1f8b84b322883d080f61775be15d4 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:26:24 -0700 Subject: [PATCH 1/2] FIX: Load HuggingFace models without blocking the event loop Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 52d2518d-ebe5-4888-8189-6146b6503448 --- doc/code/targets/11_message_normalizer.ipynb | 9 +-- .../targets/use_huggingface_chat_target.ipynb | 74 +++++++++++++++---- pyrit/common/download_hf_model.py | 31 +++----- .../hugging_face/hugging_face_chat_target.py | 14 +--- tests/unit/common/test_hf_model_downloads.py | 15 +++- .../target/test_huggingface_chat_target.py | 9 ++- 6 files changed, 95 insertions(+), 57 deletions(-) diff --git a/doc/code/targets/11_message_normalizer.ipynb b/doc/code/targets/11_message_normalizer.ipynb index 93d2489313..5c7b38e0e5 100644 --- a/doc/code/targets/11_message_normalizer.ipynb +++ b/doc/code/targets/11_message_normalizer.ipynb @@ -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", @@ -535,7 +528,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.12" + "version": "3.14.4" } }, "nbformat": 4, diff --git a/doc/code/targets/use_huggingface_chat_target.ipynb b/doc/code/targets/use_huggingface_chat_target.ipynb index 08e329dc3b..02d34c4365 100644 --- a/doc/code/targets/use_huggingface_chat_target.ipynb +++ b/doc/code/targets/use_huggingface_chat_target.ipynb @@ -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", @@ -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 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) + await asyncio.to_thread( + snapshot_download, + repo_id=model_id, + allow_patterns=file_patterns, + token=token, + local_dir=cache_dir, + ) async def download_chunk_async( diff --git a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py index d333e234d8..ad98a94969 100644 --- a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py +++ b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py @@ -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)}") diff --git a/tests/unit/common/test_hf_model_downloads.py b/tests/unit/common/test_hf_model_downloads.py index ce5d733d07..82d23cc912 100644 --- a/tests/unit/common/test_hf_model_downloads.py +++ b/tests/unit/common/test_hf_model_downloads.py @@ -40,9 +40,20 @@ def setup_environment(): async def test_download_specific_files_async(setup_environment): """Test downloading specific files""" token = setup_environment # Get the token from the fixture + cache_dir = Path("model-cache") - with patch("os.makedirs"), patch("pyrit.common.download_hf_model.download_files_async"): - await download_specific_files_async(MODEL_ID, FILE_PATTERNS, token, Path("")) + with ( + patch("pathlib.Path.mkdir"), + patch("pyrit.common.download_hf_model.snapshot_download") as snapshot_download_mock, + ): + await download_specific_files_async(MODEL_ID, FILE_PATTERNS, token, cache_dir) + + snapshot_download_mock.assert_called_once_with( + repo_id=MODEL_ID, + allow_patterns=FILE_PATTERNS, + token=token, + local_dir=cache_dir, + ) async def test_download_files_async_dispatches_one_call_per_url(): diff --git a/tests/unit/prompt_target/target/test_huggingface_chat_target.py b/tests/unit/prompt_target/target/test_huggingface_chat_target.py index 3d302c91f3..12177d7c44 100644 --- a/tests/unit/prompt_target/target/test_huggingface_chat_target.py +++ b/tests/unit/prompt_target/target/test_huggingface_chat_target.py @@ -3,6 +3,8 @@ import json from asyncio import Task +from collections.abc import Coroutine +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -92,9 +94,12 @@ async def _await(): @pytest.fixture(autouse=True) def mock_create_task(): + def _close_coroutine(coroutine: Coroutine[Any, Any, None]) -> AwaitableTask: + coroutine.close() + return AwaitableTask(spec=Task) + with patch("asyncio.create_task") as mock_task: - # Return an AwaitableTask that can be awaited - mock_task.return_value = AwaitableTask(spec=Task) + mock_task.side_effect = _close_coroutine yield mock_task From aabb189545486c4805cf254d69f240e8a0aa8937 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:43:42 -0700 Subject: [PATCH 2/2] MAINT: Remove unused custom HF download helpers The switch to huggingface_hub.snapshot_download left get_available_files and the download_files_async/download_file_async/download_chunk_async chain with no callers. Remove them and their tests. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 18c92f59-1265-4426-b826-7c407c6e8c7d --- pyrit/common/download_hf_model.py | 92 +------------------- tests/unit/common/test_hf_model_downloads.py | 70 +-------------- 2 files changed, 3 insertions(+), 159 deletions(-) diff --git a/pyrit/common/download_hf_model.py b/pyrit/common/download_hf_model.py index e7fa49b730..5463105a00 100644 --- a/pyrit/common/download_hf_model.py +++ b/pyrit/common/download_hf_model.py @@ -2,39 +2,9 @@ # Licensed under the MIT license. import asyncio -import logging from pathlib import Path -import aiofiles -import httpx -from huggingface_hub import HfApi, snapshot_download - -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( @@ -53,63 +23,3 @@ async def download_specific_files_async( token=token, local_dir=cache_dir, ) - - -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)) diff --git a/tests/unit/common/test_hf_model_downloads.py b/tests/unit/common/test_hf_model_downloads.py index 82d23cc912..54abd49ad8 100644 --- a/tests/unit/common/test_hf_model_downloads.py +++ b/tests/unit/common/test_hf_model_downloads.py @@ -3,17 +3,11 @@ import os from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import patch import pytest -# Import functions to test from local application files -from pyrit.common.download_hf_model import ( - download_chunk_async, - download_file_async, - download_files_async, - download_specific_files_async, -) +from pyrit.common.download_hf_model import download_specific_files_async # Define constants for testing MODEL_ID = "microsoft/Phi-3-mini-4k-instruct" @@ -31,7 +25,6 @@ @pytest.fixture(scope="module") def setup_environment(): """Fixture to set up the environment for Hugging Face downloads.""" - # Check for Hugging Face token with patch.dict(os.environ, {"HUGGINGFACE_TOKEN": "mocked_token"}): token = os.getenv("HUGGINGFACE_TOKEN") yield token @@ -54,62 +47,3 @@ async def test_download_specific_files_async(setup_environment): token=token, local_dir=cache_dir, ) - - -async def test_download_files_async_dispatches_one_call_per_url(): - """Exercise the nested download_with_limit_async helper plus asyncio.gather.""" - seen_urls: list[str] = [] - - async def fake_file_async(url, token, download_dir, num_splits): - seen_urls.append(url) - - with patch("pyrit.common.download_hf_model.download_file_async", new=fake_file_async): - urls = ["https://example/a", "https://example/b", "https://example/c"] - await download_files_async(urls, "token", Path("/tmp"), num_splits=2, parallel_downloads=2) - - assert sorted(seen_urls) == sorted(urls) - - -async def test_download_file_async_schedules_one_chunk_per_split(tmp_path): - """Exercise the chunk-task assembly inside download_file_async.""" - num_splits = 3 - file_size = 30 - chunk_bytes = b"abcdefghij" - - head_response = MagicMock() - head_response.raise_for_status = MagicMock() - head_response.headers = {"Content-Length": str(file_size)} - - client_instance = MagicMock() - client_instance.head = AsyncMock(return_value=head_response) - - class _ClientCM: - async def __aenter__(self): - return client_instance - - async def __aexit__(self, exc_type, exc, tb): - return False - - with ( - patch("pyrit.common.download_hf_model.httpx.AsyncClient", return_value=_ClientCM()), - patch("pyrit.common.download_hf_model.download_chunk_async", new_callable=AsyncMock) as mock_chunk_async, - ): - mock_chunk_async.return_value = chunk_bytes - await download_file_async("https://example/myfile.bin", "token", tmp_path, num_splits) - - assert mock_chunk_async.await_count == num_splits - assert (tmp_path / "myfile.bin").read_bytes() == chunk_bytes * num_splits - - -async def test_download_chunk_async_returns_response_content(): - """Sanity-check the real download_chunk_async implementation.""" - response = MagicMock() - response.raise_for_status = MagicMock() - response.content = b"chunk-payload" - client = MagicMock() - client.get = AsyncMock(return_value=response) - - result = await download_chunk_async("https://example/file", {"Authorization": "Bearer t"}, 0, 9, client) - - assert result == b"chunk-payload" - client.get.assert_awaited_once()