adding limits to get subgraph

pychunkedgraph/app/config.py

@@ -29,6 +29,21 @@ class BaseConfig(object):
             AUTH_TOKEN = json.load(f)["token"]
 
     AUTH_SERVICE_NAMESPACE = "pychunkedgraph"
+
+    # Guardrail for the subgraph endpoints, see pychunkedgraph/graph/limits.py.
+    # The "default" entry applies to every table; add an entry keyed by table id
+    # to override it, or set it to null to leave that table unrestricted.
+    # `MAX_BYTES` is how much memory one request may need. The cost per level 2
+    # chunk is derived from the chunk's physical volume, which adapts to any
+    # chunk size and resolution; a dataset that is denser or sparser than the
+    # default can pin its own `BYTES_PER_CUBIC_MICRON`, or skip the estimate
+    # entirely with a measured `BYTES_PER_L2_CHUNK`.
+    # Override or extend with the PCG_SUBGRAPH_LIMITS env var, e.g.
+    # '{"minnie3_v1": {"MAX_BYTES": 5368709120, "BYTES_PER_L2_CHUNK": 1048576}}'
+    SUBGRAPH_LIMITS = {
+        "default": {"MAX_BYTES": 5 * 1024**3},
+    }
+
     VIRTUAL_TABLES = {
         "minnie65_public_v117": {
             "table_id": "minnie3_v1",

pychunkedgraph/app/segmentation/common.py

@@ -30,6 +30,9 @@
 from pychunkedgraph.graph.misc import get_contact_sites
 from pychunkedgraph.graph.operation import GraphEditOperation
 from pychunkedgraph.graph.utils import basetypes
+from pychunkedgraph.graph import (
+    limits as cg_limits,
+)
 from pychunkedgraph.meshing import mesh_analysis
 
 __api_versions__ = [0, 1]
@@ -85,6 +88,12 @@ def _get_bounds_from_request(request):
     return bounding_box
 
 
+def _get_subgraph_limits(table_id):
+    """Memory guardrail configured for `table_id`, None if unrestricted."""
+    limits = current_app.config.get("SUBGRAPH_LIMITS") or {}
+    return limits.get(table_id, limits.get("default"))
+
+
 # -------------------
 # ------ Applications
 # -------------------
@@ -781,6 +790,10 @@ def handle_subgraph(table_id, root_id, only_internal_edges=True):
 
     # Call ChunkedGraph
     cg = app_utils.get_cg(table_id)
+    # fails with 413 before any chunk is read if the box is too expensive
+    cg_limits.check_subgraph_bounds(
+        cg, bounding_box, _get_subgraph_limits(table_id)
+    )
     l2id_agglomeration_d, edges = cg.get_subgraph(
         int(root_id),
         bbox=bounding_box,

pychunkedgraph/graph/exceptions.py

@@ -69,6 +69,15 @@ class Conflict(ClientError):
     status_code = http_client.CONFLICT
 
 
+class RequestTooLarge(ClientError):
+    """Exception mapping a ``413 Request Entity Too Large`` response.
+
+    Raised when a request is rejected up front because serving it would
+    require loading more data than the server is willing to hold in memory.
+    """
+    status_code = http_client.REQUEST_ENTITY_TOO_LARGE
+
+
 class ServerError(ChunkedGraphAPIError):
     """Base for 5xx responses."""
 

pychunkedgraph/graph/limits.py

@@ -0,0 +1,135 @@
+# pylint: disable=invalid-name, missing-docstring
+
+"""
+Guardrails for queries whose memory footprint scales with the volume of the
+requested bounding box rather than with the size of the segment.
+
+`cg.get_subgraph(...)` resolves the level 2 nodes inside the box and then calls
+`get_l2_agglomerations`, which reads the edges of *every* level 2 chunk those
+nodes live in -- `ChunkedGraph.read_chunk_edges` fetches whole chunk files, not
+just the edges of the requested supervoxels. A generous bounding box therefore
+pulls in every supervoxel and edge stored in that volume and can exhaust the
+pod's memory even when the requested segment barely touches it.
+
+Because the dominant term depends only on the box, it can be estimated before
+any data is read, which is what this module does. These limits are meant for
+user facing requests; edit operations and ingest deliberately do not go through
+them, so the check is applied by the app layer rather than by the graph itself.
+"""
+
+from typing import Dict
+from typing import Optional
+from typing import Sequence
+
+import numpy as np
+
+from . import exceptions as cg_exceptions
+from .chunks.utils import normalize_bounding_box
+
+
+DEFAULT_MAX_BYTES = 5 * 1024**3
+
+# Fallback calibration, used when a dataset has no measured
+# `BYTES_PER_L2_CHUNK`. Chunk sizes and resolutions differ between datasets, so
+# a per-chunk constant does not transfer; what does transfer reasonably well is
+# the cost per unit of *physical* volume, since supervoxel and edge density are
+# a property of the EM segmentation rather than of the chunking. The figure
+# below is derived for minnie65: ~150 edges per cubic micron at ~28 bytes per
+# edge (two node ids, an affinity and an area), times ~3 for the copies made
+# while concatenating chunk edges and splitting them into in/out/cross sets.
+# It is a rough upper-middle estimate; measure and pin per dataset when a
+# dataset turns out to be denser or sparser than that.
+DEFAULT_BYTES_PER_CUBIC_MICRON = 12_000
+
+
+def bytes_per_l2_chunk(meta, limits: Dict) -> float:
+    """Memory a single level 2 chunk of edges is expected to cost.
+
+    Uses the dataset's measured `BYTES_PER_L2_CHUNK` when configured, otherwise
+    scales `BYTES_PER_CUBIC_MICRON` by the physical volume of a chunk so that
+    one calibration applies to any chunk size and resolution.
+    """
+    measured = limits.get("BYTES_PER_L2_CHUNK")
+    if measured:
+        return float(measured)
+    chunk_nm = np.array(meta.graph_config.CHUNK_SIZE, dtype=float) * np.array(
+        meta.resolution, dtype=float
+    )
+    chunk_um3 = float(np.prod(chunk_nm)) / 1e9
+    per_um3 = limits.get("BYTES_PER_CUBIC_MICRON", DEFAULT_BYTES_PER_CUBIC_MICRON)
+    return chunk_um3 * float(per_um3)
+
+
+def level2_chunk_count(meta, bounding_box: Optional[Sequence[Sequence[int]]]) -> int:
+    """Number of level 2 chunks spanned by a bounding box in voxel coordinates.
+
+    `None` is treated as the whole dataset, which is what an omitted bounding
+    box actually asks for.
+    """
+    chunk_bounds = np.array(meta.layer_chunk_bounds[2], dtype=int)
+    if bounding_box is None:
+        bbox = np.array([[0, 0, 0], chunk_bounds], dtype=int)
+    else:
+        bbox = normalize_bounding_box(meta, np.array(bounding_box, dtype=int), True)
+    lower = np.clip(bbox[0], 0, chunk_bounds)
+    upper = np.clip(bbox[1], 0, chunk_bounds)
+    return int(np.prod(np.maximum(upper - lower, 1)))
+
+
+def _suggested_box(meta, max_chunks: int) -> Sequence[int]:
+    """A roughly cubic box, in voxels, that fits within `max_chunks`.
+
+    Chunks are usually anisotropic, so a cube of chunks would be a very
+    elongated region; this sizes each axis in physical space instead and then
+    rounds down to whole chunks, which is the granularity that actually counts.
+    """
+    chunk_size = np.array(meta.graph_config.CHUNK_SIZE, dtype=int)
+    chunk_nm = chunk_size * np.array(meta.resolution, dtype=float)
+    side_nm = (max_chunks * float(np.prod(chunk_nm))) ** (1 / 3)
+    chunks_per_axis = np.maximum(np.floor(side_nm / chunk_nm), 1).astype(int)
+    # clamping a thin axis up to one chunk can push the total over the budget
+    while np.prod(chunks_per_axis) > max_chunks and np.any(chunks_per_axis > 1):
+        chunks_per_axis[np.argmax(chunks_per_axis)] -= 1
+    return (chunks_per_axis * chunk_size).tolist()
+
+
+def check_subgraph_bounds(
+    cg,
+    bounding_box: Optional[Sequence[Sequence[int]]],
+    limits: Optional[Dict],
+) -> None:
+    """Reject a subgraph request whose bounding box is too expensive to serve.
+
+    `limits` may set `MAX_BYTES`, the memory a single request may need, and
+    either of the calibrations described on `bytes_per_l2_chunk`; both fall back
+    to the module defaults, so an empty dict still applies the default limit.
+    Pass `None` to leave the request unrestricted.
+
+    Raises `exceptions.RequestTooLarge` (HTTP 413) before any chunk is read.
+    """
+    if limits is None:
+        return
+
+    max_bytes = int(limits.get("MAX_BYTES", DEFAULT_MAX_BYTES))
+    bytes_per_chunk = bytes_per_l2_chunk(cg.meta, limits)
+    max_chunks = max(int(max_bytes // bytes_per_chunk), 1)
+
+    n_chunks = level2_chunk_count(cg.meta, bounding_box)
+    if n_chunks <= max_chunks:
+        return
+
+    x, y, z = _suggested_box(cg.meta, max_chunks)
+    rx, ry, rz = [int(r) for r in cg.meta.resolution]
+    gb = 1024.0**3
+
+    scope = "the whole dataset" if bounding_box is None else "the requested bounds"
+    raise cg_exceptions.RequestTooLarge(
+        f"Subgraph request too large: {scope} spans {n_chunks} level 2 chunks, "
+        f"which needs roughly {n_chunks * bytes_per_chunk / gb:.1f} GB of memory "
+        f"to load (limit {max_bytes / gb:.1f} GB, {max_chunks} chunks). This "
+        "query reads every edge stored in the chunks the box touches, so its "
+        "cost scales with the volume of the box rather than with the size of "
+        "the segment. Please split it into smaller boxes -- up to about "
+        f"{x}x{y}x{z} voxels at {rx}x{ry}x{rz} nm resolution each -- and "
+        "combine the results."
+    )

pychunkedgraph/tests/test_limits.py

@@ -0,0 +1,110 @@
+# pylint: disable=invalid-name, missing-docstring, redefined-outer-name
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from pychunkedgraph.graph import exceptions as cg_exceptions
+from pychunkedgraph.graph import limits
+
+
+CHUNK_SIZE = [256, 256, 512]
+# 1000 x 1000 x 1000 level 2 chunks
+CHUNK_BOUNDS = np.array([1000, 1000, 1000])
+# 8 level 2 chunks per request at most
+LIMITS = {"MAX_BYTES": 8 * 1024**2, "BYTES_PER_L2_CHUNK": 1024**2}
+
+
+@pytest.fixture
+def cg():
+    """Minimal stand-in exposing only what the bbox estimate reads."""
+    meta = SimpleNamespace(
+        graph_config=SimpleNamespace(CHUNK_SIZE=CHUNK_SIZE, FANOUT=2),
+        layer_chunk_bounds={2: CHUNK_BOUNDS},
+        resolution=np.array([8, 8, 40]),
+        voxel_bounds=np.array([[0, 0], [0, 0], [0, 0]]),
+    )
+    return SimpleNamespace(meta=meta)
+
+
+def _bbox(size):
+    """Bounding box of `size` voxels per side, anchored at the origin."""
+    return np.array([[0, 0, 0], list(size)])
+
+
+def test_small_bbox_allowed(cg):
+    # exactly 2 x 2 x 2 = 8 chunks, at the limit
+    limits.check_subgraph_bounds(cg, _bbox(np.array(CHUNK_SIZE) * 2), LIMITS)
+
+
+def test_large_bbox_rejected(cg):
+    with pytest.raises(cg_exceptions.RequestTooLarge) as exc:
+        limits.check_subgraph_bounds(cg, _bbox(np.array(CHUNK_SIZE) * 3), LIMITS)
+    assert exc.value.status_code.value == 413
+    assert "smaller boxes" in exc.value.message
+
+
+def test_missing_bbox_rejected(cg):
+    """No bounds is a request for the whole dataset."""
+    with pytest.raises(cg_exceptions.RequestTooLarge):
+        limits.check_subgraph_bounds(cg, None, LIMITS)
+
+
+def test_bbox_clipped_to_dataset(cg):
+    """A box reaching past the dataset only counts the chunks that exist."""
+    cg.meta.layer_chunk_bounds = {2: np.array([1, 1, 1])}
+    limits.check_subgraph_bounds(cg, _bbox([10**6, 10**6, 10**6]), LIMITS)
+
+
+def test_no_limits_configured(cg):
+    limits.check_subgraph_bounds(cg, None, None)
+
+
+def test_empty_limits_fall_back_to_defaults(cg):
+    """An empty entry is still guarded, with the module defaults."""
+    with pytest.raises(cg_exceptions.RequestTooLarge):
+        limits.check_subgraph_bounds(cg, None, {})
+    # ~86 um^3 per chunk at this chunk size and resolution
+    limits.check_subgraph_bounds(cg, _bbox(np.array(CHUNK_SIZE) * 2), {})
+
+
+def test_default_calibration_tracks_chunk_volume(cg):
+    """The default cost per chunk follows the chunk's physical volume."""
+    per_chunk = limits.bytes_per_l2_chunk(cg.meta, {})
+    chunk_um3 = np.prod(np.array(CHUNK_SIZE) * np.array([8, 8, 40])) / 1e9
+    assert per_chunk == pytest.approx(
+        chunk_um3 * limits.DEFAULT_BYTES_PER_CUBIC_MICRON
+    )
+
+    # a dataset chunked twice as coarsely in z costs twice as much per chunk
+    cg.meta.graph_config.CHUNK_SIZE = [256, 256, 1024]
+    assert limits.bytes_per_l2_chunk(cg.meta, {}) == pytest.approx(2 * per_chunk)
+
+
+def test_suggested_box_fits_the_budget(cg):
+    chunk_size = np.array(CHUNK_SIZE)
+    for max_chunks in [1, 8, 5208, 10**6]:
+        box = np.array(limits._suggested_box(cg.meta, max_chunks))
+        assert np.all(box % chunk_size == 0)
+        assert np.all(box >= chunk_size)
+        assert np.prod(box // chunk_size) <= max_chunks
+
+    # roughly cubic in nanometers rather than in chunks
+    box = np.array(limits._suggested_box(cg.meta, 5208)) * np.array([8, 8, 40])
+    assert box.max() / box.min() < 2
+
+
+def test_measured_calibration_wins(cg):
+    assert limits.bytes_per_l2_chunk(cg.meta, LIMITS) == 1024**2
+    assert limits.bytes_per_l2_chunk(cg.meta, {"BYTES_PER_CUBIC_MICRON": 1}) == (
+        pytest.approx(np.prod(np.array(CHUNK_SIZE) * np.array([8, 8, 40])) / 1e9)
+    )
+
+
+def test_level2_chunk_count(cg):
+    assert limits.level2_chunk_count(cg.meta, _bbox(CHUNK_SIZE)) == 1
+    assert limits.level2_chunk_count(cg.meta, _bbox(np.array(CHUNK_SIZE) * 4)) == 64
+    # sub-chunk boxes still touch one chunk
+    assert limits.level2_chunk_count(cg.meta, _bbox([1, 1, 1])) == 1
+    assert limits.level2_chunk_count(cg.meta, None) == int(np.prod(CHUNK_BOUNDS))