From 0941100b9dde6cea78298fd7b8702e8b310378d1 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 8 Aug 2026 01:53:48 -0700 Subject: [PATCH 01/10] memory improvements to changelog and uwsgi reload-on-rss fix --- pychunkedgraph/app/segmentation/common.py | 19 ++++++++++++++++++- uwsgi.ini | 8 ++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/pychunkedgraph/app/segmentation/common.py b/pychunkedgraph/app/segmentation/common.py index 3250248f2..914eecba5 100644 --- a/pychunkedgraph/app/segmentation/common.py +++ b/pychunkedgraph/app/segmentation/common.py @@ -830,7 +830,24 @@ 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) + # Only the timestamp, user, and merge/split flag are used below, so restrict the + # Bigtable read to those columns. The default (all columns) also pulls the large + # variable-length arrays (added/removed edges, coordinates, affinities) for every + # operation, which dominate row size and drive the memory footprint of this endpoint. + # AddedEdge is only existence-checked (merge vs split), so its presence is all we need. + log_rows = cg.client.read_log_entries( + start_time=start_time, + end_time=end_time, + properties=[ + attributes.OperationLogs.OperationTimeStamp, + attributes.OperationLogs.UserID, + attributes.OperationLogs.AddedEdge, + # RootID is not used directly, but read_log_entries falls back to its cell + # timestamp when OperationTimeStamp is absent on older rows; keep it so that + # fallback still works. It is a small array, unlike the edge/coord columns. + attributes.OperationLogs.RootID, + ], + ) timestamp_list = [] user_list = [] 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 From d310fff550c47bb9aeb1559703719f755952dd7d Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 8 Aug 2026 03:15:03 -0700 Subject: [PATCH 02/10] trying streaming implementation --- pychunkedgraph/app/segmentation/common.py | 53 +++++++--------- pychunkedgraph/graph/client/base.py | 19 ++++++ .../graph/client/bigtable/client.py | 62 +++++++++++++++++++ 3 files changed, 105 insertions(+), 29 deletions(-) diff --git a/pychunkedgraph/app/segmentation/common.py b/pychunkedgraph/app/segmentation/common.py index 914eecba5..5a7776f4b 100644 --- a/pychunkedgraph/app/segmentation/common.py +++ b/pychunkedgraph/app/segmentation/common.py @@ -830,48 +830,43 @@ def tabular_change_log_recent(table_id): # Call ChunkedGraph cg = app_utils.get_cg(table_id) - # Only the timestamp, user, and merge/split flag are used below, so restrict the - # Bigtable read to those columns. The default (all columns) also pulls the large - # variable-length arrays (added/removed edges, coordinates, affinities) for every - # operation, which dominate row size and drive the memory footprint of this endpoint. - # AddedEdge is only existence-checked (merge vs split), so its presence is all we need. - log_rows = cg.client.read_log_entries( + # 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 = [] + 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, - # RootID is not used directly, but read_log_entries falls back to its cell - # timestamp when OperationTimeStamp is absent on older rows; keep it so that - # fallback still works. It is a small array, unlike the edge/coord columns. attributes.OperationLogs.RootID, ], - ) - - 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) + ): + 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), } ) 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, From 96c462282696af66ea9454d0982a0300e134a2e9 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 8 Aug 2026 05:04:22 -0700 Subject: [PATCH 03/10] add pre-request log option --- pychunkedgraph/app/common.py | 35 +++++++++++++++++++++++++++++++++++ pychunkedgraph/app/config.py | 7 +++++++ 2 files changed, 42 insertions(+) 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..94c89d007 100644 --- a/pychunkedgraph/app/config.py +++ b/pychunkedgraph/app/config.py @@ -14,6 +14,13 @@ 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 + 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 From 5204518519057f164cb3b9af11fe402053d8f7bd Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 8 Aug 2026 06:05:46 -0700 Subject: [PATCH 04/10] add level 2 graph memory gaurd --- pychunkedgraph/app/config.py | 6 ++++++ pychunkedgraph/app/segmentation/common.py | 7 ++++++- pychunkedgraph/graph/analysis/pathing.py | 20 ++++++++++++++++++++ 3 files changed, 32 insertions(+), 1 deletion(-) diff --git a/pychunkedgraph/app/config.py b/pychunkedgraph/app/config.py index 94c89d007..14179634a 100644 --- a/pychunkedgraph/app/config.py +++ b/pychunkedgraph/app/config.py @@ -21,6 +21,12 @@ class BaseConfig(object): # 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 + 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 5a7776f4b..8c37a26b7 100644 --- a/pychunkedgraph/app/segmentation/common.py +++ b/pychunkedgraph/app/segmentation/common.py @@ -1154,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 From e439bf95d5a5f70550688f82b61ae30e248adef2 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 15 Aug 2026 14:42:15 -0700 Subject: [PATCH 05/10] adding subgraph gaurd --- pychunkedgraph/app/config.py | 7 +++++ pychunkedgraph/app/segmentation/common.py | 1 + pychunkedgraph/graph/chunkedgraph.py | 13 ++++++++- pychunkedgraph/graph/subgraph.py | 35 ++++++++++++++++++++++- 4 files changed, 54 insertions(+), 2 deletions(-) diff --git a/pychunkedgraph/app/config.py b/pychunkedgraph/app/config.py index 14179634a..4e8d5bd00 100644 --- a/pychunkedgraph/app/config.py +++ b/pychunkedgraph/app/config.py @@ -27,6 +27,13 @@ class BaseConfig(object): # 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 8c37a26b7..a38609c71 100644 --- a/pychunkedgraph/app/segmentation/common.py +++ b/pychunkedgraph/app/segmentation/common.py @@ -785,6 +785,7 @@ 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([], [])) diff --git a/pychunkedgraph/graph/chunkedgraph.py b/pychunkedgraph/graph/chunkedgraph.py index 210bff50b..2e06dad9a 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( 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: From 01750ee5d903a82c24119e4d51afaa70f808f81e Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 15 Aug 2026 15:25:50 -0700 Subject: [PATCH 06/10] add edges concatenate to avoid memory explosion and quadratic explosion --- pychunkedgraph/app/segmentation/common.py | 3 +- pychunkedgraph/graph/chunkedgraph.py | 7 +- pychunkedgraph/graph/edges/__init__.py | 24 +++ pychunkedgraph/tests/test_edges.py | 245 ++++++++++++++++++++++ 4 files changed, 272 insertions(+), 7 deletions(-) create mode 100644 pychunkedgraph/tests/test_edges.py diff --git a/pychunkedgraph/app/segmentation/common.py b/pychunkedgraph/app/segmentation/common.py index a38609c71..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 @@ -787,7 +786,7 @@ def handle_subgraph(table_id, root_id, only_internal_edges=True): 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( diff --git a/pychunkedgraph/graph/chunkedgraph.py b/pychunkedgraph/graph/chunkedgraph.py index 2e06dad9a..f80ef30ea 100644 --- a/pychunkedgraph/graph/chunkedgraph.py +++ b/pychunkedgraph/graph/chunkedgraph.py @@ -668,7 +668,6 @@ 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)) @@ -680,10 +679,8 @@ def get_l2_agglomerations( edges_d = self.read_chunk_edges(chunk_ids) 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: diff --git a/pychunkedgraph/graph/edges/__init__.py b/pychunkedgraph/graph/edges/__init__.py index b0e488d05..279eaadd3 100644 --- a/pychunkedgraph/graph/edges/__init__.py +++ b/pychunkedgraph/graph/edges/__init__.py @@ -65,6 +65,30 @@ def areas(self) -> np.ndarray: def areas(self, areas): self._areas = areas + @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/tests/test_edges.py b/pychunkedgraph/tests/test_edges.py new file mode 100644 index 000000000..89ecc320b --- /dev/null +++ b/pychunkedgraph/tests/test_edges.py @@ -0,0 +1,245 @@ +"""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 +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) From 204bb17492b28457506b8ef88bfaafdd5407b160 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 15 Aug 2026 17:38:38 -0700 Subject: [PATCH 07/10] filter chunk edges to the queried object as they are parsed /subgraph memory scaled with the total edge content of the chunks an object spans, not with the object. A chunk's edge file holds every object in that chunk, and get_chunk_edges retained all of it: the compressed blobs, the fully decompressed buffers, and a concatenated copy all coexisted before anything was filtered. Measured on minniev7, a ~260-chunk request peaked at 6.28 GiB of a 12 GiB pod limit to return 8.3 MB of edges. Both consumers discard edges that do not touch the queried object -- the edges_only path keeps edges with both endpoints in the supervoxel set, and categorize_edges_v2 drops any edge whose node_ids1 does not remap through sv_parent_d. So resolve the supervoxels first and pass them down, letting each chunk be filtered as it is parsed. The arrays from deserialize are np.frombuffer views into the decompressed chunk, so masking copies out the few edges that matter and lets the buffer be released instead of pinned until the end. Edges.filter_touching keeps edges with either endpoint in the set, which is a superset of both predicates, so filtering early cannot change either result. in_sorted avoids re-sorting the supervoxel set once per chunk. get_children moves above the read to supply the set; the edges_only path now reuses it instead of issuing a second identical read. The ingest caller passes no supervoxels and is unaffected. Co-Authored-By: Claude Opus 5 (1M context) --- pychunkedgraph/graph/chunkedgraph.py | 22 ++++- pychunkedgraph/graph/edges/__init__.py | 32 +++++++ pychunkedgraph/io/edges.py | 26 +++++- pychunkedgraph/tests/test_edges.py | 112 ++++++++++++++++++++++++- 4 files changed, 183 insertions(+), 9 deletions(-) diff --git a/pychunkedgraph/graph/chunkedgraph.py b/pychunkedgraph/graph/chunkedgraph.py index f80ef30ea..d9dca4be6 100644 --- a/pychunkedgraph/graph/chunkedgraph.py +++ b/pychunkedgraph/graph/chunkedgraph.py @@ -671,12 +671,25 @@ def get_l2_agglomerations( 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 = Edges.concatenate( @@ -688,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] @@ -1004,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/edges/__init__.py b/pychunkedgraph/graph/edges/__init__.py index 279eaadd3..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,26 @@ 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. diff --git a/pychunkedgraph/io/edges.py b/pychunkedgraph/io/edges.py index 82595e139..2ef1bc538 100644 --- a/pychunkedgraph/io/edges.py +++ b/pychunkedgraph/io/edges.py @@ -36,7 +36,7 @@ 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 _parse_edges(compressed: List[bytes], sorted_svs: np.ndarray = None) -> List[Dict]: result = [] if(len(compressed) == 0): return result @@ -60,18 +60,36 @@ 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.""" +def get_chunk_edges( + edges_dir: str, + chunks_coordinates: List[np.ndarray], + supervoxels: np.ndarray = 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). + """ 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)) + cf = CloudFiles(edges_dir, num_threads=4) files = cf.get(fnames, raw=True) compressed = [] @@ -79,7 +97,7 @@ def get_chunk_edges(edges_dir: str, chunks_coordinates: List[np.ndarray]) -> Dic if not f["content"]: continue compressed.append(f["content"]) - return concatenate_chunk_edges(_parse_edges(compressed)) + return concatenate_chunk_edges(_parse_edges(compressed, sorted_svs)) def put_chunk_edges( diff --git a/pychunkedgraph/tests/test_edges.py b/pychunkedgraph/tests/test_edges.py index 89ecc320b..89d09d94b 100644 --- a/pychunkedgraph/tests/test_edges.py +++ b/pychunkedgraph/tests/test_edges.py @@ -21,7 +21,7 @@ import numpy as np import pytest -from ..graph.edges import Edges +from ..graph.edges import Edges, in_sorted from ..graph.utils import basetypes @@ -243,3 +243,113 @@ def counting_concatenate(*args, **kwargs): 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) From 9d44d66391c510c298bcdaffdc1f2fd80121e42d Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 15 Aug 2026 17:42:49 -0700 Subject: [PATCH 08/10] fall back to serial decompression when multi_decompress_to_buffer is absent The fallback path already existed for builds without multi-threading support, but only caught ValueError. zstandard >= 0.23 removed multi_decompress_to_buffer altogether, so on any newer version _parse_edges raised AttributeError instead of taking the fallback. The image pins zstandard==0.21.0, which is why this has not bitten in production, but the fallback should not depend on that pin -- and it currently blocks running the edge IO tests outside the image. Co-Authored-By: Claude Opus 5 (1M context) --- pychunkedgraph/io/edges.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pychunkedgraph/io/edges.py b/pychunkedgraph/io/edges.py index 2ef1bc538..4cc88388e 100644 --- a/pychunkedgraph/io/edges.py +++ b/pychunkedgraph/io/edges.py @@ -49,7 +49,11 @@ def _parse_edges(compressed: List[bytes], sorted_svs: np.ndarray = None) -> List decompressed = [] try: decompressed = zdc.multi_decompress_to_buffer(compressed, threads=n_threads) - except ValueError: + except (ValueError, AttributeError): + # ValueError: build lacks multi-threading support. + # AttributeError: zstandard >= 0.23 removed multi_decompress_to_buffer entirely + # (the image pins 0.21.0, but the fallback should not depend on that pin). + decompressed = [] for content in compressed: decompressed.append(zdc.decompressobj().decompress(content)) From 41a3e4a7923e8d06458e9b756a596d60f4fab462 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 15 Aug 2026 17:43:11 -0700 Subject: [PATCH 09/10] fetch and decompress chunk edges in batches get_chunk_edges issued a single cf.get for every chunk in the request, so the compressed blobs for the whole query were held at once, then handed to multi_decompress_to_buffer which materialized all of the decompressed buffers at once as well. Peak scaled with the number of chunks, which is what made a large bounding box expensive regardless of the size of the object being queried. Process the files in batches instead, accumulating only the parsed (and, when supervoxels are given, already filtered) per-chunk results. Each batch's buffers are released before the next is fetched, so peak tracks batch_size rather than len(fnames). Batch size is configurable via PCG_EDGES_BATCH_SIZE, default 64. This composes with the filter commit and depends on it for most of the benefit: with no supervoxels the parsed arrays are np.frombuffer views that pin their decompressed buffers, so those cannot be released between batches. Filtering copies out the survivors, which is what lets each batch be freed. Tests cover the invariant that neither batching nor filtering changes the result: output is identical across batch sizes 1..1000, filtering matches filtering the fully accumulated set, every file is requested exactly once, and missing chunk files are still skipped. Co-Authored-By: Claude Opus 5 (1M context) --- pychunkedgraph/io/edges.py | 37 ++++-- pychunkedgraph/tests/test_io_edges.py | 173 ++++++++++++++++++++++++++ 2 files changed, 203 insertions(+), 7 deletions(-) create mode 100644 pychunkedgraph/tests/test_io_edges.py diff --git a/pychunkedgraph/io/edges.py b/pychunkedgraph/io/edges.py index 4cc88388e..3c7805da0 100644 --- a/pychunkedgraph/io/edges.py +++ b/pychunkedgraph/io/edges.py @@ -71,10 +71,17 @@ def _parse_edges(compressed: List[bytes], sorted_svs: np.ndarray = None) -> List return result +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. @@ -82,6 +89,10 @@ def get_chunk_edges( 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: @@ -94,14 +105,26 @@ def get_chunk_edges( 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, sorted_svs)) + # 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_io_edges.py b/pychunkedgraph/tests/test_io_edges.py new file mode 100644 index 000000000..1b43e1961 --- /dev/null +++ b/pychunkedgraph/tests/test_io_edges.py @@ -0,0 +1,173 @@ +"""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) + ) + ) From d9427fcef4918fa333f233bdf8fbbd127bf2cbd0 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 15 Aug 2026 18:08:42 -0700 Subject: [PATCH 10/10] decompress chunk edges with a thread pool instead of a version-specific API multi_decompress_to_buffer was removed in zstandard 0.23, and the serial fallback added for that case is roughly 4x slower. That matters: the deployment sets ZSTD_THREADS=4, so the fast path really is parallel today and falling back to serial would be a real regression on any zstandard upgrade. Decompression releases the GIL, so a thread pool recovers all of it. Measured on zstandard 0.21.0, 192 MB across 64 blobs, 4 threads: serial decompressobj() [previous fallback] 164.8 ms 1164 MB/s serial dctx.decompress() [reuse] 150.1 ms 1278 MB/s ThreadPoolExecutor(4) 41.7 ms 4593 MB/s multi_decompress_to_buffer 42.5 ms 4507 MB/s The pool matches the removed API, so this drops the version branch entirely rather than choosing between a fast path and a slow one. ZstdDecompressor is not thread-safe -- sharing one across workers silently produces corrupt output rather than raising, which cost a debugging round here -- so each worker keeps its own via threading.local. dctx.decompress needs the content size in the frame header, which put_chunk_edges writes and which multi_decompress_to_buffer also required; decompressobj covers any frame lacking it. Verified against both zstandard 0.21.0 (the image pin) and 0.25.0: identical output for 1/2/4/8 threads, threaded matches serial, order preserved, empty input, and frames written with write_content_size=False. Co-Authored-By: Claude Opus 5 (1M context) --- pychunkedgraph/io/edges.py | 52 ++++++++++++++++++++------ pychunkedgraph/tests/test_io_edges.py | 54 +++++++++++++++++++++++++++ 2 files changed, 95 insertions(+), 11 deletions(-) diff --git a/pychunkedgraph/io/edges.py b/pychunkedgraph/io/edges.py index 3c7805da0..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,26 +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 _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, AttributeError): - # ValueError: build lacks multi-threading support. - # AttributeError: zstandard >= 0.23 removed multi_decompress_to_buffer entirely - # (the image pins 0.21.0, but the fallback should not depend on that pin). - decompressed = [] - for content in compressed: - decompressed.append(zdc.decompressobj().decompress(content)) + decompressed = _decompress(compressed, n_threads) for content in decompressed: chunk_edges = ChunkEdgesMsg() diff --git a/pychunkedgraph/tests/test_io_edges.py b/pychunkedgraph/tests/test_io_edges.py index 1b43e1961..de1480c8d 100644 --- a/pychunkedgraph/tests/test_io_edges.py +++ b/pychunkedgraph/tests/test_io_edges.py @@ -171,3 +171,57 @@ def test_unsorted_and_duplicated_supervoxels_are_normalized(self, chunks): "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