From 075fe559673ff2343fda60ac1bb29a3250f35b5d Mon Sep 17 00:00:00 2001 From: prbatero <42007693+prbatero@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:28:56 -0400 Subject: [PATCH] perf(data-layer): add bounded batch reads and client reuse Add native and fallback keyed reads across metadata backends, exact metadata-type matching, one process-wide blocking I/O budget, reusable Blob clients and delegation keys, and atomic artifact downloads with regression coverage. --- .../azure_blob_artifact_storage.py | 109 ++++-- .../core/data_layer/abstract_data_layer.py | 49 ++- .../azure_blob_storage_data_layer.py | 216 ++++++----- .../data_layer/azure_cosmos_db_data_layer.py | 81 +++- .../data_layer/azure_data_lake_data_layer.py | 44 ++- .../data_layer/azure_postgresql_data_layer.py | 70 +++- .../local_file_system_data_layer.py | 28 +- .../src/hastegeo/core/data_layer/unified.py | 30 +- .../src/hastegeo/core/processors/metadata.py | 97 +++++ hastelib/src/hastegeo/core/utils/blob.py | 81 +++- hastelib/src/hastegeo/core/utils/metadata.py | 36 ++ hastelib/src/hastegeo/core/utils/parallel.py | 124 ++++++ hastelib/src/hastegeo/core/utils/perf.py | 34 +- .../test_azure_blob_artifact_storage.py | 238 ++++++++++++ .../test_azure_blob_storage_data_layer.py | 298 ++++++++++++++ .../core/data_layer/test_read_contracts.py | 367 ++++++++++++++++++ .../tests/core/processors/test_artifacts.py | 60 +-- .../core/processors/test_metadata_batch.py | 171 ++++++++ hastelib/tests/core/utils/test_blob.py | 128 ++++++ .../core/utils/test_metadata_type_matching.py | 26 ++ hastelib/tests/core/utils/test_parallel.py | 113 ++++++ hastelib/tests/core/utils/test_perf.py | 48 +++ 22 files changed, 2212 insertions(+), 236 deletions(-) create mode 100644 hastelib/src/hastegeo/core/utils/parallel.py create mode 100644 hastelib/tests/core/artifact_storage/test_azure_blob_artifact_storage.py create mode 100644 hastelib/tests/core/data_layer/test_azure_blob_storage_data_layer.py create mode 100644 hastelib/tests/core/data_layer/test_read_contracts.py create mode 100644 hastelib/tests/core/processors/test_metadata_batch.py create mode 100644 hastelib/tests/core/utils/test_metadata_type_matching.py create mode 100644 hastelib/tests/core/utils/test_parallel.py create mode 100644 hastelib/tests/core/utils/test_perf.py diff --git a/hastelib/src/hastegeo/core/artifact_storage/azure_blob_artifact_storage.py b/hastelib/src/hastegeo/core/artifact_storage/azure_blob_artifact_storage.py index adefea4c..76e23ee7 100644 --- a/hastelib/src/hastegeo/core/artifact_storage/azure_blob_artifact_storage.py +++ b/hastelib/src/hastegeo/core/artifact_storage/azure_blob_artifact_storage.py @@ -2,6 +2,7 @@ # Licensed under the MIT License. import json import os +import tempfile import time from datetime import datetime, timedelta, timezone from hashlib import sha256 @@ -12,17 +13,20 @@ import yaml from azure.core import MatchConditions from azure.core.exceptions import ResourceExistsError -from azure.identity import DefaultAzureCredential # type: ignore from azure.storage.blob import BlobClient # type: ignore from azure.storage.blob import ( AccessPolicy, BlobSasPermissions, - BlobServiceClient, ContainerSasPermissions, generate_blob_sas, generate_container_sas, ) +from hastegeo.core.utils.blob import ( + get_blob_service_client, + get_cached_user_delegation_key, +) from hastegeo.core.utils.logs import Logger +from hastegeo.core.utils.parallel import configured_worker_count, parallel_map from .abstract_artifact_storage import AbstractArtifactStorage @@ -48,34 +52,22 @@ def __init__( # which need Storage Blob Delegator / Owner. Write access alone suffices. self.serves_read_sas = serves_read_sas if connection_string: - credential = connection_string - self.blob_service_client = ( - BlobServiceClient.from_connection_string(connection_string) + self.blob_service_client = get_blob_service_client( + connection_string=connection_string ) self.user_delegation_key = None self.account_key = self.blob_service_client.credential.account_key self.identity_blob_service_client = ( - BlobServiceClient( - account_url=account_url, - credential=DefaultAzureCredential(), - ) + get_blob_service_client(account_url=account_url) if account_url and urlparse(account_url).scheme == "https" else None ) else: - credential = DefaultAzureCredential() - self.blob_service_client = BlobServiceClient( - account_url=account_url, credential=credential + self.blob_service_client = get_blob_service_client( + account_url=account_url ) self.identity_blob_service_client = self.blob_service_client - self.user_delegation_key = ( - self.blob_service_client.get_user_delegation_key( - datetime.now(timezone.utc), - datetime.now(timezone.utc) + timedelta(hours=1), - ) - if serves_read_sas - else None - ) + self.user_delegation_key = None self.account_key = None self.container_read_policy = container_read_policy_name @@ -97,7 +89,9 @@ def __init__( f"Container '{container}' created successfully." ) except ResourceExistsError: - self.logger.info(f"Container '{container}' already exists.") + self.logger.info( + f"Container '{container}' already exists." + ) if self.serves_read_sas: self._create_or_update_managed_access_policy() _INITIALIZED_CONTAINERS.add(cache_key) @@ -184,12 +178,18 @@ def get_download_url( # Otherwise generate SAS token + user_delegation_key = self.user_delegation_key + if self.account_key is None: + user_delegation_key = get_cached_user_delegation_key( + self.blob_service_client + ) + sas_token = generate_container_sas( account_name=self.container_client.account_name, container_name=self.container_client.container_name, policy_id=self.container_read_policy, account_key=self.account_key, - user_delegation_key=self.user_delegation_key, + user_delegation_key=user_delegation_key, ) return str(f"{blob_client.url}?{sas_token}") @@ -210,15 +210,45 @@ def fetch_artifact( src_path = self.get_file_path( identifier, extra_partition_keys=extra_partition_keys ) - blobs = self.container_client.list_blobs(name_starts_with=src_path) - for blob in blobs: - file_path = os.path.join(dst_path, blob.name) + if not src_path: + raise ValueError("A source artifact path is required") + if not dst_path: + raise ValueError("A destination path is required") + src_path = self.resolve_artifact_path(src_path) + blob_names = [ + blob.name + for blob in self.container_client.list_blobs( + name_starts_with=src_path + ) + ] + + def _download_one(blob_name): + relative_path = self.resolve_artifact_path(blob_name) + file_path = os.path.join( + os.path.abspath(dst_path), *PurePosixPath(relative_path).parts + ) os.makedirs(os.path.dirname(file_path), exist_ok=True) - blob_client = self.container_client.get_blob_client(blob.name) + blob_client = self.container_client.get_blob_client(blob_name) stream = blob_client.download_blob() - with open(file_path, "wb") as f: - for chunk in stream.chunks(): - f.write(chunk) + temp_path = None + try: + with tempfile.NamedTemporaryFile( + dir=os.path.dirname(file_path), delete=False + ) as temp_file: + temp_path = temp_file.name + for chunk in stream.chunks(): + temp_file.write(chunk) + os.replace(temp_path, file_path) + except Exception: + if temp_path and os.path.exists(temp_path): + os.unlink(temp_path) + raise + + if blob_names: + workers = configured_worker_count( + "HASTE_ARTIFACT_DOWNLOAD_WORKERS", 8 + ) + parallel_map(_download_one, blob_names, max_workers=workers) self.logger.info(f"Downloaded {src_path} to {dst_path}") return dst_path @@ -302,14 +332,22 @@ def resolve_artifact_path(self, location: str) -> str: if parsed.scheme: container_url = urlparse(self.container_client.url) if parsed.netloc.lower() != container_url.netloc.lower(): - raise ValueError("Artifact URL does not belong to configured storage") + raise ValueError( + "Artifact URL does not belong to configured storage" + ) container_path = container_url.path.rstrip("/") + "/" if not parsed.path.startswith(container_path): - raise ValueError("Artifact URL does not belong to configured container") + raise ValueError( + "Artifact URL does not belong to configured container" + ) location = unquote(parsed.path[len(container_path) :]) normalized = str(PurePosixPath(location.lstrip("/"))) - if not normalized or normalized == "." or ".." in PurePosixPath(normalized).parts: + if ( + not normalized + or normalized == "." + or ".." in PurePosixPath(normalized).parts + ): raise ValueError("Invalid artifact path") return normalized @@ -440,11 +478,8 @@ def get_scoped_download_url( ) account_key = None if account_key is None: - user_delegation_key = ( - delegation_client.get_user_delegation_key( - now - timedelta(minutes=5), - expiry + timedelta(minutes=5), - ) + user_delegation_key = get_cached_user_delegation_key( + delegation_client, now=now ) sas_token = generate_blob_sas( diff --git a/hastelib/src/hastegeo/core/data_layer/abstract_data_layer.py b/hastelib/src/hastegeo/core/data_layer/abstract_data_layer.py index b6f743fb..92ad312b 100644 --- a/hastelib/src/hastegeo/core/data_layer/abstract_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/abstract_data_layer.py @@ -184,14 +184,57 @@ def load_all_from_partition(self, data_type, data_format="json"): """ pass - def load_bounded( - self, data_type, max_records, data_format="json" - ): + def load_bounded(self, data_type, max_records, data_format="json"): """Load no more than ``max_records`` or fail before full materialization.""" raise NotImplementedError( f"{self.__class__.__name__} does not support bounded reads" ) + def load_map( + self, + identifiers, + data_type, + data_format="json", + max_workers=None, + ): + """Load keyed records in one backend-native operation when supported.""" + raise NotImplementedError( + f"{self.__class__.__name__} does not implement batch reads" + ) + + def list_identifiers(self, data_type, data_format="json"): + """List the identifiers of records of a type in the current partition. + + Unlike :meth:`load_all_from_partition`, this returns only the keys and + does not download record contents — a cheap way to test existence in + bulk. Not abstract so existing backends keep working; concrete backends + that support cheap listing (blob, local filesystem) override it. + + Args: + data_type (str): Type/category of the data to list. + data_format (str, optional): File format of the records. Defaults + to "json". + + Returns: + List[str]: Identifiers present in the current partition. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not implement list_identifiers" + ) + + def get_file_remote_path( + self, + identifier=None, + data_type=None, + data_format="json", + extra_partition_keys=None, + check_exists=True, + ): + """Build a remotely accessible path when the backend supports one.""" + raise NotImplementedError( + f"{self.__class__.__name__} does not expose remote file paths" + ) + @abstractmethod def delete(self, identifier, data_type, data_format="json"): """Delete a specific data record from the storage backend. diff --git a/hastelib/src/hastegeo/core/data_layer/azure_blob_storage_data_layer.py b/hastelib/src/hastegeo/core/data_layer/azure_blob_storage_data_layer.py index 582bcc2c..4876608b 100644 --- a/hastelib/src/hastegeo/core/data_layer/azure_blob_storage_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/azure_blob_storage_data_layer.py @@ -4,21 +4,21 @@ import json import logging import os -from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone from threading import Lock import yaml from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError -from azure.identity import DefaultAzureCredential # type: ignore from azure.storage.blob import AccessPolicy # type: ignore from azure.storage.blob import BlobBlock # type: ignore -from azure.storage.blob import ( - BlobServiceClient, - ContainerSasPermissions, - generate_container_sas, -) +from azure.storage.blob import ContainerSasPermissions, generate_container_sas +from ..utils.blob import ( + get_blob_service_client, + get_cached_user_delegation_key, +) +from ..utils.metadata import matches_metadata_type +from ..utils.parallel import parallel_map from .abstract_data_layer import AbstractDataLayer _INITIALIZED_CONTAINERS = set() @@ -36,23 +36,16 @@ def __init__( ): super().__init__(partition_key) if connection_string: - credential = connection_string - self.blob_service_client = ( - BlobServiceClient.from_connection_string(connection_string) + self.blob_service_client = get_blob_service_client( + connection_string=connection_string ) self.user_delegation_key = None self.account_key = self.blob_service_client.credential.account_key else: - credential = DefaultAzureCredential() - self.blob_service_client = BlobServiceClient( - account_url=account_url, credential=credential - ) - self.user_delegation_key = ( - self.blob_service_client.get_user_delegation_key( - datetime.now(timezone.utc), - datetime.now(timezone.utc) + timedelta(hours=1), - ) + self.blob_service_client = get_blob_service_client( + account_url=account_url ) + self.user_delegation_key = None self.account_key = None self.container_read_policy = container_read_policy_name @@ -141,6 +134,7 @@ def get_file_remote_path( data_type=None, data_format="json", extra_partition_keys=None, + check_exists=True, ): blob_name = self.get_file_path( identifier, @@ -150,13 +144,22 @@ def get_file_remote_path( ) blob_client = self.container_client.get_blob_client(blob_name) - if blob_client.exists(): + # ``check_exists=False`` skips the per-blob existence HEAD request. Use + # it when the caller has already confirmed presence in bulk (e.g. via + # ``list_identifiers``) — the container read SAS is identical for every + # blob, so building the URL is a purely local operation. + if not check_exists or blob_client.exists(): + user_delegation_key = self.user_delegation_key + if self.account_key is None: + user_delegation_key = get_cached_user_delegation_key( + self.blob_service_client + ) # Generate SAS token with the policy sas_token = generate_container_sas( account_name=self.container_client.account_name, container_name=self.container_client.container_name, policy_id=self.container_read_policy, - user_delegation_key=self.user_delegation_key, + user_delegation_key=user_delegation_key, account_key=self.account_key, ) @@ -276,67 +279,64 @@ def load(self, identifier, data_type, data_format="json"): return contents elif data_format == "yaml": return yaml.safe_load(downloader.readall()) - except Exception as e: + raise ValueError(f"Unsupported data_format: {data_format}") + except ResourceNotFoundError as error: raise FileNotFoundError( f"{self.__class__.__name__}.load: No data found for identifier: {identifier} and data_type: {data_type}" - ) from e + ) from error def load_all(self, data_type, data_format="json"): - data = [] + # Collect the matching blobs first (a cheap listing pass), then download + # + deserialize them concurrently. + matched = [] + + def matches(blob_name): + in_partition = not self.partition_key or blob_name.startswith( + f"{self.partition_key}/" + ) + return in_partition and matches_metadata_type(blob_name, data_type) + blobs = self.container_client.walk_blobs() for blob in blobs: - logging.info(f"Blob name: {blob.name}") # Ignore stats file if "stats" in blob.name: continue # Check if the blob is a directory if blob.name.endswith("/"): - logging.info(f"Blob is a directory: {blob.name}") sub_blobs = self.container_client.walk_blobs( name_starts_with=blob.name ) for sub_blob in sub_blobs: - logging.info(f"SubBlob name: {sub_blob.name}") - if ( - sub_blob.name.startswith( - f"{self.partition_key}/{data_type}_" - ) - if self.partition_key - else sub_blob.name.startswith( - f"{blob.name}{data_type}_" - ) - ): - sub_blob_client = ( - self.container_client.get_blob_client(sub_blob) - ) - downloader = sub_blob_client.download_blob() - if data_format == "json": - contents = json.loads(downloader.readall()) - if isinstance(contents, str): - # Need to do this conversion again. TODO: Investigate why image_layers needs this converted twice - # but models does not - contents = json.loads(contents) - data.append(contents) - elif data_format == "yaml": - data.append(yaml.safe_load(downloader.readall())) + if matches(sub_blob.name): + matched.append(sub_blob) else: - if ( - blob.name.startswith(f"{self.partition_key}/{data_type}_") - if self.partition_key - else blob.name.startswith(f"{data_type}_") - ): - blob_client = self.container_client.get_blob_client(blob) - downloader = blob_client.download_blob() - if data_format == "json": - contents = json.loads(downloader.readall()) - if isinstance(contents, str): - # Need to do this conversion again. TODO: Investigate why image_layers needs this converted twice - # but models does not - contents = json.loads(contents) - data.append(contents) - elif data_format == "yaml": - data.append(yaml.safe_load(downloader.readall())) - return data + if matches(blob.name): + matched.append(blob) + return self._read_blobs_parallel(matched, data_format) + + def _read_blob_content(self, blob, data_format): + """Download and deserialize a single blob's content. + + Tolerates records that were double-serialized on save (parse again when + the first parse yields a string) so both legacy and current blobs read + correctly. + """ + blob_client = self.container_client.get_blob_client(blob) + raw = blob_client.download_blob().readall() + if data_format == "json": + contents = json.loads(raw) + if isinstance(contents, str): + contents = json.loads(contents) + return contents + elif data_format == "yaml": + return yaml.safe_load(raw) + raise ValueError(f"Unsupported data_format: {data_format}") + + def _read_blobs_parallel(self, blobs, data_format): + """Download+parse a list of blobs concurrently, preserving order.""" + return parallel_map( + lambda blob: self._read_blob_content(blob, data_format), blobs + ) def load_all_from_partition(self, data_type, data_format="json"): if not self.partition_key: @@ -344,23 +344,36 @@ def load_all_from_partition(self, data_type, data_format="json"): f"{self.__class__.__name__}.load_all_from_partition: Partition key is not set." ) - data = [] blobs = self.container_client.walk_blobs( name_starts_with=f"{self.partition_key}/{data_type}_" ) - for blob in blobs: - blob_client = self.container_client.get_blob_client(blob) - downloader = blob_client.download_blob() - if data_format == "json": - contents = json.loads(downloader.readall()) - if isinstance(contents, str): - # Need to do this conversion again. TODO: Investigate why image_layers needs this converted twice - # but models does not - contents = json.loads(contents) - data.append(contents) - elif data_format == "yaml": - data.append(yaml.safe_load(downloader.readall())) - return data + matching_blobs = ( + blob + for blob in blobs + if matches_metadata_type(blob.name, data_type) + ) + return self._read_blobs_parallel(matching_blobs, data_format) + + def load_map( + self, + identifiers, + data_type, + data_format="json", + max_workers=None, + ): + identifiers = list(dict.fromkeys(identifiers)) + + def load_one(identifier): + try: + return identifier, self.load( + identifier, data_type, data_format=data_format + ) + except FileNotFoundError: + return identifier, None + + return dict( + parallel_map(load_one, identifiers, max_workers=max_workers) + ) def load_bounded(self, data_type, max_records, data_format="json"): records, _ = self.load_page( @@ -415,12 +428,7 @@ def load_page( parts = blob.name.split("/") if "stats" in blob.name or len(parts) > 2: continue - if self.partition_key: - matches = blob.name.startswith( - f"{self.partition_key}/{data_type}_" - ) - else: - matches = parts[-1].startswith(f"{data_type}_") + matches = matches_metadata_type(blob.name, data_type) if not matches or not blob.name.endswith(f".{data_format}"): continue catalog_record_count += 1 @@ -471,6 +479,8 @@ def _index_metadata(data_type, data): } def _load_blob_names(self, blob_names, data_format): + """Like :meth:`_read_blobs_parallel`, but tolerates blobs deleted + between listing and download (they are dropped from the result).""" if not blob_names: return [] @@ -478,26 +488,30 @@ def _load_blob_names(self, blob_names, data_format): def load_blob(blob_name): try: - downloader = self.container_client.get_blob_client( - blob_name - ).download_blob() + return self._read_blob_content(blob_name, data_format) except ResourceNotFoundError: return missing_blob - contents = downloader.readall() - if data_format == "json": - contents = json.loads(contents) - if isinstance(contents, str): - contents = json.loads(contents) - return contents - if data_format == "yaml": - return yaml.safe_load(contents) - raise ValueError(f"Unsupported data format: {data_format}") - workers = min(32, len(blob_names)) - with ThreadPoolExecutor(max_workers=workers) as executor: - records = list(executor.map(load_blob, blob_names)) + records = parallel_map(load_blob, blob_names) return [record for record in records if record is not missing_blob] + def list_identifiers(self, data_type, data_format="json"): + if not self.partition_key: + raise ValueError( + f"{self.__class__.__name__}.list_identifiers: Partition key is not set." + ) + prefix = f"{self.partition_key}/{data_type}_" + suffix = f".{data_format}" + identifiers = [] + for name in self.container_client.list_blob_names( + name_starts_with=prefix + ): + if name.endswith(suffix) and matches_metadata_type( + name, data_type + ): + identifiers.append(name[len(prefix) : -len(suffix)]) + return identifiers + def delete(self, identifier, data_type, data_format="json"): blob_name = self.get_file_path(identifier, data_type, data_format) blob_client = self.container_client.get_blob_client(blob_name) diff --git a/hastelib/src/hastegeo/core/data_layer/azure_cosmos_db_data_layer.py b/hastelib/src/hastegeo/core/data_layer/azure_cosmos_db_data_layer.py index 29aeba16..4c144216 100644 --- a/hastelib/src/hastegeo/core/data_layer/azure_cosmos_db_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/azure_cosmos_db_data_layer.py @@ -5,6 +5,7 @@ from azure.cosmos import CosmosClient, exceptions # type: ignore from azure.identity import DefaultAzureCredential # type: ignore +from ..utils.metadata import matches_metadata_type from .abstract_data_layer import AbstractDataLayer @@ -77,7 +78,9 @@ def finalize_save( "Method not implemented and supported for Azure Cosmos DB." ) - def load(self, identifier, data_type): + def load(self, identifier, data_type, data_format="json"): + if data_format != "json": + raise ValueError("Cosmos DB metadata supports only json") partition_key = ( self.partition_key if self.partition_key else identifier ) @@ -91,7 +94,9 @@ def load(self, identifier, data_type): f"No data found for identifier: {identifier} and data_type: {data_type}" ) - def load_all(self, data_type): + def load_all(self, data_type, data_format="json"): + if data_format != "json": + raise ValueError("Cosmos DB metadata supports only json") id_prefix = self._id_prefix(data_type) query = "SELECT * FROM c WHERE STARTSWITH(c.id, @id_prefix)" items = list( @@ -101,9 +106,15 @@ def load_all(self, data_type): enable_cross_partition_query=True, ) ) - return items + return [ + item + for item in items + if matches_metadata_type(item["id"], data_type) + ] - def load_all_from_partition(self, data_type): + def load_all_from_partition(self, data_type, data_format="json"): + if data_format != "json": + raise ValueError("Cosmos DB metadata supports only json") id_prefix = self._id_prefix(data_type) query = ( "SELECT * FROM c WHERE c.partition_key = @partition_key " @@ -123,7 +134,67 @@ def load_all_from_partition(self, data_type): partition_key=self.partition_key, ) ) - return items + return [ + item + for item in items + if matches_metadata_type(item["id"], data_type) + ] + + def list_identifiers(self, data_type, data_format="json"): + if data_format != "json": + return [] + id_prefix = self._id_prefix(data_type) + query = ( + "SELECT VALUE c.id FROM c WHERE c.partition_key = @partition_key " + "AND STARTSWITH(c.id, @id_prefix)" + ) + item_ids = self.container.query_items( + query=query, + parameters=[ + {"name": "@partition_key", "value": self.partition_key}, + {"name": "@id_prefix", "value": id_prefix}, + ], + enable_cross_partition_query=False, + partition_key=self.partition_key, + ) + return [ + item_id[len(id_prefix) :] + for item_id in item_ids + if matches_metadata_type(item_id, data_type) + ] + + def load_map( + self, + identifiers, + data_type, + data_format="json", + max_workers=None, + ): + if data_format != "json": + raise ValueError("Cosmos DB metadata supports only json") + identifiers = list(dict.fromkeys(identifiers)) + if not identifiers: + return {} + id_prefix = self._id_prefix(data_type) + item_ids = [f"{id_prefix}{identifier}" for identifier in identifiers] + query = ( + "SELECT * FROM c WHERE c.partition_key = @partition_key " + "AND ARRAY_CONTAINS(@item_ids, c.id)" + ) + items = self.container.query_items( + query=query, + parameters=[ + {"name": "@partition_key", "value": self.partition_key}, + {"name": "@item_ids", "value": item_ids}, + ], + enable_cross_partition_query=False, + partition_key=self.partition_key, + ) + by_identifier = {item["id"][len(id_prefix) :]: item for item in items} + return { + identifier: by_identifier.get(identifier) + for identifier in identifiers + } def load_bounded(self, data_type, max_records, data_format="json"): if ( diff --git a/hastelib/src/hastegeo/core/data_layer/azure_data_lake_data_layer.py b/hastelib/src/hastegeo/core/data_layer/azure_data_lake_data_layer.py index 5cf4cfa7..1ee5e2e3 100644 --- a/hastelib/src/hastegeo/core/data_layer/azure_data_lake_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/azure_data_lake_data_layer.py @@ -5,6 +5,7 @@ from azure.identity import DefaultAzureCredential # type: ignore from azure.storage.filedatalake import DataLakeServiceClient # type: ignore +from ..utils.metadata import matches_metadata_type from .abstract_data_layer import AbstractDataLayer @@ -49,11 +50,14 @@ def get_file_remote_path( data_type=None, data_format="json", extra_partition_keys=None, + check_exists=True, ): file_name = self.get_file_path( identifier, data_type, data_format, extra_partition_keys ) file_client = self.file_system_client.get_file_client(file_name) + if check_exists and not file_client.exists(): + return None sas_url = file_client.url return str(sas_url) @@ -111,22 +115,25 @@ def finalize_save( def update(self, data, identifier, data_type): self.save(data, identifier, data_type) - def load(self, identifier, data_type): - file_name = self.get_file_path(identifier, data_type) + def load(self, identifier, data_type, data_format="json"): + if data_format != "json": + raise ValueError("Data Lake metadata reads support only json") + file_name = self.get_file_path(identifier, data_type, data_format) file_client = self.file_system_client.get_file_client(file_name) download = file_client.download_file() file_contents = download.readall() return json.loads(file_contents) - def load_all(self, data_type): + def load_all(self, data_type, data_format="json"): + if data_format != "json": + raise ValueError("Data Lake metadata reads support only json") data = [] paths = self.file_system_client.get_paths() for path in paths: - if ( - path.name.startswith(f"{self.partition_key}/{data_type}_") - if self.partition_key - else path.name.startswith(f"{data_type}_") - ): + in_partition = not self.partition_key or path.name.startswith( + f"{self.partition_key}/" + ) + if in_partition and matches_metadata_type(path.name, data_type): file_client = self.file_system_client.get_file_client( path.name ) @@ -135,10 +142,23 @@ def load_all(self, data_type): data.append(json.loads(file_contents)) return data - def load_all_from_partition(self, data_type): - data = self.load_all(data_type) + def load_all_from_partition(self, data_type, data_format="json"): + data = self.load_all(data_type, data_format=data_format) return data + def list_identifiers(self, data_type, data_format="json"): + prefix = f"{self.partition_key}/{data_type}_" + suffix = f".{data_format}" + identifiers = [] + for path in self.file_system_client.get_paths(path=self.partition_key): + if ( + path.name.startswith(prefix) + and path.name.endswith(suffix) + and matches_metadata_type(path.name, data_type) + ): + identifiers.append(path.name[len(prefix) : -len(suffix)]) + return identifiers + def load_bounded(self, data_type, max_records, data_format="json"): if data_format != "json" or max_records < 1: raise ValueError("Invalid bounded Data Lake read") @@ -150,7 +170,9 @@ def load_bounded(self, data_type, max_records, data_format="json"): if scanned_paths > scan_limit: raise ValueError("Metadata scan exceeds the bounded envelope") parts = path.name.split("/") - if len(parts) > 2 or not parts[-1].startswith(f"{data_type}_"): + if len(parts) > 2 or not matches_metadata_type( + path.name, data_type + ): continue file_contents = ( self.file_system_client.get_file_client(path.name) diff --git a/hastelib/src/hastegeo/core/data_layer/azure_postgresql_data_layer.py b/hastelib/src/hastegeo/core/data_layer/azure_postgresql_data_layer.py index 9481dcd8..61ee3018 100644 --- a/hastelib/src/hastegeo/core/data_layer/azure_postgresql_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/azure_postgresql_data_layer.py @@ -55,6 +55,15 @@ def _build_table_identifier(table_name): def _table_identifier(self): return self._qualified_table_identifier + @staticmethod + def _require_json(data_format): + if data_format != "json": + raise ValueError("PostgreSQL metadata supports only json") + + @staticmethod + def _deserialize_json(value): + return value if isinstance(value, (dict, list)) else json.loads(value) + def _create_table_if_not_exists(self): connection_string = f"host={self.server_name} dbname={self.database_name} user={self.postgres_user} password={self.token} sslmode=require" with psycopg2.connect(connection_string) as connection: @@ -183,7 +192,8 @@ def finalize_save( def update(self, data, identifier, data_type): self.save(data, identifier, data_type) - def load(self, identifier, data_type): + def load(self, identifier, data_type, data_format="json"): + self._require_json(data_format) partition_key = ( self.partition_key if self.partition_key else identifier ) @@ -201,9 +211,10 @@ def load(self, identifier, data_type): raise FileNotFoundError( f"No data found for identifier: {identifier} and data_type: {data_type}" ) - return json.loads(result[0]) + return self._deserialize_json(result[0]) - def load_all(self, data_type): + def load_all(self, data_type, data_format="json"): + self._require_json(data_format) connection_string = f"host={self.server_name} dbname={self.database_name} user={self.postgres_user} password={self.token} sslmode=require" with psycopg2.connect(connection_string) as connection: with connection.cursor() as cursor: @@ -214,9 +225,12 @@ def load_all(self, data_type): (data_type,), ) results = cursor.fetchall() - return [json.loads(result[0]) for result in results] + return [ + self._deserialize_json(result[0]) for result in results + ] - def load_all_from_partition(self, data_type): + def load_all_from_partition(self, data_type, data_format="json"): + self._require_json(data_format) partition_key = self.partition_key connection_string = f"host={self.server_name} dbname={self.database_name} user={self.postgres_user} password={self.token} sslmode=require" with psycopg2.connect(connection_string) as connection: @@ -228,7 +242,51 @@ def load_all_from_partition(self, data_type): (data_type, partition_key), ) results = cursor.fetchall() - return [json.loads(result[0]) for result in results] + return [ + self._deserialize_json(result[0]) for result in results + ] + + def list_identifiers(self, data_type, data_format="json"): + if data_format != "json": + return [] + connection_string = f"host={self.server_name} dbname={self.database_name} user={self.postgres_user} password={self.token} sslmode=require" + with psycopg2.connect(connection_string) as connection: + with connection.cursor() as cursor: + cursor.execute( + sql.SQL( + "SELECT identifier FROM {} WHERE data_type = %s AND partition_key = %s" + ).format(self._table_identifier()), + (data_type, self.partition_key), + ) + return [result[0] for result in cursor.fetchall()] + + def load_map( + self, + identifiers, + data_type, + data_format="json", + max_workers=None, + ): + self._require_json(data_format) + identifiers = list(dict.fromkeys(identifiers)) + if not identifiers: + return {} + connection_string = f"host={self.server_name} dbname={self.database_name} user={self.postgres_user} password={self.token} sslmode=require" + with psycopg2.connect(connection_string) as connection: + with connection.cursor() as cursor: + cursor.execute( + sql.SQL( + "SELECT identifier, data FROM {} WHERE data_type = %s AND partition_key = %s AND identifier = ANY(%s)" + ).format(self._table_identifier()), + (data_type, self.partition_key, identifiers), + ) + records = { + identifier: self._deserialize_json(data) + for identifier, data in cursor.fetchall() + } + return { + identifier: records.get(identifier) for identifier in identifiers + } def load_bounded(self, data_type, max_records, data_format="json"): if data_format != "json" or max_records < 1: diff --git a/hastelib/src/hastegeo/core/data_layer/local_file_system_data_layer.py b/hastelib/src/hastegeo/core/data_layer/local_file_system_data_layer.py index a33ff2c0..d367f39c 100644 --- a/hastelib/src/hastegeo/core/data_layer/local_file_system_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/local_file_system_data_layer.py @@ -6,6 +6,7 @@ import yaml +from ..utils.metadata import matches_metadata_type from .abstract_data_layer import AbstractDataLayer @@ -90,6 +91,7 @@ def get_file_remote_path( data_type=None, data_format="json", extra_partition_keys=None, + check_exists=True, ): """Get the remote path for a file (same as local path for filesystem layer). @@ -271,9 +273,9 @@ def load(self, identifier, data_type, data_format="json"): def load_all(self, data_type, data_format="json"): data = [] for file_name in os.listdir(self.directory): - if file_name.startswith(f"{data_type}_") and file_name.endswith( - f".{data_format}" - ): + if matches_metadata_type( + file_name, data_type + ) and file_name.endswith(f".{data_format}"): with open( os.path.join(self.directory, file_name), "r" ) as file: @@ -318,9 +320,9 @@ def load_bounded(self, data_type, max_records, data_format="json"): ) if not entry.is_file(follow_symlinks=False): continue - if not entry.name.startswith(f"{data_type}_") or not entry.name.endswith( - f".{data_format}" - ): + if not matches_metadata_type( + entry.name, data_type + ) or not entry.name.endswith(f".{data_format}"): continue with open(entry.path, "r") as file: records.append( @@ -334,6 +336,20 @@ def load_bounded(self, data_type, max_records, data_format="json"): ) return records + def list_identifiers(self, data_type, data_format="json"): + prefix = f"{data_type}_" + suffix = f".{data_format}" + identifiers = [] + if os.path.exists(self.directory): + for file_name in os.listdir(self.directory): + if ( + file_name.startswith(prefix) + and file_name.endswith(suffix) + and matches_metadata_type(file_name, data_type) + ): + identifiers.append(file_name[len(prefix) : -len(suffix)]) + return identifiers + def delete(self, identifier, data_type, data_format="json"): file_path = self.get_file_path(identifier, data_type, data_format) if not os.path.exists(file_path): diff --git a/hastelib/src/hastegeo/core/data_layer/unified.py b/hastelib/src/hastegeo/core/data_layer/unified.py index 20e35fe1..cc88258b 100644 --- a/hastelib/src/hastegeo/core/data_layer/unified.py +++ b/hastelib/src/hastegeo/core/data_layer/unified.py @@ -3,6 +3,7 @@ import importlib from ..utils.metadata import MetadataUtils +from .abstract_data_layer import AbstractDataLayer class UnifiedDataLayer: @@ -35,9 +36,7 @@ def __init__(self, storage_type, partition_key=None, **kwargs): if storage_type in storage_class_map: module_name, class_name = storage_class_map[storage_type] - module = importlib.import_module( - f"{__package__}.{module_name}" - ) + module = importlib.import_module(f"{__package__}.{module_name}") data_layer_class = getattr(module, class_name) self.data_layer = data_layer_class( partition_key=self.partition_key, **kwargs @@ -132,6 +131,24 @@ def load_bounded(self, data_type, max_records, data_format="json"): data_format=data_format, ) + def supports_load_map(self): + method = type(self.data_layer).load_map + return method is not AbstractDataLayer.load_map + + def load_map( + self, + identifiers, + data_type, + data_format="json", + max_workers=None, + ): + return self.data_layer.load_map( + identifiers=identifiers, + data_type=data_type, + data_format=data_format, + max_workers=max_workers, + ) + def load_page( self, data_type, @@ -154,6 +171,11 @@ def load_page( max_records=max_records, ) + def list_identifiers(self, data_type, data_format="json"): + return self.data_layer.list_identifiers( + data_type, data_format=data_format + ) + def delete(self, identifier, data_type, data_format="json"): self.data_layer.delete(identifier, data_type, data_format=data_format) @@ -183,12 +205,14 @@ def get_file_remote_path( data_type=None, data_format="json", extra_partition_keys=None, + check_exists=True, ): return self.data_layer.get_file_remote_path( identifier, data_type, data_format=data_format, extra_partition_keys=extra_partition_keys, + check_exists=check_exists, ) def get_base_url(self): diff --git a/hastelib/src/hastegeo/core/processors/metadata.py b/hastelib/src/hastegeo/core/processors/metadata.py index 06ba67d4..87eee3be 100644 --- a/hastelib/src/hastegeo/core/processors/metadata.py +++ b/hastelib/src/hastegeo/core/processors/metadata.py @@ -5,6 +5,11 @@ from hastegeo.core.config import Config from ..data_layer.unified import UnifiedDataLayer +from ..utils.parallel import ( + configured_worker_count, + parallel_map, + validate_worker_count, +) from ..utils.perf import timed @@ -161,6 +166,98 @@ def load_page( max_records=max_records, ) + def load_map( + self, + keys: list[str], + data_format: str = "json", + max_workers: int | None = None, + ) -> dict[str, dict | None]: + """Load many records by key concurrently. + + Returns ``{key: record}``; a key whose record is missing maps to + ``None`` (mirrors a per-key ``load`` that raises ``FileNotFoundError``). + Bounded parallelism overlaps per-key storage operations while the + process-wide executor caps aggregate concurrency. Preserves perf + instrumentation across worker threads by binding the active counter + inside each task. + """ + from ..utils.perf import bind, get_counter + + keys = list(dict.fromkeys(keys)) + if not keys: + return {} + workers = ( + configured_worker_count("HASTE_METADATA_LOAD_WORKERS", 8) + if max_workers is None + else validate_worker_count(max_workers) + ) + if self.storage.supports_load_map(): + with timed("load_map"): + return self.storage.load_map( + identifiers=keys, + data_type=self.data_type, + data_format=data_format, + max_workers=workers, + ) + counter = get_counter() + + def _one(key): + with bind(counter): + try: + return key, self.load(key, data_format=data_format) + except FileNotFoundError: + return key, None + + return dict(parallel_map(_one, keys, max_workers=workers)) + + def load_filtered( + self, predicate: dict[str, object], data_format: str = "json" + ) -> list[dict]: + """Load partition records matching every key/value in ``predicate``. + + This is an explicit client-side fallback: the partition is loaded once + and then filtered in process. Backends need a separate query primitive + before this can reduce transferred records. + """ + if not isinstance(predicate, dict) or not predicate: + raise ValueError("predicate must be a non-empty dictionary") + records = self.load_all_from_partition(data_format=data_format) + return [ + record + for record in records + if isinstance(record, dict) + and all( + key in record and record[key] == value + for key, value in predicate.items() + ) + ] + + def list_keys(self, data_format="json"): + """List the identifiers present for this data type in the partition. + + A cheap, metadata-only alternative to ``load_all_from_partition`` for + bulk existence checks (does not download record contents). + """ + with timed("list_keys"): + return self.storage.list_identifiers( + data_type=self.data_type, data_format=data_format + ) + + def build_url(self, key, data_format="json"): + """Build a remote URL for a record without a per-item existence check. + + Intended for callers that have already confirmed the record exists in + bulk via :meth:`list_keys`. On the blob backend this avoids a network + round-trip per key (the container read SAS is shared), so it is a local + operation and is not counted as a storage round-trip. + """ + return self.storage.get_file_remote_path( + identifier=key, + data_type=self.data_type, + data_format=data_format, + check_exists=False, + ) + def load_and_combine_sub_data_types(self, key, data_types): """ Load and combine metadata from multiple data types. diff --git a/hastelib/src/hastegeo/core/utils/blob.py b/hastelib/src/hastegeo/core/utils/blob.py index d2f4affc..c43ccca5 100644 --- a/hastelib/src/hastegeo/core/utils/blob.py +++ b/hastelib/src/hastegeo/core/utils/blob.py @@ -20,12 +20,77 @@ from __future__ import annotations import asyncio +import functools import os import re import tempfile +import threading +from collections import OrderedDict +from datetime import datetime, timedelta, timezone from typing import NamedTuple, Optional, Tuple from urllib.parse import urlparse +_USER_DELEGATION_KEYS = OrderedDict() +_USER_DELEGATION_KEY_CACHE_SIZE = 8 +_USER_DELEGATION_KEYS_LOCK = threading.Lock() + + +@functools.lru_cache(maxsize=8) +def get_blob_service_client( + connection_string: str | None = None, + account_url: str | None = None, +): + """Return a process-wide ``BlobServiceClient`` for one storage account. + + Azure SDK clients are thread-safe and designed for reuse; creating one per + call re-parses credentials and re-establishes the connection pool. Caching + by credential target keeps that setup cost one-time per process. + """ + from azure.storage.blob import BlobServiceClient + + if connection_string: + return BlobServiceClient.from_connection_string(connection_string) + if account_url: + from azure.identity import DefaultAzureCredential + + return BlobServiceClient( + account_url=account_url, credential=DefaultAzureCredential() + ) + raise ValueError("A connection string or account URL is required") + + +def get_cached_user_delegation_key( + blob_service_client, + now: datetime | None = None, +): + """Return a reusable user-delegation key with a safe refresh margin.""" + current_time = now or datetime.now(timezone.utc) + cache_key = blob_service_client + with _USER_DELEGATION_KEYS_LOCK: + cached = _USER_DELEGATION_KEYS.get(cache_key) + if cached is not None: + delegation_key, expires_at = cached + if expires_at > current_time + timedelta(minutes=15): + _USER_DELEGATION_KEYS.move_to_end(cache_key) + return delegation_key + + expires_at = current_time + timedelta(hours=2) + delegation_key = blob_service_client.get_user_delegation_key( + current_time - timedelta(minutes=5), expires_at + ) + _USER_DELEGATION_KEYS[cache_key] = (delegation_key, expires_at) + _USER_DELEGATION_KEYS.move_to_end(cache_key) + while len(_USER_DELEGATION_KEYS) > _USER_DELEGATION_KEY_CACHE_SIZE: + _USER_DELEGATION_KEYS.popitem(last=False) + return delegation_key + + +def clear_blob_client_caches() -> None: + """Clear cached Blob clients and delegation keys for isolated tests.""" + get_blob_service_client.cache_clear() + with _USER_DELEGATION_KEYS_LOCK: + _USER_DELEGATION_KEYS.clear() + def split_blob_url(url: str) -> Tuple[str, str]: """Extract ``(container_name, blob_name)`` from a blob URL. @@ -99,14 +164,9 @@ async def download_blob_to_tempfile( caller is responsible for unlinking the returned path when done — use ``try/finally``. """ - # Imported here so this module stays cheap to import for callers that - # only need split_blob_url(): azure-storage-blob brings in tens of - # transitive imports. - from azure.storage.blob import BlobServiceClient - conn_str = os.environ.get("BLOB_CONNECTION_STRING", "") container_name, blob_name = split_blob_url(url) - bsc = BlobServiceClient.from_connection_string(conn_str) + bsc = get_blob_service_client(connection_string=conn_str) if max_bytes is not None and max_bytes < 1: raise ValueError("max_bytes must be positive") @@ -128,10 +188,7 @@ def download() -> str: downloaded_bytes = 0 for chunk in blob_client.download_blob().chunks(): downloaded_bytes += len(chunk) - if ( - max_bytes is not None - and downloaded_bytes > max_bytes - ): + if max_bytes is not None and downloaded_bytes > max_bytes: raise ValueError( "Blob exceeds the allowed download size" ) @@ -201,13 +258,11 @@ async def read_blob_range( reads to EOF. ``data`` is clamped to the blob size; an ``offset`` at or past EOF yields empty ``data`` (callers should answer ``416``). """ - from azure.storage.blob import BlobServiceClient - conn_str = os.environ.get("BLOB_CONNECTION_STRING", "") container_name, blob_name = split_blob_url(url) def _read() -> BlobRange: - bsc = BlobServiceClient.from_connection_string(conn_str) + bsc = get_blob_service_client(connection_string=conn_str) blob_client = bsc.get_container_client(container_name).get_blob_client( blob_name ) diff --git a/hastelib/src/hastegeo/core/utils/metadata.py b/hastelib/src/hastegeo/core/utils/metadata.py index 25e5570d..89804921 100644 --- a/hastelib/src/hastegeo/core/utils/metadata.py +++ b/hastelib/src/hastegeo/core/utils/metadata.py @@ -4,6 +4,42 @@ import random import uuid from datetime import datetime, timezone +from functools import lru_cache + + +@lru_cache(maxsize=1) +def _known_metadata_types() -> tuple[str, ...]: + from ..config import Config + + return tuple( + sorted( + ( + metadata_type.value + for metadata_type in Config.get_metadata_types() + ), + key=len, + reverse=True, + ) + ) + + +def matches_metadata_type(path: str, data_type: str) -> bool: + """Return whether a stored name belongs to the requested metadata type. + + Existing records use ``{type}_{identifier}``, while some type names are + prefixes of others (notably ``model`` and ``model_catalog``). Assigning a + name to the longest known matching type preserves the existing layout + without allowing broader scans to consume a narrower type. + """ + name = path.rsplit("/", 1)[-1] + matching_types = [ + known_type + for known_type in _known_metadata_types() + if name.startswith(f"{known_type}_") + ] + if not matching_types: + return name.startswith(f"{data_type}_") + return matching_types[0] == data_type class MetadataUtils: diff --git a/hastelib/src/hastegeo/core/utils/parallel.py b/hastelib/src/hastegeo/core/utils/parallel.py new file mode 100644 index 00000000..aa2b9ddc --- /dev/null +++ b/hastelib/src/hastegeo/core/utils/parallel.py @@ -0,0 +1,124 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Process-wide bounded execution for blocking I/O.""" + +import os +from collections.abc import Callable, Iterable +from concurrent.futures import ( + FIRST_COMPLETED, + Future, + ThreadPoolExecutor, + wait, +) +from threading import current_thread +from typing import TypeVar + +_MAX_CONFIGURED_WORKERS = 64 + +InputT = TypeVar("InputT") +OutputT = TypeVar("OutputT") + + +def validate_worker_count(value: int, name: str = "max_workers") -> int: + """Validate a worker count before constructing or scheduling a pool.""" + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{name} must be an integer") + if not 1 <= value <= _MAX_CONFIGURED_WORKERS: + raise ValueError( + f"{name} must be between 1 and {_MAX_CONFIGURED_WORKERS}" + ) + return value + + +def configured_worker_count(name: str, default: int) -> int: + """Read and validate a worker count from the process environment.""" + raw_value = os.environ.get(name) + if raw_value is None: + return validate_worker_count(default, name) + try: + value = int(raw_value) + except ValueError as error: + raise ValueError(f"{name} must be an integer") from error + return validate_worker_count(value, name) + + +class BoundedExecutor: + """Share one thread budget across all concurrent map operations.""" + + def __init__(self, max_workers: int) -> None: + self.max_workers = validate_worker_count(max_workers) + self._thread_prefix = f"haste-io-{id(self):x}" + self._executor = ThreadPoolExecutor( + max_workers=self.max_workers, + thread_name_prefix=self._thread_prefix, + ) + + def map( + self, + function: Callable[[InputT], OutputT], + values: Iterable[InputT], + max_workers: int | None = None, + ) -> list[OutputT]: + """Run an ordered map with bounded submissions and shared workers.""" + items = list(values) + if not items: + return [] + + requested_workers = ( + self.max_workers + if max_workers is None + else validate_worker_count(max_workers) + ) + worker_count = min(requested_workers, self.max_workers, len(items)) + if worker_count == 1 or current_thread().name.startswith( + self._thread_prefix + ): + return [function(item) for item in items] + + indexed_items = iter(enumerate(items)) + pending: dict[Future[OutputT], int] = {} + results: dict[int, OutputT] = {} + + def submit_next() -> bool: + try: + index, item = next(indexed_items) + except StopIteration: + return False + pending[self._executor.submit(function, item)] = index + return True + + for _ in range(worker_count): + submit_next() + + try: + while pending: + done, _ = wait(pending, return_when=FIRST_COMPLETED) + for future in done: + index = pending.pop(future) + results[index] = future.result() + for _ in done: + submit_next() + except Exception: + for future in pending: + future.cancel() + raise + + return [results[index] for index in range(len(items))] + + def shutdown(self) -> None: + """Release worker threads after a non-global executor is finished.""" + self._executor.shutdown(wait=True, cancel_futures=True) + + +PARALLEL_IO_EXECUTOR = BoundedExecutor( + configured_worker_count("HASTE_BLOB_DOWNLOAD_WORKERS", 16) +) + + +def parallel_map( + function: Callable[[InputT], OutputT], + values: Iterable[InputT], + max_workers: int | None = None, +) -> list[OutputT]: + """Map blocking I/O on the shared process-wide executor.""" + return PARALLEL_IO_EXECUTOR.map(function, values, max_workers=max_workers) diff --git a/hastelib/src/hastegeo/core/utils/perf.py b/hastelib/src/hastegeo/core/utils/perf.py index c454ef8a..e2456e51 100644 --- a/hastelib/src/hastegeo/core/utils/perf.py +++ b/hastelib/src/hastegeo/core/utils/perf.py @@ -2,9 +2,9 @@ # Licensed under the MIT License. """Lightweight, opt-in performance instrumentation. -Counts and times backend storage round-trips for a single logical request so we -can establish a baseline (Phase 0 of the perf-layer-loading spec) and later prove -the O(layers x models) -> O(1) improvement. +Counts and times logical data-layer operations for a single request. A bulk +operation can issue multiple backend SDK requests, so these values are not +storage transaction counts. Design notes: - A ``ContextVar`` holds a shared ``PerfCounter`` *object*. ``asyncio.to_thread`` @@ -26,7 +26,7 @@ class PerfCounter: - """Thread-safe accumulator of storage round-trip count and duration.""" + """Thread-safe accumulator of logical data-layer calls and duration.""" def __init__(self): self.calls = 0 @@ -65,6 +65,21 @@ def get_counter(): return _current.get() +@contextmanager +def bind(counter): + """Bind ``counter`` as the active counter for the current context. + + Used to propagate the active counter into worker threads (e.g. a + ``ThreadPoolExecutor``), which — unlike ``asyncio.to_thread`` — do not copy + the parent context. ``counter`` may be ``None`` (tracking disabled). + """ + token = _current.set(counter) + try: + yield + finally: + _current.reset(token) + + @contextmanager def timed(op): """Time an ``op`` and record it on the active counter, if any. @@ -83,18 +98,23 @@ def timed(op): def headers(counter, wall_start): - """Response headers exposing round-trip count/timing for benchmarking.""" + """Response headers exposing data-layer call timing for benchmarking.""" if counter is None: return {} storage_ms = counter.seconds * 1000.0 wall_ms = (time.perf_counter() - wall_start) * 1000.0 return { + "X-Haste-Data-Layer-Calls": str(counter.calls), + "X-Haste-Data-Layer-Ms": f"{storage_ms:.1f}", + # Keep the original names while benchmark consumers migrate. "X-Haste-Storage-Calls": str(counter.calls), "X-Haste-Storage-Ms": f"{storage_ms:.1f}", "X-Haste-Wall-Ms": f"{wall_ms:.1f}", "Server-Timing": ", ".join( [ - ";".join(["storage", "desc=storage", f"dur={storage_ms:.1f}"]), + ";".join( + ["data-layer", "desc=data-layer", f"dur={storage_ms:.1f}"] + ), ";".join(["wall", "desc=wall", f"dur={wall_ms:.1f}"]), ] ), @@ -112,7 +132,7 @@ def log_summary(logger, name, counter, wall_start, **fields): } extra = " ".join(f"{k}={v}" for k, v in fields.items()) logger.info( - "PERF %s %s storage_calls=%d storage_ms=%.1f wall_ms=%.1f ops=%s", + "PERF %s %s data_layer_calls=%d data_layer_ms=%.1f wall_ms=%.1f ops=%s", name, extra, counter.calls, diff --git a/hastelib/tests/core/artifact_storage/test_azure_blob_artifact_storage.py b/hastelib/tests/core/artifact_storage/test_azure_blob_artifact_storage.py new file mode 100644 index 00000000..beb23a2d --- /dev/null +++ b/hastelib/tests/core/artifact_storage/test_azure_blob_artifact_storage.py @@ -0,0 +1,238 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +import os +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from azure.core.exceptions import ResourceExistsError +from hastegeo.core.artifact_storage.azure_blob_artifact_storage import ( + _INITIALIZED_CONTAINERS, + AzureBlobArtifactStorage, +) + + +class TestAzureBlobArtifactStorageFetch(unittest.TestCase): + def setUp(self) -> None: + self.storage = AzureBlobArtifactStorage.__new__( + AzureBlobArtifactStorage + ) + self.storage.partition_key = None + self.storage.logger = Mock() + self.storage.container_client = Mock() + self.storage.container_client.url = "https://account.test/container" + self.blob_client = ( + self.storage.container_client.get_blob_client.return_value + ) + self.stream = self.blob_client.download_blob.return_value + + def _set_blob_names(self, *names: str) -> None: + self.storage.container_client.list_blobs.return_value = [ + SimpleNamespace(name=name) for name in names + ] + + def test_fetch_downloads_each_blob_atomically(self) -> None: + self._set_blob_names("project/output.txt") + self.stream.chunks.return_value = [b"hello", b" world"] + + with tempfile.TemporaryDirectory() as destination: + result = self.storage.fetch_artifact( + src_path="project", dst_path=destination + ) + + output_path = os.path.join(destination, "project", "output.txt") + with open(output_path, "rb") as output: + self.assertEqual(output.read(), b"hello world") + self.assertEqual(result, destination) + self.assertEqual( + os.listdir(os.path.dirname(output_path)), ["output.txt"] + ) + + def test_fetch_rejects_parent_path_in_blob_name(self) -> None: + self._set_blob_names("../outside.txt") + + with tempfile.TemporaryDirectory() as destination: + with self.assertRaisesRegex(ValueError, "Invalid artifact path"): + self.storage.fetch_artifact( + src_path="project", dst_path=destination + ) + self.assertFalse( + os.path.exists(os.path.join(destination, "..", "outside.txt")) + ) + + def test_fetch_removes_partial_file_when_download_fails(self) -> None: + self._set_blob_names("project/output.txt") + + def failing_chunks(): + yield b"partial" + raise RuntimeError("download failed") + + self.stream.chunks.side_effect = failing_chunks + with tempfile.TemporaryDirectory() as destination: + with self.assertRaisesRegex(RuntimeError, "download failed"): + self.storage.fetch_artifact( + src_path="project", dst_path=destination + ) + output_directory = os.path.join(destination, "project") + self.assertEqual(os.listdir(output_directory), []) + + def test_fetch_requires_source_and_destination(self) -> None: + with self.assertRaisesRegex(ValueError, "source"): + self.storage.fetch_artifact(dst_path="destination") + with self.assertRaisesRegex(ValueError, "destination"): + self.storage.fetch_artifact(src_path="source") + + def test_fetch_rejects_invalid_worker_configuration(self) -> None: + self._set_blob_names("project/output.txt") + with patch.dict(os.environ, {"HASTE_ARTIFACT_DOWNLOAD_WORKERS": "0"}): + with self.assertRaisesRegex(ValueError, "between 1 and 64"): + self.storage.fetch_artifact( + src_path="project", dst_path="destination" + ) + + @patch( + "hastegeo.core.artifact_storage.azure_blob_artifact_storage.generate_container_sas", + return_value="sas", + ) + @patch( + "hastegeo.core.artifact_storage.azure_blob_artifact_storage.get_cached_user_delegation_key", + return_value="delegation-key", + ) + def test_download_url_fetches_delegation_key_lazily( + self, delegation_key, _generate_sas + ) -> None: + self.storage.blob_service_client = Mock() + self.storage.account_key = None + self.storage.user_delegation_key = None + self.storage.container_read_policy = "policy" + self.storage.container_client.account_name = "account" + self.storage.container_client.container_name = "artifacts" + self.blob_client.url = "https://account.test/artifacts/file.txt" + + result = self.storage.get_download_url(identifier="file.txt") + + self.assertEqual(result, f"{self.blob_client.url}?sas") + delegation_key.assert_called_once_with( + self.storage.blob_service_client + ) + + @patch( + "hastegeo.core.artifact_storage.azure_blob_artifact_storage.generate_blob_sas", + return_value="sas", + ) + @patch( + "hastegeo.core.artifact_storage.azure_blob_artifact_storage.get_cached_user_delegation_key", + return_value="delegation-key", + ) + def test_scoped_url_reuses_delegation_key( + self, delegation_key, _generate_sas + ) -> None: + self.storage.blob_service_client = Mock() + self.storage.identity_blob_service_client = ( + self.storage.blob_service_client + ) + self.storage.account_key = None + self.storage.container_client.account_name = "account" + self.storage.container_client.container_name = "artifacts" + self.storage.container_client.url = "https://account.test/artifacts" + self.blob_client.url = "https://account.test/artifacts/file.txt" + self.blob_client.exists.return_value = True + + result = self.storage.get_scoped_download_url("file.txt") + + self.assertEqual(result, f"{self.blob_client.url}?sas") + delegation_key.assert_called_once() + + def test_resolve_artifact_path_rejects_other_account(self) -> None: + self.storage.container_client.url = "https://account.test/artifacts" + + with self.assertRaisesRegex(ValueError, "configured storage"): + self.storage.resolve_artifact_path( + "https://other.test/artifacts/file.txt" + ) + + def test_resolve_artifact_path_rejects_other_container(self) -> None: + self.storage.container_client.url = "https://account.test/artifacts" + + with self.assertRaisesRegex(ValueError, "configured container"): + self.storage.resolve_artifact_path( + "https://account.test/other/file.txt" + ) + + +class TestAzureBlobArtifactStorageClientReuse(unittest.TestCase): + def setUp(self) -> None: + _INITIALIZED_CONTAINERS.clear() + + def tearDown(self) -> None: + _INITIALIZED_CONTAINERS.clear() + + @patch( + "hastegeo.core.artifact_storage.azure_blob_artifact_storage.get_blob_service_client" + ) + @patch.object( + AzureBlobArtifactStorage, + "_create_or_update_managed_access_policy", + ) + def test_connection_string_uses_cached_client( + self, create_policy, factory + ): + service = Mock(url="https://account.test") + service.credential.account_key = "key" # pragma: allowlist secret + factory.return_value = service + + storage = AzureBlobArtifactStorage( + account_url="", + container="artifacts", + connection_string="connection", + ) + + self.assertIs(storage.blob_service_client, service) + factory.assert_called_once_with(connection_string="connection") + create_policy.assert_called_once_with() + + @patch( + "hastegeo.core.artifact_storage.azure_blob_artifact_storage.get_blob_service_client" + ) + def test_managed_identity_uses_cached_client(self, factory): + service = Mock(url="https://account.test") + factory.return_value = service + + storage = AzureBlobArtifactStorage( + account_url="https://account.test", + container="artifacts", + connection_string=None, + serves_read_sas=False, + ) + + self.assertIs(storage.blob_service_client, service) + factory.assert_called_once_with(account_url="https://account.test") + service.get_user_delegation_key.assert_not_called() + + @patch( + "hastegeo.core.artifact_storage.azure_blob_artifact_storage.get_blob_service_client" + ) + @patch.object( + AzureBlobArtifactStorage, + "_create_or_update_managed_access_policy", + ) + def test_existing_container_is_reused(self, create_policy, factory): + service = Mock(url="https://account.test") + service.credential.account_key = "key" # pragma: allowlist secret + container = service.get_container_client.return_value + container.create_container.side_effect = ResourceExistsError("exists") + factory.return_value = service + + storage = AzureBlobArtifactStorage( + account_url="", + container="artifacts", + connection_string="connection", + ) + + self.assertIs(storage.container_client, container) + create_policy.assert_called_once_with() + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/data_layer/test_azure_blob_storage_data_layer.py b/hastelib/tests/core/data_layer/test_azure_blob_storage_data_layer.py new file mode 100644 index 00000000..cfaac2e4 --- /dev/null +++ b/hastelib/tests/core/data_layer/test_azure_blob_storage_data_layer.py @@ -0,0 +1,298 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +import json +import unittest +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from azure.core.exceptions import ResourceNotFoundError +from hastegeo.core.data_layer.azure_blob_storage_data_layer import ( + _INITIALIZED_CONTAINERS, + AzureBlobStorageDataLayer, +) + + +class TestAzureBlobStorageDataLayerLoad(unittest.TestCase): + def setUp(self) -> None: + self.layer = AzureBlobStorageDataLayer.__new__( + AzureBlobStorageDataLayer + ) + self.layer.partition_key = "partition" + self.layer.container_client = Mock() + self.blob_client = ( + self.layer.container_client.get_blob_client.return_value + ) + self.downloader = self.blob_client.download_blob.return_value + + def test_load_deserializes_json(self) -> None: + self.downloader.readall.return_value = b'{"value": 1}' + + result = self.layer.load("record", "model") + + self.assertEqual(result, {"value": 1}) + + def test_load_tolerates_legacy_double_serialized_json(self) -> None: + self.downloader.readall.return_value = json.dumps( + json.dumps({"value": 1}) + ).encode() + + result = self.layer.load("record", "model") + + self.assertEqual(result, {"value": 1}) + + def test_load_maps_only_resource_not_found_to_file_not_found(self) -> None: + self.blob_client.download_blob.side_effect = ResourceNotFoundError( + "missing" + ) + + with self.assertRaises(FileNotFoundError): + self.layer.load("record", "model") + + def test_load_preserves_transport_errors(self) -> None: + self.blob_client.download_blob.side_effect = RuntimeError( + "transport unavailable" + ) + + with self.assertRaisesRegex(RuntimeError, "transport unavailable"): + self.layer.load("record", "model") + + def test_load_preserves_json_errors(self) -> None: + self.downloader.readall.return_value = b"{" + + with self.assertRaises(json.JSONDecodeError): + self.layer.load("record", "model") + + def test_load_rejects_unsupported_format(self) -> None: + with self.assertRaisesRegex(ValueError, "Unsupported data_format"): + self.layer.load("record", "model", data_format="xml") + + def test_parallel_read_preserves_blob_order(self) -> None: + self.layer._read_blob_content = Mock( + side_effect=lambda blob, data_format: f"{blob}:{data_format}" + ) + + result = self.layer._read_blobs_parallel(["b", "a"], "json") + + self.assertEqual(result, ["b:json", "a:json"]) + + def test_parallel_read_handles_empty_listing(self) -> None: + self.layer._read_blob_content = Mock() + + result = self.layer._read_blobs_parallel([], "json") + + self.assertEqual(result, []) + self.layer._read_blob_content.assert_not_called() + + def test_load_blob_names_drops_blobs_deleted_after_listing(self) -> None: + self.layer._read_blob_content = Mock( + side_effect=[{"value": 1}, ResourceNotFoundError("missing")] + ) + + result = self.layer._load_blob_names(["first", "missing"], "json") + + self.assertEqual(result, [{"value": 1}]) + + def test_partition_scan_excludes_longer_type(self) -> None: + self.layer.container_client.walk_blobs.return_value = [ + SimpleNamespace(name="partition/model_a.json"), + SimpleNamespace(name="partition/model_catalog_index.json"), + ] + self.layer._read_blob_content = Mock( + side_effect=lambda blob, _: blob.name + ) + + result = self.layer.load_all_from_partition("model") + + self.assertEqual(result, ["partition/model_a.json"]) + + def test_load_all_does_not_cross_configured_partition(self) -> None: + self.layer.container_client.walk_blobs.return_value = [ + SimpleNamespace(name="partition/model_a.json"), + SimpleNamespace(name="other/model_b.json"), + ] + self.layer._read_blob_content = Mock( + side_effect=lambda blob, _: blob.name + ) + + result = self.layer.load_all("model") + + self.assertEqual(result, ["partition/model_a.json"]) + + def test_load_all_reads_matching_blobs_under_directory_markers( + self, + ) -> None: + directory = SimpleNamespace(name="partition/") + nested = SimpleNamespace(name="partition/model_a.json") + self.layer.container_client.walk_blobs.side_effect = [ + [directory], + [nested], + ] + self.layer._read_blob_content = Mock( + side_effect=lambda blob, _: blob.name + ) + + result = self.layer.load_all("model") + + self.assertEqual(result, ["partition/model_a.json"]) + + def test_load_page_skips_stats_and_deep_paths(self) -> None: + blobs = [ + SimpleNamespace( + name="partition/model_stats.json", + metadata={}, + last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), + ), + SimpleNamespace( + name="partition/nested/model_a.json", + metadata={}, + last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), + ), + SimpleNamespace( + name="partition/model_a.json", + metadata={}, + last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), + ), + ] + pages = Mock() + pages.by_page.return_value = [blobs] + self.layer.container_client.list_blobs.return_value = pages + self.layer._load_blob_names = Mock(return_value=[{"modelId": "a"}]) + + records, count = self.layer.load_page("model", page=1, page_size=10) + + self.assertEqual(records, [{"modelId": "a"}]) + self.assertEqual(count, 1) + self.layer._load_blob_names.assert_called_once_with( + ["partition/model_a.json"], "json" + ) + + def test_identifier_listing_excludes_longer_type(self) -> None: + self.layer.container_client.list_blob_names.return_value = [ + "partition/model_a.json", + "partition/model_catalog_index.json", + ] + + result = self.layer.list_identifiers("model") + + self.assertEqual(result, ["a"]) + + @patch( + "hastegeo.core.data_layer.azure_blob_storage_data_layer.generate_container_sas", + return_value="sas", + ) + @patch( + "hastegeo.core.data_layer.azure_blob_storage_data_layer.get_cached_user_delegation_key", + return_value="delegation-key", + ) + def test_remote_path_fetches_delegation_key_lazily( + self, delegation_key, _generate_sas + ) -> None: + self.layer.blob_service_client = Mock() + self.layer.account_key = None + self.layer.user_delegation_key = None + self.layer.container_read_policy = "policy" + self.layer.container_client.account_name = "account" + self.layer.container_client.container_name = "metadata" + self.blob_client.url = "https://account.test/metadata/model_a.json" + + result = self.layer.get_file_remote_path( + "a", "model", check_exists=False + ) + + self.assertEqual(result, f"{self.blob_client.url}?sas") + delegation_key.assert_called_once_with(self.layer.blob_service_client) + + @patch( + "hastegeo.core.data_layer.azure_blob_storage_data_layer.generate_container_sas", + return_value="sas", + ) + @patch( + "hastegeo.core.data_layer.azure_blob_storage_data_layer.get_cached_user_delegation_key" + ) + def test_remote_path_with_account_key_skips_delegation_key( + self, delegation_key, _generate_sas + ) -> None: + self.layer.blob_service_client = Mock() + self.layer.account_key = "account-key" # pragma: allowlist secret + self.layer.user_delegation_key = None + self.layer.container_read_policy = "policy" + self.layer.container_client.account_name = "account" + self.layer.container_client.container_name = "metadata" + self.blob_client.url = "https://account.test/metadata/model_a.json" + + self.layer.get_file_remote_path("a", "model", check_exists=False) + + delegation_key.assert_not_called() + + def test_load_map_preserves_keys_and_missing_records(self) -> None: + self.layer.load = Mock( + side_effect=[{"modelId": "a"}, FileNotFoundError()] + ) + + result = self.layer.load_map( + ["a", "a", "missing"], "model", max_workers=2 + ) + + self.assertEqual(result, {"a": {"modelId": "a"}, "missing": None}) + self.assertEqual(self.layer.load.call_count, 2) + + +class TestAzureBlobStorageDataLayerClientReuse(unittest.TestCase): + def setUp(self) -> None: + _INITIALIZED_CONTAINERS.clear() + + def tearDown(self) -> None: + _INITIALIZED_CONTAINERS.clear() + + @patch( + "hastegeo.core.data_layer.azure_blob_storage_data_layer.get_blob_service_client" + ) + @patch.object( + AzureBlobStorageDataLayer, + "_create_or_update_managed_access_policy", + ) + def test_connection_string_uses_cached_client( + self, create_policy, factory + ): + service = Mock(url="https://account.test") + service.credential.account_key = "key" # pragma: allowlist secret + factory.return_value = service + + layer = AzureBlobStorageDataLayer( + account_url="", + container="metadata", + connection_string="connection", + ) + + self.assertIs(layer.blob_service_client, service) + factory.assert_called_once_with(connection_string="connection") + create_policy.assert_called_once_with() + + @patch( + "hastegeo.core.data_layer.azure_blob_storage_data_layer.get_blob_service_client" + ) + @patch.object( + AzureBlobStorageDataLayer, + "_create_or_update_managed_access_policy", + ) + def test_managed_identity_uses_cached_client(self, create_policy, factory): + service = Mock(url="https://account.test") + service.get_user_delegation_key.return_value = "delegation-key" + factory.return_value = service + + layer = AzureBlobStorageDataLayer( + account_url="https://account.test", + container="metadata", + connection_string=None, + ) + + self.assertIs(layer.blob_service_client, service) + self.assertIsNone(layer.user_delegation_key) + service.get_user_delegation_key.assert_not_called() + factory.assert_called_once_with(account_url="https://account.test") + create_policy.assert_called_once_with() + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/data_layer/test_read_contracts.py b/hastelib/tests/core/data_layer/test_read_contracts.py new file mode 100644 index 00000000..04a614f1 --- /dev/null +++ b/hastelib/tests/core/data_layer/test_read_contracts.py @@ -0,0 +1,367 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch + +from hastegeo.core.data_layer.abstract_data_layer import AbstractDataLayer +from hastegeo.core.data_layer.azure_cosmos_db_data_layer import ( + AzureCosmosDBDataLayer, +) +from hastegeo.core.data_layer.azure_data_lake_data_layer import ( + AzureDataLakeDataLayer, +) +from hastegeo.core.data_layer.azure_postgresql_data_layer import ( + AzurePostgreSQLDataLayer, +) +from hastegeo.core.data_layer.local_file_system_data_layer import ( + LocalFileSystemDataLayer, +) +from hastegeo.core.data_layer.unified import UnifiedDataLayer +from psycopg2 import sql + + +class TestCosmosReadContract(unittest.TestCase): + def setUp(self) -> None: + self.layer = AzureCosmosDBDataLayer.__new__(AzureCosmosDBDataLayer) + self.layer.partition_key = "partition" + self.layer.container = Mock() + + def test_load_accepts_unified_data_format_keyword(self) -> None: + self.layer.container.read_item.return_value = {"value": 1} + + result = self.layer.load("record", "model", data_format="json") + + self.assertEqual(result, {"value": 1}) + + def test_list_identifiers_uses_partition_query(self) -> None: + self.layer.container.query_items.return_value = [ + "model_a", + "model_b", + "model_catalog_index", + ] + + result = self.layer.list_identifiers("model") + + self.assertEqual(result, ["a", "b"]) + call = self.layer.container.query_items.call_args.kwargs + self.assertEqual(call["partition_key"], "partition") + self.assertFalse(call["enable_cross_partition_query"]) + + def test_non_json_identifier_listing_is_empty(self) -> None: + self.assertEqual( + self.layer.list_identifiers("train_labels", "geojson"), [] + ) + self.layer.container.query_items.assert_not_called() + + def test_load_map_uses_one_partition_query_and_preserves_missing( + self, + ) -> None: + self.layer.container.query_items.return_value = [ + {"id": "model_a", "value": 1} + ] + + result = self.layer.load_map( + ["a", "a", "missing"], "model", max_workers=4 + ) + + self.assertEqual( + result, + {"a": {"id": "model_a", "value": 1}, "missing": None}, + ) + call = self.layer.container.query_items.call_args.kwargs + self.assertEqual(call["partition_key"], "partition") + self.assertEqual( + call["parameters"][1], + {"name": "@item_ids", "value": ["model_a", "model_missing"]}, + ) + + def test_partition_load_excludes_longer_metadata_type(self) -> None: + self.layer.container.query_items.return_value = [ + {"id": "model_a"}, + {"id": "model_catalog_index"}, + ] + + result = self.layer.load_all_from_partition("model") + + self.assertEqual(result, [{"id": "model_a"}]) + + def test_global_load_excludes_longer_metadata_type(self) -> None: + self.layer.container.query_items.return_value = [ + {"id": "model_a"}, + {"id": "model_catalog_index"}, + ] + + result = self.layer.load_all("model") + + self.assertEqual(result, [{"id": "model_a"}]) + + def test_load_map_handles_empty_and_non_json_inputs(self) -> None: + self.assertEqual(self.layer.load_map([], "model"), {}) + with self.assertRaisesRegex(ValueError, "only json"): + self.layer.load_map(["a"], "model", data_format="geojson") + + def test_read_methods_reject_non_json_format(self) -> None: + for method, args in ( + (self.layer.load, ("record", "model")), + (self.layer.load_all, ("model",)), + (self.layer.load_all_from_partition, ("model",)), + ): + with self.subTest(method=method.__name__): + with self.assertRaisesRegex(ValueError, "only json"): + method(*args, data_format="yaml") + + +class TestDataLakeReadContract(unittest.TestCase): + def setUp(self) -> None: + self.layer = AzureDataLakeDataLayer.__new__(AzureDataLakeDataLayer) + self.layer.partition_key = "partition" + self.layer.file_system_client = Mock() + + def test_remote_path_can_skip_exists_request(self) -> None: + file_client = ( + self.layer.file_system_client.get_file_client.return_value + ) + file_client.url = "https://account.test/file" + + result = self.layer.get_file_remote_path( + "record", "model", check_exists=False + ) + + self.assertEqual(result, "https://account.test/file") + file_client.exists.assert_not_called() + + def test_missing_remote_path_returns_none(self) -> None: + file_client = ( + self.layer.file_system_client.get_file_client.return_value + ) + file_client.exists.return_value = False + + result = self.layer.get_file_remote_path("record", "model") + + self.assertIsNone(result) + + def test_list_identifiers_strips_prefix_and_suffix(self) -> None: + self.layer.file_system_client.get_paths.return_value = [ + SimpleNamespace(name="partition/model_a.json"), + SimpleNamespace(name="partition/model_b.json"), + SimpleNamespace(name="partition/model_catalog_index.json"), + SimpleNamespace(name="partition/labels_c.json"), + ] + + result = self.layer.list_identifiers("model") + + self.assertEqual(result, ["a", "b"]) + self.layer.file_system_client.get_paths.assert_called_once_with( + path="partition" + ) + + def test_load_accepts_unified_data_format_keyword(self) -> None: + file_client = ( + self.layer.file_system_client.get_file_client.return_value + ) + file_client.download_file.return_value.readall.return_value = ( + b'{"value": 1}' + ) + + result = self.layer.load("record", "model", data_format="json") + + self.assertEqual(result, {"value": 1}) + + def test_load_all_and_partition_forward_json_format(self) -> None: + self.layer.file_system_client.get_paths.return_value = [] + self.assertEqual(self.layer.load_all("model", data_format="json"), []) + with patch.object(self.layer, "load_all", return_value=[]) as load_all: + self.assertEqual( + self.layer.load_all_from_partition( + "model", data_format="json" + ), + [], + ) + load_all.assert_called_once_with("model", data_format="json") + + def test_load_all_does_not_cross_configured_partition(self) -> None: + self.layer.file_system_client.get_paths.return_value = [ + SimpleNamespace(name="partition/model_a.json"), + SimpleNamespace(name="other/model_b.json"), + ] + file_client = ( + self.layer.file_system_client.get_file_client.return_value + ) + file_client.download_file.return_value.readall.return_value = ( + b'{"id": "a"}' + ) + + result = self.layer.load_all("model") + + self.assertEqual(result, [{"id": "a"}]) + self.layer.file_system_client.get_file_client.assert_called_once_with( + "partition/model_a.json" + ) + + def test_read_methods_reject_non_json_format(self) -> None: + for method, args in ( + (self.layer.load, ("record", "model")), + (self.layer.load_all, ("model",)), + ): + with self.subTest(method=method.__name__): + with self.assertRaisesRegex(ValueError, "only json"): + method(*args, data_format="yaml") + + def test_bounded_load_skips_deep_and_other_type_paths(self) -> None: + self.layer.file_system_client.get_paths.return_value = [ + SimpleNamespace(name="partition/nested/model_a.json"), + SimpleNamespace(name="partition/labels_a.json"), + ] + + self.assertEqual(self.layer.load_bounded("model", 2), []) + + +class TestPostgreSQLReadContract(unittest.TestCase): + def setUp(self) -> None: + self.layer = AzurePostgreSQLDataLayer.__new__(AzurePostgreSQLDataLayer) + self.layer.partition_key = "partition" + self.layer.server_name = "server" + self.layer.database_name = "database" + self.layer.postgres_user = "user" + self.layer.token = "token" + self.layer._qualified_table_identifier = sql.Identifier("metadata") + + @patch( + "hastegeo.core.data_layer.azure_postgresql_data_layer.psycopg2.connect" + ) + def test_load_accepts_jsonb_dictionary(self, connect) -> None: + cursor = self._cursor(connect) + cursor.fetchone.return_value = ({"value": 1},) + + result = self.layer.load("record", "model", data_format="json") + + self.assertEqual(result, {"value": 1}) + + @patch( + "hastegeo.core.data_layer.azure_postgresql_data_layer.psycopg2.connect" + ) + def test_list_identifiers_is_partition_scoped(self, connect) -> None: + cursor = self._cursor(connect) + cursor.fetchall.return_value = [("a",), ("b",)] + + result = self.layer.list_identifiers("model") + + self.assertEqual(result, ["a", "b"]) + self.assertEqual( + cursor.execute.call_args.args[1], ("model", "partition") + ) + + def test_non_json_identifier_listing_is_empty(self) -> None: + self.assertEqual( + self.layer.list_identifiers("train_labels", "geojson"), [] + ) + + @patch( + "hastegeo.core.data_layer.azure_postgresql_data_layer.psycopg2.connect" + ) + def test_load_map_uses_one_query_and_preserves_missing( + self, connect + ) -> None: + cursor = self._cursor(connect) + cursor.fetchall.return_value = [("a", {"value": 1})] + + result = self.layer.load_map( + ["a", "a", "missing"], "model", max_workers=4 + ) + + self.assertEqual(result, {"a": {"value": 1}, "missing": None}) + self.assertEqual( + cursor.execute.call_args.args[1], + ("model", "partition", ["a", "missing"]), + ) + + def test_load_map_handles_empty_and_non_json_inputs(self) -> None: + self.assertEqual(self.layer.load_map([], "model"), {}) + with self.assertRaisesRegex(ValueError, "only json"): + self.layer.load_map(["a"], "model", data_format="geojson") + + @patch( + "hastegeo.core.data_layer.azure_postgresql_data_layer.psycopg2.connect" + ) + def test_load_all_accepts_jsonb_values(self, connect) -> None: + cursor = self._cursor(connect) + cursor.fetchall.return_value = [({"value": 1},)] + + self.assertEqual( + self.layer.load_all("model", data_format="json"), + [{"value": 1}], + ) + + @patch( + "hastegeo.core.data_layer.azure_postgresql_data_layer.psycopg2.connect" + ) + def test_load_partition_accepts_serialized_values(self, connect) -> None: + cursor = self._cursor(connect) + cursor.fetchall.return_value = [('{"value": 1}',)] + + self.assertEqual( + self.layer.load_all_from_partition("model", data_format="json"), + [{"value": 1}], + ) + + def test_read_methods_reject_non_json_format(self) -> None: + for method, args in ( + (self.layer.load, ("record", "model")), + (self.layer.load_all, ("model",)), + (self.layer.load_all_from_partition, ("model",)), + ): + with self.subTest(method=method.__name__): + with self.assertRaisesRegex(ValueError, "only json"): + method(*args, data_format="yaml") + + @staticmethod + def _cursor(connect) -> MagicMock: + connection = MagicMock() + cursor = MagicMock() + connect.return_value.__enter__.return_value = connection + connection.cursor.return_value.__enter__.return_value = cursor + return cursor + + +class TestOptionalRemotePathContract(unittest.TestCase): + def test_default_remote_path_is_explicitly_unsupported(self) -> None: + with self.assertRaisesRegex(NotImplementedError, "remote file paths"): + AbstractDataLayer.get_file_remote_path(object()) + + def test_default_batch_read_is_explicitly_unsupported(self) -> None: + with self.assertRaisesRegex(NotImplementedError, "batch reads"): + AbstractDataLayer.load_map(object(), [], "model") + + +class TestLocalReadContract(unittest.TestCase): + def test_bounded_load_skips_nonmatching_files(self) -> None: + with tempfile.TemporaryDirectory() as directory: + layer = LocalFileSystemDataLayer(directory) + layer.save("index", "model_catalog", {"models": []}) + + self.assertEqual(layer.load_bounded("model", 1), []) + + +class TestUnifiedReadContract(unittest.TestCase): + def test_load_map_delegates_all_arguments(self) -> None: + unified = UnifiedDataLayer.__new__(UnifiedDataLayer) + unified.data_layer = Mock() + unified.data_layer.load_map.return_value = {"a": {"value": 1}} + + result = unified.load_map( + ["a"], "model", data_format="json", max_workers=4 + ) + + self.assertEqual(result, {"a": {"value": 1}}) + unified.data_layer.load_map.assert_called_once_with( + identifiers=["a"], + data_type="model", + data_format="json", + max_workers=4, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/processors/test_artifacts.py b/hastelib/tests/core/processors/test_artifacts.py index 999a56a1..eae42d9a 100644 --- a/hastelib/tests/core/processors/test_artifacts.py +++ b/hastelib/tests/core/processors/test_artifacts.py @@ -1,53 +1,25 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -import os - -from azure.core.exceptions import ResourceExistsError -from azure.storage.blob import BlobServiceClient from hastegeo.core.processors.artifacts import ArtifactProcessor class TestArtifactProcessor: - def test_zip(self, mocker): - # Arrange - blob_service_client = BlobServiceClient.from_connection_string( - os.environ.get("BLOB_CONNECTION_STRING") + def test_fetch_artifact_delegates_to_storage(self, mocker): + processor = ArtifactProcessor.__new__(ArtifactProcessor) + processor.storage = mocker.Mock() + processor.storage.fetch_artifact.return_value = "/tmp/output" + + result = processor.fetch_artifact( + identifier="artifact", + extra_partition_keys=["model"], + src_path="source", + dst_path="/tmp/output", ) - try: - container_client = blob_service_client.create_container( - os.getenv("BLOB_CONTAINER") - ) - except ResourceExistsError: - container_client = blob_service_client.get_container_client( - os.getenv("BLOB_CONTAINER") - ) - except Exception as e: - print(e) - - test_artifacts = { - "folder1": ["test1.txt", "test2.txt"], - "folder2": ["test3.txt", "test4.txt"], - } - - for folder in test_artifacts.keys(): - for file in test_artifacts[folder]: - file_path = os.path.join(folder, file) - blob_client = container_client.get_blob_client(file_path) - blob_client.upload_blob( - r"This is a test file.", overwrite=True - ) - model_id = "1234" - model_name = "test_model_name" - - # Act - processor = ArtifactProcessor() - result = processor.zip( - artifact_paths=["folder1", "folder2"], - zip_path=f"model_{model_id}_artifacts/{model_name}.zip", + assert result == "/tmp/output" + processor.storage.fetch_artifact.assert_called_once_with( + identifier="artifact", + extra_partition_keys=["model"], + src_path="source", + dst_path="/tmp/output", ) - - # Assertions - assert result == f"model_{model_id}_artifacts/{model_name}.zip" - expected_blob_client = container_client.get_blob_client(result) - assert expected_blob_client.exists() is True diff --git a/hastelib/tests/core/processors/test_metadata_batch.py b/hastelib/tests/core/processors/test_metadata_batch.py new file mode 100644 index 00000000..fa6e0c98 --- /dev/null +++ b/hastelib/tests/core/processors/test_metadata_batch.py @@ -0,0 +1,171 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Tests for the batch metadata primitives added in the perf-layer-loading work: +``MetadataProcessor.load_map`` / ``load_filtered`` / ``list_keys`` / ``build_url``. + +Runs against the local filesystem backend so no Azure/Azurite is required. +""" +import importlib +from unittest.mock import Mock + +import pytest + + +@pytest.fixture() +def local_metadata(tmp_path, monkeypatch): + monkeypatch.setenv("METADATA_STORAGE_TYPE", "local") + monkeypatch.setenv("DATA_PATH", str(tmp_path)) + # Config reads env at construction; import fresh each test. + metadata = importlib.import_module("hastegeo.core.processors.metadata") + return metadata.MetadataProcessor + + +def _seed(MetadataProcessor, partition): + for i in range(5): + MetadataProcessor("model", partition).save( + f"m{i}", + {"modelId": f"m{i}", "imageLayerId": f"layer-{i % 2}"}, + ) + + +def test_list_keys_returns_all_identifiers(local_metadata): + MetadataProcessor = local_metadata + _seed(MetadataProcessor, "p1") + keys = set(MetadataProcessor("model", "p1").list_keys()) + assert keys == {"m0", "m1", "m2", "m3", "m4"} + + +def test_partition_scan_excludes_longer_metadata_type(local_metadata): + MetadataProcessor = local_metadata + MetadataProcessor("model", "prefix").save("1", {"modelId": "1"}) + MetadataProcessor("model_catalog", "prefix").save( + "index", {"modelCatalog": []} + ) + + models = MetadataProcessor("model", "prefix") + + assert models.load_all_from_partition() == [{"modelId": "1"}] + assert models.list_keys() == ["1"] + + +def test_load_map_parallel_matches_sequential(local_metadata): + MetadataProcessor = local_metadata + _seed(MetadataProcessor, "p2") + mp = MetadataProcessor("model", "p2") + keys = mp.list_keys() + + mapped = mp.load_map(keys, max_workers=4) + sequential = {k: mp.load(k) for k in keys} + assert mapped == sequential + + +def test_load_map_missing_key_is_none(local_metadata): + MetadataProcessor = local_metadata + _seed(MetadataProcessor, "p3") + mp = MetadataProcessor("model", "p3") + result = mp.load_map(["m0", "does-not-exist"]) + assert result["m0"]["modelId"] == "m0" + assert result["does-not-exist"] is None + + +def test_load_map_empty(local_metadata): + MetadataProcessor = local_metadata + assert MetadataProcessor("model", "p4").load_map([]) == {} + + +def test_load_map_rejects_invalid_worker_count(local_metadata): + MetadataProcessor = local_metadata + _seed(MetadataProcessor, "invalid-workers") + + with pytest.raises(ValueError, match="max_workers"): + MetadataProcessor("model", "invalid-workers").load_map( + ["m0"], max_workers=0 + ) + + +def test_load_map_deduplicates_keys(local_metadata, mocker): + MetadataProcessor = local_metadata + _seed(MetadataProcessor, "duplicate-keys") + processor = MetadataProcessor("model", "duplicate-keys") + load = mocker.spy(processor, "load") + + result = processor.load_map(["m0", "m0"], max_workers=2) + + assert result["m0"]["modelId"] == "m0" + load.assert_called_once_with("m0", data_format="json") + + +def test_load_map_prefers_backend_native_batch(local_metadata): + MetadataProcessor = local_metadata + processor = MetadataProcessor.__new__(MetadataProcessor) + processor.data_type = "model" + processor.storage = Mock() + processor.storage.supports_load_map.return_value = True + processor.storage.load_map.return_value = { + "m0": {"modelId": "m0"}, + "missing": None, + } + + result = processor.load_map(["m0", "m0", "missing"]) + + assert result == {"m0": {"modelId": "m0"}, "missing": None} + processor.storage.load_map.assert_called_once_with( + identifiers=["m0", "missing"], + data_type="model", + data_format="json", + max_workers=8, + ) + + +def test_native_load_map_rejects_invalid_worker_count(local_metadata): + MetadataProcessor = local_metadata + processor = MetadataProcessor.__new__(MetadataProcessor) + processor.data_type = "model" + processor.storage = Mock() + processor.storage.supports_load_map.return_value = True + + with pytest.raises(ValueError, match="max_workers"): + processor.load_map(["m0"], max_workers=0) + + processor.storage.load_map.assert_not_called() + + +def test_load_filtered_by_field(local_metadata): + MetadataProcessor = local_metadata + _seed(MetadataProcessor, "p5") + matched = MetadataProcessor("model", "p5").load_filtered( + {"imageLayerId": "layer-0"} + ) + assert {m["modelId"] for m in matched} == {"m0", "m2", "m4"} + + +def test_load_filtered_rejects_empty_predicate(local_metadata): + MetadataProcessor = local_metadata + + with pytest.raises(ValueError, match="non-empty"): + MetadataProcessor("model", "empty-predicate").load_filtered({}) + + +def test_load_filtered_does_not_treat_missing_field_as_none(local_metadata): + MetadataProcessor = local_metadata + processor = MetadataProcessor("model", "missing-field") + processor.save("missing", {"modelId": "missing"}) + processor.save("explicit", {"modelId": "explicit", "status": None}) + + matched = processor.load_filtered({"status": None}) + + assert [record["modelId"] for record in matched] == ["explicit"] + + +def test_load_map_counts_round_trips_across_threads(local_metadata): + MetadataProcessor = local_metadata + _seed(MetadataProcessor, "p6") + perf = importlib.import_module("hastegeo.core.utils.perf") + mp = MetadataProcessor("model", "p6") + keys = mp.list_keys() + + counter = perf.begin(True) + mp.load_map(keys, max_workers=4) + perf.end() + # Each threaded load records against the shared counter (context bound). + assert counter.calls >= len(keys) diff --git a/hastelib/tests/core/utils/test_blob.py b/hastelib/tests/core/utils/test_blob.py index 52f8a122..5bc88a56 100644 --- a/hastelib/tests/core/utils/test_blob.py +++ b/hastelib/tests/core/utils/test_blob.py @@ -13,16 +13,144 @@ existing test_artifacts.py. """ +import os import unittest from unittest.mock import MagicMock, patch from hastegeo.core.utils.blob import ( + clear_blob_client_caches, + download_blob_to_tempfile, fetch_url_text, + get_blob_service_client, + get_cached_user_delegation_key, parse_byte_range, + read_blob_range, split_blob_url, ) +class TestBlobServiceClientCache(unittest.TestCase): + def setUp(self): + clear_blob_client_caches() + + def tearDown(self): + clear_blob_client_caches() + + @patch("azure.storage.blob.BlobServiceClient.from_connection_string") + def test_reuses_client_for_connection_string(self, from_connection_string): + first = get_blob_service_client(connection_string="connection") + second = get_blob_service_client(connection_string="connection") + + self.assertIs(first, second) + from_connection_string.assert_called_once_with("connection") + + @patch("azure.storage.blob.BlobServiceClient") + @patch("azure.identity.DefaultAzureCredential") + def test_reuses_client_for_account_url(self, credential, client_class): + first = get_blob_service_client(account_url="https://account.test") + second = get_blob_service_client(account_url="https://account.test") + + self.assertIs(first, second) + credential.assert_called_once_with() + client_class.assert_called_once_with( + account_url="https://account.test", + credential=credential.return_value, + ) + + def test_requires_connection_target(self): + with self.assertRaises(ValueError): + get_blob_service_client() + + def test_reuses_unexpired_user_delegation_key(self): + from datetime import datetime, timedelta, timezone + + client = MagicMock(url="https://account.test") + client.get_user_delegation_key.return_value = "delegation-key" + now = datetime(2026, 9, 1, tzinfo=timezone.utc) + + first = get_cached_user_delegation_key(client, now=now) + second = get_cached_user_delegation_key( + client, now=now + timedelta(hours=1) + ) + + self.assertEqual(first, "delegation-key") + self.assertEqual(second, "delegation-key") + client.get_user_delegation_key.assert_called_once() + + def test_refreshes_expiring_user_delegation_key(self): + from datetime import datetime, timedelta, timezone + + client = MagicMock(url="https://account.test") + client.get_user_delegation_key.side_effect = ["first", "second"] + now = datetime(2026, 9, 1, tzinfo=timezone.utc) + get_cached_user_delegation_key(client, now=now) + + result = get_cached_user_delegation_key( + client, now=now + timedelta(hours=1, minutes=46) + ) + + self.assertEqual(result, "second") + self.assertEqual(client.get_user_delegation_key.call_count, 2) + + def test_user_delegation_cache_evicts_oldest_client(self): + from datetime import datetime, timezone + + now = datetime(2026, 9, 1, tzinfo=timezone.utc) + clients = [] + for index in range(9): + client = MagicMock(url=f"https://account-{index}.test") + client.get_user_delegation_key.return_value = f"key-{index}" + clients.append(client) + get_cached_user_delegation_key(client, now=now) + + get_cached_user_delegation_key(clients[0], now=now) + + self.assertEqual(clients[0].get_user_delegation_key.call_count, 2) + self.assertEqual(clients[-1].get_user_delegation_key.call_count, 1) + + +class TestAsyncBlobHelpers(unittest.IsolatedAsyncioTestCase): + @patch("hastegeo.core.utils.blob.get_blob_service_client") + async def test_download_blob_to_tempfile_uses_shared_client(self, factory): + blob_client = ( + factory.return_value.get_container_client.return_value.get_blob_client.return_value + ) + blob_client.download_blob.return_value.chunks.return_value = [ + b"hello", + b" world", + ] + + path = await download_blob_to_tempfile( + "https://account.blob.core.windows.net/container/file.txt" + ) + self.addCleanup(lambda: os.path.exists(path) and os.unlink(path)) + + with open(path, "rb") as downloaded: + self.assertEqual(downloaded.read(), b"hello world") + factory.assert_called_once() + + @patch("hastegeo.core.utils.blob.get_blob_service_client") + async def test_read_blob_range_uses_shared_client(self, factory): + blob_client = ( + factory.return_value.get_container_client.return_value.get_blob_client.return_value + ) + properties = blob_client.get_blob_properties.return_value + properties.size = 5 + properties.content_settings.content_type = "text/plain" + properties.etag = '"etag"' + blob_client.download_blob.return_value.readall.return_value = b"ell" + + result = await read_blob_range( + "https://account.blob.core.windows.net/container/file.txt", + offset=1, + length=3, + ) + + self.assertEqual(result.data, b"ell") + self.assertEqual(result.total_size, 5) + blob_client.download_blob.assert_called_once_with(offset=1, length=3) + + class TestSplitBlobUrl(unittest.TestCase): def test_azurite_docker_internal_host(self): # The case that motivated the helper in the first place: the diff --git a/hastelib/tests/core/utils/test_metadata_type_matching.py b/hastelib/tests/core/utils/test_metadata_type_matching.py new file mode 100644 index 00000000..18bfc37a --- /dev/null +++ b/hastelib/tests/core/utils/test_metadata_type_matching.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +import unittest + +from hastegeo.core.utils.metadata import matches_metadata_type + + +class TestMetadataTypeMatching(unittest.TestCase): + def test_longest_known_metadata_type_wins(self) -> None: + self.assertTrue(matches_metadata_type("model_123.json", "model")) + self.assertFalse( + matches_metadata_type("model_catalog_index.json", "model") + ) + self.assertTrue( + matches_metadata_type( + "partition/model_catalog_index.json", "model_catalog" + ) + ) + + def test_unknown_type_uses_requested_prefix(self) -> None: + self.assertTrue(matches_metadata_type("custom_123.json", "custom")) + self.assertFalse(matches_metadata_type("other_123.json", "custom")) + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/utils/test_parallel.py b/hastelib/tests/core/utils/test_parallel.py new file mode 100644 index 00000000..602e5753 --- /dev/null +++ b/hastelib/tests/core/utils/test_parallel.py @@ -0,0 +1,113 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +import os +import unittest +from concurrent.futures import ThreadPoolExecutor +from threading import Event, Lock +from unittest.mock import patch + +from hastegeo.core.utils.parallel import ( + BoundedExecutor, + configured_worker_count, +) + + +class TestConfiguredWorkerCount(unittest.TestCase): + def test_uses_default_when_environment_is_missing(self) -> None: + with patch.dict(os.environ, {}, clear=True): + self.assertEqual(configured_worker_count("TEST_WORKERS", 4), 4) + + def test_reads_valid_environment_value(self) -> None: + with patch.dict(os.environ, {"TEST_WORKERS": "7"}): + self.assertEqual(configured_worker_count("TEST_WORKERS", 4), 7) + + def test_rejects_invalid_environment_values(self) -> None: + for value in ("not-an-int", "0", "65"): + with self.subTest(value=value): + with patch.dict(os.environ, {"TEST_WORKERS": value}): + with self.assertRaises(ValueError): + configured_worker_count("TEST_WORKERS", 4) + + def test_rejects_non_integer_direct_value(self) -> None: + for value in (True, "4"): + with self.subTest(value=value): + with self.assertRaisesRegex(ValueError, "integer"): + BoundedExecutor(max_workers=value) + + +class TestBoundedExecutor(unittest.TestCase): + def setUp(self) -> None: + self.executor = BoundedExecutor(max_workers=2) + + def tearDown(self) -> None: + self.executor.shutdown() + + def test_preserves_input_order(self) -> None: + result = self.executor.map(lambda value: value * 2, [3, 1, 2]) + + self.assertEqual(result, [6, 2, 4]) + + def test_propagates_worker_exceptions(self) -> None: + def fail_on_two(value: int) -> int: + if value == 2: + raise RuntimeError("failed") + return value + + with self.assertRaisesRegex(RuntimeError, "failed"): + self.executor.map(fail_on_two, [1, 2, 3]) + + def test_cancels_queued_work_after_failure(self) -> None: + release = Event() + + def fail_with_pending_work(value: int) -> int: + if value == 1: + raise RuntimeError("failed") + release.wait(timeout=1) + return value + + try: + with self.assertRaisesRegex(RuntimeError, "failed"): + self.executor.map(fail_with_pending_work, [1, 2, 3, 4]) + finally: + release.set() + + def test_rejects_invalid_per_call_limit(self) -> None: + with self.assertRaises(ValueError): + self.executor.map(str, [1], max_workers=0) + + def test_nested_map_does_not_deadlock(self) -> None: + single_worker = BoundedExecutor(max_workers=1) + self.addCleanup(single_worker.shutdown) + + result = single_worker.map( + lambda value: single_worker.map(lambda item: item, [value])[0], + [1], + ) + + self.assertEqual(result, [1]) + + def test_concurrent_maps_share_the_process_budget(self) -> None: + active = 0 + peak = 0 + lock = Lock() + two_workers_started = Event() + + def work(value: int) -> int: + nonlocal active, peak + with lock: + active += 1 + peak = max(peak, active) + if active == 2: + two_workers_started.set() + self.assertTrue(two_workers_started.wait(timeout=1)) + with lock: + active -= 1 + return value + + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(self.executor.map, work, [1, 2, 3]) + second = callers.submit(self.executor.map, work, [4, 5, 6]) + self.assertEqual(first.result(), [1, 2, 3]) + self.assertEqual(second.result(), [4, 5, 6]) + + self.assertEqual(peak, 2) diff --git a/hastelib/tests/core/utils/test_perf.py b/hastelib/tests/core/utils/test_perf.py new file mode 100644 index 00000000..6a4bbdc7 --- /dev/null +++ b/hastelib/tests/core/utils/test_perf.py @@ -0,0 +1,48 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +import unittest +from unittest.mock import Mock + +from hastegeo.core.utils import perf + + +class TestPerfInstrumentation(unittest.TestCase): + def tearDown(self) -> None: + perf.end() + + def test_timed_records_logical_data_layer_operation(self) -> None: + counter = perf.begin(True) + + with perf.timed("load"): + pass + + self.assertEqual(counter.calls, 1) + self.assertEqual(counter.by_op["load"]["calls"], 1) + + def test_headers_include_new_and_legacy_names(self) -> None: + counter = perf.begin(True) + + headers = perf.headers(counter, 0) + + self.assertEqual(headers["X-Haste-Data-Layer-Calls"], "0") + self.assertEqual(headers["X-Haste-Storage-Calls"], "0") + self.assertIn("data-layer", headers["Server-Timing"]) + + def test_disabled_instrumentation_emits_no_headers(self) -> None: + counter = perf.begin(False) + + self.assertIsNone(counter) + self.assertEqual(perf.headers(counter, 0), {}) + + def test_log_summary_clears_active_counter(self) -> None: + counter = perf.begin(True) + logger = Mock() + + perf.log_summary(logger, "operation", counter, 0, key="value") + + self.assertIsNone(perf.get_counter()) + logger.info.assert_called_once() + + +if __name__ == "__main__": + unittest.main()