memory improvements to changelog and uwsgi reload-on-rss fix

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")

pychunkedgraph/app/config.py

@@ -14,6 +14,26 @@ class BaseConfig(object):
     LOGGING_DATEFORMAT = "%Y-%m-%dT%H:%M:%S.0Z"
     LOGGING_LEVEL = logging.DEBUG
 
+    # Opt-in start-of-request logging (see pychunkedgraph.app.common._log_request_start).
+    # When True, every non-probe request emits a REQUEST_START line to stdout before any work
+    # runs, so requests that OOM-kill their worker mid-flight (and thus never reach
+    # after_request) are still visible in Cloud Logging. Verbose; enable only for temporary
+    # diagnosis, e.g. by setting LOG_REQUEST_START = True in the instance config.cfg.
+    LOG_REQUEST_START = False
+
+    # Reject /lvl2_graph requests whose node resolves to more than this many level 2 nodes
+    # (see pychunkedgraph.graph.analysis.pathing.get_lvl2_edge_list). Such objects — typically
+    # erroneous mega-merges — produce a multi-GB induced edge list that can OOM the worker.
+    # None disables the guard; set a concrete integer in the instance config.cfg to enable.
+    LVL2_GRAPH_MAX_NODES = None
+
+    # Guard for /subgraph (see pychunkedgraph.graph.subgraph.get_subgraph_edges_and_leaves).
+    # Counts chunks rather than level 2 nodes: the endpoint reads every edge in every chunk the
+    # object touches (all objects in the chunk, not just the requested one), so cost tracks the
+    # volume queried, not the object. A large 'bounds' is expensive even for a small object.
+    # None disables the guard; set a concrete integer in the instance config.cfg to enable.
+    SUBGRAPH_MAX_CHUNKS = None
+
     CHUNKGRAPH_INSTANCE_ID = "pychunkedgraph"
     PROJECT_ID = os.environ.get("PROJECT_ID", None)
     CG_READ_ONLY = os.environ.get("CG_READ_ONLY", None) is not None

pychunkedgraph/app/segmentation/common.py

@@ -4,7 +4,6 @@
 import os
 import time
 from datetime import datetime
-from functools import reduce
 from collections import deque, defaultdict
 
 import numpy as np
@@ -785,8 +784,9 @@ def handle_subgraph(table_id, root_id, only_internal_edges=True):
         int(root_id),
         bbox=bounding_box,
         bbox_is_coordinate=True,
+        max_num_chunks=current_app.config.get("SUBGRAPH_MAX_CHUNKS"),
     )
