diff --git a/pychunkedgraph/app/config.py b/pychunkedgraph/app/config.py index 2f2a92e47..83b52b4be 100644 --- a/pychunkedgraph/app/config.py +++ b/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", diff --git a/pychunkedgraph/app/segmentation/common.py b/pychunkedgraph/app/segmentation/common.py index 3250248f2..b3cd5775d 100644 --- a/pychunkedgraph/app/segmentation/common.py +++ b/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, diff --git a/pychunkedgraph/graph/exceptions.py b/pychunkedgraph/graph/exceptions.py index 45aa57fc7..df21aaeae 100644 --- a/pychunkedgraph/graph/exceptions.py +++ b/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.""" diff --git a/pychunkedgraph/graph/limits.py b/pychunkedgraph/graph/limits.py new file mode 100644 index 000000000..a7b325759 --- /dev/null +++ b/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." + ) diff --git a/pychunkedgraph/tests/test_limits.py b/pychunkedgraph/tests/test_limits.py new file mode 100644 index 000000000..25658f9f2 --- /dev/null +++ b/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))