From 13070b9797fd4ec0c8e833f2f8b65ea2a16d910f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Wed, 4 Nov 2026 15:25:19 +0000 Subject: [PATCH 1/5] Parallelize typed object uploads and harden mutations Restore the original upload changes on current main, retaining upstream error context and covering retry exhaustion diagnostics. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- packages/evo-objects/README.md | 8 + .../src/evo/objects/typed/_data.py | 24 +- .../src/evo/objects/typed/_model.py | 33 +- .../src/evo/objects/typed/attributes.py | 85 +++-- .../evo/objects/typed/downhole_collection.py | 22 +- .../src/evo/objects/typed/pointset.py | 16 +- .../src/evo/objects/utils/_concurrency.py | 51 +++ .../evo-objects/src/evo/objects/utils/data.py | 65 +++- .../evo-objects/tests/test_data_client.py | 313 +++++++++++++++++- .../tests/test_upload_concurrency.py | 178 ++++++++++ .../tests/typed/test_attributes.py | 217 +++++++++++- .../tests/typed/test_downhole_collection.py | 115 ++++++- .../evo-objects/tests/typed/test_model.py | 178 +++++++++- .../evo-objects/tests/typed/test_pointset.py | 32 ++ 14 files changed, 1250 insertions(+), 87 deletions(-) create mode 100644 packages/evo-objects/src/evo/objects/utils/_concurrency.py create mode 100644 packages/evo-objects/tests/test_upload_concurrency.py diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index ab93e4d37..07da4feab 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -105,6 +105,14 @@ 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. +Object creation and replacement upload independent collections, geometry, and attributes concurrently, preserving +their input order. Nested uploads share a limit of eight tables per serialization operation, including Parquet writing; +direct concurrent table uploads share a per-client limit of eight. Transient upload failures receive up to three +attempts with exponential backoff, reusing cached Parquet data. Permanent failures or exhausted retries cancel and wait +for sibling work before returning. Multi-column attribute and combined table/attribute mutations upload independent +data concurrently and update their in-memory documents only after all uploads succeed. Separately awaited +`collections.add()` calls remain sequential. + ## 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/src/evo/objects/typed/_data.py b/packages/evo-objects/src/evo/objects/typed/_data.py index 8868c9afe..28504d03c 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_uploads from evo.objects.utils.table_formats import KnownTableFormat from ._model import SchemaBuilder, SchemaLocation, SchemaModel @@ -49,9 +50,11 @@ 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._data_to_schema(df, self._obj, fb=fb)) - # Mark the context as modified so loading data is not allowed + def _apply_uploaded_table(self, df: pd.DataFrame, document: dict[str, Any]) -> None: + """Commit uploaded values and any subclass-specific derived metadata.""" + self._document.update(document) self._context.mark_modified(self._data) @classmethod @@ -118,10 +121,18 @@ 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) + # Validate local metadata before either branch starts remote work. + self.attributes._attribute_descriptions(attr_df) + table_fb, attr_fb = split_feedback(fb, [1.0, float(len(attr_df.columns))]) + table_document, attribute_documents = await gather_uploads( + self._table._data_to_schema(table_df, self._obj, fb=table_fb), + self.attributes._prepare_attributes(attr_df, attr_fb), + ) + self._table._apply_uploaded_table(table_df, table_document) + self.attributes._replace_uploaded_attributes(attribute_documents) else: + await self._table.from_dataframe(table_df, fb=fb) self.attributes.clear() def validate(self) -> None: @@ -152,8 +163,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..854653f3e 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_uploads 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_uploads( + *(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_uploads(*(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..18f19e968 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -21,8 +21,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_uploads from evo.objects.utils.table_formats import ( BOOL_ARRAY_1, DATE_TIME_ARRAY, @@ -305,6 +306,18 @@ async def _data_to_schema( await cls._upload_attributes_to_list(result, data, context) return result + @staticmethod + def _attribute_descriptions(df: pd.DataFrame) -> dict[Any, dict[str, Any]]: + descriptions = df.attrs.get("attribute_descriptions", {}) + result = {} + for col in df.columns: + description = descriptions.get(col) + if description is not None: + if not isinstance(description, AttributeDescription): + raise TypeError("attribute description must be a AttributeDescription.") + result[col] = description.to_schema() + return result + @staticmethod async def _upload_attributes_to_list( attributes_list: list[dict[str, Any]], @@ -321,9 +334,14 @@ async def _upload_attributes_to_list( :param context: The context used for data upload operations. :param fb: Optional feedback object to report progress. """ + descriptions = Attributes._attribute_descriptions(df) + if len(df.columns) == 0: + fb.progress(1) + return data_client = get_data_client(context) + column_fbs = split_feedback(fb, [1.0] * len(df.columns)) - for col in df.columns: + async def upload_column(col: Any, col_fb: IFeedback) -> dict[str, Any]: series = df[col] attribute_type = _infer_attribute_type_from_series(df[col]) @@ -337,16 +355,16 @@ async def _upload_attributes_to_list( 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) + await Attribute._upload_attribute_values(attr_doc, col_df, attribute_type, data_client, fb=col_fb) - 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 + if description := descriptions.get(col): + attr_doc["attribute_description"] = description + return attr_doc - attributes_list.append(attr_doc) + # Columns are independent, so upload them concurrently while preserving column order in the result. + attributes_list.extend( + await gather_uploads(*(upload_column(col, 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. @@ -395,9 +413,11 @@ 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: list[dict[str, Any]] = [] + await self._upload_attributes_to_list(documents, 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 +426,43 @@ async def set_attributes(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): :param fb: Optional feedback object to report upload progress. """ + self._replace_uploaded_attributes(await self._prepare_attributes(df, fb)) + + async def _prepare_attributes(self, df: pd.DataFrame, fb: IFeedback) -> list[dict[str, Any]]: + """Upload replacement values without changing live documents or modification tracking.""" + self._attribute_descriptions(df) + if len(df.columns) == 0: + fb.progress(1) + return [] attributes_by_name = {attr.name: attr for attr in self} - self.clear() - for col in df.columns: + 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) -> dict[str, Any]: attribute_df = df[[col]] attribute = attributes_by_name.get(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, attribute_df, attribute.attribute_type, data_client, fb=col_fb + ) + return document + documents: list[dict[str, Any]] = [] + await self._upload_attributes_to_list(documents, attribute_df, self._obj, fb=col_fb) + return documents[0] + + return await gather_uploads(*(upload_column(col, col_fb) for col, col_fb in zip(df.columns, column_fbs))) + + def _replace_uploaded_attributes(self, documents: list[dict[str, Any]]) -> None: + existing = {document["name"]: document for document in self._document} + for index, document in enumerate(documents): + if document["name"] in existing: + original = existing[document["name"]] + original.update(document) + documents[index] = original + self._document[:] = documents + for document in documents: + self._context.mark_modified(document["values"]["data"]) 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..24665e31a 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_uploads 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_uploads( + 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_uploads(*(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/_concurrency.py b/packages/evo-objects/src/evo/objects/utils/_concurrency.py new file mode 100644 index 000000000..cc0397b58 --- /dev/null +++ b/packages/evo-objects/src/evo/objects/utils/_concurrency.py @@ -0,0 +1,51 @@ +# 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 nested table-serialization operations.""" + +import asyncio +from collections.abc import Coroutine +from contextvars import ContextVar +from typing import Any, TypeVar + +_T = TypeVar("_T") +_MAX_CONCURRENT_UPLOADS = 8 +_upload_semaphore: ContextVar[asyncio.Semaphore | None] = ContextVar("object_upload_semaphore", default=None) + + +def get_upload_semaphore() -> asyncio.Semaphore | None: + """Return the serialization's shared limit, if a serialization is active.""" + return _upload_semaphore.get() + + +async def gather_uploads(*coroutines: Coroutine[Any, Any, _T]) -> list[_T]: + """Collect ordered results, cancelling and draining siblings on failure or cancellation.""" + token = None + tasks: list[asyncio.Task[_T]] = [] + try: + if _upload_semaphore.get() is None: + token = _upload_semaphore.set(asyncio.Semaphore(_MAX_CONCURRENT_UPLOADS)) + for coroutine in coroutines: + tasks.append(asyncio.create_task(coroutine)) + # Unlike gather(), this does not cancel children before our own cleanup runs. + for completed in asyncio.as_completed(tasks): + await completed + return [task.result() for task in tasks] + except BaseException: + for coroutine in coroutines[len(tasks) :]: + coroutine.close() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + finally: + if token is not None: + _upload_semaphore.reset(token) diff --git a/packages/evo-objects/src/evo/objects/utils/data.py b/packages/evo-objects/src/evo/objects/utils/data.py index e698663c9..7482f37b0 100644 --- a/packages/evo-objects/src/evo/objects/utils/data.py +++ b/packages/evo-objects/src/evo/objects/utils/data.py @@ -19,12 +19,13 @@ 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, 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 _MAX_CONCURRENT_UPLOADS, gather_uploads, get_upload_semaphore 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 @@ -48,6 +49,15 @@ _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, 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, (TransportError, asyncio.TimeoutError, TimeoutError, ConnectionError)) def _iter_refs(target: Any, _key: str | None = None) -> Iterator[str]: @@ -109,6 +119,7 @@ def __init__(self, environment: Environment, connector: APIConnector, cache: ICa self._environment = environment self._connector = connector self._cache = cache + self._upload_semaphore = asyncio.Semaphore(_MAX_CONCURRENT_UPLOADS) @property def cache_location(self) -> Path: @@ -196,6 +207,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,21 +217,46 @@ 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 URL or storage failures. Defaults to three attempts with + exponential backoff. Retries reuse the cached table; exhausted attempts raise RetryError. :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): + if retry is None: + retry = Retry(logger=logger) + async with get_upload_semaphore() or self._upload_semaphore: + # Compression and hashing must not block other uploads on the event loop. + writer = asyncio.create_task(asyncio.to_thread(self.save_table, table, table_format)) 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) - return table_info + table_info = await asyncio.shield(writer) + except asyncio.CancelledError: + # A running worker cannot be cancelled; finish its cache write before allowing a retry. + await asyncio.gather(writer, return_exceptions=True) + raise + 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): + async for attempt in retry: + 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) + except ( + EvoAPIException, + TransportError, + asyncio.TimeoutError, + TimeoutError, + ConnectionError, + ChunkedIOError, + ) as exc: + if not _is_transient_upload_error(exc): + raise + attempt.set_exception(exc) + fb.progress(0) + return table_info async def upload_category_table(self, table: pa.Table, fb: IFeedback = NoFeedback) -> CategoryInfo: """Upload a pyarrow table with category data to the geoscience object service, returning the reference to @@ -278,8 +315,10 @@ async def upload_category_table(self, table: pa.Table, fb: IFeedback = NoFeedbac 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) + value_table, lookup_table = await gather_uploads( + self.upload_table(values, fb=values_fb, table_format=value_table_format), + self.upload_table(lookup, fb=lookup_fb, table_format=LOOKUP_TABLE_INT32), + ) return CategoryInfo(values=cast(ArrayTableInfo, value_table), table=cast(LookupTableInfo, lookup_table)) async def download_table( diff --git a/packages/evo-objects/tests/test_data_client.py b/packages/evo-objects/tests/test_data_client.py index 0fa0109f6..31ab48854 100644 --- a/packages/evo-objects/tests/test_data_client.py +++ b/packages/evo-objects/tests/test_data_client.py @@ -9,8 +9,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import json import sys +import threading from unittest import mock from uuid import UUID @@ -21,12 +23,14 @@ 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.common.exceptions import BadRequestException, EvoAPIException, RetryError, TransportError +from evo.common.io.exceptions import ChunkedIOError, DataExistsError 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._concurrency import gather_uploads from helpers import NoImport, UnloadModule, get_sample_table_and_bytes @@ -200,17 +204,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(RetryError if upload_type == "table_transient" else error_type) as cm, ): destination.upload_file.side_effect = error if upload_type == "referenced": @@ -218,10 +226,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.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 +416,286 @@ async def _mock_upload_file_side_effect(*args, **kwargs): ) self.assertIs(mock_table_info, actual_table_info) + @parameterized.expand([("direct", False), ("nested_group", True)]) + async def test_cancelled_upload_waits_for_parquet_writer(self, _name: str, nested_group: bool) -> None: + started = asyncio.Event() + release = threading.Event() + finished = threading.Event() + loop = asyncio.get_running_loop() + + def save_table(*_args): + loop.call_soon_threadsafe(started.set) + try: + if not release.wait(timeout=5): + 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(): + operation = self.data_client.upload_table(pa.table({"value": [1.0]})) + if nested_group: + return await gather_uploads(operation) + return await operation + + task = asyncio.create_task(run()) + cancelled = False + try: + await asyncio.wait_for(started.wait(), timeout=2) + 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() + + async def test_direct_concurrent_uploads_share_client_limit(self) -> None: + 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 == 8: + limit_reached.set() + try: + await release.wait() + finally: + active -= 1 + + with ( + mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}) 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), + ): + task = asyncio.gather(*(self.data_client.upload_table(mock.Mock()) for _ in range(24))) + try: + await asyncio.wait_for(limit_reached.wait(), timeout=2) + await asyncio.sleep(0) + self.assertEqual(8, save.call_count, "Queued uploads must not start Parquet serialization") + release.set() + await asyncio.wait_for(task, timeout=2) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + self.assertEqual(8, peak) + self.assertEqual(0, active) + self.assertEqual(24, save.call_count) + + @parameterized.expand( + [(str(status), EvoAPIException(status, "transient", None, None)) for status in (408, 429, 500, 502, 503, 504)] + + [ + ("timeout", TimeoutError()), + ("transport", TransportError("connection lost")), + ("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_uploads( + 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([("chunk", "comp=block&"), ("commit", "comp=blocklist")]) + async def test_retry_covers_wrapped_storage_transport_failure(self, _name, query): + response = load_test_data("put_data.json") + original_request = self.transport.request + failed = False + + async def request(**kwargs): + nonlocal failed + if query in kwargs["url"] and not failed: + failed = True + raise TransportError("connection lost") + 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.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.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_retry_exhaustion_is_bounded_and_preserves_errors(self, _name, storage): + cause = EvoAPIException(503, "unavailable", None, None) + error = ChunkedIOError("storage unavailable") if storage else cause + if storage: + error.__cause__ = cause + 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=error) as upload, + ): + with self.assertRaises(RetryError) as caught: + await self.data_client.upload_table( + mock.Mock(), retry=Retry(logger=mock.Mock(), backoff_method=BackoffLinear(0)) + ) + self.assertEqual((error, error, error), caught.exception.exceptions) + self.assertEqual(3, upload.await_count) + save.assert_called_once() + + @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)), + ("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(table, **_kwargs): + started.append(table.column_names) + if len(started) == 2: + all_started.set() + await release.wait() + return {"data": table.column_names[0], "length": table.num_rows} + + with mock.patch.object(self.data_client, "upload_table", side_effect=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.assertEqual("category", result["values"]["data"]) + self.assertEqual("key", result["table"]["data"]) + 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") 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..ea12dfbc6 --- /dev/null +++ b/packages/evo-objects/tests/test_upload_concurrency.py @@ -0,0 +1,178 @@ +# 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 +from unittest import IsolatedAsyncioTestCase +from unittest.mock import patch + +from evo.objects.utils._concurrency import gather_uploads, get_upload_semaphore + + +class TestUploadConcurrency(IsolatedAsyncioTestCase): + async def test_nested_groups_share_limit_preserve_order_and_reset_scope(self): + 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 == 8: + limit_reached.set() + try: + await release.wait() + return index + finally: + active -= 1 + + async def group(start): + return await gather_uploads(*(upload(index) for index in range(start, start + 12))) + + async def run(): + return await gather_uploads(group(0), group(12), group(24)) + + task = asyncio.create_task(run()) + try: + await asyncio.wait_for(limit_reached.wait(), timeout=2) + self.assertEqual(8, 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(8, peak) + self.assertEqual(0, active) + self.assertTrue(all(semaphore is semaphores[0] for semaphore in semaphores)) + self.assertIsNot(get_upload_semaphore(), semaphores[0]) + + 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_uploads(*(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_uploads(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_uploads(*(upload(index) for index in range(start, start + 3))) + + async def run(): + return await gather_uploads(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_scope_is_reset_after_success_and_failure(self): + semaphores = [] + + async def upload(*, fail=False): + semaphores.append(get_upload_semaphore()) + if fail: + raise ValueError("invalid input") + + await gather_uploads(upload()) + self.assertIsNot(get_upload_semaphore(), semaphores[-1]) + with self.assertRaisesRegex(ValueError, "invalid input"): + await gather_uploads(upload(fail=True)) + self.assertIsNot(get_upload_semaphore(), semaphores[-1]) + self.assertIsNot(semaphores[0], semaphores[1]) + + async def test_empty_group_returns_empty_list(self): + self.assertEqual([], await gather_uploads()) + + 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_uploads(*coroutines) + + self.assertTrue(all(task.done() for task in tasks)) + self.assertTrue(all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in coroutines)) + self.assertIsNone(get_upload_semaphore()) + + async def test_semaphore_creation_failure_closes_all_coroutines(self): + async def upload(): + 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_uploads(*coroutines) + + self.assertTrue(all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in coroutines)) + self.assertIsNone(get_upload_semaphore()) diff --git a/packages/evo-objects/tests/typed/test_attributes.py b/packages/evo-objects/tests/typed/test_attributes.py index 1a33c32c9..dc0b6496d 100644 --- a/packages/evo-objects/tests/typed/test_attributes.py +++ b/packages/evo-objects/tests/typed/test_attributes.py @@ -11,20 +11,235 @@ from __future__ import annotations +import asyncio +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, ) +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) + + class TestAttributeTypeInference(TestCase): """Tests for attribute type inference from pandas Series.""" diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index e31edc47a..f886a11e9 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 ( @@ -49,7 +54,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 +230,112 @@ 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), + ) + + async def test_nested_table_uploads_share_a_concurrency_limit(self): + 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): + 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)) + + self.assertGreater(peak, 1) + self.assertLessEqual(peak, 8) + self.assertEqual(0, active) + self.assertGreater(len(uploaded), 50) + 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..9dddca1c6 100644 --- a/packages/evo-objects/tests/typed/test_model.py +++ b/packages/evo-objects/tests/typed/test_model.py @@ -13,17 +13,21 @@ 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.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.utils.table_formats import FLOAT_ARRAY_3, KnownTableFormat from .helpers import MockClient @@ -71,6 +75,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 +128,130 @@ 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_invalid_attribute_metadata_prevents_table_upload(self): + self.frame.attrs["attribute_descriptions"] = {"first": "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() + + @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.""" @@ -111,7 +281,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: From 78e98c7593f5120977d919402f1d4df9108ac982 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 10 Nov 2026 12:22:27 +0000 Subject: [PATCH 2/5] Add upload_concurrency context manager to limit concurrent uploads --- packages/evo-objects/README.md | 16 +++++ .../src/evo/objects/utils/__init__.py | 2 + .../src/evo/objects/utils/_concurrency.py | 35 ++++++++++- .../evo-objects/src/evo/objects/utils/data.py | 4 +- .../evo-objects/tests/test_data_client.py | 28 +++++++-- .../tests/test_upload_concurrency.py | 63 +++++++++++++++++-- .../tests/typed/test_downhole_collection.py | 17 +++-- 7 files changed, 146 insertions(+), 19 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index 07da4feab..1f5e0ebe8 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -113,6 +113,22 @@ for sibling work before returning. Multi-column attribute and combined table/att data concurrently and update their in-memory documents only after all uploads succeed. Separately awaited `collections.add()` calls remain sequential. +To tune CPU and memory usage for table sizes and available resources, use the public +`upload_concurrency` scope. It applies the same limit to nested typed-object operations and direct +uploads. Each slot covers serialization through upload and retries. + +```python +from evo.objects.utils import upload_concurrency + +with upload_concurrency(max_concurrent=4): + dhc = await DownholeCollection.create(context, data) +``` + +The default remains eight. The scope also applies to existing data clients; clients created inside it +retain the configured per-client limit for direct uploads after leaving the scope. Settings are local +to the current context and inherited by child asyncio tasks. Await all uploads before leaving the scope. +Nested scopes restore the previous setting on exit, including when an operation fails. + ## 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/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 index cc0397b58..c096ee4fc 100644 --- a/packages/evo-objects/src/evo/objects/utils/_concurrency.py +++ b/packages/evo-objects/src/evo/objects/utils/_concurrency.py @@ -12,15 +12,46 @@ """Concurrency shared by nested table-serialization operations.""" import asyncio -from collections.abc import Coroutine +from collections.abc import Coroutine, Iterator +from contextlib import contextmanager from contextvars import ContextVar from typing import Any, TypeVar _T = TypeVar("_T") _MAX_CONCURRENT_UPLOADS = 8 +_max_concurrent_uploads: ContextVar[int] = ContextVar("max_concurrent_uploads", default=_MAX_CONCURRENT_UPLOADS) _upload_semaphore: ContextVar[asyncio.Semaphore | None] = ContextVar("object_upload_semaphore", default=None) +def get_max_concurrent_uploads() -> int: + """Return the configured limit for newly created upload semaphores.""" + return _max_concurrent_uploads.get() + + +@contextmanager +def upload_concurrency(max_concurrent: int = _MAX_CONCURRENT_UPLOADS) -> Iterator[None]: + """Limit concurrent table serialization and uploads within this scope. + + Applies to typed-object operations and direct ObjectDataClient uploads, including + clients created internally. Each slot is held from Parquet serialization through + upload and retries. Settings are inherited by child asyncio tasks, not unrelated tasks. + Clients created in this scope retain this limit for subsequent direct uploads. + + :param max_concurrent: Maximum number of concurrent tables. Defaults to eight. + :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") + semaphore = asyncio.Semaphore(max_concurrent) + limit_token = _max_concurrent_uploads.set(max_concurrent) + semaphore_token = _upload_semaphore.set(semaphore) + try: + yield + finally: + _upload_semaphore.reset(semaphore_token) + _max_concurrent_uploads.reset(limit_token) + + def get_upload_semaphore() -> asyncio.Semaphore | None: """Return the serialization's shared limit, if a serialization is active.""" return _upload_semaphore.get() @@ -32,7 +63,7 @@ async def gather_uploads(*coroutines: Coroutine[Any, Any, _T]) -> list[_T]: tasks: list[asyncio.Task[_T]] = [] try: if _upload_semaphore.get() is None: - token = _upload_semaphore.set(asyncio.Semaphore(_MAX_CONCURRENT_UPLOADS)) + token = _upload_semaphore.set(asyncio.Semaphore(get_max_concurrent_uploads())) for coroutine in coroutines: tasks.append(asyncio.create_task(coroutine)) # Unlike gather(), this does not cancel children before our own cleanup runs. diff --git a/packages/evo-objects/src/evo/objects/utils/data.py b/packages/evo-objects/src/evo/objects/utils/data.py index 7482f37b0..232d63e99 100644 --- a/packages/evo-objects/src/evo/objects/utils/data.py +++ b/packages/evo-objects/src/evo/objects/utils/data.py @@ -25,7 +25,7 @@ from ..exceptions import TableFormatError from ..io import _CACHE_SCOPE, ObjectDataDownload, ObjectDataUpload -from ._concurrency import _MAX_CONCURRENT_UPLOADS, gather_uploads, get_upload_semaphore +from ._concurrency import gather_uploads, get_max_concurrent_uploads, get_upload_semaphore 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 @@ -119,7 +119,7 @@ def __init__(self, environment: Environment, connector: APIConnector, cache: ICa self._environment = environment self._connector = connector self._cache = cache - self._upload_semaphore = asyncio.Semaphore(_MAX_CONCURRENT_UPLOADS) + self._upload_semaphore = asyncio.Semaphore(get_max_concurrent_uploads()) @property def cache_location(self) -> Path: diff --git a/packages/evo-objects/tests/test_data_client.py b/packages/evo-objects/tests/test_data_client.py index 31ab48854..9bc14f9e0 100644 --- a/packages/evo-objects/tests/test_data_client.py +++ b/packages/evo-objects/tests/test_data_client.py @@ -13,6 +13,7 @@ import json import sys import threading +from contextlib import nullcontext from unittest import mock from uuid import UUID @@ -29,7 +30,7 @@ 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 gather_uploads from helpers import NoImport, UnloadModule, get_sample_table_and_bytes @@ -490,7 +491,23 @@ async def test_parquet_writer_failure_is_propagated_without_uploading(self) -> N upload.assert_not_awaited() - async def test_direct_concurrent_uploads_share_client_limit(self) -> None: + @parameterized.expand( + [ + ("default", None, False), + ("serial_client", 1, False), + ("lower_client", 3, False), + ("higher_client", 12, False), + ("existing_client_scope", 3, True), + ("existing_client_higher_scope", 12, True), + ] + ) + async def test_direct_concurrent_uploads_share_client_limit(self, _name, configured_limit, scoped) -> None: + limit = configured_limit if configured_limit is not None else 8 + if configured_limit is not None and not scoped: + with upload_concurrency(configured_limit): + self.data_client = ObjectDataClient( + environment=self.environment, connector=self.connector, cache=self.cache + ) release = asyncio.Event() limit_reached = asyncio.Event() active = 0 @@ -500,7 +517,7 @@ async def upload(**_kwargs): nonlocal active, peak active += 1 peak = max(peak, active) - if active == 8: + if active == limit: limit_reached.set() try: await release.wait() @@ -511,19 +528,20 @@ async def upload(**_kwargs): mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}) 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 scoped else nullcontext(), ): task = asyncio.gather(*(self.data_client.upload_table(mock.Mock()) for _ in range(24))) try: await asyncio.wait_for(limit_reached.wait(), timeout=2) await asyncio.sleep(0) - self.assertEqual(8, save.call_count, "Queued uploads must not start Parquet serialization") + self.assertEqual(limit, save.call_count, "Queued uploads must not start Parquet serialization") release.set() await asyncio.wait_for(task, timeout=2) finally: release.set() task.cancel() await asyncio.gather(task, return_exceptions=True) - self.assertEqual(8, peak) + self.assertEqual(limit, peak) self.assertEqual(0, active) self.assertEqual(24, save.call_count) diff --git a/packages/evo-objects/tests/test_upload_concurrency.py b/packages/evo-objects/tests/test_upload_concurrency.py index ea12dfbc6..20a283cb7 100644 --- a/packages/evo-objects/tests/test_upload_concurrency.py +++ b/packages/evo-objects/tests/test_upload_concurrency.py @@ -11,14 +11,20 @@ import asyncio import inspect +from contextlib import nullcontext from unittest import IsolatedAsyncioTestCase from unittest.mock import patch -from evo.objects.utils._concurrency import gather_uploads, get_upload_semaphore +from parameterized import parameterized + +from evo.objects.utils import upload_concurrency +from evo.objects.utils._concurrency import gather_uploads, get_max_concurrent_uploads, get_upload_semaphore class TestUploadConcurrency(IsolatedAsyncioTestCase): - async def test_nested_groups_share_limit_preserve_order_and_reset_scope(self): + @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 8 release = asyncio.Event() limit_reached = asyncio.Event() active = 0 @@ -32,7 +38,7 @@ async def upload(index): async with semaphore: active += 1 peak = max(peak, active) - if active == 8: + if active == limit: limit_reached.set() try: await release.wait() @@ -46,10 +52,11 @@ async def group(start): async def run(): return await gather_uploads(group(0), group(12), group(24)) - task = asyncio.create_task(run()) + 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(8, active) + self.assertEqual(limit, active) release.set() result = await asyncio.wait_for(task, timeout=2) finally: @@ -58,11 +65,55 @@ async def run(): await asyncio.gather(task, return_exceptions=True) self.assertEqual([list(range(0, 12)), list(range(12, 24)), list(range(24, 36))], result) - self.assertEqual(8, peak) + self.assertEqual(limit, peak) self.assertEqual(0, active) self.assertTrue(all(semaphore is semaphores[0] for semaphore in semaphores)) self.assertIsNot(get_upload_semaphore(), 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): + with self.assertRaisesRegex(ValueError, "max_concurrent must be a positive integer"): + with upload_concurrency(value): + self.fail("Invalid configuration was accepted") + self.assertEqual(8, get_max_concurrent_uploads()) + self.assertIsNone(get_upload_semaphore()) + + async def test_nested_configuration_restores_scope_after_failure(self): + with upload_concurrency(3): + outer = get_upload_semaphore() + self.assertEqual(3, get_max_concurrent_uploads()) + with self.assertRaisesRegex(RuntimeError, "failed"): + with upload_concurrency(1): + self.assertEqual(1, get_max_concurrent_uploads()) + self.assertIsNot(outer, get_upload_semaphore()) + raise RuntimeError("failed") + self.assertEqual(3, get_max_concurrent_uploads()) + self.assertIs(outer, get_upload_semaphore()) + self.assertEqual(8, get_max_concurrent_uploads()) + self.assertIsNone(get_upload_semaphore()) + + async def test_configuration_is_isolated_between_tasks(self): + ready = asyncio.Event() + release = asyncio.Event() + + async def configured(): + with upload_concurrency(2): + ready.set() + await release.wait() + return get_max_concurrent_uploads() + + task = asyncio.create_task(configured()) + try: + await asyncio.wait_for(ready.wait(), timeout=2) + self.assertEqual(8, get_max_concurrent_uploads()) + self.assertIsNone(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() diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index f886a11e9..e4baa3ae9 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -45,6 +45,7 @@ IntervalCollection, ) from evo.objects.typed.exceptions import ObjectValidationError +from evo.objects.utils import upload_concurrency from .helpers import MockClient @@ -279,7 +280,9 @@ async def build(builder, payload, skip_sub_models=None): list((await result.collections.get("last").to_dataframe()).columns), ) - async def test_nested_table_uploads_share_a_concurrency_limit(self): + @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)}) @@ -314,7 +317,10 @@ async def upload(upload, cache, transport, **kwargs): context = StaticContext.from_environment( environment=self.environment, connector=self.connector, cache=cache ) - with patch.object(ObjectDataUpload, "upload_from_cache", upload): + 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"]: @@ -330,8 +336,11 @@ async def upload(upload, cache, transport, **kwargs): self.assertEqual(category["values"]["length"], len(table)) self.assertEqual(distance["values"]["length"], len(table)) - self.assertGreater(peak, 1) - self.assertLessEqual(peak, 8) + if limit > 1: + self.assertGreater(peak, 1) + else: + self.assertEqual(1, peak) + self.assertLessEqual(peak, limit) self.assertEqual(0, active) self.assertGreater(len(uploaded), 50) self.assertEqual([c.name for c in data.collections], [c["name"] for c in document["collections"]]) From 6bd63d1c0375271898c47a2b34cebf5d0bdc30aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 10 Nov 2026 15:45:51 +0000 Subject: [PATCH 3/5] Refactor upload concurrency management to support scoped limits across event loops --- packages/evo-objects/README.md | 20 +- .../src/evo/objects/typed/attributes.py | 25 +- .../src/evo/objects/utils/_concurrency.py | 58 +-- .../evo-objects/src/evo/objects/utils/data.py | 87 +++-- .../evo-objects/tests/test_data_client.py | 355 ++++++++++++++++-- .../tests/test_upload_concurrency.py | 59 ++- .../tests/typed/test_attributes.py | 34 ++ .../tests/typed/test_downhole_collection.py | 5 +- .../src/evo/common/utils/retry.py | 6 +- .../tests/common/utils/test_retry.py | 10 + 10 files changed, 548 insertions(+), 111 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index 1f5e0ebe8..ad83717e2 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -106,16 +106,20 @@ Collection coordinate units come only from `DistanceCollection.unit` or `Interva distance or `from` column is ignored. Object creation and replacement upload independent collections, geometry, and attributes concurrently, preserving -their input order. Nested uploads share a limit of eight tables per serialization operation, including Parquet writing; -direct concurrent table uploads share a per-client limit of eight. Transient upload failures receive up to three -attempts with exponential backoff, reusing cached Parquet data. Permanent failures or exhausted retries cancel and wait +their input order. All typed operations and direct table uploads share a default limit of eight active serialization +or upload operations per event loop, including across clients. Transient upload failures receive up to three +attempts with exponential backoff, reusing cached Parquet data. Intermediate failures are logged at WARNING; exhausted retries +re-raise the last original exception, with retry history available through its exception context. Concurrent uploads of +identical data using the same connector and environment share one in-flight upload and its retry policy. +Permanent failures or exhausted retries cancel and wait for sibling work before returning. Multi-column attribute and combined table/attribute mutations upload independent data concurrently and update their in-memory documents only after all uploads succeed. Separately awaited `collections.add()` calls remain sequential. To tune CPU and memory usage for table sizes and available resources, use the public `upload_concurrency` scope. It applies the same limit to nested typed-object operations and direct -uploads. Each slot covers serialization through upload and retries. +uploads. Serialization and each upload attempt acquire a slot separately; backoff sleeps release the slot so other +tables can proceed. ```python from evo.objects.utils import upload_concurrency @@ -124,10 +128,10 @@ with upload_concurrency(max_concurrent=4): dhc = await DownholeCollection.create(context, data) ``` -The default remains eight. The scope also applies to existing data clients; clients created inside it -retain the configured per-client limit for direct uploads after leaving the scope. Settings are local -to the current context and inherited by child asyncio tasks. Await all uploads before leaving the scope. -Nested scopes restore the previous setting on exit, including when an operation fails. +The default remains eight. The scope applies to existing and newly created data clients, without changing their +behavior outside the scope. Settings are local to the current context and inherited by child asyncio tasks, which +retain their scope even after the parent exits it. Each scope has a separate limit; nested scopes restore the previous +setting on exit, including when an operation fails. Clients and scopes can be reused across event loops. ## Contributing diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index 18f19e968..c96aa91f8 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -428,7 +428,9 @@ async def set_attributes(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): self._replace_uploaded_attributes(await self._prepare_attributes(df, fb)) - async def _prepare_attributes(self, df: pd.DataFrame, fb: IFeedback) -> list[dict[str, Any]]: + async def _prepare_attributes( + self, df: pd.DataFrame, fb: IFeedback + ) -> list[tuple[dict[str, Any] | None, dict[str, Any]]]: """Upload replacement values without changing live documents or modification tracking.""" self._attribute_descriptions(df) if len(df.columns) == 0: @@ -438,28 +440,29 @@ async def _prepare_attributes(self, df: pd.DataFrame, fb: IFeedback) -> list[dic 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) -> dict[str, Any]: + async def upload_column(col: Any, col_fb: IFeedback) -> tuple[dict[str, Any] | None, dict[str, Any]]: attribute_df = df[[col]] - attribute = attributes_by_name.get(col) + attribute = attributes_by_name.get(str(col)) if attribute is not None: document = attribute.as_dict() await Attribute._upload_attribute_values( document, attribute_df, attribute.attribute_type, data_client, fb=col_fb ) - return document + return attribute._document, document documents: list[dict[str, Any]] = [] await self._upload_attributes_to_list(documents, attribute_df, self._obj, fb=col_fb) - return documents[0] + return None, documents[0] return await gather_uploads(*(upload_column(col, col_fb) for col, col_fb in zip(df.columns, column_fbs))) - def _replace_uploaded_attributes(self, documents: list[dict[str, Any]]) -> None: - existing = {document["name"]: document for document in self._document} - for index, document in enumerate(documents): - if document["name"] in existing: - original = existing[document["name"]] + def _replace_uploaded_attributes(self, prepared: list[tuple[dict[str, Any] | None, dict[str, Any]]]) -> None: + documents = [] + for original, document in prepared: + if original is not None: + original.clear() original.update(document) - documents[index] = original + document = original + documents.append(document) self._document[:] = documents for document in documents: self._context.mark_modified(document["values"]["data"]) diff --git a/packages/evo-objects/src/evo/objects/utils/_concurrency.py b/packages/evo-objects/src/evo/objects/utils/_concurrency.py index c096ee4fc..f5f19bdae 100644 --- a/packages/evo-objects/src/evo/objects/utils/_concurrency.py +++ b/packages/evo-objects/src/evo/objects/utils/_concurrency.py @@ -9,23 +9,45 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Concurrency shared by nested table-serialization operations.""" +"""Concurrency shared by table serialization and uploads on each event loop.""" import asyncio from collections.abc import Coroutine, Iterator from contextlib import contextmanager from contextvars import ContextVar from typing import Any, TypeVar +from weakref import ReferenceType, WeakKeyDictionary, ref _T = TypeVar("_T") _MAX_CONCURRENT_UPLOADS = 8 -_max_concurrent_uploads: ContextVar[int] = ContextVar("max_concurrent_uploads", default=_MAX_CONCURRENT_UPLOADS) -_upload_semaphore: ContextVar[asyncio.Semaphore | None] = ContextVar("object_upload_semaphore", default=None) + + +class _UploadScope: + def __init__(self, limit: int) -> None: + self.limit = limit + # A contended semaphore retains its loop; weak values avoid retaining closed loops. + self.semaphores: WeakKeyDictionary[asyncio.AbstractEventLoop, ReferenceType[asyncio.Semaphore]] = ( + WeakKeyDictionary() + ) + + def get_semaphore(self) -> asyncio.Semaphore: + loop = asyncio.get_running_loop() + 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(_MAX_CONCURRENT_UPLOADS) +) def get_max_concurrent_uploads() -> int: - """Return the configured limit for newly created upload semaphores.""" - return _max_concurrent_uploads.get() + """Return the active upload scope's limit.""" + return _upload_scope.get().limit @contextmanager @@ -33,37 +55,32 @@ def upload_concurrency(max_concurrent: int = _MAX_CONCURRENT_UPLOADS) -> Iterato """Limit concurrent table serialization and uploads within this scope. Applies to typed-object operations and direct ObjectDataClient uploads, including - clients created internally. Each slot is held from Parquet serialization through - upload and retries. Settings are inherited by child asyncio tasks, not unrelated tasks. - Clients created in this scope retain this limit for subsequent direct uploads. + clients created internally. Serialization and each upload attempt acquire a slot; + backoff sleeps do not. 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 tables. Defaults to eight. :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") - semaphore = asyncio.Semaphore(max_concurrent) - limit_token = _max_concurrent_uploads.set(max_concurrent) - semaphore_token = _upload_semaphore.set(semaphore) + token = _upload_scope.set(_UploadScope(max_concurrent)) try: yield finally: - _upload_semaphore.reset(semaphore_token) - _max_concurrent_uploads.reset(limit_token) + _upload_scope.reset(token) -def get_upload_semaphore() -> asyncio.Semaphore | None: - """Return the serialization's shared limit, if a serialization is active.""" - return _upload_semaphore.get() +def get_upload_semaphore() -> asyncio.Semaphore: + """Return the active upload scope's shared semaphore for the running loop.""" + return _upload_scope.get().get_semaphore() async def gather_uploads(*coroutines: Coroutine[Any, Any, _T]) -> list[_T]: """Collect ordered results, cancelling and draining siblings on failure or cancellation.""" - token = None tasks: list[asyncio.Task[_T]] = [] try: - if _upload_semaphore.get() is None: - token = _upload_semaphore.set(asyncio.Semaphore(get_max_concurrent_uploads())) for coroutine in coroutines: tasks.append(asyncio.create_task(coroutine)) # Unlike gather(), this does not cancel children before our own cleanup runs. @@ -77,6 +94,3 @@ async def gather_uploads(*coroutines: Coroutine[Any, Any, _T]) -> list[_T]: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) raise - finally: - if token is not None: - _upload_semaphore.reset(token) diff --git a/packages/evo-objects/src/evo/objects/utils/data.py b/packages/evo-objects/src/evo/objects/utils/data.py index 232d63e99..8e3c509cf 100644 --- a/packages/evo-objects/src/evo/objects/utils/data.py +++ b/packages/evo-objects/src/evo/objects/utils/data.py @@ -14,18 +14,19 @@ 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 EvoAPIException, StorageFileNotFoundError, TransportError, error_context +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 gather_uploads, get_max_concurrent_uploads, get_upload_semaphore +from ._concurrency import gather_uploads, get_upload_semaphore 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 @@ -50,6 +51,9 @@ _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}) +_inflight_uploads: WeakKeyDictionary[ + asyncio.AbstractEventLoop, dict[tuple[int, Environment, str], asyncio.Future[None]] +] = WeakKeyDictionary() def _is_transient_upload_error(error: Exception) -> bool: @@ -119,7 +123,6 @@ def __init__(self, environment: Environment, connector: APIConnector, cache: ICa self._environment = environment self._connector = connector self._cache = cache - self._upload_semaphore = asyncio.Semaphore(get_max_concurrent_uploads()) @property def cache_location(self) -> Path: @@ -218,7 +221,9 @@ async def upload_table( 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 URL or storage failures. Defaults to three attempts with - exponential backoff. Retries reuse the cached table; exhausted attempts raise RetryError. + exponential backoff. Retries reuse the cached table; exhausted attempts re-raise the last original + exception. Concurrent uploads of identical data with the same connector and environment share one + upload and its retry policy. :return: A description of the uploaded data. @@ -227,7 +232,8 @@ async def upload_table( """ if retry is None: retry = Retry(logger=logger) - async with get_upload_semaphore() or self._upload_semaphore: + semaphore = get_upload_semaphore() + async with semaphore: # Compression and hashing must not block other uploads on the event loop. writer = asyncio.create_task(asyncio.to_thread(self.save_table, table, table_format)) try: @@ -236,27 +242,60 @@ async def upload_table( # A running worker cannot be cancelled; finish its cache write before allowing a retry. await asyncio.gather(writer, return_exceptions=True) raise - 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): - async for attempt in retry: - try: + + loop = asyncio.get_running_loop() + pending = _inflight_uploads.setdefault(loop, {}) + key = (id(self._connector), self._environment, table_info["data"]) + if (existing := pending.get(key)) is not None: + await asyncio.shield(existing) + fb.progress(1) + return table_info + + pending[key] = completion = loop.create_future() + try: + await self._upload_cached_table(table_info["data"], fb, retry, semaphore) + except asyncio.CancelledError: + completion.cancel() + raise + except BaseException as exc: + completion.set_exception(exc) + # The owner propagates the failure even when no duplicate is waiting. + completion.exception() + raise + else: + completion.set_result(None) + finally: + del pending[key] + if not pending: + del _inflight_uploads[loop] + return table_info + + async def _upload_cached_table(self, name: str, fb: IFeedback, retry: Retry, semaphore: asyncio.Semaphore) -> None: + upload = ObjectDataUpload(connector=self._connector, environment=self._environment, name=name) + with error_context(f"While uploading table data '{name}'", logger=logger): + async for attempt in retry: + try: + async with semaphore: 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) - except ( - EvoAPIException, - TransportError, - asyncio.TimeoutError, - TimeoutError, - ConnectionError, - ChunkedIOError, - ) as exc: - if not _is_transient_upload_error(exc): - raise + except DataExistsError: + logger.debug(f"Data not uploaded because data already exists (label: {name})") + fb.progress(1) + except ( + EvoAPIException, + TransportError, + asyncio.TimeoutError, + TimeoutError, + ConnectionError, + ChunkedIOError, + ) as exc: + if not _is_transient_upload_error(exc): + raise + try: attempt.set_exception(exc) - fb.progress(0) - return table_info + except RetryError: + # Keep the original type, instance and storage cause; retry history is the context. + raise exc + fb.progress(0) async def upload_category_table(self, table: pa.Table, fb: IFeedback = NoFeedback) -> CategoryInfo: """Upload a pyarrow table with category data to the geoscience object service, returning the reference to diff --git a/packages/evo-objects/tests/test_data_client.py b/packages/evo-objects/tests/test_data_client.py index 9bc14f9e0..75047c379 100644 --- a/packages/evo-objects/tests/test_data_client.py +++ b/packages/evo-objects/tests/test_data_client.py @@ -11,11 +11,13 @@ import asyncio import json +import logging import sys import threading from contextlib import nullcontext +from dataclasses import replace from unittest import mock -from uuid import UUID +from uuid import UUID, uuid4 import pandas as pd import pyarrow as pa @@ -219,7 +221,7 @@ async def test_upload_failure_preserves_exception_and_identifies_data( 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(RetryError if upload_type == "table_transient" else error_type) as cm, + self.assertRaises(error_type) as cm, ): destination.upload_file.side_effect = error if upload_type == "referenced": @@ -231,9 +233,9 @@ async def test_upload_failure_preserves_exception_and_identifies_data( table, retry=Retry(logger=logger, backoff_method=BackoffLinear(0)) ) if upload_type == "table_transient": - self.assertEqual((error, error, error), cm.exception.exceptions) - else: - self.assertIs(error, cm.exception) + self.assertIsInstance(cm.exception.__context__, RetryError) + self.assertEqual((error, error, error), cm.exception.__context__.exceptions) + self.assertIs(error, cm.exception) context = ( "\n".join(getattr(cm.exception, "__notes__", [])) if sys.version_info >= (3, 11) @@ -494,20 +496,21 @@ async def test_parquet_writer_failure_is_propagated_without_uploading(self) -> N @parameterized.expand( [ ("default", None, False), - ("serial_client", 1, False), - ("lower_client", 3, False), - ("higher_client", 12, False), - ("existing_client_scope", 3, True), - ("existing_client_higher_scope", 12, True), + ("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_client_limit(self, _name, configured_limit, scoped) -> None: + 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 8 - if configured_limit is not None and not scoped: - with upload_concurrency(configured_limit): - self.data_client = ObjectDataClient( - environment=self.environment, connector=self.connector, cache=self.cache - ) + 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 @@ -525,16 +528,16 @@ async def upload(**_kwargs): active -= 1 with ( - mock.patch.object(self.data_client, "save_table", return_value={"data": "saved"}) as save, + 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 scoped else nullcontext(), + upload_concurrency(configured_limit) if configured_limit is not None else nullcontext(), ): - task = asyncio.gather(*(self.data_client.upload_table(mock.Mock()) for _ in range(24))) + 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, save.call_count, "Queued uploads must not start Parquet serialization") + self.assertEqual(limit, active) release.set() await asyncio.wait_for(task, timeout=2) finally: @@ -545,6 +548,91 @@ async def upload(**_kwargs): 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(run()) + self.assertEqual([str(i) for i in range(16)], [info["data"] for info in result]) + + @parameterized.expand([("default", None), ("scope", 3)]) + async def test_concurrent_category_operations_share_limit(self, _name, configured_limit): + limit = configured_limit if configured_limit is not None else 8 + 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)] + [ @@ -633,24 +721,233 @@ async def request(**kwargs): self.assertEqual(2, request_url.await_count) save.assert_called_once() - @parameterized.expand([("url", False), ("storage", True)]) - async def test_retry_exhaustion_is_bounded_and_preserves_errors(self, _name, storage): - cause = EvoAPIException(503, "unavailable", None, None) - error = ChunkedIOError("storage unavailable") if storage else cause - if storage: - error.__cause__ = cause + @parameterized.expand( + [ + ("url", EvoAPIException), + ("storage", ChunkedIOError), + ("connection", ConnectionError), + ("transport", TransportError), + ] + ) + 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}") + if error_type is ChunkedIOError: + error.__cause__ = cause + errors.append(error) + original_cause = errors[-1].__cause__ 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=error) as upload, + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=errors) as upload, ): - with self.assertRaises(RetryError) as caught: + 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.assertEqual((error, error, error), caught.exception.exceptions) + self.assertIs(errors[-1], caught.exception) + self.assertIs(original_cause, caught.exception.__cause__) + self.assertIsInstance(caught.exception.__context__, RetryError) + self.assertEqual(tuple(errors), caught.exception.__context__.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([("success", False), ("failure", True)]) + async def test_identical_tables_share_inflight_upload_across_clients(self, _name, fail): + started = asyncio.Event() + joined = asyncio.Event() + release = asyncio.Event() + error = RuntimeError("upload failed") + shield = asyncio.shield + feedback = mock.Mock(spec=IFeedback) + other_client = ObjectDataClient(environment=self.environment, connector=self.connector, cache=self.cache) + table = pa.table({"value": [1.0]}) + + def watch_join(awaitable): + if not isinstance(awaitable, asyncio.Task): + joined.set() + return shield(awaitable) + + async def upload(**_kwargs): + started.set() + await release.wait() + if fail: + raise error + + with ( + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload) as upload_mock, + mock.patch("evo.objects.utils.data.asyncio.shield", side_effect=watch_join), + ): + owner = asyncio.create_task(self.data_client.upload_table(table)) + tasks = [owner] + try: + await asyncio.wait_for(started.wait(), timeout=2) + tasks.append(asyncio.create_task(other_client.upload_table(table, fb=feedback))) + await asyncio.wait_for(joined.wait(), timeout=2) + upload_mock.assert_awaited_once() + self.assertFalse(tasks[1].done()) + release.set() + results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=2) + if fail: + self.assertTrue(all(result is error for result in results)) + feedback.progress.assert_not_called() + else: + self.assertEqual(results[0], results[1]) + self.assertIsNot(results[0], results[1]) + feedback.progress.assert_called_with(1) + + # Completion (including failure) must remove the in-flight entry. + upload_mock.side_effect = None + await self.data_client.upload_table(table) + self.assertEqual(2, upload_mock.await_count) + finally: + release.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @parameterized.expand([("waiter", False), ("owner", True)]) + async def test_duplicate_upload_cancellation_cleans_up_without_orphaned_work(self, _name, cancel_owner): + started = asyncio.Event() + joined = asyncio.Event() + release = asyncio.Event() + shield = asyncio.shield + + def watch_join(awaitable): + if not isinstance(awaitable, asyncio.Task): + joined.set() + return shield(awaitable) + + async def upload(**_kwargs): + started.set() + await release.wait() + + 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 upload_mock, + mock.patch("evo.objects.utils.data.asyncio.shield", side_effect=watch_join), + ): + owner = asyncio.create_task(self.data_client.upload_table(mock.Mock())) + tasks = [owner] + try: + await asyncio.wait_for(started.wait(), timeout=2) + waiter = asyncio.create_task(self.data_client.upload_table(mock.Mock())) + tasks.append(waiter) + await asyncio.wait_for(joined.wait(), timeout=2) + cancelled = owner if cancel_owner else waiter + cancelled.cancel() + with self.assertRaises(asyncio.CancelledError): + await cancelled + if cancel_owner: + with self.assertRaises(asyncio.CancelledError): + await asyncio.wait_for(waiter, timeout=2) + else: + self.assertFalse(owner.done()) + release.set() + self.assertEqual({"data": "saved"}, await asyncio.wait_for(owner, timeout=2)) + upload_mock.assert_awaited_once() + release.set() + await self.data_client.upload_table(mock.Mock()) + self.assertEqual(2, upload_mock.await_count) + finally: + release.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @parameterized.expand([("hub_url",), ("org_id",), ("workspace_id",), ("connector",)]) + async def test_inflight_uploads_do_not_share_across_destinations_or_connectors(self, field): + environment = self.environment + connector = self.connector + if field == "connector": + connector = mock.Mock(wraps=self.connector) + else: + environment = replace(environment, **{field: "https://other.example" if field == "hub_url" else uuid4()}) + other_client = ObjectDataClient(environment=environment, connector=connector, cache=self.cache) + started = asyncio.Event() + both_started = asyncio.Event() + release = asyncio.Event() + count = 0 + + async def upload(**_kwargs): + nonlocal count + count += 1 + started.set() + if count == 2: + both_started.set() + await release.wait() + + with ( + mock.patch.object(ObjectDataClient, "save_table", return_value={"data": "saved"}), + mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload), + ): + tasks = [asyncio.create_task(self.data_client.upload_table(mock.Mock()))] + try: + await asyncio.wait_for(started.wait(), timeout=2) + tasks.append(asyncio.create_task(other_client.upload_table(mock.Mock()))) + await asyncio.wait_for(both_started.wait(), timeout=2) + release.set() + await asyncio.wait_for(asyncio.gather(*tasks), timeout=2) + finally: + release.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + self.assertEqual(2, count) + @parameterized.expand( [ ("bad_request", EvoAPIException(400, "bad request", None, None)), diff --git a/packages/evo-objects/tests/test_upload_concurrency.py b/packages/evo-objects/tests/test_upload_concurrency.py index 20a283cb7..796237530 100644 --- a/packages/evo-objects/tests/test_upload_concurrency.py +++ b/packages/evo-objects/tests/test_upload_concurrency.py @@ -12,7 +12,7 @@ import asyncio import inspect from contextlib import nullcontext -from unittest import IsolatedAsyncioTestCase +from unittest import IsolatedAsyncioTestCase, TestCase from unittest.mock import patch from parameterized import parameterized @@ -25,6 +25,7 @@ class TestUploadConcurrency(IsolatedAsyncioTestCase): @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 8 + default = get_upload_semaphore() release = asyncio.Event() limit_reached = asyncio.Event() active = 0 @@ -50,7 +51,7 @@ async def group(start): return await gather_uploads(*(upload(index) for index in range(start, start + 12))) async def run(): - return await gather_uploads(group(0), group(12), group(24)) + 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()) @@ -68,17 +69,23 @@ async def run(): self.assertEqual(limit, peak) self.assertEqual(0, active) self.assertTrue(all(semaphore is semaphores[0] for semaphore in semaphores)) - self.assertIsNot(get_upload_semaphore(), semaphores[0]) + 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(8, get_max_concurrent_uploads()) - self.assertIsNone(get_upload_semaphore()) + 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, get_max_concurrent_uploads()) @@ -90,9 +97,10 @@ async def test_nested_configuration_restores_scope_after_failure(self): self.assertEqual(3, get_max_concurrent_uploads()) self.assertIs(outer, get_upload_semaphore()) self.assertEqual(8, get_max_concurrent_uploads()) - self.assertIsNone(get_upload_semaphore()) + 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() @@ -106,7 +114,7 @@ async def configured(): try: await asyncio.wait_for(ready.wait(), timeout=2) self.assertEqual(8, get_max_concurrent_uploads()) - self.assertIsNone(get_upload_semaphore()) + self.assertIs(default, get_upload_semaphore()) release.set() self.assertEqual(2, await asyncio.wait_for(task, timeout=2)) finally: @@ -175,7 +183,7 @@ async def run(): task.cancel() await asyncio.gather(task, return_exceptions=True) - async def test_scope_is_reset_after_success_and_failure(self): + async def test_groups_reuse_default_limit_after_success_and_failure(self): semaphores = [] async def upload(*, fail=False): @@ -184,11 +192,11 @@ async def upload(*, fail=False): raise ValueError("invalid input") await gather_uploads(upload()) - self.assertIsNot(get_upload_semaphore(), semaphores[-1]) + self.assertIs(get_upload_semaphore(), semaphores[-1]) with self.assertRaisesRegex(ValueError, "invalid input"): await gather_uploads(upload(fail=True)) - self.assertIsNot(get_upload_semaphore(), semaphores[-1]) - self.assertIsNot(semaphores[0], semaphores[1]) + 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_uploads()) @@ -214,11 +222,12 @@ def schedule(coroutine): self.assertTrue(all(task.done() for task in tasks)) self.assertTrue(all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in coroutines)) - self.assertIsNone(get_upload_semaphore()) + self.assertEqual(8, get_max_concurrent_uploads()) - async def test_semaphore_creation_failure_closes_all_coroutines(self): + async def test_semaphore_creation_failure_drains_all_coroutines(self): async def upload(): - return 1 + async with get_upload_semaphore(): + return 1 coroutines = [upload(), upload()] with patch("evo.objects.utils._concurrency.asyncio.Semaphore", side_effect=RuntimeError("no semaphore")): @@ -226,4 +235,26 @@ async def upload(): await gather_uploads(*coroutines) self.assertTrue(all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in coroutines)) - self.assertIsNone(get_upload_semaphore()) + self.assertEqual(8, get_max_concurrent_uploads()) + + +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(run()) + asyncio.run(run()) + + 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 dc0b6496d..91958a96f 100644 --- a/packages/evo-objects/tests/typed/test_attributes.py +++ b/packages/evo-objects/tests/typed/test_attributes.py @@ -239,6 +239,40 @@ async def test_empty_mutations_do_not_upload(self): get_client.assert_not_called() self.assertEqual({"previous-change"}, self.context.data_modified) + @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): """Tests for attribute type inference from pandas Series.""" diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index e4baa3ae9..e57cb665b 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -46,6 +46,7 @@ ) 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 @@ -342,7 +343,9 @@ async def upload(upload, cache, transport, **kwargs): self.assertEqual(1, peak) self.assertLessEqual(peak, limit) self.assertEqual(0, active) - self.assertGreater(len(uploaded), 50) + 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): 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): From b507b11bbbf6c236843c1964d6b1f5bfe366d1dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 10 Nov 2026 16:24:19 +0000 Subject: [PATCH 4/5] Enhance upload concurrency management with new sub-model handling and improved attribute uploads --- packages/evo-objects/README.md | 57 ++- .../src/evo/objects/typed/_data.py | 22 +- .../src/evo/objects/typed/_model.py | 4 + .../src/evo/objects/typed/attributes.py | 110 ++--- .../src/evo/objects/utils/_concurrency.py | 53 ++- .../evo-objects/src/evo/objects/utils/data.py | 134 +++--- .../evo-objects/tests/test_data_client.py | 411 ++++++++++-------- .../tests/test_upload_concurrency.py | 55 ++- .../evo-objects/tests/typed/test_model.py | 25 ++ .../evo-sdk-common/src/evo/aio/_helpers.py | 4 +- .../evo-sdk-common/src/evo/aio/transport.py | 3 +- 11 files changed, 513 insertions(+), 365 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index ad83717e2..1240694de 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -105,33 +105,50 @@ 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. -Object creation and replacement upload independent collections, geometry, and attributes concurrently, preserving -their input order. All typed operations and direct table uploads share a default limit of eight active serialization -or upload operations per event loop, including across clients. Transient upload failures receive up to three -attempts with exponential backoff, reusing cached Parquet data. Intermediate failures are logged at WARNING; exhausted retries -re-raise the last original exception, with retry history available through its exception context. Concurrent uploads of -identical data using the same connector and environment share one in-flight upload and its retry policy. -Permanent failures or exhausted retries cancel and wait -for sibling work before returning. Multi-column attribute and combined table/attribute mutations upload independent -data concurrently and update their in-memory documents only after all uploads succeed. Separately awaited -`collections.add()` calls remain sequential. - -To tune CPU and memory usage for table sizes and available resources, use the public -`upload_concurrency` scope. It applies the same limit to nested typed-object operations and direct -uploads. Serialization and each upload attempt acquire a slot separately; backoff sleeps release the slot so other -tables can proceed. +### Upload concurrency + +Typed objects upload independent tables and attributes concurrently, preserving input order and applying +multi-column and combined table/attribute mutations only after every upload succeeds. Typed operations and direct +`ObjectDataClient` uploads share a default concurrency of **4**, matching downloads. Use `upload_concurrency` +to change the limit for a block, including its child tasks: ```python from evo.objects.utils import upload_concurrency -with upload_concurrency(max_concurrent=4): +with upload_concurrency(max_concurrent=8): dhc = await DownholeCollection.create(context, data) ``` -The default remains eight. The scope applies to existing and newly created data clients, without changing their -behavior outside the scope. Settings are local to the current context and inherited by child asyncio tasks, which -retain their scope even after the parent exits it. Each scope has a separate limit; nested scopes restore the previous -setting on exit, including when an operation fails. Clients and scopes can be reused across event loops. +The limit bounds concurrent Arrow conversion, category preparation, Parquet serialization, and upload attempts, +not total memory usage or the size of caller-owned data. Preparation runs in worker threads, and temporary Arrow +tables are released before network uploads. Nested scopes restore the previous setting on exit. + +Uploads retry transient HTTP statuses (408/429/500/502/503/504) and raw connection/timeouts up to three attempts +with exponential backoff, reusing cached data. Exhaustion raises `RetryError` with the attempt history. +Use the optional `retry=Retry(...)` argument on any `ObjectDataClient` table, DataFrame, or category upload method +to customize this policy. Already-exhausted `TransportError` failures are not retried again. Storage has its own +retry policy; exhausted storage retries are retried only when all their errors are transient, so storage request +counts can exceed the table-attempt limit. Failed operations cancel and drain their own siblings, not independent uploads. + +### HTTP connections + +HTTP connections are configured separately: `AioTransport(num_pools=...)` sets the total connection limit +(default **4**) across requests sharing that transport. `upload_concurrency` does not resize it. For example, +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 diff --git a/packages/evo-objects/src/evo/objects/typed/_data.py b/packages/evo-objects/src/evo/objects/typed/_data.py index 28504d03c..03b613b93 100644 --- a/packages/evo-objects/src/evo/objects/typed/_data.py +++ b/packages/evo-objects/src/evo/objects/typed/_data.py @@ -50,7 +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._apply_uploaded_table(df, await self._data_to_schema(df, self._obj, fb=fb)) + self.commit(df, await self.prepare(df, fb)) + + async def prepare(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 commit(self, df: pd.DataFrame, document: dict[str, Any]) -> None: + """Apply a successfully prepared table, including subclass-specific derived metadata.""" + self._apply_uploaded_table(df, document) def _apply_uploaded_table(self, df: pd.DataFrame, document: dict[str, Any]) -> None: """Commit uploaded values and any subclass-specific derived metadata.""" @@ -122,15 +130,15 @@ async def from_dataframe(self, df: pd.DataFrame, fb: IFeedback = NoFeedback) -> table_df, attr_df = self._split_dataframe(df, self._table.data_columns) if attr_df is not None: - # Validate local metadata before either branch starts remote work. - self.attributes._attribute_descriptions(attr_df) table_fb, attr_fb = split_feedback(fb, [1.0, float(len(attr_df.columns))]) + # prepare validates metadata before either branch starts remote work. + attribute_upload = self.attributes.prepare(attr_df, attr_fb) table_document, attribute_documents = await gather_uploads( - self._table._data_to_schema(table_df, self._obj, fb=table_fb), - self.attributes._prepare_attributes(attr_df, attr_fb), + self._table.prepare(table_df, table_fb), + attribute_upload, ) - self._table._apply_uploaded_table(table_df, table_document) - self.attributes._replace_uploaded_attributes(attribute_documents) + self._table.commit(table_df, table_document) + self.attributes.commit(attribute_documents) else: await self._table.from_dataframe(table_df, fb=fb) self.attributes.clear() diff --git a/packages/evo-objects/src/evo/objects/typed/_model.py b/packages/evo-objects/src/evo/objects/typed/_model.py index 854653f3e..56a5a783d 100644 --- a/packages/evo-objects/src/evo/objects/typed/_model.py +++ b/packages/evo-objects/src/evo/objects/typed/_model.py @@ -208,6 +208,10 @@ async def build_from_data(self, data: Any, skip_sub_models: set[str] | None = No return self.document + async def set_sub_model_value(self, name: str, data: Any) -> None: + """Convert and assign a single sub-model.""" + await self.set_sub_model_values([(name, data)]) + 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. diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index c96aa91f8..19c713bfd 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 collections.abc import Coroutine from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any, Protocol, runtime_checkable from uuid import UUID @@ -38,6 +39,8 @@ from .exceptions import DataLoaderError, ObjectValidationError if TYPE_CHECKING: + from evo.objects.utils import ObjectDataClient + from .block_model_ref import BlockModel __all__ = [ @@ -257,6 +260,27 @@ def __repr__(self) -> str: return f"PendingAttribute(name={self._name!r}, exists=False)" +async def _build_new_attribute_document( + col: Any, + series: pd.Series, + description: dict[str, Any] | None, + data_client: ObjectDataClient, + fb: IFeedback, +) -> dict[str, Any]: + attribute_type = _infer_attribute_type_from_series(series) + document: 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") + await Attribute._upload_attribute_values(document, pd.DataFrame({col: series}), attribute_type, data_client, fb=fb) + if description is not None: + document["attribute_description"] = description + return document + + class Attributes(SchemaList[Attribute]): """A collection of Geoscience Object Attributes @@ -296,18 +320,15 @@ 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 def _attribute_descriptions(df: pd.DataFrame) -> dict[Any, dict[str, Any]]: + """Validate descriptions and convert them to schema dictionaries once.""" descriptions = df.attrs.get("attribute_descriptions", {}) result = {} for col in df.columns: @@ -319,17 +340,15 @@ def _attribute_descriptions(df: pd.DataFrame) -> dict[Any, dict[str, Any]]: return result @staticmethod - async def _upload_attributes_to_list( - attributes_list: list[dict[str, Any]], + 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. @@ -337,33 +356,15 @@ async def _upload_attributes_to_list( descriptions = Attributes._attribute_descriptions(df) if len(df.columns) == 0: fb.progress(1) - return + return [] data_client = get_data_client(context) column_fbs = split_feedback(fb, [1.0] * len(df.columns)) - async def upload_column(col: Any, col_fb: IFeedback) -> dict[str, Any]: - 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, fb=col_fb) - - if description := descriptions.get(col): - attr_doc["attribute_description"] = description - return attr_doc - - # Columns are independent, so upload them concurrently while preserving column order in the result. - attributes_list.extend( - await gather_uploads(*(upload_column(col, col_fb) for col, col_fb in zip(df.columns, column_fbs))) + return await gather_uploads( + *( + _build_new_attribute_document(col, df[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: @@ -400,12 +401,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. @@ -413,8 +411,7 @@ 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. """ - documents: list[dict[str, Any]] = [] - await self._upload_attributes_to_list(documents, df, self._obj, 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"]) @@ -426,13 +423,23 @@ async def set_attributes(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): :param fb: Optional feedback object to report upload progress. """ - self._replace_uploaded_attributes(await self._prepare_attributes(df, fb)) + self.commit(await self.prepare(df, fb)) + + def prepare( + self, df: pd.DataFrame, fb: IFeedback = NoFeedback + ) -> Coroutine[Any, Any, list[tuple[dict[str, Any] | None, dict[str, Any]]]]: + """Validate metadata now and return an upload coroutine, without changing live attributes. + + Await the returned coroutine and pass its result to commit only after all + related uploads succeed. + """ + descriptions = self._attribute_descriptions(df) + return self._prepare_attributes(df, descriptions, fb) async def _prepare_attributes( - self, df: pd.DataFrame, fb: IFeedback + self, df: pd.DataFrame, descriptions: dict[Any, dict[str, Any]], fb: IFeedback ) -> list[tuple[dict[str, Any] | None, dict[str, Any]]]: """Upload replacement values without changing live documents or modification tracking.""" - self._attribute_descriptions(df) if len(df.columns) == 0: fb.progress(1) return [] @@ -441,21 +448,20 @@ async def _prepare_attributes( 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_df = df[[col]] attribute = attributes_by_name.get(str(col)) if attribute is not None: document = attribute.as_dict() await Attribute._upload_attribute_values( - document, attribute_df, attribute.attribute_type, data_client, fb=col_fb + document, df[[col]], attribute.attribute_type, data_client, fb=col_fb ) return attribute._document, document - documents: list[dict[str, Any]] = [] - await self._upload_attributes_to_list(documents, attribute_df, self._obj, fb=col_fb) - return None, documents[0] + document = await _build_new_attribute_document(col, df[col], descriptions.get(col), data_client, col_fb) + return None, document return await gather_uploads(*(upload_column(col, col_fb) for col, col_fb in zip(df.columns, column_fbs))) - def _replace_uploaded_attributes(self, prepared: list[tuple[dict[str, Any] | None, dict[str, Any]]]) -> None: + def commit(self, prepared: list[tuple[dict[str, Any] | None, dict[str, Any]]]) -> None: + """Replace live attributes with a successfully prepared batch, preserving existing wrappers.""" documents = [] for original, document in prepared: if original is not None: diff --git a/packages/evo-objects/src/evo/objects/utils/_concurrency.py b/packages/evo-objects/src/evo/objects/utils/_concurrency.py index f5f19bdae..56788734f 100644 --- a/packages/evo-objects/src/evo/objects/utils/_concurrency.py +++ b/packages/evo-objects/src/evo/objects/utils/_concurrency.py @@ -12,20 +12,29 @@ """Concurrency shared by table serialization and uploads on each event loop.""" import asyncio -from collections.abc import Coroutine, Iterator +from collections.abc import Callable, Coroutine, Iterator from contextlib import contextmanager from contextvars import ContextVar -from typing import Any, TypeVar +from typing import Any, ParamSpec, TypeVar from weakref import ReferenceType, WeakKeyDictionary, ref +from evo import logging + _T = TypeVar("_T") -_MAX_CONCURRENT_UPLOADS = 8 +_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 - # A contended semaphore retains its loop; weak values avoid retaining closed loops. self.semaphores: WeakKeyDictionary[asyncio.AbstractEventLoop, ReferenceType[asyncio.Semaphore]] = ( WeakKeyDictionary() ) @@ -41,26 +50,21 @@ def get_semaphore(self) -> asyncio.Semaphore: _upload_scope: ContextVar[_UploadScope] = ContextVar( - "object_upload_scope", default=_UploadScope(_MAX_CONCURRENT_UPLOADS) + "object_upload_scope", default=_UploadScope(_DEFAULT_MAX_CONCURRENCY) ) -def get_max_concurrent_uploads() -> int: - """Return the active upload scope's limit.""" - return _upload_scope.get().limit - - @contextmanager -def upload_concurrency(max_concurrent: int = _MAX_CONCURRENT_UPLOADS) -> Iterator[None]: - """Limit concurrent table serialization and uploads within this scope. +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. Serialization and each upload attempt acquire a slot; + clients created internally. Conversion/serialization and each upload attempt acquire a slot; backoff sleeps do not. 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 tables. Defaults to eight. + :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: @@ -73,10 +77,22 @@ def upload_concurrency(max_concurrent: int = _MAX_CONCURRENT_UPLOADS) -> Iterato def get_upload_semaphore() -> asyncio.Semaphore: - """Return the active upload scope's shared semaphore for the running loop.""" + """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: + # Running threads cannot be cancelled; finish their cache writes before returning. + await asyncio.gather(worker, return_exceptions=True) + raise + + async def gather_uploads(*coroutines: Coroutine[Any, Any, _T]) -> list[_T]: """Collect ordered results, cancelling and draining siblings on failure or cancellation.""" tasks: list[asyncio.Task[_T]] = [] @@ -87,10 +103,13 @@ async def gather_uploads(*coroutines: Coroutine[Any, Any, _T]) -> list[_T]: for completed in asyncio.as_completed(tasks): await completed return [task.result() for task in tasks] - except BaseException: + except BaseException as error: for coroutine in coroutines[len(tasks) :]: coroutine.close() for task in tasks: task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) + 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 8e3c509cf..767688b84 100644 --- a/packages/evo-objects/src/evo/objects/utils/data.py +++ b/packages/evo-objects/src/evo/objects/utils/data.py @@ -14,7 +14,6 @@ from pathlib import Path from typing import Any, Iterable, cast from uuid import UUID -from weakref import WeakKeyDictionary import numpy as np @@ -26,7 +25,7 @@ from ..exceptions import TableFormatError from ..io import _CACHE_SCOPE, ObjectDataDownload, ObjectDataUpload -from ._concurrency import gather_uploads, get_upload_semaphore +from ._concurrency import _DEFAULT_MAX_CONCURRENCY, gather_uploads, 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 @@ -49,19 +48,20 @@ 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}) -_inflight_uploads: WeakKeyDictionary[ - asyncio.AbstractEventLoop, dict[tuple[int, Environment, str], asyncio.Future[None]] -] = WeakKeyDictionary() 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, (TransportError, asyncio.TimeoutError, TimeoutError, ConnectionError)) + return isinstance(error, (asyncio.TimeoutError, TimeoutError, ConnectionError)) def _iter_refs(target: Any, _key: str | None = None) -> Iterator[str]: @@ -220,57 +220,24 @@ 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 URL or storage failures. Defaults to three attempts with - exponential backoff. Retries reuse the cached table; exhausted attempts re-raise the last original - exception. Concurrent uploads of identical data with the same connector and environment share one - upload and its retry policy. + :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. Exhaustion raises RetryError with the attempt history. TransportError is + not retried again. Storage retries can multiply the total requests per upload attempt. :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 = 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_cached_table(self, name: str, fb: IFeedback, retry: Retry | None) -> None: if retry is None: retry = Retry(logger=logger) semaphore = get_upload_semaphore() - async with semaphore: - # Compression and hashing must not block other uploads on the event loop. - writer = asyncio.create_task(asyncio.to_thread(self.save_table, table, table_format)) - try: - table_info = await asyncio.shield(writer) - except asyncio.CancelledError: - # A running worker cannot be cancelled; finish its cache write before allowing a retry. - await asyncio.gather(writer, return_exceptions=True) - raise - - loop = asyncio.get_running_loop() - pending = _inflight_uploads.setdefault(loop, {}) - key = (id(self._connector), self._environment, table_info["data"]) - if (existing := pending.get(key)) is not None: - await asyncio.shield(existing) - fb.progress(1) - return table_info - - pending[key] = completion = loop.create_future() - try: - await self._upload_cached_table(table_info["data"], fb, retry, semaphore) - except asyncio.CancelledError: - completion.cancel() - raise - except BaseException as exc: - completion.set_exception(exc) - # The owner propagates the failure even when no duplicate is waiting. - completion.exception() - raise - else: - completion.set_result(None) - finally: - del pending[key] - if not pending: - del _inflight_uploads[loop] - return table_info - - async def _upload_cached_table(self, name: str, fb: IFeedback, retry: Retry, semaphore: asyncio.Semaphore) -> None: upload = ObjectDataUpload(connector=self._connector, environment=self._environment, name=name) with error_context(f"While uploading table data '{name}'", logger=logger): async for attempt in retry: @@ -280,24 +247,15 @@ async def _upload_cached_table(self, name: str, fb: IFeedback, retry: Retry, sem except DataExistsError: logger.debug(f"Data not uploaded because data already exists (label: {name})") fb.progress(1) - except ( - EvoAPIException, - TransportError, - asyncio.TimeoutError, - TimeoutError, - ConnectionError, - ChunkedIOError, - ) as exc: + except Exception as exc: if not _is_transient_upload_error(exc): raise - try: - attempt.set_exception(exc) - except RetryError: - # Keep the original type, instance and storage cause; retry history is the context. - raise exc + attempt.set_exception(exc) fb.progress(0) - async def upload_category_table(self, table: pa.Table, fb: IFeedback = NoFeedback) -> CategoryInfo: + 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. @@ -306,6 +264,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. @@ -313,6 +272,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(): @@ -351,14 +315,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, lookup_table = await gather_uploads( - self.upload_table(values, fb=values_fb, table_format=value_table_format), - self.upload_table(lookup, fb=lookup_fb, table_format=LOOKUP_TABLE_INT32), + 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), ) - return CategoryInfo(values=cast(ArrayTableInfo, value_table), table=cast(LookupTableInfo, lookup_table)) + + 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_uploads( + 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 @@ -513,6 +486,7 @@ async def upload_dataframe( dataframe: pd.DataFrame, fb: IFeedback = NoFeedback, table_format: KnownTableFormat | Iterable[KnownTableFormat] | None = None, + retry: Retry | None = None, ) -> dict: """Upload pandas dataframe to the geoscience object service, returning a GO model of the uploaded data. @@ -522,18 +496,20 @@ 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. :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) + table_info = await run_upload_preparation(self.save_dataframe, dataframe, table_format) + 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. @@ -544,16 +520,20 @@ 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. :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 75047c379..1940d4236 100644 --- a/packages/evo-objects/tests/test_data_client.py +++ b/packages/evo-objects/tests/test_data_client.py @@ -14,8 +14,8 @@ import logging import sys import threading +import weakref from contextlib import nullcontext -from dataclasses import replace from unittest import mock from uuid import UUID, uuid4 @@ -25,15 +25,17 @@ from parameterized import parameterized from data import load_test_data -from evo.common import IFeedback, RequestMethod +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 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, upload_concurrency -from evo.objects.utils._concurrency import gather_uploads +from evo.objects.utils._concurrency import _DEFAULT_MAX_CONCURRENCY, gather_uploads from helpers import NoImport, UnloadModule, get_sample_table_and_bytes @@ -221,7 +223,7 @@ async def test_upload_failure_preserves_exception_and_identifies_data( 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(error_type) as cm, + self.assertRaises(RetryError if upload_type == "table_transient" else error_type) as cm, ): destination.upload_file.side_effect = error if upload_type == "referenced": @@ -233,9 +235,9 @@ async def test_upload_failure_preserves_exception_and_identifies_data( table, retry=Retry(logger=logger, backoff_method=BackoffLinear(0)) ) if upload_type == "table_transient": - self.assertIsInstance(cm.exception.__context__, RetryError) - self.assertEqual((error, error, error), cm.exception.__context__.exceptions) - self.assertIs(error, cm.exception) + self.assertEqual((error, error, error), cm.exception.exceptions) + else: + self.assertIs(error, cm.exception) context = ( "\n".join(getattr(cm.exception, "__notes__", [])) if sys.version_info >= (3, 11) @@ -419,17 +421,24 @@ async def _mock_upload_file_side_effect(*args, **kwargs): ) self.assertIs(mock_table_info, actual_table_info) - @parameterized.expand([("direct", False), ("nested_group", True)]) - async def test_cancelled_upload_waits_for_parquet_writer(self, _name: str, nested_group: bool) -> None: + @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): + def save_table(*_args, **_kwargs): loop.call_soon_threadsafe(started.set) try: - if not release.wait(timeout=5): + if not release.wait(timeout=15): raise TimeoutError("The test did not release the parquet writer") return {"data": "saved-table"} finally: @@ -441,7 +450,10 @@ def save_table(*_args): ): async def run(): - operation = self.data_client.upload_table(pa.table({"value": [1.0]})) + 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_uploads(operation) return await operation @@ -449,7 +461,7 @@ async def run(): task = asyncio.create_task(run()) cancelled = False try: - await asyncio.wait_for(started.wait(), timeout=2) + await asyncio.wait_for(started.wait(), timeout=10) task.cancel() cancelled = True await asyncio.sleep(0) @@ -493,6 +505,60 @@ async def test_parquet_writer_failure_is_propagated_without_uploading(self) -> N 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), @@ -504,7 +570,7 @@ async def test_parquet_writer_failure_is_propagated_without_uploading(self) -> N ] ) 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 8 + 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 @@ -591,12 +657,12 @@ async def run(): upload_concurrency(limit) if limit is not None else nullcontext(), ): for _ in range(2): - result = asyncio.run(run()) + 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)]) + @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 8 + limit = configured_limit if configured_limit is not None else _DEFAULT_MAX_CONCURRENCY release = asyncio.Event() limit_reached = asyncio.Event() active = 0 @@ -637,7 +703,6 @@ async def upload(**_kwargs): [(str(status), EvoAPIException(status, "transient", None, None)) for status in (408, 429, 500, 502, 503, 504)] + [ ("timeout", TimeoutError()), - ("transport", TransportError("connection lost")), ("connection", ConnectionError()), ] ) @@ -690,22 +755,85 @@ async def test_retry_covers_actual_upload_url_acquisition(self): 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(RetryError) as caught: + await getattr(self.data_client, method)( + data, retry=Retry(logger=mock.Mock(), max_attempts=1, backoff_method=BackoffLinear(0)) + ) + self.assertEqual(1, len(caught.exception.exceptions)) + self.assertEqual(2 if category else 1, upload.await_count) + @parameterized.expand([("chunk", "comp=block&"), ("commit", "comp=blocklist")]) - async def test_retry_covers_wrapped_storage_transport_failure(self, _name, query): + 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 - raise TransportError("connection lost") + 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", @@ -717,6 +845,10 @@ async def request(**kwargs): 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() @@ -724,9 +856,7 @@ async def request(**kwargs): @parameterized.expand( [ ("url", EvoAPIException), - ("storage", ChunkedIOError), ("connection", ConnectionError), - ("transport", TransportError), ] ) async def test_retry_exhaustion_is_bounded_and_preserves_errors(self, _name, error_type): @@ -734,22 +864,17 @@ async def test_retry_exhaustion_is_bounded_and_preserves_errors(self, _name, err 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}") - if error_type is ChunkedIOError: - error.__cause__ = cause errors.append(error) - original_cause = errors[-1].__cause__ 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: + with self.assertRaises(RetryError) 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.assertIs(original_cause, caught.exception.__cause__) - self.assertIsInstance(caught.exception.__context__, RetryError) - self.assertEqual(tuple(errors), caught.exception.__context__.exceptions) + self.assertEqual(tuple(errors), caught.exception.exceptions) + self.assertIsNone(errors[-1].__context__) self.assertEqual(3, upload.await_count) save.assert_called_once() @@ -805,148 +930,73 @@ async def upload(upload, **_kwargs): await asyncio.gather(task, return_exceptions=True) self.assertEqual({"first": 2, "second": 1}, attempts) - @parameterized.expand([("success", False), ("failure", True)]) - async def test_identical_tables_share_inflight_upload_across_clients(self, _name, fail): - started = asyncio.Event() - joined = asyncio.Event() - release = asyncio.Event() - error = RuntimeError("upload failed") - shield = asyncio.shield - feedback = mock.Mock(spec=IFeedback) + @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]}) - def watch_join(awaitable): - if not isinstance(awaitable, asyncio.Task): - joined.set() - return shield(awaitable) - async def upload(**_kwargs): - started.set() - await release.wait() - if fail: - raise error - - with ( - mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload) as upload_mock, - mock.patch("evo.objects.utils.data.asyncio.shield", side_effect=watch_join), - ): - owner = asyncio.create_task(self.data_client.upload_table(table)) - tasks = [owner] + 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.wait(), timeout=2) - tasks.append(asyncio.create_task(other_client.upload_table(table, fb=feedback))) - await asyncio.wait_for(joined.wait(), timeout=2) - upload_mock.assert_awaited_once() - self.assertFalse(tasks[1].done()) - release.set() - results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=2) - if fail: - self.assertTrue(all(result is error for result in results)) - feedback.progress.assert_not_called() - else: - self.assertEqual(results[0], results[1]) - self.assertIsNot(results[0], results[1]) - feedback.progress.assert_called_with(1) - - # Completion (including failure) must remove the in-flight entry. - upload_mock.side_effect = None - await self.data_client.upload_table(table) - self.assertEqual(2, upload_mock.await_count) - finally: - release.set() - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - - @parameterized.expand([("waiter", False), ("owner", True)]) - async def test_duplicate_upload_cancellation_cleans_up_without_orphaned_work(self, _name, cancel_owner): - started = asyncio.Event() - joined = asyncio.Event() - release = asyncio.Event() - shield = asyncio.shield - - def watch_join(awaitable): - if not isinstance(awaitable, asyncio.Task): - joined.set() - return shield(awaitable) - - async def upload(**_kwargs): - started.set() - await release.wait() - - 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 upload_mock, - mock.patch("evo.objects.utils.data.asyncio.shield", side_effect=watch_join), - ): - owner = asyncio.create_task(self.data_client.upload_table(mock.Mock())) - tasks = [owner] - try: - await asyncio.wait_for(started.wait(), timeout=2) - waiter = asyncio.create_task(self.data_client.upload_table(mock.Mock())) - tasks.append(waiter) - await asyncio.wait_for(joined.wait(), timeout=2) - cancelled = owner if cancel_owner else waiter - cancelled.cancel() - with self.assertRaises(asyncio.CancelledError): - await cancelled - if cancel_owner: + 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 asyncio.wait_for(waiter, timeout=2) + await first else: - self.assertFalse(owner.done()) - release.set() - self.assertEqual({"data": "saved"}, await asyncio.wait_for(owner, timeout=2)) - upload_mock.assert_awaited_once() - release.set() - await self.data_client.upload_table(mock.Mock()) - self.assertEqual(2, upload_mock.await_count) + 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: - release.set() + for event in release: + event.set() for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) + self.assertEqual(2, calls) - @parameterized.expand([("hub_url",), ("org_id",), ("workspace_id",), ("connector",)]) - async def test_inflight_uploads_do_not_share_across_destinations_or_connectors(self, field): - environment = self.environment - connector = self.connector - if field == "connector": - connector = mock.Mock(wraps=self.connector) - else: - environment = replace(environment, **{field: "https://other.example" if field == "hub_url" else uuid4()}) - other_client = ObjectDataClient(environment=environment, connector=connector, cache=self.cache) - started = asyncio.Event() - both_started = asyncio.Event() - release = asyncio.Event() - count = 0 - + async def test_identical_uploads_use_independent_retry_policies(self): async def upload(**_kwargs): - nonlocal count - count += 1 - started.set() - if count == 2: - both_started.set() - await release.wait() + await asyncio.sleep(0) + raise ConnectionError("retry") with ( - mock.patch.object(ObjectDataClient, "save_table", return_value={"data": "saved"}), - mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=upload), + 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, ): - tasks = [asyncio.create_task(self.data_client.upload_table(mock.Mock()))] - try: - await asyncio.wait_for(started.wait(), timeout=2) - tasks.append(asyncio.create_task(other_client.upload_table(mock.Mock()))) - await asyncio.wait_for(both_started.wait(), timeout=2) - release.set() - await asyncio.wait_for(asyncio.gather(*tasks), timeout=2) - finally: - release.set() - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - self.assertEqual(2, count) + 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, RetryError) for result in results)) + self.assertEqual([1, 2], [len(result.exceptions) for result in results]) + self.assertEqual(3, transfer.await_count) @parameterized.expand( [ @@ -954,6 +1004,13 @@ async def upload(**_kwargs): ("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()), @@ -990,14 +1047,13 @@ async def test_category_values_and_lookup_upload_concurrently(self) -> None: all_started = asyncio.Event() release = asyncio.Event() - async def upload(table, **_kwargs): - started.append(table.column_names) + async def upload(upload, **_kwargs): + started.append(upload.name) if len(started) == 2: all_started.set() await release.wait() - return {"data": table.column_names[0], "length": table.num_rows} - with mock.patch.object(self.data_client, "upload_table", side_effect=upload): + 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) @@ -1008,8 +1064,7 @@ async def upload(table, **_kwargs): task.cancel() await asyncio.gather(task, return_exceptions=True) - self.assertEqual("category", result["values"]["data"]) - self.assertEqual("key", result["table"]["data"]) + 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.""" @@ -1134,16 +1189,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) @@ -1205,16 +1262,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) @@ -1223,11 +1282,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 index 796237530..ba3ac8f88 100644 --- a/packages/evo-objects/tests/test_upload_concurrency.py +++ b/packages/evo-objects/tests/test_upload_concurrency.py @@ -17,14 +17,28 @@ from parameterized import parameterized -from evo.objects.utils import upload_concurrency -from evo.objects.utils._concurrency import gather_uploads, get_max_concurrent_uploads, get_upload_semaphore +from evo.objects.utils import ObjectDataClient, upload_concurrency +from evo.objects.utils._concurrency import ( + _DEFAULT_MAX_CONCURRENCY, + _upload_scope, + gather_uploads, + get_upload_semaphore, +) class TestUploadConcurrency(IsolatedAsyncioTestCase): + async def test_upload_and_download_defaults_match(self): + self.assertEqual(4, _DEFAULT_MAX_CONCURRENCY) + 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 8 + limit = configured_limit if configured_limit is not None else _DEFAULT_MAX_CONCURRENCY default = get_upload_semaphore() release = asyncio.Event() limit_reached = asyncio.Event() @@ -81,22 +95,22 @@ async def test_invalid_configuration_leaves_scope_unchanged(self, value): with self.assertRaisesRegex(ValueError, "max_concurrent must be a positive integer"): with upload_concurrency(value): self.fail("Invalid configuration was accepted") - self.assertEqual(8, get_max_concurrent_uploads()) + 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, get_max_concurrent_uploads()) + self.assertEqual(3, _upload_scope.get().limit) with self.assertRaisesRegex(RuntimeError, "failed"): with upload_concurrency(1): - self.assertEqual(1, get_max_concurrent_uploads()) + self.assertEqual(1, _upload_scope.get().limit) self.assertIsNot(outer, get_upload_semaphore()) raise RuntimeError("failed") - self.assertEqual(3, get_max_concurrent_uploads()) + self.assertEqual(3, _upload_scope.get().limit) self.assertIs(outer, get_upload_semaphore()) - self.assertEqual(8, get_max_concurrent_uploads()) + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) self.assertIs(default, get_upload_semaphore()) async def test_configuration_is_isolated_between_tasks(self): @@ -108,12 +122,12 @@ async def configured(): with upload_concurrency(2): ready.set() await release.wait() - return get_max_concurrent_uploads() + return _upload_scope.get().limit task = asyncio.create_task(configured()) try: await asyncio.wait_for(ready.wait(), timeout=2) - self.assertEqual(8, get_max_concurrent_uploads()) + 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)) @@ -201,6 +215,19 @@ async def upload(*, fail=False): async def test_empty_group_returns_empty_list(self): self.assertEqual([], await gather_uploads()) + 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_uploads(*(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) @@ -222,7 +249,7 @@ def schedule(coroutine): self.assertTrue(all(task.done() for task in tasks)) self.assertTrue(all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in coroutines)) - self.assertEqual(8, get_max_concurrent_uploads()) + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) async def test_semaphore_creation_failure_drains_all_coroutines(self): async def upload(): @@ -235,7 +262,7 @@ async def upload(): await gather_uploads(*coroutines) self.assertTrue(all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in coroutines)) - self.assertEqual(8, get_max_concurrent_uploads()) + self.assertEqual(_DEFAULT_MAX_CONCURRENCY, _upload_scope.get().limit) class TestUploadConcurrencyAcrossLoops(TestCase): @@ -254,7 +281,7 @@ async def upload(): await asyncio.gather(*(upload() for _ in range(16))) with upload_concurrency(limit) if limit is not None else nullcontext(): - asyncio.run(run()) - asyncio.run(run()) + asyncio.run(asyncio.wait_for(run(), timeout=10)) + asyncio.run(asyncio.wait_for(run(), timeout=10)) self.assertIsNot(semaphores[0], semaphores[1]) diff --git a/packages/evo-objects/tests/typed/test_model.py b/packages/evo-objects/tests/typed/test_model.py index 9dddca1c6..b75ece9fd 100644 --- a/packages/evo-objects/tests/typed/test_model.py +++ b/packages/evo-objects/tests/typed/test_model.py @@ -28,6 +28,7 @@ from evo.common.test_tools import BASE_URL, ORG, WORKSPACE_ID, TestWithConnector from evo.objects.typed._data import DataTable, DataTableAndAttributes 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 @@ -63,6 +64,14 @@ class FakeData: self.assertEqual(result["format_version"], "1.0.0") self.assertEqual(result["name"], "test") + async def test_set_sub_model_value_remains_compatible(self): + builder = SchemaBuilder(ExtendedLocations, self.context) + frame = pd.DataFrame({"x": [1.0], "y": [2.0], "z": [3.0]}) + with patch.object(DataTable, "_data_to_schema", return_value={"data": "coordinates"}) as convert: + await builder.set_sub_model_value(name="_table", data=frame) + convert.assert_awaited_once_with(frame, context=self.context) + self.assertEqual({"coordinates": {"data": "coordinates"}}, builder.document) + def test_union_schema_list_requires_polymorphic_overrides(self): class First(SchemaModel): pass @@ -221,6 +230,22 @@ async def upload(frame, **_kwargs): 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(2, 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.attrs["attribute_descriptions"] = {"first": "invalid"} client = Mock(upload_dataframe=AsyncMock()) 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. From 5409813f672ad7cf5bfd6f14a85d83e2214ce915 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 10 Nov 2026 17:25:09 +0000 Subject: [PATCH 5/5] Refactor upload concurrency management to use gather_cancelling for improved error handling and task cancellation --- packages/evo-objects/README.md | 31 ++--- packages/evo-objects/pyproject.toml | 2 +- .../src/evo/objects/typed/_data.py | 36 +++--- .../src/evo/objects/typed/_model.py | 10 +- .../src/evo/objects/typed/attributes.py | 121 ++++++++++++------ .../evo/objects/typed/downhole_collection.py | 6 +- .../src/evo/objects/utils/_concurrency.py | 35 +++-- .../evo-objects/src/evo/objects/utils/data.py | 101 +++++++++++---- .../evo-objects/tests/test_data_client.py | 67 ++++++++-- .../tests/test_upload_concurrency.py | 76 +++++++++-- .../tests/typed/test_attributes.py | 61 +++++++++ .../evo-objects/tests/typed/test_model.py | 49 +++++-- packages/evo-sdk-common/pyproject.toml | 2 +- uv.lock | 2 +- 14 files changed, 436 insertions(+), 163 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index 1240694de..d09c1aa88 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -107,10 +107,8 @@ distance or `from` column is ignored. ### Upload concurrency -Typed objects upload independent tables and attributes concurrently, preserving input order and applying -multi-column and combined table/attribute mutations only after every upload succeeds. Typed operations and direct -`ObjectDataClient` uploads share a default concurrency of **4**, matching downloads. Use `upload_concurrency` -to change the limit for a block, including its child tasks: +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 @@ -119,22 +117,19 @@ with upload_concurrency(max_concurrent=8): dhc = await DownholeCollection.create(context, data) ``` -The limit bounds concurrent Arrow conversion, category preparation, Parquet serialization, and upload attempts, -not total memory usage or the size of caller-owned data. Preparation runs in worker threads, and temporary Arrow -tables are released before network uploads. Nested scopes restore the previous setting on exit. +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 HTTP statuses (408/429/500/502/503/504) and raw connection/timeouts up to three attempts -with exponential backoff, reusing cached data. Exhaustion raises `RetryError` with the attempt history. -Use the optional `retry=Retry(...)` argument on any `ObjectDataClient` table, DataFrame, or category upload method -to customize this policy. Already-exhausted `TransportError` failures are not retried again. Storage has its own -retry policy; exhausted storage retries are retried only when all their errors are transient, so storage request -counts can exceed the table-attempt limit. Failed operations cancel and drain their own siblings, not independent uploads. +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. -### HTTP connections - -HTTP connections are configured separately: `AioTransport(num_pools=...)` sets the total connection limit -(default **4**) across requests sharing that transport. `upload_concurrency` does not resize it. For example, -pass a configured transport to your context using your existing authorizer, cache, and environment: +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 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 03b613b93..cabee40d9 100644 --- a/packages/evo-objects/src/evo/objects/typed/_data.py +++ b/packages/evo-objects/src/evo/objects/typed/_data.py @@ -19,7 +19,7 @@ from evo.common.interfaces import IContext from evo.common.utils import NoFeedback, split_feedback from evo.objects.typed.attributes import Attributes -from evo.objects.utils._concurrency import gather_uploads +from evo.objects.utils._concurrency import gather_cancelling from evo.objects.utils.table_formats import KnownTableFormat from ._model import SchemaBuilder, SchemaLocation, SchemaModel @@ -50,18 +50,14 @@ 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.commit(df, await self.prepare(df, fb)) + self._apply_uploaded_table(df, await self._upload(df, fb)) - async def prepare(self, df: pd.DataFrame, fb: IFeedback = NoFeedback) -> dict[str, Any]: + 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 commit(self, df: pd.DataFrame, document: dict[str, Any]) -> None: - """Apply a successfully prepared table, including subclass-specific derived metadata.""" - self._apply_uploaded_table(df, document) - def _apply_uploaded_table(self, df: pd.DataFrame, document: dict[str, Any]) -> None: - """Commit uploaded values and any subclass-specific derived metadata.""" + """Commit uploaded values and derived metadata; override to extend table replacement behavior.""" self._document.update(document) self._context.mark_modified(self._data) @@ -129,19 +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) - if attr_df is not None: - table_fb, attr_fb = split_feedback(fb, [1.0, float(len(attr_df.columns))]) - # prepare validates metadata before either branch starts remote work. - attribute_upload = self.attributes.prepare(attr_df, attr_fb) - table_document, attribute_documents = await gather_uploads( - self._table.prepare(table_df, table_fb), - attribute_upload, - ) - self._table.commit(table_df, table_document) - self.attributes.commit(attribute_documents) - else: - await self._table.from_dataframe(table_df, fb=fb) - 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.""" diff --git a/packages/evo-objects/src/evo/objects/typed/_model.py b/packages/evo-objects/src/evo/objects/typed/_model.py index 56a5a783d..b971ab03b 100644 --- a/packages/evo-objects/src/evo/objects/typed/_model.py +++ b/packages/evo-objects/src/evo/objects/typed/_model.py @@ -21,7 +21,7 @@ from evo import jmespath from evo.common import IContext from evo.objects import DownloadedObject -from evo.objects.utils._concurrency import gather_uploads +from evo.objects.utils._concurrency import gather_cancelling from ._utils import ( assign_jmespath_value, @@ -208,16 +208,12 @@ async def build_from_data(self, data: Any, skip_sub_models: set[str] | None = No return self.document - async def set_sub_model_value(self, name: str, data: Any) -> None: - """Convert and assign a single sub-model.""" - await self.set_sub_model_values([(name, data)]) - 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_uploads( + 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): @@ -514,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 gather_uploads(*(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 19c713bfd..457bb3679 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -13,7 +13,7 @@ import typing import uuid -from collections.abc import Coroutine +from copy import deepcopy from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any, Protocol, runtime_checkable from uuid import UUID @@ -24,7 +24,7 @@ from evo.common import IContext, IFeedback from evo.common.utils import NoFeedback, iter_with_fb, split_feedback from evo.objects import DownloadedObject -from evo.objects.utils._concurrency import gather_uploads +from evo.objects.utils._concurrency import gather_cancelling, run_upload_preparation from evo.objects.utils.table_formats import ( BOOL_ARRAY_1, DATE_TIME_ARRAY, @@ -214,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": []} @@ -263,24 +268,34 @@ def __repr__(self) -> str: 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]: - attribute_type = _infer_attribute_type_from_series(series) document: 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") - await Attribute._upload_attribute_values(document, pd.DataFrame({col: series}), attribute_type, data_client, fb=fb) + 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 @@ -327,17 +342,30 @@ async def _data_to_schema( return await cls._upload_attributes(data, context) if data is not None else [] @staticmethod - def _attribute_descriptions(df: pd.DataFrame) -> dict[Any, dict[str, Any]]: - """Validate descriptions and convert them to schema dictionaries once.""" + 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", {}) - result = {} + 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.") - result[col] = description.to_schema() - return result + 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( @@ -353,16 +381,18 @@ async def _upload_attributes( :param context: The context used for data upload operations. :param fb: Optional feedback object to report progress. """ - descriptions = Attributes._attribute_descriptions(df) 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)) - return await gather_uploads( + return await gather_cancelling( *( - _build_new_attribute_document(col, df[col], descriptions.get(col), data_client, col_fb) + _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) ) ) @@ -423,26 +453,26 @@ async def set_attributes(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): :param fb: Optional feedback object to report upload progress. """ - self.commit(await self.prepare(df, fb)) + descriptions, attribute_types = await self._validate(df) + prepared = await self._upload(df, descriptions, attribute_types, fb) + self._apply(prepared) - def prepare( - self, df: pd.DataFrame, fb: IFeedback = NoFeedback - ) -> Coroutine[Any, Any, list[tuple[dict[str, Any] | None, dict[str, Any]]]]: - """Validate metadata now and return an upload coroutine, without changing live attributes. + 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}) - Await the returned coroutine and pass its result to commit only after all - related uploads succeed. - """ - descriptions = self._attribute_descriptions(df) - return self._prepare_attributes(df, descriptions, fb) - - async def _prepare_attributes( - self, df: pd.DataFrame, descriptions: dict[Any, dict[str, Any]], fb: IFeedback - ) -> list[tuple[dict[str, Any] | None, dict[str, Any]]]: + 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 [] + return _PreparedAttributes([], originals, original_snapshot) attributes_by_name = {attr.name: attr for attr in self} data_client = get_data_client(self._obj) column_fbs = split_feedback(fb, [1.0] * len(df.columns)) @@ -452,26 +482,39 @@ async def upload_column(col: Any, col_fb: IFeedback) -> tuple[dict[str, Any] | N if attribute is not None: document = attribute.as_dict() await Attribute._upload_attribute_values( - document, df[[col]], attribute.attribute_type, data_client, fb=col_fb + 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], descriptions.get(col), data_client, col_fb) + document = await _build_new_attribute_document( + col, df[col], attribute_types[col], descriptions.get(col), data_client, col_fb + ) return None, document - return await gather_uploads(*(upload_column(col, col_fb) for col, col_fb in zip(df.columns, column_fbs))) - - def commit(self, prepared: list[tuple[dict[str, Any] | None, dict[str, Any]]]) -> None: - """Replace live attributes with a successfully prepared batch, preserving existing wrappers.""" + 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: + 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) - self._document[:] = documents - for document in documents: + 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 24665e31a..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,7 +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_uploads +from evo.objects.utils._concurrency import gather_cancelling from evo.objects.utils.table_formats import ( DOWNHOLE_COLLECTION_LOCATION_HOLES, FLOAT_ARRAY_1, @@ -354,7 +354,7 @@ 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) - values, table = await gather_uploads( + 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 @@ -489,7 +489,7 @@ async def convert(collection: DownholeCollectionEntry) -> dict[str, Any]: return schema # Collections are independent, so upload them concurrently while preserving input order. - return await gather_uploads(*(convert(collection) for collection in data)) + 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/utils/_concurrency.py b/packages/evo-objects/src/evo/objects/utils/_concurrency.py index 56788734f..09a6b35e6 100644 --- a/packages/evo-objects/src/evo/objects/utils/_concurrency.py +++ b/packages/evo-objects/src/evo/objects/utils/_concurrency.py @@ -15,6 +15,7 @@ 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 @@ -35,18 +36,21 @@ class _UploadScope: 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() - 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 + # 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( @@ -60,7 +64,7 @@ def upload_concurrency(max_concurrent: int = _DEFAULT_MAX_CONCURRENCY) -> Iterat Applies to typed-object operations and direct ObjectDataClient uploads, including clients created internally. Conversion/serialization and each upload attempt acquire a slot; - backoff sleeps do not. Settings are inherited by child asyncio tasks, not unrelated + 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. @@ -88,21 +92,24 @@ async def run_upload_preparation(function: Callable[_P, _T], *args: _P.args, **k try: return await asyncio.shield(worker) except asyncio.CancelledError: - # Running threads cannot be cancelled; finish their cache writes before returning. - await asyncio.gather(worker, return_exceptions=True) + # 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_uploads(*coroutines: Coroutine[Any, Any, _T]) -> list[_T]: +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)) - # Unlike gather(), this does not cancel children before our own cleanup runs. - for completed in asyncio.as_completed(tasks): - await completed - return [task.result() for task in tasks] + return list(await asyncio.gather(*tasks)) except BaseException as error: for coroutine in coroutines[len(tasks) :]: coroutine.close() diff --git a/packages/evo-objects/src/evo/objects/utils/data.py b/packages/evo-objects/src/evo/objects/utils/data.py index 767688b84..0c7a8cfca 100644 --- a/packages/evo-objects/src/evo/objects/utils/data.py +++ b/packages/evo-objects/src/evo/objects/utils/data.py @@ -10,10 +10,11 @@ # 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 @@ -25,7 +26,7 @@ from ..exceptions import TableFormatError from ..io import _CACHE_SCOPE, ObjectDataDownload, ObjectDataUpload -from ._concurrency import _DEFAULT_MAX_CONCURRENCY, gather_uploads, get_upload_semaphore, run_upload_preparation +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 @@ -64,6 +65,18 @@ def _is_transient_upload_error(error: Exception) -> bool: 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]: """Iterate over all data references in a geoscience object. @@ -123,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: @@ -222,8 +238,8 @@ async def upload_table( 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. Exhaustion raises RetryError with the attempt history. TransportError is - not retried again. Storage retries can multiply the total requests per upload attempt. + 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. @@ -235,23 +251,61 @@ async def upload_table( return table_info 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: - await upload.upload_from_cache(cache=self._cache, transport=self._connector.transport, fb=fb) + 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 - attempt.set_exception(exc) - fb.progress(0) + 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 @@ -327,7 +381,7 @@ 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_uploads( + 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), ) @@ -443,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. @@ -487,6 +536,7 @@ async def upload_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. @@ -497,13 +547,19 @@ async def upload_dataframe( 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. """ - table_info = await run_upload_preparation(self.save_dataframe, dataframe, table_format) + + 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 @@ -521,6 +577,7 @@ async def upload_category_dataframe( :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. diff --git a/packages/evo-objects/tests/test_data_client.py b/packages/evo-objects/tests/test_data_client.py index 1940d4236..32c468303 100644 --- a/packages/evo-objects/tests/test_data_client.py +++ b/packages/evo-objects/tests/test_data_client.py @@ -35,7 +35,7 @@ from evo.objects.endpoints.models import DataUploadResponseBody from evo.objects.exceptions import TableFormatError from evo.objects.utils import KnownTableFormat, ObjectDataClient, table_formats, upload_concurrency -from evo.objects.utils._concurrency import _DEFAULT_MAX_CONCURRENCY, gather_uploads +from evo.objects.utils._concurrency import _DEFAULT_MAX_CONCURRENCY, gather_cancelling from helpers import NoImport, UnloadModule, get_sample_table_and_bytes @@ -223,7 +223,7 @@ async def test_upload_failure_preserves_exception_and_identifies_data( 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(RetryError if upload_type == "table_transient" else error_type) as cm, + self.assertRaises(error_type) as cm, ): destination.upload_file.side_effect = error if upload_type == "referenced": @@ -235,7 +235,7 @@ async def test_upload_failure_preserves_exception_and_identifies_data( table, retry=Retry(logger=logger, backoff_method=BackoffLinear(0)) ) if upload_type == "table_transient": - self.assertEqual((error, error, error), cm.exception.exceptions) + self.assertEqual((error, error, error), cm.exception.__cause__.exceptions) else: self.assertIs(error, cm.exception) context = ( @@ -455,7 +455,7 @@ async def run(): data = pa.Table.from_pandas(data) operation = getattr(self.data_client, method)(data) if nested_group: - return await gather_uploads(operation) + return await gather_cancelling(operation) return await operation task = asyncio.create_task(run()) @@ -726,7 +726,7 @@ async def upload(upload_context, **_kwargs): mock.patch("evo.objects.io.ObjectDataUpload.upload_from_cache", new=upload), ): results = await asyncio.wait_for( - gather_uploads( + gather_cancelling( self.data_client.upload_table("first", retry=retry), self.data_client.upload_table("second", retry=retry), ), @@ -800,11 +800,12 @@ async def test_public_upload_methods_honor_retry_policy(self, method, category): "evo.objects.io.ObjectDataUpload.upload_from_cache", side_effect=EvoAPIException(503, "unavailable", None, None), ) as upload: - with self.assertRaises(RetryError) as caught: + 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.assertEqual(1, len(caught.exception.exceptions)) + 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")]) @@ -869,12 +870,12 @@ async def test_retry_exhaustion_is_bounded_and_preserves_errors(self, _name, err 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(RetryError) as caught: + 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.assertEqual(tuple(errors), caught.exception.exceptions) - self.assertIsNone(errors[-1].__context__) + self.assertIs(errors[-1], caught.exception) + self.assertEqual(tuple(errors), caught.exception.__cause__.exceptions) self.assertEqual(3, upload.await_count) save.assert_called_once() @@ -994,10 +995,52 @@ async def upload(**_kwargs): ), return_exceptions=True, ) - self.assertTrue(all(isinstance(result, RetryError) for result in results)) - self.assertEqual([1, 2], [len(result.exceptions) for result in results]) + 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)), diff --git a/packages/evo-objects/tests/test_upload_concurrency.py b/packages/evo-objects/tests/test_upload_concurrency.py index ba3ac8f88..8c4d01391 100644 --- a/packages/evo-objects/tests/test_upload_concurrency.py +++ b/packages/evo-objects/tests/test_upload_concurrency.py @@ -11,6 +11,8 @@ 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 @@ -21,14 +23,14 @@ from evo.objects.utils._concurrency import ( _DEFAULT_MAX_CONCURRENCY, _upload_scope, - gather_uploads, + gather_cancelling, get_upload_semaphore, + run_upload_preparation, ) class TestUploadConcurrency(IsolatedAsyncioTestCase): - async def test_upload_and_download_defaults_match(self): - self.assertEqual(4, _DEFAULT_MAX_CONCURRENCY) + def test_upload_and_download_defaults_match(self): self.assertEqual( _DEFAULT_MAX_CONCURRENCY, inspect.signature(ObjectDataClient.download_tables).parameters["max_concurrency"].default, @@ -62,7 +64,7 @@ async def upload(index): active -= 1 async def group(start): - return await gather_uploads(*(upload(index) for index in range(start, start + 12))) + 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)) @@ -152,7 +154,7 @@ async def upload(index): cleaned_up.add(index) async def group(start): - return await gather_uploads(*(upload(index) for index in range(start, start + 12))) + return await gather_cancelling(*(upload(index) for index in range(start, start + 12))) async def fail(): await started.wait() @@ -160,7 +162,7 @@ async def fail(): try: with self.assertRaises(RuntimeError) as caught: - await gather_uploads(group(0), group(12), fail()) + await gather_cancelling(group(0), group(12), fail()) self.assertIs(failure, caught.exception) self.assertEqual(set(range(24)), cleaned_up) finally: @@ -180,10 +182,10 @@ async def upload(index): cleaned_up.add(index) async def group(start): - return await gather_uploads(*(upload(index) for index in range(start, start + 3))) + return await gather_cancelling(*(upload(index) for index in range(start, start + 3))) async def run(): - return await gather_uploads(group(0), group(3)) + return await gather_cancelling(group(0), group(3)) task = asyncio.create_task(run()) try: @@ -205,15 +207,49 @@ async def upload(*, fail=False): if fail: raise ValueError("invalid input") - await gather_uploads(upload()) + await gather_cancelling(upload()) self.assertIs(get_upload_semaphore(), semaphores[-1]) with self.assertRaisesRegex(ValueError, "invalid input"): - await gather_uploads(upload(fail=True)) + 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_uploads()) + 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")] @@ -223,7 +259,7 @@ async def fail(error): with self.assertLogs("evo.object.uploads", level="DEBUG") as logs: with self.assertRaises(RuntimeError) as caught: - await gather_uploads(*(fail(error) for error in errors)) + 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]) @@ -245,7 +281,7 @@ def schedule(coroutine): with patch("evo.objects.utils._concurrency.asyncio.create_task", side_effect=schedule): with self.assertRaisesRegex(RuntimeError, "task factory failed"): - await gather_uploads(*coroutines) + 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)) @@ -259,7 +295,7 @@ async def upload(): coroutines = [upload(), upload()] with patch("evo.objects.utils._concurrency.asyncio.Semaphore", side_effect=RuntimeError("no semaphore")): with self.assertRaisesRegex(RuntimeError, "no semaphore"): - await gather_uploads(*coroutines) + 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) @@ -284,4 +320,16 @@ async def upload(): 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 91958a96f..9bbf0929c 100644 --- a/packages/evo-objects/tests/typed/test_attributes.py +++ b/packages/evo-objects/tests/typed/test_attributes.py @@ -12,6 +12,7 @@ from __future__ import annotations import asyncio +import threading from copy import deepcopy from typing import Any, cast from unittest import IsolatedAsyncioTestCase, TestCase @@ -30,6 +31,7 @@ UnSupportedDataTypeError, _infer_attribute_type_from_series, ) +from evo.objects.utils._concurrency import get_upload_semaphore, upload_concurrency class TestConcurrentAttributeUploads(IsolatedAsyncioTestCase): @@ -239,6 +241,65 @@ async def test_empty_mutations_do_not_upload(self): 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[:] = [ diff --git a/packages/evo-objects/tests/typed/test_model.py b/packages/evo-objects/tests/typed/test_model.py index b75ece9fd..6cc229bf5 100644 --- a/packages/evo-objects/tests/typed/test_model.py +++ b/packages/evo-objects/tests/typed/test_model.py @@ -26,6 +26,7 @@ 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 ModelContext, SchemaBuilder, SchemaList, SchemaLocation, SchemaModel from evo.objects.typed.attributes import AttributeDescription @@ -64,14 +65,6 @@ class FakeData: self.assertEqual(result["format_version"], "1.0.0") self.assertEqual(result["name"], "test") - async def test_set_sub_model_value_remains_compatible(self): - builder = SchemaBuilder(ExtendedLocations, self.context) - frame = pd.DataFrame({"x": [1.0], "y": [2.0], "z": [3.0]}) - with patch.object(DataTable, "_data_to_schema", return_value={"data": "coordinates"}) as convert: - await builder.set_sub_model_value(name="_table", data=frame) - convert.assert_awaited_once_with(frame, context=self.context) - self.assertEqual({"coordinates": {"data": "coordinates"}}, builder.document) - def test_union_schema_list_requires_polymorphic_overrides(self): class First(SchemaModel): pass @@ -243,11 +236,12 @@ async def test_attribute_descriptions_are_prepared_once_before_combined_uploads( patch.object(AttributeDescription, "to_schema", autospec=True, return_value={"unit": "s"}) as describe, ): await self.model.from_dataframe(self.frame) - self.assertEqual(2, describe.call_count) + 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.attrs["attribute_descriptions"] = {"first": "invalid"} + 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), @@ -257,6 +251,21 @@ async def test_invalid_attribute_metadata_prevents_table_upload(self): 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) @@ -288,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) 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/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" },