From 0941100b9dde6cea78298fd7b8702e8b310378d1 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sat, 8 Aug 2026 01:53:48 -0700 Subject: [PATCH 01/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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 From 9fe7faec1d595234b3e9bb751c035449a481b9e7 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sun, 23 Aug 2026 11:12:39 -0700 Subject: [PATCH 11/14] move remeshing to queue as an edit with no operation id --- pychunkedgraph/app/meshing/common.py | 97 +++++++++++++++++++--------- 1 file changed, 67 insertions(+), 30 deletions(-) diff --git a/pychunkedgraph/app/meshing/common.py b/pychunkedgraph/app/meshing/common.py index 8f1a0c20a..7b4657316 100644 --- a/pychunkedgraph/app/meshing/common.py +++ b/pychunkedgraph/app/meshing/common.py @@ -1,18 +1,15 @@ # 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.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 @@ -142,9 +139,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 +218,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) From 4277b8eaa7484fb3b23cb530942e949a753bc40e Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Sun, 23 Aug 2026 21:06:09 -0700 Subject: [PATCH 12/14] add catch for start layer parsing --- pychunkedgraph/app/meshing/common.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/pychunkedgraph/app/meshing/common.py b/pychunkedgraph/app/meshing/common.py index 7b4657316..43972350d 100644 --- a/pychunkedgraph/app/meshing/common.py +++ b/pychunkedgraph/app/meshing/common.py @@ -9,6 +9,7 @@ from pychunkedgraph import __version__ from pychunkedgraph.app import app_utils +from pychunkedgraph.graph import exceptions as cg_exceptions from pychunkedgraph.app.meshing import tasks as meshing_tasks from pychunkedgraph.meshing.manifest import get_highest_child_nodes_with_meshes from pychunkedgraph.meshing.manifest import get_children_before_start_layer @@ -72,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: From d31d96fe2f56e6dfbe2ae838dc330a0247caeb89 Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Mon, 24 Aug 2026 06:40:27 -0700 Subject: [PATCH 13/14] improve memry and speed of level2graph finding --- pychunkedgraph/graph/analysis/pathing.py | 92 ++++++++---- pychunkedgraph/tests/test_lvl2_edges.py | 179 +++++++++++++++++++++++ 2 files changed, 245 insertions(+), 26 deletions(-) create mode 100644 pychunkedgraph/tests/test_lvl2_edges.py diff --git a/pychunkedgraph/graph/analysis/pathing.py b/pychunkedgraph/graph/analysis/pathing.py index ca22683c4..0f073f6a9 100644 --- a/pychunkedgraph/graph/analysis/pathing.py +++ b/pychunkedgraph/graph/analysis/pathing.py @@ -7,6 +7,7 @@ from pychunkedgraph.graph.utils import flatgraph from .. import exceptions as cg_exceptions +from ..edges import in_sorted from ..subgraph import get_subgraph_nodes @@ -123,46 +124,74 @@ def get_lvl2_edge_list( def _get_edges_for_lvl2_ids(cg, lvl2_ids, induced=False): + """Cross-chunk edges between `lvl2_ids`, remapped from supervoxels to level 2 ids. + + Memory, not speed, is the binding constraint here: on a large object the intermediates + dwarf the answer. On a synthetic object the size of the ones that trip uwsgi's + reload-on-rss in production (250k level 2 ids, 2M cross edges, 300k induced edges), + this function peaked at 888MB to return 4.6MB of edges. The four numbered comments + below mark what accounted for that; the same case now peaks at 512MB and runs in 0.8s + instead of 3.1s, returning byte-identical output. test_lvl2_edges.py pins that + equivalence against a reference copy of the previous implementation. + """ # 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) - # 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]) + # Flatten the per-(lvl2 id, layer) arrays into one edge array, recording which lvl2 id + # owns each block. Keeping the block lengths lets the supervoxel -> lvl2 mapping below + # be rebuilt from column views of the concatenated array, so the dict only has to be + # walked once. + 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)) # protect in case there are no edges - if len(edge_array) == 0: + if len(blocks) == 0: 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) + + edge_array = np.concatenate(blocks) + # (1) np.concatenate copied every block, so drop the originals before the peak rather + # than holding them to the end of the function. This dict is one small array per + # (lvl2 id, layer) -- 500k of them in the measured case, 213MB for 31MB of edges, + # nearly all of it numpy object overhead. + del blocks + cce_dict.clear() + + # Column 0 of each block is a supervoxel whose lvl2 parent is known (the block's + # owner); column 1 is its partner across the chunk boundary, whose parent may not be. + # (2) Both are views into edge_array. The previous version rebuilt them as two more + # lists of per-block arrays plus an np.full() per block for the owner, then + # concatenated all three -- 390MB of the old peak, for data already sitting in + # edge_array. One np.repeat replaces the per-block np.full calls. + owner_column = np.repeat( + np.array(owners, dtype=edge_array.dtype), np.array(lengths, dtype=np.int64) + ) + del owners, lengths # 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 + edge_array[:, 0], return_index=True ) - known_l2_array = (np.concatenate(known_l2_list))[unique_indices] - unknown_supervoxel_array = np.unique(np.concatenate(unknown_supervoxel_list)) + known_l2_array = owner_column[unique_indices] + del owner_column, unique_indices + unknown_supervoxel_array = np.unique(edge_array[:, 1]) # 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 ) + del unknown_supervoxel_array if len(supervoxels_to_query_parent) > 0: missing_l2_ids = cg.get_parents(supervoxels_to_query_parent) known_supervoxel_array = np.concatenate( @@ -174,17 +203,28 @@ def _get_edges_for_lvl2_ids(cg, lvl2_ids, induced=False): edge_view = edge_array.view() edge_view.shape = -1 fastremap.remap_from_array_kv(edge_view, known_supervoxel_array, known_l2_array) + del known_supervoxel_array, known_l2_array - edge_array = np.unique(np.sort(edge_array, axis=1), axis=0) + # (3) In place: edge_array is ours (np.concatenate always copies) and was just + # rewritten in place by the remap above, so sorting it in place avoids duplicating the + # largest array in the function. + edge_array.sort(axis=1) if induced: # make this an induced subgraph - # keep only the edges that are between the lvl2 ids asked for + # keep only the edges that are between the lvl2 ids asked for. + # (4) Filter BEFORE the dedup below, not after. Most of these edges leave the + # object -- 2M in, 300k out in the measured case -- and np.unique(axis=0) is the + # single most expensive step here, so it should see the small array. Deduplication + # and this filter commute, so the result is unchanged. in_sorted rather than + # np.isin because np.isin re-sorts lvl2_ids on each of the two calls. + sorted_lvl2_ids = np.unique(lvl2_ids) edge_array = edge_array[ - np.isin(edge_array[:, 0], lvl2_ids) & np.isin(edge_array[:, 1], lvl2_ids) + in_sorted(edge_array[:, 0], sorted_lvl2_ids) + & in_sorted(edge_array[:, 1], sorted_lvl2_ids) ] - return edge_array + return np.unique(edge_array, axis=0) def find_l2_shortest_path( diff --git a/pychunkedgraph/tests/test_lvl2_edges.py b/pychunkedgraph/tests/test_lvl2_edges.py new file mode 100644 index 000000000..9b6ef7ad8 --- /dev/null +++ b/pychunkedgraph/tests/test_lvl2_edges.py @@ -0,0 +1,179 @@ +"""Equivalence tests for ``_get_edges_for_lvl2_ids``. + +The function was reworked to cut its peak memory (725MB -> 304MB on a synthetic object of +250k level 2 ids / 2M cross edges). Every change was supposed to be behaviour-preserving, +so what these tests pin is exactly that: 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. +""" + +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 + + def get_atomic_cross_edges(self, l2_ids): + # a fresh dict per call, as both real implementations return + return {k: dict(v) for k, v in self._cce.items()} + + 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) From a9097a847988011da5221fe43ed540357716267c Mon Sep 17 00:00:00 2001 From: Forrest Collman Date: Mon, 24 Aug 2026 08:46:55 -0700 Subject: [PATCH 14/14] improving memory profile of l2graph --- pychunkedgraph/graph/analysis/pathing.py | 245 +++++++++++++++-------- pychunkedgraph/tests/test_lvl2_edges.py | 133 +++++++++++- 2 files changed, 287 insertions(+), 91 deletions(-) diff --git a/pychunkedgraph/graph/analysis/pathing.py b/pychunkedgraph/graph/analysis/pathing.py index 0f073f6a9..e2ccffddc 100644 --- a/pychunkedgraph/graph/analysis/pathing.py +++ b/pychunkedgraph/graph/analysis/pathing.py @@ -1,6 +1,6 @@ +import os import typing -import fastremap import graph_tool import numpy as np @@ -123,108 +123,187 @@ def get_lvl2_edge_list( 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 here: on a large object the intermediates - dwarf the answer. On a synthetic object the size of the ones that trip uwsgi's - reload-on-rss in production (250k level 2 ids, 2M cross edges, 300k induced edges), - this function peaked at 888MB to return 4.6MB of edges. The four numbered comments - below mark what accounted for that; the same case now peaks at 512MB and runs in 0.8s - instead of 3.1s, returning byte-identical output. test_lvl2_edges.py pins that - equivalence against a reference copy of the previous implementation. + 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) - - # Flatten the per-(lvl2 id, layer) arrays into one edge array, recording which lvl2 id - # owns each block. Keeping the block lengths lets the supervoxel -> lvl2 mapping below - # be rebuilt from column views of the concatenated array, so the dict only has to be - # walked once. - 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)) + # 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 + + # 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(blocks) == 0: + if not batches: return np.empty((0, 2), dtype=np.uint64) - edge_array = np.concatenate(blocks) - # (1) np.concatenate copied every block, so drop the originals before the peak rather - # than holding them to the end of the function. This dict is one small array per - # (lvl2 id, layer) -- 500k of them in the measured case, 213MB for 31MB of edges, - # nearly all of it numpy object overhead. - del blocks - cce_dict.clear() - - # Column 0 of each block is a supervoxel whose lvl2 parent is known (the block's - # owner); column 1 is its partner across the chunk boundary, whose parent may not be. - # (2) Both are views into edge_array. The previous version rebuilt them as two more - # lists of per-block arrays plus an np.full() per block for the owner, then - # concatenated all three -- 390MB of the old peak, for data already sitting in - # edge_array. One np.repeat replaces the per-block np.full calls. - owner_column = np.repeat( - np.array(owners, dtype=edge_array.dtype), np.array(lengths, dtype=np.int64) - ) - del owners, lengths - - # Create two arrays to map supervoxels for which we know their parents - known_supervoxel_array, unique_indices = np.unique( - edge_array[:, 0], return_index=True + # 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 ) - known_l2_array = owner_column[unique_indices] - del owner_column, unique_indices - unknown_supervoxel_array = np.unique(edge_array[:, 1]) - # 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 - ) - del unknown_supervoxel_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) - del known_supervoxel_array, known_l2_array + # 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 - # (3) In place: edge_array is ours (np.concatenate always copies) and was just - # rewritten in place by the remap above, so sorting it in place avoids duplicating the - # largest array in the function. - edge_array.sort(axis=1) - - if induced: - # make this an induced subgraph - # keep only the edges that are between the lvl2 ids asked for. - # (4) Filter BEFORE the dedup below, not after. Most of these edges leave the - # object -- 2M in, 300k out in the measured case -- and np.unique(axis=0) is the - # single most expensive step here, so it should see the small array. Deduplication - # and this filter commute, so the result is unchanged. in_sorted rather than - # np.isin because np.isin re-sorts lvl2_ids on each of the two calls. - sorted_lvl2_ids = np.unique(lvl2_ids) - edge_array = edge_array[ - in_sorted(edge_array[:, 0], sorted_lvl2_ids) - & in_sorted(edge_array[:, 1], sorted_lvl2_ids) - ] - - return np.unique(edge_array, axis=0) + 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/tests/test_lvl2_edges.py b/pychunkedgraph/tests/test_lvl2_edges.py index 9b6ef7ad8..b12e0c720 100644 --- a/pychunkedgraph/tests/test_lvl2_edges.py +++ b/pychunkedgraph/tests/test_lvl2_edges.py @@ -1,13 +1,20 @@ """Equivalence tests for ``_get_edges_for_lvl2_ids``. -The function was reworked to cut its peak memory (725MB -> 304MB on a synthetic object of -250k level 2 ids / 2M cross edges). Every change was supposed to be behaviour-preserving, -so what these tests pin is exactly that: 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 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 @@ -21,10 +28,18 @@ 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, as both real implementations return - return {k: dict(v) for k, v in self._cce.items()} + # 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 @@ -177,3 +192,105 @@ def test_duplicate_lvl2_ids_in_input(): 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)