diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index ab93e4d37..d09c1aa88 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -105,6 +105,46 @@ Attribute descriptions for ordinary attributes round-trip through `DataFrame.att Collection coordinate units come only from `DistanceCollection.unit` or `IntervalCollection.unit`; unit metadata on the distance or `from` column is ignored. +### Upload concurrency + +Typed objects upload independent tables and attributes concurrently, preserving input order and applying mutations only +after all required uploads succeed. Configure the upload limit for a block, including its child tasks: + +```python +from evo.objects.utils import upload_concurrency + +with upload_concurrency(max_concurrent=8): + dhc = await DownholeCollection.create(context, data) +``` + +The default limit is **4** and is shared by uploads on the same event loop, including direct `ObjectDataClient` calls. +Unlike uploads, `download_tables(max_concurrency=...)` applies a separate limit to each call. The upload limit covers +preparation and transfers; a storage retry may hold a slot. Nested scopes restore the previous setting. +HTTP connection limits are configured separately with `AioTransport(num_pools=...)`. + +Uploads retry transient failures up to three times by default; pass `retry=Retry(...)` to an `ObjectDataClient` upload +method to customize that policy. On exhaustion, the last failure is raised and the retry history is chained as its +cause. Storage transfers make one attempt per outer retry, preventing retry layers from multiplying requests. +For custom interception, override `upload_dataframe` or `upload_category_dataframe` for DataFrame inputs; these methods +do not dispatch through the table upload methods. + +To configure a separate HTTP connection limit, pass a configured transport to your context using your existing +authorizer, cache, and environment: + +```python +from evo.aio import AioTransport +from evo.common import StaticContext + +transport = AioTransport(user_agent="my-app", num_pools=8) +context = StaticContext( + transport=transport, authorizer=authorizer, cache=cache, + hub_url=environment.hub_url, org_id=environment.org_id, workspace_id=environment.workspace_id, +) +``` + +`ServiceManager` and `ServiceManagerWidget` also accept a transport in their constructors. +The `ServiceManagerWidget.with_auth_code(...)` convenience factory currently uses the default connection limit. + ## Contributing For instructions on contributing to the development of this library, please refer to the [evo-python-sdk documentation](https://github.com/seequentevo/evo-python-sdk). diff --git a/packages/evo-objects/pyproject.toml b/packages/evo-objects/pyproject.toml index 84d5b2f00..27d5dfb0f 100644 --- a/packages/evo-objects/pyproject.toml +++ b/packages/evo-objects/pyproject.toml @@ -10,7 +10,7 @@ authors = [ ] dependencies = [ - "evo-sdk-common[jmespath]>=0.5.31", + "evo-sdk-common[jmespath]>=0.5.32", "pydantic>=2,<3", ] diff --git a/packages/evo-objects/src/evo/objects/typed/_data.py b/packages/evo-objects/src/evo/objects/typed/_data.py index 8868c9afe..cabee40d9 100644 --- a/packages/evo-objects/src/evo/objects/typed/_data.py +++ b/packages/evo-objects/src/evo/objects/typed/_data.py @@ -17,8 +17,9 @@ from evo.common import IFeedback from evo.common.interfaces import IContext -from evo.common.utils import NoFeedback +from evo.common.utils import NoFeedback, split_feedback from evo.objects.typed.attributes import Attributes +from evo.objects.utils._concurrency import gather_cancelling from evo.objects.utils.table_formats import KnownTableFormat from ._model import SchemaBuilder, SchemaLocation, SchemaModel @@ -49,9 +50,15 @@ async def from_dataframe(self, df: pd.DataFrame, fb: IFeedback = NoFeedback) -> :param df: DataFrame containing the new values for this table. :param fb: Optional feedback object to report upload progress. """ - self._document.update(await self._data_to_schema(df, self._context, fb=fb)) + self._apply_uploaded_table(df, await self._upload(df, fb)) - # Mark the context as modified so loading data is not allowed + async def _upload(self, df: pd.DataFrame, fb: IFeedback = NoFeedback) -> dict[str, Any]: + """Upload replacement values without changing the table or its derived metadata.""" + return await self._data_to_schema(df, self._obj, fb=fb) + + def _apply_uploaded_table(self, df: pd.DataFrame, document: dict[str, Any]) -> None: + """Commit uploaded values and derived metadata; override to extend table replacement behavior.""" + self._document.update(document) self._context.mark_modified(self._data) @classmethod @@ -118,11 +125,17 @@ async def from_dataframe(self, df: pd.DataFrame, fb: IFeedback = NoFeedback) -> """ table_df, attr_df = self._split_dataframe(df, self._table.data_columns) - await self._table.from_dataframe(table_df, fb=fb) - if attr_df is not None: - await self.attributes.set_attributes(attr_df, fb=fb) - else: - self.attributes.clear() + if attr_df is None: + attr_df = pd.DataFrame(index=table_df.index) + table_fb, attr_fb = split_feedback(fb, [1.0, float(len(attr_df.columns))]) + descriptions, attribute_types = await self.attributes._validate(attr_df) + table_upload, attribute_upload = ( + self._table._upload(table_df, table_fb), + self.attributes._upload(attr_df, descriptions, attribute_types, attr_fb), + ) + table_document, attribute_documents = await gather_cancelling(table_upload, attribute_upload) + self._table._apply_uploaded_table(table_df, table_document) + self.attributes._apply(attribute_documents) def validate(self) -> None: """Validate that all attributes have the correct length.""" @@ -152,8 +165,7 @@ async def _data_to_schema(cls, data: Any, context: IContext) -> Any: table_df, attr_df = cls._split_dataframe(data, table_type.data_columns) builder = SchemaBuilder(cls, context) - await builder.set_sub_model_value("_table", table_df) - await builder.set_sub_model_value("attributes", attr_df) + await builder.set_sub_model_values([("_table", table_df), ("attributes", attr_df)]) # As we specially handle the _table and attributes sub-models above, we skip them in the default data processing of build_from_data. return await builder.build_from_data(data, skip_sub_models={"_table", "attributes"}) diff --git a/packages/evo-objects/src/evo/objects/typed/_model.py b/packages/evo-objects/src/evo/objects/typed/_model.py index 85d0fb610..b971ab03b 100644 --- a/packages/evo-objects/src/evo/objects/typed/_model.py +++ b/packages/evo-objects/src/evo/objects/typed/_model.py @@ -21,6 +21,7 @@ from evo import jmespath from evo.common import IContext from evo.objects import DownloadedObject +from evo.objects.utils._concurrency import gather_cancelling from ._utils import ( assign_jmespath_value, @@ -185,14 +186,6 @@ def set_property(self, name: str, value: Any) -> None: schema_property = self._properties[name] _set_property_value(schema_property, self.document, value) - async def set_sub_model_value(self, name: str, data: Any) -> None: - metadata = self._sub_models[name] - sub_document = await metadata.model_type._data_to_schema(data, context=self._context) - if metadata.jmespath_expr: - assign_jmespath_value(self.document, metadata.jmespath_expr, sub_document) - else: - self.document.update(sub_document) - async def build_from_data(self, data: Any, skip_sub_models: set[str] | None = None) -> dict[str, Any]: """Build the schema document by applying all properties and sub-models.""" for key, prop in self._properties.items(): @@ -202,6 +195,7 @@ async def build_from_data(self, data: Any, skip_sub_models: set[str] | None = No value = getattr(data, key, None) self.set_property(key, value) + pending: list[tuple[str, Any]] = [] for name, metadata in self._sub_models.items(): if skip_sub_models and name in skip_sub_models: continue @@ -209,10 +203,29 @@ async def build_from_data(self, data: Any, skip_sub_models: set[str] | None = No sub_data = getattr(data, metadata.data_field, None) else: sub_data = data - await self.set_sub_model_value(name, sub_data) + pending.append((name, sub_data)) + await self.set_sub_model_values(pending) return self.document + async def set_sub_model_values(self, values: Sequence[tuple[str, Any]]) -> None: + """Convert several sub-models concurrently, then assign the results in the given order. + + Sub-models are independent, so their data uploads can run in parallel. + """ + sub_documents = await gather_cancelling( + *(self._sub_models[name].model_type._data_to_schema(data, context=self._context) for name, data in values) + ) + for (name, _), sub_document in zip(values, sub_documents): + self._assign_sub_document(name, sub_document) + + def _assign_sub_document(self, name: str, sub_document: Any) -> None: + metadata = self._sub_models[name] + if metadata.jmespath_expr: + assign_jmespath_value(self.document, metadata.jmespath_expr, sub_document) + else: + self.document.update(sub_document) + def _get_base_type(annotation: Any) -> tuple[Any, SchemaLocation | None, DataLocation | None]: """Extract the base type and SchemaLocation from an annotation. @@ -497,4 +510,4 @@ async def _data_to_schema(cls, data: Any, context: IContext) -> list[Any]: return [] if not isinstance(data, Sequence): raise TypeError(f"Expected a sequence for SchemaList data, got {type(data).__name__}") - return [await cls._item_type._data_to_schema(item, context) for item in data] + return await gather_cancelling(*(cls._item_type._data_to_schema(item, context) for item in data)) diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index 3f35a4ff8..457bb3679 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -13,6 +13,7 @@ import typing import uuid +from copy import deepcopy from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any, Protocol, runtime_checkable from uuid import UUID @@ -21,8 +22,9 @@ from evo import jmespath from evo.common import IContext, IFeedback -from evo.common.utils import NoFeedback, iter_with_fb +from evo.common.utils import NoFeedback, iter_with_fb, split_feedback from evo.objects import DownloadedObject +from evo.objects.utils._concurrency import gather_cancelling, run_upload_preparation from evo.objects.utils.table_formats import ( BOOL_ARRAY_1, DATE_TIME_ARRAY, @@ -37,6 +39,8 @@ from .exceptions import DataLoaderError, ObjectValidationError if TYPE_CHECKING: + from evo.objects.utils import ObjectDataClient + from .block_model_ref import BlockModel __all__ = [ @@ -210,7 +214,12 @@ async def _upload_attribute_values( attr_doc.update(await data_client.upload_category_dataframe(df, fb=fb)) else: table_formats = _attribute_table_formats.get(attribute_type) - attr_doc["values"] = await data_client.upload_dataframe(df, table_format=table_formats, fb=fb) + attr_doc["values"] = await data_client.upload_dataframe( + df, + table_format=table_formats, + fb=fb, + _prepare_dataframe=_prepare_datetime_dataframe if attribute_type == "date_time" else None, + ) if attribute_type in ["scalar", "integer", "category", "date_time"]: attr_doc["nan_description"] = {"values": []} @@ -256,6 +265,37 @@ def __repr__(self) -> str: return f"PendingAttribute(name={self._name!r}, exists=False)" +async def _build_new_attribute_document( + col: Any, + series: pd.Series, + attribute_type: str, + description: dict[str, Any] | None, + data_client: ObjectDataClient, + fb: IFeedback, +) -> dict[str, Any]: + document: dict[str, Any] = { + "name": str(col), + "key": str(uuid.uuid4()), + "attribute_type": attribute_type, + } + await Attribute._upload_attribute_values(document, series.to_frame(name=col), attribute_type, data_client, fb=fb) + if description is not None: + document["attribute_description"] = description + return document + + +def _prepare_datetime_dataframe(dataframe: pd.DataFrame) -> pd.DataFrame: + series = pd.to_datetime(dataframe.iloc[:, 0], utc=True).dt.as_unit("us") + return series.to_frame(name=dataframe.columns[0]) + + +@dataclass +class _PreparedAttributes: + documents: list[tuple[dict[str, Any] | None, dict[str, Any]]] + original_documents: tuple[dict[str, Any], ...] + original_snapshot: tuple[dict[str, Any], ...] + + class Attributes(SchemaList[Attribute]): """A collection of Geoscience Object Attributes @@ -295,58 +335,67 @@ async def _data_to_schema( ) -> list[dict[str, Any]]: """Convert a DataFrame to a list of attribute dictionaries for object creation. - :param df: The DataFrame with columns to convert to attributes, or None. + :param data: The DataFrame with columns to convert to attributes, or None. :param context: The context used for data upload operations. - :param fb: Optional feedback object to report progress. :return: A list of attribute dictionaries suitable for the object document. """ - result: list[dict[str, Any]] = [] - if data is not None: - await cls._upload_attributes_to_list(result, data, context) - return result + return await cls._upload_attributes(data, context) if data is not None else [] @staticmethod - async def _upload_attributes_to_list( - attributes_list: list[dict[str, Any]], + async def _validate_dataframe( + df: pd.DataFrame, existing: dict[str, Attribute] | None = None + ) -> tuple[dict[Any, dict[str, Any]], dict[Any, str]]: + """Validate metadata and infer every new column type before starting uploads.""" + existing = existing or {} + descriptions = df.attrs.get("attribute_descriptions", {}) + validated_descriptions = {} + for col in df.columns: + if str(col) in existing: + continue + description = descriptions.get(col) + if description is not None: + if not isinstance(description, AttributeDescription): + raise TypeError("attribute description must be a AttributeDescription.") + validated_descriptions[col] = description.to_schema() + new_columns = [col for col in df.columns if str(col) not in existing] + inferred = await gather_cancelling( + *(run_upload_preparation(_infer_attribute_type_from_series, df[col]) for col in new_columns) + ) + new_types = dict(zip(new_columns, inferred)) + attribute_types = { + col: existing[str(col)].attribute_type if str(col) in existing else new_types[col] for col in df.columns + } + return validated_descriptions, attribute_types + + @staticmethod + async def _upload_attributes( df: pd.DataFrame, context: IContext | DownloadedObject, fb: IFeedback = NoFeedback, - ) -> None: - """Upload DataFrame columns as attributes and append to the attributes list. + ) -> list[dict[str, Any]]: + """Upload DataFrame columns and return new attribute documents in column order. This is a static operation that doesn't require a model instance. - :param attributes_list: The list in the document to append attribute dicts to. :param df: The DataFrame with columns to upload as attributes. :param context: The context used for data upload operations. :param fb: Optional feedback object to report progress. """ + if len(df.columns) == 0: + fb.progress(1) + return [] + descriptions, attribute_types = await Attributes._validate_dataframe(df) data_client = get_data_client(context) + column_fbs = split_feedback(fb, [1.0] * len(df.columns)) - for col in df.columns: - series = df[col] - attribute_type = _infer_attribute_type_from_series(df[col]) - - attr_doc: dict[str, Any] = { - "name": str(col), - "key": str(uuid.uuid4()), - "attribute_type": attribute_type, - } - - if attribute_type == "date_time": - series = pd.to_datetime(series, utc=True).dt.as_unit("us") - - col_df = pd.DataFrame({col: series}) - await Attribute._upload_attribute_values(attr_doc, col_df, attribute_type, data_client) - - attr_desc: AttributeDescription | None = df.attrs.get("attribute_descriptions", {}).get(col) - if attr_desc is not None: - if not isinstance(attr_desc, AttributeDescription): - raise TypeError("attribute description must be a AttributeDescription.") - if description := attr_desc.to_schema(): - attr_doc["attribute_description"] = description - - attributes_list.append(attr_doc) + return await gather_cancelling( + *( + _build_new_attribute_document( + col, df[col], attribute_types[col], descriptions.get(col), data_client, col_fb + ) + for col, col_fb in zip(df.columns, column_fbs) + ) + ) async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: """Load a DataFrame containing the values from the specified attributes in the object. @@ -382,12 +431,9 @@ async def append_attribute(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): if df.shape[1] != 1: raise ValueError("DataFrame must contain exactly one column to append as an attribute.") - # Use the static method to upload and append to document - await self._upload_attributes_to_list(self._document, df, self._obj, fb) - - # Mark context as modified and create the Attribute wrapper - attribute = self[-1] - self._context.mark_modified(attribute._data) + documents = await self._upload_attributes(df, self._obj, fb) + self._document.extend(documents) + self._context.mark_modified(documents[0]["values"]["data"]) async def append_attributes(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): """Add new attributes to the object. @@ -395,9 +441,10 @@ async def append_attributes(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): :param df: DataFrame containing the values for the new attributes. :param fb: Optional feedback object to report upload progress. """ - for attribute in df.columns: - attribute_df = df[[attribute]] - await self.append_attribute(attribute_df, fb) + documents = await self._upload_attributes(df, self._obj, fb) + self._document.extend(documents) + for document in documents: + self._context.mark_modified(document["values"]["data"]) async def set_attributes(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): """Set the attributes of the object to match the provided DataFrame. @@ -406,16 +453,68 @@ async def set_attributes(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): :param fb: Optional feedback object to report upload progress. """ + descriptions, attribute_types = await self._validate(df) + prepared = await self._upload(df, descriptions, attribute_types, fb) + self._apply(prepared) + + async def _validate(self, df: pd.DataFrame) -> tuple[dict[Any, dict[str, Any]], dict[Any, str]]: + return await self._validate_dataframe(df, {attr.name: attr for attr in self}) + + async def _upload( + self, + df: pd.DataFrame, + descriptions: dict[Any, dict[str, Any]], + attribute_types: dict[Any, str], + fb: IFeedback, + ) -> _PreparedAttributes: + """Upload replacement values without changing live documents or modification tracking.""" + originals = tuple(self._document) + original_snapshot = tuple(deepcopy(originals)) + if len(df.columns) == 0: + fb.progress(1) + return _PreparedAttributes([], originals, original_snapshot) attributes_by_name = {attr.name: attr for attr in self} - self.clear() - for col in df.columns: - attribute_df = df[[col]] - attribute = attributes_by_name.get(col) + data_client = get_data_client(self._obj) + column_fbs = split_feedback(fb, [1.0] * len(df.columns)) + + async def upload_column(col: Any, col_fb: IFeedback) -> tuple[dict[str, Any] | None, dict[str, Any]]: + attribute = attributes_by_name.get(str(col)) if attribute is not None: - await attribute.set_attribute_values(attribute_df, fb=fb) - self._append(attribute) - else: - await self.append_attribute(attribute_df, fb=fb) + document = attribute.as_dict() + await Attribute._upload_attribute_values( + document, df[col].to_frame(name=col), attribute_types[col], data_client, fb=col_fb + ) + return attribute._document, document + document = await _build_new_attribute_document( + col, df[col], attribute_types[col], descriptions.get(col), data_client, col_fb + ) + return None, document + + documents = await gather_cancelling( + *(upload_column(col, col_fb) for col, col_fb in zip(df.columns, column_fbs)) + ) + return _PreparedAttributes(documents, originals, original_snapshot) + + def _apply(self, prepared: _PreparedAttributes) -> None: + """Replace live attributes after uploads, preserving wrappers and rejecting concurrent edits.""" + if ( + len(self._document) != len(prepared.original_documents) + or any(current is not original for current, original in zip(self._document, prepared.original_documents)) + or tuple(self._document) != prepared.original_snapshot + ): + raise RuntimeError("Attributes changed while replacement values were uploading") + documents = [] + for original, document in prepared.documents: + if original is not None: + previous_data = original.get("values", {}).get("data") + original.clear() + original.update(document) + document = original + documents.append(document) + if original is not None and previous_data == document.get("values", {}).get("data"): + continue + self._context.mark_modified(document["values"]["data"]) + self._document[:] = documents def validate_lengths(self, expected_length: int) -> None: """Validate that all attributes have the expected length. diff --git a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py index d5d2c83e6..9eea3a847 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -32,6 +32,7 @@ from evo.objects.typed.exceptions import ObjectValidationError from evo.objects.typed.spatial import BaseSpatialObject, BaseSpatialObjectData from evo.objects.typed.types import BoundingBox +from evo.objects.utils._concurrency import gather_cancelling from evo.objects.utils.table_formats import ( DOWNHOLE_COLLECTION_LOCATION_HOLES, FLOAT_ARRAY_1, @@ -353,14 +354,13 @@ async def _data_to_schema(cls, data: pd.DataFrame, context: IContext) -> Any: raise ObjectValidationError("properties hole_id values must be unique") keys = _property_hole_indices(data).astype("int32") data_client = _attributes.get_data_client(context) - return { - "values": await data_client.upload_dataframe( - pd.DataFrame({"hole_id": keys}), table_format=INTEGER_ARRAY_1_INT32 - ), - "table": await data_client.upload_dataframe( + values, table = await gather_cancelling( + data_client.upload_dataframe(pd.DataFrame({"hole_id": keys}), table_format=INTEGER_ARRAY_1_INT32), + data_client.upload_dataframe( pd.DataFrame({"key": keys, "value": hole_ids}), table_format=LOOKUP_TABLE_INT32 ), - } + ) + return {"values": values, "table": table} class DownholeLocation(SchemaModel): @@ -477,8 +477,8 @@ def _resolve_item_type(cls, document: dict[str, Any]) -> type[DownholeDistanceTa async def _data_to_schema(cls, data: Any, context: IContext) -> list[Any]: if data is None: return [] - result = [] - for collection in data: + + async def convert(collection: DownholeCollectionEntry) -> dict[str, Any]: model = DownholeIntervalTable if isinstance(collection, IntervalCollection) else DownholeDistanceTable schema = await model._data_to_schema(collection, context) if collection.unit is not None: @@ -486,8 +486,10 @@ async def _data_to_schema(cls, data: Any, context: IContext) -> list[Any]: schema["from_to"]["unit"] = collection.unit else: schema["distance"]["unit"] = collection.unit - result.append(schema) - return result + return schema + + # Collections are independent, so upload them concurrently while preserving input order. + return await gather_cancelling(*(convert(collection) for collection in data)) def _indices_for_name(self, name: str) -> list[int]: return [index for index, document in enumerate(self._document) if document.get("name") == name] diff --git a/packages/evo-objects/src/evo/objects/typed/pointset.py b/packages/evo-objects/src/evo/objects/typed/pointset.py index 3141a8fc5..1140c6e5c 100644 --- a/packages/evo-objects/src/evo/objects/typed/pointset.py +++ b/packages/evo-objects/src/evo/objects/typed/pointset.py @@ -12,7 +12,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Annotated, ClassVar +from typing import Annotated, Any, ClassVar import pandas as pd @@ -78,16 +78,10 @@ class CoordinateTable(DataTable): table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 data_columns: ClassVar[list[str]] = _COORDINATE_COLUMNS - async def from_dataframe(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): - """Update the coordinate values and recalculate the bounding box. - - :param df: DataFrame containing x, y, z coordinate columns. - :param fb: Optional feedback object to report upload progress. - """ - await super().from_dataframe(df, fb) - - # Update the bounding box in the parent object context - self._context.root_model.bounding_box = _bounding_box_from_dataframe(df) + def _apply_uploaded_table(self, df: pd.DataFrame, document: dict[str, Any]) -> None: + bounding_box = _bounding_box_from_dataframe(df) + super()._apply_uploaded_table(df, document) + self._context.root_model.bounding_box = bounding_box class Locations(DataTableAndAttributes): diff --git a/packages/evo-objects/src/evo/objects/utils/__init__.py b/packages/evo-objects/src/evo/objects/utils/__init__.py index af0879537..4c25df89d 100644 --- a/packages/evo-objects/src/evo/objects/utils/__init__.py +++ b/packages/evo-objects/src/evo/objects/utils/__init__.py @@ -22,6 +22,7 @@ except ImportError: DataFrame = None # type: ignore +from ._concurrency import upload_concurrency from .data import ObjectDataClient from .downhole import expand_hole_index, hole_chunks_from_ids from .table_formats import all_known_formats, get_known_format @@ -47,4 +48,5 @@ "expand_hole_index", "get_known_format", "hole_chunks_from_ids", + "upload_concurrency", ] diff --git a/packages/evo-objects/src/evo/objects/utils/_concurrency.py b/packages/evo-objects/src/evo/objects/utils/_concurrency.py new file mode 100644 index 000000000..09a6b35e6 --- /dev/null +++ b/packages/evo-objects/src/evo/objects/utils/_concurrency.py @@ -0,0 +1,122 @@ +# Copyright © 2026 Bentley Systems, Incorporated +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Concurrency shared by table serialization and uploads on each event loop.""" + +import asyncio +from collections.abc import Callable, Coroutine, Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from threading import Lock +from typing import Any, ParamSpec, TypeVar +from weakref import ReferenceType, WeakKeyDictionary, ref + +from evo import logging + +_T = TypeVar("_T") +_P = ParamSpec("_P") +_DEFAULT_MAX_CONCURRENCY = 4 # Downloads and uploads share the same default. +logger = logging.getLogger("object.uploads") + + +class _UploadScope: + """Share a limit across clients without retaining closed event loops. + + Contended semaphores retain their loop, so both loop keys and semaphore values + must be weak references. Active work keeps its semaphore alive. + """ + + def __init__(self, limit: int) -> None: + self.limit = limit + self._lock = Lock() + self.semaphores: WeakKeyDictionary[asyncio.AbstractEventLoop, ReferenceType[asyncio.Semaphore]] = ( + WeakKeyDictionary() + ) + + def get_semaphore(self) -> asyncio.Semaphore: + loop = asyncio.get_running_loop() + # A default scope can be inherited by tasks running on loops in separate threads. + with self._lock: + reference = self.semaphores.get(loop) + semaphore = reference() if reference is not None else None + if semaphore is None: + semaphore = asyncio.Semaphore(self.limit) + self.semaphores[loop] = ref(semaphore) + return semaphore + + +_upload_scope: ContextVar[_UploadScope] = ContextVar( + "object_upload_scope", default=_UploadScope(_DEFAULT_MAX_CONCURRENCY) +) + + +@contextmanager +def upload_concurrency(max_concurrent: int = _DEFAULT_MAX_CONCURRENCY) -> Iterator[None]: + """Limit concurrent table preparation and uploads within this scope. + + Applies to typed-object operations and direct ObjectDataClient uploads, including + clients created internally. Conversion/serialization and each upload attempt acquire a slot; + storage-level backoff may hold a slot. Settings are inherited by child asyncio tasks, not unrelated + tasks. Each event loop has its own semaphore, so a scope can span asyncio.run calls. + Outside an explicit scope, all clients and operations on a loop share the default limit. + + :param max_concurrent: Maximum number of concurrent preparation or upload operations. Defaults to four. + :raises ValueError: If max_concurrent is not a positive integer. + """ + if isinstance(max_concurrent, bool) or not isinstance(max_concurrent, int) or max_concurrent < 1: + raise ValueError("max_concurrent must be a positive integer") + token = _upload_scope.set(_UploadScope(max_concurrent)) + try: + yield + finally: + _upload_scope.reset(token) + + +def get_upload_semaphore() -> asyncio.Semaphore: + """Return the shared semaphore; only leaf work may hold a slot, never a parent awaiting child uploads.""" + return _upload_scope.get().get_semaphore() + + +async def run_upload_preparation(function: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs) -> _T: + """Convert and serialize in a bounded worker, draining it before releasing a cancelled slot.""" + async with get_upload_semaphore(): + worker = asyncio.create_task(asyncio.to_thread(function, *args, **kwargs)) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + # Repeated cancellation must not release the slot while the thread writes. + while not worker.done(): + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + continue + except Exception: + break + raise + + +async def gather_cancelling(*coroutines: Coroutine[Any, Any, _T]) -> list[_T]: + """Collect ordered results, cancelling and draining siblings on failure or cancellation.""" + tasks: list[asyncio.Task[_T]] = [] + try: + for coroutine in coroutines: + tasks.append(asyncio.create_task(coroutine)) + return list(await asyncio.gather(*tasks)) + except BaseException as error: + for coroutine in coroutines[len(tasks) :]: + coroutine.close() + for task in tasks: + task.cancel() + results = await asyncio.gather(*tasks, return_exceptions=True) + for result in results: + if isinstance(result, Exception) and result is not error: + logger.debug("Additional upload failure", exc_info=(type(result), result, result.__traceback__)) + raise diff --git a/packages/evo-objects/src/evo/objects/utils/data.py b/packages/evo-objects/src/evo/objects/utils/data.py index e698663c9..0c7a8cfca 100644 --- a/packages/evo-objects/src/evo/objects/utils/data.py +++ b/packages/evo-objects/src/evo/objects/utils/data.py @@ -10,21 +10,23 @@ # limitations under the License. import asyncio -from collections.abc import Iterator, Sequence +from collections.abc import Callable, Iterator, Sequence from pathlib import Path from typing import Any, Iterable, cast from uuid import UUID +from weakref import WeakKeyDictionary import numpy as np from evo import logging from evo.common import APIConnector, Environment, ICache, IFeedback, StaticContext -from evo.common.exceptions import StorageFileNotFoundError, error_context -from evo.common.io.exceptions import DataExistsError -from evo.common.utils import NoFeedback, PartialFeedback, split_feedback +from evo.common.exceptions import EvoAPIException, RetryError, StorageFileNotFoundError, TransportError, error_context +from evo.common.io.exceptions import ChunkedIOError, DataExistsError +from evo.common.utils import NoFeedback, PartialFeedback, Retry, split_feedback from ..exceptions import TableFormatError from ..io import _CACHE_SCOPE, ObjectDataDownload, ObjectDataUpload +from ._concurrency import _DEFAULT_MAX_CONCURRENCY, gather_cancelling, get_upload_semaphore, run_upload_preparation from .table_formats import INTEGER_ARRAY_1_INT32, INTEGER_ARRAY_MD_INT32, LOOKUP_TABLE_INT32 from .tables import KnownTableFormat from .types import ArrayTableInfo, CategoryInfo, LookupTableInfo @@ -47,7 +49,32 @@ logger = logging.getLogger("object.data") _DATA_ID_KEY = "data" # The key used to identify data references in geoscience objects. -_DEFAULT_MAX_CONCURRENCY = 4 # Decent performance improvement going from 1 to 4 but beyond that not so much +_TRANSIENT_UPLOAD_STATUSES = frozenset({408, 429, 500, 502, 503, 504}) + + +def _is_transient_upload_error(error: Exception) -> bool: + if isinstance(error, TransportError): + # Connection failures have already exhausted the transport's retry policy. + return False + if isinstance(error, RetryError): + return bool(error.exceptions) and all(_is_transient_upload_error(exc) for exc in error.exceptions) + if isinstance(error, ChunkedIOError): + return isinstance(error.__cause__, Exception) and _is_transient_upload_error(error.__cause__) + if isinstance(error, EvoAPIException): + return error.status in _TRANSIENT_UPLOAD_STATUSES + return isinstance(error, (asyncio.TimeoutError, TimeoutError, ConnectionError)) + + +class _MonotonicFeedback(IFeedback): + """Prevent nested storage retries from reporting progress below an earlier attempt.""" + + def __init__(self, feedback: IFeedback) -> None: + self._feedback = feedback + self._progress = 0.0 + + def progress(self, progress: float, message: str | None = None) -> None: + self._progress = max(self._progress, progress) + self._feedback.progress(self._progress, message) def _iter_refs(target: Any, _key: str | None = None) -> Iterator[str]: @@ -109,6 +136,9 @@ def __init__(self, environment: Environment, connector: APIConnector, cache: ICa self._environment = environment self._connector = connector self._cache = cache + self._inflight_uploads: WeakKeyDictionary[ + asyncio.AbstractEventLoop, dict[tuple[str, int | None], asyncio.Task[None]] + ] = WeakKeyDictionary() @property def cache_location(self) -> Path: @@ -196,6 +226,7 @@ async def upload_table( table: pa.Table, fb: IFeedback = NoFeedback, table_format: KnownTableFormat | Iterable[KnownTableFormat] | None = None, + retry: Retry | None = None, ) -> dict: """Upload pyarrow table to the geoscience object service, returning a GO model of the uploaded data. @@ -205,23 +236,80 @@ async def upload_table( If a single KnownTableFormat is provided, that format will be used. Otherwise, the format will be inferred from the table, either by checking against the provided formats or by checking against all known formats. + :param retry: Retry policy for transient HTTP statuses (408, 429, 500, 502, 503, 504), raw connection/timeouts, + and exhausted storage retries containing only such failures. Defaults to three attempts with exponential + backoff, reusing cached data. On exhaustion, the last failure is raised with the retry history chained as + its cause. TransportError is not retried again. Storage transfer attempts are bounded by this retry policy. :return: A description of the uploaded data. :raises TableFormatError: If the provided table does not match any of the specified formats. If no table formats are specified, raised when the table does not match any known format. """ - table_info = self.save_table(table, table_format) - upload = ObjectDataUpload(connector=self._connector, environment=self._environment, name=table_info["data"]) - with error_context(f"While uploading table data '{table_info['data']}'", logger=logger): - try: - await upload.upload_from_cache(cache=self._cache, transport=self._connector.transport, fb=fb) - except DataExistsError: - logger.debug(f"Data not uploaded because data already exists (label: {table_info['data']})") - fb.progress(1) + table_info = await run_upload_preparation(self.save_table, table, table_format) + await self._upload_cached_table(table_info["data"], fb, retry) return table_info - async def upload_category_table(self, table: pa.Table, fb: IFeedback = NoFeedback) -> CategoryInfo: + async def _upload_cached_table(self, name: str, fb: IFeedback, retry: Retry | None) -> None: + loop = asyncio.get_running_loop() + uploads = self._inflight_uploads.setdefault(loop, {}) + key = (name, id(retry) if retry is not None else None) + task = uploads.get(key) + is_owner = task is None + if task is None: + task = asyncio.create_task(self._upload_cached_table_once(name, fb, retry)) + uploads[key] = task + + def remove_finished_upload(finished: asyncio.Task[None]) -> None: + if uploads.get(key) is finished: + uploads.pop(key, None) + if not finished.cancelled(): + finished.exception() + + task.add_done_callback(remove_finished_upload) + if is_owner: + await task + else: + await asyncio.shield(task) + fb.progress(1) + + async def _upload_cached_table_once(self, name: str, fb: IFeedback, retry: Retry | None) -> None: + if retry is None: + retry = Retry(logger=logger) + semaphore = get_upload_semaphore() + upload = ObjectDataUpload(connector=self._connector, environment=self._environment, name=name) + upload_fb = _MonotonicFeedback(fb) if fb is not NoFeedback else NoFeedback + with error_context(f"While uploading table data '{name}'", logger=logger): + async for attempt in retry: + try: + async with semaphore: + try: + await upload.upload_from_cache( + cache=self._cache, + transport=self._connector.transport, + retry=Retry(logger=logger, max_attempts=1), + fb=upload_fb, + ) + except RetryError as storage_error: + if len(storage_error.exceptions) == 1 and isinstance( + storage_error.exceptions[0], ChunkedIOError + ): + raise storage_error.exceptions[0] + raise + except DataExistsError: + logger.debug(f"Data not uploaded because data already exists (label: {name})") + fb.progress(1) + except Exception as exc: + if not _is_transient_upload_error(exc): + raise + try: + attempt.set_exception(exc) + except RetryError as retry_error: + raise exc from retry_error + + async def upload_category_table( + self, table: pa.Table, fb: IFeedback = NoFeedback, retry: Retry | None = None + ) -> CategoryInfo: """Upload a pyarrow table with category data to the geoscience object service, returning the reference to both the value and lookup tables. @@ -230,6 +318,7 @@ async def upload_category_table(self, table: pa.Table, fb: IFeedback = NoFeedbac :param table: The table to be uploaded. :param fb: A feedback object for tracking upload progress. + :param retry: Retry policy for each uploaded table, as described by upload_table. :return: A description of the uploaded data. @@ -237,6 +326,11 @@ async def upload_category_table(self, table: pa.Table, fb: IFeedback = NoFeedbac what int32 type can represent. """ + category_info, sizes = await run_upload_preparation(self._save_category_table, table) + return await self._upload_category_info(category_info, sizes, fb, retry) + + def _save_category_table(self, table: pa.Table) -> tuple[CategoryInfo, tuple[int, int]]: + """Build and cache both category tables without retaining Arrow copies during upload.""" all_chunks = [] columns_chunk_range = [] for column in table.itercolumns(): @@ -275,12 +369,23 @@ async def upload_category_table(self, table: pa.Table, fb: IFeedback = NoFeedbac v_size = values.num_rows * values.num_columns t_size = lookup.num_rows * lookup.num_columns - values_fb, lookup_fb = split_feedback(fb, [v_size, t_size]) - value_table_format = INTEGER_ARRAY_1_INT32 if values.num_columns == 1 else INTEGER_ARRAY_MD_INT32 - value_table = await self.upload_table(values, fb=values_fb, table_format=value_table_format) - lookup_table = await self.upload_table(lookup, fb=lookup_fb, table_format=LOOKUP_TABLE_INT32) - return CategoryInfo(values=cast(ArrayTableInfo, value_table), table=cast(LookupTableInfo, lookup_table)) + value_table = self.save_table(values, value_table_format) + lookup_table = self.save_table(lookup, LOOKUP_TABLE_INT32) + return ( + CategoryInfo(values=cast(ArrayTableInfo, value_table), table=cast(LookupTableInfo, lookup_table)), + (v_size, t_size), + ) + + async def _upload_category_info( + self, category_info: CategoryInfo, sizes: tuple[int, int], fb: IFeedback, retry: Retry | None + ) -> CategoryInfo: + values_fb, lookup_fb = split_feedback(fb, sizes) + await gather_cancelling( + self._upload_cached_table(category_info["values"]["data"], values_fb, retry), + self._upload_cached_table(category_info["table"]["data"], lookup_fb, retry), + ) + return category_info async def download_table( self, object_id: UUID, version_id: str, table_info: dict, fb: IFeedback = NoFeedback @@ -392,19 +497,14 @@ async def download_one(data_id: str, table_info: dict, table_fb: IFeedback) -> p async with semaphore: return await self._download_prepared_table(downloads[data_id], table_info, table_fb) - tasks = { - data_id: asyncio.create_task(download_one(data_id, table_info, table_fb)) - for (data_id, table_info), table_fb in zip(table_infos_by_data_id.items(), feedbacks) - } - try: - await asyncio.gather(*tasks.values()) - except BaseException: - # asyncio.gather doesn't cancel sibling tasks if one fails, manually cancel other tasks if one fails - for task in tasks.values(): - task.cancel() - await asyncio.gather(*tasks.values(), return_exceptions=True) - raise - return {data_id: task.result() for data_id, task in tasks.items()} + data_id_items = list(table_infos_by_data_id.items()) + tables = await gather_cancelling( + *( + download_one(data_id, table_info, table_fb) + for (data_id, table_info), table_fb in zip(data_id_items, feedbacks) + ) + ) + return {data_id: table for (data_id, _), table in zip(data_id_items, tables)} if _PD_AVAILABLE: # Optional support for pandas dataframes. Depends on both pyarrow and pandas. @@ -435,6 +535,8 @@ async def upload_dataframe( dataframe: pd.DataFrame, fb: IFeedback = NoFeedback, table_format: KnownTableFormat | Iterable[KnownTableFormat] | None = None, + retry: Retry | None = None, + _prepare_dataframe: Callable[[pd.DataFrame], pd.DataFrame] | None = None, ) -> dict: """Upload pandas dataframe to the geoscience object service, returning a GO model of the uploaded data. @@ -444,18 +546,26 @@ async def upload_dataframe( If a single KnownTableFormat is provided, that format will be used. Otherwise, the format will be inferred from the table, either by checking against the provided formats or by checking against all known formats. + :param retry: Retry policy for the uploaded table, as described by upload_table. + Override this method, not upload_table, to intercept DataFrame uploads. :return: A description of the uploaded data. :raises TableFormatError: If the provided table does not match any of the specified formats. If no table formats are specified, raised when the table does not match any known format. """ - schema = _get_schema_from_dataframe(dataframe) - table = pa.Table.from_pandas(dataframe, schema) - table_info = await self.upload_table(table, table_format=table_format, fb=fb) + + def save() -> dict: + prepared = _prepare_dataframe(dataframe) if _prepare_dataframe is not None else dataframe + return self.save_dataframe(prepared, table_format) + + table_info = await run_upload_preparation(save) + await self._upload_cached_table(table_info["data"], fb, retry) return table_info - async def upload_category_dataframe(self, dataframe: pd.DataFrame, fb: IFeedback = NoFeedback) -> CategoryInfo: + async def upload_category_dataframe( + self, dataframe: pd.DataFrame, fb: IFeedback = NoFeedback, retry: Retry | None = None + ) -> CategoryInfo: """Upload a pandas dataframe with category data to the geoscience object service, returning the reference to both the value and lookup tables. @@ -466,16 +576,21 @@ async def upload_category_dataframe(self, dataframe: pd.DataFrame, fb: IFeedback :param dataframe: The pandas dataframe to be uploaded. :param fb: A feedback object for tracking upload progress. + :param retry: Retry policy for each uploaded table, as described by upload_table. + Override this method, not upload_category_table, to intercept category DataFrame uploads. :return: A description of the uploaded data. :raises TableFormatError: If the table isn't a valid category table, or if the number of categories exceeds what int32 type can represent. """ - schema = _get_schema_from_dataframe(dataframe) - table = pa.Table.from_pandas(dataframe, schema) - category_info = await self.upload_category_table(table, fb=fb) - return category_info + + def save() -> tuple[CategoryInfo, tuple[int, int]]: + schema = _get_schema_from_dataframe(dataframe) + return self._save_category_table(pa.Table.from_pandas(dataframe, schema)) + + category_info, sizes = await run_upload_preparation(save) + return await self._upload_category_info(category_info, sizes, fb, retry) async def download_dataframe( self, object_id: UUID, version_id: str, table_info: dict, fb: IFeedback = NoFeedback diff --git a/packages/evo-objects/tests/test_data_client.py b/packages/evo-objects/tests/test_data_client.py index 0fa0109f6..32c468303 100644 --- a/packages/evo-objects/tests/test_data_client.py +++ b/packages/evo-objects/tests/test_data_client.py @@ -9,10 +9,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import json +import logging import sys +import threading +import weakref +from contextlib import nullcontext from unittest import mock -from uuid import UUID +from uuid import UUID, uuid4 import pandas as pd import pyarrow as pa @@ -20,13 +25,17 @@ from parameterized import parameterized from data import load_test_data -from evo.common import IFeedback, RequestMethod -from evo.common.exceptions import BadRequestException -from evo.common.io.exceptions import DataExistsError +from evo.aio import AioTransport +from evo.common import APIConnector, HTTPHeaderDict, HTTPResponse, IFeedback, RequestMethod +from evo.common.exceptions import BadRequestException, EvoAPIException, RetryError, TransportError +from evo.common.io.exceptions import ChunkedIOError, DataExistsError +from evo.common.io.storage import StorageDestination from evo.common.test_tools import TestWithConnector, TestWithStorage -from evo.common.utils import NoFeedback, PartialFeedback, get_header_metadata +from evo.common.utils import BackoffLinear, NoFeedback, PartialFeedback, Retry, get_header_metadata +from evo.objects.endpoints.models import DataUploadResponseBody from evo.objects.exceptions import TableFormatError -from evo.objects.utils import KnownTableFormat, ObjectDataClient, table_formats +from evo.objects.utils import KnownTableFormat, ObjectDataClient, table_formats, upload_concurrency +from evo.objects.utils._concurrency import _DEFAULT_MAX_CONCURRENCY, gather_cancelling from helpers import NoImport, UnloadModule, get_sample_table_and_bytes @@ -200,17 +209,21 @@ async def _mock_upload_file_side_effect(*args, **kwargs): body=[{"name": data["name"]} for data in put_data_response[1:]] + [{"name": put_data_response[0]["name"]}], ) - @parameterized.expand(["referenced", "table"]) - async def test_upload_failure_preserves_exception_and_identifies_data(self, upload_type: str) -> None: + @parameterized.expand( + [("referenced", ConnectionError), ("table", RuntimeError), ("table_transient", ConnectionError)] + ) + async def test_upload_failure_preserves_exception_and_identifies_data( + self, upload_type: str, error_type: type[Exception] + ) -> None: response = load_test_data("put_data_batch.json")[:1] name = response[0]["name"] (self.data_client.cache_location / name).touch() - error = ConnectionError("Storage unavailable") + error = error_type("Storage unavailable") with ( self.transport.set_http_response(status_code=200, content=json.dumps(response)), mock.patch("evo.common.io.upload.StorageDestination", autospec=True) as destination, mock.patch("evo.objects.utils.data.logger") as logger, - self.assertRaises(ConnectionError) as cm, + self.assertRaises(error_type) as cm, ): destination.upload_file.side_effect = error if upload_type == "referenced": @@ -218,10 +231,15 @@ async def test_upload_failure_preserves_exception_and_identifies_data(self, uplo else: table, _ = get_sample_table_and_bytes(table_formats.FLOAT_ARRAY_2, 1) with mock.patch.object(self.data_client, "save_table", return_value={"data": name}): - await self.data_client.upload_table(table) - self.assertIs(error, cm.exception) + await self.data_client.upload_table( + table, retry=Retry(logger=logger, backoff_method=BackoffLinear(0)) + ) + if upload_type == "table_transient": + self.assertEqual((error, error, error), cm.exception.__cause__.exceptions) + else: + self.assertIs(error, cm.exception) context = ( - "\n".join(getattr(error, "__notes__", [])) + "\n".join(getattr(cm.exception, "__notes__", [])) if sys.version_info >= (3, 11) else "\n".join(call.args[0] for call in logger.error.call_args_list) ) @@ -403,6 +421,694 @@ async def _mock_upload_file_side_effect(*args, **kwargs): ) self.assertIs(mock_table_info, actual_table_info) + @parameterized.expand( + [ + ("table", "upload_table", False), + ("nested_group", "upload_table", True), + ("dataframe", "upload_dataframe", False), + ("category", "upload_category_dataframe", False), + ] + ) + async def test_cancelled_upload_waits_for_parquet_writer(self, _name, method, nested_group) -> None: + started = asyncio.Event() + release = threading.Event() + finished = threading.Event() + loop = asyncio.get_running_loop() + + def save_table(*_args, **_kwargs): + loop.call_soon_threadsafe(started.set) + try: + if not release.wait(timeout=15): + raise TimeoutError("The test did not release the parquet writer") + return {"data": "saved-table"} + finally: + finished.set() + + with ( + mock.patch.object(self.data_client, "save_table", side_effect=save_table), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache") as upload, + ): + + async def run(): + data = pd.DataFrame({"value": ["sand"] if "category" in method else [1.0]}) + if method == "upload_table": + data = pa.Table.from_pandas(data) + operation = getattr(self.data_client, method)(data) + if nested_group: + return await gather_cancelling(operation) + return await operation + + task = asyncio.create_task(run()) + cancelled = False + try: + await asyncio.wait_for(started.wait(), timeout=10) + task.cancel() + cancelled = True + await asyncio.sleep(0) + self.assertFalse(task.done(), "Cancellation returned while the parquet writer was still running") + finally: + release.set() + if not cancelled: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.assertTrue(task.cancelled()) + self.assertTrue(finished.is_set()) + upload.assert_not_awaited() + + async def test_parquet_serialization_runs_off_the_event_loop(self) -> None: + threads = [] + main_thread = threading.get_ident() + + def save_table(*_args): + threads.append(threading.get_ident()) + return {"data": "saved-table"} + + with ( + mock.patch.object(self.data_client, "save_table", side_effect=save_table), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache") as upload, + ): + result = await self.data_client.upload_table(pa.table({"value": [1.0]})) + + self.assertEqual({"data": "saved-table"}, result) + self.assertEqual(1, len(threads)) + self.assertNotEqual(main_thread, threads[0]) + upload.assert_awaited_once() + + async def test_parquet_writer_failure_is_propagated_without_uploading(self) -> None: + with ( + mock.patch.object(self.data_client, "save_table", side_effect=RuntimeError("parquet write failed")), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache") as upload, + ): + with self.assertRaisesRegex(RuntimeError, "parquet write failed"): + await self.data_client.upload_table(pa.table({"value": [1.0]})) + + upload.assert_not_awaited() + + @parameterized.expand([("plain_serial", False, 1), ("plain", False, 3), ("category", True, 3)]) + async def test_dataframe_conversion_is_off_loop_and_live_arrow_tables_are_bounded(self, _name, category, limit): + loop = asyncio.get_running_loop() + main_thread = threading.get_ident() + reached_limit = asyncio.Event() + release = threading.Event() + lock = threading.Lock() + tables = [] + conversion_threads = [] + peak = 0 + from_pandas = pa.Table.from_pandas + + def convert(*args, **kwargs): + nonlocal peak + table = from_pandas(*args, **kwargs) + with lock: + tables.append(weakref.ref(table)) + conversion_threads.append(threading.get_ident()) + live = sum(reference() is not None for reference in tables) + peak = max(peak, live) + if live == limit: + loop.call_soon_threadsafe(reached_limit.set) + return table + + def save_table(table, table_format=None): + if not release.wait(timeout=15): + raise TimeoutError("The test did not release serialization") + return {"data": str(uuid4())} + + with ( + mock.patch("evo.objects.utils.data.pa.Table", wraps=pa.Table) as arrow, + mock.patch.object(self.data_client, "save_table", new=save_table), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache"), + upload_concurrency(limit), + ): + arrow.from_pandas.side_effect = convert + upload = self.data_client.upload_category_dataframe if category else self.data_client.upload_dataframe + frame = pd.DataFrame({"value": pd.Categorical(["sand", "clay"]) if category else [1.0, 2.0]}) + task = asyncio.gather(*(upload(frame) for _ in range(12))) + try: + await asyncio.wait_for(reached_limit.wait(), timeout=10) + self.assertEqual(limit, peak) + release.set() + await asyncio.wait_for(task, timeout=10) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.assertEqual(12, len(tables)) + self.assertEqual(limit, peak) + self.assertTrue(all(thread != main_thread for thread in conversion_threads)) + self.assertTrue(all(reference() is None for reference in tables)) + + @parameterized.expand( + [ + ("default", None, False), + ("serial", 1, False), + ("lower", 3, False), + ("higher", 12, False), + ("multiple_clients_default", None, True), + ("multiple_clients_scope", 3, True), + ] + ) + async def test_direct_concurrent_uploads_share_loop_limit(self, _name, configured_limit, multiple_clients) -> None: + limit = configured_limit if configured_limit is not None else _DEFAULT_MAX_CONCURRENCY + clients = [ + ObjectDataClient(environment=self.environment, connector=self.connector, cache=self.cache) + if multiple_clients + else self.data_client + for _ in range(24) + ] + release = asyncio.Event() + limit_reached = asyncio.Event() + active = 0 + peak = 0 + + async def upload(**_kwargs): + nonlocal active, peak + active += 1 + peak = max(peak, active) + if active == limit: + limit_reached.set() + try: + await release.wait() + finally: + active -= 1 + + with ( + mock.patch.object(ObjectDataClient, "save_table", side_effect=lambda table, _: {"data": table}) as save, + mock.patch("evo.objects.utils.data.asyncio.to_thread", side_effect=lambda fn, *args: fn(*args)), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload), + upload_concurrency(configured_limit) if configured_limit is not None else nullcontext(), + ): + task = asyncio.gather(*(client.upload_table(str(i)) for i, client in enumerate(clients))) + try: + await asyncio.wait_for(limit_reached.wait(), timeout=2) + await asyncio.sleep(0) + self.assertEqual(limit, active) + release.set() + await asyncio.wait_for(task, timeout=2) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + self.assertEqual(limit, peak) + self.assertEqual(0, active) + self.assertEqual(24, save.call_count) + + async def test_serialization_and_upload_attempts_share_the_same_limit(self): + active = 0 + peak = 0 + + async def work(function, *args): + nonlocal active, peak + active += 1 + peak = max(peak, active) + try: + await asyncio.sleep(0) + return function(*args) + finally: + active -= 1 + + async def upload(**_kwargs): + await work(lambda: None) + + with ( + mock.patch.object(self.data_client, "save_table", side_effect=lambda table, _: {"data": table}), + mock.patch("evo.objects.utils.data.asyncio.to_thread", side_effect=work) as serialize, + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload) as transfer, + upload_concurrency(3), + ): + await asyncio.gather(*(self.data_client.upload_table(str(i)) for i in range(24))) + self.assertEqual(3, peak) + self.assertEqual(0, active) + self.assertEqual(24, serialize.await_count) + self.assertEqual(24, transfer.await_count) + + @parameterized.expand([("default", None), ("scope", 1)]) + def test_client_can_be_reused_across_contended_event_loops(self, _name, limit): + async def upload(**_kwargs): + await asyncio.sleep(0) + + async def run(): + return await asyncio.gather(*(self.data_client.upload_table(str(i)) for i in range(16))) + + with ( + mock.patch.object(self.data_client, "save_table", side_effect=lambda table, _: {"data": table}), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload), + upload_concurrency(limit) if limit is not None else nullcontext(), + ): + for _ in range(2): + result = asyncio.run(asyncio.wait_for(run(), timeout=10)) + self.assertEqual([str(i) for i in range(16)], [info["data"] for info in result]) + + @parameterized.expand([("default", None), ("scope", 3), ("serial", 1)]) + async def test_concurrent_category_operations_share_limit(self, _name, configured_limit): + limit = configured_limit if configured_limit is not None else _DEFAULT_MAX_CONCURRENCY + release = asyncio.Event() + limit_reached = asyncio.Event() + active = 0 + peak = 0 + + async def upload(**_kwargs): + nonlocal active, peak + active += 1 + peak = max(peak, active) + if active == limit: + limit_reached.set() + try: + await release.wait() + finally: + active -= 1 + + with ( + mock.patch.object(self.data_client, "save_table", side_effect=lambda *_: {"data": str(uuid4())}), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload), + upload_concurrency(configured_limit) if configured_limit is not None else nullcontext(), + ): + task = asyncio.gather( + *(self.data_client.upload_category_table(pa.table({"value": [str(i)]})) for i in range(10)) + ) + try: + await asyncio.wait_for(limit_reached.wait(), timeout=2) + await asyncio.sleep(0) + release.set() + await asyncio.wait_for(task, timeout=2) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + self.assertEqual(limit, peak) + self.assertEqual(0, active) + + @parameterized.expand( + [(str(status), EvoAPIException(status, "transient", None, None)) for status in (408, 429, 500, 502, 503, 504)] + + [ + ("timeout", TimeoutError()), + ("connection", ConnectionError()), + ] + ) + async def test_transient_upload_retry_preserves_order_and_cached_tables(self, _name, error): + attempts = {} + completed = [] + sibling_finished = asyncio.Event() + + async def upload(upload_context, **_kwargs): + name = upload_context.name + attempts[name] = attempts.get(name, 0) + 1 + if name == "first" and attempts[name] <= 2: + await sibling_finished.wait() + raise error + completed.append(name) + sibling_finished.set() + + retry = Retry(logger=mock.Mock(), backoff_method=BackoffLinear(0)) + with ( + mock.patch.object(self.data_client, "save_table", side_effect=lambda table, _: {"data": table}) as save, + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", new=upload), + ): + results = await asyncio.wait_for( + gather_cancelling( + self.data_client.upload_table("first", retry=retry), + self.data_client.upload_table("second", retry=retry), + ), + timeout=2, + ) + self.assertEqual(["first", "second"], [result["data"] for result in results]) + self.assertEqual(["second", "first"], completed) + self.assertEqual({"first": 3, "second": 1}, attempts) + self.assertEqual(2, save.call_count) + + async def test_retry_covers_actual_upload_url_acquisition(self): + response = load_test_data("put_data.json") + with ( + self.transport.set_http_response(status_code=201, content=""), + mock.patch.object(self.data_client, "save_table", wraps=self.data_client.save_table) as save, + mock.patch( + "evo.objects.io.DataApi.put_data_in_workspace", + side_effect=[TimeoutError(), [DataUploadResponseBody(**response[0])]], + ) as request_url, + ): + result = await self.data_client.upload_table( + pa.table({"value": [1.0]}), + retry=Retry(logger=mock.Mock(), backoff_method=BackoffLinear(0)), + ) + self.assertTrue((self.data_client.cache_location / result["data"]).is_file()) + self.assertEqual(2, request_url.await_count) + save.assert_called_once() + + @parameterized.expand([("url", False), ("storage", True)]) + async def test_transport_retries_are_not_multiplied_by_table_retries(self, _name, storage): + transport = AioTransport( + user_agent="test", max_attempts=3, backoff_method=BackoffLinear(0), close_grace_period_ms=0 + ) + client = ObjectDataClient( + environment=self.environment, + connector=APIConnector(base_url=self.connector.base_url, transport=transport), + cache=self.cache, + ) + response = load_test_data("put_data.json") + with ( + mock.patch("aiohttp.ClientSession.request", side_effect=ConnectionError("offline")) as request, + mock.patch( + "evo.objects.io.DataApi.put_data_in_workspace", + return_value=[DataUploadResponseBody(**response[0])], + ) + if storage + else nullcontext(), + ): + async with transport: + with self.assertRaises(ChunkedIOError if storage else TransportError) as caught: + await client.upload_table( + pa.table({"value": [1.0]}), + retry=Retry(logger=mock.Mock(), backoff_method=BackoffLinear(0)), + ) + self.assertEqual(3, request.call_count) + if storage: + self.assertIsInstance(caught.exception.__cause__, TransportError) + + @parameterized.expand( + [ + ("upload_table", False), + ("upload_dataframe", False), + ("upload_category_table", True), + ("upload_category_dataframe", True), + ] + ) + async def test_public_upload_methods_honor_retry_policy(self, method, category): + frame = pd.DataFrame({"value": ["sand", "clay"] if category else [1.0, 2.0]}) + data = frame if "dataframe" in method else pa.table({"value": frame["value"].tolist()}) + with mock.patch( + "evo.objects.io.ObjectDataUpload.upload_from_cache", + side_effect=EvoAPIException(503, "unavailable", None, None), + ) as upload: + with self.assertRaises(EvoAPIException) as caught: + await getattr(self.data_client, method)( + data, retry=Retry(logger=mock.Mock(), max_attempts=1, backoff_method=BackoffLinear(0)) + ) + self.assertIsInstance(caught.exception.__cause__, RetryError) + self.assertEqual(1, len(caught.exception.__cause__.exceptions)) + self.assertEqual(2 if category else 1, upload.await_count) + + @parameterized.expand([("chunk", "comp=block&"), ("commit", "comp=blocklist")]) + async def test_retry_covers_real_storage_retry_exhaustion(self, _name, query): + response = load_test_data("put_data.json") + original_request = self.transport.request + failed = False + storage_errors = [] + upload_file = StorageDestination.upload_file + + async def upload_with_single_attempt(**kwargs): + kwargs["retry"] = Retry(logger=mock.Mock(), max_attempts=1, backoff_method=BackoffLinear(0)) + try: + await upload_file(**kwargs) + except RetryError as error: + storage_errors.append(error) + raise + + async def request(**kwargs): + nonlocal failed + if query in kwargs["url"] and not failed: + failed = True + return HTTPResponse(status=503, data=b"", reason="unavailable", headers=HTTPHeaderDict()) + return await original_request(**kwargs) + + with ( + self.transport.set_http_response(status_code=201, content=""), + mock.patch.object(self.transport, "request", side_effect=request), + mock.patch("evo.common.io.upload.StorageDestination.upload_file", side_effect=upload_with_single_attempt), + mock.patch.object(self.data_client, "save_table", wraps=self.data_client.save_table) as save, + mock.patch( + "evo.objects.io.DataApi.put_data_in_workspace", + return_value=[DataUploadResponseBody(**response[0])], + ) as request_url, + ): + result = await self.data_client.upload_table( + pa.table({"value": [1.0]}), + retry=Retry(logger=mock.Mock(), backoff_method=BackoffLinear(0)), + ) + self.assertTrue(failed) + self.assertEqual(1, len(storage_errors)) + self.assertIsInstance(storage_errors[0].exceptions[0], ChunkedIOError) + self.assertIsInstance(storage_errors[0].exceptions[0].__cause__, EvoAPIException) + self.assertEqual(503, storage_errors[0].exceptions[0].__cause__.status) + self.assertTrue((self.data_client.cache_location / result["data"]).is_file()) + self.assertEqual(2, request_url.await_count) + save.assert_called_once() + + @parameterized.expand( + [ + ("url", EvoAPIException), + ("connection", ConnectionError), + ] + ) + async def test_retry_exhaustion_is_bounded_and_preserves_errors(self, _name, error_type): + errors = [] + for i in range(3): + cause = EvoAPIException(503, f"unavailable-{i}", None, None) + error = cause if error_type is EvoAPIException else error_type(f"unavailable-{i}") + errors.append(error) + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}) as save, + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=errors) as upload, + ): + with self.assertRaises(error_type) as caught: + await self.data_client.upload_table( + mock.Mock(), retry=Retry(logger=mock.Mock(), backoff_method=BackoffLinear(0)) + ) + self.assertIs(errors[-1], caught.exception) + self.assertEqual(tuple(errors), caught.exception.__cause__.exceptions) + self.assertEqual(3, upload.await_count) + save.assert_called_once() + + async def test_recovered_upload_logs_warning_not_error(self): + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}), + mock.patch( + "evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=[ConnectionError("retry"), None] + ), + self.assertLogs("evo.object.data", level=logging.WARNING) as logs, + ): + await self.data_client.upload_table( + mock.Mock(), retry=Retry(logger=logging.getLogger("evo.object.data"), backoff_method=BackoffLinear(0)) + ) + self.assertEqual([logging.WARNING], [record.levelno for record in logs.records]) + + async def test_other_uploads_proceed_during_backoff(self): + backing_off = asyncio.Event() + release = asyncio.Event() + attempts = {} + sleep = asyncio.sleep + + async def backoff(delay): + if delay == 1: + backing_off.set() + await release.wait() + else: + await sleep(delay) + + async def upload(upload, **_kwargs): + attempts[upload.name] = attempts.get(upload.name, 0) + 1 + if upload.name == "first" and attempts[upload.name] == 1: + raise ConnectionError("retry") + + with ( + mock.patch.object(self.data_client, "save_table", side_effect=lambda table, _: {"data": table}), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", new=upload), + mock.patch("evo.common.utils.retry.asyncio.sleep", side_effect=backoff), + upload_concurrency(1), + ): + task = asyncio.create_task( + self.data_client.upload_table("first", retry=Retry(logger=mock.Mock(), backoff_method=BackoffLinear(1))) + ) + try: + await asyncio.wait_for(backing_off.wait(), timeout=2) + self.assertEqual({"data": "second"}, await asyncio.wait_for(self.data_client.upload_table("second"), 2)) + self.assertFalse(task.done()) + release.set() + await asyncio.wait_for(task, timeout=2) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + self.assertEqual({"first": 2, "second": 1}, attempts) + + @parameterized.expand([("cancelled",), ("failed",)]) + async def test_identical_uploads_do_not_share_cancellation_or_failures(self, outcome): + started = [asyncio.Event(), asyncio.Event()] + release = [asyncio.Event(), asyncio.Event()] + calls = 0 + other_client = ObjectDataClient(environment=self.environment, connector=self.connector, cache=self.cache) + table = pa.table({"value": [1.0]}) + + async def upload(**_kwargs): + nonlocal calls + index = calls + calls += 1 + started[index].set() + await release[index].wait() + if index == 0: + raise RuntimeError("upload failed") + + with mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload): + first = asyncio.create_task(self.data_client.upload_table(table)) + tasks = [first] + try: + await asyncio.wait_for(started[0].wait(), timeout=10) + second = asyncio.create_task(other_client.upload_table(table)) + tasks.append(second) + await asyncio.wait_for(started[1].wait(), timeout=10) + if outcome == "cancelled": + first.cancel() + with self.assertRaises(asyncio.CancelledError): + await first + else: + release[0].set() + with self.assertRaisesRegex(RuntimeError, "upload failed"): + await first + self.assertFalse(second.done()) + release[1].set() + result = await asyncio.wait_for(second, timeout=10) + self.assertEqual(self.data_client.save_table(table), result) + finally: + for event in release: + event.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + self.assertEqual(2, calls) + + async def test_identical_uploads_use_independent_retry_policies(self): + async def upload(**_kwargs): + await asyncio.sleep(0) + raise ConnectionError("retry") + + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload) as transfer, + ): + results = await asyncio.gather( + *( + self.data_client.upload_table( + mock.Mock(), + retry=Retry(logger=mock.Mock(), max_attempts=limit, backoff_method=BackoffLinear(0)), + ) + for limit in (1, 2) + ), + return_exceptions=True, + ) + self.assertTrue(all(isinstance(result, ConnectionError) for result in results)) + self.assertEqual([1, 2], [len(result.__cause__.exceptions) for result in results]) + self.assertEqual(3, transfer.await_count) + + async def test_concurrent_uploads_share_identical_cached_data(self): + async def upload(**_kwargs): + await asyncio.sleep(0) + + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "same"}), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload) as transfer, + ): + results = await asyncio.gather( + self.data_client.upload_table(mock.Mock()), + self.data_client.upload_table(mock.Mock()), + ) + + self.assertEqual([{"data": "same"}, {"data": "same"}], results) + transfer.assert_awaited_once() + + async def test_upload_progress_does_not_decrease_after_retry(self): + feedback = mock.Mock(spec=IFeedback) + calls = 0 + + async def upload(**kwargs): + nonlocal calls + calls += 1 + kwargs["fb"].progress(0.8 if calls == 1 else 0.2) + if calls == 1: + raise ConnectionError("retry") + kwargs["fb"].progress(1) + + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload), + ): + await self.data_client.upload_table( + mock.Mock(), + fb=feedback, + retry=Retry(logger=mock.Mock(), max_attempts=2, backoff_method=BackoffLinear(0)), + ) + + progress = [call.args[0] for call in feedback.progress.call_args_list] + self.assertEqual(sorted(progress), progress) + self.assertEqual(1, progress[-1]) + + @parameterized.expand( + [ + ("bad_request", EvoAPIException(400, "bad request", None, None)), + ("unauthorized", EvoAPIException(401, "unauthorized", None, None)), + ("forbidden", EvoAPIException(403, "forbidden", None, None)), + ("not_found", EvoAPIException(404, "not found", None, None)), + ("transport", TransportError("already retried")), + ("storage_transport", RetryError("storage", [TransportError("already retried")])), + ("storage_permanent", RetryError("storage", [EvoAPIException(401, "unauthorized", None, None)])), + ( + "storage_mixed", + RetryError("storage", [EvoAPIException(503, "unavailable", None, None), ValueError("invalid")]), + ), + ("unclassified_chunk_error", ChunkedIOError("unknown cause")), + ("invalid_data", ValueError("invalid")), + ("cancelled", asyncio.CancelledError()), + ] + ) + async def test_permanent_errors_and_cancellation_are_not_retried(self, _name, error): + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=error) as upload, + ): + with self.assertRaises(type(error)) as caught: + await self.data_client.upload_table(mock.Mock()) + self.assertIs(error, caught.exception) + upload.assert_awaited_once() + + async def test_already_existing_data_after_retry_is_success(self): + feedback = mock.Mock(spec=IFeedback) + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}) as save, + mock.patch( + "evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=[TimeoutError(), DataExistsError()] + ) as upload, + ): + result = await self.data_client.upload_table( + mock.Mock(), fb=feedback, retry=Retry(logger=mock.Mock(), backoff_method=BackoffLinear(0)) + ) + self.assertEqual({"data": "saved"}, result) + self.assertEqual(2, upload.await_count) + save.assert_called_once() + feedback.progress.assert_called_with(1) + + async def test_category_values_and_lookup_upload_concurrently(self) -> None: + started = [] + all_started = asyncio.Event() + release = asyncio.Event() + + async def upload(upload, **_kwargs): + started.append(upload.name) + if len(started) == 2: + all_started.set() + await release.wait() + + with mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", new=upload): + task = asyncio.create_task(self.data_client.upload_category_table(pa.table({"category": ["sand", "clay"]}))) + try: + await asyncio.wait_for(all_started.wait(), timeout=2) + release.set() + result = await asyncio.wait_for(task, timeout=2) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.assertCountEqual([result["values"]["data"], result["table"]["data"]], started) + async def test_upload_dataframe_exists(self) -> None: """Test uploading tabular data using pandas when the table exists.""" put_data_response = load_test_data("put_data_exists.json") @@ -526,16 +1232,18 @@ async def test_upload_category_table( ) -> None: """Test uploading categorical data using pandas.""" with ( - mock.patch("evo.objects.utils.ObjectDataClient.upload_table") as mock_upload_table, + mock.patch.object(self.data_client, "save_table") as save, + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache") as upload, ): - def side_effect(table_to_upload, table_format=None, fb=None): - return {"table": table_to_upload, "table_format": table_format} + def side_effect(table_to_upload, table_format=None): + return {"data": str(uuid4()), "table": table_to_upload, "table_format": table_format} - mock_upload_table.side_effect = side_effect + save.side_effect = side_effect category_info = await self.data_client.upload_category_table(table) - self.assertEqual(mock_upload_table.call_count, 2) + self.assertEqual(save.call_count, 2) + self.assertEqual(upload.await_count, 2) self.assertEqual(category_info["table"]["table"], lookup_table) self.assertEqual(category_info["table"]["table_format"], table_formats.LOOKUP_TABLE_INT32) self.assertEqual(category_info["values"]["table"], values_table) @@ -597,16 +1305,18 @@ async def test_upload_category_dataframe( ) -> None: """Test uploading categorical data using pandas.""" with ( - mock.patch("evo.objects.utils.ObjectDataClient.upload_table") as mock_upload_table, + mock.patch.object(self.data_client, "save_table") as save, + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache") as upload, ): - def side_effect(table, table_format=None, fb=None): - return {"table": table, "table_format": table_format} + def side_effect(table, table_format=None): + return {"data": str(uuid4()), "table": table, "table_format": table_format} - mock_upload_table.side_effect = side_effect + save.side_effect = side_effect category_info = await self.data_client.upload_category_dataframe(dataframe) - self.assertEqual(mock_upload_table.call_count, 2) + self.assertEqual(save.call_count, 2) + self.assertEqual(upload.await_count, 2) self.assertEqual(category_info["table"]["table"], lookup_table) self.assertEqual(category_info["table"]["table_format"], table_formats.LOOKUP_TABLE_INT32) self.assertEqual(category_info["values"]["table"], values_table) @@ -615,11 +1325,13 @@ def side_effect(table, table_format=None, fb=None): async def test_upload_dataframe_converts_large_string_via_schema(self) -> None: """Test that upload_dataframe converts large_string columns (from pandas string[pyarrow]) to string via schema.""" dataframe = pd.DataFrame({"col": pd.Series(["A", "B", "C"], dtype="string[pyarrow]")}) - with mock.patch.object(self.data_client, "upload_table", new_callable=mock.AsyncMock) as mock_upload: - mock_upload.return_value = {} + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}) as save, + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache"), + ): await self.data_client.upload_dataframe(dataframe) - (table,), _ = mock_upload.await_args + (table,), _ = save.call_args self.assertTrue(pa.types.is_string(table.schema.field("col").type)) async def test_download_table(self) -> None: diff --git a/packages/evo-objects/tests/test_upload_concurrency.py b/packages/evo-objects/tests/test_upload_concurrency.py new file mode 100644 index 000000000..8c4d01391 --- /dev/null +++ b/packages/evo-objects/tests/test_upload_concurrency.py @@ -0,0 +1,335 @@ +# Copyright © 2026 Bentley Systems, Incorporated +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import inspect +import threading +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext +from unittest import IsolatedAsyncioTestCase, TestCase +from unittest.mock import patch + +from parameterized import parameterized + +from evo.objects.utils import ObjectDataClient, upload_concurrency +from evo.objects.utils._concurrency import ( + _DEFAULT_MAX_CONCURRENCY, + _upload_scope, + gather_cancelling, + get_upload_semaphore, + run_upload_preparation, +) + + +class TestUploadConcurrency(IsolatedAsyncioTestCase): + def test_upload_and_download_defaults_match(self): + self.assertEqual( + _DEFAULT_MAX_CONCURRENCY, + inspect.signature(ObjectDataClient.download_tables).parameters["max_concurrency"].default, + ) + with upload_concurrency(): + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) + + @parameterized.expand([("default", None), ("serial", 1), ("lower", 3), ("higher", 12)]) + async def test_nested_groups_share_limit_preserve_order_and_reset_scope(self, _name, configured_limit): + limit = configured_limit if configured_limit is not None else _DEFAULT_MAX_CONCURRENCY + default = get_upload_semaphore() + release = asyncio.Event() + limit_reached = asyncio.Event() + active = 0 + peak = 0 + semaphores = [] + + async def upload(index): + nonlocal active, peak + semaphore = get_upload_semaphore() + semaphores.append(semaphore) + async with semaphore: + active += 1 + peak = max(peak, active) + if active == limit: + limit_reached.set() + try: + await release.wait() + return index + finally: + active -= 1 + + async def group(start): + return await gather_cancelling(*(upload(index) for index in range(start, start + 12))) + + async def run(): + return await asyncio.gather(group(0), group(12), group(24)) + + with upload_concurrency(configured_limit) if configured_limit is not None else nullcontext(): + task = asyncio.create_task(run()) + try: + await asyncio.wait_for(limit_reached.wait(), timeout=2) + self.assertEqual(limit, active) + release.set() + result = await asyncio.wait_for(task, timeout=2) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.assertEqual([list(range(0, 12)), list(range(12, 24)), list(range(24, 36))], result) + self.assertEqual(limit, peak) + self.assertEqual(0, active) + self.assertTrue(all(semaphore is semaphores[0] for semaphore in semaphores)) + self.assertIs(default, get_upload_semaphore()) + if configured_limit is None: + self.assertIs(default, semaphores[0]) + else: + self.assertIsNot(default, semaphores[0]) + + @parameterized.expand([(value,) for value in (0, -1, True, False, 1.5, "3", None)]) + async def test_invalid_configuration_leaves_scope_unchanged(self, value): + default = get_upload_semaphore() + with self.assertRaisesRegex(ValueError, "max_concurrent must be a positive integer"): + with upload_concurrency(value): + self.fail("Invalid configuration was accepted") + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) + self.assertIs(default, get_upload_semaphore()) + + async def test_nested_configuration_restores_scope_after_failure(self): + default = get_upload_semaphore() + with upload_concurrency(3): + outer = get_upload_semaphore() + self.assertEqual(3, _upload_scope.get().limit) + with self.assertRaisesRegex(RuntimeError, "failed"): + with upload_concurrency(1): + self.assertEqual(1, _upload_scope.get().limit) + self.assertIsNot(outer, get_upload_semaphore()) + raise RuntimeError("failed") + self.assertEqual(3, _upload_scope.get().limit) + self.assertIs(outer, get_upload_semaphore()) + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) + self.assertIs(default, get_upload_semaphore()) + + async def test_configuration_is_isolated_between_tasks(self): + default = get_upload_semaphore() + ready = asyncio.Event() + release = asyncio.Event() + + async def configured(): + with upload_concurrency(2): + ready.set() + await release.wait() + return _upload_scope.get().limit + + task = asyncio.create_task(configured()) + try: + await asyncio.wait_for(ready.wait(), timeout=2) + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) + self.assertIs(default, get_upload_semaphore()) + release.set() + self.assertEqual(2, await asyncio.wait_for(task, timeout=2)) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + async def test_failure_cancels_and_drains_nested_uploads_and_waiters(self): + started = asyncio.Event() + release = asyncio.Event() + cleaned_up = set() + failure = RuntimeError("upload failed") + + async def upload(index): + try: + async with get_upload_semaphore(): + started.set() + await release.wait() + finally: + await asyncio.sleep(0) + cleaned_up.add(index) + + async def group(start): + return await gather_cancelling(*(upload(index) for index in range(start, start + 12))) + + async def fail(): + await started.wait() + raise failure + + try: + with self.assertRaises(RuntimeError) as caught: + await gather_cancelling(group(0), group(12), fail()) + self.assertIs(failure, caught.exception) + self.assertEqual(set(range(24)), cleaned_up) + finally: + release.set() + + async def test_caller_cancellation_waits_for_nested_cleanup(self): + started = asyncio.Event() + release = asyncio.Event() + cleaned_up = set() + + async def upload(index): + started.set() + try: + await release.wait() + finally: + await asyncio.sleep(0) + cleaned_up.add(index) + + async def group(start): + return await gather_cancelling(*(upload(index) for index in range(start, start + 3))) + + async def run(): + return await gather_cancelling(group(0), group(3)) + + task = asyncio.create_task(run()) + try: + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + self.assertEqual(set(range(6)), cleaned_up) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + async def test_groups_reuse_default_limit_after_success_and_failure(self): + semaphores = [] + + async def upload(*, fail=False): + semaphores.append(get_upload_semaphore()) + if fail: + raise ValueError("invalid input") + + await gather_cancelling(upload()) + self.assertIs(get_upload_semaphore(), semaphores[-1]) + with self.assertRaisesRegex(ValueError, "invalid input"): + await gather_cancelling(upload(fail=True)) + self.assertIs(get_upload_semaphore(), semaphores[-1]) + self.assertIs(semaphores[0], semaphores[1]) + + async def test_empty_group_returns_empty_list(self): + self.assertEqual([], await gather_cancelling()) + + async def test_repeated_cancellation_drains_preparation_before_releasing_slot(self): + started = threading.Event() + release = threading.Event() + acquired = asyncio.Event() + + def worker(): + started.set() + release.wait() + return 42 + + async def occupy_slot(): + async with get_upload_semaphore(): + acquired.set() + + with upload_concurrency(1): + preparation = asyncio.create_task(run_upload_preparation(worker)) + try: + await asyncio.wait_for(asyncio.to_thread(started.wait), timeout=2) + preparation.cancel() + await asyncio.sleep(0) + preparation.cancel() + waiter = asyncio.create_task(occupy_slot()) + await asyncio.sleep(0) + self.assertFalse(acquired.is_set()) + release.set() + with self.assertRaises(asyncio.CancelledError): + await preparation + await waiter + self.assertTrue(acquired.is_set()) + finally: + release.set() + preparation.cancel() + await asyncio.gather(preparation, return_exceptions=True) + + async def test_secondary_failures_are_logged_without_replacing_the_first(self): + errors = [RuntimeError("first"), ValueError("second")] + + async def fail(error): + raise error + + with self.assertLogs("evo.object.uploads", level="DEBUG") as logs: + with self.assertRaises(RuntimeError) as caught: + await gather_cancelling(*(fail(error) for error in errors)) + self.assertIs(errors[0], caught.exception) + self.assertEqual(1, len(logs.records)) + self.assertIs(errors[1], logs.records[0].exc_info[1]) + + async def test_task_creation_failure_closes_unscheduled_coroutines_and_drains_tasks(self): + async def upload(): + await asyncio.sleep(0) + + coroutines = [upload() for _ in range(3)] + tasks = [] + create_task = asyncio.create_task + + def schedule(coroutine): + if tasks: + raise RuntimeError("task factory failed") + task = create_task(coroutine) + tasks.append(task) + return task + + with patch("evo.objects.utils._concurrency.asyncio.create_task", side_effect=schedule): + with self.assertRaisesRegex(RuntimeError, "task factory failed"): + await gather_cancelling(*coroutines) + + self.assertTrue(all(task.done() for task in tasks)) + self.assertTrue(all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in coroutines)) + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) + + async def test_semaphore_creation_failure_drains_all_coroutines(self): + async def upload(): + async with get_upload_semaphore(): + return 1 + + coroutines = [upload(), upload()] + with patch("evo.objects.utils._concurrency.asyncio.Semaphore", side_effect=RuntimeError("no semaphore")): + with self.assertRaisesRegex(RuntimeError, "no semaphore"): + await gather_cancelling(*coroutines) + + self.assertTrue(all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in coroutines)) + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) + + +class TestUploadConcurrencyAcrossLoops(TestCase): + @parameterized.expand([("default", None), ("scope", 1)]) + def test_contended_scope_can_span_event_loops(self, _name, limit): + semaphores = [] + + async def run(): + semaphore = get_upload_semaphore() + semaphores.append(semaphore) + + async def upload(): + async with semaphore: + await asyncio.sleep(0) + + await asyncio.gather(*(upload() for _ in range(16))) + + with upload_concurrency(limit) if limit is not None else nullcontext(): + asyncio.run(asyncio.wait_for(run(), timeout=10)) + asyncio.run(asyncio.wait_for(run(), timeout=10)) + + def test_default_scope_is_safe_across_threads_and_event_loops(self): + def get_semaphore(): + async def run(): + return get_upload_semaphore() + + return asyncio.run(run()) + + with ThreadPoolExecutor(max_workers=8) as executor: + semaphores = list(executor.map(lambda _: get_semaphore(), range(32))) + + self.assertEqual(32, len({id(semaphore) for semaphore in semaphores})) + + self.assertIsNot(semaphores[0], semaphores[1]) diff --git a/packages/evo-objects/tests/typed/test_attributes.py b/packages/evo-objects/tests/typed/test_attributes.py index 1a33c32c9..9bbf0929c 100644 --- a/packages/evo-objects/tests/typed/test_attributes.py +++ b/packages/evo-objects/tests/typed/test_attributes.py @@ -11,18 +11,328 @@ from __future__ import annotations +import asyncio +import threading +from copy import deepcopy from typing import Any, cast -from unittest import TestCase +from unittest import IsolatedAsyncioTestCase, TestCase +from unittest.mock import AsyncMock, Mock, patch import pandas as pd from parameterized import parameterized +from evo.common import IFeedback +from evo.objects.typed._model import ModelContext from evo.objects.typed.attributes import ( + Attribute, AttributeDescription, + Attributes, PendingAttribute, UnSupportedDataTypeError, _infer_attribute_type_from_series, ) +from evo.objects.utils._concurrency import get_upload_semaphore, upload_concurrency + + +class TestConcurrentAttributeUploads(IsolatedAsyncioTestCase): + def setUp(self): + self.documents = [ + { + "name": "first", + "key": "stable-key", + "attribute_type": "scalar", + "attribute_description": {"discipline": "geology", "type": "length", "unit": "m"}, + "values": {"data": "old-first", "length": 1}, + }, + {"name": "removed", "attribute_type": "integer", "values": {"data": "old-removed", "length": 1}}, + ] + self.context = ModelContext(obj=Mock(), root_model=None, data_modified={"previous-change"}) + self.attributes = Attributes(self.context, self.documents) + + async def test_attribute_order_and_descriptions_survive_out_of_order_completion(self): + second_uploaded = asyncio.Event() + frame = pd.DataFrame({"first": [1.0], "second": [2.0]}) + frame.attrs["attribute_descriptions"] = {"first": AttributeDescription(unit="m")} + + async def upload(document, *_args, **_kwargs): + if document["name"] == "first": + await second_uploaded.wait() + else: + second_uploaded.set() + document["values"] = {"data": document["name"], "length": 1} + + with ( + patch("evo.objects.typed.attributes.get_data_client"), + patch.object(Attribute, "_upload_attribute_values", side_effect=upload), + ): + result = await asyncio.wait_for(Attributes._data_to_schema(frame, Mock()), timeout=2) + + self.assertEqual(["first", "second"], [document["name"] for document in result]) + self.assertEqual("m", result[0]["attribute_description"]["unit"]) + self.assertNotIn("attribute_description", result[1]) + + async def test_failure_cancels_and_waits_for_other_columns(self): + started = asyncio.Event() + release = asyncio.Event() + cleaned_up = asyncio.Event() + + async def upload(document, *_args, **_kwargs): + if document["name"] == "slow": + started.set() + try: + await release.wait() + finally: + await asyncio.sleep(0) + cleaned_up.set() + else: + await started.wait() + raise RuntimeError("attribute upload failed") + + with ( + patch("evo.objects.typed.attributes.get_data_client"), + patch.object(Attribute, "_upload_attribute_values", side_effect=upload), + ): + try: + with self.assertRaisesRegex(RuntimeError, "attribute upload failed"): + await Attributes._data_to_schema(pd.DataFrame({"slow": [1.0], "failed": [2.0]}), Mock()) + self.assertTrue(cleaned_up.is_set(), "A failed upload left another column running") + finally: + release.set() + await asyncio.wait_for(cleaned_up.wait(), timeout=2) + + @parameterized.expand([("append_attribute",), ("append_attributes",), ("set_attributes",)]) + async def test_mutations_forward_aggregated_feedback(self, method): + frame = pd.DataFrame({"first": [2.0]}) + if method != "append_attribute": + frame["second"] = pd.Categorical(["sand"]) + feedback = Mock(spec=IFeedback) + client = Mock() + + async def upload(df, *, fb, **_kwargs): + fb.progress(0.5) + await asyncio.sleep(0) + fb.progress(1) + return {"data": f"new-{df.columns[0]}", "length": 1} + + async def upload_category(df, *, fb): + values = await upload(df, fb=fb) + return {"values": values, "table": {"data": "lookup"}} + + client.upload_dataframe = AsyncMock(side_effect=upload) + client.upload_category_dataframe = AsyncMock(side_effect=upload_category) + with patch("evo.objects.typed.attributes.get_data_client", return_value=client): + await getattr(self.attributes, method)(frame, fb=feedback) + + progress = [call.args[0] for call in feedback.progress.call_args_list] + self.assertEqual(1, progress[-1]) + self.assertEqual(sorted(progress), progress) + self.assertTrue(all(0 <= value <= 1 for value in progress)) + self.assertIn("new-first", self.context.data_modified) + if method != "append_attribute": + self.assertIn("new-second", self.context.data_modified) + self.assertEqual(0.25, progress[0]) + + @parameterized.expand([("create",), ("append_attributes",), ("set_attributes",)]) + async def test_invalid_description_fails_before_any_upload(self, method): + frame = pd.DataFrame({"first": [1.0], "invalid": [2.0]}) + frame.attrs["attribute_descriptions"] = {"invalid": "not an AttributeDescription"} + original = deepcopy(self.documents) + with ( + patch("evo.objects.typed.attributes.get_data_client") as get_client, + patch.object(Attribute, "_upload_attribute_values") as upload, + ): + with self.assertRaisesRegex(TypeError, "attribute description"): + if method == "create": + await Attributes._data_to_schema(frame, Mock()) + else: + await getattr(self.attributes, method)(frame) + get_client.assert_not_called() + upload.assert_not_awaited() + self.assertEqual(original, self.documents) + self.assertEqual({"previous-change"}, self.context.data_modified) + + @parameterized.expand([("append_attributes",), ("set_attributes",)]) + async def test_mutations_overlap_and_preserve_column_order(self, method): + first_finished = asyncio.Event() + second_started = asyncio.Event() + existing_attribute = self.attributes["first"] + frame = pd.DataFrame({"second": [3.0], "first": [2.0]}) + frame.attrs["attribute_descriptions"] = {"second": AttributeDescription(unit="s")} + + async def upload(document, *_args, **_kwargs): + if document["name"] == "second": + second_started.set() + await first_finished.wait() + else: + await second_started.wait() + first_finished.set() + document["values"] = {"data": f"new-{document['name']}", "length": 1} + + with ( + patch("evo.objects.typed.attributes.get_data_client"), + patch.object(Attribute, "_upload_attribute_values", side_effect=upload), + ): + await asyncio.wait_for(getattr(self.attributes, method)(frame), timeout=2) + + self.assertEqual(["second", "first"], [document["name"] for document in self.documents[-2:]]) + self.assertEqual("s", self.documents[-2]["attribute_description"]["unit"]) + self.assertTrue({"new-first", "new-second"} <= self.context.data_modified) + if method == "set_attributes": + self.assertEqual(2, len(self.documents)) + self.assertEqual("stable-key", self.documents[1]["key"]) + self.assertEqual("m", self.documents[1]["attribute_description"]["unit"]) + self.assertEqual("scalar", self.documents[1]["attribute_type"]) + self.assertEqual("new-first", existing_attribute._data) + + @parameterized.expand( + [ + ("append_failure", "append_attributes", False), + ("set_failure", "set_attributes", False), + ("append_cancel", "append_attributes", True), + ("set_cancel", "set_attributes", True), + ] + ) + async def test_failed_or_cancelled_mutations_preserve_live_state(self, _name, method, cancel): + original = deepcopy(self.documents) + existing_attribute = self.attributes["first"] + first_finished = asyncio.Event() + second_started = asyncio.Event() + release = asyncio.Event() + frame = pd.DataFrame({"first": [2.0], "second": [3.0]}) + + async def upload(document, *_args, **_kwargs): + document["values"] = {"data": f"new-{document['name']}", "length": 1} + if document["name"] == "first": + first_finished.set() + else: + await first_finished.wait() + second_started.set() + if cancel: + await release.wait() + else: + raise RuntimeError("upload failed") + + with ( + patch("evo.objects.typed.attributes.get_data_client"), + patch.object(Attribute, "_upload_attribute_values", side_effect=upload), + ): + task = asyncio.create_task(getattr(self.attributes, method)(frame)) + try: + await asyncio.wait_for(second_started.wait(), timeout=2) + if cancel: + task.cancel() + with self.assertRaises(asyncio.CancelledError if cancel else RuntimeError): + await task + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.assertEqual(original, self.documents) + self.assertEqual("old-first", existing_attribute._data) + self.assertEqual({"previous-change"}, self.context.data_modified) + + async def test_empty_mutations_do_not_upload(self): + with patch("evo.objects.typed.attributes.get_data_client") as get_client: + await self.attributes.append_attributes(pd.DataFrame()) + self.assertEqual(2, len(self.documents)) + await self.attributes.set_attributes(pd.DataFrame()) + self.assertEqual([], self.documents) + get_client.assert_not_called() + self.assertEqual({"previous-change"}, self.context.data_modified) + + async def test_unsupported_type_is_rejected_before_upload(self): + frame = pd.DataFrame({"first": [1.0], "unsupported": [1 + 2j]}) + with patch.object(Attribute, "_upload_attribute_values") as upload: + with self.assertRaises(UnSupportedDataTypeError): + await self.attributes.set_attributes(frame) + upload.assert_not_awaited() + self.assertEqual("old-first", self.documents[0]["values"]["data"]) + + async def test_type_inference_runs_in_worker_with_upload_slot_held(self): + started = threading.Event() + release = threading.Event() + worker_threads = [] + infer = _infer_attribute_type_from_series + + def blocking_infer(series): + worker_threads.append(threading.get_ident()) + started.set() + release.wait() + return infer(series) + + with ( + upload_concurrency(1), + patch("evo.objects.typed.attributes._infer_attribute_type_from_series", side_effect=blocking_infer), + patch("evo.objects.typed.attributes.get_data_client"), + patch.object(Attribute, "_upload_attribute_values", new_callable=AsyncMock), + ): + task = asyncio.create_task(Attributes._data_to_schema(pd.DataFrame({"new": [2.0]}), Mock())) + try: + await asyncio.wait_for(asyncio.to_thread(started.wait), timeout=2) + self.assertEqual(0, get_upload_semaphore()._value) + self.assertNotEqual(threading.get_ident(), worker_threads[0]) + finally: + release.set() + await task + + async def test_set_attributes_rejects_concurrent_document_edits(self): + started = asyncio.Event() + release = asyncio.Event() + + async def upload(*_args, **_kwargs): + started.set() + await release.wait() + + with patch.object(Attribute, "_upload_attribute_values", side_effect=upload): + task = asyncio.create_task(self.attributes.set_attributes(pd.DataFrame({"first": [2.0]}))) + try: + await asyncio.wait_for(started.wait(), timeout=2) + self.documents[0]["external_edit"] = True + release.set() + with self.assertRaisesRegex(RuntimeError, "changed while replacement"): + await task + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.assertEqual("old-first", self.documents[0]["values"]["data"]) + self.assertTrue(self.documents[0]["external_edit"]) + + @parameterized.expand([("string_label", "0"), ("integer_label", 0)]) + async def test_replacement_matches_stringified_names_and_preserves_existing_type(self, _name, label): + self.documents[:] = [ + { + "name": "0", + "key": "stable-key", + "attribute_type": "category", + "table": {"data": "old-lookup"}, + "values": {"data": "old-values", "length": 1}, + } + ] + original = self.documents[0] + attribute = self.attributes["0"] + client = Mock() + client.upload_category_dataframe = AsyncMock( + return_value={ + "table": {"data": "new-lookup"}, + "values": {"data": "new-values", "length": 1}, + } + ) + client.upload_dataframe = AsyncMock() + with patch("evo.objects.typed.attributes.get_data_client", return_value=client): + # A float column would infer scalar for a new attribute, but replacements keep their existing type. + await self.attributes.set_attributes(pd.DataFrame({label: [2.0]})) + + self.assertIs(original, self.documents[0]) + self.assertEqual("category", attribute.attribute_type) + self.assertEqual("stable-key", attribute.key) + self.assertEqual("new-values", attribute._data) + self.assertEqual({"data": "new-lookup"}, original["table"]) + client.upload_category_dataframe.assert_awaited_once() + client.upload_dataframe.assert_not_awaited() + self.assertEqual({"previous-change", "new-values"}, self.context.data_modified) class TestAttributeTypeInference(TestCase): diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index e31edc47a..e57cb665b 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -11,11 +11,13 @@ from __future__ import annotations +import asyncio import contextlib import copy import dataclasses import inspect import math +import tempfile import uuid import warnings from datetime import date @@ -24,13 +26,16 @@ import numpy as np import numpy.testing as npt import pandas as pd +import pyarrow.parquet as pq from parameterized import parameterized from evo.common import Environment, StaticContext from evo.common.test_tools import BASE_URL, ORG, WORKSPACE_ID, TestWithConnector -from evo.common.utils import NoFeedback +from evo.common.utils import Cache, NoFeedback from evo.objects import ObjectReference +from evo.objects.io import ObjectDataUpload from evo.objects.typed import BoundingBox +from evo.objects.typed._model import SchemaBuilder from evo.objects.typed.attributes import AttributeDescription from evo.objects.typed.base import BaseObject from evo.objects.typed.downhole_collection import ( @@ -40,6 +45,8 @@ IntervalCollection, ) from evo.objects.typed.exceptions import ObjectValidationError +from evo.objects.utils import upload_concurrency +from evo.objects.utils.data import _iter_refs from .helpers import MockClient @@ -49,7 +56,7 @@ def _make_example_data( description: str | None = None, tags: dict[str, str] | None = None, attributes: pd.DataFrame | None = None, - collections: list[DistanceCollection] | None = None, + collections: list[DistanceCollection | IntervalCollection] | None = None, ) -> DownholeCollectionData: """Helper to build a simple two-hole DownholeCollectionData.""" # Concatenated path table: hole 0 has 4 rows, hole 1 has 3 rows @@ -225,6 +232,122 @@ async def test_create_with_empty_collections(self): result = await DownholeCollection.create(context=self.context, data=data) self.assertIsInstance(result, DownholeCollection) + async def test_collection_uploads_run_concurrently_and_preserve_order(self): + template = _make_example_data().collections[0] + interval = IntervalCollection( + name="intervals", + holes=template.holes.copy(), + table=pd.DataFrame({"from": [0.0, 1.0, 2.0, 3.0], "to": [1.0, 2.0, 3.0, 4.0]}), + unit="m", + ) + collections = [ + dataclasses.replace(template, name="first"), + interval, + dataclasses.replace(template, name="last"), + ] + data = _make_example_data(collections=collections) + started = set() + all_started = asyncio.Event() + release = asyncio.Event() + build_from_data = SchemaBuilder.build_from_data + + async def build(builder, payload, skip_sub_models=None): + if isinstance(payload, (DistanceCollection, IntervalCollection)): + started.add(payload.name) + if len(started) == len(collections): + all_started.set() + await release.wait() + return await build_from_data(builder, payload, skip_sub_models=skip_sub_models) + + with ( + self._mock_geoscience_objects(), + patch.object(SchemaBuilder, "build_from_data", new=build), + ): + creation = asyncio.create_task(DownholeCollection.create(context=self.context, data=data)) + try: + await asyncio.wait_for(all_started.wait(), timeout=2) + release.set() + result = await asyncio.wait_for(creation, timeout=2) + finally: + release.set() + creation.cancel() + await asyncio.gather(creation, return_exceptions=True) + + self.assertEqual(["first", "intervals", "last"], result.collections.names()) + self.assertEqual("m", result.collections.get("intervals").from_to.unit) + npt.assert_array_equal(interval.table, await result.collections.get("intervals").to_dataframe()) + self.assertEqual( + list(template.table.columns), + list((await result.collections.get("last").to_dataframe()).columns), + ) + + @parameterized.expand([("default", None), ("serial", 1), ("lower", 3), ("higher", 12)]) + async def test_nested_table_uploads_share_a_concurrency_limit(self, _name, configured_limit): + limit = configured_limit if configured_limit is not None else 8 + template = _make_example_data().collections[0] + holes = template.holes.astype({"hole_index": "int32", "offset": "uint64", "count": "uint64"}) + table = template.table.assign(**{f"reading_{i}": np.arange(4, dtype=float) for i in range(16)}) + table["category"] = pd.Categorical(["sand", "clay", "sand", "silt"]) + data = _make_example_data( + collections=[ + dataclasses.replace(template, name=f"collection_{i}", holes=holes, table=table) for i in range(3) + ] + ) + data = dataclasses.replace( + data, holes=data.holes.astype({"hole_index": "int32", "offset": "uint64", "count": "uint64"}) + ) + active = 0 + peak = 0 + uploaded = [] + + async def upload(upload, cache, transport, **kwargs): + nonlocal active, peak + self.assertTrue( + (cache.get_location(environment=self.environment, scope="geoscience-object") / upload.name).is_file() + ) + active += 1 + peak = max(peak, active) + try: + await asyncio.sleep(0.01) + uploaded.append(upload.name) + finally: + active -= 1 + + with tempfile.TemporaryDirectory() as directory: + cache = Cache(directory) + context = StaticContext.from_environment( + environment=self.environment, connector=self.connector, cache=cache + ) + with ( + patch.object(ObjectDataUpload, "upload_from_cache", upload), + upload_concurrency(configured_limit) if configured_limit is not None else contextlib.nullcontext(), + ): + document = await DownholeCollection._data_to_schema(data, context) + cache_location = cache.get_location(environment=self.environment, scope="geoscience-object") + for collection in document["collections"]: + distance = collection["distance"] + geometry = pq.read_table(cache_location / distance["values"]["data"]) + npt.assert_array_equal(table["distance"], geometry.column(0).to_numpy()) + self.assertEqual(list(table.columns[1:]), [a["name"] for a in distance["attributes"]]) + category = distance["attributes"][-1] + codes = pq.read_table(cache_location / category["values"]["data"]).column(0).to_numpy() + lookup = pq.read_table(cache_location / category["table"]["data"]).to_pandas() + decoded = pd.Series(codes).map(dict(zip(lookup["key"], lookup["value"], strict=True))) + npt.assert_array_equal(table["category"], decoded) + self.assertEqual(category["values"]["length"], len(table)) + self.assertEqual(distance["values"]["length"], len(table)) + + if limit > 1: + self.assertGreater(peak, 1) + else: + self.assertEqual(1, peak) + self.assertLessEqual(peak, limit) + self.assertEqual(0, active) + references = list(_iter_refs(document)) + self.assertGreater(len(references), 50) + self.assertEqual(set(references), set(uploaded)) + self.assertEqual([c.name for c in data.collections], [c["name"] for c in document["collections"]]) + async def test_mixed_collections_round_trip_and_mutation(self): interval = IntervalCollection( name="intervals", diff --git a/packages/evo-objects/tests/typed/test_model.py b/packages/evo-objects/tests/typed/test_model.py index a432590b4..6cc229bf5 100644 --- a/packages/evo-objects/tests/typed/test_model.py +++ b/packages/evo-objects/tests/typed/test_model.py @@ -13,17 +13,23 @@ from __future__ import annotations +import asyncio import contextlib +from copy import deepcopy from dataclasses import dataclass from typing import Annotated, ClassVar -from unittest.mock import patch +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, Mock, patch import pandas as pd +from parameterized import parameterized -from evo.common import Environment, StaticContext +from evo.common import Environment, IFeedback, StaticContext from evo.common.test_tools import BASE_URL, ORG, WORKSPACE_ID, TestWithConnector +from evo.objects import DownloadedObject from evo.objects.typed._data import DataTable, DataTableAndAttributes -from evo.objects.typed._model import SchemaBuilder, SchemaList, SchemaLocation, SchemaModel +from evo.objects.typed._model import ModelContext, SchemaBuilder, SchemaList, SchemaLocation, SchemaModel +from evo.objects.typed.attributes import AttributeDescription from evo.objects.utils.table_formats import FLOAT_ARRAY_3, KnownTableFormat from .helpers import MockClient @@ -71,6 +77,48 @@ class Second(SchemaModel): class InvalidUnionList(SchemaList[First | Second]): pass + async def test_schema_list_uploads_items_concurrently_and_preserves_order(self): + class Models(SchemaList[SchemaModel]): + pass + + started = set() + all_started = asyncio.Event() + release = asyncio.Event() + second_finished = asyncio.Event() + + async def convert(value, context): + started.add(value) + if len(started) == 2: + all_started.set() + await release.wait() + if value == 0: + await second_finished.wait() + else: + second_finished.set() + return {"value": value} + + with patch.object(SchemaModel, "_data_to_schema", side_effect=convert): + task = asyncio.create_task(Models._data_to_schema([0, 1], self.context)) + try: + await asyncio.wait_for(all_started.wait(), timeout=2) + release.set() + result = await asyncio.wait_for(task, timeout=2) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.assertEqual([{"value": 0}, {"value": 1}], result) + + async def test_schema_list_accepts_missing_and_empty_data(self): + for data in (None, [], ()): + with self.subTest(data=data): + self.assertEqual([], await SchemaList._data_to_schema(data, self.context)) + + async def test_schema_list_rejects_non_sequences(self): + with self.assertRaisesRegex(TypeError, "Expected a sequence"): + await SchemaList._data_to_schema(42, self.context) + class CoordinateTable(DataTable): table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 @@ -82,6 +130,162 @@ class ExtendedLocations(DataTableAndAttributes): point_count: Annotated[int, SchemaLocation("point_count")] +class TestDataTableAndAttributesMutation(IsolatedAsyncioTestCase): + def setUp(self): + self.document = { + "coordinates": {"data": "old-coordinates", "length": 1}, + "attributes": [ + { + "name": "first", + "key": "stable", + "attribute_type": "scalar", + "values": {"data": "old-first", "length": 1}, + } + ], + } + self.context = ModelContext(obj=Mock(), root_model=None, data_modified={"previous-change"}) + self.model = ExtendedLocations(self.context, self.document) + self.frame = pd.DataFrame({"x": [1.0], "y": [2.0], "z": [3.0], "first": [4.0]}) + + @parameterized.expand([("success", None), ("table_fails", "x"), ("attribute_fails", "first")]) + async def test_concurrent_mutations_commit_only_after_all_uploads_succeed(self, _name, failing_column): + started = set() + all_started = asyncio.Event() + successful_finished = asyncio.Event() + original = deepcopy(self.document) + original_table = self.model._table + original_attributes = self.model.attributes + feedback = Mock(spec=IFeedback) + + async def upload(frame, *, fb, **_kwargs): + column = frame.columns[0] + started.add(column) + if len(started) == 2: + all_started.set() + await all_started.wait() + if column == failing_column: + await successful_finished.wait() + raise RuntimeError("upload failed") + fb.progress(1) + successful_finished.set() + return {"data": f"new-{column}", "length": 1} + + client = Mock(upload_dataframe=AsyncMock(side_effect=upload)) + with ( + patch("evo.objects.typed.attributes.get_data_client", return_value=client), + patch("evo.objects.typed._data.get_data_client", return_value=client) as get_client, + ): + if failing_column: + with self.assertRaisesRegex(RuntimeError, "upload failed"): + await asyncio.wait_for(self.model.from_dataframe(self.frame, fb=feedback), timeout=2) + else: + await asyncio.wait_for(self.model.from_dataframe(self.frame, fb=feedback), timeout=2) + + if failing_column: + self.assertEqual(original, self.document) + self.assertEqual({"previous-change"}, self.context.data_modified) + else: + self.assertEqual("new-x", self.document["coordinates"]["data"]) + self.assertEqual("new-first", self.document["attributes"][0]["values"]["data"]) + self.assertEqual("stable", self.document["attributes"][0]["key"]) + self.assertEqual({"previous-change", "new-x", "new-first"}, self.context.data_modified) + self.assertEqual([0.5, 1], [call.args[0] for call in feedback.progress.call_args_list]) + self.assertIs(original_table, self.model._table) + self.assertIs(original_attributes, self.model.attributes) + get_client.assert_called_once_with(self.context.obj) + + async def test_cancellation_keeps_table_and_attributes_unchanged(self): + original = deepcopy(self.document) + started = asyncio.Event() + release = asyncio.Event() + + async def upload(frame, **_kwargs): + if frame.columns[0] == "first": + started.set() + await release.wait() + return {"data": "new", "length": 1} + + client = Mock(upload_dataframe=AsyncMock(side_effect=upload)) + with ( + patch("evo.objects.typed.attributes.get_data_client", return_value=client), + patch("evo.objects.typed._data.get_data_client", return_value=client), + ): + task = asyncio.create_task(self.model.from_dataframe(self.frame)) + try: + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + self.assertEqual(original, self.document) + self.assertEqual({"previous-change"}, self.context.data_modified) + + async def test_attribute_descriptions_are_prepared_once_before_combined_uploads(self): + self.frame["second"] = [5.0] + self.frame.attrs["attribute_descriptions"] = { + "first": AttributeDescription(unit="m"), + "second": AttributeDescription(unit="s"), + } + client = Mock(upload_dataframe=AsyncMock(return_value={"data": "new", "length": 1})) + with ( + patch("evo.objects.typed.attributes.get_data_client", return_value=client), + patch("evo.objects.typed._data.get_data_client", return_value=client), + patch.object(AttributeDescription, "to_schema", autospec=True, return_value={"unit": "s"}) as describe, + ): + await self.model.from_dataframe(self.frame) + self.assertEqual(1, describe.call_count) + self.assertEqual({"unit": "s"}, self.document["attributes"][1]["attribute_description"]) + + async def test_invalid_attribute_metadata_prevents_table_upload(self): + self.frame["second"] = [5.0] + self.frame.attrs["attribute_descriptions"] = {"second": "invalid"} + client = Mock(upload_dataframe=AsyncMock()) + with ( + patch("evo.objects.typed.attributes.get_data_client", return_value=client), + patch("evo.objects.typed._data.get_data_client", return_value=client), + ): + with self.assertRaisesRegex(TypeError, "attribute description"): + await self.model.from_dataframe(self.frame) + client.upload_dataframe.assert_not_awaited() + + async def test_table_apply_hook_runs_with_attributes(self): + client = Mock(upload_dataframe=AsyncMock(return_value={"data": "new", "length": 1})) + with ( + patch("evo.objects.typed.attributes.get_data_client", return_value=client), + patch("evo.objects.typed._data.get_data_client", return_value=client), + patch.object( + self.model._table, + "_apply_uploaded_table", + wraps=self.model._table._apply_uploaded_table, + ) as apply_hook, + ): + await self.model.from_dataframe(self.frame) + + apply_hook.assert_called_once() + + @parameterized.expand([("success", False), ("failure", True)]) + async def test_table_only_update_clears_attributes_only_on_success(self, _name, fail): + original = deepcopy(self.document) + upload = ( + AsyncMock(side_effect=RuntimeError("failed")) + if fail + else AsyncMock(return_value={"data": "new-table", "length": 1}) + ) + with patch("evo.objects.typed._data.get_data_client", return_value=Mock(upload_dataframe=upload)) as get_client: + if fail: + with self.assertRaisesRegex(RuntimeError, "failed"): + await self.model.from_dataframe(self.frame[["x", "y", "z"]]) + self.assertEqual(original, self.document) + else: + await self.model.from_dataframe(self.frame[["x", "y", "z"]]) + self.assertEqual([], self.document["attributes"]) + self.assertEqual("new-table", self.document["coordinates"]["data"]) + get_client.assert_called_once_with(self.context.obj) + + class TestDataTableAndAttributesBuildFromData(TestWithConnector): """Tests for DataTableAndAttributes._data_to_schema using build_from_data with skip_sub_models.""" @@ -93,6 +297,26 @@ def setUp(self) -> None: connector=self.connector, ) + async def test_data_table_uses_real_data_client_context(self): + obj = DownloadedObject( + object_=Mock(), + metadata=Mock(environment=self.environment), + urls_by_name={}, + connector=self.connector, + cache=Mock(), + ) + table = CoordinateTable(obj, {"data": "old", "length": 1}) + frame = pd.DataFrame({"x": [1.0], "y": [2.0], "z": [3.0]}) + with patch( + "evo.objects.utils.data.ObjectDataClient.upload_dataframe", + new_callable=AsyncMock, + return_value={"data": "new", "length": 1}, + ) as upload: + await table.from_dataframe(frame) + + upload.assert_awaited_once() + self.assertEqual("new", table._data) + @contextlib.contextmanager def _mock_geoscience_objects(self): mock_client = MockClient(self.environment) @@ -111,7 +335,7 @@ async def test_data_to_schema_extra_properties(self): with self._mock_geoscience_objects(): result = await ExtendedLocations._data_to_schema(df, self.context) - # Coordinates and attributes should be handled by the manual set_sub_model_value calls + # Coordinates and attributes are handled before the remaining properties. self.assertIn("coordinates", result) self.assertEqual(result["attributes"], []) diff --git a/packages/evo-objects/tests/typed/test_pointset.py b/packages/evo-objects/tests/typed/test_pointset.py index e77cee99c..95571f1ac 100644 --- a/packages/evo-objects/tests/typed/test_pointset.py +++ b/packages/evo-objects/tests/typed/test_pointset.py @@ -11,6 +11,7 @@ from __future__ import annotations +import asyncio import contextlib import sys import uuid @@ -276,6 +277,37 @@ async def test_validate_attribute_length(self): obj.validate() self.assertIn("does not match expected length", str(cm.exception)) + async def test_failed_attribute_update_preserves_coordinates_and_bounding_box(self): + with self._mock_geoscience_objects() as client: + obj = await PointSet.create(context=self.context, data=self.example_pointset) + original = obj.as_dict() + modified = obj._context.data_modified.copy() + coordinates_uploaded = asyncio.Event() + original_upload = client.upload_dataframe + + async def upload(frame, *args, **kwargs): + if list(frame.columns) == ["x", "y", "z"]: + result = await original_upload(frame, *args, **kwargs) + coordinates_uploaded.set() + return result + await coordinates_uploaded.wait() + raise RuntimeError("attribute upload failed") + + frame = pd.DataFrame({"x": [10.0], "y": [-5.0], "z": [20.0], "value": [7.0]}) + with patch.object(client, "upload_dataframe", side_effect=upload): + with self.assertRaisesRegex(RuntimeError, "attribute upload failed"): + await asyncio.wait_for(obj.locations.from_dataframe(frame), timeout=2) + + self.assertEqual(original, obj.as_dict()) + self.assertEqual(modified, obj._context.data_modified) + + async def test_table_only_mutation_updates_bounding_box(self): + with self._mock_geoscience_objects(): + obj = await PointSet.create(context=self.context, data=self.example_pointset) + await obj.locations.from_dataframe(pd.DataFrame({"x": [10.0], "y": [-5.0], "z": [20.0]})) + self._assert_bounding_box_equal(obj.bounding_box, 10.0, 10.0, -5.0, -5.0, 20.0, 20.0) + self.assertEqual(0, len(obj.locations.attributes)) + async def test_json(self): """Test the JSON structure of the created object.""" with self._mock_geoscience_objects() as mock_client: diff --git a/packages/evo-sdk-common/pyproject.toml b/packages/evo-sdk-common/pyproject.toml index 12099a1a3..e121e5e6c 100644 --- a/packages/evo-sdk-common/pyproject.toml +++ b/packages/evo-sdk-common/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "evo-sdk-common" description = "Python package that establishes a common framework for use by client libraries that interact with Seequent Evo APIs" -version = "0.5.31" +version = "0.5.32" requires-python = ">=3.10" license-files = ["LICENSE.md"] dynamic = ["readme"] diff --git a/packages/evo-sdk-common/src/evo/aio/_helpers.py b/packages/evo-sdk-common/src/evo/aio/_helpers.py index 60443ff3b..10f01c12e 100644 --- a/packages/evo-sdk-common/src/evo/aio/_helpers.py +++ b/packages/evo-sdk-common/src/evo/aio/_helpers.py @@ -32,7 +32,7 @@ class Config: """The value to provide in the `User-Agent` header.""" num_pools: int - """Number of connection pools to cache before discarding the least recently used pool.""" + """Maximum simultaneous HTTP connections across all hosts; 0 disables the limit.""" verify_ssl: bool """Verify SSL certificates.""" @@ -63,7 +63,7 @@ def __init__( self, num_pools: int, verify_ssl: bool, retry_handler: Retry, proxy: StrOrURL | None, close_grace_period_ms: int ) -> None: """ - :param num_pools: Number of connection pools to cache before discarding the least recently used pool. + :param num_pools: Maximum simultaneous HTTP connections across all hosts; 0 disables the limit. :param verify_ssl: Verify SSL certificates. :param retry_handler: Retry handler. """ diff --git a/packages/evo-sdk-common/src/evo/aio/transport.py b/packages/evo-sdk-common/src/evo/aio/transport.py index 26f4965f9..796b77ee2 100644 --- a/packages/evo-sdk-common/src/evo/aio/transport.py +++ b/packages/evo-sdk-common/src/evo/aio/transport.py @@ -58,7 +58,8 @@ def __init__( :param max_attempts: Configure the number of retries to allow before raising a `MaxRetryError` exception. :param backoff_method: Configure the retry backoff implementation. - :param num_pools: Number of connection pools to cache before discarding the least recently used pool. + :param num_pools: Maximum simultaneous HTTP connections across all hosts, shared by requests using this + transport. Defaults to 4; 0 disables the connection limit. Independent of table upload concurrency. :param verify_ssl: Verify SSL certificates. This should not be disabled unless you really absolutely have to, and never in production environments. :param close_grace_period_ms: Grace period (in milliseconds) to wait for connections to close. diff --git a/packages/evo-sdk-common/src/evo/common/utils/retry.py b/packages/evo-sdk-common/src/evo/common/utils/retry.py index 9ae492ce0..9232ccc70 100644 --- a/packages/evo-sdk-common/src/evo/common/utils/retry.py +++ b/packages/evo-sdk-common/src/evo/common/utils/retry.py @@ -95,10 +95,11 @@ def __init__(self, logger: logging.Logger, max_attempts: int, backoff_method: Ba self.__errors = [] def add_error(self, error: Exception) -> None: - self.__logger.error(error) + exhausted = self.__current_handler._number >= self.__max_attempts + self.__logger.log(logging.ERROR if exhausted else logging.WARNING, error) self.__logger.debug(f"{self.__current_handler} failed") self.__errors.append(error) - if self.__current_handler._number >= self.__max_attempts: + if exhausted: raise RetryError("Retry failed", self.__errors) async def __anext__(self) -> RetryHandler: @@ -193,6 +194,7 @@ class Retry: """AsyncIterable retry implementation for handling exceptions in retryable logic. A single Retry object can be used to handle multiple requests. The only state in this class is the configuration. + Failed attempts are logged at WARNING while retries remain, and at ERROR when attempts are exhausted. Usage:: retry = Retry(logger=logging.getLogger(__name__), max_attempts=3, backoff_method=BackoffLinear(1)) diff --git a/packages/evo-sdk-common/tests/common/utils/test_retry.py b/packages/evo-sdk-common/tests/common/utils/test_retry.py index b1e0e75d9..c2d8c7f8a 100644 --- a/packages/evo-sdk-common/tests/common/utils/test_retry.py +++ b/packages/evo-sdk-common/tests/common/utils/test_retry.py @@ -80,6 +80,16 @@ async def test_successful_attempt(self, mock_sleep: mock.MagicMock) -> None: mock_sleep.assert_not_called() + async def test_only_exhausted_attempts_log_at_error(self): + retry_logger = mock.Mock(spec=logging.Logger) + error = ConnectionError("retry") + with self.assertRaises(RetryError): + async for handler in Retry(retry_logger, max_attempts=2, backoff_method=BackoffLinear(0)): + handler.set_exception(error) + self.assertEqual( + [mock.call(logging.WARNING, error), mock.call(logging.ERROR, error)], retry_logger.log.call_args_list + ) + @mock.patch("asyncio.sleep", spec_set=True) async def test_max_attempts(self, mock_sleep: mock.MagicMock) -> None: with self.assertRaises(RetryError): diff --git a/uv.lock b/uv.lock index af9afe08d..65ade73a4 100644 --- a/uv.lock +++ b/uv.lock @@ -1321,7 +1321,7 @@ test = [ [[package]] name = "evo-sdk-common" -version = "0.5.31" +version = "0.5.32" source = { editable = "packages/evo-sdk-common" } dependencies = [ { name = "pure-interface" },