diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 5015795657f..b7e67fffd57 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -1815,13 +1815,14 @@ class RustCore( def _rebuild(self) -> None: """A new Rust core for changed static settings, over the same deployments and the providers the old one's builds resolved for them.""" - old: Final = self._handle - model_list: Final = list(self._hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # see __init__ - self._handle = self._new_handle(model_list) - self._handle.adopt_providers(old) - old.retire() - # The live list and indexes read the new handle's (the same entries, kept as the same objects). - self._hosted.refresh() + with self._hosted._deployment_index_lock: # pyright: ignore[reportPrivateUsage] # the registry's lock, held across the swap and the read back + old: Final = self._handle + model_list: Final = list(self._hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # see __init__ + self._handle = self._new_handle(model_list) + self._handle.adopt_providers(old) + old.retire() + # The live list and indexes read the new handle's (the same entries, kept as the same objects). + self._hosted.refresh() # The new handle builds its indexes from the deployments: written ones are handed again. self._index_sent = None diff --git a/litellm/router_rust/surface/model_list.py b/litellm/router_rust/surface/model_list.py index bb7fb53acea..32ab8f1cb76 100644 --- a/litellm/router_rust/surface/model_list.py +++ b/litellm/router_rust/surface/model_list.py @@ -26,7 +26,9 @@ model info caches) mean what they did: the same object until a write changes tha import copy import datetime +import threading from collections.abc import Callable, Iterable, Mapping, Sequence +from contextlib import AbstractContextManager from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -189,9 +191,12 @@ class HostedModelList(list[Entry]): apply: Callable[[Mapping[str, object]], None], live: Callable[[object], Mapping[str, object]] = lambda _: {}, live_info: Callable[[object], Mapping[str, object]] = lambda _: {}, + *, + lock: AbstractContextManager[object] | None = None, ) -> None: super().__init__() self._handle: Final = handle + self._lock: Final[AbstractContextManager[object]] = threading.RLock() if lock is None else lock self._apply: Final = apply # A deployment's litellm_params values its record carries as text, by id (``live_params``). self._live: Final = live @@ -209,6 +214,10 @@ class HostedModelList(list[Entry]): def refresh(self) -> None: """Read router-core's deployments again when its registry moved since the last read.""" + with self._lock: + self._refresh() + + def _refresh(self) -> None: handle: Final = self._handle() if handle is None or self._building is not None: # While a build runs, the list shows what it accepted until router-core serves it. @@ -260,28 +269,31 @@ class HostedModelList(list[Entry]): """A build starts: router-core serves what it adds once it ends, so until then the list shows each deployment it accepts here (``built``), after the entries it keeps, or alone when it replaces them, as the Python router's list grows while it adds them.""" - self.refresh() - if replacing: - self._show([]) - self._id_counts = {} - self._building = (len(self), self._id_counts, self._read_at) + with self._lock: + self._refresh() + if replacing: + self._show([]) + self._id_counts = {} + self._building = (len(self), self._id_counts, self._read_at) def built(self, entry: Mapping[str, object]) -> None: """A deployment the build accepted, shown until the build ends.""" - if self._building is not None: - super().extend(self._kept([entry], appended=True)) + with self._lock: + if self._building is not None: + super().extend(self._kept([entry], appended=True)) def end_build(self) -> None: """The build ended (or failed): router-core serves what it added, read back in place of the entries shown meanwhile, each kept as the object it was shown as while unchanged.""" - if self._building is None: - return - start, counts, read_at = self._building - self._building = None - super().__delitem__(slice(start, None)) - self._id_counts = counts - self._read_at = read_at - self.refresh() + with self._lock: + if self._building is None: + return + start, counts, read_at = self._building + self._building = None + super().__delitem__(slice(start, None)) + self._id_counts = counts + self._read_at = read_at + self._refresh() def _seed(self, entries: Sequence[object]) -> None: """Keep ``entries``, the list router-core now holds, as the objects its records for them @@ -311,19 +323,21 @@ class HostedModelList(list[Entry]): def _replace(self, entries: list[Entry]) -> None: # mutable-ok: the list router-core takes # Seeded first: router-core's records for them read back as these objects however soon # the list is read again. - self._seed(entries) - self._apply({"replace": entries}) - self.refresh() + with self._lock: + self._seed(entries) + self._apply({"replace": entries}) + self._refresh() def replace(self, entries: Iterable[object]) -> None: """Every deployment at once (``router.model_list = [...]``).""" self._replace(cast(list[Entry], list(entries))) # cast-ok: the entries as assigned, which router-core validates def append(self, value: Entry) -> None: - self.refresh() - self._seed_appended([value]) - self._apply({"append": [value]}) - self.refresh() + with self._lock: + self._refresh() + self._seed_appended([value]) + self._apply({"append": [value]}) + self._refresh() def extend(self, values: Iterable[Entry]) -> None: for value in tuple(values): @@ -362,8 +376,9 @@ class HostedModelList(list[Entry]): # A delete removes every deployment with the id; the list removes only this one. self._replace(changed) return - self._apply({"delete": [removed_id]}) - self.refresh() + with self._lock: + self._apply({"delete": [removed_id]}) + self._refresh() def pop(self, index: SupportsIndex = -1) -> Entry: entry: Final = self[index] diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index e48569c4c9d..83008bdbb94 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -743,7 +743,7 @@ class HostRegistry(HostModelInfo): self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = model_group_alias # mutable-ok: the live alias map # fmt: skip # The live model_list: router-core's deployments read back, every write a changeset. self._core_model_list: Final[HostedModelList] = HostedModelList( - self._live_handle, self._hand_changes, self.live_params, self.live_info + self._live_handle, self._hand_changes, self.live_params, self.live_info, lock=self._deployment_index_lock ) # The litellm_params values a deployment's JSON cannot carry (an Exception or model object as # ``mock_response``, a callable ``azure_ad_token_provider``), by deployment id: an attempt is called with them, as Python's list held them. @@ -1472,9 +1472,10 @@ class HostRegistry(HostModelInfo): def _hand_changes(self, changes: Mapping[str, object]) -> None: if self.changes_runner is None: raise RuntimeError("the host registry writes deployments through its Rust core") - self.changes_runner(changes) - self._core_model_list.refresh() - self._drop_removed_live_values() + with self._deployment_index_lock: + self.changes_runner(changes) + self._core_model_list.refresh() + self._drop_removed_live_values() def _drop_removed_live_values(self) -> None: """Forget the live values of ids no deployment carries any more, so a router that rotates @@ -1678,10 +1679,11 @@ class HostRegistry(HostModelInfo): return handle is not None and key is not None and handle.has_model_id(key) def _core_get_model_info(self, id: str) -> UntypedEntry | None: - handle: Final = self._live_handle() - key: Final = _index_key(id) - idx: Final[object] = None if handle is None or key is None else handle.id_position(key) - return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(id) + idx: Final[object] = None if handle is None or key is None else handle.id_position(key) + return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it def live_params(self, model_id: object) -> Mapping[str, object]: """The litellm_params values of deployment ``model_id`` that its JSON record cannot carry.""" @@ -1696,21 +1698,23 @@ class HostRegistry(HostModelInfo): return None if entry is None else Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry def _core_get_deployment_by_model_group_name(self, model_group_name: str) -> Deployment | None: - handle: Final = self._live_handle() - key: Final = _index_key(model_group_name) - indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) - model: Final = self._core_model_list[indices[0]] if indices else None # pyright: ignore[reportIndexIssue, reportArgumentType, reportCallIssue] # positions, or what a caller wrote, read as Python reads them - return None if model is None else Deployment(**model) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(model_group_name) + indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) + model: Final = self._core_model_list[indices[0]] if indices else None # pyright: ignore[reportIndexIssue, reportArgumentType, reportCallIssue] # positions, or what a caller wrote, read as Python reads them + return None if model is None else Deployment(**model) # pyright: ignore[reportArgumentType] # a stored entry def _at(self, indices: object) -> tuple[UntypedEntry, ...]: """The entries at ``indices`` (what an index holds), read as Python's list reads them.""" return tuple(self._core_model_list[idx] for idx in indices) # pyright: ignore[reportGeneralTypeIssues, reportArgumentType, reportCallIssue] # positions, or what a caller wrote def _core_deployments_named(self, model_name: str) -> tuple[UntypedEntry, ...] | None: - handle: Final = self._live_handle() - key: Final = _index_key(model_name) - indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) - return None if indices is None else self._at(indices) + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(model_name) + indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) + return None if indices is None else self._at(indices) def team_public_model_names_of(self, team_id: object) -> frozenset[str]: handle: Final = self._live_handle() @@ -1719,13 +1723,14 @@ class HostRegistry(HostModelInfo): return frozenset(handle.team_public_names_of(team_id)) def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[UntypedEntry, ...] | None: - handle: Final = self._live_handle() - team: Final = _index_key(team_id) - public: Final = _index_key(public_model_name) - indices: Final[object] = ( - None if handle is None or team is None or public is None else handle.team_positions(team, public) - ) - return None if indices is None else self._at(indices) + with self._deployment_index_lock: + handle: Final = self._live_handle() + team: Final = _index_key(team_id) + public: Final = _index_key(public_model_name) + indices: Final[object] = ( + None if handle is None or team is None or public is None else handle.team_positions(team, public) + ) + return None if indices is None else self._at(indices) # -- lookups (router-core `lookup.rs`): rows read at the stored entries -------------------- @@ -1841,24 +1846,26 @@ class HostRegistry(HostModelInfo): def _team_deployments_across_teams(self, model: str) -> list[UntypedEntry]: # mutable-ok: a fresh list """Every team's deployments under public name ``model``, for a proxy admin calling without a team.""" - rows, teams = self._lookup().lookup_across_teams(model) - team_deployments: Final = self._rows(rows) - if teams > 1: - raise litellm_exceptions.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) - return team_deployments + with self._deployment_index_lock: + rows, teams = self._lookup().lookup_across_teams(model) + team_deployments: Final = self._rows(rows) + if teams > 1: + raise litellm_exceptions.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments def _core_get_all_deployments( self, model_name: str, model_alias: str | None = None, team_id: str | None = None ) -> list[UntypedEntry]: # mutable-ok: a fresh list, as the core returns it - key: Final = _index_key(model_name) - return [] if key is None else self._rows(self._lookup().lookup_all(key, team_id, model_alias)) + with self._deployment_index_lock: + key: Final = _index_key(model_name) + return [] if key is None else self._rows(self._lookup().lookup_all(key, team_id, model_alias)) def get_routing_group(self, model_name: str) -> RoutingGroup | None: """The routing group callable as ``model_name``: a deployment model_name or an alias of the @@ -1877,61 +1884,70 @@ class HostRegistry(HostModelInfo): def _core_routing_group_has_alternatives(self, model_group: str | None) -> bool: """A callable routing group whose member union spans more than one deployment.""" - if model_group is None: - return False - resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group - group: Final = self.get_routing_group(resolved) - if group is None: - return False - handle: Final = self._lookup() + with self._deployment_index_lock: + if model_group is None: + return False + resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group + group: Final = self.get_routing_group(resolved) + if group is None: + return False + handle: Final = self._lookup() - def held(member: str) -> Sequence[object]: - return cast(Sequence[object], handle.named_positions(member) or ()) # cast-ok: positions, or junk + def held(member: str) -> Sequence[object]: + return cast(Sequence[object], handle.named_positions(member) or ()) # cast-ok: positions, or junk - return sum(len(held(member)) for member in group.models) > 1 + return sum(len(held(member)) for member in group.models) > 1 def _core_get_model_list_from_model_alias( self, model_name: str | None = None ) -> list[UntypedEntry]: # mutable-ok: a fresh list - return self._rows( - self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="aliases")) - ) + with self._deployment_index_lock: + return self._rows( + self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="aliases")) + ) def _core_get_model_list_from_routing_groups(self, model_name: str | None = None) -> Sequence[UntypedEntry]: """Callable routing groups as model-list rows: each member deployment under the group's name.""" - if model_name is None and self._routing_group_rows is not None: - return self._routing_group_rows - rows: Final = tuple( - self._rows(self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="groups"))) - ) - if model_name is None: - self._routing_group_rows = rows - return rows + with self._deployment_index_lock: + if model_name is None and self._routing_group_rows is not None: + return self._routing_group_rows + rows: Final = tuple( + self._rows( + self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="groups")) + ) + ) + if model_name is None: + self._routing_group_rows = rows + return rows def _core_get_model_list( self, model_name: str | None = None, team_id: str | None = None ) -> list[UntypedEntry] | None: # mutable-ok: a fresh list, as get_model_list returns - if model_name is not None and _index_key(model_name) is None: - # No index holds it, so `_get_wildcard_deployments` reads it: what the pattern routers raise. - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - return [ - *self._core_pattern_router.get_deployments_by_pattern(model=model_name), - *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), - ] - return self._rows(self._lookup().lookup_model_list(model_name, team_id, self._list_scope(model_name))) + with self._deployment_index_lock: + if model_name is not None and _index_key(model_name) is None: + # No index holds it, so `_get_wildcard_deployments` reads it: what the pattern routers raise. + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + return [ + *self._core_pattern_router.get_deployments_by_pattern(model=model_name), + *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), + ] + return self._rows(self._lookup().lookup_model_list(model_name, team_id, self._list_scope(model_name))) def _core_wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: - from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + with self._deployment_index_lock: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider - if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model reads the base routers - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_owned: Final = team_router.get_deployments_by_pattern(model=model) if team_router is not None else () - return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) - rows: Final = self._lookup().lookup_wildcard_route( - model, team_id, [model, f"{inferred_provider(model)}/{model}"] - ) - routed: Final = cast(Sequence[Mapping[str, object]], rows) # cast-ok: router-core's wildcard rows - return [cast(Mapping[str, object], row["pattern"]) for row in routed] # cast-ok: each carries its routed copy + if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model reads the base routers + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + team_owned: Final = ( + team_router.get_deployments_by_pattern(model=model) if team_router is not None else () + ) + return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) + rows: Final = self._lookup().lookup_wildcard_route( + model, team_id, [model, f"{inferred_provider(model)}/{model}"] + ) + routed: Final = cast(Sequence[Mapping[str, object]], rows) # cast-ok: router-core's wildcard rows + return [cast(Mapping[str, object], row["pattern"]) for row in routed] # cast-ok: its routed copy def _core_wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None @@ -1942,23 +1958,27 @@ class HostRegistry(HostModelInfo): def _core_config_deployments(self) -> Iterator[Mapping[str, object]]: """A generator, as the Python core's: the rows are read when iteration starts.""" - yield from self._rows(self._lookup().lookup_config()) + with self._deployment_index_lock: + rows: Final = self._rows(self._lookup().lookup_config()) + yield from rows # -- names ------------------------------------------------------------------------------------ def _core_get_model_names(self, team_id: str | None = None) -> list[str]: # mutable-ok: a fresh list - names, raw = self._lookup().lookup_model_names(team_id, self._list_scope(None)) - self._raise_at(raw) - return names + with self._deployment_index_lock: + names, raw = self._lookup().lookup_model_names(team_id, self._list_scope(None)) + self._raise_at(raw) + return names def _core_get_model_ids( self, model_name: str | None = None, exclude_team_models: bool = False ) -> list[str]: # mutable-ok: a fresh list - if model_name is not None and _index_key(model_name) is None: - return [] - ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models) - self._raise_at(raw) - return ids + with self._deployment_index_lock: + if model_name is not None and _index_key(model_name) is None: + return [] + ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models) + self._raise_at(raw) + return ids # -- resolving a name ------------------------------------------------------------------------- @@ -2017,12 +2037,6 @@ class HostRegistry(HostModelInfo): return provider_prefixed return None - def _get_deployment_by_litellm_model(self, model: object) -> list[UntypedEntry]: # mutable-ok: a fresh list - handle: Final = self._live_handle() - if handle is None or not isinstance(model, str): - return [] - return list(self._at(handle.litellm_model_positions(model))) - # -- credentials ------------------------------------------------------------------------------ @staticmethod @@ -2041,18 +2055,20 @@ class HostRegistry(HostModelInfo): def _get_model_group_deployment_usable_by_team( self, model_group_name: str, team_id: str | None ) -> Deployment | None: - if _index_key(model_group_name) is None: - return None - position, raw = self._lookup().lookup_group_usable_by_team(model_group_name, team_id) - self._raise_at(raw) - return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + if _index_key(model_group_name) is None: + return None + position, raw = self._lookup().lookup_group_usable_by_team(model_group_name, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: - if team_id is None or _index_key(model_id) is None: - return None - position, raw = self._lookup().lookup_team_public_name(model_id, team_id) - self._raise_at(raw) - return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + if team_id is None or _index_key(model_id) is None: + return None + position, raw = self._lookup().lookup_team_public_name(model_id, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None @@ -2105,7 +2121,10 @@ class HostRegistry(HostModelInfo): # -- routable names --------------------------------------------------------------------------- def _routable_deployments(self, model_name: str, team_id: str | None) -> tuple[UntypedEntry, ...]: - return tuple(self._rows(self._lookup().lookup_routable(self._core_routable_model_group(model_name), team_id))) + with self._deployment_index_lock: + return tuple( + self._rows(self._lookup().lookup_routable(self._core_routable_model_group(model_name), team_id)) + ) def _core_team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, @@ -2169,24 +2188,26 @@ class HostRegistry(HostModelInfo): """The deployment ids that could serve ``model`` for ``team_id``: the alias target, then a routing group, then a team public name or wildcard route for a name that is no model group, then the model group (``get_candidate_model_ids_for_route``). Read-only.""" - from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + with self._deployment_index_lock: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider - resolved: Final = self._core_get_model_from_alias(model=model) or model - group: Final = self.get_routing_group(resolved) if self._core_routing_groups else None - ids, raw = self._lookup().lookup_candidate_ids( - resolved, - team_id, - None if group is None else list(group.models), - [resolved, f"{inferred_provider(resolved)}/{resolved}"], - ) - self._raise_at(raw) - return frozenset(ids) + resolved: Final = self._core_get_model_from_alias(model=model) or model + group: Final = self.get_routing_group(resolved) if self._core_routing_groups else None + ids, raw = self._lookup().lookup_candidate_ids( + resolved, + team_id, + None if group is None else list(group.models), + [resolved, f"{inferred_provider(resolved)}/{resolved}"], + ) + self._raise_at(raw) + return frozenset(ids) def _core_team_model_has_alternatives(self, deployment_id: str) -> bool: """Whether a team deployment's public name has another unpaused deployment of that team.""" - found, raw = self._lookup().lookup_team_model_has_alternatives(deployment_id) - self._raise_at(raw) - return found + with self._deployment_index_lock: + found, raw = self._lookup().lookup_team_model_has_alternatives(deployment_id) + self._raise_at(raw) + return found def _core_resolve_served_model_name(self, model_name: str, served_model: str) -> str: """``served_model`` when ``model_name`` routes to it (the alias target, a routing group @@ -2315,17 +2336,18 @@ class HostRegistry(HostModelInfo): def _core_get_model_access_groups( self, model_name: str | None = None, model_access_group: str | None = None, team_id: str | None = None ) -> Mapping[str, Sequence[str]]: - use_cache: Final = model_name is None and model_access_group is None and team_id is None - if use_cache and self._access_groups_cache is not None: - return self._access_groups_cache - pairs, raw = self._lookup().lookup_access_groups( - model_name, model_access_group, team_id, self._list_scope(model_name) - ) - self._raise_at(raw) - groups: Final[dict[str, list[str]]] = {g: list(n) for g, n in pairs} # mutable-ok: as a defaultdict - if use_cache: - self._access_groups_cache = groups - return groups + with self._deployment_index_lock: + use_cache: Final = model_name is None and model_access_group is None and team_id is None + if use_cache and self._access_groups_cache is not None: + return self._access_groups_cache + pairs, raw = self._lookup().lookup_access_groups( + model_name, model_access_group, team_id, self._list_scope(model_name) + ) + self._raise_at(raw) + groups: Final[dict[str, list[str]]] = {g: list(n) for g, n in pairs} # mutable-ok: as a defaultdict + if use_cache: + self._access_groups_cache = groups + return groups def _core_is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: groups: Final = self._core_get_model_access_groups(model_access_group=model_access_group) diff --git a/tests/unit/test_router/core/test_registry_characterization_lookups.py b/tests/unit/test_router/core/test_registry_characterization_lookups.py index 19d58ff10d4..94c88de73fe 100644 --- a/tests/unit/test_router/core/test_registry_characterization_lookups.py +++ b/tests/unit/test_router/core/test_registry_characterization_lookups.py @@ -4,6 +4,7 @@ access groups; blocked and unhealthy groups; credentials lookups. """ import time +from typing import Final import pytest @@ -680,3 +681,17 @@ def test_assigning_model_list_swaps_the_list_without_reindexing(make_core: MakeC found = core.registry().deployment("a") assert found is not None assert found.model_name == "n" + + +def test_the_routing_groups_source_reads_its_rows_once(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + routing_groups=[{"group_name": "grp", "models": ["g"], "routing_strategy": "simple-shuffle"}], + ) + first: Final = core.registry().deployments(source="routing_groups") + assert ids(first) == ["a", "b"] + assert core.registry().deployments(source="routing_groups") is first + + +def test_the_credential_deployment_of_a_non_string_id_is_none(make_core: MakeCore) -> None: + assert _core(make_core).registry().credentials(123, form="deployment") is None # pyright: ignore[reportArgumentType] # a caller's non-string id diff --git a/tests/unit/test_router/core/test_registry_characterization_writes.py b/tests/unit/test_router/core/test_registry_characterization_writes.py index 0558a16fe9d..23fa2a8e0e9 100644 --- a/tests/unit/test_router/core/test_registry_characterization_writes.py +++ b/tests/unit/test_router/core/test_registry_characterization_writes.py @@ -5,7 +5,11 @@ what ``add_deployment``, ``set_model_list``, upserts and deletes leave behind in and ``registry()`` views. """ +import threading import warnings +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from typing import Final import pytest @@ -514,3 +518,80 @@ async def test_the_deprecated_arefresh_model_info_warns(make_core: MakeCore) -> await core.arefresh_model_info() # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public assert any("arefresh_model_info is deprecated" in str(w.message) for w in caught) + + +# -- reads racing writes ---------------------------------------------------------------------------- + +_GROUPS: Final = 4 + + +def _two_per_group(make_core: MakeCore) -> RouterCore: + return make_core(model_list=[deployment(f"m{index % _GROUPS}", f"d{index}") for index in range(2 * _GROUPS)]) + + +def _group_ids(core: RouterCore) -> Mapping[int, Sequence[str]]: + return {group: ids(core.registry().deployments(f"m{group}", source="all")) for group in range(_GROUPS)} + + +@pytest.mark.parametrize("upserted", [0, 3, _GROUPS + 1]) +def test_a_read_on_another_thread_waits_for_an_upsert_to_read_the_list_back(make_core: MakeCore, upserted: int) -> None: + """A lookup's positions are read against the mirror of the same registry version: a read that + starts once router-core holds a write waits until the mirror is read back, so it never reads + past the list nor, after a removal, the neighbouring group's entries.""" + core: Final = _two_per_group(make_core) + hosted: Final = core._hosted # pyright: ignore[reportAttributeAccessIssue] # the registry under test + write: Final[Callable[[Mapping[str, object]], None]] = hosted.changes_runner + started: Final = threading.Event() + pool: Final = ThreadPoolExecutor(max_workers=1) + reads: Final[list[object]] = [] # mutable-ok: one per write + + def read() -> Mapping[int, Sequence[str]]: + started.set() + return _group_ids(core) + + def write_with_a_read_started(changes: Mapping[str, object]) -> None: + write(changes) + started.clear() + reads.append(pool.submit(read)) + assert started.wait(timeout=5) + # The reader runs while router-core holds the write and the mirror does not yet. + _ = threading.Event().wait(0.02) + + hosted.changes_runner = write_with_a_read_started + try: + _ = core.apply(ChangeSet(upsert=[typed(f"m{upserted % _GROUPS}", f"d{upserted}", timeout=11)])) + finally: + hosted.changes_runner = write + + assert reads + for future in reads: + seen = future.result(timeout=5) # pyright: ignore[reportAttributeAccessIssue] # a Future + for group in range(_GROUPS): + # Between the upsert's removal and its add, its group is served by the other deployment. + assert seen[group], (group, seen) + assert set(seen[group]) <= {f"d{group}", f"d{group + _GROUPS}"}, (group, seen) + pool.shutdown() + assert sorted(_group_ids(core)[upserted % _GROUPS]) == sorted( + [f"d{upserted % _GROUPS}", f"d{upserted % _GROUPS + _GROUPS}"] + ) + + +def test_appending_to_the_live_model_list_reads_the_entry_back(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + live.append(deployment("m0", "d8")) + assert ids(live) == [*(f"d{index}" for index in range(2 * _GROUPS)), "d8"] + + +def test_deleting_a_live_model_list_entry_reads_the_list_back(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + del live[0] + assert ids(live) == [f"d{index}" for index in range(1, 2 * _GROUPS)] + + +def test_ending_a_build_that_never_began_leaves_the_list(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + live.end_build() + assert ids(live) == [f"d{index}" for index in range(2 * _GROUPS)]