From a879aea80c66aa280bc847bffde36da3c6d15b31 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:34:43 -0700 Subject: [PATCH 1/2] FIX: Keep cold target catalog responsive Defer Hugging Face inference imports, coordinate registry metadata discovery, offload cold catalog work, and keep fallback target choices interactive while metadata loads. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../Config/CreateTargetDialog.test.tsx | 37 +++++++++- .../components/Config/CreateTargetDialog.tsx | 11 +-- pyrit/backend/services/target_service.py | 4 +- pyrit/prompt_target/__init__.py | 4 +- .../hugging_face/hugging_face_chat_target.py | 22 +++--- pyrit/registry/registry.py | 74 +++++++++++-------- .../test_target_catalog_concurrency.py | 44 +++++++++++ tests/unit/backend/test_target_service.py | 18 +++++ tests/unit/cli/test_import_guards.py | 21 ++++++ .../target/test_huggingface_chat_target.py | 2 +- tests/unit/registry/test_registry.py | 72 ++++++++++++++++++ tests/unit/registry/test_target_registry.py | 9 +++ 12 files changed, 264 insertions(+), 54 deletions(-) create mode 100644 tests/unit/backend/test_target_catalog_concurrency.py diff --git a/frontend/src/components/Config/CreateTargetDialog.test.tsx b/frontend/src/components/Config/CreateTargetDialog.test.tsx index a137ea254e..d19eea8a09 100644 --- a/frontend/src/components/Config/CreateTargetDialog.test.tsx +++ b/frontend/src/components/Config/CreateTargetDialog.test.tsx @@ -1,4 +1,4 @@ -import { render, screen, waitFor, fireEvent, within } from "@testing-library/react"; +import { act, render, screen, waitFor, fireEvent, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { FluentProvider, webLightTheme } from "@fluentui/react-components"; import { makeTarget } from "@/test-utils/targetFixtures"; @@ -313,7 +313,7 @@ describe("CreateTargetDialog", () => { expect(picker).not.toHaveTextContent("Select a target type"); }); - it("should disable target selection while catalog details are loading", () => { + it("should expose fallback target choices while catalog details are loading", async () => { mockedTargetsApi.listTargetCatalog.mockReturnValue( new Promise(() => {}), ); @@ -325,7 +325,38 @@ describe("CreateTargetDialog", () => { ); expect(screen.getByText("Loading target details...")).toBeInTheDocument(); - expect(screen.getByRole("combobox", { name: /target type/i })).toBeDisabled(); + await openTargetTypePicker(); + expect(screen.getAllByRole("option")).toHaveLength(8); + }); + + it("should preserve a fallback selection when catalog details arrive", async () => { + let resolveCatalog: ((catalog: TargetCatalogResponse) => void) | null = null; + mockedTargetsApi.listTargetCatalog.mockReturnValue( + new Promise((resolve) => { + resolveCatalog = resolve; + }), + ); + + render( + + + + ); + + await selectTargetType("OpenAIChatTarget"); + expect(screen.getByRole("combobox", { name: /target type/i })).toHaveTextContent("OpenAI chat"); + + if (resolveCatalog === null) { + throw new Error("Catalog resolver was not initialized"); + } + await act(async () => { + resolveCatalog(TARGET_CATALOG); + }); + + await waitFor(() => { + expect(screen.queryByText("Loading target details...")).not.toBeInTheDocument(); + }); + expect(screen.getByRole("combobox", { name: /target type/i })).toHaveTextContent("OpenAI chat"); }); it("should keep all target types selectable and explain when catalog details fail to load", async () => { diff --git a/frontend/src/components/Config/CreateTargetDialog.tsx b/frontend/src/components/Config/CreateTargetDialog.tsx index 6fa4afd62c..bcfb2a3e1c 100644 --- a/frontend/src/components/Config/CreateTargetDialog.tsx +++ b/frontend/src/components/Config/CreateTargetDialog.tsx @@ -269,11 +269,9 @@ export default function CreateTargetDialog({ open, onClose, onCreated, existingT }, [catalogEntries]) const catalogMetadataAvailable = catalogTargetTypeOptions.length > 0 const catalogUnavailable = catalogStatus !== 'loading' && !catalogMetadataAvailable - const targetTypeOptions = catalogStatus === 'loading' - ? [] - : catalogMetadataAvailable - ? catalogTargetTypeOptions - : FALLBACK_TARGET_CATALOG_ENTRIES + const targetTypeOptions = catalogMetadataAvailable + ? catalogTargetTypeOptions + : FALLBACK_TARGET_CATALOG_ENTRIES const formShape = TARGET_FORM_SHAPES[targetType] const isRoundRobin = formShape === 'roundrobin' @@ -531,9 +529,8 @@ export default function CreateTargetDialog({ open, onClose, onCreated, existingT TargetCatalogResponse: Returns: TargetCatalogResponse containing all available target classes. """ + metadata_items = await asyncio.to_thread(self._registry.get_all_registered_class_metadata) items: list[TargetCatalogEntry] = [ TargetCatalogEntry( target_type=metadata.class_name, @@ -144,7 +146,7 @@ async def list_target_catalog_async(self) -> TargetCatalogResponse: supported_auth_modes=cast("list[Literal['api_key', 'identity']]", list(metadata.supported_auth_modes)), description=metadata.class_description or None, ) - for metadata in self._registry.get_all_registered_class_metadata() + for metadata in metadata_items ] return TargetCatalogResponse(items=items) diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 9f45cb4ee5..91b6cce7f8 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -56,8 +56,8 @@ if TYPE_CHECKING: from pyrit.prompt_target.hugging_face.hugging_face_chat_target import HuggingFaceChatTarget -# Lazy imports for modules with heavy third-party dependencies (PEP 562). -# HuggingFaceChatTarget imports `transformers` which adds ~4s to startup. +# Keep optional inference targets lazy so package imports do not load their +# target-specific runtime modules. _LAZY_IMPORTS: dict[str, str] = { "HuggingFaceChatTarget": "pyrit.prompt_target.hugging_face.hugging_face_chat_target", } 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 ad98a94969..bb5c39be0c 100644 --- a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py +++ b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py @@ -6,17 +6,9 @@ import logging import warnings from pathlib import Path -from typing import Any, cast - -from transformers import ( - AutoModelForCausalLM, # type: ignore[ty:possibly-missing-import] - AutoTokenizer, # type: ignore[ty:possibly-missing-import] - BatchEncoding, - PretrainedConfig, -) +from typing import TYPE_CHECKING, Any, cast from pyrit.common import default_values -from pyrit.common.download_hf_model import download_specific_files_async from pyrit.exceptions import EmptyResponseException, pyrit_target_retry from pyrit.models import ComponentIdentifier, Message, construct_response_from_request from pyrit.prompt_target.common.prompt_target import PromptTarget @@ -24,6 +16,9 @@ from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.utils import limit_requests_per_minute +if TYPE_CHECKING: + from transformers import BatchEncoding + logger = logging.getLogger(__name__) @@ -212,6 +207,11 @@ def _load_from_path(self, path: str, **kwargs: Any) -> None: path: The path to load the model and tokenizer from. **kwargs: Additional keyword arguments to pass to the model loader. """ + from transformers import ( + AutoModelForCausalLM, # type: ignore[ty:possibly-missing-import] + AutoTokenizer, # type: ignore[ty:possibly-missing-import] + ) + logger.info(f"Loading model and tokenizer from path: {path}...") self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=self.trust_remote_code) self.model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=self.trust_remote_code, **kwargs) @@ -223,6 +223,8 @@ def is_model_id_valid(self) -> bool: Returns: bool: True if valid, False otherwise. """ + from transformers import PretrainedConfig # type: ignore[ty:possibly-missing-import] + try: # Attempt to load the configuration of the model PretrainedConfig.from_pretrained(self.model_id or "") @@ -267,6 +269,8 @@ async def load_model_and_tokenizer_async(self) -> None: logger.info(f"Loading model from local path: {self.model_path}...") self._load_from_path(self.model_path, **optional_model_kwargs) else: + from pyrit.common.download_hf_model import download_specific_files_async + # Define the default Hugging Face cache directory cache_dir = ( Path.home() diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index f79f1f20c7..5f370ac081 100644 --- a/pyrit/registry/registry.py +++ b/pyrit/registry/registry.py @@ -28,6 +28,7 @@ import inspect import logging +import threading from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar @@ -170,6 +171,7 @@ class Registry(ABC, Generic[T, MetadataT]): # Class-level singleton instances, keyed by registry class. _singletons: dict[type, Registry[Any, Any]] = {} + _singletons_lock = threading.RLock() def __init__(self, *, lazy_discovery: bool = True) -> None: """ @@ -179,6 +181,7 @@ def __init__(self, *, lazy_discovery: bool = True) -> None: lazy_discovery (bool): If True, discovery is deferred until first access. If False, discovery runs immediately in the constructor. """ + self._catalog_lock = threading.RLock() self._classes: dict[str, type[T]] = {} self._metadata_cache: dict[str, MetadataT] | None = None self._discovered = False @@ -198,9 +201,10 @@ def get_registry_singleton(cls) -> Self: Returns: The singleton instance of this registry class. """ - if cls not in cls._singletons: - cls._singletons[cls] = cls() - return cls._singletons[cls] # type: ignore[ty:invalid-return-type] + with cls._singletons_lock: + if cls not in cls._singletons: + cls._singletons[cls] = cls() + return cls._singletons[cls] # type: ignore[ty:invalid-return-type] @classmethod def reset_registry_singleton(cls) -> None: @@ -209,14 +213,16 @@ def reset_registry_singleton(cls) -> None: Useful for testing or when re-discovery is needed. """ - if cls in cls._singletons: - del cls._singletons[cls] + with cls._singletons_lock: + if cls in cls._singletons: + del cls._singletons[cls] def _ensure_discovered(self) -> None: """Ensure discovery has been performed. Runs discovery on first access.""" - if not self._discovered: - self._discover() - self._discovered = True + with self._catalog_lock: + if not self._discovered: + self._discover() + self._discovered = True def _base_type(self) -> type[T]: """ @@ -490,11 +496,12 @@ def register_class(self, cls: type[T], *, name: str | None = None) -> None: Raises: ValueError: If the class fails validation. """ - if name is None: - name = self._get_registry_name(cls) - self._validate_class(cls) - self._classes[name] = cls - self._metadata_cache = None + with self._catalog_lock: + if name is None: + name = self._get_registry_name(cls) + self._validate_class(cls) + self._classes[name] = cls + self._metadata_cache = None def get_class(self, name: str) -> type[T]: """ @@ -509,12 +516,13 @@ def get_class(self, name: str) -> type[T]: Raises: KeyError: If the name is not registered. """ - self._ensure_discovered() - cls = self._classes.get(name) - if cls is None: - available = ", ".join(self.get_class_names()) - raise KeyError(f"'{name}' not found in registry. Available: {available}") - return cls + with self._catalog_lock: + self._ensure_discovered() + cls = self._classes.get(name) + if cls is None: + available = ", ".join(self.get_class_names()) + raise KeyError(f"'{name}' not found in registry. Available: {available}") + return cls def get_class_names(self) -> list[str]: """ @@ -523,8 +531,9 @@ def get_class_names(self) -> list[str]: Returns: list[str]: Sorted catalog names. """ - self._ensure_discovered() - return sorted(self._classes.keys()) + with self._catalog_lock: + self._ensure_discovered() + return sorted(self._classes.keys()) def _ensure_metadata(self) -> dict[str, MetadataT]: """ @@ -533,12 +542,13 @@ def _ensure_metadata(self) -> dict[str, MetadataT]: Returns: dict[str, MetadataT]: Metadata for every registered class, keyed by name. """ - self._ensure_discovered() - if self._metadata_cache is None: - self._metadata_cache = { - name: self._build_metadata(name, cls) for name, cls in sorted(self._classes.items()) - } - return self._metadata_cache + with self._catalog_lock: + self._ensure_discovered() + if self._metadata_cache is None: + self._metadata_cache = { + name: self._build_metadata(name, cls) for name, cls in sorted(self._classes.items()) + } + return self._metadata_cache def get_all_registered_class_metadata( self, @@ -634,8 +644,9 @@ def __contains__(self, name: str) -> bool: Returns: bool: True if the name is registered, False otherwise. """ - self._ensure_discovered() - return name in self._classes + with self._catalog_lock: + self._ensure_discovered() + return name in self._classes def __len__(self) -> int: """ @@ -644,8 +655,9 @@ def __len__(self) -> int: Returns: int: The number of registered classes. """ - self._ensure_discovered() - return len(self._classes) + with self._catalog_lock: + self._ensure_discovered() + return len(self._classes) def __iter__(self) -> Iterator[str]: """ diff --git a/tests/unit/backend/test_target_catalog_concurrency.py b/tests/unit/backend/test_target_catalog_concurrency.py new file mode 100644 index 0000000000..d2b7cbe46e --- /dev/null +++ b/tests/unit/backend/test_target_catalog_concurrency.py @@ -0,0 +1,44 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Concurrency regressions for target catalog routes.""" + +import asyncio +from threading import Event +from unittest.mock import patch + +from httpx import ASGITransport, AsyncClient + +from pyrit.backend.main import app +from pyrit.backend.services.target_service import TargetService + + +async def test_health_remains_schedulable_during_cold_target_catalog() -> None: + discovery_started = Event() + discovery_release = Event() + discovery_finished = Event() + service = TargetService() + + def _blocking_metadata_discovery() -> list[object]: + discovery_started.set() + discovery_release.wait(timeout=5) + discovery_finished.set() + return [] + + transport = ASGITransport(app=app) + with ( + patch.object(service._registry, "get_all_registered_class_metadata", side_effect=_blocking_metadata_discovery), + patch("pyrit.backend.routes.targets.get_target_service", return_value=service), + ): + async with AsyncClient(transport=transport, base_url="http://test") as client: + catalog_request = asyncio.create_task(client.get("/api/targets/catalog")) + assert await asyncio.to_thread(discovery_started.wait, 5) + + health_response = await asyncio.wait_for(client.get("/api/health"), timeout=2) + + assert health_response.status_code == 200 + assert not discovery_finished.is_set() + discovery_release.set() + catalog_response = await asyncio.wait_for(catalog_request, timeout=2) + + assert catalog_response.status_code == 200 diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py index d6fd6df784..e226d33416 100644 --- a/tests/unit/backend/test_target_service.py +++ b/tests/unit/backend/test_target_service.py @@ -253,6 +253,24 @@ async def test_catalog_includes_declarative_auth_facts(self) -> None: assert "api_key" in openai_entry.supported_auth_modes assert "identity" in openai_entry.supported_auth_modes + async def test_catalog_cold_and_warm_results_are_equal(self) -> None: + service = TargetService() + + cold = await service.list_target_catalog_async() + warm = await service.list_target_catalog_async() + + assert cold == warm + + async def test_catalog_refreshes_after_runtime_class_registration(self) -> None: + service = TargetService() + initial = await service.list_target_catalog_async() + + service._registry.register_class(MockPromptTarget) + refreshed = await service.list_target_catalog_async() + + assert all(item.target_type != "MockPromptTarget" for item in initial.items) + assert any(item.target_type == "MockPromptTarget" for item in refreshed.items) + @pytest.mark.parametrize( ("target_type", "parameter_name", "type_name", "required", "choices"), [ diff --git a/tests/unit/cli/test_import_guards.py b/tests/unit/cli/test_import_guards.py index cf5b062165..b5f95572d2 100644 --- a/tests/unit/cli/test_import_guards.py +++ b/tests/unit/cli/test_import_guards.py @@ -89,6 +89,12 @@ def _check_forbidden_imports(*, import_statement: str, forbidden: list[str]) -> "transformers", ] +_TARGET_CATALOG_FORBIDDEN = [ + "huggingface_hub", + "torch", + "transformers", +] + class TestImportGuards: """Verify heavy modules are not eagerly loaded at key import points.""" @@ -133,3 +139,18 @@ def test_prompt_target_base_does_not_load_ml_modules(self): f"PromptTarget base class loaded ML modules: {loaded}. " f"Ensure heavy subclass imports use __getattr__ lazy loading in __init__.py." ) + + def test_target_catalog_discovery_does_not_load_inference_frameworks(self) -> None: + """Full target discovery includes Hugging Face without importing its runtime frameworks.""" + loaded = _check_forbidden_imports( + import_statement=( + "from pyrit.registry import TargetRegistry\n" + "metadata = TargetRegistry.get_registry_singleton().get_all_registered_class_metadata()\n" + "assert any(item.class_name == 'HuggingFaceChatTarget' for item in metadata)" + ), + forbidden=_TARGET_CATALOG_FORBIDDEN, + ) + assert not loaded, ( + f"Target catalog discovery loaded inference frameworks: {loaded}. " + f"Move target-specific runtime imports to construction or execution paths." + ) 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 12177d7c44..b81a23e917 100644 --- a/tests/unit/prompt_target/target/test_huggingface_chat_target.py +++ b/tests/unit/prompt_target/target/test_huggingface_chat_target.py @@ -106,7 +106,7 @@ def _close_coroutine(coroutine: Coroutine[Any, Any, None]) -> AwaitableTask: @pytest.fixture(autouse=True) def mock_download_specific_files_async(): with patch( - "pyrit.prompt_target.hugging_face.hugging_face_chat_target.download_specific_files_async", + "pyrit.common.download_hf_model.download_specific_files_async", new_callable=AsyncMock, ) as mock: yield mock diff --git a/tests/unit/registry/test_registry.py b/tests/unit/registry/test_registry.py index 8f0aef0388..9d0263e06e 100644 --- a/tests/unit/registry/test_registry.py +++ b/tests/unit/registry/test_registry.py @@ -11,7 +11,9 @@ every base default. """ +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field +from threading import Event from types import ModuleType from unittest.mock import MagicMock @@ -43,6 +45,13 @@ def __init__(self, *, size: int = 1) -> None: self.size = size +class PluginWidget: + """A widget registered after discovery.""" + + def __init__(self, *, size: int = 1) -> None: + self.size = size + + class WidgetRegistry(Registry[object, RegistryMetadata]): """Minimal Registry subclass that keeps every base default.""" @@ -59,6 +68,24 @@ def _metadata_class(self) -> type[RegistryMetadata]: return RegistryMetadata +class CoordinatedWidgetRegistry(WidgetRegistry): + """Widget registry whose first metadata build can be held for concurrency tests.""" + + def __init__(self) -> None: + self.metadata_started = Event() + self.metadata_release = Event() + self.metadata_build_calls = 0 + super().__init__() + + def _build_metadata(self, name: str, cls: type[object]) -> RegistryMetadata: + self.metadata_build_calls += 1 + if self.metadata_build_calls == 1: + self.metadata_started.set() + if not self.metadata_release.wait(timeout=5): + raise TimeoutError("Metadata test release was not signaled.") + return super()._build_metadata(name, cls) + + @dataclass(frozen=True) class _TaggedMetadata(RegistryMetadata): tags: tuple[str, ...] = field(kw_only=True, default=()) @@ -170,6 +197,51 @@ def test_get_all_metadata_no_filters_returns_all(): assert {m.registry_name for m in all_meta} == {"SampleWidget", "UndocumentedWidget"} +def test_concurrent_metadata_callers_share_one_cache_build() -> None: + registry = CoordinatedWidgetRegistry() + + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(registry.get_all_registered_class_metadata) + assert registry.metadata_started.wait(timeout=5) + second = executor.submit(registry.get_all_registered_class_metadata) + registry.metadata_release.set() + + assert first.result(timeout=5) == second.result(timeout=5) + + assert registry.metadata_build_calls == 2 + + +def test_class_registration_waits_for_metadata_build_and_invalidates_cache() -> None: + registry = CoordinatedWidgetRegistry() + registration_started = Event() + + def _register_plugin() -> None: + registration_started.set() + registry.register_class(PluginWidget) + + with ThreadPoolExecutor(max_workers=2) as executor: + initial_metadata = executor.submit(registry.get_all_registered_class_metadata) + assert registry.metadata_started.wait(timeout=5) + registration = executor.submit(_register_plugin) + assert registration_started.wait(timeout=5) + assert not registration.done() + + registry.metadata_release.set() + assert {item.class_name for item in initial_metadata.result(timeout=5)} == { + "SampleWidget", + "UndocumentedWidget", + } + registration.result(timeout=5) + + refreshed = registry.get_all_registered_class_metadata() + assert {item.class_name for item in refreshed} == { + "PluginWidget", + "SampleWidget", + "UndocumentedWidget", + } + assert registry.metadata_build_calls == 5 + + def test_get_all_metadata_include_filter_matches_subset(): registry = WidgetRegistry() diff --git a/tests/unit/registry/test_target_registry.py b/tests/unit/registry/test_target_registry.py index da753d5076..60a78a4e09 100644 --- a/tests/unit/registry/test_target_registry.py +++ b/tests/unit/registry/test_target_registry.py @@ -347,6 +347,15 @@ def test_metadata_supported_auth_modes_sourced_from_class_attributes(self, regis assert "supported_auth_modes" in meta.class_attributes assert meta.class_attributes["supported_auth_modes"] == ("api_key", "identity") + @pytest.mark.usefixtures("patch_central_database") + def test_instance_registration_does_not_invalidate_class_metadata(self, registry: TargetRegistry) -> None: + registry.get_all_registered_class_metadata() + initial_cache = registry._metadata_cache + + registry.instances.register(MockPromptTarget(), name="runtime-instance") + + assert registry._metadata_cache is initial_cache + def test_openai_metadata_includes_forwarded_base_parameters(self, registry: TargetRegistry) -> None: params = {param.name: param for param in self._metadata_for(registry, "OpenAIChatTarget").parameters} From e598358d6d1d70296d1c2e58c65fcb5d46f7e042 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:32:15 -0700 Subject: [PATCH 2/2] FIX: Address target catalog review feedback Narrow registry lock scope, coordinate singleton construction per registry class, keep Hugging Face model loading off the event loop, and eliminate frontend act warnings. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 130110b1-e9bb-4dc5-9f71-c2405871b1e4 --- .../Config/CreateTargetDialog.test.tsx | 25 +++- .../hugging_face/hugging_face_chat_target.py | 17 ++- pyrit/registry/registry.py | 71 ++++++++++-- .../target/test_huggingface_chat_target.py | 49 ++++++++ tests/unit/registry/test_registry.py | 109 ++++++++++++++++-- 5 files changed, 242 insertions(+), 29 deletions(-) diff --git a/frontend/src/components/Config/CreateTargetDialog.test.tsx b/frontend/src/components/Config/CreateTargetDialog.test.tsx index d19eea8a09..319bed9128 100644 --- a/frontend/src/components/Config/CreateTargetDialog.test.tsx +++ b/frontend/src/components/Config/CreateTargetDialog.test.tsx @@ -144,6 +144,19 @@ async function selectTargetType(value: string): Promise { restoreDialogAccessibility(); } +// The catalog fetch mock (see beforeEach) resolves on mount, and its +// setCatalogEntries/setCatalogStatus updates land on the next microtask +// tick. Tests that follow up with an `await` (selectTargetType, +// openTargetTypePicker, userEvent, ...) give React a chance to settle that +// update inside their own act()-wrapped waiting. Tests that only make +// synchronous assertions after `render` never yield, so the update fires +// after the test body returns and React reports it as outside act(...). +// Call this right after `render` in those synchronous tests to flush it +// deterministically. +async function flushCatalogFetch(): Promise { + await act(async () => {}); +} + describe("parseWeight", () => { it("rejects empty input", () => { expect(parseWeight("")).toEqual({ ok: false, error: "Weight is required" }); @@ -233,12 +246,13 @@ describe("CreateTargetDialog", () => { dialogAccessibilityObserver.disconnect(); }); - it("should render dialog when open", () => { + it("should render dialog when open", async () => { render( ); + await flushCatalogFetch(); expect(screen.getByText("Create New Target")).toBeInTheDocument(); expect(screen.getByText("Create Target")).toBeInTheDocument(); @@ -386,12 +400,13 @@ describe("CreateTargetDialog", () => { expect(screen.queryByText("Create New Target")).not.toBeInTheDocument(); }); - it("should have Create button disabled until type and endpoint filled", () => { + it("should have Create button disabled until type and endpoint filled", async () => { render( ); + await flushCatalogFetch(); const createButton = screen.getByText("Create Target"); expect(createButton.closest("button")).toBeDisabled(); @@ -634,12 +649,13 @@ describe("CreateTargetDialog", () => { }); }); - it("should display supported target initializer guidance", () => { + it("should display supported target initializer guidance", async () => { render( ); + await flushCatalogFetch(); expect(screen.getByText("target", { selector: "code" })).toBeInTheDocument(); expect( @@ -650,12 +666,13 @@ describe("CreateTargetDialog", () => { ).not.toBeInTheDocument(); }); - it("should render .pyrit_conf_example as an accessible link", () => { + it("should render .pyrit_conf_example as an accessible link", async () => { render( ); + await flushCatalogFetch(); const link = screen.getByRole("link", { name: ".pyrit_conf_example" }); expect(link).toBeInTheDocument(); 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 bb5c39be0c..492f34f21d 100644 --- a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py +++ b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py @@ -265,9 +265,11 @@ async def load_model_and_tokenizer_async(self) -> None: return if self.model_path: - # Load the tokenizer and model from the local directory + # Load the tokenizer and model from the local directory. This imports `transformers` + # and performs blocking disk I/O, so it is offloaded to a worker thread to keep the + # event loop responsive. logger.info(f"Loading model from local path: {self.model_path}...") - self._load_from_path(self.model_path, **optional_model_kwargs) + await asyncio.to_thread(self._load_from_path, self.model_path, **optional_model_kwargs) else: from pyrit.common.download_hf_model import download_specific_files_async @@ -299,12 +301,15 @@ async def load_model_and_tokenizer_async(self) -> None: Path(cache_dir), ) - # Load the tokenizer and model from the downloaded local snapshot. + # Load the tokenizer and model from the downloaded local snapshot. This imports + # `transformers` and performs blocking disk I/O, so it is offloaded to a worker + # thread to keep the event loop responsive. logger.info(f"Loading model {self.model_id} from cache path: {cache_dir}...") - self._load_from_path(str(cache_dir), **optional_model_kwargs) + await asyncio.to_thread(self._load_from_path, str(cache_dir), **optional_model_kwargs) - # Move the model to the correct device - self.model = self.model.to(self.device) + # Move the model to the correct device. This can be a slow, blocking operation + # (e.g., copying weights to a GPU), so it is offloaded to a worker thread as well. + self.model = await asyncio.to_thread(self.model.to, self.device) # Debug prints to check types logger.info(f"Model loaded: {type(self.model)}") diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index 5f370ac081..33b7c196f5 100644 --- a/pyrit/registry/registry.py +++ b/pyrit/registry/registry.py @@ -171,7 +171,13 @@ class Registry(ABC, Generic[T, MetadataT]): # Class-level singleton instances, keyed by registry class. _singletons: dict[type, Registry[Any, Any]] = {} + # Guards only the two class-level dicts below (fast dict lookups); it is never held + # while a registry class is being constructed, so unrelated registry classes never + # queue behind one another's constructor. _singletons_lock = threading.RLock() + # Per-registry-class construction locks, so concurrent same-class callers converge + # onto a single construction while different registry classes construct independently. + _singleton_construction_locks: dict[type, threading.Lock] = {} def __init__(self, *, lazy_discovery: bool = True) -> None: """ @@ -182,8 +188,15 @@ def __init__(self, *, lazy_discovery: bool = True) -> None: If False, discovery runs immediately in the constructor. """ self._catalog_lock = threading.RLock() + # Guards the (potentially slow) metadata build so concurrent callers single-flight + # onto one build instead of each redoing it; never held while _catalog_lock is held, + # so catalog reads (__contains__, get_class, ...) are never parked behind a build. + self._metadata_build_lock = threading.RLock() self._classes: dict[str, type[T]] = {} self._metadata_cache: dict[str, MetadataT] | None = None + # Bumped on every catalog mutation so an in-flight metadata build can detect a + # registration that landed mid-build and retry instead of caching a stale snapshot. + self._catalog_version = 0 self._discovered = False self._lazy_discovery = lazy_discovery @@ -196,15 +209,30 @@ def get_registry_singleton(cls) -> Self: """ Get the singleton instance of this registry. - Creates the instance on first call with default parameters. + Creates the instance on first call with default parameters. Construction + happens under a per-class lock rather than the shared map lock, so a slow + eager constructor (e.g. one that runs discovery synchronously) blocks only + other callers of the *same* registry class; unrelated registry classes + construct independently. Returns: The singleton instance of this registry class. """ with cls._singletons_lock: - if cls not in cls._singletons: - cls._singletons[cls] = cls() - return cls._singletons[cls] # type: ignore[ty:invalid-return-type] + instance = cls._singletons.get(cls) + if instance is not None: + return instance # type: ignore[ty:invalid-return-type] + construction_lock = cls._singleton_construction_locks.setdefault(cls, threading.Lock()) + + with construction_lock: + with cls._singletons_lock: + instance = cls._singletons.get(cls) + if instance is not None: + return instance # type: ignore[ty:invalid-return-type] + instance = cls() + with cls._singletons_lock: + cls._singletons[cls] = instance + return instance # type: ignore[ty:invalid-return-type] @classmethod def reset_registry_singleton(cls) -> None: @@ -502,6 +530,7 @@ def register_class(self, cls: type[T], *, name: str | None = None) -> None: self._validate_class(cls) self._classes[name] = cls self._metadata_cache = None + self._catalog_version += 1 def get_class(self, name: str) -> type[T]: """ @@ -539,16 +568,40 @@ def _ensure_metadata(self) -> dict[str, MetadataT]: """ Build (once) and return the metadata cache keyed by catalog name. + Building metadata re-derives every registered class's build contract (the + same introspection cost as discovery) and can take a while for a large + catalog. That work runs with ``_catalog_lock`` released so it never parks + concurrent catalog reads (``__contains__``, ``get_class``, ...) — those + only need the lock for a quick dict lookup. ``_metadata_build_lock`` still + makes concurrent metadata callers single-flight onto one build, and the + ``_catalog_version`` stamp detects a registration that lands mid-build so a + stale snapshot is never cached over a newer one; the build simply retries. + Returns: dict[str, MetadataT]: Metadata for every registered class, keyed by name. """ with self._catalog_lock: self._ensure_discovered() - if self._metadata_cache is None: - self._metadata_cache = { - name: self._build_metadata(name, cls) for name, cls in sorted(self._classes.items()) - } - return self._metadata_cache + if self._metadata_cache is not None: + return self._metadata_cache + + with self._metadata_build_lock: + while True: + with self._catalog_lock: + if self._metadata_cache is not None: + return self._metadata_cache + classes_snapshot = dict(self._classes) + version = self._catalog_version + + built = {name: self._build_metadata(name, cls) for name, cls in sorted(classes_snapshot.items())} + + with self._catalog_lock: + if self._metadata_cache is not None: + return self._metadata_cache + if self._catalog_version == version: + self._metadata_cache = built + return self._metadata_cache + # A registration landed mid-build; retry with a fresh snapshot. def get_all_registered_class_metadata( self, 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 b81a23e917..ad030403b6 100644 --- a/tests/unit/prompt_target/target/test_huggingface_chat_target.py +++ b/tests/unit/prompt_target/target/test_huggingface_chat_target.py @@ -1,7 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import json +import threading from asyncio import Task from collections.abc import Coroutine from typing import Any @@ -179,6 +181,53 @@ async def test_load_model_and_tokenizer(): assert hf_chat.tokenizer is not None +@pytest.mark.skipif(not is_torch_installed(), reason="torch is not installed") +async def test_load_model_and_tokenizer_keeps_event_loop_schedulable(patch_central_database): + """The blocking `transformers` import/model load must run off the event loop. + + `_load_from_path` is patched to block a real OS thread (via `threading.Event`) rather than + sleeping, so if it ran directly on the event loop this test would deadlock/timeout instead of + merely running slow -- a deterministic failure signal rather than a flaky timing assertion. + """ + HuggingFaceChatTarget.disable_cache() + try: + hf_chat = HuggingFaceChatTarget(model_id="test_model_event_loop_probe", use_cuda=False) + + load_started = threading.Event() + load_release = threading.Event() + + def _blocking_load(path: str, **kwargs: Any) -> None: + load_started.set() + assert load_release.wait(timeout=5), "load_release was never set; test would hang otherwise" + hf_chat.tokenizer = MagicMock() + hf_chat.model = MagicMock() + hf_chat.model.to.return_value = hf_chat.model + + with patch.object(hf_chat, "_load_from_path", side_effect=_blocking_load) as mock_load_from_path: + load_task = asyncio.ensure_future(hf_chat.load_model_and_tokenizer_async()) + + # Confirm the blocking call actually started on a worker thread before probing. + assert await asyncio.to_thread(load_started.wait, 5) + assert not load_task.done() + + # While the worker thread is parked on `load_release`, the event loop itself must + # still be able to schedule and complete unrelated work. If `_load_from_path` (and the + # `transformers` import it performs) ran directly on the event loop, this would never + # get a chance to run and `asyncio.wait_for` would raise `TimeoutError`. + probe_result = await asyncio.wait_for(asyncio.sleep(0, result="probe-completed"), timeout=2) + assert probe_result == "probe-completed" + assert not load_task.done() + + load_release.set() + await asyncio.wait_for(load_task, timeout=5) + + mock_load_from_path.assert_called_once() + assert hf_chat.model is not None + assert hf_chat.tokenizer is not None + finally: + HuggingFaceChatTarget.enable_cache() + + @pytest.mark.skipif(not is_torch_installed(), reason="torch is not installed") @pytest.mark.usefixtures("patch_central_database") async def test_send_prompt_async(): diff --git a/tests/unit/registry/test_registry.py b/tests/unit/registry/test_registry.py index 9d0263e06e..5f5172fc82 100644 --- a/tests/unit/registry/test_registry.py +++ b/tests/unit/registry/test_registry.py @@ -86,6 +86,26 @@ def _build_metadata(self, name: str, cls: type[object]) -> RegistryMetadata: return super()._build_metadata(name, cls) +class SlowConstructingRegistry(WidgetRegistry): + """Registry subclass whose constructor pauses so singleton construction concurrency + can be tested. Uses class-level events because ``get_registry_singleton`` calls + ``cls()`` with no arguments.""" + + construction_started = Event() + construction_release = Event() + + def __init__(self) -> None: + self.construction_started.set() + if not self.construction_release.wait(timeout=5): + raise TimeoutError("Slow-construction test release was not signaled.") + super().__init__() + + +class OtherSlowConstructingRegistry(WidgetRegistry): + """A distinct Registry subclass with an independent singleton key, used to prove + that constructing SlowConstructingRegistry's singleton does not block this one.""" + + @dataclass(frozen=True) class _TaggedMetadata(RegistryMetadata): tags: tuple[str, ...] = field(kw_only=True, default=()) @@ -211,27 +231,31 @@ def test_concurrent_metadata_callers_share_one_cache_build() -> None: assert registry.metadata_build_calls == 2 -def test_class_registration_waits_for_metadata_build_and_invalidates_cache() -> None: +def test_class_registration_does_not_block_behind_metadata_build_and_converges() -> None: + """ + A concurrent metadata build must not park a fast catalog mutation like + ``register_class`` behind its (potentially slow) work. The build must instead + detect, via the catalog version stamp, that the catalog changed mid-build and + retry so the returned metadata converges on a snapshot that includes the new + class rather than caching a stale one. + """ registry = CoordinatedWidgetRegistry() - registration_started = Event() - - def _register_plugin() -> None: - registration_started.set() - registry.register_class(PluginWidget) with ThreadPoolExecutor(max_workers=2) as executor: initial_metadata = executor.submit(registry.get_all_registered_class_metadata) assert registry.metadata_started.wait(timeout=5) - registration = executor.submit(_register_plugin) - assert registration_started.wait(timeout=5) - assert not registration.done() + + registration = executor.submit(registry.register_class, PluginWidget) + # Registration only needs the (uncontended) catalog lock, so it must + # complete promptly even though the metadata build is still paused. + registration.result(timeout=1) registry.metadata_release.set() assert {item.class_name for item in initial_metadata.result(timeout=5)} == { + "PluginWidget", "SampleWidget", "UndocumentedWidget", } - registration.result(timeout=5) refreshed = registry.get_all_registered_class_metadata() assert {item.class_name for item in refreshed} == { @@ -242,6 +266,71 @@ def _register_plugin() -> None: assert registry.metadata_build_calls == 5 +def test_contains_and_get_class_do_not_block_behind_metadata_build() -> None: + """ + ``__contains__`` and ``get_class`` must return promptly even while a metadata + build is in flight and paused on another thread: they must never be parked + behind the lock that guards the (slow) metadata build. + """ + registry = CoordinatedWidgetRegistry() + registry.get_class_names() # Trigger discovery up front to isolate the build pause. + + with ThreadPoolExecutor(max_workers=3) as executor: + metadata_future = executor.submit(registry.get_all_registered_class_metadata) + assert registry.metadata_started.wait(timeout=5) + + contains_future = executor.submit(lambda: "SampleWidget" in registry) + get_class_future = executor.submit(registry.get_class, "SampleWidget") + + assert contains_future.result(timeout=1) is True + assert get_class_future.result(timeout=1) is SampleWidget + + registry.metadata_release.set() + metadata_future.result(timeout=5) + + +def test_get_registry_singleton_converges_for_same_class() -> None: + """Concurrent callers requesting the same registry class's singleton must + converge onto a single construction.""" + SlowConstructingRegistry.construction_started.clear() + SlowConstructingRegistry.construction_release.clear() + SlowConstructingRegistry.reset_registry_singleton() + try: + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(SlowConstructingRegistry.get_registry_singleton) + assert SlowConstructingRegistry.construction_started.wait(timeout=5) + second = executor.submit(SlowConstructingRegistry.get_registry_singleton) + SlowConstructingRegistry.construction_release.set() + + assert first.result(timeout=5) is second.result(timeout=5) + finally: + SlowConstructingRegistry.reset_registry_singleton() + + +def test_get_registry_singleton_does_not_block_unrelated_registry_class() -> None: + """A slow, in-flight singleton construction for one registry class must not + block singleton construction of an unrelated registry class.""" + SlowConstructingRegistry.construction_started.clear() + SlowConstructingRegistry.construction_release.clear() + SlowConstructingRegistry.reset_registry_singleton() + OtherSlowConstructingRegistry.reset_registry_singleton() + try: + with ThreadPoolExecutor(max_workers=2) as executor: + slow = executor.submit(SlowConstructingRegistry.get_registry_singleton) + assert SlowConstructingRegistry.construction_started.wait(timeout=5) + + # Must complete promptly: different registry classes must not share a + # construction lock. + other = executor.submit(OtherSlowConstructingRegistry.get_registry_singleton) + assert isinstance(other.result(timeout=1), OtherSlowConstructingRegistry) + + SlowConstructingRegistry.construction_release.set() + assert isinstance(slow.result(timeout=5), SlowConstructingRegistry) + finally: + SlowConstructingRegistry.reset_registry_singleton() + OtherSlowConstructingRegistry.reset_registry_singleton() + + def test_get_all_metadata_include_filter_matches_subset(): registry = WidgetRegistry()