-    edges = reduce(lambda x, y: x + y, edges, cg_edges.Edges([], []))
+    edges = cg_edges.Edges.concatenate(edges)
 
     if only_internal_edges:
         supervoxels = np.concatenate(
@@ -830,31 +830,43 @@ def tabular_change_log_recent(table_id):
     # Call ChunkedGraph
     cg = app_utils.get_cg(table_id)
 
-    log_rows = cg.client.read_log_entries(start_time=start_time, end_time=end_time)
-
+    # Stream the operation-log rows instead of materializing them all at once. Only the
+    # timestamp, user, and merge/split flag are needed, so the read is restricted to those
+    # columns (the default pulls the large variable-length added/removed-edge, coordinate,
+    # and affinity arrays for every operation, which dominate row size). AddedEdge is only
+    # existence-checked (merge vs split). RootID is not used directly but is kept so the
+    # streaming reader's timestamp fallback still works on older rows that predate the
+    # OperationTimeStamp column.
+    #
+    # read_log_entries_streaming yields one (operation_id, record) at a time and frees each
+    # decoded row before reading the next, so peak memory is bounded by the compact output
+    # columns below rather than by the full set of Bigtable cell objects for the window.
+    # Rows arrive already ordered by operation id (fixed-width keys), so no sort is needed.
+    operation_ids = []
     timestamp_list = []
     user_list = []
     is_merge_list = []
-
-    operation_ids = np.sort(list(log_rows.keys()))
-    for operation_id in operation_ids:
-        operation = log_rows[operation_id]
-
-        timestamp = operation["timestamp"]
-        timestamp_list.append(timestamp)
-
-        user_id = operation[attributes.OperationLogs.UserID]
-        user_list.append(user_id)
-
-        is_merge = attributes.OperationLogs.AddedEdge in operation
-        is_merge_list.append(is_merge)
+    for operation_id, operation in cg.client.read_log_entries_streaming(
+        start_time=start_time,
+        end_time=end_time,
+        properties=[
+            attributes.OperationLogs.OperationTimeStamp,
+            attributes.OperationLogs.UserID,
+            attributes.OperationLogs.AddedEdge,
+            attributes.OperationLogs.RootID,
+        ],
+    ):
+        operation_ids.append(operation_id)
+        timestamp_list.append(operation["timestamp"])
+        user_list.append(operation[attributes.OperationLogs.UserID])
+        is_merge_list.append(attributes.OperationLogs.AddedEdge in operation)
 
     return pd.DataFrame.from_dict(
         {
-            "operation_id": operation_ids,
+            "operation_id": np.array(operation_ids, dtype=np.uint64),
             "timestamp": timestamp_list,
             "user_id": user_list,
-            "is_merge": is_merge_list,
+            "is_merge": np.array(is_merge_list, dtype=bool),
         }
     )
 
@@ -1142,7 +1154,12 @@ def handle_get_layer2_graph(table_id, node_id):
 
     cg = app_utils.get_cg(table_id)
     print("Finding edge graph...")
-    edge_graph = pathing.get_lvl2_edge_list(cg, int(node_id), bbox=bounding_box)
+    edge_graph = pathing.get_lvl2_edge_list(
+        cg,
+        int(node_id),
+        bbox=bounding_box,
+        max_num_lvl2_ids=current_app.config.get("LVL2_GRAPH_MAX_NODES"),
+    )
     print("Edge graph found len: {}".format(len(edge_graph)))
     return {"edge_graph": edge_graph}
 

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
 

pychunkedgraph/graph/chunkedgraph.py

@@ -556,9 +556,14 @@ def get_subgraph(
         edges_only: bool = False,
         leaves_only: bool = False,
         return_flattened: bool = False,
+        max_num_chunks: typing.Optional[int] = None,
     ) -> typing.Tuple[typing.Dict, typing.Dict, Edges]:
         """
         Generic subgraph method.
+
+        :param max_num_chunks: Optional[int] reject the request (raising BadRequest) when the
+            node ids span more than this many chunks. ``None`` disables the guard. Only
+            applies to the edges/leaves path, which is the one that can OOM.
         """
         from .subgraph import get_subgraph_nodes
         from .subgraph import get_subgraph_edges_and_leaves
@@ -573,7 +578,13 @@ def get_subgraph(
                 return_flattened=return_flattened,
             )
         return get_subgraph_edges_and_leaves(
-            self, node_id_or_ids, bbox, bbox_is_coordinate, edges_only, leaves_only
+            self,
+            node_id_or_ids,
+            bbox,
+            bbox_is_coordinate,
+            edges_only,
+            leaves_only,
+            max_num_chunks=max_num_chunks,
         )
 
     def get_subgraph_nodes(
@@ -657,35 +668,43 @@ def get_l2_agglomerations(
         Edges are read from cloud storage.
         """
         from itertools import chain
-        from functools import reduce
         from .misc import get_agglomerations
 
         chunk_ids = np.unique(self.get_chunk_ids_from_node_ids(level2_ids))
+
+        # Resolve the object's supervoxels before reading any edges. Both consumers below
+        # discard every edge that does not touch one of them, so passing them down lets each
+        # chunk be filtered as it is parsed instead of after the whole set is materialized.
+        # A chunk's edge file holds every object in that chunk, so for a single neuron this
+        # is the difference between retaining the chunk content and retaining the object.
+        l2id_children_d = self.get_children(level2_ids)
+        supervoxels = (
+            np.concatenate(list(l2id_children_d.values()))
+            if l2id_children_d
+            else types.empty_1d.copy()
+        )
+
         # google does not provide a storage emulator at the moment
         # this is an ugly hack to avoid permission issues in tests
         # find a better way to test
         edges_d = {}
         if self.mock_edges is None:
-            edges_d = self.read_chunk_edges(chunk_ids)
+            edges_d = self.read_chunk_edges(chunk_ids, supervoxels=supervoxels)
 
         fake_edges = self.get_fake_edges(chunk_ids)
-        all_chunk_edges = reduce(
-            lambda x, y: x + y,
-            chain(edges_d.values(), fake_edges.values()),
-            Edges([], []),
+        all_chunk_edges = Edges.concatenate(
+            chain(edges_d.values(), fake_edges.values())
         )
 
         if edges_only:
             if self.mock_edges is not None:
                 all_chunk_edges = self.mock_edges.get_pairs()
             else:
                 all_chunk_edges = all_chunk_edges.get_pairs()
-            supervoxels = self.get_children(level2_ids, flatten=True)
             mask0 = np.in1d(all_chunk_edges[:, 0], supervoxels)
             mask1 = np.in1d(all_chunk_edges[:, 1], supervoxels)
             return all_chunk_edges[mask0 & mask1]
 
-        l2id_children_d = self.get_children(level2_ids)
         sv_parent_d = {}
         for l2id in l2id_children_d:
             svs = l2id_children_d[l2id]
@@ -996,12 +1015,15 @@ def get_parent_chunk_id_dict(self, node_or_chunk_id: basetypes.NODE_ID):
     def get_cross_chunk_edges_layer(self, cross_edges: typing.Iterable):
         return edge_utils.get_cross_chunk_edges_layer(self.meta, cross_edges)
 
-    def read_chunk_edges(self, chunk_ids: typing.Iterable) -> typing.Dict:
+    def read_chunk_edges(
+        self, chunk_ids: typing.Iterable, supervoxels: np.ndarray = None
+    ) -> typing.Dict:
         from ..io.edges import get_chunk_edges
 
         return get_chunk_edges(
             self.meta.data_source.EDGES,
             self.get_chunk_coordinates_multiple(chunk_ids),
+            supervoxels=supervoxels,
         )
 
     def get_proofread_root_ids(

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).
+        """