Skip to content

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

Merged
romanlutz merged 3 commits into
microsoft:mainfrom
romanlutz:romanlutz-validate-c-hf-load
Jul 17, 2026
Merged

FIX: Load HuggingFace models without blocking the event loop#2211
romanlutz merged 3 commits into
microsoft:mainfrom
romanlutz:romanlutz-validate-c-hf-load

Conversation

@romanlutz

@romanlutz romanlutz commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

Description

HuggingFaceChatTarget previously downloaded and loaded models using blocking calls that ran directly on the asyncio event loop. For large or slow model pulls this stalled every other concurrent task until the download and weight loading finished. This change makes model loading non-blocking.

Key changes:

  • download_specific_files_async now delegates to huggingface_hub.snapshot_download run via asyncio.to_thread(...), so the download happens off the event loop. This also replaces the bespoke URL/chunked-download logic with the canonical HF snapshot API, which handles allow-patterns, auth tokens, and resume for us.
  • HuggingFaceChatTarget.load_model_and_tokenizer_async now loads the downloaded snapshot through the existing _load_from_path helper instead of duplicating from_pretrained calls, consolidating the two load paths (and removing a no-longer-needed cast).
  • Removed the now-dead custom download helpers (get_available_files and the download_files_async / download_file_async / download_chunk_async chain) and their unused imports, since snapshot_download supersedes them. download_specific_files_async is the only remaining helper.

Tests and Documentation

Tests:

  • Updated tests/unit/common/test_hf_model_downloads.py to assert snapshot_download is invoked with the expected repo_id, allow_patterns, token, and local_dir, and removed the tests covering the deleted helpers.
  • Updated tests/unit/prompt_target/target/test_huggingface_chat_target.py so the asyncio.create_task mock closes the coroutine (avoids "coroutine was never awaited" warnings).
  • uv run pytest on both files passes.

JupyText:

  • Ran uv run jupytext --execute --set-kernel python3 --to notebook doc/code/targets/use_huggingface_chat_target.py -- this downloads a real Qwen/Qwen2-0.5B-Instruct model through the new snapshot_download + _load_from_path path and runs inference end to end.
  • Ran the same for doc/code/targets/11_message_normalizer.py.
  • Both executed cleanly; the regenerated .ipynb outputs are included in this PR.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 52d2518d-ebe5-4888-8189-6146b6503448

@behnam-o behnam-o left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

little question/nit but otherwise great. thanks!

Comment thread pyrit/common/download_hf_model.py
Copilot AI added 2 commits July 16, 2026 20:39
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
@romanlutz
romanlutz enabled auto-merge July 17, 2026 03:49
@romanlutz
romanlutz added this pull request to the merge queue Jul 17, 2026
Merged via the queue into microsoft:main with commit 27d20ee Jul 17, 2026
53 checks passed
@romanlutz
romanlutz deleted the romanlutz-validate-c-hf-load branch July 17, 2026 04:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants