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/meshing/common.py b/pychunkedgraph/app/meshing/common.py index 8f1a0c20a..43972350d 100644 --- a/pychunkedgraph/app/meshing/common.py +++ b/pychunkedgraph/app/meshing/common.py @@ -1,18 +1,16 @@ # pylint: disable=invalid-name, missing-docstring import json import os -import threading import numpy as np import redis from rq import Queue, Connection, Retry -from flask import Response, current_app, jsonify, make_response, request +from flask import Response, current_app, g, jsonify, make_response, request from pychunkedgraph import __version__ from pychunkedgraph.app import app_utils -from pychunkedgraph.graph import chunkedgraph +from pychunkedgraph.graph import exceptions as cg_exceptions from pychunkedgraph.app.meshing import tasks as meshing_tasks -from pychunkedgraph.meshing import meshgen from pychunkedgraph.meshing.manifest import get_highest_child_nodes_with_meshes from pychunkedgraph.meshing.manifest import get_children_before_start_layer from pychunkedgraph.meshing.manifest import ManifestCache @@ -75,9 +73,20 @@ def handle_get_manifest(table_id, node_id): return_seg_ids = return_seg_ids in ["True", "true", "1", True] prepend_seg_ids = prepend_seg_ids in ["True", "true", "1", True] start_layer = cg.meta.custom_data.get("mesh", {}).get("max_layer", 2) - start_layer = int(request.args.get("start_layer", start_layer)) - if "start_layer" in data: - start_layer = int(data["start_layer"]) + raw_start_layer = data.get("start_layer", request.args.get("start_layer", start_layer)) + try: + start_layer = int(raw_start_layer) + except (TypeError, ValueError): + raise cg_exceptions.BadRequest( + f"start_layer must be an integer, got {raw_start_layer!r}." + ) + + # Meshes only exist from layer 2 upwards. Below that this endpoint cannot return anything: + if start_layer < 2: + raise cg_exceptions.BadRequest( + f"start_layer must be at least 2, got {start_layer}. Meshes exist from layer 2 " + "upwards, so a lower value can only ever produce an empty manifest." + ) flexible_start_layer = None if "flexible_start_layer" in data: @@ -142,9 +151,64 @@ def _check_post_options(cg, resp, data, seg_ids): ## REMESHING ----------------------------------------------------- +def publish_remesh(table_id: str, user_id: str, lvl2_ids, is_priority: bool = True): + """Enqueue a remesh onto the same Pub/Sub topic the edit path publishes to. + + Mirrors segmentation.common.publish_edit deliberately: one topic, and the + `remesh_priority` attribute is what routes a message to a subscription. The + infrastructure defines those subscriptions with attribute filters + (terraform-google-cave/modules/local_cluster/pubsub.tf): + + _PCG_HIGH_PRIORITY_REMESH remesh_priority="true" -> meshworker + _PCG_LOW_PRIORITY_REMESH remesh_priority="false" -> remeshworker + + so priority here is not a hint, it selects the consumer fleet. + + Note the same topic also feeds __L2CACHE_{HIGH,LOW}_PRIORITY_TRIGGER, so a + manual remesh now also refreshes the l2 cache for these ids. That is intended -- a manual + remesh usually follows a data problem, and the l2 cache derives from the same chunks -- but + it is a real fan-out, not a no-op. + """ + import pickle + + from messagingclient import MessagingClient + + attributes = { + "table_id": table_id, + "user_id": user_id, + "remesh_priority": "true" if is_priority else "false", + "remesh": "true", + } + payload = { + # 0 means "no operation". A manual remesh has no GraphEditOperation behind it, and the + # graph does not record which operation created a given level 2 node -- OperationID is + # written only onto root-id rows (edits.py::_update_root_id_lineage), so the best available + # answer is the latest operation on the whole object, which is not this node's provenance. + # A wrong id in the worker's log line is worse than an honest unknown. + # + # The key must still exist and be int-convertible: mesh_worker.callback does + # int(data["operation_id"]) unconditionally. + "operation_id": 0, + "new_lvl2_ids": np.asarray(lvl2_ids, dtype=np.uint64).tolist(), + # Neither consumer reads these (mesh_worker uses new_lvl2_ids, the l2cache trigger uses + # new_lvl2_ids); present so the payload shape stays identical to publish_edit's. + "new_root_ids": [], + "old_root_ids": [], + } + + exchange = os.getenv("PYCHUNKEDGRAPH_EDITS_EXCHANGE", "pychunkedgraph") + c = MessagingClient() + c.publish(exchange, pickle.dumps(payload), attributes) + + def handle_remesh(table_id): current_app.request_type = "remesh_enque" current_app.table_id = table_id + # Same `priority` parameter, default, and semantics as every edit endpoint in + # segmentation.common, so the two paths cannot drift. Unset means high priority, which is + # the right default for an interactive request; a programmatic caller (caveclient, a + # backfill script) should pass priority=false so bulk work lands on the low-priority + # subscription and cannot starve human-triggered remeshes. is_priority = request.args.get("priority", True, type=str2bool) is_redisjob = request.args.get("use_redis", False, type=str2bool) @@ -166,38 +230,23 @@ def handle_remesh(table_id): return jsonify(response_object), 202 else: + # Publish, don't mesh here. This used to run meshgen.remeshing in a threading.Thread + # inside the api pod, which put an unbounded, unretryable, invisible workload in a + # request-serving process: the 202 was already returned, so a failure left no trace, and + # a worker recycle, rollout or HPA scale-down silently discarded the work. Measured on + # api6 2026-08-23, one remesh took the meshing pod from 198Mi to a 491Mi peak and left + # it at 420Mi -- rss does not fall back -- so pods ratcheted up until they died: one had + # reached 847Mi after 14h and was OOMKilled (whole cgroup, supervisord included) when a + # remesh pushed it past its 1536Mi limit. The mesh workers exist for exactly this work, + # request 3000Mi, and get retries and dead-lettering from Pub/Sub. new_lvl2_ids = np.array(new_lvl2_ids, dtype=np.uint64) - cg = app_utils.get_cg(table_id) if len(new_lvl2_ids) > 0: - t = threading.Thread( - target=_remeshing, args=(cg.get_serialized_info(), new_lvl2_ids) - ) - t.start() + user_id = str(g.auth_user.get("id", current_app.user_id)) + publish_remesh(table_id, user_id, new_lvl2_ids, is_priority=is_priority) return Response(status=202) - -def _remeshing(serialized_cg_info, lvl2_nodes): - cg = chunkedgraph.ChunkedGraph(**serialized_cg_info) - cv_mesh_dir = cg.meta.dataset_info["mesh"] - cv_unsharded_mesh_dir = cg.meta.dataset_info["mesh_metadata"]["unsharded_mesh_dir"] - cv_unsharded_mesh_path = os.path.join( - cg.meta.data_source.WATERSHED, cv_mesh_dir, cv_unsharded_mesh_dir - ) - mesh_data = cg.meta.custom_data["mesh"] - - # TODO: stop_layer and mip should be configurable by dataset - meshgen.remeshing( - cg, - lvl2_nodes, - stop_layer=mesh_data["max_layer"], - mip=mesh_data["mip"], - max_err=mesh_data["max_error"], - cv_sharded_mesh_dir=cv_mesh_dir, - cv_unsharded_mesh_path=cv_unsharded_mesh_path, - ) - return Response(status=200) 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..e2ccffddc 100644 --- a/pychunkedgraph/graph/analysis/pathing.py +++ b/pychunkedgraph/graph/analysis/pathing.py @@ -1,11 +1,13 @@ +import os import typing -import fastremap import graph_tool import numpy as np from pychunkedgraph.graph.utils import flatgraph +from .. import exceptions as cg_exceptions +from ..edges import in_sorted from ..subgraph import get_subgraph_nodes @@ -77,12 +79,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,73 +105,205 @@ 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 -def _get_edges_for_lvl2_ids(cg, lvl2_ids, induced=False): +# Level 2 ids per call to cg.get_atomic_cross_edges. The read is the one thing here that +# scales with object size rather than with the answer, so it is sliced. Swept on the +# profiled object below (peak rss): 3k/5k/10k all land within noise of each other, 25k +# drifts up. Flat enough that this is not a knife edge, so the default rarely needs +# changing -- PCG_CROSS_EDGE_READ_BATCH is there for tuning a deployment whose objects or +# pod memory differ, without a redeploy of the code. +try: + CROSS_EDGE_READ_BATCH = int(os.environ.get("PCG_CROSS_EDGE_READ_BATCH", 10_000)) +except ValueError: + CROSS_EDGE_READ_BATCH = 10_000 +if CROSS_EDGE_READ_BATCH < 1: + # Zero or negative would make the range() below yield no slices at all, so the + # function would return an empty edge list instead of failing: a wrong answer rather + # than an error. Ignore it and keep the default. + CROSS_EDGE_READ_BATCH = 10_000 + + +def _sorted_by_key(keys, values): + """Sort a key/value pair of arrays together, by key. + + Each slice contributes a sorted fragment but their ranges interleave, so the + concatenation is not sorted. Both the miss scan and the remap binary-search it. + """ + order = np.argsort(keys, kind="stable") + return keys[order], values[order] + + +def _get_edges_for_lvl2_ids(cg, lvl2_ids, induced=False, batch_size=None): + """Cross-chunk edges between `lvl2_ids`, remapped from supervoxels to level 2 ids. + + Memory, not speed, is the binding constraint: the intermediates dwarf the answer. + Profiled on api6 for aibs_v1dd root 864691132764660103, the object that trips uwsgi's + reload-on-rss: 240,928 level 2 ids whose cross edges are 18,349,472 supervoxel pairs + (280MB), collapsing to 276,472 level 2 edges (4.4MB). + + The costly part is not the edges but the dict they arrive in -- one small array per + (level 2 id, layer), ~480k of them, 930MB for 280MB of data, nearly all numpy object + overhead. So the read is sliced and each slice is concatenated and its dict dropped + before the next is fetched, which keeps only one slice's worth of that overhead alive. + The edges themselves are all still held, so this is still a single pass over Bigtable. + + Measured in the read pod against the previous single-shot version: peak rss + 1861-1887MB -> 933-952MB, wall clock 33.3-35.3s -> 30.3-31.0s, output byte-identical. + It gets faster rather than slower because the working set shrinks and because of the + miss scan below. + + The read slice size defaults to CROSS_EDGE_READ_BATCH (the PCG_CROSS_EDGE_READ_BATCH + environment variable, 10,000) and can be overridden per call with `batch_size`. It + changes the peak and nothing else -- the result is independent of it. + + Two things that look like obvious wins here and are not, both measured: + - Slicing the *reduce* as well, so the edges need not be held either, drops the peak + to ~700MB but needs a second pass over Bigtable and costs 7s. Not worth it. + - Swapping fastremap for searchsorted on its own, without slicing the read, is a + wash (1834-1869MB). It pays only in combination, because the reduce loop below + would otherwise rebuild fastremap's 6.95M-key table once per slice. + """ # protect in case there are no lvl2 ids if len(lvl2_ids) == 0: return np.empty((0, 2), dtype=np.uint64) - cce_dict = cg.get_atomic_cross_edges(lvl2_ids) + # Sorted and deduplicated once, up front. It is the reference set for the induced + # filter below, and slicing it keeps each supervoxel's mapping entry in exactly one + # slice. Reading a dict keyed by id already collapsed duplicates, so this does not + # change which edges are fetched. + lvl2_ids = np.unique(lvl2_ids) + if batch_size is None: + batch_size = CROSS_EDGE_READ_BATCH + elif batch_size < 1: + # Same trap as above, reached explicitly instead of through the environment. + raise ValueError(f"batch_size must be positive, got {batch_size}") + + batches = [] + map_keys = [] + map_values = [] + partner_ids = [] + for start in range(0, len(lvl2_ids), batch_size): + cce_dict = cg.get_atomic_cross_edges(lvl2_ids[start : start + batch_size]) + blocks = [] + owners = [] + lengths = [] + for lvl2_id in cce_dict: + for level in cce_dict[lvl2_id]: + block = cce_dict[lvl2_id][level] + if len(block) == 0: + continue + blocks.append(block) + owners.append(lvl2_id) + lengths.append(len(block)) + # The whole point: the slice's edges are copied out by the concatenate, so its + # dict can go before the next read allocates another one. + cce_dict.clear() + if not blocks: + continue + edges = np.concatenate(blocks) + del blocks - # Gather all of the supervoxel ids into two lists, we will map them to - # their parent lvl2 ids - edge_array = [] - for l2_id in cce_dict: - for level in cce_dict[l2_id]: - edge_array.append(cce_dict[l2_id][level]) + # Column 0 is a supervoxel of the block's owner, column 1 its partner across the + # chunk boundary; both are views into `edges`. One np.repeat gives the parent of + # every row, rather than an np.full per block. + owner_column = np.repeat( + np.array(owners, dtype=edges.dtype), np.array(lengths, dtype=np.int64) + ) + keys, first = np.unique(edges[:, 0], return_index=True) + map_keys.append(keys) + map_values.append(owner_column[first]) + partner_ids.append(np.unique(edges[:, 1])) + batches.append(edges) + del edges, owner_column, keys, first # protect in case there are no edges - if len(edge_array) == 0: + if not batches: return np.empty((0, 2), dtype=np.uint64) - - edge_array = np.concatenate(edge_array) - known_supervoxels_list = [] - known_l2_list = [] - unknown_supervoxel_list = [] - for lvl2_id in cce_dict: - for level in cce_dict[lvl2_id]: - known_supervoxels_for_lv2_id = cce_dict[lvl2_id][level][:, 0] - unknown_supervoxels_for_lv2_id = cce_dict[lvl2_id][level][:, 1] - known_supervoxels_list.append(known_supervoxels_for_lv2_id) - known_l2_list.append(np.full(known_supervoxels_for_lv2_id.shape, lvl2_id)) - unknown_supervoxel_list.append(unknown_supervoxels_for_lv2_id) - - # Create two arrays to map supervoxels for which we know their parents - known_supervoxel_array, unique_indices = np.unique( - np.concatenate(known_supervoxels_list), return_index=True - ) - known_l2_array = (np.concatenate(known_l2_list))[unique_indices] - unknown_supervoxel_array = np.unique(np.concatenate(unknown_supervoxel_list)) - # Call get_parents on any supervoxels for which we don't know their parents - supervoxels_to_query_parent = np.setdiff1d( - unknown_supervoxel_array, known_supervoxel_array + # A supervoxel belongs to exactly one level 2 node, so the fragments are disjoint. + known_supervoxel_array = np.concatenate(map_keys) + known_l2_array = np.concatenate(map_values) + del map_keys, map_values + known_supervoxel_array, known_l2_array = _sorted_by_key( + known_supervoxel_array, known_l2_array ) - if len(supervoxels_to_query_parent) > 0: + + # Partners with no mapping entry belong to level 2 nodes outside `lvl2_ids` and have + # to be looked up. For a whole object there are none -- cross edges are stored from + # both sides, so every partner also appears in some slice's column 0 -- but for a + # subset (a 'bounds' query, or a shared parent below the root) there are. + # + # Scanned one slice at a time on purpose. Concatenating every slice's partners and + # handing that to np.setdiff1d, which sorts both sides again, cost 477MB on the + # profiled object to produce an empty answer. + misses = [] + for partners in partner_ids: + missed = partners[~in_sorted(partners, known_supervoxel_array)] + if len(missed): + misses.append(missed) + del partner_ids + if misses: + supervoxels_to_query_parent = np.unique(np.concatenate(misses)) + del misses missing_l2_ids = cg.get_parents(supervoxels_to_query_parent) known_supervoxel_array = np.concatenate( (known_supervoxel_array, supervoxels_to_query_parent) ) known_l2_array = np.concatenate((known_l2_array, missing_l2_ids)) + del supervoxels_to_query_parent, missing_l2_ids + known_supervoxel_array, known_l2_array = _sorted_by_key( + known_supervoxel_array, known_l2_array + ) - # Map the cross-chunk edges from supervoxels to lvl2 ids - edge_view = edge_array.view() - edge_view.shape = -1 - fastremap.remap_from_array_kv(edge_view, known_supervoxel_array, known_l2_array) - - edge_array = np.unique(np.sort(edge_array, axis=1), axis=0) - - if induced: - # make this an induced subgraph - # keep only the edges that are between the lvl2 ids asked for - edge_array = edge_array[ - np.isin(edge_array[:, 0], lvl2_ids) & np.isin(edge_array[:, 1], lvl2_ids) - ] + # Reduce each slice to level 2 pairs and free it before moving on, so the 66:1 + # collapse happens per slice and the accumulator stays at answer size. + reduced = [] + while batches: + edges = batches.pop() + # Every supervoxel here is in the mapping -- column 0 by construction, column 1 + # via the get_parents step above -- so a binary search is exact. searchsorted + # rather than fastremap.remap_from_array_kv, which would rebuild a lookup over + # the whole mapping on every iteration of this loop. + flat = edges.reshape(-1) + flat[:] = known_l2_array[np.searchsorted(known_supervoxel_array, flat)] + del flat + # In place: `edges` is ours, np.concatenate having copied it out of the dict. + edges.sort(axis=1) + if induced: + # make this an induced subgraph + # keep only the edges that are between the lvl2 ids asked for. Filtering + # before the dedup is worth it when it bites -- a subset query -- and costs + # little when it does not. in_sorted rather than np.isin, which would re-sort + # lvl2_ids on each of the two calls. + edges = edges[ + in_sorted(edges[:, 0], lvl2_ids) & in_sorted(edges[:, 1], lvl2_ids) + ] + reduced.append(np.unique(edges, axis=0)) + del edges - return edge_array + del known_supervoxel_array, known_l2_array + if len(reduced) == 1: + return reduced[0] + # Slices reduce independently, so the same level 2 pair can come out of more than one + # of them; this is what makes the result independent of batch_size. + return np.unique(np.concatenate(reduced), axis=0) def find_l2_shortest_path( 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/pychunkedgraph/tests/test_lvl2_edges.py b/pychunkedgraph/tests/test_lvl2_edges.py new file mode 100644 index 000000000..b12e0c720 --- /dev/null +++ b/pychunkedgraph/tests/test_lvl2_edges.py @@ -0,0 +1,296 @@ +"""Equivalence tests for ``_get_edges_for_lvl2_ids``. + +The function was reworked to cut its peak memory -- measured in the api6 read pod on +aibs_v1dd root 864691132764660103, peak rss 1868-1878MB -> 991-1031MB, and slightly +faster with it. Every change was meant to be behaviour-preserving, so that is what these +tests pin: a reference implementation of the previous algorithm runs beside the current +one over randomized inputs, and the two must agree element for element. The reference is +deliberately written the slow, obvious way -- it is the specification, not an +optimization. + +The read is now sliced, so the tests also cover the slicing itself: that the result does +not depend on the slice size, that every id is read exactly once (one Bigtable pass, not +two), and that a partner living in a different slice is still resolved. +""" + +import importlib + +import numpy as np +import pytest + +from ..graph.analysis.pathing import _get_edges_for_lvl2_ids + + +class FakeCg: + """Just the two methods the function calls on a ChunkedGraph.""" + + def __init__(self, cce_dict, parents=None): + self._cce = cce_dict + self._parents = parents or {} + self.get_parents_calls = 0 + self.read_calls = 0 + self.ids_read = [] + + def get_atomic_cross_edges(self, l2_ids): + # A fresh dict per call holding only the ids asked for, as the real one does. + # Honouring l2_ids matters: the implementation calls this once per slice, and a + # fake that ignored the argument would hand every slice the whole object, so no + # slicing bug could ever fail a test. + wanted = {int(i) for i in l2_ids} + self.read_calls += 1 + self.ids_read.extend(sorted(wanted)) + return {k: dict(v) for k, v in self._cce.items() if int(k) in wanted} + + def get_parents(self, supervoxels): + self.get_parents_calls += 1 + return np.array([self._parents[int(s)] for s in supervoxels], dtype=np.uint64) + + +def reference(cg, lvl2_ids, induced=False): + """The previous implementation, verbatim in structure.""" + import fastremap + + if len(lvl2_ids) == 0: + return np.empty((0, 2), dtype=np.uint64) + cce_dict = cg.get_atomic_cross_edges(lvl2_ids) + edge_array = [] + for l2_id in cce_dict: + for level in cce_dict[l2_id]: + edge_array.append(cce_dict[l2_id][level]) + if len(edge_array) == 0: + return np.empty((0, 2), dtype=np.uint64) + edge_array = np.concatenate(edge_array) + known_sv, known_l2, unknown_sv = [], [], [] + for lvl2_id in cce_dict: + for level in cce_dict[lvl2_id]: + k = cce_dict[lvl2_id][level][:, 0] + u = cce_dict[lvl2_id][level][:, 1] + known_sv.append(k) + known_l2.append(np.full(k.shape, lvl2_id)) + unknown_sv.append(u) + known_supervoxel_array, unique_indices = np.unique( + np.concatenate(known_sv), return_index=True + ) + known_l2_array = (np.concatenate(known_l2))[unique_indices] + unknown_supervoxel_array = np.unique(np.concatenate(unknown_sv)) + to_query = np.setdiff1d(unknown_supervoxel_array, known_supervoxel_array) + if len(to_query) > 0: + missing = cg.get_parents(to_query) + known_supervoxel_array = np.concatenate((known_supervoxel_array, to_query)) + known_l2_array = np.concatenate((known_l2_array, missing)) + ev = edge_array.view() + ev.shape = -1 + fastremap.remap_from_array_kv(ev, known_supervoxel_array, known_l2_array) + edge_array = np.unique(np.sort(edge_array, axis=1), axis=0) + if induced: + edge_array = edge_array[ + np.isin(edge_array[:, 0], lvl2_ids) & np.isin(edge_array[:, 1], lvl2_ids) + ] + return edge_array + + +def build_case(seed, n_l2=60, layers=(2, 3), max_edges=4, outside_frac=0.3): + """Random cross-edge dict plus the parent lookup for supervoxels outside it. + + Supervoxel ids are laid out so each belongs to exactly one level 2 node, which is what + the real data guarantees and what the supervoxel -> parent mapping relies on. + """ + rng = np.random.default_rng(seed) + lvl2_ids = np.arange(1, n_l2 + 1, dtype=np.uint64) * 1000 + # supervoxel s belongs to lvl2 node (s // 1000) * 1000 + def svs_of(l2, count): + return (np.uint64(l2) + rng.integers(1, 999, size=count)).astype(np.uint64) + + cce, parents = {}, {} + outside_l2 = np.arange(n_l2 + 1, n_l2 + 21, dtype=np.uint64) * 1000 + for l2 in lvl2_ids: + d = {} + for layer in layers: + n = int(rng.integers(0, max_edges + 1)) + if n == 0: + d[layer] = np.empty((0, 2), dtype=np.uint64) + continue + col0 = svs_of(l2, n) + partner_l2 = np.where( + rng.random(n) < outside_frac, + rng.choice(outside_l2, size=n), + rng.choice(lvl2_ids, size=n), + ).astype(np.uint64) + col1 = np.array( + [int(p) + int(rng.integers(1, 999)) for p in partner_l2], dtype=np.uint64 + ) + for sv, p in zip(col1.tolist(), partner_l2.tolist()): + parents[int(sv)] = np.uint64(p) + d[layer] = np.column_stack([col0, col1]) + cce[l2] = d + # a supervoxel that appears in column 0 has a known parent; make sure the reference + # and the implementation agree about those too + for l2, d in cce.items(): + for layer, block in d.items(): + for sv in block[:, 0].tolist(): + parents.setdefault(int(sv), np.uint64(l2)) + return cce, parents, lvl2_ids + + +@pytest.mark.parametrize("seed", range(12)) +@pytest.mark.parametrize("induced", [True, False]) +def test_matches_previous_implementation(seed, induced): + cce, parents, lvl2_ids = build_case(seed) + got = _get_edges_for_lvl2_ids(FakeCg(cce, parents), lvl2_ids, induced=induced) + want = reference(FakeCg(cce, parents), lvl2_ids, induced=induced) + assert got.dtype == want.dtype + assert np.array_equal(got, want), f"seed={seed} induced={induced}" + + +@pytest.mark.parametrize("seed", [0, 5]) +def test_induced_edges_stay_inside_the_object(seed): + cce, parents, lvl2_ids = build_case(seed) + got = _get_edges_for_lvl2_ids(FakeCg(cce, parents), lvl2_ids, induced=True) + assert np.isin(got, lvl2_ids).all() + + +def test_result_is_deduplicated_and_sorted_within_each_edge(): + cce, parents, lvl2_ids = build_case(1) + got = _get_edges_for_lvl2_ids(FakeCg(cce, parents), lvl2_ids, induced=True) + assert (got[:, 0] <= got[:, 1]).all(), "each edge should be sorted low, high" + assert len(np.unique(got, axis=0)) == len(got), "no duplicate edges" + + +def test_no_lvl2_ids(): + out = _get_edges_for_lvl2_ids(FakeCg({}, {}), np.array([], dtype=np.uint64)) + assert out.shape == (0, 2) and out.dtype == np.uint64 + + +def test_lvl2_ids_but_no_edges(): + lvl2_ids = np.array([1000, 2000], dtype=np.uint64) + cce = {np.uint64(1000): {2: np.empty((0, 2), dtype=np.uint64)}, np.uint64(2000): {}} + out = _get_edges_for_lvl2_ids(FakeCg(cce, {}), lvl2_ids, induced=True) + assert out.shape == (0, 2) and out.dtype == np.uint64 + + +def test_parents_are_fetched_only_for_unknown_supervoxels(): + """Supervoxels seen in column 0 must not be re-queried against the graph.""" + cce, parents, lvl2_ids = build_case(3) + cg = FakeCg(cce, parents) + _get_edges_for_lvl2_ids(cg, lvl2_ids, induced=True) + assert cg.get_parents_calls == 1 + + +def test_caller_lvl2_ids_not_mutated(): + """np.unique / sort inside must not reorder the array the caller passed in.""" + cce, parents, lvl2_ids = build_case(7) + shuffled = lvl2_ids.copy() + np.random.default_rng(0).shuffle(shuffled) + before = shuffled.copy() + _get_edges_for_lvl2_ids(FakeCg(cce, parents), shuffled, induced=True) + assert np.array_equal(shuffled, before) + + +def test_duplicate_lvl2_ids_in_input(): + cce, parents, lvl2_ids = build_case(2) + dup = np.concatenate([lvl2_ids, lvl2_ids[:10]]) + got = _get_edges_for_lvl2_ids(FakeCg(cce, parents), dup, induced=True) + want = _get_edges_for_lvl2_ids(FakeCg(cce, parents), lvl2_ids, induced=True) + assert np.array_equal(got, want) + + +@pytest.mark.parametrize("batch_size", [1, 2, 7, 59, 10_000]) +@pytest.mark.parametrize("induced", [True, False]) +def test_result_is_independent_of_batch_size(batch_size, induced): + """Slicing the read must change the peak memory and nothing else.""" + cce, parents, lvl2_ids = build_case(4) + want = reference(FakeCg(cce, parents), lvl2_ids, induced=induced) + got = _get_edges_for_lvl2_ids( + FakeCg(cce, parents), lvl2_ids, induced=induced, batch_size=batch_size + ) + assert np.array_equal(got, want), f"batch_size={batch_size} induced={induced}" + + +def test_read_is_sliced_and_covers_every_id_once(): + """One read per slice, every id read exactly once -- a single pass over Bigtable. + + Asserted directly because a slicing bug that skipped or repeated ids would still + usually produce plausible-looking output. + """ + cce, parents, lvl2_ids = build_case(6) + cg = FakeCg(cce, parents) + _get_edges_for_lvl2_ids(cg, lvl2_ids, induced=True, batch_size=7) + assert cg.read_calls == int(np.ceil(len(lvl2_ids) / 7)) + ids, counts = np.unique(np.array(cg.ids_read), return_counts=True) + assert np.array_equal(ids, np.unique(lvl2_ids)), "every id must be read" + assert set(counts.tolist()) == {1}, "and read exactly once -- one pass, not two" + + +def test_partner_in_a_different_slice_is_still_resolved(): + """The mapping must be complete before anything is remapped; with batch_size=1 every + partner necessarily belongs to a different slice.""" + cce, parents, lvl2_ids = build_case(8) + want = reference(FakeCg(cce, parents), lvl2_ids, induced=True) + got = _get_edges_for_lvl2_ids(FakeCg(cce, parents), lvl2_ids, induced=True, batch_size=1) + assert np.array_equal(got, want) + assert len(got) > 0, "fixture should produce cross-slice edges" + + +def test_subset_query_still_looks_up_outside_parents(): + """A bounds-style query: ask for half the object, so partners outside it have no + mapping entry and have to come from get_parents.""" + cce, parents, lvl2_ids = build_case(9) + subset = lvl2_ids[: len(lvl2_ids) // 2] + cg = FakeCg(cce, parents) + got = _get_edges_for_lvl2_ids(cg, subset, induced=True, batch_size=5) + want = reference(FakeCg(cce, parents), subset, induced=True) + assert np.array_equal(got, want) + assert cg.get_parents_calls == 1, "partners outside the subset must be looked up" + assert np.isin(got, subset).all() + + +def _reload_pathing(): + """Re-import the module so its module-level env parsing runs again.""" + from pychunkedgraph.graph.analysis import pathing + + return importlib.reload(pathing) + + +@pytest.fixture +def pathing_module(): + """Reload around the test so an env override cannot leak into the rest of the suite.""" + yield _reload_pathing() + _reload_pathing() + + +@pytest.mark.parametrize( + "value,expected", + [ + ("2500", 2500), + ("1", 1), + (None, 10_000), # unset -> default + ("not-a-number", 10_000), # unparseable -> default, not a crash at import + ("0", 10_000), # would yield no slices at all -> refused + ("-5", 10_000), # same + ], +) +def test_batch_size_env_var(monkeypatch, pathing_module, value, expected): + if value is None: + monkeypatch.delenv("PCG_CROSS_EDGE_READ_BATCH", raising=False) + else: + monkeypatch.setenv("PCG_CROSS_EDGE_READ_BATCH", value) + assert _reload_pathing().CROSS_EDGE_READ_BATCH == expected + + +def test_env_batch_size_is_actually_used(monkeypatch, pathing_module): + """The constant has to reach the read loop, not just sit in the module.""" + monkeypatch.setenv("PCG_CROSS_EDGE_READ_BATCH", "9") + mod = _reload_pathing() + cce, parents, lvl2_ids = build_case(6) + cg = FakeCg(cce, parents) + got = mod._get_edges_for_lvl2_ids(cg, lvl2_ids, induced=True) + assert cg.read_calls == int(np.ceil(len(lvl2_ids) / 9)) + assert np.array_equal(got, reference(FakeCg(cce, parents), lvl2_ids, induced=True)) + + +def test_explicit_non_positive_batch_size_raises(): + """Silently returning zero edges would be far worse than failing.""" + cce, parents, lvl2_ids = build_case(0) + for bad in (0, -1): + with pytest.raises(ValueError, match="batch_size must be positive"): + _get_edges_for_lvl2_ids(FakeCg(cce, parents), lvl2_ids, batch_size=bad) 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