diff --git a/pychunkedgraph/app/common.py b/pychunkedgraph/app/common.py index 237e11fc0..73ba9a22a 100644 --- a/pychunkedgraph/app/common.py +++ b/pychunkedgraph/app/common.py @@ -17,6 +17,37 @@ ENABLE_LOGS = os.environ.get("PCG_SERVER_ENABLE_LOGS", "") != "" LOG_LEAVES_MANY = os.environ.get("PCG_SERVER_LOGS_LEAVES_MANY", "") != "" +# Health-check paths to skip, matched EXACTLY (not as prefixes) so the start-log signal +# isn't flooded by probes. These are the literal probe paths from the pychunkedgraph chart: +# the read/write deployments' readiness+liveness probes hit "/segmentation" and the GCP load +# balancer health check hits "/". Real API traffic lives under "/segmentation/api/..." and +# "/meshing/api/...", which are NOT equal to these entries and so are still logged. +_REQUEST_START_SKIP_PATHS = frozenset(("/", "/segmentation")) + + +def _log_request_start(): + # Emit a line to stdout at the *start* of a request, before any work runs, so it is + # captured by Cloud Logging even if the request goes on to OOM-kill or otherwise crash + # its worker before completing (such requests never reach after_request and so are + # invisible in the Datastore server_logs completion logs). The `content_length` field is + # the on-the-wire request body size (e.g. for a roots_binary POST, ~8 bytes per node id) + # and `pid` is the uwsgi worker, so a spike/OOM can be traced to the specific in-flight + # request and worker. Gated by the LOG_REQUEST_START Flask config value (default False); + # verbose, intended for temporary diagnosis. + try: + user_id = g.auth_user["id"] + except (AttributeError, KeyError): + user_id = USER_NOT_FOUND + current_app.logger.info( + "REQUEST_START pid=%s method=%s path=%s content_length=%s user=%s remote=%s", + os.getpid(), + request.method, + request.path, + request.content_length, + user_id, + request.remote_addr, + ) + def _log_request(response_time): try: @@ -58,6 +89,10 @@ def before_request(): current_app.table_id = None current_app.operation_id = None current_app.request_type = None + if current_app.config.get("LOG_REQUEST_START", False) and ( + request.path not in _REQUEST_START_SKIP_PATHS + ): + _log_request_start() content_encoding = request.headers.get("Content-Encoding", "") if "gzip" in content_encoding.lower(): request.data = compression.decompress(request.data, "gzip") diff --git a/pychunkedgraph/app/config.py b/pychunkedgraph/app/config.py index 2f2a92e47..4e8d5bd00 100644 --- a/pychunkedgraph/app/config.py +++ b/pychunkedgraph/app/config.py @@ -14,6 +14,26 @@ class BaseConfig(object): LOGGING_DATEFORMAT = "%Y-%m-%dT%H:%M:%S.0Z" LOGGING_LEVEL = logging.DEBUG + # Opt-in start-of-request logging (see pychunkedgraph.app.common._log_request_start). + # When True, every non-probe request emits a REQUEST_START line to stdout before any work + # runs, so requests that OOM-kill their worker mid-flight (and thus never reach + # after_request) are still visible in Cloud Logging. Verbose; enable only for temporary + # diagnosis, e.g. by setting LOG_REQUEST_START = True in the instance config.cfg. + LOG_REQUEST_START = False + + # Reject /lvl2_graph requests whose node resolves to more than this many level 2 nodes + # (see pychunkedgraph.graph.analysis.pathing.get_lvl2_edge_list). Such objects — typically + # erroneous mega-merges — produce a multi-GB induced edge list that can OOM the worker. + # None disables the guard; set a concrete integer in the instance config.cfg to enable. + LVL2_GRAPH_MAX_NODES = None + + # Guard for /subgraph (see pychunkedgraph.graph.subgraph.get_subgraph_edges_and_leaves). + # Counts chunks rather than level 2 nodes: the endpoint reads every edge in every chunk the + # object touches (all objects in the chunk, not just the requested one), so cost tracks the + # volume queried, not the object. A large 'bounds' is expensive even for a small object. + # None disables the guard; set a concrete integer in the instance config.cfg to enable. + SUBGRAPH_MAX_CHUNKS = None + CHUNKGRAPH_INSTANCE_ID = "pychunkedgraph" PROJECT_ID = os.environ.get("PROJECT_ID", None) CG_READ_ONLY = os.environ.get("CG_READ_ONLY", None) is not None diff --git a/pychunkedgraph/app/segmentation/common.py b/pychunkedgraph/app/segmentation/common.py index 3250248f2..3fd06975e 100644 --- a/pychunkedgraph/app/segmentation/common.py +++ b/pychunkedgraph/app/segmentation/common.py @@ -4,7 +4,6 @@ import os import time from datetime import datetime -from functools import reduce from collections import deque, defaultdict import numpy as np @@ -785,8 +784,9 @@ def handle_subgraph(table_id, root_id, only_internal_edges=True): int(root_id), bbox=bounding_box, bbox_is_coordinate=True, + max_num_chunks=current_app.config.get("SUBGRAPH_MAX_CHUNKS"), ) - edges = reduce(lambda x, y: x + y, edges, cg_edges.Edges([], [])) + edges = cg_edges.Edges.concatenate(edges) if only_internal_edges: supervoxels = np.concatenate( @@ -830,31 +830,43 @@ def tabular_change_log_recent(table_id): # Call ChunkedGraph cg = app_utils.get_cg(table_id) - log_rows = cg.client.read_log_entries(start_time=start_time, end_time=end_time) - + # Stream the operation-log rows instead of materializing them all at once. Only the + # timestamp, user, and merge/split flag are needed, so the read is restricted to those + # columns (the default pulls the large variable-length added/removed-edge, coordinate, + # and affinity arrays for every operation, which dominate row size). AddedEdge is only + # existence-checked (merge vs split). RootID is not used directly but is kept so the + # streaming reader's timestamp fallback still works on older rows that predate the + # OperationTimeStamp column. + # + # read_log_entries_streaming yields one (operation_id, record) at a time and frees each + # decoded row before reading the next, so peak memory is bounded by the compact output + # columns below rather than by the full set of Bigtable cell objects for the window. + # Rows arrive already ordered by operation id (fixed-width keys), so no sort is needed. + operation_ids = [] timestamp_list = [] user_list = [] is_merge_list = [] - - operation_ids = np.sort(list(log_rows.keys())) - for operation_id in operation_ids: - operation = log_rows[operation_id] - - timestamp = operation["timestamp"] - timestamp_list.append(timestamp) - - user_id = operation[attributes.OperationLogs.UserID] - user_list.append(user_id) - - is_merge = attributes.OperationLogs.AddedEdge in operation - is_merge_list.append(is_merge) + for operation_id, operation in cg.client.read_log_entries_streaming( + start_time=start_time, + end_time=end_time, + properties=[ + attributes.OperationLogs.OperationTimeStamp, + attributes.OperationLogs.UserID, + attributes.OperationLogs.AddedEdge, + attributes.OperationLogs.RootID, + ], + ): + operation_ids.append(operation_id) + timestamp_list.append(operation["timestamp"]) + user_list.append(operation[attributes.OperationLogs.UserID]) + is_merge_list.append(attributes.OperationLogs.AddedEdge in operation) return pd.DataFrame.from_dict( { - "operation_id": operation_ids, + "operation_id": np.array(operation_ids, dtype=np.uint64), "timestamp": timestamp_list, "user_id": user_list, - "is_merge": is_merge_list, + "is_merge": np.array(is_merge_list, dtype=bool), } ) @@ -1142,7 +1154,12 @@ def handle_get_layer2_graph(table_id, node_id): cg = app_utils.get_cg(table_id) print("Finding edge graph...") - edge_graph = pathing.get_lvl2_edge_list(cg, int(node_id), bbox=bounding_box) + edge_graph = pathing.get_lvl2_edge_list( + cg, + int(node_id), + bbox=bounding_box, + max_num_lvl2_ids=current_app.config.get("LVL2_GRAPH_MAX_NODES"), + ) print("Edge graph found len: {}".format(len(edge_graph))) return {"edge_graph": edge_graph} diff --git a/pychunkedgraph/graph/analysis/pathing.py b/pychunkedgraph/graph/analysis/pathing.py index 062b7a1c3..ca22683c4 100644 --- a/pychunkedgraph/graph/analysis/pathing.py +++ b/pychunkedgraph/graph/analysis/pathing.py @@ -6,6 +6,7 @@ from pychunkedgraph.graph.utils import flatgraph +from .. import exceptions as cg_exceptions from ..subgraph import get_subgraph_nodes @@ -77,12 +78,17 @@ def get_lvl2_edge_list( cg, node_id: np.uint64, bbox: typing.Optional[typing.Sequence[typing.Sequence[int]]] = None, + max_num_lvl2_ids: typing.Optional[int] = None, ): """get an edge list of lvl2 ids for a particular node :param cg: ChunkedGraph object :param node_id: np.uint64 that you want the edge list for :param bbox: Optional[Sequence[Sequence[int]]] a bounding box to limit the search + :param max_num_lvl2_ids: Optional[int] reject the request (raising BadRequest) when the + node resolves to more than this many level 2 ids. Guards against pathologically large + objects (e.g. erroneous mega-merges) whose induced level 2 edge list would be many GB + and can OOM the worker. ``None`` disables the guard. """ if bbox is None: @@ -98,6 +104,20 @@ def get_lvl2_edge_list( return_flattened=True, ) + # Enforce the size guard *before* the (potentially multi-GB) induced-edge computation + # below. The level 2 id count is the cheap proxy we already have in hand; the edge read + # in _get_edges_for_lvl2_ids scales with it and is what actually exhausts memory. + if max_num_lvl2_ids is not None and len(lvl2_ids) > max_num_lvl2_ids: + hint = ( + "Provide a smaller bounding box ('bounds')." + if bbox is not None + else "Provide a bounding box ('bounds') to restrict the query to a sub-region." + ) + raise cg_exceptions.BadRequest( + f"The level 2 graph for {node_id} has {len(lvl2_ids)} level 2 nodes, which exceeds " + f"the maximum of {max_num_lvl2_ids}. {hint}" + ) + edges = _get_edges_for_lvl2_ids(cg, lvl2_ids, induced=True) return edges diff --git a/pychunkedgraph/graph/chunkedgraph.py b/pychunkedgraph/graph/chunkedgraph.py index 210bff50b..d9dca4be6 100644 --- a/pychunkedgraph/graph/chunkedgraph.py +++ b/pychunkedgraph/graph/chunkedgraph.py @@ -556,9 +556,14 @@ def get_subgraph( edges_only: bool = False, leaves_only: bool = False, return_flattened: bool = False, + max_num_chunks: typing.Optional[int] = None, ) -> typing.Tuple[typing.Dict, typing.Dict, Edges]: """ Generic subgraph method. + + :param max_num_chunks: Optional[int] reject the request (raising BadRequest) when the + node ids span more than this many chunks. ``None`` disables the guard. Only + applies to the edges/leaves path, which is the one that can OOM. """ from .subgraph import get_subgraph_nodes from .subgraph import get_subgraph_edges_and_leaves @@ -573,7 +578,13 @@ def get_subgraph( return_flattened=return_flattened, ) return get_subgraph_edges_and_leaves( - self, node_id_or_ids, bbox, bbox_is_coordinate, edges_only, leaves_only + self, + node_id_or_ids, + bbox, + bbox_is_coordinate, + edges_only, + leaves_only, + max_num_chunks=max_num_chunks, ) def get_subgraph_nodes( @@ -657,22 +668,32 @@ def get_l2_agglomerations( Edges are read from cloud storage. """ from itertools import chain - from functools import reduce from .misc import get_agglomerations chunk_ids = np.unique(self.get_chunk_ids_from_node_ids(level2_ids)) + + # Resolve the object's supervoxels before reading any edges. Both consumers below + # discard every edge that does not touch one of them, so passing them down lets each + # chunk be filtered as it is parsed instead of after the whole set is materialized. + # A chunk's edge file holds every object in that chunk, so for a single neuron this + # is the difference between retaining the chunk content and retaining the object. + l2id_children_d = self.get_children(level2_ids) + supervoxels = ( + np.concatenate(list(l2id_children_d.values())) + if l2id_children_d + else types.empty_1d.copy() + ) + # google does not provide a storage emulator at the moment # this is an ugly hack to avoid permission issues in tests # find a better way to test edges_d = {} if self.mock_edges is None: - edges_d = self.read_chunk_edges(chunk_ids) + edges_d = self.read_chunk_edges(chunk_ids, supervoxels=supervoxels) fake_edges = self.get_fake_edges(chunk_ids) - all_chunk_edges = reduce( - lambda x, y: x + y, - chain(edges_d.values(), fake_edges.values()), - Edges([], []), + all_chunk_edges = Edges.concatenate( + chain(edges_d.values(), fake_edges.values()) ) if edges_only: @@ -680,12 +701,10 @@ def get_l2_agglomerations( all_chunk_edges = self.mock_edges.get_pairs() else: all_chunk_edges = all_chunk_edges.get_pairs() - supervoxels = self.get_children(level2_ids, flatten=True) mask0 = np.in1d(all_chunk_edges[:, 0], supervoxels) mask1 = np.in1d(all_chunk_edges[:, 1], supervoxels) return all_chunk_edges[mask0 & mask1] - l2id_children_d = self.get_children(level2_ids) sv_parent_d = {} for l2id in l2id_children_d: svs = l2id_children_d[l2id] @@ -996,12 +1015,15 @@ def get_parent_chunk_id_dict(self, node_or_chunk_id: basetypes.NODE_ID): def get_cross_chunk_edges_layer(self, cross_edges: typing.Iterable): return edge_utils.get_cross_chunk_edges_layer(self.meta, cross_edges) - def read_chunk_edges(self, chunk_ids: typing.Iterable) -> typing.Dict: + def read_chunk_edges( + self, chunk_ids: typing.Iterable, supervoxels: np.ndarray = None + ) -> typing.Dict: from ..io.edges import get_chunk_edges return get_chunk_edges( self.meta.data_source.EDGES, self.get_chunk_coordinates_multiple(chunk_ids), + supervoxels=supervoxels, ) def get_proofread_root_ids( diff --git a/pychunkedgraph/graph/client/base.py b/pychunkedgraph/graph/client/base.py index a66602a6a..aa77a9391 100644 --- a/pychunkedgraph/graph/client/base.py +++ b/pychunkedgraph/graph/client/base.py @@ -150,3 +150,22 @@ def read_log_entry(self, operation_id: int) -> None: @abstractmethod def read_log_entries(self, operation_ids) -> None: """Read log entries for given operation IDs.""" + + @abstractmethod + def read_log_entries_streaming( + self, + properties=None, + start_time=None, + end_time=None, + end_time_inclusive=False, + user_id=None, + ): + """Yield ``(operation_id, log_record)`` for every operation in a time range. + + Streaming counterpart to :meth:`read_log_entries` for the "all operations in a time + range" case (``operation_ids=None``). Implementations should iterate the backend's + result lazily and yield one operation at a time so peak memory is bounded by what the + caller accumulates rather than by the full result set. Each ``log_record`` must match + the per-operation shape returned by :meth:`read_log_entries` (columns unwrapped to their + value, plus a derived ``"timestamp"`` key). + """ diff --git a/pychunkedgraph/graph/client/bigtable/client.py b/pychunkedgraph/graph/client/bigtable/client.py index 5b86826bd..2748d8d33 100644 --- a/pychunkedgraph/graph/client/bigtable/client.py +++ b/pychunkedgraph/graph/client/bigtable/client.py @@ -282,6 +282,68 @@ def read_log_entries( log_record["timestamp"] = timestamp return logs_d + def read_log_entries_streaming( + self, + properties: typing.Optional[typing.Iterable[attributes._Attribute]] = None, + start_time: typing.Optional[datetime] = None, + end_time: typing.Optional[datetime] = None, + end_time_inclusive: bool = False, + user_id: typing.Optional[str] = None, + ): + """Streaming counterpart to :meth:`read_log_entries` for the "all operations in a + time range" case (i.e. ``operation_ids=None``). + + :meth:`read_log_entries` materializes every matching operation-log row into a single + dict up front. For a wide time window that dict holds tens of thousands of heavyweight + Bigtable cell objects in memory simultaneously, which dominates the request's peak RSS. + This method instead iterates the underlying ``read_rows`` stream and yields one + ``(operation_id, log_record)`` pair at a time, letting each decoded row be freed before + the next is read. Peak memory is then bounded by whatever the caller accumulates, not by + the full row set. + + ``log_record`` has the same shape as the per-operation values produced by + :meth:`read_log_entries`: columns unwrapped to their first cell's deserialized value, + plus a derived ``"timestamp"`` key. + + The operation-log key space is a single contiguous range (0 -> max operation id) of + fixed-width, zero-padded keys, so the range read returns rows already ordered by + operation id; callers need not sort. + """ + if properties is None: + properties = attributes.OperationLogs.all() + + row_set = RowSet() + row_set.add_row_range_from_keys( + start_key=serialize_uint64(np.uint64(0)), + start_inclusive=True, + end_key=serialize_uint64(self.get_max_operation_id()), + end_inclusive=True, + ) + row_filter = utils.get_time_range_and_column_filter( + columns=properties, + start_time=start_time, + end_time=end_time, + end_inclusive=end_time_inclusive, + user_id=user_id, + ) + + for row in self._table.read_rows(row_set=row_set, filter_=row_filter): + column_dict = utils.partial_row_data_to_column_dict(row) + # Deserialize cell values in place (mirrors the post-read loop in _read_byte_rows). + for column, cells in column_dict.items(): + for cell in cells: + cell.value = column.deserialize(cell.value) + # Derive the operation timestamp exactly as read_log_entries does: prefer the + # explicit OperationTimeStamp value, falling back to the RootID cell's timestamp + # on older rows that predate that column. + try: + timestamp = column_dict[attributes.OperationLogs.OperationTimeStamp][0].value + except KeyError: + timestamp = column_dict[attributes.OperationLogs.RootID][0].timestamp + log_record = {column: cells[0].value for column, cells in column_dict.items()} + log_record["timestamp"] = timestamp + yield deserialize_uint64(row.row_key), log_record + # Helpers def write( self, diff --git a/pychunkedgraph/graph/edges/__init__.py b/pychunkedgraph/graph/edges/__init__.py index b0e488d05..8f3c8a7ad 100644 --- a/pychunkedgraph/graph/edges/__init__.py +++ b/pychunkedgraph/graph/edges/__init__.py @@ -20,6 +20,18 @@ DEFAULT_AREA = np.finfo(np.float32).tiny +def in_sorted(values: np.ndarray, sorted_unique: np.ndarray) -> np.ndarray: + """Boolean mask of `values` present in `sorted_unique` (sorted, deduplicated). + + Same result as np.isin, but does not re-sort the reference set on every call. + """ + if values.size == 0 or sorted_unique.size == 0: + return np.zeros(values.size, dtype=bool) + idx = np.searchsorted(sorted_unique, values) + idx[idx == sorted_unique.size] = 0 + return sorted_unique[idx] == values + + class Edges: def __init__( self, @@ -65,6 +77,50 @@ def areas(self) -> np.ndarray: def areas(self, areas): self._areas = areas + def filter_touching(self, sorted_ids: np.ndarray) -> "Edges": + """Keep only edges with at least one endpoint in `sorted_ids`. + + `sorted_ids` must be sorted and deduplicated; sorting once at the call site + matters because this runs per chunk against the same set. + + This is deliberately a superset of what consumers keep -- categorize_edges_v2 + drops any edge whose node_ids1 does not remap through sv_parent_d, and the + edges_only path keeps only edges with *both* endpoints in the set -- so applying + it early cannot change their results. Filtering before the per-chunk edges are + accumulated is what bounds memory: the arrays produced by io.edges.deserialize + are np.frombuffer views into the decompressed chunk, so masking copies out the + few edges that matter and lets the whole decompressed buffer be released. + """ + if len(self) == 0 or sorted_ids.size == 0: + return self if len(self) == 0 else self[np.zeros(len(self), dtype=bool)] + mask = in_sorted(self.node_ids1, sorted_ids) + mask |= in_sorted(self.node_ids2, sorted_ids) + return self[mask] + + @classmethod + def concatenate(cls, edges_iterable) -> "Edges": + """Combine any number of Edges in a single pass. + + Equivalent to ``reduce(lambda x, y: x + y, edges_iterable, Edges([], []))`` but + allocates each output array once instead of once per element. Folding with ``+`` + is quadratic in allocation: combining n chunk edge sets copies every edge already + accumulated on each step, so peak memory runs well above the size of the result. + Callers that combine per-chunk edges (see ChunkedGraph.get_l2_agglomerations) + should use this instead. + """ + parts = list(edges_iterable) + if not parts: + return cls( + np.array([], dtype=basetypes.NODE_ID), + np.array([], dtype=basetypes.NODE_ID), + ) + return cls( + np.concatenate([p.node_ids1 for p in parts]), + np.concatenate([p.node_ids2 for p in parts]), + affinities=np.concatenate([p.affinities for p in parts]), + areas=np.concatenate([p.areas for p in parts]), + ) + def __add__(self, other): """add two Edges instances""" node_ids1 = np.concatenate([self.node_ids1, other.node_ids1]) diff --git a/pychunkedgraph/graph/subgraph.py b/pychunkedgraph/graph/subgraph.py index ab2593175..cca9e610c 100644 --- a/pychunkedgraph/graph/subgraph.py +++ b/pychunkedgraph/graph/subgraph.py @@ -155,9 +155,23 @@ def get_subgraph_edges_and_leaves( bbox_is_coordinate: bool = False, edges_only: bool = False, leaves_only: bool = False, + max_num_chunks: Optional[int] = None, ) -> Tuple[Dict, Dict, Edges]: - """Get the edges and/or leaves of the specified node_ids within the specified bounding box.""" + """Get the edges and/or leaves of the specified node_ids within the specified bounding box. + + :param max_num_chunks: Optional[int] reject the request (raising BadRequest) when the node + ids span more than this many chunks. ``None`` disables the guard. + + The guard counts *chunks*, not level 2 ids or supervoxels, because that is what the + memory actually scales with: get_l2_agglomerations below maps the level 2 ids to their + chunks and then reads every edge in each of those chunks from cloud storage — all + objects in the chunk, not just the requested one. So a request over a large bounding + box is expensive even when the object itself is small, and the level 2 count is a poor + predictor of the byte count. Deriving the chunk ids is pure bit manipulation on ids we + already hold, so the check costs nothing. + """ from .types import empty_1d + from . import exceptions as cg_exceptions node_ids = node_id_or_ids bbox = normalize_bounding_box(cg.meta, bbox, bbox_is_coordinate) @@ -170,6 +184,25 @@ def get_subgraph_edges_and_leaves( for node_id in node_ids: level2_ids.append(layer_nodes_d[node_id]) level2_ids = np.concatenate(level2_ids) + + # Enforce the size guard *before* the (potentially multi-GB) agglomeration read below. + # Same chunk id derivation get_l2_agglomerations does, but without the reads that follow. + if max_num_chunks is not None: + num_chunks = np.unique(cg.get_chunk_ids_from_node_ids(level2_ids)).size + if num_chunks > max_num_chunks: + hint = ( + "Provide a smaller bounding box ('bounds')." + if bbox is not None + else "Provide a bounding box ('bounds') to restrict the query to a sub-region." + ) + nodes_str = ", ".join(str(node_id) for node_id in node_ids) + raise cg_exceptions.BadRequest( + f"The subgraph for {nodes_str} spans {num_chunks} chunks " + f"({len(level2_ids)} level 2 nodes), which exceeds the maximum of " + f"{max_num_chunks}. Every edge in each chunk is read, so the cost scales with " + f"the volume queried rather than the size of the object. {hint}" + ) + if leaves_only: return cg.get_children(level2_ids, flatten=True) if edges_only: diff --git a/pychunkedgraph/io/edges.py b/pychunkedgraph/io/edges.py index 82595e139..bf9c92876 100644 --- a/pychunkedgraph/io/edges.py +++ b/pychunkedgraph/io/edges.py @@ -3,6 +3,8 @@ Functions for reading and writing edges from cloud storage. """ import os +import threading +from concurrent.futures import ThreadPoolExecutor from typing import Dict from typing import List from typing import Tuple @@ -36,22 +38,54 @@ def deserialize(edges_message: EdgesMsg) -> Tuple[np.ndarray, np.ndarray, np.nda return Edges(sv_ids1, sv_ids2, affinities=affinities, areas=areas) -def _parse_edges(compressed: List[bytes]) -> List[Dict]: +def _decompress_one(zdc, content): + # zdc.decompress needs the content size in the frame header, which is what + # put_chunk_edges writes (and what multi_decompress_to_buffer also required). + # decompressobj streams and needs no size, so it covers any frame lacking it. + try: + return zdc.decompress(content) + except zstd.ZstdError: + return zdc.decompressobj().decompress(content) + + +def _decompress(compressed: List[bytes], n_threads: int) -> List[bytes]: + """Decompress chunk blobs, in parallel when n_threads > 1. + + Replaces multi_decompress_to_buffer, which zstandard removed in 0.23. A thread pool + matches it because decompression releases the GIL: measured on zstandard 0.21.0 with + 192 MB across 64 blobs and 4 threads, 41.7 ms for the pool vs 42.5 ms for + multi_decompress_to_buffer, against 164.8 ms serial. This keeps that ~4x without + depending on a specific zstandard version. + + ZstdDecompressor is not thread-safe -- sharing one across threads silently produces + corrupt output rather than raising -- so each worker keeps its own. + """ + if n_threads <= 1 or len(compressed) < 2: + zdc = zstd.ZstdDecompressor() + return [_decompress_one(zdc, content) for content in compressed] + + local = threading.local() + + def _one(content): + zdc = getattr(local, "zdc", None) + if zdc is None: + zdc = local.zdc = zstd.ZstdDecompressor() + return _decompress_one(zdc, content) + + with ThreadPoolExecutor(max_workers=n_threads) as pool: + return list(pool.map(_one, compressed)) + + +def _parse_edges(compressed: List[bytes], sorted_svs: np.ndarray = None) -> List[Dict]: result = [] if(len(compressed) == 0): return result - zdc = zstd.ZstdDecompressor() try: n_threads = int(os.environ.get("ZSTD_THREADS", 1)) except ValueError: n_threads = 1 - decompressed = [] - try: - decompressed = zdc.multi_decompress_to_buffer(compressed, threads=n_threads) - except ValueError: - for content in compressed: - decompressed.append(zdc.decompressobj().decompress(content)) + decompressed = _decompress(compressed, n_threads) for content in decompressed: chunk_edges = ChunkEdgesMsg() @@ -60,26 +94,67 @@ def _parse_edges(compressed: List[bytes]) -> List[Dict]: edges_dict[EDGE_TYPES.in_chunk] = deserialize(chunk_edges.in_chunk) edges_dict[EDGE_TYPES.between_chunk] = deserialize(chunk_edges.between_chunk) edges_dict[EDGE_TYPES.cross_chunk] = deserialize(chunk_edges.cross_chunk) + if sorted_svs is not None: + for edge_type, edges in edges_dict.items(): + edges_dict[edge_type] = edges.filter_touching(sorted_svs) result.append(edges_dict) return result -def get_chunk_edges(edges_dir: str, chunks_coordinates: List[np.ndarray]) -> Dict: - """Read edges from GCS.""" +try: + EDGES_BATCH_SIZE = int(os.environ.get("PCG_EDGES_BATCH_SIZE", 64)) +except ValueError: + EDGES_BATCH_SIZE = 64 + + +def get_chunk_edges( + edges_dir: str, + chunks_coordinates: List[np.ndarray], + supervoxels: np.ndarray = None, + batch_size: int = None, +) -> Dict: + """Read edges from GCS. + + :param supervoxels: optional supervoxel ids of the object being queried. When given, + each chunk is filtered to edges touching one of them before anything is retained, + so peak memory tracks the size of the object rather than the total edge content of + the chunks it spans. ``None`` reads every edge (the ingest path relies on this). + :param batch_size: how many chunk files to fetch and decompress at a time. Bounds the + transient buffers to the batch instead of the whole request, so a query spanning + many chunks costs no more per moment than one spanning a few. Defaults to + PCG_EDGES_BATCH_SIZE (64). + """ fnames = [] for chunk_coords in chunks_coordinates: chunk_str = "_".join(str(coord) for coord in chunk_coords) # filename format - edges_x_y_z.serialization.compression fnames.append(f"edges_{chunk_str}.proto.zst") + # sort once here rather than per chunk inside the parse loop + sorted_svs = None + if supervoxels is not None: + sorted_svs = np.unique(np.asarray(supervoxels, dtype=basetypes.NODE_ID)) + + if batch_size is None: + batch_size = EDGES_BATCH_SIZE + batch_size = max(1, batch_size) + cf = CloudFiles(edges_dir, num_threads=4) - files = cf.get(fnames, raw=True) - compressed = [] - for f in files: - if not f["content"]: - continue - compressed.append(f["content"]) - return concatenate_chunk_edges(_parse_edges(compressed)) + # Accumulate the per-chunk dicts batch by batch. Each batch's compressed and + # decompressed buffers are released before the next is fetched; only the filtered + # survivors are carried forward, so peak tracks batch_size rather than len(fnames). + parsed = [] + for start in range(0, len(fnames), batch_size): + files = cf.get(fnames[start : start + batch_size], raw=True) + compressed = [] + for f in files: + if not f["content"]: + continue + compressed.append(f["content"]) + del files + parsed.extend(_parse_edges(compressed, sorted_svs)) + del compressed + return concatenate_chunk_edges(parsed) def put_chunk_edges( diff --git a/pychunkedgraph/tests/test_edges.py b/pychunkedgraph/tests/test_edges.py new file mode 100644 index 000000000..89d09d94b --- /dev/null +++ b/pychunkedgraph/tests/test_edges.py @@ -0,0 +1,355 @@ +"""Characterization tests for combining ``Edges`` instances. + +These pin the behavior of folding with ``+`` (``Edges.__add__``), which is how +``ChunkedGraph.get_l2_agglomerations`` combines per-chunk edge sets: + + all_chunk_edges = reduce( + lambda x, y: x + y, chain(edges_d.values(), fake_edges.values()), Edges([], []) + ) + +Every step of that fold reallocates all four arrays, so it is a candidate for +replacement by a single bulk concatenation. Nothing in the suite covered it before: +the tests that reach ``get_l2_agglomerations`` set ``mock_edges``, which both skips +the chunk read (leaving the fold with an empty sequence) and discards the fold's +result. These tests exist so that such a replacement is verifiable -- they describe +the behavior any bulk implementation has to reproduce, not any particular one. +""" + +from functools import reduce +from itertools import chain + +import numpy as np +import pytest + +from ..graph.edges import Edges, in_sorted +from ..graph.utils import basetypes + + +def _edges(start, count, *, with_attrs=True): + """Build a deterministic Edges of ``count`` edges, ids offset by ``start``.""" + node_ids1 = np.arange(start, start + count, dtype=basetypes.NODE_ID) + node_ids2 = np.arange(start + 1000, start + 1000 + count, dtype=basetypes.NODE_ID) + if not with_attrs: + return Edges(node_ids1, node_ids2) + return Edges( + node_ids1, + node_ids2, + affinities=np.arange(count, dtype=basetypes.EDGE_AFFINITY) + 0.5, + areas=np.arange(count, dtype=basetypes.EDGE_AREA) + 3, + ) + + +def _fold(parts): + """The exact fold used by get_l2_agglomerations.""" + return reduce(lambda x, y: x + y, chain(parts), Edges([], [])) + + +def _bulk(parts): + """Reference bulk concatenation: one np.concatenate per attribute.""" + return Edges( + np.concatenate([p.node_ids1 for p in parts] or [np.array([], dtype=basetypes.NODE_ID)]), + np.concatenate([p.node_ids2 for p in parts] or [np.array([], dtype=basetypes.NODE_ID)]), + affinities=np.concatenate( + [p.affinities for p in parts] or [np.array([], dtype=basetypes.EDGE_AFFINITY)] + ), + areas=np.concatenate( + [p.areas for p in parts] or [np.array([], dtype=basetypes.EDGE_AREA)] + ), + ) + + +def _assert_same(actual, expected): + np.testing.assert_array_equal(actual.node_ids1, expected.node_ids1) + np.testing.assert_array_equal(actual.node_ids2, expected.node_ids2) + np.testing.assert_array_equal(actual.affinities, expected.affinities) + np.testing.assert_array_equal(actual.areas, expected.areas) + + +class TestEdgesConcatenation: + def test_add_combines_all_four_arrays(self): + """``+`` must carry affinities and areas, not just the node ids.""" + a, b = _edges(0, 3), _edges(100, 2) + combined = a + b + + assert len(combined) == 5 + np.testing.assert_array_equal( + combined.node_ids1, np.concatenate([a.node_ids1, b.node_ids1]) + ) + np.testing.assert_array_equal( + combined.node_ids2, np.concatenate([a.node_ids2, b.node_ids2]) + ) + np.testing.assert_array_equal( + combined.affinities, np.concatenate([a.affinities, b.affinities]) + ) + np.testing.assert_array_equal(combined.areas, np.concatenate([a.areas, b.areas])) + + def test_add_preserves_order(self): + """Order is positional: consumers zip edges against affinities/areas.""" + a, b = _edges(0, 2), _edges(100, 2) + + assert (a + b).node_ids1.tolist() == a.node_ids1.tolist() + b.node_ids1.tolist() + assert (b + a).node_ids1.tolist() == b.node_ids1.tolist() + a.node_ids1.tolist() + + def test_add_leaves_operands_unmodified(self): + a, b = _edges(0, 3), _edges(100, 2) + before = a.node_ids1.copy() + + _ = a + b + + np.testing.assert_array_equal(a.node_ids1, before) + assert len(a) == 3 and len(b) == 2 + + def test_fold_matches_bulk_concatenation(self): + """The invariant a bulk replacement has to satisfy.""" + parts = [_edges(i * 100, i + 1) for i in range(6)] + + _assert_same(_fold(parts), _bulk(parts)) + assert len(_fold(parts)) == sum(len(p) for p in parts) + + def test_fold_of_empty_sequence(self): + """get_l2_agglomerations folds an empty chain whenever mock_edges is set.""" + folded = _fold([]) + + assert len(folded) == 0 + assert folded.node_ids1.size == 0 + assert folded.get_pairs().shape == (0, 2) + + def test_fold_of_single_element(self): + part = _edges(0, 4) + + _assert_same(_fold([part]), part) + + def test_fold_with_empty_parts_interleaved(self): + """Chunks with no edges are common; they must not perturb the result.""" + parts = [_edges(0, 2), Edges([], []), _edges(100, 3), Edges([], [])] + + _assert_same(_fold(parts), _bulk([p for p in parts if len(p)])) + assert len(_fold(parts)) == 5 + + def test_fold_materializes_defaults_for_parts_without_attrs(self): + """Edges built without affinities/areas still contribute full arrays.""" + with_attrs, without = _edges(0, 2), _edges(100, 3, with_attrs=False) + + folded = _fold([with_attrs, without]) + + assert folded.affinities.size == len(folded) + assert folded.areas.size == len(folded) + np.testing.assert_array_equal(folded.affinities[:2], with_attrs.affinities) + np.testing.assert_array_equal(folded.affinities[2:], without.affinities) + + def test_fold_preserves_dtypes(self): + """Downstream code indexes these as id/affinity/area types.""" + folded = _fold([_edges(0, 2), _edges(100, 3)]) + + assert folded.node_ids1.dtype == basetypes.NODE_ID + assert folded.node_ids2.dtype == basetypes.NODE_ID + assert folded.affinities.dtype == _edges(0, 1).affinities.dtype + assert folded.areas.dtype == _edges(0, 1).areas.dtype + + def test_get_pairs_after_fold(self): + """get_l2_agglomerations passes the folded result on as pairs.""" + parts = [_edges(0, 2), _edges(100, 3)] + + pairs = _fold(parts).get_pairs() + + assert pairs.shape == (5, 2) + np.testing.assert_array_equal(pairs[:, 0], _bulk(parts).node_ids1) + np.testing.assert_array_equal(pairs[:, 1], _bulk(parts).node_ids2) + + @pytest.mark.parametrize("count", [0, 1, 2, 10]) + def test_fold_matches_bulk_for_various_lengths(self, count): + parts = [_edges(i * 100, 2) for i in range(count)] + + folded = _fold(parts) + + assert len(folded) == 2 * count + if count: + _assert_same(folded, _bulk(parts)) + + +class TestEdgesConcatenateReplacesFold: + """Edges.concatenate replaced the reduce in get_l2_agglomerations. + + Every case above that pins the fold is re-asserted here against concatenate, so the + two are interchangeable. If they ever diverge these fail rather than the change + silently altering what get_l2_agglomerations hands to categorize_edges_v2. + """ + + @pytest.mark.parametrize("count", [0, 1, 2, 3, 10]) + def test_matches_fold_for_various_lengths(self, count): + parts = [_edges(i * 100, i + 1) for i in range(count)] + + _assert_same(Edges.concatenate(parts), _fold(parts)) + + def test_matches_fold_with_empty_parts_interleaved(self): + parts = [_edges(0, 2), Edges([], []), _edges(100, 3), Edges([], [])] + + _assert_same(Edges.concatenate(parts), _fold(parts)) + + def test_matches_fold_for_parts_without_attrs(self): + parts = [_edges(0, 2), _edges(100, 3, with_attrs=False)] + + _assert_same(Edges.concatenate(parts), _fold(parts)) + + def test_empty_input_matches_fold(self): + result = Edges.concatenate([]) + + assert len(result) == 0 + assert result.get_pairs().shape == (0, 2) + _assert_same(result, _fold([])) + + def test_preserves_dtypes(self): + result = Edges.concatenate([_edges(0, 2), _edges(100, 3)]) + + assert result.node_ids1.dtype == basetypes.NODE_ID + assert result.node_ids2.dtype == basetypes.NODE_ID + assert result.affinities.dtype == _edges(0, 1).affinities.dtype + assert result.areas.dtype == _edges(0, 1).areas.dtype + + def test_accepts_a_generator(self): + """get_l2_agglomerations passes an itertools.chain, not a list.""" + parts = [_edges(0, 2), _edges(100, 3)] + + _assert_same(Edges.concatenate(chain(parts)), _fold(parts)) + + def test_leaves_inputs_unmodified(self): + parts = [_edges(0, 3), _edges(100, 2)] + before = [p.node_ids1.copy() for p in parts] + + _ = Edges.concatenate(parts) + + for part, original in zip(parts, before): + np.testing.assert_array_equal(part.node_ids1, original) + + def test_allocates_each_output_array_once(self): + """The point of the change: n parts must not cost n concatenations.""" + parts = [_edges(i * 100, 2) for i in range(20)] + calls = [] + original = np.concatenate + + def counting_concatenate(*args, **kwargs): + calls.append(1) + return original(*args, **kwargs) + + np.concatenate = counting_concatenate + try: + Edges.concatenate(parts) + bulk_calls = len(calls) + calls.clear() + _fold(parts) + fold_calls = len(calls) + finally: + np.concatenate = original + + assert bulk_calls == 4, f"expected one concatenate per array, got {bulk_calls}" + assert fold_calls == 4 * len(parts) + + +def _svs(ids): + return np.unique(np.array(ids, dtype=basetypes.NODE_ID)) + + +class TestInSorted: + def test_matches_np_isin(self): + rng = np.random.default_rng(0) + values = rng.integers(0, 200, size=500).astype(basetypes.NODE_ID) + ref = _svs(rng.integers(0, 200, size=40)) + + np.testing.assert_array_equal(in_sorted(values, ref), np.isin(values, ref)) + + @pytest.mark.parametrize( + "values,ref", + [([], [1, 2]), ([1, 2], []), ([], []), ([5], [5]), ([5], [6])], + ) + def test_edge_cases(self, values, ref): + v, r = np.array(values, dtype=basetypes.NODE_ID), _svs(ref) + np.testing.assert_array_equal(in_sorted(v, r), np.isin(v, r)) + + def test_values_beyond_reference_range(self): + """searchsorted returns len(ref) for values past the end; must not index out of bounds.""" + v = np.array([0, 999999], dtype=basetypes.NODE_ID) + r = _svs([10, 20]) + + np.testing.assert_array_equal(in_sorted(v, r), np.array([False, False])) + + +class TestFilterTouching: + """Edges.filter_touching runs per chunk before edges are accumulated. + + The property that makes that sound: it must keep a superset of what the two + consumers keep, so filtering early cannot change their output. + """ + + def _edges(self): + # endpoints chosen to cover: both in, only first in, only second in, neither in + return Edges( + np.array([10, 20, 99, 98], dtype=basetypes.NODE_ID), + np.array([11, 97, 30, 96], dtype=basetypes.NODE_ID), + affinities=np.array([1, 2, 3, 4], dtype=basetypes.EDGE_AFFINITY), + areas=np.array([5, 6, 7, 8], dtype=basetypes.EDGE_AREA), + ) + + def test_keeps_edges_touching_the_set(self): + kept = self._edges().filter_touching(_svs([10, 11, 20, 30])) + + assert kept.node_ids1.tolist() == [10, 20, 99] + assert kept.node_ids2.tolist() == [11, 97, 30] + + def test_carries_affinities_and_areas(self): + kept = self._edges().filter_touching(_svs([10, 11, 20, 30])) + + assert kept.affinities.tolist() == [1, 2, 3] + assert kept.areas.tolist() == [5, 6, 7] + + def test_is_superset_of_categorize_predicate(self): + """categorize_edges_v2 only keeps edges whose node_ids1 is in the set.""" + e, svs = self._edges(), _svs([10, 11, 20, 30]) + kept = set(map(tuple, e.filter_touching(svs).get_pairs().tolist())) + + needed = { + tuple(p) for p in e.get_pairs().tolist() if in_sorted(np.array([p[0]], dtype=basetypes.NODE_ID), svs)[0] + } + assert needed <= kept + + def test_is_superset_of_edges_only_predicate(self): + """The edges_only path keeps edges with BOTH endpoints in the set.""" + e, svs = self._edges(), _svs([10, 11, 20, 30]) + kept = set(map(tuple, e.filter_touching(svs).get_pairs().tolist())) + + pairs = e.get_pairs() + both = pairs[np.isin(pairs[:, 0], svs) & np.isin(pairs[:, 1], svs)] + assert {tuple(p) for p in both.tolist()} <= kept + + def test_empty_set_drops_everything(self): + kept = self._edges().filter_touching(_svs([])) + + assert len(kept) == 0 + assert kept.get_pairs().shape == (0, 2) + + def test_empty_edges(self): + assert len(Edges([], []).filter_touching(_svs([1, 2]))) == 0 + + def test_all_matching_is_identity(self): + e = self._edges() + kept = e.filter_touching(_svs([10, 11, 20, 97, 99, 30, 98, 96])) + + np.testing.assert_array_equal(kept.node_ids1, e.node_ids1) + np.testing.assert_array_equal(kept.node_ids2, e.node_ids2) + + def test_filter_then_concatenate_equals_concatenate_then_filter(self): + """Per-chunk filtering must equal filtering the fully accumulated set.""" + rng = np.random.default_rng(7) + svs = _svs(rng.integers(0, 50, size=12)) + chunks = [ + Edges( + rng.integers(0, 100, size=30).astype(basetypes.NODE_ID), + rng.integers(0, 100, size=30).astype(basetypes.NODE_ID), + ) + for _ in range(5) + ] + + early = Edges.concatenate([c.filter_touching(svs) for c in chunks]) + late = Edges.concatenate(chunks).filter_touching(svs) + + np.testing.assert_array_equal(early.node_ids1, late.node_ids1) + np.testing.assert_array_equal(early.node_ids2, late.node_ids2) diff --git a/pychunkedgraph/tests/test_io_edges.py b/pychunkedgraph/tests/test_io_edges.py new file mode 100644 index 000000000..de1480c8d --- /dev/null +++ b/pychunkedgraph/tests/test_io_edges.py @@ -0,0 +1,227 @@ +"""Tests for the chunk-edge read path. + +get_chunk_edges fetches and decompresses chunk edge files in batches and, when given +the queried object's supervoxels, filters each chunk as it is parsed. Both behaviors +exist to bound peak memory, so what matters is that neither changes the result: the +output must be identical to reading everything at once and filtering at the end. +""" + +import numpy as np +import pytest +import zstandard as zstd + +from ..io import edges as io_edges +from ..io.protobuf.chunkEdges_pb2 import ChunkEdgesMsg +from ..graph.edges import Edges +from ..graph.edges import EDGE_TYPES +from ..graph.utils import basetypes + + +def _edges(pairs): + a = np.array([p[0] for p in pairs], dtype=basetypes.NODE_ID) + b = np.array([p[1] for p in pairs], dtype=basetypes.NODE_ID) + return Edges( + a, + b, + affinities=np.arange(len(pairs), dtype=basetypes.EDGE_AFFINITY) + 1, + areas=np.arange(len(pairs), dtype=basetypes.EDGE_AREA) + 2, + ) + + +def _blob(in_chunk, between, cross): + msg = ChunkEdgesMsg() + msg.in_chunk.CopyFrom(io_edges.serialize(in_chunk)) + msg.between_chunk.CopyFrom(io_edges.serialize(between)) + msg.cross_chunk.CopyFrom(io_edges.serialize(cross)) + return zstd.ZstdCompressor().compress(msg.SerializeToString()) + + +class FakeCloudFiles: + """Records each get() so batching is observable.""" + + def __init__(self, blobs): + self.blobs = blobs + self.calls = [] + + def __call__(self, *args, **kwargs): + return self + + def get(self, fnames, raw=False): + self.calls.append(list(fnames)) + return [{"content": self.blobs.get(f)} for f in fnames] + + +@pytest.fixture +def chunks(monkeypatch): + """10 chunks; each has one edge touching sv 5 and two that do not.""" + blobs, coords = {}, [] + for i in range(10): + base = 1000 * (i + 1) + blobs[f"edges_{i}_0_0.proto.zst"] = _blob( + _edges([(5, base), (base + 1, base + 2)]), + _edges([(base + 3, 5)]), + _edges([(base + 4, base + 5)]), + ) + coords.append(np.array([i, 0, 0])) + fake = FakeCloudFiles(blobs) + monkeypatch.setattr(io_edges, "CloudFiles", fake) + return fake, coords + + +def _all_pairs(result): + return sorted( + tuple(p) for t in EDGE_TYPES for p in result[t].get_pairs().tolist() + ) + + +class TestBatching: + @pytest.mark.parametrize("batch_size", [1, 3, 7, 64, 1000]) + def test_result_is_independent_of_batch_size(self, chunks, batch_size): + fake, coords = chunks + expected = _all_pairs(io_edges.get_chunk_edges("gs://x", coords, batch_size=1000)) + + got = _all_pairs(io_edges.get_chunk_edges("gs://x", coords, batch_size=batch_size)) + + assert got == expected + + @pytest.mark.parametrize("batch_size,expected_calls", [(1, 10), (3, 4), (5, 2), (64, 1)]) + def test_fetches_in_batches(self, chunks, batch_size, expected_calls): + fake, coords = chunks + fake.calls.clear() + + io_edges.get_chunk_edges("gs://x", coords, batch_size=batch_size) + + assert len(fake.calls) == expected_calls + assert max(len(c) for c in fake.calls) <= batch_size + assert sum(len(c) for c in fake.calls) == len(coords) + + def test_every_file_requested_exactly_once(self, chunks): + fake, coords = chunks + fake.calls.clear() + + io_edges.get_chunk_edges("gs://x", coords, batch_size=3) + + requested = [f for call in fake.calls for f in call] + assert sorted(requested) == sorted(fake.blobs) + + def test_missing_chunk_files_are_skipped(self, chunks): + fake, coords = chunks + fake.blobs["edges_2_0_0.proto.zst"] = None + + result = io_edges.get_chunk_edges("gs://x", coords, batch_size=3) + + assert all((3000, 3001) != p for p in _all_pairs(result)) + + +class TestFiltering: + def test_filter_matches_filtering_after_the_fact(self, chunks): + fake, coords = chunks + svs = np.array([5], dtype=basetypes.NODE_ID) + + filtered = io_edges.get_chunk_edges("gs://x", coords, supervoxels=svs) + unfiltered = io_edges.get_chunk_edges("gs://x", coords) + expected = { + t: unfiltered[t].filter_touching(np.unique(svs)) for t in EDGE_TYPES + } + + for t in EDGE_TYPES: + np.testing.assert_array_equal(filtered[t].node_ids1, expected[t].node_ids1) + np.testing.assert_array_equal(filtered[t].node_ids2, expected[t].node_ids2) + np.testing.assert_array_equal(filtered[t].affinities, expected[t].affinities) + np.testing.assert_array_equal(filtered[t].areas, expected[t].areas) + + def test_filter_keeps_only_touching_edges(self, chunks): + fake, coords = chunks + svs = np.array([5], dtype=basetypes.NODE_ID) + + result = io_edges.get_chunk_edges("gs://x", coords, supervoxels=svs) + + pairs = _all_pairs(result) + assert len(pairs) == 20 # one in_chunk + one between_chunk per chunk + assert all(5 in p for p in pairs) + + @pytest.mark.parametrize("batch_size", [1, 3, 64]) + def test_filter_is_independent_of_batch_size(self, chunks, batch_size): + fake, coords = chunks + svs = np.array([5], dtype=basetypes.NODE_ID) + + got = _all_pairs( + io_edges.get_chunk_edges("gs://x", coords, supervoxels=svs, batch_size=batch_size) + ) + + assert got == _all_pairs( + io_edges.get_chunk_edges("gs://x", coords, supervoxels=svs, batch_size=1000) + ) + + def test_no_supervoxels_reads_everything(self, chunks): + fake, coords = chunks + + result = io_edges.get_chunk_edges("gs://x", coords) + + assert len(_all_pairs(result)) == 40 # 4 edges x 10 chunks + + def test_unsorted_and_duplicated_supervoxels_are_normalized(self, chunks): + fake, coords = chunks + messy = np.array([5, 5, 5], dtype=basetypes.NODE_ID) + + got = _all_pairs(io_edges.get_chunk_edges("gs://x", coords, supervoxels=messy)) + + assert got == _all_pairs( + io_edges.get_chunk_edges( + "gs://x", coords, supervoxels=np.array([5], dtype=basetypes.NODE_ID) + ) + ) + + +class TestDecompression: + """_decompress replaced multi_decompress_to_buffer (removed in zstandard 0.23). + + A thread pool matches its throughput because decompression releases the GIL, but + ZstdDecompressor is not thread-safe, so correctness under threads is what these pin. + """ + + def _blobs(self, n=32, size=200_000): + rng = np.random.default_rng(3) + cctx = zstd.ZstdCompressor(level=3) + raw = [rng.integers(0, 255, size=size, dtype=np.uint8).tobytes() for _ in range(n)] + return raw, [cctx.compress(r) for r in raw] + + @pytest.mark.parametrize("n_threads", [1, 2, 4, 8]) + def test_matches_input_for_any_thread_count(self, n_threads): + raw, blobs = self._blobs() + + out = io_edges._decompress(blobs, n_threads) + + assert [bytes(o) for o in out] == raw + + def test_threaded_matches_serial(self): + """Sharing one ZstdDecompressor across threads corrupts silently; this catches it.""" + raw, blobs = self._blobs() + + assert [bytes(o) for o in io_edges._decompress(blobs, 4)] == [ + bytes(o) for o in io_edges._decompress(blobs, 1) + ] + + def test_preserves_order(self): + raw, blobs = self._blobs(n=16) + + out = [bytes(o) for o in io_edges._decompress(blobs, 4)] + + assert out == raw # pool.map must not reorder + + @pytest.mark.parametrize("n_threads", [1, 4]) + def test_empty_and_single(self, n_threads): + raw, blobs = self._blobs(n=1) + + assert io_edges._decompress([], n_threads) == [] + assert [bytes(o) for o in io_edges._decompress(blobs, n_threads)] == raw + + @pytest.mark.parametrize("n_threads", [1, 4]) + def test_frames_without_content_size(self, n_threads): + """dctx.decompress needs the size in the header; decompressobj covers frames without it.""" + rng = np.random.default_rng(4) + raw = [rng.integers(0, 255, size=50_000, dtype=np.uint8).tobytes() for _ in range(4)] + cctx = zstd.ZstdCompressor(level=3, write_content_size=False) + blobs = [cctx.compress(r) for r in raw] + + assert [bytes(o) for o in io_edges._decompress(blobs, n_threads)] == raw diff --git a/uwsgi.ini b/uwsgi.ini index 776e2ff00..c6ba30ba5 100644 --- a/uwsgi.ini +++ b/uwsgi.ini @@ -57,6 +57,14 @@ buffer-size = 65535 # Don't spawn new workers if total memory over 6 GiB cheaper-rss-limit-soft = 6442450944 +# Gracefully recycle a worker once its RSS exceeds this many MB: uwsgi lets it finish the +# current request, then respawns it. Bounds per-worker memory growth so a bloated worker +# can't accumulate toward the pod memory limit (set in helm to ~1.2x the request). +# NOTE: this is graceful (post-request); it does NOT stop a single request that balloons +# memory mid-flight -- the pod memory limit (OOM) is the backstop for that. For a forceful +# mid-request kill instead, use `evil-reload-on-rss`. Tune to observed per-worker RSS. +reload-on-rss = 768 + # Reload worker after serving X requests max-requests = 5000