diff --git a/.bumpversion.cfg b/.bumpversion.cfg deleted file mode 100644 index b6b4de269..000000000 --- a/.bumpversion.cfg +++ /dev/null @@ -1,6 +0,0 @@ -[bumpversion] -current_version = 2.21.1 -commit = True -tag = True - -[bumpversion:file:pychunkedgraph/__init__.py] diff --git a/.coveragerc b/.coveragerc index a38e1c392..d351f3e7e 100644 --- a/.coveragerc +++ b/.coveragerc @@ -5,6 +5,13 @@ source = pychunkedgraph omit = *test* *benchmarking/* + pychunkedgraph/debug/* + pychunkedgraph/export/* + pychunkedgraph/jobs/* + pychunkedgraph/logging/* + pychunkedgraph/repair/* + pychunkedgraph/meshing/* + pychunkedgraph/app/* [report] # Regexes for lines to exclude from consideration diff --git a/.dockerignore b/.dockerignore index 66349c8a4..09e319073 100644 --- a/.dockerignore +++ b/.dockerignore @@ -106,6 +106,7 @@ venv.bak/ # Visual Code .vscode/ +*.code-workspace # terraform .terraform/ diff --git a/.github/workflows/README.md b/.github/workflows/README.md new file mode 100644 index 000000000..8c56c2a21 --- /dev/null +++ b/.github/workflows/README.md @@ -0,0 +1,51 @@ +# GitHub Actions workflows + +Two workflows drive CI and releases for PyChunkedGraph. + +| File | Workflow name | Trigger | Purpose | +|---|---|---|---| +| `main.yml` | PyChunkedGraph | push / PR to `main` or `pcgv3` | Build the image and run the test suite with coverage | +| `release.yml` | publish release | manual (`workflow_dispatch`) | Bump the version, tag, and create the GitHub Release; optionally bump the Helm chart | + +## `main.yml` — CI + +Runs on every push and pull request targeting `main` or `pcgv3`. One `unit-tests` job: + +1. Builds the Docker image locally with Buildx (`load: true`, not pushed), tagged with the commit SHA, using the GitHub Actions layer cache. +2. Runs `pytest` with coverage inside the container against `pychunkedgraph/tests`. +3. Copies `coverage.xml` out and uploads it to Codecov — runs even when tests fail, and a Codecov upload error does not fail the build. +4. Removes the test container. + +Secrets: `CODECOV_TOKEN`. + +## `release.yml` — release + +Manual only — dispatch it from the branch you want to release. **The branch is the major line**: `main` carries 2.x, `pcgv3` carries 3.x, because the version is a committed literal in `pychunkedgraph/_version.py` (the repo-root README has the full versioning flow). The tag and the in-code literal are written by the same job, so they never drift. + +### Inputs + +| Input | Default | Effect | +|---|---|---| +| `part` | `patch` | which semver component to bump — `major` / `minor` / `patch` | +| `dry-run` | `false` | `true` computes the next version and stops: no commit, tag, release, or chart bump | +| `skip-tests` | `false` | currently unused — no step references it | +| `update-chart` | `false` | `true` also bumps the Helm chart `appVersion` (needs `HELM_CHART_UPDATE_TOKEN`) | + +### Jobs + +- **`bump`** — reads the version from `pychunkedgraph/_version.py`, bumps `part`, writes it back. Unless dry-run: commits `release vX.Y.Z`, tags `vX.Y.Z`, pushes the branch and the tag, then creates the GitHub Release. Needs `contents: write`. +- **`update-chart`** — opt-in: runs only when `update-chart` is `true` (and not dry-run). Checks out `CAVEconnectome/cave-helm-charts`, sets `charts/pychunkedgraph/Chart.yaml` `appVersion` to the new version, bumps the chart's own `version` by a patch, and pushes. Needs `HELM_CHART_UPDATE_TOKEN` (write access to the chart repo). + +## Cutting a release + +``` +# preview the next version, no writes +gh workflow run release.yml --ref pcgv3 -f part=patch -f dry-run=true + +# cut the release +gh workflow run release.yml --ref pcgv3 -f part=patch +``` + +Or via the Actions UI: **publish release → Run workflow →** pick the branch and `part`. + +The pushed tag is what the image build (`cloudbuild.yaml`) builds from. New tables created by that image are stamped with this version, and the server only serves tables whose major matches. The Helm chart `appVersion` bump that rolls the image out is opt-in via `update-chart`. diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 899f0431f..bd3f83cc9 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -4,19 +4,49 @@ on: push: branches: - "main" + - "pcgv3" pull_request: branches: - "main" + - "pcgv3" jobs: unit-tests: runs-on: ubuntu-latest steps: - name: Check out code - uses: actions/checkout@v2 + uses: actions/checkout@v4 - - name: Build image and run tests + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Build image + uses: docker/build-push-action@v6 + with: + context: . + load: true + tags: seunglab/pychunkedgraph:${{ github.sha }} + cache-from: type=gha + cache-to: type=gha,mode=max + + - name: Run tests with coverage run: | - docker build --tag seunglab/pychunkedgraph:$GITHUB_SHA . - docker run --rm seunglab/pychunkedgraph:$GITHUB_SHA /bin/sh -c "pytest --cov-config .coveragerc --cov=pychunkedgraph ./pychunkedgraph/tests && codecov" + docker run --name pcg-tests seunglab/pychunkedgraph:${{ github.sha }} \ + /bin/sh -c "pytest --cov-config .coveragerc --cov=pychunkedgraph --cov-report=xml:/app/coverage.xml ./pychunkedgraph/tests" + + - name: Copy coverage report from container + if: always() + run: docker cp pcg-tests:/app/coverage.xml ./coverage.xml + + - name: Upload coverage to Codecov + if: always() + uses: codecov/codecov-action@v5 + with: + files: ./coverage.xml + token: ${{ secrets.CODECOV_TOKEN }} + slug: CAVEconnectome/PyChunkedGraph + fail_ci_if_error: false + - name: Cleanup + if: always() + run: docker rm pcg-tests || true diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 6ee89f6c6..e56be70ff 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -20,6 +20,11 @@ on: type: boolean required: true default: false + update-chart: + description: "Also bump the pychunkedgraph Helm chart appVersion (needs HELM_CHART_UPDATE_TOKEN)" + type: boolean + required: false + default: false jobs: bump: @@ -36,37 +41,34 @@ jobs: steps: - name: Checkout uses: actions/checkout@v4 - - name: Get tags - run: git fetch --tags origin - name: Configure git for github-actions[bot] run: | git config --global user.name "github-actions[bot]" git config --global user.email "github-actions[bot]@users.noreply.github.com" - - name: Install Python - uses: actions/setup-python@v5 - with: - python-version: "3.11" - - name: Install bumpversion - run: pip install bumpversion - - name: Bump version with bumpversion - run: | - bumpversion ${{ github.event.inputs.part }} - - name: Commit and push with tags - if: ${{ github.event.inputs.dry-run == 'false' }} - run: git push --follow-tags - - name: Get version + - name: Bump version in pychunkedgraph/_version.py id: get-version run: | - version="$(git describe --tags)" - # remove the leading v from version - version="${version:1}" + cur=$(sed -n 's/^__version__ = "\(.*\)"/\1/p' pychunkedgraph/_version.py) + IFS=. read -r major minor patch <<< "$cur" + case "${{ github.event.inputs.part }}" in + major) major=$((major + 1)); minor=0; patch=0 ;; + minor) minor=$((minor + 1)); patch=0 ;; + patch) patch=$((patch + 1)) ;; + esac + version="$major.$minor.$patch" + sed -i "s/^__version__ = .*/__version__ = \"$version\"/" pychunkedgraph/_version.py echo "VERSION=$version" >> $GITHUB_OUTPUT - major_version="$(cut -d '.' -f 1 <<< $version)" - echo "MAJOR_VERSION=$major_version" >> $GITHUB_OUTPUT - minor_version="$(cut -d '.' -f 2 <<< $version)" - echo "MINOR_VERSION=$minor_version" >> $GITHUB_OUTPUT - short_version="$major_version.$minor_version" - echo "SHORT_VERSION=$short_version" >> $GITHUB_OUTPUT + echo "MAJOR_VERSION=$major" >> $GITHUB_OUTPUT + echo "MINOR_VERSION=$minor" >> $GITHUB_OUTPUT + echo "SHORT_VERSION=$major.$minor" >> $GITHUB_OUTPUT + - name: Commit the bump and push the tag + if: ${{ github.event.inputs.dry-run == 'false' }} + run: | + git add pychunkedgraph/_version.py + git commit -m "release v${{ steps.get-version.outputs.VERSION }}" + git tag "v${{ steps.get-version.outputs.VERSION }}" + git push origin HEAD + git push origin "v${{ steps.get-version.outputs.VERSION }}" - name: Show short version run: echo ${{ steps.get-version.outputs.SHORT_VERSION }} @@ -92,7 +94,7 @@ jobs: name: Update pychunkedgraph Helm chart runs-on: ubuntu-latest needs: bump - if: ${{ github.event.inputs.dry-run == 'false' }} + if: ${{ github.event.inputs.dry-run == 'false' && github.event.inputs.update-chart == 'true' }} permissions: contents: read id-token: write diff --git a/.gitignore b/.gitignore index 498253791..61b81289b 100644 --- a/.gitignore +++ b/.gitignore @@ -115,8 +115,11 @@ venv.bak/ # local dev stuff +*.code-workspace +.claude/ .devcontainer/ *.ipynb *.rdb /protobuf* -.DS_Store \ No newline at end of file +.DS_Store +*debug.py \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 000000000..70ceaed90 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,5 @@ +repos: + - repo: https://github.com/psf/black + rev: 26.1.0 + hooks: + - id: black diff --git a/.travis.yml b/.travis.yml deleted file mode 100644 index a5e33242d..000000000 --- a/.travis.yml +++ /dev/null @@ -1,60 +0,0 @@ -sudo: true -services: - docker - -env: - global: - - CLOUDSDK_CORE_DISABLE_PROMPTS=1 - -stages: - - test - - name: merge-deploy -python: 3.6 -notifications: - email: - on_success: change - on_failure: always - -jobs: - include: - - stage: test - name: "Running Tests" - language: minimal - before_script: - # request codecov to detect CI environment to pass through to docker - - ci_env=`bash <(curl -s https://codecov.io/env)` - - script: - - openssl aes-256-cbc -K $encrypted_506e835c2891_key -iv $encrypted_506e835c2891_iv -in key.json.enc -out key.json -d - - curl https://sdk.cloud.google.com | bash > /dev/null - - source "$HOME/google-cloud-sdk/path.bash.inc" - - gcloud auth activate-service-account --key-file=key.json - - gcloud auth configure-docker - - docker build --tag seunglab/pychunkedgraph:$TRAVIS_BRANCH . || travis_terminate 1 - - docker run $ci_env --rm seunglab/pychunkedgraph:$TRAVIS_BRANCH /bin/sh -c "tox -v -- --cov-config .coveragerc --cov=pychunkedgraph && codecov" - - - stage: merge-deploy - name: "version bump and merge into master" - language: python - install: - - pip install bumpversion - - before_script: - - "git clone https://gist.github.com/2c04596a45ccac57fe8dde0718ad58ee.git /tmp/travis-automerge" - - "chmod a+x /tmp/travis-automerge/auto_merge_travis_with_bumpversion.sh" - - script: - - "BRANCHES_TO_MERGE_REGEX='develop' BRANCH_TO_MERGE_INTO=master /tmp/travis-automerge/auto_merge_travis_with_bumpversion.sh" - - - stage: merge-deploy - name: "deploy to pypi" - language: python - install: - - pip install twine - - before_script: - - "git clone https://gist.github.com/cf9b261f26a1bf3fae6b59e7047f007a.git /tmp/travis-autodist" - - "chmod a+x /tmp/travis-autodist/pypi_dist.sh" - - script: - - "BRANCHES_TO_DIST='develop' /tmp/travis-autodist/pypi_dist.sh" diff --git a/Dockerfile b/Dockerfile index 2b7eeb151..9faebac84 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,11 +1,93 @@ -FROM caveconnectome/pychunkedgraph:base_042124 +# syntax=docker/dockerfile:1 +ARG PYTHON_VERSION=3.14 +# python:X-slim is the official upstream image. Pin by digest once a +# known-good build is identified so cache invalidation stays explicit; +# until then, the tag follows the latest 3.14 patch release. +ARG BASE_IMAGE=python:${PYTHON_VERSION}-slim + + +###################################################### +# Stage 1: Conda environment +###################################################### +FROM ${BASE_IMAGE} AS conda-deps +ENV PATH="/root/miniconda3/bin:${PATH}" + +RUN apt-get update && apt-get install build-essential wget -y \ + && wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \ + && bash Miniconda3-latest-Linux-x86_64.sh -b \ + && rm Miniconda3-latest-Linux-x86_64.sh \ + && conda config --add channels conda-forge \ + && conda update -y --override-channels -c conda-forge conda \ + && conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main \ + && conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r \ + && conda install -y --override-channels -c conda-forge conda-pack + +COPY requirements.yml requirements.txt requirements-dev.txt ./ + +# Solve against baseline x86-64. conda picks microarch-optimized builds from the +# *builder's* cpu (graph-tool ships v1/v3/v4 variants), so an image built on an +# AVX-512 host raises SIGILL on import wherever it later lands on an older node. +RUN --mount=type=cache,target=/opt/conda/pkgs \ + CONDA_OVERRIDE_ARCHSPEC=x86_64 conda env create -n pcg -f requirements.yml + +RUN conda-pack -n pcg --ignore-missing-files -o /tmp/env.tar \ + && mkdir -p /app/venv && cd /app/venv \ + && tar xf /tmp/env.tar && rm /tmp/env.tar \ + && /app/venv/bin/conda-unpack + + +###################################################### +# Stage 2: Bigtable emulator +###################################################### +FROM golang:bullseye AS bigtable-emulator +ARG GOOGLE_CLOUD_GO_VERSION=bigtable/v1.19.0 +RUN --mount=type=cache,target=/go/pkg/mod \ + --mount=type=cache,target=/root/.cache/go-build \ + git clone --depth=1 --branch="$GOOGLE_CLOUD_GO_VERSION" \ + https://github.com/googleapis/google-cloud-go.git /usr/src \ + && cd /usr/src/bigtable && go install -v ./cmd/emulator + + +###################################################### +# Stage 3: Production +###################################################### +FROM ${BASE_IMAGE} ENV VIRTUAL_ENV=/app/venv ENV PATH="$VIRTUAL_ENV/bin:$PATH" +# Force ld to resolve libpython3.x.so to the conda venv's copy. The +# slim base ships its own /usr/local/lib/libpython, which the loader +# would otherwise pair with conda-built C extensions, segfaulting in +# PyObject_Hash during the first import. +ENV LD_LIBRARY_PATH="$VIRTUAL_ENV/lib" +RUN apt-get update && apt-get install -y --no-install-recommends \ + nginx supervisor redis-tools procps \ + && (id nginx >/dev/null 2>&1 || useradd -r -d /home/nginx -s /bin/bash nginx) \ + && mkdir -p /etc/uwsgi /home/nginx/.cloudvolume/secrets \ + && chown -R nginx /home/nginx \ + && rm -rf /var/lib/apt/lists/* + +COPY --from=conda-deps /app/venv /app/venv +COPY --from=bigtable-emulator /go/bin/emulator /app/venv/bin/cbtemulator COPY override/gcloud /app/venv/bin/gcloud COPY override/timeout.conf /etc/nginx/conf.d/timeout.conf +COPY override/nginx.conf /etc/nginx/nginx.conf COPY override/supervisord.conf /etc/supervisor/conf.d/supervisord.conf +COPY uwsgi.ini /etc/uwsgi/uwsgi.ini + +# PyPI wheel bundles the zstd C source; conda-forge's system-linked +# build lacks `multi_decompress_to_buffer`, used by io/edges.py. +RUN pip install --no-cache-dir --no-deps --force-reinstall zstandard>=0.23.0 COPY requirements.txt . -RUN pip install --upgrade -r requirements.txt +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install --upgrade -r requirements.txt + COPY . /app +WORKDIR /app + +# --no-deps: graph-tool/cloudvolume etc. are already in the venv. __version__ comes from +# the committed pychunkedgraph/_version.py literal (bumped by the release workflow). +RUN pip install --no-deps -e . + +CMD ["/usr/bin/supervisord", "-n", "-c", "/etc/supervisor/supervisord.conf"] diff --git a/README.md b/README.md index ef888b3c6..c4325a65e 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ # PyChunkedGraph -[![Build Status](https://travis-ci.org/seung-lab/PyChunkedGraph.svg?branch=master)](https://travis-ci.org/seung-lab/PyChunkedGraph) -[![codecov](https://codecov.io/gh/seung-lab/PyChunkedGraph/branch/master/graph/badge.svg)](https://codecov.io/gh/seung-lab/PyChunkedGraph) +[![Tests](https://github.com/CAVEconnectome/PyChunkedGraph/actions/workflows/main.yml/badge.svg)](https://github.com/CAVEconnectome/PyChunkedGraph/actions/workflows/main.yml) +[![codecov](https://codecov.io/gh/CAVEconnectome/PyChunkedGraph/branch/main/graph/badge.svg)](https://codecov.io/gh/CAVEconnectome/PyChunkedGraph) The PyChunkedGraph is a proofreading and segmentation data management backend powering FlyWire and other proofreading platforms. It builds on an initial agglomeration of supervoxels and facilitates fast and parallel editing of connected components in the agglomeration graph by many users. @@ -35,6 +35,21 @@ The PyChunkedGraph can be locally deployed (`run_dev.py`), imported in a python As a backend the PyChunkedGraph can be combined with any frontend that adheres to its API. We use an adapted version of [neuroglancer](https://github.com/seung-lab/neuroglancer/tree/nkem-multicut) which is publicly available. +## Release + +PyChunkedGraph is versioned by a committed `pychunkedgraph/_version.py` literal, bumped by a +one-click workflow — no manual edit. Versioning is per branch: `main` is 2.x, `pcgv3` is 3.x. + +- **Release:** Actions → **publish release** → **Run workflow** (from the target branch) → choose + `part` (`major`/`minor`/`patch`), or `gh workflow run release.yml --ref -f part=patch`. + It bumps `_version.py`, commits, tags `vX.Y.Z`, pushes, and creates a GitHub Release; the Cloud + Build trigger builds the image from the tag. Set `update-chart=true` to also bump the Helm chart. +- **Preview:** `dry-run=true` prints the next version without committing or tagging. + +New tables are stamped with this version at creation; the server only serves tables whose major +matches. See [.github/workflows/README.md](.github/workflows/README.md) for full CI/release detail. + + ## Publication When using or referencing the PyChunkedGraph, please use the citation below. The FlyWire paper described and published the PyChunkedGraph v1. diff --git a/base.Dockerfile b/base.Dockerfile deleted file mode 100644 index b5123e137..000000000 --- a/base.Dockerfile +++ /dev/null @@ -1,70 +0,0 @@ -ARG PYTHON_VERSION=3.11 -ARG BASE_IMAGE=tiangolo/uwsgi-nginx-flask:python${PYTHON_VERSION} - - -###################################################### -# Build Image - PCG dependencies -###################################################### -FROM ${BASE_IMAGE} AS pcg-build -ENV PATH="/root/miniconda3/bin:${PATH}" -ENV CONDA_ENV="pychunkedgraph" - -# Setup Miniconda -RUN apt-get update && apt-get install build-essential wget -y -RUN wget \ - https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \ - && mkdir /root/.conda \ - && bash Miniconda3-latest-Linux-x86_64.sh -b \ - && rm -f Miniconda3-latest-Linux-x86_64.sh \ - && conda update conda - -# Install PCG dependencies - especially graph-tool -# Note: uwsgi has trouble with pip and python3.11, so adding this with conda, too -COPY requirements.txt . -COPY requirements.yml . -COPY requirements-dev.txt . -RUN conda env create -n ${CONDA_ENV} -f requirements.yml - -# Shrink conda environment into portable non-conda env -RUN conda install conda-pack -c conda-forge - -RUN conda-pack -n ${CONDA_ENV} --ignore-missing-files -o /tmp/env.tar \ - && mkdir -p /app/venv \ - && cd /app/venv \ - && tar xf /tmp/env.tar \ - && rm /tmp/env.tar -RUN /app/venv/bin/conda-unpack - - -###################################################### -# Build Image - Bigtable Emulator (without Google SDK) -###################################################### -FROM golang:bullseye as bigtable-emulator-build -RUN mkdir -p /usr/src -WORKDIR /usr/src -ENV GOOGLE_CLOUD_GO_VERSION bigtable/v1.19.0 -RUN apt-get update && apt-get install git -y -RUN git clone --depth=1 --branch="$GOOGLE_CLOUD_GO_VERSION" https://github.com/googleapis/google-cloud-go.git . \ - && cd bigtable \ - && go install -v ./cmd/emulator - - -###################################################### -# Production Image -###################################################### -FROM ${BASE_IMAGE} -ENV VIRTUAL_ENV=/app/venv -ENV PATH="$VIRTUAL_ENV/bin:$PATH" - -COPY --from=pcg-build /app/venv /app/venv -COPY --from=bigtable-emulator-build /go/bin/emulator /app/venv/bin/cbtemulator -COPY override/gcloud /app/venv/bin/gcloud -COPY override/timeout.conf /etc/nginx/conf.d/timeout.conf -COPY override/supervisord.conf /etc/supervisor/conf.d/supervisord.conf -# Hack to get zstandard from PyPI - remove if conda-forge linked lib issue is resolved -RUN pip install --no-cache-dir --no-deps --force-reinstall zstandard==0.21.0 -COPY . /app - -RUN mkdir -p /home/nginx/.cloudvolume/secrets \ - && chown -R nginx /home/nginx \ - && usermod -d /home/nginx -s /bin/bash nginx diff --git a/build_pypi.sh b/build_pypi.sh deleted file mode 100644 index c952f5cb4..000000000 --- a/build_pypi.sh +++ /dev/null @@ -1,4 +0,0 @@ -#!/bin/sh - -python setup.py sdist -twine upload dist/* diff --git a/cloudbuild.yaml b/cloudbuild.yaml index 21f4cc58d..434248a9b 100644 --- a/cloudbuild.yaml +++ b/cloudbuild.yaml @@ -5,19 +5,26 @@ steps: args: ["-c", "docker login --username=$$USERNAME --password=$$PASSWORD"] secretEnv: ["USERNAME", "PASSWORD"] + # Build + push in one BuildKit invocation using a docker-container + # builder (required for registry-type cache export). The builder + # `--use` setting is client-side and doesn't persist across cloudbuild + # steps, so create + use + build must happen in a single step. + # Registry cache at the fixed :buildcache tag lets unchanged stages + # (conda env, bigtable emulator, pip install) reuse the previous + # build's exact layer artifacts, so already-warm nodes only download + # what actually changed on pull. - name: "gcr.io/cloud-builders/docker" entrypoint: "bash" args: - "-c" - | - docker build -t $$USERNAME/pychunkedgraph:$TAG_NAME . - timeout: 600s - secretEnv: ["USERNAME"] - - # Push the final image to Dockerhub - - name: "gcr.io/cloud-builders/docker" - entrypoint: "bash" - args: ["-c", "docker push $$USERNAME/pychunkedgraph:$TAG_NAME"] + docker buildx create --use --name pcg-builder --driver docker-container + docker buildx build \ + --cache-from type=registry,ref=$$USERNAME/pychunkedgraph:buildcache \ + --cache-to type=registry,ref=$$USERNAME/pychunkedgraph:buildcache,mode=max \ + --push \ + -t $$USERNAME/pychunkedgraph:$TAG_NAME . + timeout: 1800s secretEnv: ["USERNAME"] availableSecrets: diff --git a/codecov.yml b/codecov.yml new file mode 100644 index 000000000..fc04b242e --- /dev/null +++ b/codecov.yml @@ -0,0 +1,19 @@ +codecov: + require_ci_to_pass: true + +coverage: + status: + project: + default: + target: auto + threshold: 1% + informational: true + patch: + default: + target: 1% + informational: true + +comment: + layout: "reach,diff,flags,files" + behavior: default + require_changes: false diff --git a/compile_reqs.sh b/compile_reqs.sh deleted file mode 100755 index 2d74c225d..000000000 --- a/compile_reqs.sh +++ /dev/null @@ -1 +0,0 @@ -docker run -v ${PWD}:/app caveconnectome/pychunkedgraph:v2.4.0 /bin/bash -c "pip install pip-tools && pip-compile requirements.in --resolver=backtracking -v --output-file requirements.txt" \ No newline at end of file diff --git a/docs/Readme.md b/docs/Readme.md index 45799326e..c05ad6979 100644 --- a/docs/Readme.md +++ b/docs/Readme.md @@ -10,7 +10,7 @@ pip install -r requirements.txt ## Multiprocessing -Check out [multiprocessing.md](https://github.com/seung-lab/PyChunkedGraph/blob/master/src/pychunkedgraph/multiprocessing.md) for how to use the multiprocessing tools implemented for the ChunkedGraph +Check out [multiprocessing.md](https://github.com/CAVEconnectome/PyChunkedGraph/blob/master/src/pychunkedgraph/multiprocessing.md) for how to use the multiprocessing tools implemented for the ChunkedGraph ## Credentials @@ -30,7 +30,7 @@ The current version of the ChunkedGraph contains supervoxels from `gs://nkem/bas ### Building the graph -[buildgraph.md](https://github.com/seung-lab/PyChunkedGraph/blob/master/src/pychunkedgraph/buildgraph.md) explains how to build a graph from scratch. +[buildgraph.md](https://github.com/CAVEconnectome/PyChunkedGraph/blob/master/src/pychunkedgraph/buildgraph.md) explains how to build a graph from scratch. ### Initialization diff --git a/docs/edges.md b/docs/edges.md index 9dc15a98b..ccda4205b 100644 --- a/docs/edges.md +++ b/docs/edges.md @@ -2,7 +2,7 @@ PyChunkedgraph uses protobuf for serialization and zstandard for compression. -Edges and connected components per chunk are stored using the protobuf definitions in [`pychunkedgraph.io.protobuf`](https://github.com/seung-lab/PyChunkedGraph/pychunkedgraph/io/protobuf/chunkEdges.proto). +Edges and connected components per chunk are stored using the protobuf definitions in [`pychunkedgraph.io.protobuf`](https://github.com/CAVEconnectome/PyChunkedGraph/pychunkedgraph/io/protobuf/chunkEdges.proto). This format is a result of performance tests. It provided the best tradeoff between deserialzation speed and storage size. diff --git a/docs/edges_and_components.md b/docs/edges_and_components.md index 9dc15a98b..ccda4205b 100644 --- a/docs/edges_and_components.md +++ b/docs/edges_and_components.md @@ -2,7 +2,7 @@ PyChunkedgraph uses protobuf for serialization and zstandard for compression. -Edges and connected components per chunk are stored using the protobuf definitions in [`pychunkedgraph.io.protobuf`](https://github.com/seung-lab/PyChunkedGraph/pychunkedgraph/io/protobuf/chunkEdges.proto). +Edges and connected components per chunk are stored using the protobuf definitions in [`pychunkedgraph.io.protobuf`](https://github.com/CAVEconnectome/PyChunkedGraph/pychunkedgraph/io/protobuf/chunkEdges.proto). This format is a result of performance tests. It provided the best tradeoff between deserialzation speed and storage size. diff --git a/docs/precomputed_ocdbt_hybrid.md b/docs/precomputed_ocdbt_hybrid.md new file mode 100644 index 000000000..d1caafd5f --- /dev/null +++ b/docs/precomputed_ocdbt_hybrid.md @@ -0,0 +1,101 @@ +# Hybrid base: precomputed + OCDBT fork (proposal) + +Status: proposal, not implemented. Open question is whether storage and ingest-compute savings justify the read-path complexity. + +## Problem + +PCG ingest copies the entire watershed segmentation into `/ocdbt/base/` in OCDBT format before any CG edit can happen. Per-CG forks at `/ocdbt//` store only the deltas from SV splits. Two costs follow: + +- **Storage**: roughly 2× the segmentation footprint per dataset — original precomputed plus full OCDBT copy. +- **Ingest compute**: a per-chunk pass that reads the precomputed and writes it through the OCDBT driver. Hours of cluster time on TB-scale datasets. + +Both costs are paid up-front, before any user has done a single edit. The proposal here: skip the base copy and serve unedited chunks directly from the raw precomputed directory. Per-CG OCDBT forks remain as the delta store. + +## Why the current architecture has the base copy + +Today's per-CG read spec is: + +``` +neuroglancer_precomputed + └─ kvstore: ocdbt + ├─ base: kvstack [base_layer, fork_manifest, fork_data] + └─ config: { compression, max_inline_value_bytes, ... } +``` + +When a reader asks for chunk key `8_8_40/1024-..._0-128`: + +1. The `neuroglancer_precomputed` driver passes the chunk key to its kvstore (the OCDBT driver). +2. OCDBT looks up the key **in its B+tree**. The B+tree's leaves map chunk keys to values. +3. If the key isn't in the B+tree, OCDBT returns not-found. It does not consult the kvstack any further. + +The three kvstack layers serve OCDBT's *internal* storage (B+tree manifest + node blobs + leaf blobs) — they have no visibility into chunk-key lookups. So the OCDBT B+tree must contain every chunk key the reader will ever ask for, and that's why ingest copies the whole watershed: to populate the B+tree. + +## What tensorstore primitives provide + +Confirmed against tensorstore docs: + +- **`kvstack` routes by exact / prefix match, with no fallthrough on miss.** A layer that claims a key range absorbs misses — they return `state='missing'` and do not cascade to the next layer. So we can't put raw precomputed below an OCDBT layer in a kvstack and expect kvstack to fall through when OCDBT doesn't have a key. +- **No native overlay/fallback kvstore driver.** `kvstack` is the only composition primitive at the kvstore level; it's precedence-based, not fallthrough. +- **OCDBT has no external-blob references.** B+tree leaves either inline the value or point to a data file under the OCDBT directory. There's no way to make a leaf reference a raw GCS precomputed file. +- **Array-level `stack` / `ts.overlay`** layers arrays by spatial domain. In overlapping regions, the later layer takes absolute precedence — missing-in-later does not fall back to earlier. + +No single tensorstore primitive provides "try OCDBT delta first, fall through to raw precomputed on miss." + +## Architectural options + +### A — Two-stage read at the pcg layer + +PCG reads open two handles: the OCDBT fork for the delta, and a raw `neuroglancer_precomputed` reader for the watershed base. For any voxel region, issue both reads and merge with "delta wins where present, base fills the rest." + +- **Pros**: works inside pcg (`lookup_svs_from_seg`, sanity checks, debug tools) without any tensorstore changes. +- **Cons**: every pcg caller that uses `meta.ws_ocdbt` needs to route through a new merging reader. Neuroglancer doesn't benefit — it still gets a single kvstore spec from `dataset_info`. Either NG runs two layers itself (Option B) or we stand up a server-side proxy that does the merge before serving. + +### B — NG-side layer stack + +`dataset_info` publishes two precomputed layers: the raw watershed (read-only base) and the per-CG OCDBT fork (delta). NG composites them — visible segmentation is whichever has data at a given chunk. + +- **Pros**: no change to pcg's read path. Pushes the architecture complexity into the viewer. +- **Cons**: requires NG to treat "missing chunk in delta" as "fall through to base," not "render as background." Default NG behavior is the latter, so a viewer-side or proxy-side shim is likely needed. + +### C — Custom tensorstore kvstore driver + +A new "fallthrough" kvstore driver: read tries layer N, falls through on miss to layer N−1. Implement upstream in tensorstore or fork-and-maintain. + +- **Pros**: cleanest consumer-facing story — pcg and NG both keep using a single kvstore spec. +- **Cons**: tensorstore kvstore drivers are C++. Non-trivial maintenance surface; review/merge timeline if upstreaming. + +### D — Lazy base population (not a win on its own) + +Skip the ingest copy; copy a chunk from precomputed to OCDBT on first edit. Saves ingest compute. Does **not** save storage for reads — unedited chunks still 404 in OCDBT for a reader that doesn't have a fallback. Only useful in combination with A/B/C. + +## Recommendation + +Measure first. Confirm the actual storage and ingest-compute savings on a real dataset and weigh against the engineering cost of A/B/C. + +If the savings justify the work, **A + B together** is the most pragmatic path: +- A gives pcg a single merged-read API. Edits, sanity checks, debug tooling keep working. +- B avoids standing up a proxy service for the viewer by letting NG handle the overlay. + +Both require upstream verification: +- **For A**: confirm that `(x0:x1, y0:y1, z0:z1)` reads on an OCDBT with sparse keys surface missing-ness *per chunk* at the `neuroglancer_precomputed` array layer (not per-region, not silently fill-valued). +- **For B**: confirm NG's segmentation loader can be configured to fall through gaps in one layer to another. If it can't, build a small server-side merging shim — at which point Option A's reader becomes that shim and B reduces to "publish two specs." + +C is the cleanest design but carries the highest cost. Pursue only if A/B turn out to have unworkable semantics. + +## Open questions before any implementation + +1. Does OCDBT's `read_result.state == 'missing'` surface per-chunk at the `neuroglancer_precomputed` array layer, or does the array silently fill missing chunks with fill-value? Verifiable by opening an OCDBT with sparse keys and reading a region that spans present + missing chunks. +2. Does NG distinguish "chunk returned as missing" from "chunk is all fill-value"? If not, a viewer-side overlay needs a shim regardless. +3. What's the actual delta volume per CG over its lifetime? If SV splits eventually touch a significant fraction of chunks, the storage win shrinks toward zero — at which point the simpler architecture (today's full base copy) wins on engineering cost. + +## Files to start from when implementing + +- `pychunkedgraph/graph/ocdbt.py` — spec construction (`build_cg_ocdbt_spec`), base population (`create_base_ocdbt`), fork setup (`fork_base_manifest`). +- `pychunkedgraph/ingest/cli.py`, `pychunkedgraph/ingest/cluster.py` — current base-copy flow. +- `pychunkedgraph/graph/utils/generic.py::get_local_segmentation` — single pcg read entry point that would need the two-stage merge in Option A. + +## Verification (per chosen option) + +- **A**: unit test that simulates a partial-delta OCDBT + raw precomputed and confirms the pcg reader returns the correct labels for spans crossing both. +- **B**: configure an NG link with both layers against a test dataset; compare the rendered segmentation to a known-good reference at edited and unedited regions. +- **C**: a tensorstore build with the new driver passes a fallthrough test (missing key in upper layer resolves from lower layer). diff --git a/docs/segmentation_preprocessing.md b/docs/segmentation_preprocessing.md index 3fb1bf59b..028419a3a 100644 --- a/docs/segmentation_preprocessing.md +++ b/docs/segmentation_preprocessing.md @@ -32,10 +32,10 @@ There are three types of edges: 2. `cross_chunk`: edges between parts of "the same" supervoxel in the unchunked segmentation that has been split across chunk boundary 3. `between_chunk`: edges between supervoxels across chunks -Every pair of touching supervoxels has an edge between them. All edges are stored using [protobuf](https://github.com/seung-lab/PyChunkedGraph/blob/pcgv2/pychunkedgraph/io/protobuf/chunkEdges.proto). During ingest only edges of type 2. and 3. are copied into BigTable, whereas edges of type 1. are always read from storage to reduce cost. Similar to the supervoxel segmentation, we recommed storing these on GCloud in the same zone the ChunkedGraph server will be deployed in to reduce latency. +Every pair of touching supervoxels has an edge between them. All edges are stored using [protobuf](https://github.com/CAVEconnectome/PyChunkedGraph/blob/pcgv2/pychunkedgraph/io/protobuf/chunkEdges.proto). During ingest only edges of type 2. and 3. are copied into BigTable, whereas edges of type 1. are always read from storage to reduce cost. Similar to the supervoxel segmentation, we recommed storing these on GCloud in the same zone the ChunkedGraph server will be deployed in to reduce latency. To denote which edges form a connected component within a chunk, a component mapping needs to be created. This mapping is only used during ingest. -More details on how to create these protobuf files can be found [here](https://github.com/seung-lab/PyChunkedGraph/blob/pcgv2/docs/storage.md). +More details on how to create these protobuf files can be found [here](https://github.com/CAVEconnectome/PyChunkedGraph/blob/pcgv2/docs/storage.md). diff --git a/override/nginx.conf b/override/nginx.conf new file mode 100644 index 000000000..3659d6977 --- /dev/null +++ b/override/nginx.conf @@ -0,0 +1,43 @@ +user nginx; +worker_processes 1; + +error_log /var/log/nginx/error.log warn; +pid /var/run/nginx.pid; + +daemon off; + +events { + worker_connections 1024; +} + +http { + include /etc/nginx/mime.types; + default_type application/octet-stream; + + log_format main '$remote_addr - $remote_user [$time_local] "$request" ' + '$status $body_bytes_sent "$http_referer" ' + '"$http_user_agent" "$http_x_forwarded_for"'; + + map $http_user_agent $loggable { + default 1; + ~*kube-probe 0; + } + + access_log /var/log/nginx/access.log main if=$loggable; + + sendfile on; + + keepalive_timeout 65; + + client_max_body_size 0; + + include /etc/nginx/conf.d/*.conf; + + server { + listen 80; + location / { + include uwsgi_params; + uwsgi_pass unix:///tmp/uwsgi.sock; + } + } +} diff --git a/pychunkedgraph/__init__.py b/pychunkedgraph/__init__.py index e615ea2b7..0f0587232 100644 --- a/pychunkedgraph/__init__.py +++ b/pychunkedgraph/__init__.py @@ -1 +1,98 @@ -__version__ = "2.21.1" +from pychunkedgraph._version import __version__ # noqa: F401 + +import sys +import warnings +import logging as stdlib_logging # Use alias to avoid conflict with pychunkedgraph.logging + +# Suppress annoying warning from python_jsonschema_objects dependency +warnings.filterwarnings( + "ignore", message="Schema id not specified", module="python_jsonschema_objects" +) + +# Custom log level between INFO (20) and WARNING (30) +# Use logger.notice() for pychunkedgraph logs that should always show +# even when third-party INFO is suppressed +NOTICE = 25 +stdlib_logging.addLevelName(NOTICE, "NOTICE") + +# Diagnostic level above DEBUG (10) but below INFO (20). Lets the user +# enable per-stage timing/count summaries without the much noisier DEBUG +# tracing — set the logger to VERBOSE for performance/correctness diagnosis. +VERBOSE = 15 +stdlib_logging.addLevelName(VERBOSE, "VERBOSE") + + +class PCGLogger(stdlib_logging.Logger): + def note(self, message, *args, **kwargs): + if self.isEnabledFor(NOTICE): + self._log(NOTICE, message, args, stacklevel=2, **kwargs) + + def verbose(self, message, *args, **kwargs): + if self.isEnabledFor(VERBOSE): + self._log(VERBOSE, message, args, stacklevel=2, **kwargs) + + +stdlib_logging.setLoggerClass(PCGLogger) + + +def get_logger(name: str) -> PCGLogger: + return stdlib_logging.getLogger(name) # type: ignore[return-value] + + +# Export logging levels for convenience +DEBUG = stdlib_logging.DEBUG +INFO = stdlib_logging.INFO +WARNING = stdlib_logging.WARNING +ERROR = stdlib_logging.ERROR + +# Set up library-level logger with NullHandler (Python logging best practice) +stdlib_logging.getLogger(__name__).addHandler(stdlib_logging.NullHandler()) + + +def configure_logging(level=stdlib_logging.INFO, format_str=None, stream=None): + """ + Configure logging for pychunkedgraph. Call this to enable log output. + + Works in Jupyter notebooks and scripts. + + Args: + level: Logging level (default: INFO). Use pychunkedgraph.DEBUG, .INFO, .WARNING, .ERROR + format_str: Custom format string (optional) + stream: Output stream (default: sys.stdout for Jupyter compatibility) + + Example: + import pychunkedgraph + pychunkedgraph.configure_logging() # Enable INFO level logging + pychunkedgraph.configure_logging(pychunkedgraph.DEBUG) # Enable DEBUG level + """ + if format_str is None: + format_str = "%(asctime)s [%(module)s:%(funcName)s:%(lineno)d] %(message)s" + if stream is None: + stream = sys.stdout + + # Get root logger for pychunkedgraph + logger = stdlib_logging.getLogger(__name__) + logger.setLevel(level) + + # Remove existing handlers and add fresh StreamHandler + # This allows reconfiguring with different levels/formats + for h in logger.handlers[:]: + if isinstance(h, stdlib_logging.StreamHandler) and not isinstance( + h, stdlib_logging.NullHandler + ): + logger.removeHandler(h) + + handler = stdlib_logging.StreamHandler(stream) + handler.setLevel(level) + formatter = stdlib_logging.Formatter(format_str) + formatter.default_msec_format = "%s.%03d" + handler.setFormatter(formatter) + logger.addHandler(handler) + # the package has its own handler; propagating would print every record a + # second time through any root handler (e.g. an entrypoint's basicConfig) + logger.propagate = False + + return logger + + +configure_logging(level=NOTICE) diff --git a/pychunkedgraph/_version.py b/pychunkedgraph/_version.py new file mode 100644 index 000000000..348e212c1 --- /dev/null +++ b/pychunkedgraph/_version.py @@ -0,0 +1,3 @@ +# Single source of truth for the package version, bumped by the release workflow on a +# semver release. Versioning is per branch: main is 2.x, pcgv3 is 3.x. +__version__ = "3.1.7" diff --git a/pychunkedgraph/app/__init__.py b/pychunkedgraph/app/__init__.py index 3e938628b..d724fbb84 100644 --- a/pychunkedgraph/app/__init__.py +++ b/pychunkedgraph/app/__init__.py @@ -14,6 +14,7 @@ from flask_cors import CORS from rq import Queue +from pychunkedgraph import NOTICE, configure_logging from pychunkedgraph.logging import jsonformatter from . import config @@ -95,16 +96,38 @@ def configure_app(app): formatter.converter = time.gmtime handler.setFormatter(formatter) app.logger.removeHandler(default_handler) + logging.getLogger().removeHandler(default_handler) app.logger.addHandler(handler) app.logger.setLevel(app.config["LOGGING_LEVEL"]) app.logger.propagate = False + # Ensure pychunkedgraph logger always works at NOTICE level + # regardless of app config or environment log level + configure_logging(level=NOTICE) + pcg_logger = logging.getLogger("pychunkedgraph") + # Root logger on the server image has a BASIC_FORMAT StreamHandler + # (installed by uwsgi/gunicorn or an upstream basicConfig); propagating + # past our own handler would re-emit every record in the + # `LEVELNAME:logger.name:message` form. + pcg_logger.propagate = False + # app.logger.propagate = False blocks children under pychunkedgraph.app + # from reaching the pychunkedgraph handler — attach it directly + app_ns_logger = logging.getLogger("pychunkedgraph.app") + for h in pcg_logger.handlers: + if isinstance(h, logging.StreamHandler) and not isinstance( + h, logging.NullHandler + ): + app_ns_logger.addHandler(h) + break + if app.config["USE_REDIS_JOBS"]: app.redis = redis.Redis.from_url(app.config["REDIS_URL"]) app.test_q = Queue("test", connection=app.redis) with app.app_context(): from ..ingest.rq_cli import init_rq_cmds from ..ingest.cli import init_ingest_cmds + from ..ingest.cli_upgrade import init_upgrade_cmds init_rq_cmds(app) init_ingest_cmds(app) + init_upgrade_cmds(app) diff --git a/pychunkedgraph/app/app_utils.py b/pychunkedgraph/app/app_utils.py index b46e4b192..870e9859d 100644 --- a/pychunkedgraph/app/app_utils.py +++ b/pychunkedgraph/app/app_utils.py @@ -6,17 +6,15 @@ from functools import wraps import numpy as np -import networkx as nx import requests from flask import current_app, json, request -from scipy import spatial from werkzeug.datastructures import ImmutableMultiDict from pychunkedgraph import __version__ from pychunkedgraph.graph import ChunkedGraph -from pychunkedgraph.graph.client import get_default_client_info +from pychunkedgraph.graph import get_default_client_info from pychunkedgraph.graph import exceptions as cg_exceptions - +from pychunkedgraph.graph.sv_lookup import resolve_supervoxels_at_coords PCG_CACHE = {} @@ -215,49 +213,12 @@ def tobinary_multiples(arr): def handle_supervoxel_id_lookup( cg, coordinates: Sequence[Sequence[int]], node_ids: Sequence[np.uint64] ) -> Sequence[np.uint64]: - """ - Helper to lookup supervoxel ids. - This takes care of grouping coordinates. - """ + """Resolve voxel coordinates to current supervoxel ids. - def ccs(coordinates_nm_): - graph = nx.Graph() - dist_mat = spatial.distance.cdist(coordinates_nm_, coordinates_nm_) - for edge in np.array(np.where(dist_mat < 1000)).T: - graph.add_edge(*edge) - ccs = [np.array(list(cc)) for cc in nx.connected_components(graph)] - return ccs - - coordinates = np.array(coordinates, dtype=int) - coordinates_nm = coordinates * cg.meta.resolution - max_dist_steps = np.array([4, 8, 14, 28], dtype=float) * np.mean(cg.meta.resolution) - - node_ids = np.array(node_ids, dtype=np.uint64) - if len(coordinates.shape) != 2: - raise cg_exceptions.BadRequest( - f"Could not determine supervoxel ID for coordinates " - f"{coordinates} - Validation stage." - ) - - atomic_ids = np.zeros(len(coordinates), dtype=np.uint64) - for node_id in np.unique(node_ids): - node_id_m = node_ids == node_id - for cc in ccs(coordinates_nm[node_id_m]): - m_ids = np.where(node_id_m)[0][cc] - - for max_dist_nm in max_dist_steps: - atomic_ids_sub = cg.get_atomic_ids_from_coords( - coordinates[m_ids], parent_id=node_id, max_dist_nm=max_dist_nm - ) - if atomic_ids_sub is not None: - break - if atomic_ids_sub is None: - raise cg_exceptions.BadRequest( - f"Could not determine supervoxel ID for coordinates " - f"{coordinates} - Lookup stage." - ) - atomic_ids[m_ids] = atomic_ids_sub - return atomic_ids + Thin app-layer wrapper. The 2D/3D resolution contract lives in + :func:`pychunkedgraph.graph.sv_lookup.resolve_supervoxels_at_coords`. + """ + return resolve_supervoxels_at_coords(cg, coordinates, node_ids) def get_username_dict(user_ids, auth_token) -> dict: diff --git a/pychunkedgraph/app/common.py b/pychunkedgraph/app/common.py index 237e11fc0..f29b482ac 100644 --- a/pychunkedgraph/app/common.py +++ b/pychunkedgraph/app/common.py @@ -4,9 +4,8 @@ import json import time import traceback -from datetime import datetime +from datetime import datetime, timezone -from cloudvolume import compression from google.api_core.exceptions import GoogleAPIError from flask import current_app, g, jsonify, request @@ -50,7 +49,7 @@ def _log_request(response_time): def before_request(): current_app.request_start_time = time.time() - current_app.request_start_date = datetime.utcnow() + current_app.request_start_date = datetime.now(timezone.utc) try: current_app.user_id = g.auth_user["id"] except (AttributeError, KeyError): @@ -60,6 +59,8 @@ def before_request(): current_app.request_type = None content_encoding = request.headers.get("Content-Encoding", "") if "gzip" in content_encoding.lower(): + from cloudvolume import compression + request.data = compression.decompress(request.data, "gzip") @@ -80,6 +81,8 @@ def after_request(response): ): return response + from cloudvolume import compression + response.data = compression.gzip_compress(response.data) response.headers["Content-Encoding"] = "gzip" response.headers["Vary"] = "Accept-Encoding" diff --git a/pychunkedgraph/app/meshing/common.py b/pychunkedgraph/app/meshing/common.py index 8f1a0c20a..af7f6f7a6 100644 --- a/pychunkedgraph/app/meshing/common.py +++ b/pychunkedgraph/app/meshing/common.py @@ -4,19 +4,11 @@ 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 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 - __meshing_url_prefix__ = os.environ.get("MESHING_URL_PREFIX", "meshing") @@ -43,6 +35,9 @@ def home(): def handle_valid_frags(table_id, node_id): + # nested: pulls meshing/cloudvolume, only needed at call time + from pychunkedgraph.meshing.manifest import get_highest_child_nodes_with_meshes + current_app.table_id = table_id cg = app_utils.get_cg(table_id) seg_ids = get_highest_child_nodes_with_meshes( @@ -97,6 +92,7 @@ def handle_get_manifest(table_id, node_id): def manifest_response(cg, args): from pychunkedgraph.meshing.manifest import speculative_manifest_sharded + from pychunkedgraph.meshing.manifest import get_highest_child_nodes_with_meshes ( node_id, @@ -145,40 +141,21 @@ def _check_post_options(cg, resp, data, seg_ids): def handle_remesh(table_id): current_app.request_type = "remesh_enque" current_app.table_id = table_id - is_priority = request.args.get("priority", True, type=str2bool) - is_redisjob = request.args.get("use_redis", False, type=str2bool) - new_lvl2_ids = json.loads(request.data)["new_lvl2_ids"] - - if is_redisjob: - with Connection(redis.from_url(current_app.config["REDIS_URL"])): - - if is_priority: - retry = Retry(max=3, interval=[1, 10, 60]) - queue_name = "mesh-chunks" - else: - retry = Retry(max=3, interval=[60, 60, 60]) - queue_name = "mesh-chunks-low-priority" - q = Queue(queue_name, retry=retry, default_timeout=1200) - task = q.enqueue(meshing_tasks.remeshing, table_id, new_lvl2_ids) - - response_object = {"status": "success", "data": {"task_id": task.get_id()}} - - return jsonify(response_object), 202 - else: - 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() - - return Response(status=202) + 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() + return Response(status=202) def _remeshing(serialized_cg_info, lvl2_nodes): + # nested: pulls meshing/cloudvolume, only needed at call time + from pychunkedgraph.meshing import meshgen + 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"] @@ -202,5 +179,21 @@ def _remeshing(serialized_cg_info, lvl2_nodes): def clear_manifest_cache(cg, node_id): + # nested: pulls meshing/cloudvolume, only needed at call time + from pychunkedgraph.meshing.manifest import get_children_before_start_layer + from pychunkedgraph.meshing.manifest import ManifestCache + node_ids = get_children_before_start_layer(cg, node_id, start_layer=2) ManifestCache(cg.graph_id).clear_fragments(node_ids) + + +def clear_manifest_cache_all(cg) -> int: + """Delete every cached manifest fragment for this graph. + + Returns the number of redis keys deleted across both initial and + dynamic caches (they share the ``:`` namespace). + """ + # nested: pulls meshing/cloudvolume, only needed at call time + from pychunkedgraph.meshing.manifest import ManifestCache + + return ManifestCache(cg.graph_id).clear_namespace() diff --git a/pychunkedgraph/app/meshing/tasks.py b/pychunkedgraph/app/meshing/tasks.py index a1f11ca68..e6550ed0d 100644 --- a/pychunkedgraph/app/meshing/tasks.py +++ b/pychunkedgraph/app/meshing/tasks.py @@ -1,10 +1,12 @@ from pychunkedgraph.app import app_utils -from pychunkedgraph.meshing import meshgen, meshgen_utils import numpy as np import os def remeshing(table_id, lvl2_nodes): + # nested: pulls meshing/cloudvolume, only needed at call time + from pychunkedgraph.meshing import meshgen + lvl2_nodes = np.array(lvl2_nodes, dtype=np.uint64) cg = app_utils.get_cg(table_id, skip_cache=True) diff --git a/pychunkedgraph/app/meshing/v1/routes.py b/pychunkedgraph/app/meshing/v1/routes.py index dda067e90..14b286e3a 100644 --- a/pychunkedgraph/app/meshing/v1/routes.py +++ b/pychunkedgraph/app/meshing/v1/routes.py @@ -9,7 +9,6 @@ from pychunkedgraph.app.app_utils import get_cg from pychunkedgraph.app.app_utils import remap_public - bp = Blueprint( "pcg_meshing_v1", __name__, url_prefix=f"/{common.__meshing_url_prefix__}/api/v1" ) @@ -98,3 +97,12 @@ def handle_remesh(table_id): def handle_clear_manifest_cache(table_id, node_id): cg = get_cg(table_id) common.clear_manifest_cache(cg, node_id) + + +@bp.route("/table//clear_manifest_cache", methods=["GET"]) +@auth_requires_permission("edit") +def handle_clear_manifest_cache_all(table_id): + """Drop every cached manifest fragment for this graph.""" + cg = get_cg(table_id) + deleted = common.clear_manifest_cache_all(cg) + return {"deleted": deleted} diff --git a/pychunkedgraph/app/segmentation/common.py b/pychunkedgraph/app/segmentation/common.py index 3250248f2..4c0f131a3 100644 --- a/pychunkedgraph/app/segmentation/common.py +++ b/pychunkedgraph/app/segmentation/common.py @@ -2,23 +2,23 @@ import json import os +import pickle import time -from datetime import datetime +from datetime import datetime, timezone from functools import reduce from collections import deque, defaultdict import numpy as np import pandas as pd from flask import current_app, g, jsonify, make_response, request +from messagingclient import MessagingClient from pytz import UTC -from pychunkedgraph import __version__ +from pychunkedgraph import __version__, get_logger + +logger = get_logger(__name__) from pychunkedgraph.app import app_utils -from pychunkedgraph.graph import ( - attributes, - cutting, - segmenthistory, -) +from pychunkedgraph.graph import attributes, cutting, segmenthistory, ChunkedGraph from pychunkedgraph.graph import ( edges as cg_edges, ) @@ -26,11 +26,9 @@ exceptions as cg_exceptions, ) from pychunkedgraph.graph.analysis import pathing -from pychunkedgraph.graph.attributes import OperationLogs from pychunkedgraph.graph.misc import get_contact_sites from pychunkedgraph.graph.operation import GraphEditOperation -from pychunkedgraph.graph.utils import basetypes -from pychunkedgraph.meshing import mesh_analysis +from pychunkedgraph.graph import basetypes __api_versions__ = [0, 1] __segmentation_url_prefix__ = os.environ.get("SEGMENTATION_URL_PREFIX", "segmentation") @@ -106,11 +104,25 @@ def handle_info(table_id): combined_info["verify_mesh"] = cg.meta.custom_data.get("mesh", {}).get( "verify", False ) - mesh_dir = cg.meta.custom_data.get("mesh", {}).get("dir", None) + mesh_meta = cg.meta.custom_data.get("mesh", {}) + mesh_dir = mesh_meta.get("dir", None) if mesh_dir is not None: combined_info["mesh_dir"] = mesh_dir elif combined_info.get("mesh_dir", None) is not None: combined_info["mesh_dir"] = "graphene_meshes" + # `dynamic_mesh_dir` lets a dataset name the unsharded dynamic-mesh + # subdir explicitly. Default `"dynamic"` matches mesh_worker.py's + # fallback and NG's current hardcoded subdir name in graphene + # backend.ts (`${fragmentUrl}dynamic/`). NG must read + # this info field before non-default values route correctly. + dynamic_dir = mesh_meta.get("dynamic_mesh_dir", "dynamic") + combined_info["dynamic_mesh_dir"] = dynamic_dir + # cloud-volume reads the dynamic dir from mesh_metadata.unsharded_mesh_dir, not + # dynamic_mesh_dir; mirror it so an unpatched client fetches dynamic meshes from + # the right dir. Copy the dict so cg.meta.dataset_info is untouched. + mesh_metadata = dict(combined_info.get("mesh_metadata", {})) + mesh_metadata["unsharded_mesh_dir"] = dynamic_dir + combined_info["mesh_metadata"] = mesh_metadata return jsonify(combined_info) @@ -229,7 +241,9 @@ def handle_find_minimal_covering_nodes(table_id, is_binary=True): node_queue[layer].clear() # Return the download list - download_list = np.concatenate([np.array(list(v)) for v in download_list.values()]) + download_list = np.concatenate( + [np.array(list(v), dtype=np.uint64) for v in download_list.values()] + ) return download_list @@ -320,15 +334,13 @@ def publish_edit( is_priority=True, remesh: bool = True, ): - import pickle - - from messagingclient import MessagingClient - + downsample = bool(result.seg_bbox) attributes = { "table_id": table_id, "user_id": user_id, "remesh_priority": "true" if is_priority else "false", "remesh": "true" if remesh else "false", + "downsample": "true" if downsample else "false", } payload = { "operation_id": int(result.operation_id), @@ -336,6 +348,13 @@ def publish_edit( "new_root_ids": result.new_root_ids.tolist(), "old_root_ids": result.old_root_ids.tolist(), } + if downsample: + # Each entry is the base-resolution bbox of one supervoxel split's + # writes. Kept as a list (not merged) so the worker only rewrites + # tiles whose base footprint actually changed. + payload["seg_bboxes"] = [ + [bbs.tolist(), bbe.tolist()] for bbs, bbe in result.seg_bbox + ] exchange = os.getenv("PYCHUNKEDGRAPH_EDITS_EXCHANGE", "pychunkedgraph") c = MessagingClient() @@ -394,7 +413,7 @@ def handle_merge(table_id, allow_same_segment_merge=False): current_app.operation_id = ret.operation_id if ret.new_root_ids is None: raise cg_exceptions.InternalServerError( - "Could not merge selected " "supervoxel." + f"{ret.operation_id}: Could not merge selected supervoxels." ) current_app.logger.debug(("lvl2_nodes:", ret.new_lvl2_ids)) @@ -408,24 +427,9 @@ def handle_merge(table_id, allow_same_segment_merge=False): ### SPLIT ---------------------------------------------------------------------- -def handle_split(table_id): - current_app.table_id = table_id - user_id = str(g.auth_user.get("id", current_app.user_id)) - - data = json.loads(request.data) - is_priority = request.args.get("priority", True, type=str2bool) - remesh = request.args.get("remesh", True, type=str2bool) - mincut = request.args.get("mincut", True, type=str2bool) - - current_app.logger.debug(data) - - # Call ChunkedGraph - cg = app_utils.get_cg(table_id, skip_cache=True) +def _get_sources_and_sinks(cg: ChunkedGraph, data): node_idents = [] - node_ident_map = { - "sources": 0, - "sinks": 1, - } + node_ident_map = {"sources": 0, "sinks": 1} coords = [] node_ids = [] @@ -438,18 +442,35 @@ def handle_split(table_id): node_ids = np.array(node_ids, dtype=np.uint64) coords = np.array(coords) node_idents = np.array(node_idents) + sv_ids = app_utils.handle_supervoxel_id_lookup(cg, coords, node_ids) - current_app.logger.debug( - {"node_id": node_ids, "sv_id": sv_ids, "node_ident": node_idents} - ) + source_ids = sv_ids[node_idents == 0] + sink_ids = sv_ids[node_idents == 1] + source_coords = coords[node_idents == 0] + sink_coords = coords[node_idents == 1] + return (source_ids, sink_ids, source_coords, sink_coords) + +def handle_split(table_id): + current_app.table_id = table_id + user_id = str(g.auth_user.get("id", current_app.user_id)) + + data = json.loads(request.data) + is_priority = request.args.get("priority", True, type=str2bool) + remesh = request.args.get("remesh", True, type=str2bool) + mincut = request.args.get("mincut", True, type=str2bool) + + cg = app_utils.get_cg(table_id, skip_cache=True) + current_app.logger.debug(data) + sources, sinks, source_coords, sink_coords = _get_sources_and_sinks(cg, data) + logger.note(f"split inputs: sources={sources}, sinks={sinks}") try: ret = cg.remove_edges( user_id=user_id, - source_ids=sv_ids[node_idents == 0], - sink_ids=sv_ids[node_idents == 1], - source_coords=coords[node_idents == 0], - sink_coords=coords[node_idents == 1], + source_ids=sources, + sink_ids=sinks, + source_coords=source_coords, + sink_coords=sink_coords, mincut=mincut, ) except cg_exceptions.LockingError as e: @@ -460,7 +481,7 @@ def handle_split(table_id): current_app.operation_id = ret.operation_id if ret.new_root_ids is None: raise cg_exceptions.InternalServerError( - "Could not split selected segment groups." + f"{ret.operation_id}: Could not split selected segment groups." ) current_app.logger.debug(("after split:", ret.new_root_ids)) @@ -601,7 +622,9 @@ def all_user_operations( target_user_id = request.args.get("user_id", None) start_time = _parse_timestamp("start_time", 0, return_datetime=True) - end_time = _parse_timestamp("end_time", datetime.utcnow(), return_datetime=True) + end_time = _parse_timestamp( + "end_time", datetime.now(timezone.utc), return_datetime=True + ) # Call ChunkedGraph cg = app_utils.get_cg(table_id) @@ -611,23 +634,24 @@ def all_user_operations( valid_entry_ids = [] timestamp_list = [] - undone_ids = np.array([]) + undone_ids = np.array([], dtype=np.uint64) entry_ids = np.sort(list(log_rows.keys())) for entry_id in entry_ids: entry = log_rows[entry_id] - user_id = entry[OperationLogs.UserID] + user_id = entry[attributes.OperationLogs.UserID] should_check = ( - OperationLogs.Status not in entry - or entry[OperationLogs.Status] == OperationLogs.StatusCodes.SUCCESS.value + attributes.OperationLogs.Status not in entry + or entry[attributes.OperationLogs.Status] + == attributes.OperationLogs.StatusCodes.SUCCESS.value ) split_valid = ( include_partial_splits - or (OperationLogs.AddedEdge in entry) - or (OperationLogs.RootID not in entry) - or (len(entry[OperationLogs.RootID]) > 1) + or (attributes.OperationLogs.AddedEdge in entry) + or (attributes.OperationLogs.RootID not in entry) + or (len(entry[attributes.OperationLogs.RootID]) > 1) ) if not split_valid: print("excluding partial split", entry_id) @@ -641,13 +665,13 @@ def all_user_operations( if should_check: # if it is an undo of another operation, mark it as undone - if OperationLogs.UndoOperationID in entry: - undone_id = entry[OperationLogs.UndoOperationID] + if attributes.OperationLogs.UndoOperationID in entry: + undone_id = entry[attributes.OperationLogs.UndoOperationID] undone_ids = np.append(undone_ids, undone_id) # if it is a redo of another operation, unmark it as undone - if OperationLogs.RedoOperationID in entry: - redone_id = entry[OperationLogs.RedoOperationID] + if attributes.OperationLogs.RedoOperationID in entry: + redone_id = entry[attributes.OperationLogs.RedoOperationID] undone_ids = np.delete(undone_ids, np.argwhere(undone_ids == redone_id)) if include_undone: @@ -660,8 +684,8 @@ def all_user_operations( entry = log_rows[entry_id] if ( - OperationLogs.UndoOperationID in entry - or OperationLogs.RedoOperationID in entry + attributes.OperationLogs.UndoOperationID in entry + or attributes.OperationLogs.RedoOperationID in entry ): continue @@ -689,7 +713,7 @@ def handle_children(table_id, parent_id): if layer > 1: children = cg.get_children(parent_id) else: - children = np.array([]) + children = np.array([], dtype=np.uint64) return children @@ -792,8 +816,8 @@ def handle_subgraph(table_id, root_id, only_internal_edges=True): supervoxels = np.concatenate( [agg.supervoxels for agg in l2id_agglomeration_d.values()] ) - mask0 = np.in1d(edges.node_ids1, supervoxels) - mask1 = np.in1d(edges.node_ids2, supervoxels) + mask0 = np.isin(edges.node_ids1, supervoxels) + mask1 = np.isin(edges.node_ids2, supervoxels) edges = edges[mask0 & mask1] return edges @@ -1081,6 +1105,9 @@ def handle_split_preview(table_id): def handle_find_path(table_id, precision_mode): + # nested: pulls meshing/cloudvolume, only needed at call time + from pychunkedgraph.meshing import mesh_analysis + current_app.table_id = table_id user_id = str(g.auth_user.get("id", current_app.user_id)) diff --git a/pychunkedgraph/debug/cross_edge_test.py b/pychunkedgraph/debug/cross_edge_test.py deleted file mode 100644 index 25bacfa0b..000000000 --- a/pychunkedgraph/debug/cross_edge_test.py +++ /dev/null @@ -1,60 +0,0 @@ -import os -from datetime import datetime -import numpy as np - -from pychunkedgraph.graph import chunkedgraph -from pychunkedgraph.graph import attributes - -#os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/home/svenmd/.cloudvolume/secrets/google-secret.json" - -layer = 2 -n_chunks = 1000 -n_segments_per_chunk = 200 -# timestamp = datetime.datetime.fromtimestamp(1588875769) -timestamp = datetime.utcnow() - -cg = chunkedgraph.ChunkedGraph(graph_id="pinky_nf_v2") - -np.random.seed(42) - -node_ids = [] -for _ in range(n_chunks): - c_x = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][0]) - c_y = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][1]) - c_z = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][2]) - - chunk_id = cg.get_chunk_id(layer=layer, x=c_x, y=c_y, z=c_z) - - max_segment_id = cg.get_segment_id(cg.id_client.get_max_node_id(chunk_id)) - - if max_segment_id < 10: - continue - - segment_ids = np.random.randint(1, max_segment_id, n_segments_per_chunk) - - for segment_id in segment_ids: - node_ids.append(cg.get_node_id(np.uint64(segment_id), np.uint64(chunk_id))) - -rows = cg.client.read_nodes(node_ids=node_ids, end_time=timestamp, - properties=attributes.Hierarchy.Parent) -valid_node_ids = [] -non_valid_node_ids = [] -for k in rows.keys(): - if len(rows[k]) > 0: - valid_node_ids.append(k) - else: - non_valid_node_ids.append(k) - -cc_edges = cg.get_atomic_cross_edges(valid_node_ids) -cc_ids = np.unique(np.concatenate([np.concatenate(list(v.values())) for v in list(cc_edges.values()) if len(v.values())])) - -roots = cg.get_roots(cc_ids) -root_dict = dict(zip(cc_ids, roots)) -root_dict_vec = np.vectorize(root_dict.get) - -for k in cc_edges: - if len(cc_edges[k]) == 0: - continue - local_ids = np.unique(np.concatenate(list(cc_edges[k].values()))) - - assert len(np.unique(root_dict_vec(local_ids))) \ No newline at end of file diff --git a/pychunkedgraph/debug/existence_test.py b/pychunkedgraph/debug/existence_test.py deleted file mode 100644 index 757d3d542..000000000 --- a/pychunkedgraph/debug/existence_test.py +++ /dev/null @@ -1,78 +0,0 @@ -import os -from datetime import datetime -import numpy as np - -from pychunkedgraph.graph import chunkedgraph -from pychunkedgraph.graph import attributes - -#os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/home/svenmd/.cloudvolume/secrets/google-secret.json" - -layer = 2 -n_chunks = 100 -n_segments_per_chunk = 200 -# timestamp = datetime.datetime.fromtimestamp(1588875769) -timestamp = datetime.utcnow() - -cg = chunkedgraph.ChunkedGraph(graph_id="pinky_nf_v2") - -np.random.seed(42) - -node_ids = [] -for _ in range(n_chunks): - c_x = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][0]) - c_y = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][1]) - c_z = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][2]) - - chunk_id = cg.get_chunk_id(layer=layer, x=c_x, y=c_y, z=c_z) - - max_segment_id = cg.get_segment_id(cg.id_client.get_max_node_id(chunk_id)) - - if max_segment_id < 10: - continue - - segment_ids = np.random.randint(1, max_segment_id, n_segments_per_chunk) - - for segment_id in segment_ids: - node_ids.append(cg.get_node_id(np.uint64(segment_id), np.uint64(chunk_id))) - -rows = cg.client.read_nodes(node_ids=node_ids, end_time=timestamp, - properties=attributes.Hierarchy.Parent) -valid_node_ids = [] -non_valid_node_ids = [] -for k in rows.keys(): - if len(rows[k]) > 0: - valid_node_ids.append(k) - else: - non_valid_node_ids.append(k) - -roots = cg.get_roots(valid_node_ids, time_stamp=timestamp) - -roots = [] -try: - roots = cg.get_roots(valid_node_ids) - assert len(roots) == len(valid_node_ids) - print(f"ALL {len(roots)} have been successful!") -except: - print("At least one node failed. Checking nodes one by one now") - -if len(roots) != len(valid_node_ids): - log_dict = {} - success_dict = {} - for node_id in valid_node_ids: - try: - root = cg.get_root(node_id, time_stamp=timestamp) - print(f"Success: {node_id} from chunk {cg.get_chunk_id(node_id)}") - success_dict[node_id] = True - except Exception as e: - print(f"{node_id} from chunk {cg.get_chunk_id(node_id)} failed with {e}") - success_dict[node_id] = False - - t_id = node_id - - while t_id is not None: - last_working_chunk = cg.get_chunk_id(t_id) - t_id = cg.get_parent(t_id) - - print(f"Failed on layer {cg.get_chunk_layer(last_working_chunk)} in chunk {last_working_chunk}") - log_dict[node_id] = last_working_chunk - diff --git a/pychunkedgraph/debug/family_test.py b/pychunkedgraph/debug/family_test.py deleted file mode 100644 index 198351e74..000000000 --- a/pychunkedgraph/debug/family_test.py +++ /dev/null @@ -1,54 +0,0 @@ -import os -from datetime import datetime -import numpy as np - -from pychunkedgraph.graph import chunkedgraph -from pychunkedgraph.graph import attributes - -# os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/home/svenmd/.cloudvolume/secrets/google-secret.json" - -layers = [2, 3, 4, 5, 6, 7] -n_chunks = 10 -n_segments_per_chunk = 200 -# timestamp = datetime.datetime.fromtimestamp(1588875769) -timestamp = datetime.utcnow() - -cg = chunkedgraph.ChunkedGraph(graph_id="pinky_nf_v2") - -np.random.seed(42) - -node_ids = [] - -for layer in layers: - for _ in range(n_chunks): - c_x = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][0]) - c_y = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][1]) - c_z = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][2]) - - chunk_id = cg.get_chunk_id(layer=layer, x=c_x, y=c_y, z=c_z) - - max_segment_id = cg.get_segment_id(cg.id_client.get_max_node_id(chunk_id)) - - if max_segment_id < 10: - continue - - segment_ids = np.random.randint(1, max_segment_id, n_segments_per_chunk) - - for segment_id in segment_ids: - node_ids.append(cg.get_node_id(np.uint64(segment_id), np.uint64(chunk_id))) - -rows = cg.client.read_nodes(node_ids=node_ids, end_time=timestamp, - properties=attributes.Hierarchy.Parent) -valid_node_ids = [] -non_valid_node_ids = [] -for k in rows.keys(): - if len(rows[k]) > 0: - valid_node_ids.append(k) - else: - non_valid_node_ids.append(k) - -parents = cg.get_parents(valid_node_ids, time_stamp=timestamp) -children_dict = cg.get_children(parents) - -for child, parent in zip(valid_node_ids, parents): - assert child in children_dict[parent] \ No newline at end of file diff --git a/pychunkedgraph/debug/utils.py b/pychunkedgraph/debug/utils.py index 179f50aef..ad12103b2 100644 --- a/pychunkedgraph/debug/utils.py +++ b/pychunkedgraph/debug/utils.py @@ -1,7 +1,8 @@ +# pylint: disable=invalid-name, missing-docstring, bare-except, unidiomatic-typecheck + import numpy as np -from ..graph import ChunkedGraph -from ..graph.utils.basetypes import NODE_ID +from pychunkedgraph.graph.meta import ChunkedGraphMeta, GraphConfig def print_attrs(d): @@ -16,28 +17,59 @@ def print_attrs(d): print(v) -def print_node( - cg: ChunkedGraph, - node: NODE_ID, - indent: int = 0, - stop_layer: int = 2, -) -> None: +def print_node(cg, node: np.uint64, indent: int = 0, stop_layer: int = 2) -> None: children = cg.get_children(node) print(f"{' ' * indent}{node}[{len(children)}]") if cg.get_chunk_layer(node) <= stop_layer: return for child in children: - print_node(cg, child, indent=indent + 1, stop_layer=stop_layer) - - -def get_l2children(cg: ChunkedGraph, node: NODE_ID) -> np.ndarray: - nodes = np.array([node], dtype=NODE_ID) - layers = cg.get_chunk_layers(nodes) - assert np.all(layers > 2), "nodes must be at layers > 2" - l2children = [] - while nodes.size: - children = cg.get_children(nodes, flatten=True) - layers = cg.get_chunk_layers(children) - l2children.append(children[layers == 2]) - nodes = children[layers > 2] - return np.concatenate(l2children) + print_node(cg, child, indent=indent + 4, stop_layer=stop_layer) + + +def sanity_check(cg, new_roots, operation_id): + """ + Check for duplicates in hierarchy, useful for debugging. + """ + # print(f"{len(new_roots)} new ids from {operation_id}") + l2c_d = {} + for new_root in new_roots: + l2c_d[new_root] = cg.get_l2children([new_root]) + success = True + for k, v in l2c_d.items(): + success = success and (len(v) == np.unique(v).size) + # print(f"{k}: {np.unique(v).size}, {len(v)}") + if not success: + raise RuntimeError(f"{operation_id}: some ids are not valid.") + + +def sanity_check_single(cg, node, operation_id): + v = cg.get_l2children([node]) + msg = f"invalid node {node}:" + msg += f" found {len(v)} l2 ids, must be {np.unique(v).size}" + assert np.unique(v).size == len(v), f"{msg}, from {operation_id}." + return v + + +def update_graph_id(cg, new_graph_id:str): + old_gc = cg.meta.graph_config._asdict() + old_gc["ID"] = new_graph_id + new_gc = GraphConfig(**old_gc) + new_meta = ChunkedGraphMeta(new_gc, cg.meta.data_source, cg.meta.custom_data) + cg.update_meta(new_meta, overwrite=True) + + +def get_random_l1_ids(cg, n_chunks=100, n_per_chunk=10, seed=None): + """Generate random layer 1 IDs from different chunks.""" + if seed: + np.random.seed(seed) + bounds = cg.meta.layer_chunk_bounds[2] + ids = [] + for _ in range(n_chunks): + cx, cy, cz = [np.random.randint(0, b) for b in bounds] + chunk_id = cg.get_chunk_id(layer=2, x=cx, y=cy, z=cz) + max_seg = cg.get_segment_id(cg.id_client.get_max_node_id(chunk_id)) + if max_seg < 2: + continue + for seg in np.random.randint(1, max_seg + 1, n_per_chunk): + ids.append(cg.get_node_id(np.uint64(seg), np.uint64(chunk_id))) + return np.array(ids, dtype=np.uint64) diff --git a/pychunkedgraph/export/operation_logs.py b/pychunkedgraph/export/operation_logs.py index ec7141ce7..1ee22a5a1 100644 --- a/pychunkedgraph/export/operation_logs.py +++ b/pychunkedgraph/export/operation_logs.py @@ -2,9 +2,11 @@ from typing import Iterable from datetime import datetime +from kvdbclient import attributes, basetypes +from kvdbclient.attributes import OperationLogs + from .models import OperationLog from ..graph import ChunkedGraph -from ..graph.attributes import OperationLogs def parse_attr(attr, val) -> str: @@ -54,7 +56,8 @@ def get_logs_with_previous_roots( from numpy import concatenate from ..graph.types import empty_1d from ..graph.lineage import get_previous_root_ids - from ..graph.utils.basetypes import NODE_ID + + NODE_ID = basetypes.NODE_ID print(f"getting olg roots for {len(parsed_logs)} logs.") roots = [empty_1d] diff --git a/pychunkedgraph/graph/__init__.py b/pychunkedgraph/graph/__init__.py index 96b342427..2be4fa1d6 100644 --- a/pychunkedgraph/graph/__init__.py +++ b/pychunkedgraph/graph/__init__.py @@ -1,2 +1,19 @@ +import sys + +from kvdbclient import attributes +from kvdbclient import serializers +from kvdbclient import base as client_base +from kvdbclient import ( + BackendClientInfo, + ClientType, + get_client_class, + get_default_client_info, +) +from kvdbclient.utils import get_valid_timestamp, get_min_time, get_max_time + +# Register submodule aliases so `from pychunkedgraph.graph.attributes import X` works. +sys.modules[f"{__name__}.attributes"] = attributes +sys.modules[f"{__name__}.serializers"] = serializers + from .chunkedgraph import ChunkedGraph from .meta import ChunkedGraphMeta diff --git a/pychunkedgraph/graph/analysis/pathing.py b/pychunkedgraph/graph/analysis/pathing.py index 062b7a1c3..38715ba49 100644 --- a/pychunkedgraph/graph/analysis/pathing.py +++ b/pychunkedgraph/graph/analysis/pathing.py @@ -218,10 +218,11 @@ def compute_rough_coordinate_path(cg, l2_ids): coordinate_path = [] for l2_id in l2_ids: chunk_center = cg.get_chunk_coordinates(l2_id) + np.array([0.5, 0.5, 0.5]) - coordinate = chunk_center * np.array( - cg.meta.graph_config.CHUNK_SIZE - ) + np.array(cg.meta.cv.mip_voxel_offset(0)) - coordinate = coordinate * np.array(cg.meta.cv.mip_resolution(0)) + coordinate = ( + chunk_center * np.array(cg.meta.graph_config.CHUNK_SIZE) + + cg.meta.voxel_bounds[:, 0] + ) + coordinate = coordinate * cg.meta.resolution coordinate = coordinate.astype(np.float32) coordinate_path.append(coordinate) return coordinate_path diff --git a/pychunkedgraph/graph/attributes.py b/pychunkedgraph/graph/attributes.py deleted file mode 100644 index 3e48d204a..000000000 --- a/pychunkedgraph/graph/attributes.py +++ /dev/null @@ -1,284 +0,0 @@ -# TODO design to use these attributes across different clients -# `family_id` is specific to bigtable - -from typing import NamedTuple - -from .utils import serializers -from .utils import basetypes - - -class _AttributeType(NamedTuple): - key: bytes - family_id: str - serializer: serializers._Serializer - - -class _Attribute(_AttributeType): - __slots__ = () - _attributes = {} - - def __init__(self, **kwargs): - super().__init__() - _Attribute._attributes[(kwargs["family_id"], kwargs["key"])] = self - - def serialize(self, obj): - return self.serializer.serialize(obj) - - def deserialize(self, stream): - return self.serializer.deserialize(stream) - - @property - def basetype(self): - return self.serializer.basetype - - @property - def index(self): - return int(self.key.decode("utf-8").split("_")[-1]) - - -class _AttributeArray: - _attributearrays = {} - - def __init__(self, pattern, family_id, serializer): - self._pattern = pattern - self._family_id = family_id - self._serializer = serializer - _AttributeArray._attributearrays[(family_id, pattern)] = self - - # TODO: Add missing check in `fromkey(family_id, key)` and remove this - # loop (pre-creates `_Attributes`, so that the inverse lookup works) - for i in range(20): - self[i] # pylint: disable=W0104 - - def __getitem__(self, item): - return _Attribute( - key=self.pattern % item, - family_id=self._family_id, - serializer=self._serializer, - ) - - @property - def pattern(self): - return self._pattern - - @property - def serialize(self): - return self._serializer.serialize - - @property - def deserialize(self): - return self._serializer.deserialize - - @property - def basetype(self): - return self._serializer.basetype - - -class Concurrency: - Counter = _Attribute( - key=b"counter", - family_id="1", - serializer=serializers.NumPyValue(dtype=basetypes.COUNTER), - ) - - Lock = _Attribute(key=b"lock", family_id="0", serializer=serializers.UInt64String()) - - IndefiniteLock = _Attribute( - key=b"indefinite_lock", family_id="0", serializer=serializers.UInt64String() - ) - - -class Connectivity: - Affinity = _Attribute( - key=b"affinities", - family_id="0", - serializer=serializers.NumPyArray(dtype=basetypes.EDGE_AFFINITY), - ) - - Area = _Attribute( - key=b"areas", - family_id="0", - serializer=serializers.NumPyArray(dtype=basetypes.EDGE_AREA), - ) - - CrossChunkEdge = _AttributeArray( - pattern=b"atomic_cross_edges_%d", - family_id="3", - serializer=serializers.NumPyArray( - dtype=basetypes.NODE_ID, shape=(-1, 2), compression_level=22 - ), - ) - - FakeEdges = _Attribute( - key=b"fake_edges", - family_id="3", - serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID, shape=(-1, 2)), - ) - - -class Hierarchy: - Child = _Attribute( - key=b"children", - family_id="0", - serializer=serializers.NumPyArray( - dtype=basetypes.NODE_ID, compression_level=22 - ), - ) - - FormerParent = _Attribute( - key=b"former_parents", - family_id="0", - serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID), - ) - - NewParent = _Attribute( - key=b"new_parents", - family_id="0", - serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID), - ) - - Parent = _Attribute( - key=b"parents", - family_id="0", - serializer=serializers.NumPyValue(dtype=basetypes.NODE_ID), - ) - - -class GraphMeta: - key = b"meta" - Meta = _Attribute(key=key, family_id="0", serializer=serializers.Pickle()) - - -class GraphVersion: - key = b"version" - Version = _Attribute(key=key, family_id="0", serializer=serializers.String("utf-8")) - - -class OperationLogs: - key = b"ioperations" - - from enum import Enum - - class StatusCodes(Enum): - SUCCESS = 0 # all is well, new changes persisted - CREATED = 1 # log record created in storage - EXCEPTION = 2 # edit unsuccessful, unknown error - WRITE_STARTED = 3 # edit successful, start persisting changes - WRITE_FAILED = 4 # edit successful, but changes not persisted - - OperationID = _Attribute( - key=b"operation_id", family_id="0", serializer=serializers.UInt64String() - ) - - UndoOperationID = _Attribute( - key=b"undo_operation_id", family_id="2", serializer=serializers.UInt64String() - ) - - RedoOperationID = _Attribute( - key=b"redo_operation_id", family_id="2", serializer=serializers.UInt64String() - ) - - UserID = _Attribute( - key=b"user", family_id="2", serializer=serializers.String("utf-8") - ) - - RootID = _Attribute( - key=b"roots", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID), - ) - - SourceID = _Attribute( - key=b"source_ids", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID), - ) - - SinkID = _Attribute( - key=b"sink_ids", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID), - ) - - SourceCoordinate = _Attribute( - key=b"source_coords", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.COORDINATES, shape=(-1, 3)), - ) - - SinkCoordinate = _Attribute( - key=b"sink_coords", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.COORDINATES, shape=(-1, 3)), - ) - - BoundingBoxOffset = _Attribute( - key=b"bb_offset", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.COORDINATES), - ) - - AddedEdge = _Attribute( - key=b"added_edges", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID, shape=(-1, 2)), - ) - - RemovedEdge = _Attribute( - key=b"removed_edges", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID, shape=(-1, 2)), - ) - - Affinity = _Attribute( - key=b"affinities", - family_id="2", - serializer=serializers.NumPyArray(dtype=basetypes.EDGE_AFFINITY), - ) - - Status = _Attribute( - key=b"operation_status", family_id="0", serializer=serializers.Pickle() - ) - - OperationException = _Attribute( - key=b"operation_exception", - family_id="0", - serializer=serializers.String("utf-8"), - ) - - # timestamp at which the new IDs were created during the operation - # this is needed because the timestamp of the operation log - # will change with change in status - OperationTimeStamp = _Attribute( - key=b"operation_ts", family_id="0", serializer=serializers.Pickle() - ) - - @staticmethod - def all(): - return [ - OperationLogs.OperationID, - OperationLogs.UndoOperationID, - OperationLogs.RedoOperationID, - OperationLogs.UserID, - OperationLogs.RootID, - OperationLogs.SourceID, - OperationLogs.SinkID, - OperationLogs.SourceCoordinate, - OperationLogs.SinkCoordinate, - OperationLogs.BoundingBoxOffset, - OperationLogs.AddedEdge, - OperationLogs.RemovedEdge, - OperationLogs.Affinity, - OperationLogs.Status, - OperationLogs.OperationException, - OperationLogs.OperationTimeStamp, - ] - - -def from_key(family_id: str, key: bytes): - try: - return _Attribute._attributes[(family_id, key)] - except KeyError: - # FIXME: Look if the key matches a columnarray pattern and - # remove loop initialization in _AttributeArray.__init__() - raise KeyError(f"Unknown key {family_id}:{key.decode()}") diff --git a/pychunkedgraph/graph/basetypes.py b/pychunkedgraph/graph/basetypes.py new file mode 100644 index 000000000..ff7963363 --- /dev/null +++ b/pychunkedgraph/graph/basetypes.py @@ -0,0 +1 @@ +from kvdbclient.basetypes import * # noqa: F401,F403 diff --git a/pychunkedgraph/graph/cache.py b/pychunkedgraph/graph/cache.py index f60b6ca92..430a998c5 100644 --- a/pychunkedgraph/graph/cache.py +++ b/pychunkedgraph/graph/cache.py @@ -1,6 +1,10 @@ +# pylint: disable=invalid-name, missing-docstring, import-outside-toplevel """ Cache nodes, parents, children and cross edges. """ + +import traceback +from collections import defaultdict as defaultd from sys import maxsize from datetime import datetime @@ -10,7 +14,7 @@ import numpy as np -from .utils.basetypes import NODE_ID +from pychunkedgraph.graph import basetypes def update(cache, keys, vals): @@ -30,28 +34,84 @@ def __init__(self, cg): self._parent_vec = np.vectorize(self.parent, otypes=[np.uint64]) self._children_vec = np.vectorize(self.children, otypes=[np.ndarray]) - self._atomic_cross_edges_vec = np.vectorize( - self.atomic_cross_edges, otypes=[dict] + self._cross_chunk_edges_vec = np.vectorize( + self.cross_chunk_edges, otypes=[dict] ) # no limit because we don't want to lose new IDs self.parents_cache = LRUCache(maxsize=maxsize) self.children_cache = LRUCache(maxsize=maxsize) - self.atomic_cx_edges_cache = LRUCache(maxsize=maxsize) + self.cross_chunk_edges_cache = LRUCache(maxsize=maxsize) + + self.new_ids = set() + + # Stats tracking for cache hits/misses + self.stats = { + "parents": {"hits": 0, "misses": 0, "calls": 0}, + "children": {"hits": 0, "misses": 0, "calls": 0}, + "cross_chunk_edges": {"hits": 0, "misses": 0, "calls": 0}, + } + # Track where calls/misses come from + self.sources = defaultd(lambda: defaultd(lambda: {"calls": 0, "misses": 0})) + + def _get_caller(self, skip_frames=2): + """Get caller info (filename:line:function).""" + stack = traceback.extract_stack() + # Skip frames: _get_caller, the cache method, and go to actual caller + if len(stack) > skip_frames: + frame = stack[-(skip_frames + 1)] + return f"{frame.filename.split('/')[-1]}:{frame.lineno}:{frame.name}" + return "unknown" + + def _record_call(self, cache_type, misses=0): + """Record a call and its source.""" + caller = self._get_caller(skip_frames=3) + self.sources[cache_type][caller]["calls"] += 1 + self.sources[cache_type][caller]["misses"] += misses def __len__(self): return ( len(self.parents_cache) + len(self.children_cache) - + len(self.atomic_cx_edges_cache) + + len(self.cross_chunk_edges_cache) ) def clear(self): self.parents_cache.clear() self.children_cache.clear() - self.atomic_cx_edges_cache.clear() + self.cross_chunk_edges_cache.clear() + + def get_stats(self): + """Return stats with hit rates calculated.""" + result = {} + for name, s in self.stats.items(): + total = s["hits"] + s["misses"] + hit_rate = s["hits"] / total if total > 0 else 0 + result[name] = { + **s, + "total": total, + "hit_rate": f"{hit_rate:.1%}", + "sources": dict(self.sources[name]), + } + return result + + def reset_stats(self): + for s in self.stats.values(): + s["hits"] = 0 + s["misses"] = 0 + s["calls"] = 0 + self.sources.clear() def parent(self, node_id: np.uint64, *, time_stamp: datetime = None): + self.stats["parents"]["calls"] += 1 + is_cached = node_id in self.parents_cache + miss_count = 0 if is_cached else 1 + if is_cached: + self.stats["parents"]["hits"] += 1 + else: + self.stats["parents"]["misses"] += 1 + self._record_call("parents", misses=miss_count) + @cached(cache=self.parents_cache, key=lambda node_id: node_id) def parent_decorated(node_id): return self._cg.get_parent(node_id, raw_only=True, time_stamp=time_stamp) @@ -59,6 +119,15 @@ def parent_decorated(node_id): return parent_decorated(node_id) def children(self, node_id): + self.stats["children"]["calls"] += 1 + is_cached = node_id in self.children_cache + miss_count = 0 if is_cached else 1 + if is_cached: + self.stats["children"]["hits"] += 1 + else: + self.stats["children"]["misses"] += 1 + self._record_call("children", misses=miss_count) + @cached(cache=self.children_cache, key=lambda node_id: node_id) def children_decorated(node_id): children = self._cg.get_children(node_id, raw_only=True) @@ -67,33 +136,72 @@ def children_decorated(node_id): return children_decorated(node_id) - def atomic_cross_edges(self, node_id): - @cached(cache=self.atomic_cx_edges_cache, key=lambda node_id: node_id) - def atomic_cross_edges_decorated(node_id): - edges = self._cg.get_atomic_cross_edges( - np.array([node_id], dtype=NODE_ID), raw_only=True + def cross_chunk_edges(self, node_id, *, time_stamp: datetime = None): + self.stats["cross_chunk_edges"]["calls"] += 1 + is_cached = node_id in self.cross_chunk_edges_cache + miss_count = 0 if is_cached else 1 + if is_cached: + self.stats["cross_chunk_edges"]["hits"] += 1 + else: + self.stats["cross_chunk_edges"]["misses"] += 1 + self._record_call("cross_chunk_edges", misses=miss_count) + + @cached(cache=self.cross_chunk_edges_cache, key=lambda node_id: node_id) + def cross_edges_decorated(node_id): + edges = self._cg.get_cross_chunk_edges( + np.array([node_id], dtype=basetypes.NODE_ID), + raw_only=True, + time_stamp=time_stamp, ) return edges[node_id] - return atomic_cross_edges_decorated(node_id) + return cross_edges_decorated(node_id) - def parents_multiple(self, node_ids: np.ndarray, *, time_stamp: datetime = None): + def parents_multiple( + self, + node_ids: np.ndarray, + *, + time_stamp: datetime = None, + fail_to_zero: bool = False, + ): + node_ids = np.asarray(node_ids, dtype=basetypes.NODE_ID) if not node_ids.size: return node_ids - mask = np.in1d(node_ids, np.fromiter(self.parents_cache.keys(), dtype=NODE_ID)) + self.stats["parents"]["calls"] += 1 + mask = np.isin( + node_ids, np.fromiter(self.parents_cache.keys(), dtype=basetypes.NODE_ID) + ) + hits = int(np.sum(mask)) + misses = len(node_ids) - hits + self.stats["parents"]["hits"] += hits + self.stats["parents"]["misses"] += misses + self._record_call("parents", misses=misses) parents = node_ids.copy() parents[mask] = self._parent_vec(node_ids[mask]) parents[~mask] = self._cg.get_parents( - node_ids[~mask], raw_only=True, time_stamp=time_stamp + node_ids[~mask], + raw_only=True, + time_stamp=time_stamp, + fail_to_zero=fail_to_zero, ) + mask = mask | (parents == 0) update(self.parents_cache, node_ids[~mask], parents[~mask]) return parents def children_multiple(self, node_ids: np.ndarray, *, flatten=False): result = {} + node_ids = np.asarray(node_ids, dtype=basetypes.NODE_ID) if not node_ids.size: return result - mask = np.in1d(node_ids, np.fromiter(self.children_cache.keys(), dtype=NODE_ID)) + self.stats["children"]["calls"] += 1 + mask = np.isin( + node_ids, np.fromiter(self.children_cache.keys(), dtype=basetypes.NODE_ID) + ) + hits = int(np.sum(mask)) + misses = len(node_ids) - hits + self.stats["children"]["hits"] += hits + self.stats["children"]["misses"] += misses + self._record_call("children", misses=misses) cached_children_ = self._children_vec(node_ids[mask]) result.update({id_: c_ for id_, c_ in zip(node_ids[mask], cached_children_)}) result.update(self._cg.get_children(node_ids[~mask], raw_only=True)) @@ -104,20 +212,34 @@ def children_multiple(self, node_ids: np.ndarray, *, flatten=False): return np.concatenate([*result.values()]) return result - def atomic_cross_edges_multiple(self, node_ids: np.ndarray): + def cross_chunk_edges_multiple( + self, node_ids: np.ndarray, *, time_stamp: datetime = None + ): result = {} + node_ids = np.asarray(node_ids, dtype=basetypes.NODE_ID) if not node_ids.size: return result - mask = np.in1d( - node_ids, np.fromiter(self.atomic_cx_edges_cache.keys(), dtype=NODE_ID) + self.stats["cross_chunk_edges"]["calls"] += 1 + mask = np.isin( + node_ids, + np.fromiter(self.cross_chunk_edges_cache.keys(), dtype=basetypes.NODE_ID), ) - cached_edges_ = self._atomic_cross_edges_vec(node_ids[mask]) + hits = int(np.sum(mask)) + misses = len(node_ids) - hits + self.stats["cross_chunk_edges"]["hits"] += hits + self.stats["cross_chunk_edges"]["misses"] += misses + self._record_call("cross_chunk_edges", misses=misses) + cached_edges_ = self._cross_chunk_edges_vec(node_ids[mask]) result.update( {id_: edges_ for id_, edges_ in zip(node_ids[mask], cached_edges_)} ) - result.update(self._cg.get_atomic_cross_edges(node_ids[~mask], raw_only=True)) + result.update( + self._cg.get_cross_chunk_edges( + node_ids[~mask], raw_only=True, time_stamp=time_stamp + ) + ) update( - self.atomic_cx_edges_cache, + self.cross_chunk_edges_cache, node_ids[~mask], [result[k] for k in node_ids[~mask]], ) diff --git a/pychunkedgraph/graph/chunkedgraph.py b/pychunkedgraph/graph/chunkedgraph.py index 210bff50b..02bac55c9 100644 --- a/pychunkedgraph/graph/chunkedgraph.py +++ b/pychunkedgraph/graph/chunkedgraph.py @@ -1,29 +1,37 @@ -# pylint: disable=invalid-name, missing-docstring, too-many-lines, import-outside-toplevel +# pylint: disable=invalid-name, missing-docstring, too-many-lines, import-outside-toplevel, unsupported-binary-operation import time import typing import datetime +from itertools import chain +from functools import reduce import numpy as np from pychunkedgraph import __version__ from . import types from . import operation -from . import attributes +from pychunkedgraph.graph import attributes from . import exceptions -from .client import base -from .client import BigTableClient -from .client import BackendClientInfo -from .client import get_default_client_info +from pychunkedgraph.graph import client_base as base +from pychunkedgraph.graph import BackendClientInfo +from pychunkedgraph.graph import ClientType +from pychunkedgraph.graph import get_client_class +from pychunkedgraph.graph import get_default_client_info from .cache import CacheService from .meta import ChunkedGraphMeta -from .utils import basetypes +from pychunkedgraph.graph import basetypes +from .sv_lookup import utils as sv_lookup_utils from .utils import id_helpers +from pychunkedgraph.graph import serializers +from pychunkedgraph.graph import get_valid_timestamp from .utils import generic as misc_utils from .edges import Edges from .edges import utils as edge_utils from .chunks import utils as chunk_utils from .chunks import hierarchy as chunk_hierarchy +from .subgraph import get_subgraph_nodes +from .subgraph import get_subgraph_edges_and_leaves class ChunkedGraph: @@ -34,33 +42,36 @@ def __init__( meta: ChunkedGraphMeta = None, client_info: BackendClientInfo = get_default_client_info(), ): - """ - 1. New graph - Requires `meta`; if `client_info` is not passed the default client is used. - After creating `ChunkedGraph` instance, run instance.create(). - 2. Existing graph in default client - Requires `graph_id`. - 3. Existing graphs in other projects/clients, - Requires `graph_id` and `client_info`. - """ - # create client based on type - # for now, just use BigTableClient + """Open a chunked graph: `meta` for a new graph (then `.create()`), else `graph_id` + (+ `client_info` for other projects/clients). A graph_id naming a table copied from + another graph gets its graph-id-bearing meta (id, mesh dirs) rewritten here.""" + ClientClass = get_client_class(client_info.TYPE) if meta: graph_id = meta.graph_config.ID_PREFIX + meta.graph_config.ID - bt_client = BigTableClient( - graph_id, config=client_info.CONFIG, graph_meta=meta + _client = ClientClass( + graph_id, + config=client_info.CONFIG, + table_meta=meta, + lock_expiry=meta.graph_config.ROOT_LOCK_EXPIRY, ) self._meta = meta else: - bt_client = BigTableClient(graph_id, config=client_info.CONFIG) - self._meta = bt_client.read_graph_meta() + _client = ClientClass(graph_id, config=client_info.CONFIG) + self._meta = _client.read_table_meta() + _client._lock_expiry = self._meta.graph_config.ROOT_LOCK_EXPIRY - self._client = bt_client - self._id_client = bt_client + self._client = _client + self._id_client = _client self._cache_service = None self.mock_edges = None # hack for unit tests + # A copied/restored table carries the source's graph-id-bearing meta; + # on first access under a new id, rewrite + persist it once (later + # instantiations match this id and no-op). + if graph_id != self.graph_id: + self.update_meta(self.meta.for_copied_graph(graph_id), overwrite=True) + @property def meta(self) -> ChunkedGraphMeta: return self._meta @@ -71,10 +82,10 @@ def graph_id(self) -> str: @property def version(self) -> str: - return self.client.read_graph_version() + return self.client.read_table_version() @property - def client(self) -> base.SimpleClient: + def client(self) -> ClientType: return self._client @property @@ -87,7 +98,7 @@ def cache(self): @property def segmentation_resolution(self) -> np.ndarray: - return np.array(self.meta.ws_cv.scale["resolution"]) + return self.meta.resolution @cache.setter def cache(self, cache_service: CacheService): @@ -95,11 +106,11 @@ def cache(self, cache_service: CacheService): def create(self): """Creates the graph in storage client and stores meta.""" - self._client.create_graph(self._meta, version=__version__) + self._client.create_table(self._meta, version=__version__) def update_meta(self, meta: ChunkedGraphMeta, overwrite: bool): """Update meta of an already existing graph.""" - self.client.update_graph_meta(meta, overwrite=overwrite) + self.client.update_table_meta(meta, overwrite=overwrite) def range_read_chunk( self, @@ -112,13 +123,15 @@ def range_read_chunk( """Read all nodes in a chunk.""" layer = self.get_chunk_layer(chunk_id) root_chunk = layer == self.meta.layer_count - max_node_id = self.id_client.get_max_node_id(chunk_id=chunk_id, root_chunk=root_chunk) + max_id = self.id_client.get_max_node_id( + chunk_id=chunk_id, root_chunk=root_chunk + ) if layer == 1: - max_node_id = chunk_id | self.get_segment_id_limit(chunk_id) # pylint: disable=unsupported-binary-operation + max_id = chunk_id | self.get_segment_id_limit(chunk_id) return self.client.read_nodes( start_id=self.get_node_id(np.uint64(0), chunk_id=chunk_id), - end_id=max_node_id, + end_id=max_id, end_id_inclusive=True, properties=properties, end_time=time_stamp, @@ -137,7 +150,7 @@ def get_atomic_id_from_coord( """Determines atomic id given a coordinate.""" if self.get_chunk_layer(parent_id) == 1: return parent_id - return id_helpers.get_atomic_id_from_coord( + return sv_lookup_utils.get_atomic_id_from_coord( self.meta, self.get_root, x, @@ -161,16 +174,21 @@ def get_atomic_ids_from_coords( :param max_dist_nm: max distance explored :return: supervoxel ids; returns None if no solution was found """ - if self.get_chunk_layer(parent_id) == 1: + if self.get_chunk_layer(parent_id) == 1 and not self.meta.ocdbt_seg: return np.array([parent_id] * len(coordinates), dtype=np.uint64) - # Enable search with old parent by using its timestamp and map to parents - parent_ts = self.get_node_timestamps([parent_id], return_numpy=False)[0] - return id_helpers.get_atomic_ids_from_coords( + layer = self.get_chunk_layer(parent_id) + # L1 nodes don't have children, skip timestamp lookup + parent_ts = ( + None + if layer == 1 + else self.get_node_timestamps([parent_id], return_numpy=False)[0] + ) + return sv_lookup_utils.get_atomic_ids_from_coords( self.meta, coordinates, parent_id, - self.get_chunk_layer(parent_id), + layer, parent_ts, self.get_roots, max_dist_nm, @@ -190,14 +208,14 @@ def get_parents( Else all parents along with timestamps. """ if raw_only or not self.cache: - time_stamp = misc_utils.get_valid_timestamp(time_stamp) + time_stamp = get_valid_timestamp(time_stamp) parent_rows = self.client.read_nodes( node_ids=node_ids, properties=attributes.Hierarchy.Parent, end_time=time_stamp, end_time_inclusive=True, ) - if not parent_rows: + if not parent_rows and not fail_to_zero: return types.empty_1d parents = [] @@ -209,6 +227,7 @@ def get_parents( if fail_to_zero: parents.append(0) else: + exc.add_note(f"timestamp: {time_stamp}") raise KeyError from exc parents = np.array(parents, dtype=basetypes.NODE_ID) else: @@ -223,7 +242,9 @@ def get_parents( else: raise KeyError from exc return parents - return self.cache.parents_multiple(node_ids, time_stamp=time_stamp) + return self.cache.parents_multiple( + node_ids, time_stamp=time_stamp, fail_to_zero=fail_to_zero + ) def get_parent( self, @@ -234,7 +255,7 @@ def get_parent( time_stamp: typing.Optional[datetime.datetime] = None, ) -> typing.Union[typing.List[typing.Tuple], np.uint64]: if raw_only or not self.cache: - time_stamp = misc_utils.get_valid_timestamp(time_stamp) + time_stamp = get_valid_timestamp(time_stamp) parents = self.client.read_node( node_id, properties=attributes.Hierarchy.Parent, @@ -283,97 +304,79 @@ def _get_children_multiple( node_ids=node_ids, properties=attributes.Hierarchy.Child ) return { - x: node_children_d[x][0].value - if x in node_children_d - else types.empty_1d.copy() + x: ( + node_children_d[x][0].value + if x in node_children_d + else types.empty_1d.copy() + ) for x in node_ids } return self.cache.children_multiple(node_ids) - def get_atomic_cross_edges( - self, l2_ids: typing.Iterable, *, raw_only=False - ) -> typing.Dict[np.uint64, typing.Dict[int, typing.Iterable]]: - """Returns cross edges for level 2 IDs.""" + def get_atomic_cross_edges(self, l2_ids: typing.Iterable) -> typing.Dict: + """ + Returns atomic cross edges for level 2 IDs. + A dict of the form `{l2id: {layer: atomic_cross_edges}}`. + """ + node_edges_d_d = self.client.read_nodes( + node_ids=l2_ids, + properties=[ + attributes.Connectivity.AtomicCrossChunkEdge[l] + for l in range(2, max(3, self.meta.layer_count)) + ], + ) + result = {} + for id_ in l2_ids: + try: + result[id_] = { + prop.index: val[0].value.copy() + for prop, val in node_edges_d_d[id_].items() + } + except KeyError: + result[id_] = {} + return result + + def get_cross_chunk_edges( + self, + node_ids: typing.Iterable, + *, + raw_only=False, + all_layers=True, + time_stamp: typing.Optional[datetime.datetime] = None, + ) -> typing.Dict: + """ + Returns cross edges for `node_ids`. + A dict of the form `{node_id: {layer: cross_edges}}`. + """ + time_stamp = get_valid_timestamp(time_stamp) if raw_only or not self.cache: + result = {} + node_ids = np.array(node_ids, dtype=basetypes.NODE_ID) + if node_ids.size == 0: + return result + layers = range(2, max(3, self.meta.layer_count)) + attrs = [attributes.Connectivity.CrossChunkEdge[l] for l in layers] node_edges_d_d = self.client.read_nodes( - node_ids=l2_ids, - properties=[ - attributes.Connectivity.CrossChunkEdge[l] - for l in range(2, max(3, self.meta.layer_count)) - ], + node_ids=node_ids, + properties=attrs, + end_time=time_stamp, + end_time_inclusive=True, ) - result = {} - for id_ in l2_ids: + layers = self.get_chunk_layers(node_ids) + valid_layer = lambda x, y: x >= y + if not all_layers: + valid_layer = lambda x, y: x == y + for layer, id_ in zip(layers, node_ids): try: result[id_] = { prop.index: val[0].value.copy() for prop, val in node_edges_d_d[id_].items() + if valid_layer(prop.index, layer) } except KeyError: result[id_] = {} return result - return self.cache.atomic_cross_edges_multiple(l2_ids) - - def get_cross_chunk_edges( - self, node_ids: typing.Iterable, uplift=True, all_layers=False - ) -> typing.Dict[np.uint64, typing.Dict[int, typing.Iterable]]: - """ - Cross chunk edges for `node_id` at `node_layer`. - The edges are between node IDs at the `node_layer`, not atomic cross edges. - Returns dict {layer_id: cross_edges} - The first layer (>= `node_layer`) with atleast one cross chunk edge. - For current use-cases, other layers are not relevant. - - For performance, only children that lie along chunk boundary are considered. - Cross edges that belong to inner level 2 IDs are subsumed within the chunk. - This is because cross edges are stored only in level 2 IDs. - """ - result = {} - node_ids = np.array(node_ids, dtype=basetypes.NODE_ID) - if not node_ids.size: - return result - - node_l2ids_d = {} - layers_ = self.get_chunk_layers(node_ids) - for l in set(layers_): - node_l2ids_d.update(self._get_bounding_l2_children(node_ids[layers_ == l])) - l2_edges_d_d = self.get_atomic_cross_edges( - np.concatenate(list(node_l2ids_d.values())) - ) - for node_id in node_ids: - l2_edges_ds = [l2_edges_d_d[l2_id] for l2_id in node_l2ids_d[node_id]] - if all_layers: - result[node_id] = edge_utils.concatenate_cross_edge_dicts(l2_edges_ds) - else: - result[node_id] = self._get_min_layer_cross_edges( - node_id, l2_edges_ds, uplift=uplift - ) - return result - - def _get_min_layer_cross_edges( - self, - node_id: basetypes.NODE_ID, - l2id_atomic_cross_edges_ds: typing.Iterable, - uplift=True, - ) -> typing.Dict[int, typing.Iterable]: - """ - Find edges at relevant min_layer >= node_layer. - `l2id_atomic_cross_edges_ds` is a list of atomic cross edges of - level 2 IDs that are descendants of `node_id`. - """ - min_layer, edges = edge_utils.filter_min_layer_cross_edges_multiple( - self.meta, l2id_atomic_cross_edges_ds, self.get_chunk_layer(node_id) - ) - if self.get_chunk_layer(node_id) < min_layer: - # cross edges irrelevant - return {self.get_chunk_layer(node_id): types.empty_2d} - if not uplift: - return {min_layer: edges} - node_root_id = node_id - node_root_id = self.get_root(node_id, stop_layer=min_layer, ceil=False) - edges[:, 0] = node_root_id - edges[:, 1] = self.get_roots(edges[:, 1], stop_layer=min_layer, ceil=False) - return {min_layer: np.unique(edges, axis=0) if edges.size else types.empty_2d} + return self.cache.cross_chunk_edges_multiple(node_ids, time_stamp=time_stamp) def get_roots( self, @@ -384,6 +387,7 @@ def get_roots( stop_layer: int = None, ceil: bool = True, fail_to_zero: bool = False, + raw_only=False, n_tries: int = 1, ) -> typing.Union[np.ndarray, typing.Dict[int, np.ndarray]]: """ @@ -392,7 +396,7 @@ def get_roots( When `assert_roots=False`, returns highest available IDs and cases where there are no root IDs are silently ignored. """ - time_stamp = misc_utils.get_valid_timestamp(time_stamp) + time_stamp = get_valid_timestamp(time_stamp) stop_layer = self.meta.layer_count if not stop_layer else stop_layer assert stop_layer <= self.meta.layer_count layer_mask = np.ones(len(node_ids), dtype=bool) @@ -407,7 +411,10 @@ def get_roots( filtered_ids = parent_ids[layer_mask] unique_ids, inverse = np.unique(filtered_ids, return_inverse=True) temp_ids = self.get_parents( - unique_ids, time_stamp=time_stamp, fail_to_zero=fail_to_zero + unique_ids, + time_stamp=time_stamp, + fail_to_zero=fail_to_zero, + raw_only=raw_only, ) if not temp_ids.size: break @@ -462,10 +469,11 @@ def get_root( get_all_parents: bool = False, stop_layer: int = None, ceil: bool = True, + raw_only: bool = False, n_tries: int = 1, ) -> typing.Union[typing.List[np.uint64], np.uint64]: """Takes a node id and returns the associated agglomeration ids.""" - time_stamp = misc_utils.get_valid_timestamp(time_stamp) + time_stamp = get_valid_timestamp(time_stamp) parent_id = node_id all_parent_ids = [] stop_layer = self.meta.layer_count if not stop_layer else stop_layer @@ -479,7 +487,9 @@ def get_root( for _ in range(n_tries): parent_id = node_id for _ in range(self.get_chunk_layer(node_id), int(stop_layer + 1)): - temp_parent_id = self.get_parent(parent_id, time_stamp=time_stamp) + temp_parent_id = self.get_parent( + parent_id, time_stamp=time_stamp, raw_only=raw_only + ) if temp_parent_id is None: break else: @@ -499,7 +509,7 @@ def get_root( else: time.sleep(0.5) - if self.get_chunk_layer(parent_id) < stop_layer: + if ceil and self.get_chunk_layer(parent_id) < stop_layer: raise exceptions.ChunkedGraphError( f"Cannot find root id {node_id}, {stop_layer}, {time_stamp}" ) @@ -518,7 +528,7 @@ def is_latest_roots( time_stamp: typing.Optional[datetime.datetime] = None, ) -> typing.Iterable: """Determines whether root ids are superseded.""" - time_stamp = misc_utils.get_valid_timestamp(time_stamp) + time_stamp = get_valid_timestamp(time_stamp) row_dict = self.client.read_nodes( node_ids=root_ids, @@ -546,22 +556,52 @@ def get_all_parents_dict( ) return dict(zip(self.get_chunk_layers(parent_ids), parent_ids)) + def get_all_parents_dict_multiple(self, node_ids, *, time_stamp=None): + """Batch fetch all parent hierarchies layer by layer.""" + result = {node: {} for node in node_ids} + nodes = np.array(node_ids, dtype=basetypes.NODE_ID) + layers_map = {} + child_parent_map = {} + + while nodes.size > 0: + parents = self.get_parents(nodes, time_stamp=time_stamp) + parent_layers = self.get_chunk_layers(parents) + for node, parent, layer in zip(nodes, parents, parent_layers): + layers_map[parent] = layer + child_parent_map[node] = parent + nodes = parents[parent_layers < self.meta.layer_count] + + for node in node_ids: + current = node + node_result = {} + while True: + try: + parent = child_parent_map[current] + except KeyError: + break + parent_layer = layers_map[parent] + node_result[parent_layer] = parent + current = parent + result[node] = node_result + return result + def get_subgraph( self, node_id_or_ids: typing.Union[np.uint64, typing.Iterable], bbox: typing.Optional[typing.Sequence[typing.Sequence[int]]] = None, bbox_is_coordinate: bool = False, - return_layers: typing.List = [2], + return_layers: typing.List = None, nodes_only: bool = False, edges_only: bool = False, leaves_only: bool = False, return_flattened: bool = False, - ) -> typing.Tuple[typing.Dict, typing.Dict, Edges]: + ) -> typing.Tuple[typing.Dict, typing.Tuple[Edges]]: """ Generic subgraph method. """ - from .subgraph import get_subgraph_nodes - from .subgraph import get_subgraph_edges_and_leaves + + if return_layers is None: + return_layers = [2] if nodes_only: return get_subgraph_nodes( @@ -581,7 +621,7 @@ def get_subgraph_nodes( node_id_or_ids: typing.Union[np.uint64, typing.Iterable], bbox: typing.Optional[typing.Sequence[typing.Sequence[int]]] = None, bbox_is_coordinate: bool = False, - return_layers: typing.List = [2], + return_layers: typing.List = None, serializable: bool = False, return_flattened: bool = False, ) -> typing.Tuple[typing.Dict, typing.Dict, Edges]: @@ -589,7 +629,8 @@ def get_subgraph_nodes( Get the children of `node_ids` that are at each of return_layers within the specified bounding box. """ - from .subgraph import get_subgraph_nodes + if return_layers is None: + return_layers = [2] return get_subgraph_nodes( self, @@ -610,8 +651,6 @@ def get_subgraph_edges( """ Get the atomic edges of the `node_ids` within the specified bounding box. """ - from .subgraph import get_subgraph_edges_and_leaves - return get_subgraph_edges_and_leaves( self, node_id_or_ids, bbox, bbox_is_coordinate, True, False ) @@ -625,39 +664,98 @@ def get_subgraph_leaves( """ Get the supervoxels of the `node_ids` within the specified bounding box. """ - from .subgraph import get_subgraph_edges_and_leaves - return get_subgraph_edges_and_leaves( self, node_id_or_ids, bbox, bbox_is_coordinate, False, True ) - def get_fake_edges( + def get_edges_from_edits( self, chunk_ids: np.ndarray, time_stamp: datetime.datetime = None ) -> typing.Dict: + """ + Edges stored within a pcg that were created as a result of edits. + Either 'fake' edges that were adding for a merge edit; + Or 'split' edges resulting from a supervoxel split. + + SplitEdges accumulate across operations (append-only, preserves history). + CompactedSplitEdges is a single cell per chunk with only currently-valid + edges — updated by add_new_edges on each SV split (reads existing, filters + out edges referencing replaced SVs, merges with new edges, overwrites). + + For current-time queries (time_stamp=None), reads CompactedSplitEdges + for O(1) cells per chunk. For historical queries, reads all SplitEdges + up to that timestamp. + """ + use_compacted = time_stamp is None and self.meta.ocdbt_seg + if use_compacted: + properties = [ + attributes.Connectivity.FakeEdges, + attributes.Connectivity.CompactedSplitEdges, + attributes.Connectivity.CompactedAffinity, + attributes.Connectivity.CompactedArea, + ] + else: + properties = [ + attributes.Connectivity.FakeEdges, + attributes.Connectivity.SplitEdges, + attributes.Connectivity.Affinity, + attributes.Connectivity.Area, + ] result = {} - fake_edges_d = self.client.read_nodes( + _edges_d = self.client.read_nodes( node_ids=chunk_ids, - properties=attributes.Connectivity.FakeEdges, + properties=properties, end_time=time_stamp, end_time_inclusive=True, fake_edges=True, ) - for id_, val in fake_edges_d.items(): - edges = np.concatenate( - [np.array(e.value, dtype=basetypes.NODE_ID) for e in val] - ) - result[id_] = Edges(edges[:, 0], edges[:, 1], fake_edges=True) + for id_, val in _edges_d.items(): + edges = val.get(attributes.Connectivity.FakeEdges, []) + edges = np.concatenate([types.empty_2d, *[e.value for e in edges]]) + fake_edges_ = Edges(edges[:, 0], edges[:, 1]) + + if use_compacted: + se = val.get(attributes.Connectivity.CompactedSplitEdges, []) + af = val.get(attributes.Connectivity.CompactedAffinity, []) + ar = val.get(attributes.Connectivity.CompactedArea, []) + else: + se = val.get(attributes.Connectivity.SplitEdges, []) + af = val.get(attributes.Connectivity.Affinity, []) + ar = val.get(attributes.Connectivity.Area, []) + + edges = np.concatenate([types.empty_2d, *[e.value for e in se]]) + aff = np.concatenate([types.empty_affinities, *[e.value for e in af]]) + areas = np.concatenate([types.empty_areas, *[e.value for e in ar]]) + split_edges_ = Edges(edges[:, 0], edges[:, 1], affinities=aff, areas=areas) + + result[id_] = fake_edges_ + split_edges_ return result + def copy_fake_edges(self, chunk_id: np.uint64) -> None: + _edges = self.client.read_node( + node_id=chunk_id, + properties=attributes.Connectivity.FakeEdgesCF3, + end_time_inclusive=True, + fake_edges=True, + ) + mutations = [] + _id = serializers.serialize_uint64(chunk_id, fake_edges=True) + for e in _edges: + val_dict = {attributes.Connectivity.FakeEdges: e.value} + row = self.client.mutate_row(_id, val_dict, time_stamp=e.timestamp) + mutations.append(row) + self.client.write(mutations) + def get_l2_agglomerations( - self, level2_ids: np.ndarray, edges_only: bool = False - ) -> typing.Tuple[typing.Dict[int, types.Agglomeration], np.ndarray]: + self, + level2_ids: np.ndarray, + edges_only: bool = False, + active: bool = False, + time_stamp: typing.Optional[datetime.datetime] = None, + ) -> typing.Tuple[typing.Dict[int, types.Agglomeration], typing.Tuple[Edges]]: """ Children of Level 2 Node IDs and edges. 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)) @@ -668,12 +766,14 @@ def get_l2_agglomerations( if self.mock_edges is None: edges_d = self.read_chunk_edges(chunk_ids) - fake_edges = self.get_fake_edges(chunk_ids) + edited_edges = self.get_edges_from_edits(chunk_ids) all_chunk_edges = reduce( lambda x, y: x + y, - chain(edges_d.values(), fake_edges.values()), + chain(edges_d.values(), edited_edges.values()), Edges([], []), ) + if self.mock_edges is not None: + all_chunk_edges += self.mock_edges if edges_only: if self.mock_edges is not None: @@ -681,20 +781,27 @@ def get_l2_agglomerations( 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) + mask0 = np.isin(all_chunk_edges[:, 0], supervoxels) + mask1 = np.isin(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] + for sv in svs: + if sv in sv_parent_d: + raise ValueError("Found conflicting parents.") sv_parent_d.update(dict(zip(svs.tolist(), [l2id] * len(svs)))) + all_chunk_edges = self._filter_stale_svs(all_chunk_edges, sv_parent_d) + if active: + all_chunk_edges = edge_utils.filter_inactive_cross_edges( + self, all_chunk_edges, time_stamp=time_stamp + ) + in_edges, out_edges, cross_edges = edge_utils.categorize_edges_v2( - self.meta, - all_chunk_edges, - sv_parent_d + self.meta, all_chunk_edges, sv_parent_d ) agglomeration_d = get_agglomerations( @@ -702,13 +809,44 @@ def get_l2_agglomerations( ) return ( agglomeration_d, - (self.mock_edges,) - if self.mock_edges is not None - else (in_edges, out_edges, cross_edges), + ( + (self.mock_edges,) + if self.mock_edges is not None + else (in_edges, out_edges, cross_edges) + ), + ) + + def _filter_stale_svs(self, edges: Edges, sv_parent_d: dict) -> Edges: + """Filter edges referencing SVs replaced by prior SV splits. + + Stale SVs have NewIdentity set (replaced by split fragments) but are + not in sv_parent_d. Cross-root SVs also aren't in sv_parent_d but + don't have NewIdentity — those edges are legitimate and kept. + Only applies to ocdbt_seg graphs (SV splitting is not possible otherwise). + """ + if not self.meta.ocdbt_seg or len(edges) == 0: + return edges + all_svs = np.unique(np.concatenate([edges.node_ids1, edges.node_ids2])) + unknown_svs = np.array( + [sv for sv in all_svs if sv not in sv_parent_d], dtype=np.uint64 + ) + if len(unknown_svs) == 0: + return edges + new_id_cells = self.client.read_nodes( + node_ids=unknown_svs, + properties=attributes.Hierarchy.NewIdentity, + ) + stale_svs = set(int(sv) for sv in unknown_svs if new_id_cells.get(sv)) + if not stale_svs: + return edges + stale_arr = np.array(list(stale_svs), dtype=np.uint64) + keep_m = ~np.isin(edges.node_ids1, stale_arr) & ~np.isin( + edges.node_ids2, stale_arr ) + return edges[keep_m] def get_node_timestamps( - self, node_ids: typing.Sequence[np.uint64], return_numpy=True + self, node_ids: typing.Sequence[np.uint64], return_numpy=True, normalize=False ) -> typing.Iterable: """ The timestamp of the children column can be assumed @@ -722,22 +860,32 @@ def get_node_timestamps( if return_numpy: return np.array([], dtype=np.datetime64) return [] + result = [] + earliest_ts = self.get_earliest_timestamp() + for n in node_ids: + try: + ts = children[n][0].timestamp + except KeyError: + ts = datetime.datetime.now(datetime.timezone.utc) + if normalize: + ts = earliest_ts if ts < earliest_ts else ts + result.append(ts) if return_numpy: - return np.array( - [children[x][0].timestamp for x in node_ids], dtype=np.datetime64 - ) - return [children[x][0].timestamp for x in node_ids] + return np.array(result, dtype=np.datetime64) + return result # OPERATIONS def add_edges( self, user_id: str, - atomic_edges: typing.Sequence[np.uint64], + atomic_edges: typing.Sequence[typing.Sequence[np.uint64]], *, affinities: typing.Sequence[np.float32] = None, source_coords: typing.Sequence[int] = None, sink_coords: typing.Sequence[int] = None, allow_same_segment_merge: typing.Optional[bool] = False, + do_sanity_check: typing.Optional[bool] = True, + stitch_mode: typing.Optional[bool] = False, ) -> operation.GraphEditOperation.Result: """ Adds an edge to the chunkedgraph @@ -754,6 +902,8 @@ def add_edges( source_coords=source_coords, sink_coords=sink_coords, allow_same_segment_merge=allow_same_segment_merge, + do_sanity_check=do_sanity_check, + stitch_mode=stitch_mode, ).execute() def remove_edges( @@ -769,6 +919,7 @@ def remove_edges( path_augment: bool = True, disallow_isolating_cut: bool = True, bb_offset: typing.Tuple[int, int, int] = (240, 240, 24), + do_sanity_check: typing.Optional[bool] = True, ) -> operation.GraphEditOperation.Result: """ Removes edges - either directly or after applying a mincut @@ -793,6 +944,7 @@ def remove_edges( bbox_offset=bb_offset, path_augment=path_augment, disallow_isolating_cut=disallow_isolating_cut, + do_sanity_check=do_sanity_check, ).execute() if not atomic_edges: @@ -842,82 +994,7 @@ def redo_operation( multicut_as_split=True, ).execute() - # PRIVATE - - def _get_bounding_chunk_ids( - self, - parent_chunk_ids: typing.Iterable, - unique: bool = False, - ) -> typing.Dict: - """ - Returns bounding chunk IDs at layers < parent_layer for all chunk IDs. - Dict[parent_chunk_id] = np.array(bounding_chunk_ids) - """ - parent_chunk_coords = self.get_chunk_coordinates_multiple(parent_chunk_ids) - parents_layer = self.get_chunk_layer(parent_chunk_ids[0]) - chunk_id_bchunk_ids_d = {} - for i, chunk_id in enumerate(parent_chunk_ids): - if chunk_id in chunk_id_bchunk_ids_d: - # `parent_chunk_ids` can have duplicates - # avoid redundant calculations - continue - parent_coord = parent_chunk_coords[i] - chunk_ids = [types.empty_1d] - for child_layer in range(2, parents_layer): - bcoords = chunk_utils.get_bounding_children_chunks( - self.meta, - parents_layer, - parent_coord, - child_layer, - return_unique=False, - ) - bchunks_ids = chunk_utils.get_chunk_ids_from_coords( - self.meta, child_layer, bcoords - ) - chunk_ids.append(bchunks_ids) - chunk_ids = np.concatenate(chunk_ids) - if unique: - chunk_ids = np.unique(chunk_ids) - chunk_id_bchunk_ids_d[chunk_id] = chunk_ids - return chunk_id_bchunk_ids_d - - def _get_bounding_l2_children(self, parents: typing.Iterable) -> typing.Dict: - parent_chunk_ids = self.get_chunk_ids_from_node_ids(parents) - chunk_id_bchunk_ids_d = self._get_bounding_chunk_ids( - parent_chunk_ids, unique=len(parents) >= 200 - ) - - parent_descendants_d = { - _id: np.array([_id], dtype=basetypes.NODE_ID) for _id in parents - } - descendants_all = np.concatenate(list(parent_descendants_d.values())) - descendants_layers = self.get_chunk_layers(descendants_all) - layer_mask = descendants_layers > 2 - descendants_all = descendants_all[layer_mask] - - while descendants_all.size: - descendant_children_d = self.get_children(descendants_all) - for i, parent_id in enumerate(parents): - _descendants = parent_descendants_d[parent_id] - _layers = self.get_chunk_layers(_descendants) - _l2mask = _layers == 2 - descendants = [_descendants[_l2mask]] - for child in _descendants[~_l2mask]: - descendants.append(descendant_children_d[child]) - descendants = np.concatenate(descendants) - chunk_ids = self.get_chunk_ids_from_node_ids(descendants) - bchunk_ids = chunk_id_bchunk_ids_d[parent_chunk_ids[i]] - bounding_descendants = descendants[np.in1d(chunk_ids, bchunk_ids)] - parent_descendants_d[parent_id] = bounding_descendants - - descendants_all = np.concatenate(list(parent_descendants_d.values())) - descendants_layers = self.get_chunk_layers(descendants_all) - layer_mask = descendants_layers > 2 - descendants_all = descendants_all[layer_mask] - return parent_descendants_d - # HELPERS / WRAPPERS - def is_root(self, node_id: basetypes.NODE_ID) -> bool: return self.get_chunk_layer(node_id) == self.meta.layer_count @@ -955,11 +1032,26 @@ def get_chunk_coordinates(self, node_or_chunk_id: basetypes.NODE_ID): return chunk_utils.get_chunk_coordinates(self.meta, node_or_chunk_id) def get_chunk_coordinates_multiple(self, node_or_chunk_ids: typing.Sequence): - node_or_chunk_ids = np.array(node_or_chunk_ids, dtype=basetypes.NODE_ID) + node_or_chunk_ids = np.asarray(node_or_chunk_ids, dtype=basetypes.NODE_ID) layers = self.get_chunk_layers(node_or_chunk_ids) - assert np.all(layers == layers[0]), "All IDs must have the same layer." + assert len(layers) == 0 or np.all(layers == layers[0]), "must be same layer." return chunk_utils.get_chunk_coordinates_multiple(self.meta, node_or_chunk_ids) + def get_chunk_center_voxel(self, node_or_chunk_id: basetypes.NODE_ID) -> np.ndarray: + """Approximate base-resolution voxel coord at the chunk's center. + + Useful for debugging: feed the returned ``[x, y, z]`` to NGL's + position bar to navigate to where a chunk lives in the volume. + Layer L chunk side = ``CHUNK_SIZE * 2 ** (L - 2)`` base voxels. + """ + layer = int(self.get_chunk_layer(node_or_chunk_id)) + cx, cy, cz = self.get_chunk_coordinates(node_or_chunk_id) + chunk_size = np.asarray(self.meta.graph_config.CHUNK_SIZE, dtype=int) * ( + 2 ** (layer - 2) + ) + origin = self.meta.voxel_bounds[:, 0] + np.array([cx, cy, cz]) * chunk_size + return (origin + chunk_size // 2).astype(int) + def get_chunk_id( self, node_id: basetypes.NODE_ID = None, @@ -987,6 +1079,11 @@ def get_parent_chunk_id( self.meta, node_or_chunk_id, parent_layer ) + def get_parent_chunk_id_multiple(self, node_or_chunk_ids: typing.Sequence): + return chunk_hierarchy.get_parent_chunk_id_multiple( + self.meta, node_or_chunk_ids + ) + def get_parent_chunk_ids(self, node_or_chunk_id: basetypes.NODE_ID): return chunk_hierarchy.get_parent_chunk_ids(self.meta, node_or_chunk_id) @@ -1017,6 +1114,76 @@ def get_earliest_timestamp(self): from datetime import timedelta for op_id in range(100): - _, timestamp = self.client.read_log_entry(op_id) + _log, timestamp = self.client.read_log_entry(op_id) if timestamp is not None: return timestamp - timedelta(milliseconds=500) + if _log: + return self.client.read_node( + op_id, properties=attributes.OperationLogs.Status + )[-1].timestamp + # no ops: the ingest-completion boundary stamped during the root-layer build + stamped = self.meta.custom_data.get("earliest_ts") + if stamped is not None: + return datetime.datetime.fromisoformat(stamped) + return datetime.datetime.fromtimestamp(0, tz=datetime.timezone.utc) + + def get_operation_ids(self, node_ids: typing.Sequence): + response = self.client.read_nodes(node_ids=node_ids) + result = {} + for node in node_ids: + try: + operations = response[node][attributes.OperationLogs.OperationID] + result[node] = [(x.value, x.timestamp) for x in operations] + except KeyError: + ... + return result + + def get_single_leaf_multiple(self, node_ids): + """Returns the first supervoxel found for each node_id.""" + result = {} + node_ids_copy = np.copy(node_ids) + children = np.copy(node_ids) + children_d = self.get_children(node_ids) + while True: + children = [children_d[k][0] for k in children] + children = np.array(children, dtype=basetypes.NODE_ID) + mask = self.get_chunk_layers(children) == 1 + result.update( + [(node, sv) for node, sv in zip(node_ids[mask], children[mask])] + ) + node_ids = node_ids[~mask] + children = children[~mask] + if children.size == 0: + break + children_d = self.get_children(children) + return np.array([result[k] for k in node_ids_copy], dtype=basetypes.NODE_ID) + + def get_chunk_layers_and_coordinates(self, node_or_chunk_ids: typing.Sequence): + """ + Helper function that wraps get chunk layer and coordinates for nodes at any layer. + """ + node_or_chunk_ids = np.array(node_or_chunk_ids, dtype=basetypes.NODE_ID) + layers = self.get_chunk_layers(node_or_chunk_ids) + chunk_coords = np.zeros(shape=(len(node_or_chunk_ids), 3), dtype=int) + for _layer in np.unique(layers): + mask = layers == _layer + _nodes = node_or_chunk_ids[mask] + chunk_coords[mask] = chunk_utils.get_chunk_coordinates_multiple( + self.meta, _nodes + ) + return layers, chunk_coords + + def get_l2children(self, node_ids) -> np.ndarray: + """ + Get L2 children of all node_ids, returns a flat array. + """ + node_ids = np.asarray(node_ids, dtype=basetypes.NODE_ID) + layers = self.get_chunk_layers(node_ids) + assert np.all(layers >= 2), "nodes must be at layers >= 2" + l2children = [types.empty_1d] + while node_ids.size: + children = self.get_children(node_ids, flatten=True) + layers = self.get_chunk_layers(children) + l2children.append(children[layers == 2]) + node_ids = children[layers > 2] + return np.concatenate(l2children) diff --git a/pychunkedgraph/graph/chunks/atomic.py b/pychunkedgraph/graph/chunks/atomic.py index e3de065ff..ec0109c69 100644 --- a/pychunkedgraph/graph/chunks/atomic.py +++ b/pychunkedgraph/graph/chunks/atomic.py @@ -1,3 +1,5 @@ +# pylint: disable=invalid-name, missing-docstring + from typing import List from typing import Sequence from itertools import product @@ -6,8 +8,6 @@ from .utils import get_bounding_children_chunks from ..meta import ChunkedGraphMeta -from ..utils.generic import get_valid_timestamp -from ..utils import basetypes def get_touching_atomic_chunks( @@ -27,7 +27,7 @@ def get_touching_atomic_chunks( chunk_offset = chunk_coords * atomic_chunk_count mid = (atomic_chunk_count // 2) - 1 - # TODO (akhileshh) convert this for loop to numpy + # TODO (akhileshh) convert this for loop to numpy; # relevant chunks along touching planes at center for axis_1, axis_2 in product(*[range(atomic_chunk_count)] * 2): # x-y plane @@ -62,4 +62,6 @@ def get_bounding_atomic_chunks( chunkedgraph_meta: ChunkedGraphMeta, layer: int, chunk_coords: Sequence[int] ) -> List: """Atomic chunk coordinates along the boundary of a chunk""" - return get_bounding_children_chunks(chunkedgraph_meta, layer, chunk_coords, 2) + return get_bounding_children_chunks( + chunkedgraph_meta, layer, tuple(chunk_coords), 2 + ) diff --git a/pychunkedgraph/graph/chunks/hierarchy.py b/pychunkedgraph/graph/chunks/hierarchy.py index 32d6029ee..5ff7823fe 100644 --- a/pychunkedgraph/graph/chunks/hierarchy.py +++ b/pychunkedgraph/graph/chunks/hierarchy.py @@ -37,17 +37,17 @@ def get_children_chunk_ids( layer = utils.get_chunk_layer(meta, node_or_chunk_id) if layer == 1: - return np.array([]) + return np.array([], dtype=np.uint64) elif layer == 2: return np.array([utils.get_chunk_id(meta, layer=layer, x=x, y=y, z=z)]) else: children_coords = get_children_chunk_coords(meta, layer, (x, y, z)) children_chunk_ids = [] - for (x, y, z) in children_coords: + for x, y, z in children_coords: children_chunk_ids.append( utils.get_chunk_id(meta, layer=layer - 1, x=x, y=y, z=z) ) - return np.array(children_chunk_ids) + return np.array(children_chunk_ids, dtype=np.uint64) def get_parent_chunk_id( @@ -62,6 +62,19 @@ def get_parent_chunk_id( return utils.get_chunk_id(meta, layer=parent_layer, x=x, y=y, z=z) +def get_parent_chunk_id_multiple( + meta: ChunkedGraphMeta, node_or_chunk_ids: np.ndarray +) -> np.ndarray: + """Parent chunk IDs for multiple nodes. Assumes nodes at same layer.""" + + node_layers = utils.get_chunk_layers(meta, node_or_chunk_ids) + assert np.unique(node_layers).size == 1, np.unique(node_layers) + parent_layer = node_layers[0] + 1 + coords = utils.get_chunk_coordinates_multiple(meta, node_or_chunk_ids) + coords = coords // meta.graph_config.FANOUT + return utils.get_chunk_ids_from_coords(meta, layer=parent_layer, coords=coords) + + def get_parent_chunk_ids( meta: ChunkedGraphMeta, node_or_chunk_id: np.uint64 ) -> np.ndarray: diff --git a/pychunkedgraph/graph/chunks/utils.py b/pychunkedgraph/graph/chunks/utils.py index dc895bde4..cd3b96ccc 100644 --- a/pychunkedgraph/graph/chunks/utils.py +++ b/pychunkedgraph/graph/chunks/utils.py @@ -1,13 +1,17 @@ # pylint: disable=invalid-name, missing-docstring -from typing import List from typing import Union from typing import Optional from typing import Sequence +from typing import Tuple from typing import Iterable +from copy import copy +from functools import lru_cache + import numpy as np + def get_chunks_boundary(voxel_boundary, chunk_size) -> np.ndarray: """returns number of chunks in each dimension""" return np.ceil((voxel_boundary / chunk_size)).astype(int) @@ -43,7 +47,7 @@ def normalize_bounding_box( def get_chunk_layer(meta, node_or_chunk_id: np.uint64) -> int: - """ Extract Layer from Node ID or Chunk ID """ + """Extract Layer from Node ID or Chunk ID""" return int(int(node_or_chunk_id) >> 64 - meta.graph_config.LAYER_ID_BITS) @@ -75,9 +79,9 @@ def get_chunk_coordinates(meta, node_or_chunk_id: np.uint64) -> np.ndarray: y_offset = x_offset - bits_per_dim z_offset = y_offset - bits_per_dim - x = int(node_or_chunk_id) >> x_offset & 2 ** bits_per_dim - 1 - y = int(node_or_chunk_id) >> y_offset & 2 ** bits_per_dim - 1 - z = int(node_or_chunk_id) >> z_offset & 2 ** bits_per_dim - 1 + x = int(node_or_chunk_id) >> x_offset & 2**bits_per_dim - 1 + y = int(node_or_chunk_id) >> y_offset & 2**bits_per_dim - 1 + z = int(node_or_chunk_id) >> z_offset & 2**bits_per_dim - 1 return np.array([x, y, z]) @@ -86,8 +90,8 @@ def get_chunk_coordinates_multiple(meta, ids: np.ndarray) -> np.ndarray: Array version of get_chunk_coordinates. Assumes all given IDs are in same layer. """ - if not len(ids): - return np.array([]) + if len(ids) == 0: + return np.array([], dtype=int).reshape(0, 3) layer = get_chunk_layer(meta, ids[0]) bits_per_dim = meta.bitmasks[layer] @@ -96,9 +100,9 @@ def get_chunk_coordinates_multiple(meta, ids: np.ndarray) -> np.ndarray: z_offset = y_offset - bits_per_dim ids = np.array(ids, dtype=int) - X = ids >> x_offset & 2 ** bits_per_dim - 1 - Y = ids >> y_offset & 2 ** bits_per_dim - 1 - Z = ids >> z_offset & 2 ** bits_per_dim - 1 + X = ids >> x_offset & 2**bits_per_dim - 1 + Y = ids >> y_offset & 2**bits_per_dim - 1 + Z = ids >> z_offset & 2**bits_per_dim - 1 return np.column_stack((X, Y, Z)) @@ -125,6 +129,7 @@ def get_chunk_id( def get_chunk_ids_from_coords(meta, layer: int, coords: np.ndarray): + layer = int(layer) result = np.zeros(len(coords), dtype=np.uint64) s_bits_per_dim = meta.bitmasks[layer] @@ -142,14 +147,15 @@ def get_chunk_ids_from_coords(meta, layer: int, coords: np.ndarray): def get_chunk_ids_from_node_ids(meta, ids: Iterable[np.uint64]) -> np.ndarray: - """ Extract Chunk IDs from Node IDs""" + """Extract Chunk IDs from Node IDs""" if len(ids) == 0: return np.array([], dtype=np.uint64) bits_per_dims = np.array([meta.bitmasks[l] for l in get_chunk_layers(meta, ids)]) offsets = 64 - meta.graph_config.LAYER_ID_BITS - 3 * bits_per_dims - cids1 = np.array((np.array(ids, dtype=int) >> offsets) << offsets, dtype=np.uint64) + ids = np.array(ids, dtype=int) + cids1 = np.array((ids >> offsets) << offsets, dtype=np.uint64) # cids2 = np.vectorize(get_chunk_id)(meta, ids) # assert np.all(cids1 == cids2) return cids1 @@ -163,14 +169,10 @@ def _compute_chunk_id( z: int, ) -> np.uint64: s_bits_per_dim = meta.bitmasks[layer] - if not ( - x < 2 ** s_bits_per_dim and y < 2 ** s_bits_per_dim and z < 2 ** s_bits_per_dim - ): - raise ValueError( - f"Coordinate is out of range \ + if not (x < 2**s_bits_per_dim and y < 2**s_bits_per_dim and z < 2**s_bits_per_dim): + raise ValueError(f"Coordinate is out of range \ layer: {layer} bits/dim {s_bits_per_dim}. \ - [{x}, {y}, {z}]; max = {2 ** s_bits_per_dim}." - ) + [{x}, {y}, {z}]; max = {2 ** s_bits_per_dim}.") layer_offset = 64 - meta.graph_config.LAYER_ID_BITS x_offset = layer_offset - s_bits_per_dim y_offset = x_offset - s_bits_per_dim @@ -208,8 +210,9 @@ def _get_chunk_coordinates_from_vol_coordinates( return coords.astype(int) +@lru_cache() def get_bounding_children_chunks( - cg_meta, layer: int, chunk_coords: Sequence[int], children_layer, return_unique=True + cg_meta, layer: int, chunk_coords: Tuple[int], children_layer, return_unique=True ) -> np.ndarray: """Children chunk coordinates at given layer, along the boundary of a chunk""" chunk_coords = np.array(chunk_coords, dtype=int) @@ -233,3 +236,105 @@ def get_bounding_children_chunks( if return_unique: return np.unique(result, axis=0) if result.size else result return result + + +@lru_cache() +def get_l2chunkids_along_boundary( + cg_meta, mlayer: int, coord_a, coord_b, padding: int = 0 +): + """ + Gets L2 Chunk IDs along opposing faces for larger chunks. + If padding is enabled, more faces of L2 chunks are padded on both sides. + This is necessary to find fake edges that can span more than 2 L2 chunks. + """ + bounds_a = get_bounding_children_chunks(cg_meta, mlayer, tuple(coord_a), 2) + bounds_b = get_bounding_children_chunks(cg_meta, mlayer, tuple(coord_b), 2) + + coord_a, coord_b = np.array(coord_a, dtype=int), np.array(coord_b, dtype=int) + direction = coord_a - coord_b + major_axis = np.argmax(np.abs(direction)) + + l2chunk_count = 2 ** (mlayer - 2) + max_coord = coord_a if direction[major_axis] > 0 else coord_b + + skip = abs(direction[major_axis]) - 1 + l2_skip = skip * l2chunk_count + + mid = max_coord[major_axis] * l2chunk_count + face_a = mid if direction[major_axis] > 0 else (mid - l2_skip - 1) + face_b = mid if direction[major_axis] < 0 else (mid - l2_skip - 1) + + l2chunks_a = [bounds_a[bounds_a[:, major_axis] == face_a]] + l2chunks_b = [bounds_b[bounds_b[:, major_axis] == face_b]] + + step_a, step_b = (1, -1) if direction[major_axis] > 0 else (-1, 1) + for _ in range(padding): + _l2_chunks_a = copy(l2chunks_a[-1]) + _l2_chunks_b = copy(l2chunks_b[-1]) + _l2_chunks_a[:, major_axis] += step_a + _l2_chunks_b[:, major_axis] += step_b + l2chunks_a.append(_l2_chunks_a) + l2chunks_b.append(_l2_chunks_b) + + l2chunks_a = np.concatenate(l2chunks_a) + l2chunks_b = np.concatenate(l2chunks_b) + + l2chunk_ids_a = get_chunk_ids_from_coords(cg_meta, 2, l2chunks_a) + l2chunk_ids_b = get_chunk_ids_from_coords(cg_meta, 2, l2chunks_b) + return l2chunk_ids_a, l2chunk_ids_b + + +def chunks_overlapping_bbox(bbox_min, bbox_max, chunk_size, origin=0) -> dict: + """ + Find octree chunks overlapping with a bounding box in 3D + and return a dictionary mapping chunk indices to clipped bounding boxes. + + `origin` is the voxel coordinate of chunk index (0, 0, 0). Pass + `meta.voxel_bounds[:, 0]` so the lattice aligns to the dataset's + chunks; the default 0 leaves the lattice anchored at the volume + origin. + """ + bbox_min = np.asarray(bbox_min, dtype=int) + bbox_max = np.asarray(bbox_max, dtype=int) + chunk_size = np.asarray(chunk_size, dtype=int) + origin = np.asarray(origin, dtype=int) + + start_idx = np.floor_divide(bbox_min - origin, chunk_size).astype(int) + end_idx = np.floor_divide(bbox_max - origin, chunk_size).astype(int) + + ix = np.arange(start_idx[0], end_idx[0] + 1) + iy = np.arange(start_idx[1], end_idx[1] + 1) + iz = np.arange(start_idx[2], end_idx[2] + 1) + grid = np.stack(np.meshgrid(ix, iy, iz, indexing="ij"), axis=-1, dtype=int) + grid = grid.reshape(-1, 3) + + chunk_min = grid * chunk_size + origin + chunk_max = chunk_min + chunk_size + clipped_min = np.maximum(chunk_min, bbox_min) + clipped_max = np.minimum(chunk_max, bbox_max) + return { + tuple(idx): np.stack([cmin, cmax], axis=0, dtype=int) + for idx, cmin, cmax in zip(grid, clipped_min, clipped_max) + } + + +def get_neighbors(coord, inclusive: bool = True, min_coord=None, max_coord=None): + """ + Get all valid coordinates in the 3×3×3 cube around a given chunk, + including the chunk itself (if inclusive=True), + respecting bounding box constraints. + """ + offsets = np.array(np.meshgrid([-1, 0, 1], [-1, 0, 1], [-1, 0, 1])).T.reshape(-1, 3) + if not inclusive: + offsets = offsets[~np.all(offsets == 0, axis=1)] + + neighbors = np.array(coord) + offsets + if min_coord is None: + min_coord = (0, 0, 0) + min_coord = np.array(min_coord) + neighbors = neighbors[(neighbors >= min_coord).all(axis=1)] + + if max_coord is not None: + max_coord = np.array(max_coord) + neighbors = neighbors[(neighbors <= max_coord).all(axis=1)] + return neighbors diff --git a/pychunkedgraph/graph/client/__init__.py b/pychunkedgraph/graph/client/__init__.py deleted file mode 100644 index 6e025bd35..000000000 --- a/pychunkedgraph/graph/client/__init__.py +++ /dev/null @@ -1,44 +0,0 @@ -""" -Sub packages/modules for backend storage clients -Currently supports Google Big Table - -A simple client needs to be able to create the graph, -store graph meta and to write and read node information. -Also needs locking support to prevent race conditions -when modifying root/parent nodes. - -In addition, clients with more features like generating unique IDs -and logging facilities can be implemented by inherting respective base classes. - -These methods are in separate classes because they are logically related. -This also makes it possible to have different backend storage solutions, -making it possible to use any unique features these solutions may provide. - -Please see `base.py` for more details. -""" - -from collections import namedtuple - -from .bigtable.client import Client as BigTableClient - - -_backend_clientinfo_fields = ("TYPE", "CONFIG") -_backend_clientinfo_defaults = (None, None) -BackendClientInfo = namedtuple( - "BackendClientInfo", - _backend_clientinfo_fields, - defaults=_backend_clientinfo_defaults, -) - - -def get_default_client_info(): - """ - Load client from env variables. - """ - - # TODO make dynamic after multiple platform support is added - from .bigtable import get_client_info as get_bigtable_client_info - - return BackendClientInfo( - CONFIG=get_bigtable_client_info(admin=True, read_only=False) - ) diff --git a/pychunkedgraph/graph/client/base.py b/pychunkedgraph/graph/client/base.py deleted file mode 100644 index a66602a6a..000000000 --- a/pychunkedgraph/graph/client/base.py +++ /dev/null @@ -1,152 +0,0 @@ -from abc import ABC -from abc import abstractmethod - - -class SimpleClient(ABC): - """ - Abstract class for interacting with backend data store where the chunkedgraph is stored. - Eg., BigTableClient for using big table as storage. - """ - - @abstractmethod - def create_graph(self) -> None: - """Initialize the graph and store associated meta.""" - - @abstractmethod - def add_graph_version(self, version): - """Add a version to the graph.""" - - @abstractmethod - def read_graph_version(self): - """Read stored graph version.""" - - @abstractmethod - def update_graph_meta(self, meta): - """Update stored graph meta.""" - - @abstractmethod - def read_graph_meta(self): - """Read stored graph meta.""" - - @abstractmethod - def read_nodes( - self, - start_id=None, - end_id=None, - node_ids=None, - properties=None, - start_time=None, - end_time=None, - end_time_inclusive=False, - ): - """ - Read nodes and their properties. - Accepts a range of node IDs or specific node IDs. - """ - - @abstractmethod - def read_node( - self, - node_id, - properties=None, - start_time=None, - end_time=None, - end_time_inclusive=False, - ): - """Read a single node and it's properties.""" - - @abstractmethod - def write_nodes(self, nodes): - """Writes/updates nodes (IDs along with properties).""" - - @abstractmethod - def lock_root(self, node_id, operation_id): - """Locks root node with operation_id to prevent race conditions.""" - - @abstractmethod - def lock_roots(self, node_ids, operation_id): - """Locks root nodes to prevent race conditions.""" - - @abstractmethod - def lock_root_indefinitely(self, node_id, operation_id): - """Locks root node with operation_id to prevent race conditions.""" - - @abstractmethod - def lock_roots_indefinitely(self, node_ids, operation_id): - """ - Locks root nodes indefinitely to prevent structural damage to graph. - This scenario is rare and needs asynchronous fix or inspection to unlock. - """ - - @abstractmethod - def unlock_root(self, node_id, operation_id): - """Unlocks root node that is locked with operation_id.""" - - @abstractmethod - def unlock_indefinitely_locked_root(self, node_id, operation_id): - """Unlocks root node that is indefinitely locked with operation_id.""" - - @abstractmethod - def renew_lock(self, node_id, operation_id): - """Renews existing node lock with operation_id for extended time.""" - - @abstractmethod - def renew_locks(self, node_ids, operation_id): - """Renews existing node locks with operation_id for extended time.""" - - @abstractmethod - def get_lock_timestamp(self, node_ids, operation_id): - """Reads timestamp from lock row to get a consistent timestamp.""" - - @abstractmethod - def get_consolidated_lock_timestamp(self, root_ids, operation_ids): - """Minimum of multiple lock timestamps.""" - - @abstractmethod - def get_compatible_timestamp(self, time_stamp): - """Datetime time stamp compatible with client's services.""" - - -class ClientWithIDGen(SimpleClient): - """ - Abstract class for client to backend data store that has support for generating IDs. - If not, something else can be used but these methods need to be implemented. - Eg., Big Table row cells can be used to generate unique IDs. - """ - - @abstractmethod - def create_node_ids(self, chunk_id): - """Generate a range of unique IDs in the chunk.""" - - @abstractmethod - def create_node_id(self, chunk_id): - """Generate a unique ID in the chunk.""" - - @abstractmethod - def get_max_node_id(self, chunk_id): - """Gets the current maximum node ID in the chunk.""" - - @abstractmethod - def create_operation_id(self): - """Generate a unique operation ID.""" - - @abstractmethod - def get_max_operation_id(self): - """Gets the current maximum operation ID.""" - - -class OperationLogger(ABC): - """ - Abstract class for interacting with backend data store where the operation logs are stored. - Eg., BigTableClient can be used to store logs in Google BigTable. - """ - - # TODO add functions for writing - - @abstractmethod - def read_log_entry(self, operation_id: int) -> None: - """Read log entry for a given operation ID.""" - - @abstractmethod - def read_log_entries(self, operation_ids) -> None: - """Read log entries for given operation IDs.""" diff --git a/pychunkedgraph/graph/client/bigtable/__init__.py b/pychunkedgraph/graph/client/bigtable/__init__.py deleted file mode 100644 index b3dbd777b..000000000 --- a/pychunkedgraph/graph/client/bigtable/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -from collections import namedtuple -from os import environ - -DEFAULT_PROJECT = "neuromancer-seung-import" -DEFAULT_INSTANCE = "pychunkedgraph" - -_bigtableconfig_fields = ( - "PROJECT", - "INSTANCE", - "ADMIN", - "READ_ONLY", - "CREDENTIALS", - "MAX_ROW_KEY_COUNT" -) -_bigtableconfig_defaults = ( - environ.get("BIGTABLE_PROJECT", DEFAULT_PROJECT), - environ.get("BIGTABLE_INSTANCE", DEFAULT_INSTANCE), - False, - True, - None, - 1000 -) -BigTableConfig = namedtuple( - "BigTableConfig", _bigtableconfig_fields, defaults=_bigtableconfig_defaults -) - - -def get_client_info( - project: str = None, - instance: str = None, - admin: bool = False, - read_only: bool = True, -): - """Helper function to load config from env.""" - _project = environ.get("BIGTABLE_PROJECT", DEFAULT_PROJECT) - if project: - _project = project - - _instance = environ.get("BIGTABLE_INSTANCE", DEFAULT_INSTANCE) - if instance: - _instance = instance - - kwargs = { - "PROJECT": _project, - "INSTANCE": _instance, - "ADMIN": admin, - "READ_ONLY": read_only, - } - return BigTableConfig(**kwargs) diff --git a/pychunkedgraph/graph/client/bigtable/client.py b/pychunkedgraph/graph/client/bigtable/client.py deleted file mode 100644 index 5b86826bd..000000000 --- a/pychunkedgraph/graph/client/bigtable/client.py +++ /dev/null @@ -1,860 +0,0 @@ -# pylint: disable=invalid-name, missing-docstring, import-outside-toplevel, line-too-long, protected-access, arguments-differ, arguments-renamed, logging-fstring-interpolation - -import sys -import time -import typing -import logging -import datetime -from datetime import datetime - -import numpy as np -from multiwrapper import multiprocessing_utils as mu -from google.cloud import bigtable -from google.api_core.retry import Retry -from google.api_core.retry import if_exception_type -from google.api_core.exceptions import Aborted -from google.api_core.exceptions import DeadlineExceeded -from google.api_core.exceptions import ServiceUnavailable -from google.cloud.bigtable.table import Table -from google.cloud.bigtable.row_set import RowSet -from google.cloud.bigtable.row_data import PartialRowData -from google.cloud.bigtable.row_filters import RowFilter -from google.cloud.bigtable.column_family import MaxVersionsGCRule - -from . import utils -from . import BigTableConfig -from ..base import ClientWithIDGen -from ..base import OperationLogger -from ... import attributes -from ... import exceptions -from ...utils import basetypes -from ...utils.serializers import pad_node_id -from ...utils.serializers import serialize_key -from ...utils.serializers import serialize_uint64 -from ...utils.serializers import deserialize_uint64 -from ...meta import ChunkedGraphMeta -from ...utils.generic import get_valid_timestamp - - -class Client(bigtable.Client, ClientWithIDGen, OperationLogger): - def __init__( - self, - table_id: str, - config: BigTableConfig = BigTableConfig(), - graph_meta: ChunkedGraphMeta = None, - ): - if config.CREDENTIALS: - super(Client, self).__init__( - project=config.PROJECT, - read_only=config.READ_ONLY, - admin=config.ADMIN, - credentials=config.CREDENTIALS, - ) - else: - super(Client, self).__init__( - project=config.PROJECT, - read_only=config.READ_ONLY, - admin=config.ADMIN, - ) - self._instance = self.instance(config.INSTANCE) - self._table = self._instance.table(table_id) - - self.logger = logging.getLogger( - f"{config.PROJECT}/{config.INSTANCE}/{table_id}" - ) - self.logger.setLevel(logging.WARNING) - if not self.logger.handlers: - sh = logging.StreamHandler(sys.stdout) - sh.setLevel(logging.WARNING) - self.logger.addHandler(sh) - self._graph_meta = graph_meta - self._version = None - self._max_row_key_count = config.MAX_ROW_KEY_COUNT - - @property - def graph_meta(self): - return self._graph_meta - - def create_graph(self, meta: ChunkedGraphMeta, version: str) -> None: - """Initialize the graph and store associated meta.""" - if self._table.exists(): - raise ValueError(f"{self._table.table_id} already exists.") - self._table.create() - self._create_column_families() - self.add_graph_version(version) - self.update_graph_meta(meta) - - def add_graph_version(self, version: str): - assert self.read_graph_version() is None, "Graph has already been versioned." - self._version = version - row = self.mutate_row( - attributes.GraphVersion.key, - {attributes.GraphVersion.Version: version}, - ) - self.write([row]) - - def read_graph_version(self) -> str: - try: - row = self._read_byte_row(attributes.GraphVersion.key) - self._version = row[attributes.GraphVersion.Version][0].value - return self._version - except KeyError: - return None - - def _delete_meta(self): - # temprorary fix, use new column with GCRule for permanent fix - # delete existing meta before update, but compatibilty issues - meta_row = self._table.direct_row(attributes.GraphMeta.key) - meta_row.delete() - meta_row.commit() - - def update_graph_meta( - self, meta: ChunkedGraphMeta, overwrite: typing.Optional[bool] = False - ): - if overwrite: - self._delete_meta() - self._graph_meta = meta - row = self.mutate_row( - attributes.GraphMeta.key, - {attributes.GraphMeta.Meta: meta}, - ) - self.write([row]) - - def read_graph_meta(self) -> ChunkedGraphMeta: - row = self._read_byte_row(attributes.GraphMeta.key) - self._graph_meta = row[attributes.GraphMeta.Meta][0].value - return self._graph_meta - - def read_nodes( - self, - start_id=None, - end_id=None, - end_id_inclusive=False, - user_id=None, - node_ids=None, - properties=None, - start_time=None, - end_time=None, - end_time_inclusive: bool = False, - fake_edges: bool = False, - ): - """ - Read nodes and their properties. - Accepts a range of node IDs or specific node IDs. - """ - if node_ids is not None and len(node_ids) > self._max_row_key_count: - # bigtable reading is faster - # when all IDs in a block are within a range - node_ids = np.sort(node_ids) - rows = self._read_byte_rows( - start_key=serialize_uint64(start_id, fake_edges=fake_edges) - if start_id is not None - else None, - end_key=serialize_uint64(end_id, fake_edges=fake_edges) - if end_id is not None - else None, - end_key_inclusive=end_id_inclusive, - row_keys=( - serialize_uint64(node_id, fake_edges=fake_edges) for node_id in node_ids - ) - if node_ids is not None - else None, - columns=properties, - start_time=start_time, - end_time=end_time, - end_time_inclusive=end_time_inclusive, - user_id=user_id, - ) - return { - deserialize_uint64(row_key, fake_edges=fake_edges): data - for (row_key, data) in rows.items() - } - - def read_node( - self, - node_id: np.uint64, - properties: typing.Optional[ - typing.Union[typing.Iterable[attributes._Attribute], attributes._Attribute] - ] = None, - start_time: typing.Optional[datetime] = None, - end_time: typing.Optional[datetime] = None, - end_time_inclusive: bool = False, - fake_edges: bool = False, - ) -> typing.Union[ - typing.Dict[attributes._Attribute, typing.List[bigtable.row_data.Cell]], - typing.List[bigtable.row_data.Cell], - ]: - """Convenience function for reading a single node from Bigtable. - Arguments: - node_id {np.uint64} -- the NodeID of the row to be read. - Keyword Arguments: - columns {typing.Optional[typing.Union[typing.Iterable[attributes._Attribute], attributes._Attribute]]} -- - typing.Optional filtering by columns to speed up the query. If `columns` is a single - column (not iterable), the column key will be omitted from the result. - (default: {None}) - start_time {typing.Optional[datetime]} -- Ignore cells with timestamp before - `start_time`. If None, no lower bound. (default: {None}) - end_time {typing.Optional[datetime]} -- Ignore cells with timestamp after `end_time`. - If None, no upper bound. (default: {None}) - end_time_inclusive {bool} -- Whether or not `end_time` itself should be included in the - request, ignored if `end_time` is None. (default: {False}) - Returns: - typing.Union[typing.Dict[attributes._Attribute, typing.List[bigtable.row_data.Cell]], - typing.List[bigtable.row_data.Cell]] -- - Returns a mapping of columns to a typing.List of cells (one cell per timestamp). Each cell - has a `value` property, which returns the deserialized field, and a `timestamp` - property, which returns the timestamp as `datetime` object. - If only a single `attributes._Attribute` was requested, the typing.List of cells is returned - directly. - """ - return self._read_byte_row( - row_key=serialize_uint64(node_id, fake_edges=fake_edges), - columns=properties, - start_time=start_time, - end_time=end_time, - end_time_inclusive=end_time_inclusive, - ) - - def write_nodes(self, nodes, root_ids=None, operation_id=None): - """ - Writes/updates nodes (IDs along with properties) - by locking root nodes until changes are written. - """ - - def read_log_entry( - self, operation_id: np.uint64 - ) -> typing.Tuple[typing.Dict, datetime]: - log_record = self.read_node( - operation_id, properties=attributes.OperationLogs.all() - ) - if len(log_record) == 0: - return {}, None - try: - timestamp = log_record[attributes.OperationLogs.OperationTimeStamp][0].value - except KeyError: - timestamp = log_record[attributes.OperationLogs.RootID][0].timestamp - log_record.update((column, v[0].value) for column, v in log_record.items()) - return log_record, timestamp - - def read_log_entries( - self, - operation_ids: typing.Optional[typing.Iterable] = None, - user_id: typing.Optional[str] = None, - 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, - ): - if properties is None: - properties = attributes.OperationLogs.all() - - if operation_ids is None: - logs_d = self.read_nodes( - start_id=np.uint64(0), - end_id=self.get_max_operation_id(), - end_id_inclusive=True, - user_id=user_id, - properties=properties, - start_time=start_time, - end_time=end_time, - end_time_inclusive=end_time_inclusive, - ) - else: - logs_d = self.read_nodes( - node_ids=operation_ids, - properties=properties, - start_time=start_time, - end_time=end_time, - end_time_inclusive=end_time_inclusive, - user_id=user_id, - ) - if not logs_d: - return {} - for operation_id in logs_d: - log_record = logs_d[operation_id] - try: - timestamp = log_record[attributes.OperationLogs.OperationTimeStamp][ - 0 - ].value - except KeyError: - timestamp = log_record[attributes.OperationLogs.RootID][0].timestamp - log_record.update((column, v[0].value) for column, v in log_record.items()) - log_record["timestamp"] = timestamp - return logs_d - - # Helpers - def write( - self, - rows: typing.Iterable[bigtable.row.DirectRow], - root_ids: typing.Optional[ - typing.Union[np.uint64, typing.Iterable[np.uint64]] - ] = None, - operation_id: typing.Optional[np.uint64] = None, - slow_retry: bool = True, - block_size: int = 2000, - ): - """Writes a list of mutated rows in bulk - WARNING: If contains the same row (same row_key) and column - key two times only the last one is effectively written to the BigTable - (even when the mutations were applied to different columns) - --> no versioning! - :param rows: list - list of mutated rows - :param root_ids: list if uint64 - :param operation_id: uint64 or None - operation_id (or other unique id) that *was* used to lock the root - the bulk write is only executed if the root is still locked with - the same id. - :param slow_retry: bool - :param block_size: int - """ - if slow_retry: - initial = 5 - else: - initial = 1 - - exception_types = (Aborted, DeadlineExceeded, ServiceUnavailable) - retry = Retry( - predicate=if_exception_type(exception_types), - initial=initial, - maximum=15.0, - multiplier=2.0, - deadline=self.graph_meta.graph_config.ROOT_LOCK_EXPIRY.seconds, - ) - - if root_ids is not None and operation_id is not None: - if isinstance(root_ids, int): - root_ids = [root_ids] - if not self.renew_locks(root_ids, operation_id): - raise exceptions.LockingError( - f"Root lock renewal failed: operation {operation_id}" - ) - - for i in range(0, len(rows), block_size): - status = self._table.mutate_rows(rows[i : i + block_size], retry=retry) - if not all(status): - raise exceptions.ChunkedGraphError( - f"Bulk write failed: operation {operation_id}" - ) - - def mutate_row( - self, - row_key: bytes, - val_dict: typing.Dict[attributes._Attribute, typing.Any], - time_stamp: typing.Optional[datetime] = None, - ) -> bigtable.row.Row: - """Mutates a single row (doesn't write to big table).""" - row = self._table.direct_row(row_key) - for column, value in val_dict.items(): - row.set_cell( - column_family_id=column.family_id, - column=column.key, - value=column.serialize(value), - timestamp=time_stamp, - ) - return row - - # Locking - def lock_root( - self, - root_id: np.uint64, - operation_id: np.uint64, - ) -> bool: - """Attempts to lock the latest version of a root node.""" - lock_expiry = self.graph_meta.graph_config.ROOT_LOCK_EXPIRY - lock_column = attributes.Concurrency.Lock - indefinite_lock_column = attributes.Concurrency.IndefiniteLock - filter_ = utils.get_root_lock_filter( - lock_column, lock_expiry, indefinite_lock_column - ) - - root_row = self._table.conditional_row( - serialize_uint64(root_id), filter_=filter_ - ) - # Set row lock if condition returns no results (state == False) - root_row.set_cell( - lock_column.family_id, - lock_column.key, - serialize_uint64(operation_id), - state=False, - timestamp=get_valid_timestamp(None), - ) - - # The lock was acquired when set_cell returns False (state) - lock_acquired = not root_row.commit() - if not lock_acquired: - row = self._read_byte_row(serialize_uint64(root_id), columns=lock_column) - l_operation_ids = [cell.value for cell in row] - self.logger.debug(f"Locked operation ids: {l_operation_ids}") - return lock_acquired - - def lock_root_indefinitely( - self, - root_id: np.uint64, - operation_id: np.uint64, - ) -> bool: - """Attempts to indefinitely lock the latest version of a root node.""" - lock_column = attributes.Concurrency.IndefiniteLock - filter_ = utils.get_indefinite_root_lock_filter(lock_column) - root_row = self._table.conditional_row( - serialize_uint64(root_id), filter_=filter_ - ) - # Set row lock if condition returns no results (state == False) - root_row.set_cell( - lock_column.family_id, - lock_column.key, - serialize_uint64(operation_id), - state=False, - timestamp=get_valid_timestamp(None), - ) - - # The lock was acquired when set_cell returns False (state) - lock_acquired = not root_row.commit() - if not lock_acquired: - row = self._read_byte_row(serialize_uint64(root_id), columns=lock_column) - l_operation_ids = [cell.value for cell in row] - self.logger.debug(f"Indefinitely locked operation ids: {l_operation_ids}") - return lock_acquired - - def lock_roots( - self, - root_ids: typing.Sequence[np.uint64], - operation_id: np.uint64, - future_root_ids_d: typing.Dict, - max_tries: int = 1, - waittime_s: float = 0.5, - ) -> typing.Tuple[bool, typing.Iterable]: - """Attempts to lock multiple nodes with same operation id""" - i_try = 0 - while i_try < max_tries: - lock_acquired = False - # Collect latest root ids - new_root_ids: typing.List[np.uint64] = [] - for root_id in root_ids: - future_root_ids = future_root_ids_d[root_id] - if not future_root_ids.size: - new_root_ids.append(root_id) - else: - new_root_ids.extend(future_root_ids) - - # Attempt to lock all latest root ids - root_ids = np.unique(new_root_ids) - for root_id in root_ids: - lock_acquired = self.lock_root(root_id, operation_id) - # Roll back locks if one root cannot be locked - if not lock_acquired: - for id_ in root_ids: - self.unlock_root(id_, operation_id) - break - - if lock_acquired: - return True, root_ids - time.sleep(waittime_s) - i_try += 1 - self.logger.debug(f"Try {i_try}") - return False, root_ids - - def lock_roots_indefinitely( - self, - root_ids: typing.Sequence[np.uint64], - operation_id: np.uint64, - future_root_ids_d: typing.Dict, - ) -> typing.Tuple[bool, typing.Iterable]: - """Attempts to indefinitely lock multiple nodes with same operation id""" - lock_acquired = False - # Collect latest root ids - new_root_ids: typing.List[np.uint64] = [] - for _id in root_ids: - future_root_ids = future_root_ids_d.get(_id) - if not future_root_ids.size: - new_root_ids.append(_id) - else: - new_root_ids.extend(future_root_ids) - - # Attempt to lock all latest root ids - failed_to_lock_id = None - root_ids = np.unique(new_root_ids) - for _id in root_ids: - self.logger.debug(f"operation {operation_id} root_id {_id}") - lock_acquired = self.lock_root_indefinitely(_id, operation_id) - # Roll back locks if one root cannot be locked - if not lock_acquired: - failed_to_lock_id = _id - for id_ in root_ids: - self.unlock_indefinitely_locked_root(id_, operation_id) - break - if lock_acquired: - return True, root_ids, failed_to_lock_id - return False, root_ids, failed_to_lock_id - - def unlock_root(self, root_id: np.uint64, operation_id: np.uint64): - """Unlocks root node that is locked with operation_id.""" - lock_column = attributes.Concurrency.Lock - expiry = self.graph_meta.graph_config.ROOT_LOCK_EXPIRY - root_row = self._table.conditional_row( - serialize_uint64(root_id), - filter_=utils.get_unlock_root_filter(lock_column, expiry, operation_id), - ) - # Delete row if conditions are met (state == True) - root_row.delete_cell(lock_column.family_id, lock_column.key, state=True) - return root_row.commit() - - def unlock_indefinitely_locked_root( - self, root_id: np.uint64, operation_id: np.uint64 - ): - """Unlocks root node that is indefinitely locked with operation_id.""" - lock_column = attributes.Concurrency.IndefiniteLock - # Get conditional row using the chained filter - root_row = self._table.conditional_row( - serialize_uint64(root_id), - filter_=utils.get_indefinite_unlock_root_filter(lock_column, operation_id), - ) - # Delete row if conditions are met (state == True) - root_row.delete_cell(lock_column.family_id, lock_column.key, state=True) - return root_row.commit() - - def renew_lock(self, root_id: np.uint64, operation_id: np.uint64) -> bool: - """Renews existing root node lock with operation_id to extend time.""" - lock_column = attributes.Concurrency.Lock - root_row = self._table.conditional_row( - serialize_uint64(root_id), - filter_=utils.get_renew_lock_filter(lock_column, operation_id), - ) - # Set row lock if condition returns a result (state == True) - root_row.set_cell( - lock_column.family_id, - lock_column.key, - lock_column.serialize(operation_id), - state=False, - ) - # The lock was acquired when set_cell returns True (state) - return not root_row.commit() - - def renew_locks(self, root_ids: np.uint64, operation_id: np.uint64) -> bool: - """Renews existing root node locks with operation_id to extend time.""" - for root_id in root_ids: - if not self.renew_lock(root_id, operation_id): - self.logger.warning(f"renew_lock failed - {root_id}") - return False - return True - - def get_lock_timestamp( - self, root_id: np.uint64, operation_id: np.uint64 - ) -> typing.Union[datetime, None]: - """Lock timestamp for a Root ID operation.""" - row = self.read_node(root_id, properties=attributes.Concurrency.Lock) - if len(row) == 0: - self.logger.warning(f"No lock found for {root_id}") - return None - if row[0].value != operation_id: - self.logger.warning(f"{root_id} not locked with {operation_id}") - return None - return row[0].timestamp - - def get_consolidated_lock_timestamp( - self, - root_ids: typing.Sequence[np.uint64], - operation_ids: typing.Sequence[np.uint64], - ) -> typing.Union[datetime, None]: - """Minimum of multiple lock timestamps.""" - time_stamps = [] - for root_id, operation_id in zip(root_ids, operation_ids): - time_stamp = self.get_lock_timestamp(root_id, operation_id) - if time_stamp is None: - return None - time_stamps.append(time_stamp) - if len(time_stamps) == 0: - return None - return np.min(time_stamps) - - # IDs - def create_node_ids( - self, chunk_id: np.uint64, size: int, root_chunk=False - ) -> np.ndarray: - """Generates a list of unique node IDs for the given chunk.""" - if root_chunk: - new_ids = self._get_root_segment_ids_range(chunk_id, size) - else: - low, high = self._get_ids_range( - serialize_uint64(chunk_id, counter=True), size - ) - low, high = basetypes.SEGMENT_ID.type(low), basetypes.SEGMENT_ID.type(high) - new_ids = np.arange(low, high + np.uint64(1), dtype=basetypes.SEGMENT_ID) - return new_ids | chunk_id - - def create_node_id( - self, chunk_id: np.uint64, root_chunk=False - ) -> basetypes.NODE_ID: - """Generate a unique node ID in the chunk.""" - return self.create_node_ids(chunk_id, 1, root_chunk=root_chunk)[0] - - def get_max_node_id( - self, chunk_id: basetypes.CHUNK_ID, root_chunk=False - ) -> basetypes.NODE_ID: - """Gets the current maximum segment ID in the chunk.""" - if root_chunk: - n_counters = np.uint64(2**8) - max_value = 0 - for counter in range(n_counters): - row = self._read_byte_row( - serialize_key(f"i{pad_node_id(chunk_id)}_{counter}"), - columns=attributes.Concurrency.Counter, - ) - val = ( - basetypes.SEGMENT_ID.type(row[0].value if row else 0) * n_counters - + counter - ) - max_value = val if val > max_value else max_value - return chunk_id | basetypes.SEGMENT_ID.type(max_value) - column = attributes.Concurrency.Counter - row = self._read_byte_row( - serialize_uint64(chunk_id, counter=True), columns=column - ) - return chunk_id | basetypes.SEGMENT_ID.type(row[0].value if row else 0) - - def create_operation_id(self): - """Generate a unique operation ID.""" - return self._get_ids_range(attributes.OperationLogs.key, 1)[1] - - def get_max_operation_id(self): - """Gets the current maximum operation ID.""" - column = attributes.Concurrency.Counter - row = self._read_byte_row(attributes.OperationLogs.key, columns=column) - return row[0].value if row else column.basetype(0) - - def get_compatible_timestamp( - self, time_stamp: datetime, round_up: bool = False - ) -> datetime: - return utils.get_google_compatible_time_stamp(time_stamp, round_up=round_up) - - # PRIVATE METHODS - def _create_column_families(self): - f = self._table.column_family("0") - f.create() - f = self._table.column_family("1", gc_rule=MaxVersionsGCRule(1)) - f.create() - f = self._table.column_family("2") - f.create() - f = self._table.column_family("3") - f.create() - - def _get_ids_range(self, key: bytes, size: int) -> typing.Tuple: - """Returns a range (min, max) of IDs for a given `key`.""" - column = attributes.Concurrency.Counter - row = self._table.append_row(key) - row.increment_cell_value(column.family_id, column.key, size) - row = row.commit() - high = column.deserialize(row[column.family_id][column.key][0][0]) - return high + np.uint64(1) - size, high - - def _get_root_segment_ids_range( - self, chunk_id: basetypes.CHUNK_ID, size: int = 1, counter: int = None - ) -> np.ndarray: - """Return unique segment ID for the root chunk.""" - n_counters = np.uint64(2**8) - counter = ( - np.uint64(counter % n_counters) - if counter - else np.uint64(np.random.randint(0, n_counters)) - ) - key = serialize_key(f"i{pad_node_id(chunk_id)}_{counter}") - min_, max_ = self._get_ids_range(key=key, size=size) - return np.arange( - min_ * n_counters + counter, - max_ * n_counters + np.uint64(1) + counter, - n_counters, - dtype=basetypes.SEGMENT_ID, - ) - - def _read_byte_rows( - self, - start_key: typing.Optional[bytes] = None, - end_key: typing.Optional[bytes] = None, - end_key_inclusive: bool = False, - row_keys: typing.Optional[typing.Iterable[bytes]] = None, - columns: typing.Optional[ - typing.Union[typing.Iterable[attributes._Attribute], 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, - ) -> typing.Dict[ - bytes, - typing.Union[ - typing.Dict[attributes._Attribute, typing.List[bigtable.row_data.Cell]], - typing.List[bigtable.row_data.Cell], - ], - ]: - """Main function for reading a row range or non-contiguous row sets from Bigtable using - `bytes` keys. - - Keyword Arguments: - start_key {typing.Optional[bytes]} -- The first row to be read, ignored if `row_keys` is set. - If None, no lower boundary is used. (default: {None}) - end_key {typing.Optional[bytes]} -- The end of the row range, ignored if `row_keys` is set. - If None, no upper boundary is used. (default: {None}) - end_key_inclusive {bool} -- Whether or not `end_key` itself should be included in the - request, ignored if `row_keys` is set or `end_key` is None. (default: {False}) - row_keys {typing.Optional[typing.Iterable[bytes]]} -- An `typing.Iterable` containing possibly - non-contiguous row keys. Takes precedence over `start_key` and `end_key`. - (default: {None}) - columns {typing.Optional[typing.Union[typing.Iterable[attributes._Attribute], attributes._Attribute]]} -- - typing.Optional filtering by columns to speed up the query. If `columns` is a single - column (not iterable), the column key will be omitted from the result. - (default: {None}) - start_time {typing.Optional[datetime]} -- Ignore cells with timestamp before - `start_time`. If None, no lower bound. (default: {None}) - end_time {typing.Optional[datetime]} -- Ignore cells with timestamp after `end_time`. - If None, no upper bound. (default: {None}) - end_time_inclusive {bool} -- Whether or not `end_time` itself should be included in the - request, ignored if `end_time` is None. (default: {False}) - user_id {typing.Optional[str]} -- Only return cells with userID equal to this - - Returns: - typing.Dict[bytes, typing.Union[typing.Dict[attributes._Attribute, typing.List[bigtable.row_data.Cell]], - typing.List[bigtable.row_data.Cell]]] -- - Returns a dictionary of `byte` rows as keys. Their value will be a mapping of - columns to a typing.List of cells (one cell per timestamp). Each cell has a `value` - property, which returns the deserialized field, and a `timestamp` property, which - returns the timestamp as `datetime` object. - If only a single `attributes._Attribute` was requested, the typing.List of cells will be - attached to the row dictionary directly (skipping the column dictionary). - """ - - # Create filters: Rows - row_set = RowSet() - if row_keys is not None: - row_set.row_keys = list(row_keys) - elif start_key is not None and end_key is not None: - row_set.add_row_range_from_keys( - start_key=start_key, - start_inclusive=True, - end_key=end_key, - end_inclusive=end_key_inclusive, - ) - else: - raise exceptions.PreconditionError( - "Need to either provide a valid set of rows, or" - " both, a start row and an end row." - ) - filter_ = utils.get_time_range_and_column_filter( - columns=columns, - start_time=start_time, - end_time=end_time, - end_inclusive=end_time_inclusive, - user_id=user_id, - ) - # Bigtable read with retries - rows = self._read(row_set=row_set, row_filter=filter_) - - # Deserialize cells - for row_key, column_dict in rows.items(): - for column, cell_entries in column_dict.items(): - for cell_entry in cell_entries: - cell_entry.value = column.deserialize(cell_entry.value) - # If no column array was requested, reattach single column's values directly to the row - if isinstance(columns, attributes._Attribute): - rows[row_key] = cell_entries - return rows - - def _read_byte_row( - self, - row_key: bytes, - columns: typing.Optional[ - typing.Union[typing.Iterable[attributes._Attribute], attributes._Attribute] - ] = None, - start_time: typing.Optional[datetime] = None, - end_time: typing.Optional[datetime] = None, - end_time_inclusive: bool = False, - ) -> typing.Union[ - typing.Dict[attributes._Attribute, typing.List[bigtable.row_data.Cell]], - typing.List[bigtable.row_data.Cell], - ]: - """Convenience function for reading a single row from Bigtable using its `bytes` keys. - - Arguments: - row_key {bytes} -- The row to be read. - - Keyword Arguments: - columns {typing.Optional[typing.Union[typing.Iterable[attributes._Attribute], attributes._Attribute]]} -- - typing.Optional filtering by columns to speed up the query. If `columns` is a single - column (not iterable), the column key will be omitted from the result. - (default: {None}) - start_time {typing.Optional[datetime]} -- Ignore cells with timestamp before - `start_time`. If None, no lower bound. (default: {None}) - end_time {typing.Optional[datetime]} -- Ignore cells with timestamp after `end_time`. - If None, no upper bound. (default: {None}) - end_time_inclusive {bool} -- Whether or not `end_time` itself should be included in the - request, ignored if `end_time` is None. (default: {False}) - - Returns: - typing.Union[typing.Dict[attributes._Attribute, typing.List[bigtable.row_data.Cell]], - typing.List[bigtable.row_data.Cell]] -- - Returns a mapping of columns to a typing.List of cells (one cell per timestamp). Each cell - has a `value` property, which returns the deserialized field, and a `timestamp` - property, which returns the timestamp as `datetime` object. - If only a single `attributes._Attribute` was requested, the typing.List of cells is returned - directly. - """ - row = self._read_byte_rows( - row_keys=[row_key], - columns=columns, - start_time=start_time, - end_time=end_time, - end_time_inclusive=end_time_inclusive, - ) - return ( - row.get(row_key, []) - if isinstance(columns, attributes._Attribute) - else row.get(row_key, {}) - ) - - def _execute_read_thread(self, args: typing.Tuple[Table, RowSet, RowFilter]): - table, row_set, row_filter = args - if not row_set.row_keys and not row_set.row_ranges: - # Check for everything falsy, because Bigtable considers even empty - # lists of row_keys as no upper/lower bound! - return {} - range_read = table.read_rows(row_set=row_set, filter_=row_filter) - res = {v.row_key: utils.partial_row_data_to_column_dict(v) for v in range_read} - return res - - def _read( - self, row_set: RowSet, row_filter: RowFilter = None - ) -> typing.Dict[bytes, typing.Dict[attributes._Attribute, PartialRowData]]: - """Core function to read rows from Bigtable. Uses standard Bigtable retry logic - :param row_set: BigTable RowSet - :param row_filter: BigTable RowFilter - :return: typing.Dict[bytes, typing.Dict[attributes._Attribute, bigtable.row_data.PartialRowData]] - """ - # FIXME: Bigtable limits the length of the serialized request to 512 KiB. We should - # calculate this properly (range_read.request.SerializeToString()), but this estimate is - # good enough for now - - n_subrequests = max( - 1, int(np.ceil(len(row_set.row_keys) / self._max_row_key_count)) - ) - n_threads = min(n_subrequests, 2 * mu.n_cpus) - - row_sets = [] - for i in range(n_subrequests): - r = RowSet() - r.row_keys = row_set.row_keys[ - i * self._max_row_key_count : (i + 1) * self._max_row_key_count - ] - row_sets.append(r) - - # Don't forget the original RowSet's row_ranges - row_sets[0].row_ranges = row_set.row_ranges - responses = mu.multithread_func( - self._execute_read_thread, - params=((self._table, r, row_filter) for r in row_sets), - debug=n_threads == 1, - n_threads=n_threads, - ) - - combined_response = {} - for resp in responses: - combined_response.update(resp) - return combined_response diff --git a/pychunkedgraph/graph/client/bigtable/utils.py b/pychunkedgraph/graph/client/bigtable/utils.py deleted file mode 100644 index 2d30eeb32..000000000 --- a/pychunkedgraph/graph/client/bigtable/utils.py +++ /dev/null @@ -1,304 +0,0 @@ -from typing import Dict -from typing import Union -from typing import Iterable -from typing import Optional -from datetime import datetime -from datetime import timedelta - -import numpy as np -from google.cloud.bigtable.row_data import PartialRowData -from google.cloud.bigtable.row_filters import RowFilter -from google.cloud.bigtable.row_filters import PassAllFilter -from google.cloud.bigtable.row_filters import BlockAllFilter -from google.cloud.bigtable.row_filters import TimestampRange -from google.cloud.bigtable.row_filters import RowFilterChain -from google.cloud.bigtable.row_filters import RowFilterUnion -from google.cloud.bigtable.row_filters import ValueRangeFilter -from google.cloud.bigtable.row_filters import CellsRowLimitFilter -from google.cloud.bigtable.row_filters import ColumnRangeFilter -from google.cloud.bigtable.row_filters import TimestampRangeFilter -from google.cloud.bigtable.row_filters import ConditionalRowFilter -from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter - -from ... import attributes - - -def partial_row_data_to_column_dict( - partial_row_data: PartialRowData, -) -> Dict[attributes._Attribute, PartialRowData]: - new_column_dict = {} - for family_id, column_dict in partial_row_data._cells.items(): - for column_key, column_values in column_dict.items(): - column = attributes.from_key(family_id, column_key) - new_column_dict[column] = column_values - return new_column_dict - - -def get_google_compatible_time_stamp( - time_stamp: datetime, round_up: bool = False -) -> datetime: - """ - Makes a datetime time stamp compatible with googles' services. - Google restricts the accuracy of time stamps to milliseconds. Hence, the - microseconds are cut of. By default, time stamps are rounded to the lower - number. - """ - micro_s_gap = timedelta(microseconds=time_stamp.microsecond % 1000) - if micro_s_gap == 0: - return time_stamp - if round_up: - time_stamp += timedelta(microseconds=1000) - micro_s_gap - else: - time_stamp -= micro_s_gap - return time_stamp - - -def _get_column_filter( - columns: Union[Iterable[attributes._Attribute], attributes._Attribute] = None -) -> RowFilter: - """Generates a RowFilter that accepts the specified columns""" - if isinstance(columns, attributes._Attribute): - return ColumnRangeFilter( - columns.family_id, start_column=columns.key, end_column=columns.key - ) - elif len(columns) == 1: - return ColumnRangeFilter( - columns[0].family_id, start_column=columns[0].key, end_column=columns[0].key - ) - return RowFilterUnion( - [ - ColumnRangeFilter(col.family_id, start_column=col.key, end_column=col.key) - for col in columns - ] - ) - - -def _get_user_filter(user_id: str): - """generates a ColumnRegEx Filter which filters user ids - - Args: - user_id (str): userID to select for - """ - - condition = RowFilterChain( - [ - ColumnQualifierRegexFilter(attributes.OperationLogs.UserID.key), - ValueRangeFilter(str.encode(user_id), str.encode(user_id)), - CellsRowLimitFilter(1), - ] - ) - - conditional_filter = ConditionalRowFilter( - base_filter=condition, - true_filter=PassAllFilter(True), - false_filter=BlockAllFilter(True), - ) - return conditional_filter - - -def _get_time_range_filter( - start_time: Optional[datetime] = None, - end_time: Optional[datetime] = None, - end_inclusive: bool = True, -) -> RowFilter: - """Generates a TimeStampRangeFilter which is inclusive for start and (optionally) end. - - :param start: - :param end: - :return: - """ - # Comply to resolution of BigTables TimeRange - if start_time is not None: - start_time = get_google_compatible_time_stamp(start_time, round_up=False) - if end_time is not None: - end_time = get_google_compatible_time_stamp(end_time, round_up=end_inclusive) - return TimestampRangeFilter(TimestampRange(start=start_time, end=end_time)) - - -def get_time_range_and_column_filter( - columns: Optional[ - Union[Iterable[attributes._Attribute], attributes._Attribute] - ] = None, - start_time: Optional[datetime] = None, - end_time: Optional[datetime] = None, - end_inclusive: bool = False, - user_id: Optional[str] = None, -) -> RowFilter: - time_filter = _get_time_range_filter( - start_time=start_time, end_time=end_time, end_inclusive=end_inclusive - ) - filters = [time_filter] - if columns is not None: - if len(columns) == 0: - raise ValueError( - f"Empty column filter {columns} is ambiguous. Pass `None` if no column filter should be applied." - ) - column_filter = _get_column_filter(columns) - filters = [column_filter, time_filter] - if user_id is not None: - user_filter = _get_user_filter(user_id=user_id) - filters.append(user_filter) - if len(filters) > 1: - return RowFilterChain(filters) - return filters[0] - - -def get_root_lock_filter( - lock_column, lock_expiry, indefinite_lock_column -) -> ConditionalRowFilter: - time_cutoff = datetime.utcnow() - lock_expiry - # Comply to resolution of BigTables TimeRange - time_cutoff -= timedelta(microseconds=time_cutoff.microsecond % 1000) - time_filter = TimestampRangeFilter(TimestampRange(start=time_cutoff)) - - # Build a column filter which tests if a lock was set (== lock column - # exists) and if it is still valid (timestamp younger than - # LOCK_EXPIRED_TIME_DELTA) and if there is no new parent (== new_parents - # exists) - lock_key_filter = ColumnRangeFilter( - column_family_id=lock_column.family_id, - start_column=lock_column.key, - end_column=lock_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - indefinite_lock_key_filter = ColumnRangeFilter( - column_family_id=indefinite_lock_column.family_id, - start_column=indefinite_lock_column.key, - end_column=indefinite_lock_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - new_parents_column = attributes.Hierarchy.NewParent - new_parents_key_filter = ColumnRangeFilter( - column_family_id=new_parents_column.family_id, - start_column=new_parents_column.key, - end_column=new_parents_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - temporal_lock_filter = RowFilterChain([time_filter, lock_key_filter]) - return ConditionalRowFilter( - base_filter=RowFilterUnion([indefinite_lock_key_filter, temporal_lock_filter]), - true_filter=PassAllFilter(True), - false_filter=new_parents_key_filter, - ) - - -def get_indefinite_root_lock_filter(lock_column) -> ConditionalRowFilter: - lock_key_filter = ColumnRangeFilter( - column_family_id=lock_column.family_id, - start_column=lock_column.key, - end_column=lock_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - new_parents_column = attributes.Hierarchy.NewParent - new_parents_key_filter = ColumnRangeFilter( - column_family_id=new_parents_column.family_id, - start_column=new_parents_column.key, - end_column=new_parents_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - return ConditionalRowFilter( - base_filter=lock_key_filter, - true_filter=PassAllFilter(True), - false_filter=new_parents_key_filter, - ) - - -def get_renew_lock_filter( - lock_column: attributes._Attribute, operation_id: np.uint64 -) -> ConditionalRowFilter: - new_parents_column = attributes.Hierarchy.NewParent - operation_id_b = lock_column.serialize(operation_id) - - # Build a column filter which tests if a lock was set (== lock column - # exists) and if the given operation_id is still the active lock holder - # and there is no new parent (== new_parents column exists). The latter - # is not necessary but we include it as a backup to prevent things - # from going really bad. - - column_key_filter = ColumnRangeFilter( - column_family_id=lock_column.family_id, - start_column=lock_column.key, - end_column=lock_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - value_filter = ValueRangeFilter( - start_value=operation_id_b, - end_value=operation_id_b, - inclusive_start=True, - inclusive_end=True, - ) - - new_parents_key_filter = ColumnRangeFilter( - column_family_id=new_parents_column.family_id, - start_column=new_parents_column.key, - end_column=new_parents_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - return ConditionalRowFilter( - base_filter=RowFilterChain([column_key_filter, value_filter]), - true_filter=new_parents_key_filter, - false_filter=PassAllFilter(True), - ) - - -def get_unlock_root_filter(lock_column, lock_expiry, operation_id) -> RowFilterChain: - time_cutoff = datetime.utcnow() - lock_expiry - # Comply to resolution of BigTables TimeRange - time_cutoff -= timedelta(microseconds=time_cutoff.microsecond % 1000) - time_filter = TimestampRangeFilter(TimestampRange(start=time_cutoff)) - - # Build a column filter which tests if a lock was set (== lock column - # exists) and if it is still valid (timestamp younger than - # LOCK_EXPIRED_TIME_DELTA) and if the given operation_id is still - # the active lock holder - column_key_filter = ColumnRangeFilter( - column_family_id=lock_column.family_id, - start_column=lock_column.key, - end_column=lock_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - value_filter = ValueRangeFilter( - start_value=lock_column.serialize(operation_id), - end_value=lock_column.serialize(operation_id), - inclusive_start=True, - inclusive_end=True, - ) - - # Chain these filters together - return RowFilterChain([time_filter, column_key_filter, value_filter]) - - -def get_indefinite_unlock_root_filter(lock_column, operation_id) -> RowFilterChain: - column_key_filter = ColumnRangeFilter( - column_family_id=lock_column.family_id, - start_column=lock_column.key, - end_column=lock_column.key, - inclusive_start=True, - inclusive_end=True, - ) - - value_filter = ValueRangeFilter( - start_value=lock_column.serialize(operation_id), - end_value=lock_column.serialize(operation_id), - inclusive_start=True, - inclusive_end=True, - ) - - # Chain these filters together - return RowFilterChain([column_key_filter, value_filter]) diff --git a/pychunkedgraph/graph/client/utils.py b/pychunkedgraph/graph/client/utils.py deleted file mode 100644 index 12eebec82..000000000 --- a/pychunkedgraph/graph/client/utils.py +++ /dev/null @@ -1,3 +0,0 @@ -""" -Common client util functions -""" \ No newline at end of file diff --git a/pychunkedgraph/graph/connectivity/cross_edges.py b/pychunkedgraph/graph/connectivity/cross_edges.py deleted file mode 100644 index 8aa52a9f1..000000000 --- a/pychunkedgraph/graph/connectivity/cross_edges.py +++ /dev/null @@ -1,219 +0,0 @@ -import time -import math -import multiprocessing as mp -from collections import defaultdict -from typing import Optional -from typing import Sequence -from typing import List -from typing import Dict - -import numpy as np -from multiwrapper.multiprocessing_utils import multiprocess_func - -from .. import attributes -from ..types import empty_2d -from ..utils import basetypes -from ..utils import serializers -from ..chunkedgraph import ChunkedGraph -from ..utils.generic import get_valid_timestamp -from ..utils.generic import filter_failed_node_ids -from ..chunks.atomic import get_touching_atomic_chunks -from ..chunks.atomic import get_bounding_atomic_chunks -from ...utils.general import chunked - - -def get_children_chunk_cross_edges( - cg, layer, chunk_coord, *, use_threads=True -) -> np.ndarray: - """ - Cross edges that connect children chunks. - The edges are between node IDs in the given layer (not atomic). - """ - atomic_chunks = get_touching_atomic_chunks(cg.meta, layer, chunk_coord) - if not len(atomic_chunks): - return [] - - print(f"touching atomic chunk count {len(atomic_chunks)}") - if not use_threads: - return _get_children_chunk_cross_edges(cg, atomic_chunks, layer - 1) - - print("get_children_chunk_cross_edges, atomic chunks", len(atomic_chunks)) - with mp.Manager() as manager: - edge_ids_shared = manager.list() - edge_ids_shared.append(empty_2d) - - task_size = int(math.ceil(len(atomic_chunks) / mp.cpu_count() / 10)) - chunked_l2chunk_list = chunked(atomic_chunks, task_size) - multi_args = [] - for atomic_chunks in chunked_l2chunk_list: - multi_args.append( - (edge_ids_shared, cg.get_serialized_info(), atomic_chunks, layer - 1) - ) - - multiprocess_func( - _get_children_chunk_cross_edges_helper, - multi_args, - n_threads=min(len(multi_args), mp.cpu_count()), - ) - - cross_edges = np.concatenate(edge_ids_shared) - if cross_edges.size: - return np.unique(cross_edges, axis=0) - return cross_edges - - -def _get_children_chunk_cross_edges_helper(args) -> None: - edge_ids_shared, cg_info, atomic_chunks, layer = args - cg = ChunkedGraph(**cg_info) - edge_ids_shared.append(_get_children_chunk_cross_edges(cg, atomic_chunks, layer)) - - -def _get_children_chunk_cross_edges(cg, atomic_chunks, layer) -> None: - print( - f"_get_children_chunk_cross_edges {layer} atomic_chunks count {len(atomic_chunks)}" - ) - cross_edges = [empty_2d] - for layer2_chunk in atomic_chunks: - edges = _read_atomic_chunk_cross_edges(cg, layer2_chunk, layer) - cross_edges.append(edges) - - cross_edges = np.concatenate(cross_edges) - if not cross_edges.size: - return empty_2d - print(f"getting roots at stop_layer {layer} {cross_edges.shape}") - cross_edges[:, 0] = cg.get_roots(cross_edges[:, 0], stop_layer=layer, ceil=False) - cross_edges[:, 1] = cg.get_roots(cross_edges[:, 1], stop_layer=layer, ceil=False) - result = np.unique(cross_edges, axis=0) if cross_edges.size else empty_2d - print(f"_get_children_chunk_cross_edges done {result.shape}") - return result - - -def _read_atomic_chunk_cross_edges( - cg, chunk_coord: Sequence[int], cross_edge_layer: int -) -> np.ndarray: - cross_edge_col = attributes.Connectivity.CrossChunkEdge[cross_edge_layer] - range_read, l2ids = _read_atomic_chunk(cg, chunk_coord, [cross_edge_layer]) - - parent_neighboring_chunk_supervoxels_d = defaultdict(list) - for l2id in l2ids: - if not cross_edge_col in range_read[l2id]: - continue - edges = range_read[l2id][cross_edge_col][0].value - parent_neighboring_chunk_supervoxels_d[l2id] = edges[:, 1] - - cross_edges = [empty_2d] - for l2id in parent_neighboring_chunk_supervoxels_d: - nebor_svs = parent_neighboring_chunk_supervoxels_d[l2id] - chunk_parent_ids = np.array([l2id] * len(nebor_svs), dtype=basetypes.NODE_ID) - cross_edges.append(np.vstack([chunk_parent_ids, nebor_svs]).T) - cross_edges = np.concatenate(cross_edges) - return cross_edges - - -def get_chunk_nodes_cross_edge_layer( - cg, layer: int, chunk_coord: Sequence[int], use_threads=True -) -> Dict: - """ - gets nodes in a chunk that are part of cross chunk edges - return_type dict {node_id: layer} - the lowest layer (>= current layer) at which a node_id is part of a cross edge - """ - print("get_bounding_atomic_chunks") - atomic_chunks = get_bounding_atomic_chunks(cg.meta, layer, chunk_coord) - print("get_bounding_atomic_chunks complete") - if not len(atomic_chunks): - return {} - - if not use_threads: - return _get_chunk_nodes_cross_edge_layer(cg, atomic_chunks, layer) - - print("divide tasks") - cg_info = cg.get_serialized_info() - manager = mp.Manager() - ids_l_shared = manager.list() - layers_l_shared = manager.list() - task_size = int(math.ceil(len(atomic_chunks) / mp.cpu_count() / 10)) - chunked_l2chunk_list = chunked(atomic_chunks, task_size) - multi_args = [] - for atomic_chunks in chunked_l2chunk_list: - multi_args.append( - (ids_l_shared, layers_l_shared, cg_info, atomic_chunks, layer) - ) - print("divide tasks complete") - - multiprocess_func( - _get_chunk_nodes_cross_edge_layer_helper, - multi_args, - n_threads=min(len(multi_args), mp.cpu_count()), - ) - - node_layer_d_shared = manager.dict() - _find_min_layer(node_layer_d_shared, ids_l_shared, layers_l_shared) - print("_find_min_layer complete") - return node_layer_d_shared - - -def _get_chunk_nodes_cross_edge_layer_helper(args): - ids_l_shared, layers_l_shared, cg_info, atomic_chunks, layer = args - cg = ChunkedGraph(**cg_info) - node_layer_d = _get_chunk_nodes_cross_edge_layer(cg, atomic_chunks, layer) - ids_l_shared.append(np.fromiter(node_layer_d.keys(), dtype=basetypes.NODE_ID)) - layers_l_shared.append(np.fromiter(node_layer_d.values(), dtype=np.uint8)) - - -def _get_chunk_nodes_cross_edge_layer(cg, atomic_chunks, layer): - atomic_node_layer_d = {} - for atomic_chunk in atomic_chunks: - chunk_node_layer_d = _read_atomic_chunk_cross_edge_nodes( - cg, atomic_chunk, range(layer, cg.meta.layer_count + 1) - ) - atomic_node_layer_d.update(chunk_node_layer_d) - - l2ids = np.fromiter(atomic_node_layer_d.keys(), dtype=basetypes.NODE_ID) - parents = cg.get_roots(l2ids, stop_layer=layer - 1, ceil=False) - layers = np.fromiter(atomic_node_layer_d.values(), dtype=int) - - node_layer_d = defaultdict(lambda: cg.meta.layer_count) - for i, parent in enumerate(parents): - node_layer_d[parent] = min(node_layer_d[parent], layers[i]) - return node_layer_d - - -def _read_atomic_chunk_cross_edge_nodes(cg, chunk_coord, cross_edge_layers): - node_layer_d = {} - range_read, l2ids = _read_atomic_chunk(cg, chunk_coord, cross_edge_layers) - for l2id in l2ids: - for layer in cross_edge_layers: - if attributes.Connectivity.CrossChunkEdge[layer] in range_read[l2id]: - node_layer_d[l2id] = layer - break - return node_layer_d - - -def _find_min_layer(node_layer_d_shared, ids_l_shared, layers_l_shared): - node_ids = np.concatenate(ids_l_shared) - layers = np.concatenate(layers_l_shared) - for i, node_id in enumerate(node_ids): - layer = node_layer_d_shared.get(node_id, layers[i]) - node_layer_d_shared[node_id] = min(layer, layers[i]) - - -def _read_atomic_chunk(cg, chunk_coord, layers): - x, y, z = chunk_coord - child_col = attributes.Hierarchy.Child - range_read = cg.range_read_chunk( - cg.get_chunk_id(layer=2, x=x, y=y, z=z), - properties=[child_col] - + [attributes.Connectivity.CrossChunkEdge[l] for l in layers], - ) - - row_ids = [] - max_children_ids = [] - for row_id, row_data in range_read.items(): - row_ids.append(row_id) - max_children_ids.append(np.max(row_data[child_col][0].value)) - - row_ids = np.array(row_ids, dtype=basetypes.NODE_ID) - segment_ids = np.array([cg.get_segment_id(r_id) for r_id in row_ids]) - l2ids = filter_failed_node_ids(row_ids, segment_ids, max_children_ids) - return range_read, l2ids diff --git a/pychunkedgraph/graph/connectivity/search.py b/pychunkedgraph/graph/connectivity/search.py deleted file mode 100644 index bd3faf227..000000000 --- a/pychunkedgraph/graph/connectivity/search.py +++ /dev/null @@ -1,47 +0,0 @@ -import random -from typing import List - -import numpy as np -from graph_tool.search import bfs_search -from graph_tool.search import BFSVisitor -from graph_tool.search import StopSearch - -from ..utils.basetypes import NODE_ID - - -class TargetVisitor(BFSVisitor): - def __init__(self, target, reachable): - self.target = target - self.reachable = reachable - - def discover_vertex(self, u): - if u == self.target: - self.reachable[u] = 1 - raise StopSearch - - -def check_reachability(g, sv1s: np.ndarray, sv2s: np.ndarray, original_ids: np.ndarray) -> np.ndarray: - """ - g: graph tool Graph instance with ids 0 to N-1 where N = vertex count - original_ids: sorted ChunkedGraph supervoxel ids - (to identify corresponding ids in graph tool) - for each pair (sv1, sv2) check if a path exists (BFS) - """ - # mapping from original ids to graph tool ids - original_ids_d = { - sv_id: index for sv_id, index in zip(original_ids, range(len(original_ids))) - } - reachable = g.new_vertex_property("int", val=0) - - def _check_reachability(source, target): - bfs_search(g, source, TargetVisitor(target, reachable)) - return reachable[target] - - return np.array( - [ - _check_reachability(original_ids_d[source], original_ids_d[target]) - for source, target in zip(sv1s, sv2s) - ], - dtype=bool, - ) - diff --git a/pychunkedgraph/graph/cutting.py b/pychunkedgraph/graph/cutting.py index 8b1583871..ffbee9936 100644 --- a/pychunkedgraph/graph/cutting.py +++ b/pychunkedgraph/graph/cutting.py @@ -1,28 +1,57 @@ -import collections import fastremap import numpy as np import itertools -import logging import time -import graph_tool -import graph_tool.flow -from typing import Dict -from typing import Tuple -from typing import Optional +from dataclasses import dataclass +from typing import Tuple, Union from typing import Sequence from typing import Iterable from .utils import flatgraph -from .utils import basetypes -from .utils.generic import get_bounding_box +from pychunkedgraph.graph import basetypes +from .utils.generic import get_bounding_box, assert_same_root from .edges import Edges -from .exceptions import PreconditionError +from .exceptions import PreconditionError, SupervoxelSplitRequiredError from .exceptions import PostconditionError DEBUG_MODE = False +@dataclass +class Cut: + """Multicut produced a clean partition — these SV-pair edges are to be cut.""" + + atomic_edges: np.ndarray # shape (N, 2) + + +@dataclass +class PreviewCut: + """Multicut in preview mode — connected components after the proposed cut. + + `illegal_split` flags cases where the cut isolates source or sink. + """ + + supervoxel_ccs: list + illegal_split: bool + + +@dataclass +class SvSplitRequired: + """Multicut could not partition without first splitting a supervoxel. + + Carries the cross-chunk-representative remapping the caller needs to + run the actual SV split. Returned (not raised) from run_multicut; the + SupervoxelSplitRequiredError that surfaces this condition is caught + inside run_multicut and never escapes as control flow. + """ + + sv_remapping: dict # old_sv_id -> rep_sv_id + + +MulticutResult = Union[Cut, PreviewCut, SvSplitRequired] + + class IsolatingCutException(Exception): """Raised when mincut would split off one of the labeled supervoxel exactly. This is used to trigger a PostconditionError with a custom message. @@ -62,7 +91,7 @@ def merge_cross_chunk_edges_graph_tool( if len(mapping) > 0: mapping = np.concatenate(mapping) u_nodes = np.unique(edges) - u_unmapped_nodes = u_nodes[~np.in1d(u_nodes, mapping)] + u_unmapped_nodes = u_nodes[~np.isin(u_nodes, mapping)] unmapped_mapping = np.concatenate( [u_unmapped_nodes.reshape(-1, 1), u_unmapped_nodes.reshape(-1, 1)], axis=1 ) @@ -95,6 +124,7 @@ def __init__( split_preview=False, path_augment=True, disallow_isolating_cut=True, + sv_split_supported=False, logger=None, ): self.cg_edges = cg_edges @@ -102,6 +132,7 @@ def __init__( self.logger = logger self.path_augment = path_augment self.disallow_isolating_cut = disallow_isolating_cut + self.sv_split_supported = sv_split_supported time_start = time.time() @@ -116,6 +147,10 @@ def __init__( self.cross_chunk_edge_remapping, ) = merge_cross_chunk_edges_graph_tool(cg_edges, cg_affs) + # save this representative mapping for supervoxel splitting + # passed along with SupervoxelSplitRequiredError + self.sv_remapping = dict(complete_mapping) + dt = time.time() - time_start if logger is not None: logger.debug("Cross edge merging: %.2fms" % (dt * 1000)) @@ -127,9 +162,12 @@ def __init__( ) if len(cross_chunk_edge_mapping) > 0: - assert ( - np.unique(cross_chunk_edge_mapping[:, 0], return_counts=True)[1].max() - == 1 + unique_src, counts = np.unique( + cross_chunk_edge_mapping[:, 0], return_counts=True + ) + assert counts.max() == 1, ( + f"cross_chunk_edge_mapping[:, 0] must be unique; " + f"duplicates={unique_src[counts > 1].tolist()}" ) # Map cg sources and sinks with the cross chunk edge mapping @@ -140,6 +178,18 @@ def __init__( np.array(cg_sinks), complete_mapping[:, 0], complete_mapping[:, 1] ) + # Detect source/sink overlap after cross-chunk remapping + # (both sides mapped to the same representative → need SV split) + overlap = np.intersect1d(self.sources, self.sinks) + if len(overlap) > 0: + msg = ( + "Source and sink supervoxels share a cross-chunk edge representative. " + "A supervoxel split is required." + ) + if self.sv_split_supported: + raise SupervoxelSplitRequiredError(msg, self.sv_remapping) + raise PreconditionError(msg) + self._build_gt_graph(mapped_edges, mapped_affs) self.source_path_vertices = self.source_graph_ids @@ -161,9 +211,18 @@ def _build_gt_graph(self, edges, affs): self.weighted_graph_raw, self.capacities_raw, self.gt_edges_raw, - _, + self.unique_supervoxel_ids_raw, ) = flatgraph.build_gt_graph(edges, affs, make_directed=True) + # Compute vertex indices valid for the raw graph + # (these differ from source_graph_ids/sink_graph_ids which are for weighted_graph) + self.source_graph_ids_raw = np.where( + np.isin(self.unique_supervoxel_ids_raw, self.sources) + )[0] + self.sink_graph_ids_raw = np.where( + np.isin(self.unique_supervoxel_ids_raw, self.sinks) + )[0] + self.source_edges = list(itertools.product(self.sources, self.sources)) self.sink_edges = list(itertools.product(self.sinks, self.sinks)) @@ -189,9 +248,9 @@ def _build_gt_graph(self, edges, affs): ) = flatgraph.build_gt_graph(comb_edges, comb_affs, make_directed=True) self.source_graph_ids = np.where( - np.in1d(self.unique_supervoxel_ids, self.sources) + np.isin(self.unique_supervoxel_ids, self.sources) )[0] - self.sink_graph_ids = np.where(np.in1d(self.unique_supervoxel_ids, self.sinks))[ + self.sink_graph_ids = np.where(np.isin(self.unique_supervoxel_ids, self.sinks))[ 0 ] @@ -201,6 +260,8 @@ def _build_gt_graph(self, edges, affs): def _compute_mincut_direct(self): """Uses additional edges directly between source/sink points.""" + from .utils import _graph_tool as graph_tool + self._filter_graph_connected_components() src, tgt = ( self.weighted_graph.vertex(self.source_graph_ids[0]), @@ -223,20 +284,23 @@ def _augment_mincut_capacity(self): paths_v_s, paths_e_s, invaff_s = flatgraph.compute_filtered_paths( self.weighted_graph_raw, self.capacities_raw, - self.source_graph_ids, - self.sink_graph_ids, + self.source_graph_ids_raw, + self.sink_graph_ids_raw, ) paths_v_y, paths_e_y, invaff_y = flatgraph.compute_filtered_paths( self.weighted_graph_raw, self.capacities_raw, - self.sink_graph_ids, - self.source_graph_ids, + self.sink_graph_ids_raw, + self.source_graph_ids_raw, ) except AssertionError: - raise PreconditionError( + msg = ( "Paths between source or sink points irreparably overlap other labels from other side. " "Check that labels are correct and consider spreading points out farther." ) + if self.sv_split_supported: + raise SupervoxelSplitRequiredError(msg, self.sv_remapping) + raise PreconditionError(msg) paths_e_s_no, paths_e_y_no, do_check = flatgraph.remove_overlapping_edges( paths_v_s, paths_e_s, paths_v_y, paths_e_y @@ -294,7 +358,7 @@ def rerun_paths_without_overlap( _, paths_e_y_no, _ = flatgraph.compute_filtered_paths( self.weighted_graph_raw, self.capacities_raw, - self.sink_graph_ids, + self.sink_graph_ids_raw, omit_verts, ) @@ -303,7 +367,7 @@ def rerun_paths_without_overlap( _, paths_e_s_no, _ = flatgraph.compute_filtered_paths( self.weighted_graph_raw, self.capacities_raw, - self.source_graph_ids, + self.source_graph_ids_raw, omit_verts, ) paths_e_y_no = paths_e_y @@ -326,11 +390,13 @@ def rerun_paths_without_overlap( def _compute_mincut_path_augmented(self): """Compute mincut using edges found from a shortest-path search.""" + from .utils import _graph_tool as graph_tool + adj_capacity = self._augment_mincut_capacity() gr = self.weighted_graph_raw - src, tgt = gr.vertex(self.source_graph_ids[0]), gr.vertex( - self.sink_graph_ids[0] + src, tgt = gr.vertex(self.source_graph_ids_raw[0]), gr.vertex( + self.sink_graph_ids_raw[0] ) residuals = graph_tool.flow.boykov_kolmogorov_max_flow( @@ -347,7 +413,11 @@ def compute_mincut(self): time_start = time.time() - if self.path_augment: + if ( + self.path_augment + and len(self.source_graph_ids_raw) > 0 + and len(self.sink_graph_ids_raw) > 0 + ): partition = self._compute_mincut_path_augmented() else: partition = self._compute_mincut_direct() @@ -398,7 +468,9 @@ def _remap_cut_edge_set(self, cut_edge_set): remapped_cutset_flattened_view = remapped_cutset.view(dtype="u8,u8") edges_flattened_view = self.cg_edges.view(dtype="u8,u8") - cutset_mask = np.in1d(remapped_cutset_flattened_view, edges_flattened_view) + cutset_mask = np.isin( + remapped_cutset_flattened_view, edges_flattened_view + ).ravel() return remapped_cutset[cutset_mask] @@ -432,8 +504,8 @@ def _get_split_preview_connected_components(self, cut_edge_set): max_sinks = 0 i = 0 for cc in ccs_test_post_cut: - num_sources = np.count_nonzero(np.in1d(self.source_graph_ids, cc)) - num_sinks = np.count_nonzero(np.in1d(self.sink_graph_ids, cc)) + num_sources = np.count_nonzero(np.isin(self.source_graph_ids, cc)) + num_sinks = np.count_nonzero(np.isin(self.sink_graph_ids, cc)) if num_sources > max_sources: max_sources = num_sources max_source_index = i @@ -477,6 +549,8 @@ def _filter_graph_connected_components(self): Filter out connected components in the graph that are not involved in the local mincut """ + from .utils import _graph_tool as graph_tool + ccs = flatgraph.connected_components(self.weighted_graph) removed = self.weighted_graph.new_vertex_property("bool") @@ -486,13 +560,15 @@ def _filter_graph_connected_components(self): # If connected component contains no sources or no sinks, # remove its nodes from the mincut computation if not ( - np.any(np.in1d(self.source_graph_ids, cc)) - and np.any(np.in1d(self.sink_graph_ids, cc)) + np.any(np.isin(self.source_graph_ids, cc)) + and np.any(np.isin(self.sink_graph_ids, cc)) ): for node_id in cc: removed[node_id] = True - self.weighted_graph.set_vertex_filter(removed, inverted=True) + keep = self.weighted_graph.new_vertex_property("bool") + keep.a = ~removed.a.astype(bool) + self.weighted_graph.set_vertex_filter(keep) pruned_graph = graph_tool.Graph(self.weighted_graph, prune=True) # Test that there is only one connected component left ccs = flatgraph.connected_components(pruned_graph) @@ -525,13 +601,13 @@ def _gt_mincut_sanity_check(self, partition): np.array(np.where(partition.a == i_cc)[0], dtype=int) ] - if np.any(np.in1d(self.sources, cc_list)): - assert np.all(np.in1d(self.sources, cc_list)) - assert ~np.any(np.in1d(self.sinks, cc_list)) + if np.any(np.isin(self.sources, cc_list)): + assert np.all(np.isin(self.sources, cc_list)) + assert ~np.any(np.isin(self.sinks, cc_list)) - if np.any(np.in1d(self.sinks, cc_list)): - assert np.all(np.in1d(self.sinks, cc_list)) - assert ~np.any(np.in1d(self.sources, cc_list)) + if np.any(np.isin(self.sinks, cc_list)): + assert np.all(np.isin(self.sinks, cc_list)) + assert ~np.any(np.isin(self.sources, cc_list)) def _sink_and_source_connectivity_sanity_check(self, cut_edge_set): """ @@ -547,7 +623,8 @@ def _sink_and_source_connectivity_sanity_check(self, cut_edge_set): for edge_to_remove in parallel_edges: self.edges_to_remove[edge_to_remove] = True - self.weighted_graph.set_edge_filter(self.edges_to_remove, True) + self.edges_to_remove.a = ~self.edges_to_remove.a.astype(bool) + self.weighted_graph.set_edge_filter(self.edges_to_remove) ccs_test_post_cut = flatgraph.connected_components(self.weighted_graph) # Make sure sinks and sources are among each other and not in different sets @@ -555,9 +632,9 @@ def _sink_and_source_connectivity_sanity_check(self, cut_edge_set): illegal_split = False try: for cc in ccs_test_post_cut: - if np.any(np.in1d(self.source_graph_ids, cc)): - assert np.all(np.in1d(self.source_graph_ids, cc)) - assert ~np.any(np.in1d(self.sink_graph_ids, cc)) + if np.any(np.isin(self.source_graph_ids, cc)): + assert np.all(np.isin(self.source_graph_ids, cc)) + assert ~np.any(np.isin(self.sink_graph_ids, cc)) if ( len(self.source_path_vertices) == len(cc) and self.disallow_isolating_cut @@ -565,9 +642,9 @@ def _sink_and_source_connectivity_sanity_check(self, cut_edge_set): if not self.partition_edges_within_label(cc): raise IsolatingCutException("Source") - if np.any(np.in1d(self.sink_graph_ids, cc)): - assert np.all(np.in1d(self.sink_graph_ids, cc)) - assert ~np.any(np.in1d(self.source_graph_ids, cc)) + if np.any(np.isin(self.sink_graph_ids, cc)): + assert np.all(np.isin(self.sink_graph_ids, cc)) + assert ~np.any(np.isin(self.source_graph_ids, cc)) if ( len(self.sink_path_vertices) == len(cc) and self.disallow_isolating_cut @@ -581,12 +658,15 @@ def _sink_and_source_connectivity_sanity_check(self, cut_edge_set): # but return a flag to return a message to the user illegal_split = True else: - raise PreconditionError( + msg = ( "Failed to find a cut that separated the sources from the sinks. " "Please try another cut that partitions the sets cleanly if possible. " "If there is a clear path between all the supervoxels in each set, " "that helps the mincut algorithm." ) + if self.sv_split_supported: + raise SupervoxelSplitRequiredError(msg, self.sv_remapping) + raise PreconditionError(msg) except IsolatingCutException as e: if self.split_preview: illegal_split = True @@ -601,18 +681,24 @@ def _sink_and_source_connectivity_sanity_check(self, cut_edge_set): return ccs_test_post_cut, illegal_split def partition_edges_within_label(self, cc): - """Test is an isolated component has out-edges only within the original - labeled points of the cut + """Test if an isolated component has out-edges only within the original + labeled points of the cut. cc contains weighted_graph indices. + Use weighted_graph_raw to avoid fake infinite edges between sources/sinks. """ - label_graph_ids = np.concatenate((self.source_graph_ids, self.sink_graph_ids)) - + label_svs = np.concatenate((self.sources, self.sinks)) for vind in cc: - v = self.weighted_graph_raw.vertex(vind) - out_vinds = [int(x) for x in v.out_neighbors()] - if not np.all(np.isin(out_vinds, label_graph_ids)): + sv = self.unique_supervoxel_ids[vind] + raw_inds = np.where(self.unique_supervoxel_ids_raw == sv)[0] + if len(raw_inds) == 0: + # SV not in raw graph (only cross-chunk edges) — no local neighbors + continue + v = self.weighted_graph_raw.vertex(raw_inds[0]) + neighbor_svs = self.unique_supervoxel_ids_raw[ + [int(x) for x in v.out_neighbors()] + ] + if not np.all(np.isin(neighbor_svs, label_svs)): return False - else: - return True + return True def run_multicut( @@ -623,20 +709,39 @@ def run_multicut( split_preview: bool = False, path_augment: bool = True, disallow_isolating_cut: bool = True, -): - local_mincut_graph = LocalMincutGraph( - edges.get_pairs(), - edges.affinities, - source_ids, - sink_ids, - split_preview, - path_augment, - disallow_isolating_cut=disallow_isolating_cut, - ) - atomic_edges = local_mincut_graph.compute_mincut() - if len(atomic_edges) == 0: + sv_split_supported: bool = False, +) -> MulticutResult: + """Run the multicut and return either the cut edges or an SV-split request. + + When `sv_split_supported=True`, the "source and sink share a cross-chunk + rep" condition is returned as `SvSplitRequired` rather than raised — + `SupervoxelSplitRequiredError` is an implementation detail of + `LocalMincutGraph` unwinding, caught at this boundary so it never + drives control flow in callers. + """ + try: + local_mincut_graph = LocalMincutGraph( + edges.get_pairs(), + edges.affinities, + source_ids, + sink_ids, + split_preview, + path_augment, + disallow_isolating_cut=disallow_isolating_cut, + sv_split_supported=sv_split_supported, + ) + mincut_output = local_mincut_graph.compute_mincut() + except SupervoxelSplitRequiredError as err: + return SvSplitRequired(err.sv_remapping) + + if split_preview: + # compute_mincut returns (ccs, illegal_split) in preview mode. + supervoxel_ccs, illegal_split = mincut_output + return PreviewCut(supervoxel_ccs, illegal_split) + + if len(mincut_output) == 0: raise PostconditionError(f"Mincut failed. Try with a different set of points.") - return atomic_edges + return Cut(mincut_output) def run_split_preview( @@ -649,11 +754,13 @@ def run_split_preview( path_augment: bool = True, disallow_isolating_cut: bool = True, ): + sink_and_source_ids = np.concatenate([source_ids, sink_ids]) + roots = cg.get_roots(sink_and_source_ids, assert_roots=True) root_ids = set( - cg.get_roots(np.concatenate([source_ids, sink_ids]), assert_roots=True) + assert_same_root( + sink_and_source_ids, roots, source="run_split_preview" + ).tolist() ) - if len(root_ids) > 1: - raise PreconditionError("Supervoxels must belong to the same object.") bbox = get_bounding_box(source_coords, sink_coords, bb_offset) l2id_agglomeration_d, edges = cg.get_subgraph( @@ -664,19 +771,26 @@ def run_split_preview( supervoxels = np.concatenate( [agg.supervoxels for agg in l2id_agglomeration_d.values()] ) - mask0 = np.in1d(edges.node_ids1, supervoxels) - mask1 = np.in1d(edges.node_ids2, supervoxels) + mask0 = np.isin(edges.node_ids1, supervoxels) + mask1 = np.isin(edges.node_ids2, supervoxels) edges = edges[mask0 & mask1] - edges_to_remove, illegal_split = run_multicut( + result = run_multicut( edges, source_ids, sink_ids, split_preview=True, path_augment=path_augment, disallow_isolating_cut=disallow_isolating_cut, + sv_split_supported=cg.meta.ocdbt_seg, ) + if isinstance(result, SvSplitRequired): + # Preview callers can't perform an SV split; surface as a precondition. + raise PreconditionError( + "Supervoxel split required to cut these source/sink points; " + "preview is not available until an edit is applied." + ) - if len(edges_to_remove) == 0: + assert isinstance(result, PreviewCut), f"unexpected preview result type: {result!r}" + if len(result.supervoxel_ccs) == 0: raise PostconditionError("Mincut could not find any edges to remove.") - - return edges_to_remove, illegal_split + return result.supervoxel_ccs, result.illegal_split diff --git a/pychunkedgraph/graph/downsample.py b/pychunkedgraph/graph/downsample.py new file mode 100644 index 000000000..7a28305c4 --- /dev/null +++ b/pychunkedgraph/graph/downsample.py @@ -0,0 +1,342 @@ +"""Async mip-pyramid downsample worker support. + +An SV split writes at base resolution only; coarser mips are produced +afterwards by a pubsub worker that consumes this module's primitives. + +Work is organized into `pyramid_block`s. A block is a cubic physical +region sized so that at the coarsest scale in the pyramid it equals +exactly one storage chunk. Because every finer scale's chunk grid is a +power-of-2 refinement of the coarsest, a block aligned at the coarsest +scale is automatically aligned at every finer scale — so two different +blocks never share a storage chunk at any mip. That is what makes a +single lock per block safe. + +Within a block we pick one of two code paths: + 1. Fast in-memory path: read the affected base region once, call + tinybrain with `num_mips=K` (all mips at once), write each mip's + output. Used when the base read fits a memory budget — the typical + case because the SV-split bbox is bounded by the /split endpoint + (source+sink coords + small padding). + 2. Per-mip fallback: read the previous mip, tinybrain one step, write. + K storage round-trips instead of 1. Kept for pathological inputs + whose base read would exceed the memory budget. + +Uniform downsample factor (e.g. 2x2x2) across all non-base scales is +assumed and asserted. +""" + +import numpy as np +import tinybrain + +from pychunkedgraph import get_logger + +logger = get_logger(__name__) + +# Default memory budget for the in-memory path's base read. +# uint64 segmentation is 8 bytes/voxel; 1 GiB ≈ 512^3 voxels. Edits +# produced by the /split endpoint are bounded far below this. +DEFAULT_MEMORY_BUDGET_BYTES = 1 << 30 + + +def num_output_mips(meta) -> int: + """Count of non-base scales — what the worker actually writes.""" + return len(meta.ws_ocdbt_scales) - 1 + + +def uniform_factor(meta) -> tuple: + """Per-axis downsample factor between consecutive scales. + + tinybrain takes one factor tuple per call, so the factor must be + constant across the pyramid. Asserts rather than silently producing + wrong mips for a dataset with mixed factors. + """ + resolutions = [np.array(r, dtype=float) for r in meta.ws_ocdbt_resolutions] + factors = [ + tuple((resolutions[i] / resolutions[i - 1]).astype(int)) + for i in range(1, len(resolutions)) + ] + assert all( + f == factors[0] for f in factors + ), f"non-uniform downsample factors {factors}" + return factors[0] + + +def _chunk_size_at_scale(meta, scale_idx: int) -> np.ndarray: + """Storage chunk size at a given scale (excluding the channel dim).""" + return np.array( + meta.ws_ocdbt_scales[scale_idx].chunk_layout.read_chunk.shape[:3], dtype=int + ) + + +def block_shape(meta) -> np.ndarray: + """pyramid_block size in base-resolution voxels. + + Chosen so that at the coarsest scale K the block equals exactly one + storage chunk — which transitively aligns it to every finer scale's + chunk grid. + """ + K = num_output_mips(meta) + coarsest_chunk = _chunk_size_at_scale(meta, K) + factor = np.array(uniform_factor(meta), dtype=int) + return coarsest_chunk * factor**K + + +def blocks_for_bbox(meta, bbs, bbe) -> list: + """Block coords intersected by a base-resolution bbox. + + Bbox is rounded outward to the block grid — a tiny bbox inside one + block still yields that one block coord. Returns sorted list of + `(bx, by, bz)` ints for deadlock-free lock acquisition. + """ + shape = block_shape(meta) + lo = np.asarray(bbs, dtype=int) // shape + hi = -(-np.asarray(bbe, dtype=int) // shape) + coords = [ + (int(bx), int(by), int(bz)) + for bx in range(lo[0], hi[0]) + for by in range(lo[1], hi[1]) + for bz in range(lo[2], hi[2]) + ] + return sorted(coords) + + +def block_base_bbox(meta, block_coord) -> tuple: + """Inverse of `blocks_for_bbox` for a single coord — base-voxel bbox.""" + shape = block_shape(meta) + lo = np.asarray(block_coord, dtype=int) * shape + hi = lo + shape + return lo, hi + + +def _seg_bboxes_to_np(seg_bboxes): + return [ + (np.asarray(bbs, dtype=int), np.asarray(bbe, dtype=int)) + for bbs, bbe in seg_bboxes + ] + + +def _affected_region_base(meta, block_coord, seg_bboxes_np): + """Base-voxel region covering all tiles this block will write, at any mip. + + Starts from the union of (seg bbox ∩ block ∩ volume) then aligns + outward to the coarsest mip's base-voxel grid (= factor**K per axis). + That alignment both makes the region tinybrain-valid for num_mips=K + and guarantees clean chunk-aligned writes at every mip (coarsest + alignment refines down to every finer scale). + + Returns `(base_lo, base_hi)` or `None` if no overlap. + """ + K = num_output_mips(meta) + factor = np.array(uniform_factor(meta), dtype=int) + align = factor**K + + block_lo, block_hi = block_base_bbox(meta, block_coord) + vol_lo = meta.voxel_bounds[:, 0] + vol_hi = meta.voxel_bounds[:, 1] + clipped_lo = np.maximum(block_lo, vol_lo) + clipped_hi = np.minimum(block_hi, vol_hi) + if np.any(clipped_hi <= clipped_lo): + return None + + union_lo, union_hi = None, None + for sb, eb in seg_bboxes_np: + ilo = np.maximum(sb, clipped_lo) + ihi = np.minimum(eb, clipped_hi) + if np.any(ihi <= ilo): + continue + union_lo = ilo if union_lo is None else np.minimum(union_lo, ilo) + union_hi = ihi if union_hi is None else np.maximum(union_hi, ihi) + if union_lo is None: + return None + + base_lo = (union_lo // align) * align + base_hi = -(-union_hi // align) * align + # Keep within the clipped block. Block corners are factor**K-aligned + # (block_shape is a multiple of factor**K), so this clip preserves + # alignment. + base_lo = np.maximum(base_lo, clipped_lo) + base_hi = np.minimum(base_hi, clipped_hi) + if np.any(base_hi <= base_lo): + return None + return base_lo, base_hi + + +def _process_block_in_memory(meta, base_region, K, factor): + """Read base once, tinybrain all mips, write each output. + + Assumes the base region is factor**K-aligned in size (which is what + `_affected_region_base` returns) so tinybrain with num_mips=K emits + clean integer voxel counts at every mip. + """ + base_lo, base_hi = base_region + base = meta.ws_ocdbt_scales[0] + arr = ( + base[ + base_lo[0] : base_hi[0], + base_lo[1] : base_hi[1], + base_lo[2] : base_hi[2], + :, + ] + .read() + .result() + ) + mips = tinybrain.downsample_segmentation( + arr, factor=tuple(int(f) for f in factor), num_mips=K, sparse=False + ) + for m, out in enumerate(mips, start=1): + scale = factor**m + mip_lo = base_lo // scale + mip_hi = base_hi // scale + dst = meta.ws_ocdbt_scales[m] + dst[ + mip_lo[0] : mip_hi[0], + mip_lo[1] : mip_hi[1], + mip_lo[2] : mip_hi[2], + :, + ].write(out).result() + + +def _affected_region_at_mip( + block_lo_base, + block_hi_base, + vol_lo, + vol_hi, + seg_bboxes_base, + mip: int, + factor: np.ndarray, + mip_chunk: np.ndarray, +): + """Write region at this mip in mip-local voxel coords. + + Union of seg bboxes ∩ block ∩ volume, aligned outward to this mip's + storage-chunk grid. Returns `(mip_lo, mip_hi)` or None. + """ + scale = factor**mip + clipped_lo = np.maximum(block_lo_base, vol_lo) + clipped_hi = np.minimum(block_hi_base, vol_hi) + if np.any(clipped_hi <= clipped_lo): + return None + + union_lo, union_hi = None, None + for sb, eb in seg_bboxes_base: + ilo = np.maximum(sb, clipped_lo) + ihi = np.minimum(eb, clipped_hi) + if np.any(ihi <= ilo): + continue + union_lo = ilo if union_lo is None else np.minimum(union_lo, ilo) + union_hi = ihi if union_hi is None else np.maximum(union_hi, ihi) + if union_lo is None: + return None + + mip_lo = union_lo // scale + mip_hi = -(-union_hi // scale) + mip_lo = (mip_lo // mip_chunk) * mip_chunk + mip_hi = -(-mip_hi // mip_chunk) * mip_chunk + + vol_lo_mip = vol_lo // scale + vol_hi_mip = -(-vol_hi // scale) + mip_lo = np.maximum(mip_lo, vol_lo_mip) + mip_hi = np.minimum(mip_hi, vol_hi_mip) + if np.any(mip_hi <= mip_lo): + return None + return mip_lo, mip_hi + + +def _process_block_per_mip(meta, block_coord, seg_bboxes_np, K, factor): + """Fallback path: process one mip at a time. + + Used when the full in-memory base read would exceed the memory + budget. Each mip reads the prior mip from storage, does one + tinybrain step, writes. + + Safe across mip boundaries only because the caller holds the block + lock — no other task can write the storage chunks this block owns, + so reading mip N here always sees what we wrote at mip N in the + previous iteration. + """ + vol_lo = meta.voxel_bounds[:, 0] + vol_hi = meta.voxel_bounds[:, 1] + block_lo_base, block_hi_base = block_base_bbox(meta, block_coord) + + for mip in range(1, K + 1): + mip_chunk = _chunk_size_at_scale(meta, mip) + region = _affected_region_at_mip( + block_lo_base, + block_hi_base, + vol_lo, + vol_hi, + seg_bboxes_np, + mip, + factor, + mip_chunk, + ) + if region is None: + continue + mip_lo, mip_hi = region + src = meta.ws_ocdbt_scales[mip - 1] + src_lo = mip_lo * factor + src_hi = mip_hi * factor + arr = ( + src[ + src_lo[0] : src_hi[0], + src_lo[1] : src_hi[1], + src_lo[2] : src_hi[2], + :, + ] + .read() + .result() + ) + out = tinybrain.downsample_segmentation( + arr, factor=tuple(int(f) for f in factor), num_mips=1, sparse=False + )[0] + dst = meta.ws_ocdbt_scales[mip] + dst[ + mip_lo[0] : mip_hi[0], + mip_lo[1] : mip_hi[1], + mip_lo[2] : mip_hi[2], + :, + ].write(out).result() + + +def process_block( + meta, + block_coord, + seg_bboxes, + memory_budget_bytes: int = DEFAULT_MEMORY_BUDGET_BYTES, +): + """Downsample one pyramid_block through every non-base mip. + + Atomic within the block: caller must hold the block lock. Picks the + in-memory path when the base read fits the memory budget, falls + back to the per-mip path otherwise. + + Reads and writes only the aligned region covering `seg_bboxes` + inside the block; the rest of the block is untouched. Region + alignment rounds outward to the coarsest mip's grid so the aligned + region is always tinybrain-valid and chunk-aligned at every mip. + + Args: + meta: ChunkedGraphMeta with `ws_ocdbt_scales` / `ws_ocdbt_resolutions`. + block_coord: (bx, by, bz) block grid coord. + seg_bboxes: iterable of `(bbs, bbe)` base-voxel bbox pairs from + the SV splits that triggered this job. + """ + K = num_output_mips(meta) + factor = np.array(uniform_factor(meta), dtype=int) + seg_bboxes_np = _seg_bboxes_to_np(seg_bboxes) + + region = _affected_region_base(meta, block_coord, seg_bboxes_np) + if region is None: + return + base_lo, base_hi = region + + bytes_per_voxel = meta.ws_ocdbt_scales[0].dtype.numpy_dtype.itemsize + base_bytes = int(np.prod(base_hi - base_lo)) * bytes_per_voxel + if base_bytes <= memory_budget_bytes: + _process_block_in_memory(meta, region, K, factor) + else: + logger.info( + f"block {block_coord} base read {base_bytes / 1e9:.2f} GB exceeds " + f"budget {memory_budget_bytes / 1e9:.2f} GB; using per-mip path" + ) + _process_block_per_mip(meta, block_coord, seg_bboxes_np, K, factor) diff --git a/pychunkedgraph/graph/dry_run.py b/pychunkedgraph/graph/dry_run.py new file mode 100644 index 000000000..985d3e3ff --- /dev/null +++ b/pychunkedgraph/graph/dry_run.py @@ -0,0 +1,38 @@ +import os +from contextlib import contextmanager + +DRY_RUN_ENV = "PCG_DRY_RUN" + + +def is_dry_run() -> bool: + """True iff ``PCG_DRY_RUN=1`` in the environment. + + When true, every write function in the edit flow + (``operation._write``, the operation log writes, ``write_seg_chunks``, + and the lock acquire/release paths) returns early without + persisting. Used by debug tooling to re-run edits against + production BT/OCDBT state without mutating it. + + Strict ``"1"`` match so unset / empty / ``"true"`` / typos do not + accidentally trigger in production. + """ + return os.environ.get(DRY_RUN_ENV) == "1" + + +@contextmanager +def dry_run_scope(): + """Set ``PCG_DRY_RUN=1`` for the duration of the block; restore on exit. + + Single point for set/restore of the env var. The caller's + pre-existing value (including absence) is restored even if the + block raises. + """ + prev = os.environ.get(DRY_RUN_ENV) + os.environ[DRY_RUN_ENV] = "1" + try: + yield + finally: + if prev is None: + os.environ.pop(DRY_RUN_ENV, None) + else: + os.environ[DRY_RUN_ENV] = prev diff --git a/pychunkedgraph/graph/edges/__init__.py b/pychunkedgraph/graph/edges/__init__.py index b0e488d05..b8fd9d301 100644 --- a/pychunkedgraph/graph/edges/__init__.py +++ b/pychunkedgraph/graph/edges/__init__.py @@ -2,104 +2,11 @@ Classes and types for edges """ -from typing import Optional -from collections import namedtuple - -import numpy as np - -from ..utils import basetypes - - -_edge_type_fileds = ("in_chunk", "between_chunk", "cross_chunk") -_edge_type_defaults = ("in", "between", "cross") - -EdgeTypes = namedtuple("EdgeTypes", _edge_type_fileds, defaults=_edge_type_defaults) -EDGE_TYPES = EdgeTypes() - -DEFAULT_AFFINITY = np.finfo(np.float32).tiny -DEFAULT_AREA = np.finfo(np.float32).tiny - - -class Edges: - def __init__( - self, - node_ids1: np.ndarray, - node_ids2: np.ndarray, - *, - affinities: Optional[np.ndarray] = None, - areas: Optional[np.ndarray] = None, - fake_edges=False, - ): - self.node_ids1 = np.array(node_ids1, dtype=basetypes.NODE_ID, copy=False) - self.node_ids2 = np.array(node_ids2, dtype=basetypes.NODE_ID, copy=False) - assert self.node_ids1.size == self.node_ids2.size - - self._as_pairs = None - self._fake_edges = fake_edges - - if affinities is not None and len(affinities) > 0: - self._affinities = np.array(affinities, dtype=basetypes.EDGE_AFFINITY, copy=False) - assert self.node_ids1.size == self._affinities.size - else: - self._affinities = np.full(len(self.node_ids1), DEFAULT_AFFINITY) - - if areas is not None and len(areas) > 0: - self._areas = np.array(areas, dtype=basetypes.EDGE_AREA, copy=False) - assert self.node_ids1.size == self._areas.size - else: - self._areas = np.full(len(self.node_ids1), DEFAULT_AREA) - - @property - def affinities(self) -> np.ndarray: - return self._affinities - - @affinities.setter - def affinities(self, affinities): - self._affinities = affinities - - @property - def areas(self) -> np.ndarray: - return self._areas - - @areas.setter - def areas(self, areas): - self._areas = areas - - def __add__(self, other): - """add two Edges instances""" - node_ids1 = np.concatenate([self.node_ids1, other.node_ids1]) - node_ids2 = np.concatenate([self.node_ids2, other.node_ids2]) - affinities = np.concatenate([self.affinities, other.affinities]) - areas = np.concatenate([self.areas, other.areas]) - return Edges(node_ids1, node_ids2, affinities=affinities, areas=areas) - - def __iadd__(self, other): - self.node_ids1 = np.concatenate([self.node_ids1, other.node_ids1]) - self.node_ids2 = np.concatenate([self.node_ids2, other.node_ids2]) - self.affinities = np.concatenate([self.affinities, other.affinities]) - self.areas = np.concatenate([self.areas, other.areas]) - return self - - def __len__(self): - return self.node_ids1.size - - def __getitem__(self, key): - """`key` must be a boolean numpy array.""" - try: - return Edges( - self.node_ids1[key], - self.node_ids2[key], - affinities=self.affinities[key], - areas=self.areas[key], - ) - except Exception as err: - raise (err) - - def get_pairs(self) -> np.ndarray: - """ - return numpy array of edge pairs [[sv1, sv2] ... ] - """ - if not self._as_pairs is None: - return self._as_pairs - self._as_pairs = np.column_stack((self.node_ids1, self.node_ids2)) - return self._as_pairs +from .definitions import EDGE_TYPES, Edges + +from .stale import ( + get_new_nodes, + get_stale_nodes, + get_latest_edges, + get_latest_edges_wrapper, +) diff --git a/pychunkedgraph/graph/edges/definitions.py b/pychunkedgraph/graph/edges/definitions.py new file mode 100644 index 000000000..831ca9798 --- /dev/null +++ b/pychunkedgraph/graph/edges/definitions.py @@ -0,0 +1,110 @@ +""" +Edge data structures and type definitions. +""" + +from collections import namedtuple +from typing import Optional + +import numpy as np + +from pychunkedgraph.graph import basetypes + +_edge_type_fileds = ("in_chunk", "between_chunk", "cross_chunk") +_edge_type_defaults = ("in", "between", "cross") + +EdgeTypes = namedtuple("EdgeTypes", _edge_type_fileds, defaults=_edge_type_defaults) +EDGE_TYPES = EdgeTypes() + +DEFAULT_AFFINITY = np.finfo(np.float32).tiny +DEFAULT_AREA = np.finfo(np.float32).tiny +ADJACENCY_DTYPE = np.dtype( + [ + ("node", basetypes.NODE_ID), + ("aff", basetypes.EDGE_AFFINITY), + ("area", basetypes.EDGE_AREA), + ] +) +ZSTD_EDGE_COMPRESSION = 17 + + +class Edges: + def __init__( + self, + node_ids1: np.ndarray, + node_ids2: np.ndarray, + *, + affinities: Optional[np.ndarray] = None, + areas: Optional[np.ndarray] = None, + ): + self.node_ids1 = np.array(node_ids1, dtype=basetypes.NODE_ID) + self.node_ids2 = np.array(node_ids2, dtype=basetypes.NODE_ID) + assert self.node_ids1.size == self.node_ids2.size + + self._as_pairs = None + + if affinities is not None and len(affinities) > 0: + self._affinities = np.array(affinities, dtype=basetypes.EDGE_AFFINITY) + assert self.node_ids1.size == self._affinities.size + else: + self._affinities = np.full(len(self.node_ids1), DEFAULT_AFFINITY) + + if areas is not None and len(areas) > 0: + self._areas = np.array(areas, dtype=basetypes.EDGE_AREA) + assert self.node_ids1.size == self._areas.size + else: + self._areas = np.full(len(self.node_ids1), DEFAULT_AREA) + + @property + def affinities(self) -> np.ndarray: + return self._affinities + + @affinities.setter + def affinities(self, affinities): + self._affinities = affinities + + @property + def areas(self) -> np.ndarray: + return self._areas + + @areas.setter + def areas(self, areas): + self._areas = areas + + def __add__(self, other): + """add two Edges instances""" + node_ids1 = np.concatenate([self.node_ids1, other.node_ids1]) + node_ids2 = np.concatenate([self.node_ids2, other.node_ids2]) + affinities = np.concatenate([self.affinities, other.affinities]) + areas = np.concatenate([self.areas, other.areas]) + return Edges(node_ids1, node_ids2, affinities=affinities, areas=areas) + + def __iadd__(self, other): + self.node_ids1 = np.concatenate([self.node_ids1, other.node_ids1]) + self.node_ids2 = np.concatenate([self.node_ids2, other.node_ids2]) + self.affinities = np.concatenate([self.affinities, other.affinities]) + self.areas = np.concatenate([self.areas, other.areas]) + return self + + def __len__(self): + return self.node_ids1.size + + def __getitem__(self, key): + """`key` must be a boolean numpy array.""" + try: + return Edges( + self.node_ids1[key], + self.node_ids2[key], + affinities=self.affinities[key], + areas=self.areas[key], + ) + except Exception as err: + raise (err) + + def get_pairs(self) -> np.ndarray: + """ + return numpy array of edge pairs [[sv1, sv2] ... ] + """ + if not self._as_pairs is None: + return self._as_pairs + self._as_pairs = np.column_stack((self.node_ids1, self.node_ids2)) + return self._as_pairs diff --git a/pychunkedgraph/graph/edges/stale.py b/pychunkedgraph/graph/edges/stale.py new file mode 100644 index 000000000..9aab8c245 --- /dev/null +++ b/pychunkedgraph/graph/edges/stale.py @@ -0,0 +1,503 @@ +""" +Stale node detection and edge update logic. +""" + +import datetime +from os import environ + +from pychunkedgraph import get_logger + +logger = get_logger(__name__) +from typing import Iterable + +import numpy as np +from cachetools import LRUCache + +from pychunkedgraph.graph import types +from pychunkedgraph.graph.chunks.utils import get_l2chunkids_along_boundary + +from pychunkedgraph.graph import basetypes +from ..utils.generic import get_parents_at_timestamp + +PARENTS_CACHE: LRUCache = None +CHILDREN_CACHE: LRUCache = None + + +def get_new_nodes( + cg, nodes: np.ndarray, layer: int, parent_ts: datetime.datetime = None +): + unique_nodes, inverse = np.unique(nodes, return_inverse=True) + node_root_map = {n: n for n in unique_nodes} + lookup = np.ones(len(unique_nodes), dtype=unique_nodes.dtype) + while np.any(lookup): + roots = np.fromiter(node_root_map.values(), dtype=basetypes.NODE_ID) + roots = cg.get_parents(roots, time_stamp=parent_ts, fail_to_zero=True) + layers = cg.get_chunk_layers(roots) + lookup[layers > layer] = 0 + lookup[roots == 0] = 0 + + layer_mask = layers <= layer + non_zero_mask = roots != 0 + mask = layer_mask & non_zero_mask + for node, root in zip(unique_nodes[mask], roots[mask]): + node_root_map[node] = root + + unique_results = np.fromiter(node_root_map.values(), dtype=basetypes.NODE_ID) + return unique_results[inverse] + + +def get_stale_nodes( + cg, nodes: Iterable[basetypes.NODE_ID], parent_ts: datetime.datetime = None +): + """ + Checks to see if given nodes are stale. + This is done by getting a supervoxel of a node and checking + if it has a new parent at the same layer as the node. + """ + nodes = np.unique(np.array(nodes, dtype=basetypes.NODE_ID)) + new_ids = set() if cg.cache is None else cg.cache.new_ids + nodes = nodes[~np.isin(nodes, new_ids)] + supervoxels = cg.get_single_leaf_multiple(nodes) + # nodes can be at different layers due to skip connections + node_layers = cg.get_chunk_layers(nodes) + stale_nodes = [types.empty_1d] + for layer in np.unique(node_layers): + _mask = node_layers == layer + layer_nodes = nodes[_mask] + _nodes = get_new_nodes(cg, supervoxels[_mask], layer, parent_ts) + stale_mask = layer_nodes != _nodes + stale_nodes.append(layer_nodes[stale_mask]) + return np.concatenate(stale_nodes) + + +class LatestEdgesFinder: + """ + For each of stale_edges [[`node`, `partner`]], get their L2 edge equivalent. + Then get supervoxels of those L2 IDs and get parent(s) at `node` level. + These parents would be the new identities for the stale `partner`. + """ + + def __init__( + self, + cg, + stale_edges: Iterable, + edge_layers: Iterable, + parent_ts: datetime.datetime = None, + ): + self.cg = cg + self.stale_edges = stale_edges + self.edge_layers = edge_layers + self.parent_ts = parent_ts + + _nodes = np.unique(stale_edges) + self.nodes_ts_map = dict( + zip( + _nodes, + cg.get_node_timestamps(_nodes, return_numpy=False, normalize=True), + ) + ) + layers, coords = cg.get_chunk_layers_and_coordinates(_nodes) + self.layers_d = dict(zip(_nodes, layers)) + self.coords_d = dict(zip(_nodes, coords)) + + def _get_children_from_cache(self, nodes): + children = [] + non_cached = [] + for node in nodes: + try: + v = CHILDREN_CACHE[node] + children.append(v) + except KeyError: + non_cached.append(node) + + children_map = self.cg.get_children(non_cached) + for k, v in children_map.items(): + CHILDREN_CACHE[k] = v + children.append(v) + return np.concatenate(children) + + def _get_normalized_coords(self, node_a, node_b) -> tuple: + max_layer = self.layers_d[node_a] + coord_a, coord_b = self.coords_d[node_a], self.coords_d[node_b] + if self.layers_d[node_a] != self.layers_d[node_b]: + # normalize if nodes are not from the same layer + max_layer = max(self.layers_d[node_a], self.layers_d[node_b]) + chunk_a = self.cg.get_parent_chunk_id(node_a, parent_layer=max_layer) + chunk_b = self.cg.get_parent_chunk_id(node_b, parent_layer=max_layer) + coord_a, coord_b = self.cg.get_chunk_coordinates_multiple( + [chunk_a, chunk_b] + ) + return max_layer, tuple(coord_a), tuple(coord_b) + + def _get_filtered_l2ids(self, node_a, node_b, padding: int): + """ + Finds L2 IDs along opposing faces for given nodes. + Filterting is done by first finding L2 chunks along these faces. + Then get their parent chunks iteratively. + Then filter children iteratively using these chunks. + """ + chunks_map = {} + + def _filter(node): + result = [] + children = np.array([node], dtype=basetypes.NODE_ID) + while True: + chunk_ids = self.cg.get_chunk_ids_from_node_ids(children) + mask = np.isin(chunk_ids, chunks_map[node]) + children = children[mask] + + mask = self.cg.get_chunk_layers(children) == 2 + result.append(children[mask]) + + mask = self.cg.get_chunk_layers(children) > 2 + if children[mask].size == 0: + break + if PARENTS_CACHE is None: + children = self.cg.get_children(children[mask], flatten=True) + else: + children = self._get_children_from_cache(children[mask]) + return np.concatenate(result) + + mlayer, coord_a, coord_b = self._get_normalized_coords(node_a, node_b) + chunks_a, chunks_b = get_l2chunkids_along_boundary( + self.cg.meta, mlayer, coord_a, coord_b, padding + ) + + chunks_map[node_a] = [[self.cg.get_chunk_id(node_a)]] + chunks_map[node_b] = [[self.cg.get_chunk_id(node_b)]] + _layer = 2 + while _layer < mlayer: + chunks_map[node_a].append(chunks_a) + chunks_map[node_b].append(chunks_b) + chunks_a = np.unique(self.cg.get_parent_chunk_id_multiple(chunks_a)) + chunks_b = np.unique(self.cg.get_parent_chunk_id_multiple(chunks_b)) + _layer += 1 + chunks_map[node_a] = np.concatenate(chunks_map[node_a]) + chunks_map[node_b] = np.concatenate(chunks_map[node_b]) + return int(mlayer), _filter(node_a), _filter(node_b) + + def _populate_parents_cache(self, children: np.ndarray): + global PARENTS_CACHE + + not_cached = [] + for child in children: + try: + # reset lru index, these will be needed soon + _ = PARENTS_CACHE[child] + except KeyError: + not_cached.append(child) + + all_parents = self.cg.get_parents(not_cached, current=False) + for child, parents in zip(not_cached, all_parents): + PARENTS_CACHE[child] = {} + for parent, ts in parents: + PARENTS_CACHE[child][ts] = parent + + def _get_hierarchy(self, nodes, layer): + _hierarchy = [nodes] + for _a in nodes: + _hierarchy.append( + self.cg.get_root( + _a, + time_stamp=self.parent_ts, + stop_layer=layer, + get_all_parents=True, + ceil=False, + raw_only=True, + ) + ) + _children = self.cg.get_children(_a, raw_only=True) + _children_layers = self.cg.get_chunk_layers(_children) + _hierarchy.append(_children[_children_layers == 2]) + _children = _children[_children_layers > 2] + while _children.size: + _hierarchy.append(_children) + _children = self.cg.get_children(_children, flatten=True, raw_only=True) + _children_layers = self.cg.get_chunk_layers(_children) + _hierarchy.append(_children[_children_layers == 2]) + _children = _children[_children_layers > 2] + return np.concatenate(_hierarchy) + + def _check_cross_edges_from_a(self, node_b, nodes_a, layer, parent_ts): + """ + Checks to match cross edges from partners_a + to hierarchy of potential node from partner b. + """ + if len(nodes_a) == 0: + return False + + _hierarchy_b = self.cg.get_root( + node_b, + time_stamp=parent_ts, + stop_layer=layer, + get_all_parents=True, + ceil=False, + raw_only=True, + ) + _hierarchy_b = np.append(_hierarchy_b, node_b) + _cx_edges_d_from_a = self.cg.get_cross_chunk_edges( + nodes_a, time_stamp=parent_ts + ) + for _edges_d_from_a in _cx_edges_d_from_a.values(): + _edges_from_a = _edges_d_from_a.get(layer, types.empty_2d) + nodes_b_from_a = _edges_from_a[:, 1] + hierarchy_b_from_a = self._get_hierarchy(nodes_b_from_a, layer) + _mask = np.isin(hierarchy_b_from_a, _hierarchy_b) + if np.any(_mask): + return True + return False + + def _check_hierarchy_a_from_b(self, parents_a, nodes_a_from_b, layer, parent_ts): + """ + Checks for overlap between hierarchy of a, + and hierarchy of a identified from partners of b. + """ + if len(nodes_a_from_b) == 0: + return False + + _hierarchy_a = [parents_a] + for _a in parents_a: + _hierarchy_a.append( + self.cg.get_root( + _a, + time_stamp=parent_ts, + stop_layer=layer, + get_all_parents=True, + ceil=False, + raw_only=True, + ) + ) + hierarchy_a = np.concatenate(_hierarchy_a) + hierarchy_a_from_b = self._get_hierarchy(nodes_a_from_b, layer) + return np.any(np.isin(hierarchy_a_from_b, hierarchy_a)) + + def _get_parents_b(self, edges, parent_ts, layer, fallback: bool = False): + """ + Attempts to find new partner side nodes. + Gets new partners at parent_ts using supervoxels, at `parent_ts`. + Searches for new partners that may have any edges to `edges[:,0]`. + """ + if PARENTS_CACHE is None: + # this cache is set only during migration + # also, fallback is not applicable if no migration + children_b = self.cg.get_children(edges[:, 1], flatten=True) + parents_b = np.unique(self.cg.get_parents(children_b, time_stamp=parent_ts)) + fallback = False + else: + children_b = self._get_children_from_cache(edges[:, 1]) + self._populate_parents_cache(children_b) + _parents_b, missing = get_parents_at_timestamp( + children_b, PARENTS_CACHE, time_stamp=parent_ts, unique=True + ) + # handle cache miss cases + _parents_b_missed = np.unique( + self.cg.get_parents(missing, time_stamp=parent_ts) + ) + parents_b = np.concatenate([_parents_b, _parents_b_missed]) + + parents_a = np.unique(edges[:, 0]) + stale_a = get_stale_nodes(self.cg, parents_a, parent_ts=parent_ts) + if stale_a.size == parents_a.size or fallback: + # this is applicable only for v2 to v3 migration + # handle cases when source nodes in `edges[:,0]` are stale + atomic_edges_d = self.cg.get_atomic_cross_edges(stale_a) + partners = [types.empty_1d] + for _edges_d in atomic_edges_d.values(): + _edges = _edges_d.get(layer, types.empty_2d) + partners.append(_edges[:, 1]) + partners = np.concatenate(partners) + return np.unique(self.cg.get_parents(partners, time_stamp=parent_ts)) + + _cx_edges_d = self.cg.get_cross_chunk_edges(parents_b, time_stamp=parent_ts) + _parents_b = [] + for _node, _edges_d in _cx_edges_d.items(): + _edges = _edges_d.get(layer, types.empty_2d) + if self._check_cross_edges_from_a(_node, _edges[:, 1], layer, parent_ts): + _parents_b.append(_node) + elif self._check_hierarchy_a_from_b( + parents_a, _edges[:, 1], layer, parent_ts + ): + _parents_b.append(_node) + else: + _new_ids = list(self.cg.cache.new_ids) + if np.any(np.isin(_new_ids, parents_a)): + _parents_b.append(_node) + return np.array(_parents_b, dtype=basetypes.NODE_ID) + + def _get_parents_b_with_chunk_mask( + self, + l2ids_b: np.ndarray, + nodes_b_from_a: np.ndarray, + max_ts: datetime.datetime, + edge, + ): + chunks_old = self.cg.get_chunk_ids_from_node_ids(l2ids_b) + chunks_new = self.cg.get_chunk_ids_from_node_ids(nodes_b_from_a) + chunk_mask = np.isin(chunks_new, chunks_old) + nodes_b_from_a = nodes_b_from_a[chunk_mask] + _stale_nodes = get_stale_nodes(self.cg, nodes_b_from_a, parent_ts=max_ts) + assert _stale_nodes.size == 0, ( + f"stale nodes remain after latest-edge resolve; " + f"edge={edge} stale_nodes={_stale_nodes.tolist()} max_ts={max_ts}" + ) + return nodes_b_from_a + + def _get_cx_edges(self, l2ids_a, max_node_ts, edge_layer, raw_only: bool = True): + _edges_d = self.cg.get_cross_chunk_edges( + node_ids=l2ids_a, time_stamp=max_node_ts, raw_only=raw_only + ) + _edges = [] + for v in _edges_d.values(): + if edge_layer in v: + _edges.append(v[edge_layer]) + return np.concatenate(_edges) + + def _get_dilated_edges(self, edges): + layers_b = self.cg.get_chunk_layers(edges[:, 1]) + _mask = layers_b == 2 + _l2_edges = [edges[_mask]] + for _edge in edges[~_mask]: + _node_a, _node_b = _edge + _nodes_b = self.cg.get_l2children([_node_b]) + _l2_edges.append( + np.array([[_node_a, _b] for _b in _nodes_b], dtype=basetypes.NODE_ID) + ) + return np.unique(np.concatenate(_l2_edges), axis=0) + + def _get_new_edge( + self, edge, edge_layer, parent_ts, padding, fallback: bool = False + ): + """ + Attempts to find new edge(s) for the stale `edge`. + * Find L2 IDs on opposite sides of the face in L2 chunks along the face. + * Find new edges between them (before the given timestamp). + * If none found, expand search by adding another layer of L2 chunks. + """ + node_a, node_b = edge + mlayer, l2ids_a, l2ids_b = self._get_filtered_l2ids( + node_a, node_b, padding=padding + ) + if l2ids_a.size == 0 or l2ids_b.size == 0: + return types.empty_2d.copy() + + max_ts = max(self.nodes_ts_map[node_a], self.nodes_ts_map[node_b]) + is_l2_edge = node_a in l2ids_a and node_b in l2ids_b + if is_l2_edge and (l2ids_a.size == 1 and l2ids_b.size == 1): + _edges = np.array([edge], dtype=basetypes.NODE_ID) + else: + try: + _edges = self._get_cx_edges(l2ids_a, max_ts, edge_layer) + except ValueError: + _edges = self._get_cx_edges(l2ids_a, max_ts, edge_layer, raw_only=False) + except ValueError: + return types.empty_2d.copy() + + mask = np.isin(_edges[:, 1], l2ids_b) + if np.any(mask): + parents_b = self._get_parents_b(_edges[mask], parent_ts, edge_layer) + else: + # partner nodes likely lifted, dilate and retry + _edges = self._get_dilated_edges(_edges) + mask = np.isin(_edges[:, 1], l2ids_b) + if np.any(mask): + parents_b = self._get_parents_b(_edges[mask], parent_ts, edge_layer) + else: + # if none of `l2ids_b` were found in edges, `l2ids_a` already have new edges + # so get the new identities of `l2ids_b` by using chunk mask + try: + parents_b = self._get_parents_b_with_chunk_mask( + l2ids_b, _edges[:, 1], max_ts, edge + ) + except AssertionError: + parents_b = [] + if fallback: + parents_b = self._get_parents_b( + _edges, parent_ts, edge_layer, True + ) + + parents_b = np.unique(get_new_nodes(self.cg, parents_b, mlayer, parent_ts)) + parents_a = np.array([node_a] * parents_b.size, dtype=basetypes.NODE_ID) + return np.column_stack((parents_a, parents_b)) + + def run(self): + result = [types.empty_2d] + for edge_layer, _edge in zip(self.edge_layers, self.stale_edges): + max_chebyshev_distance = int(environ.get("MAX_CHEBYSHEV_DISTANCE", 3)) + for pad in range(0, max_chebyshev_distance + 1): + fallback = pad == max_chebyshev_distance + _new_edges = self._get_new_edge( + _edge, + edge_layer, + self.parent_ts, + padding=pad, + fallback=fallback, + ) + if _new_edges.size: + break + logger.note(f"{_edge}, expanding search with padding {pad+1}.") + assert ( + _new_edges.size + ), f"No new edge found {_edge}; {edge_layer}, {self.parent_ts}" + result.append(_new_edges) + return np.concatenate(result) + + +def get_latest_edges( + cg, + stale_edges: Iterable, + edge_layers: Iterable, + parent_ts: datetime.datetime = None, +) -> np.ndarray: + """ + For each of stale_edges [[`node`, `partner`]], get their L2 edge equivalent. + Then get supervoxels of those L2 IDs and get parent(s) at `node` level. + These parents would be the new identities for the stale `partner`. + """ + return LatestEdgesFinder(cg, stale_edges, edge_layers, parent_ts).run() + + +def get_latest_edges_wrapper( + cg, cx_edges_d: dict, parent_ts: datetime.datetime = None +) -> tuple[dict, np.ndarray]: + """ + Helper function to filter stale edges and replace with latest edges. + Filters out edges with nodes stale in source, edges[:,0], at given timestamp. + """ + nodes = [types.empty_1d] + new_cx_edges_d = {0: types.empty_2d} + + all_edges = np.concatenate(list(cx_edges_d.values())) + all_edge_nodes = np.unique(all_edges) + all_stale_nodes = get_stale_nodes(cg, all_edge_nodes, parent_ts=parent_ts) + if all_stale_nodes.size == 0: + return cx_edges_d, all_edge_nodes + + for layer, _cx_edges in cx_edges_d.items(): + if _cx_edges.size == 0: + continue + + _new_cx_edges = [types.empty_2d] + _edge_layers = np.array([layer] * len(_cx_edges), dtype=int) + + stale_source_mask = np.isin(_cx_edges[:, 0], all_stale_nodes) + _new_cx_edges.append(_cx_edges[stale_source_mask]) + + _cx_edges = _cx_edges[~stale_source_mask] + _edge_layers = _edge_layers[~stale_source_mask] + stale_destination_mask = np.isin(_cx_edges[:, 1], all_stale_nodes) + _new_cx_edges.append(_cx_edges[~stale_destination_mask]) + + if np.any(stale_destination_mask): + stale_edges = _cx_edges[stale_destination_mask] + stale_edge_layers = _edge_layers[stale_destination_mask] + latest_edges = get_latest_edges( + cg, + stale_edges, + stale_edge_layers, + parent_ts=parent_ts, + ) + logger.debug(f"{stale_edges} -> {latest_edges}; {parent_ts}") + _new_cx_edges.append(latest_edges) + new_cx_edges_d[layer] = np.concatenate(_new_cx_edges) + nodes.append(np.unique(new_cx_edges_d[layer])) + return new_cx_edges_d, np.concatenate(nodes) diff --git a/pychunkedgraph/graph/edges/utils.py b/pychunkedgraph/graph/edges/utils.py index 034ca6ebc..70a0ae32f 100644 --- a/pychunkedgraph/graph/edges/utils.py +++ b/pychunkedgraph/graph/edges/utils.py @@ -8,16 +8,18 @@ from typing import Tuple from typing import Iterable from typing import Optional +from collections import defaultdict +from functools import reduce import fastremap import numpy as np from . import Edges from . import EDGE_TYPES -from ..types import empty_2d -from ..utils import basetypes +from pychunkedgraph.graph import basetypes from ..chunks import utils as chunk_utils from ..meta import ChunkedGraphMeta +from ...utils.general import in2d def concatenate_chunk_edges(chunk_edge_dicts: Iterable) -> Dict: @@ -45,18 +47,21 @@ def concatenate_chunk_edges(chunk_edge_dicts: Iterable) -> Dict: return edges_dict -def concatenate_cross_edge_dicts(edges_ds: Iterable[Dict]) -> Dict: +def concatenate_cross_edge_dicts( + edges_ds: Iterable[Dict], unique: bool = False +) -> Dict: """Combines cross chunk edge dicts of form {layer id : edge list}.""" - from collections import defaultdict - result_d = defaultdict(list) - for edges_d in edges_ds: for layer, edges in edges_d.items(): result_d[layer].append(edges) for layer, edge_lists in result_d.items(): - result_d[layer] = np.concatenate(edge_lists) + edge_lists = [np.asarray(e, dtype=basetypes.NODE_ID) for e in edge_lists] + edges = np.concatenate(edge_lists) + if unique: + edges = np.unique(edges, axis=0) + result_d[layer] = edges return result_d @@ -65,7 +70,11 @@ def merge_cross_edge_dicts(x_edges_d1: Dict, x_edges_d2: Dict) -> Dict: Combines two cross chunk dictionaries of form {node_id: {layer id : edge list}}. """ - node_ids = np.unique(list(x_edges_d1.keys()) + list(x_edges_d2.keys())) + node_ids = np.unique( + np.array( + list(x_edges_d1.keys()) + list(x_edges_d2.keys()), dtype=basetypes.NODE_ID + ) + ) result_d = {} for node_id in node_ids: cross_edge_ds = [x_edges_d1.get(node_id, {}), x_edges_d2.get(node_id, {})] @@ -131,7 +140,7 @@ def categorize_edges_v2( def get_cross_chunk_edges_layer(meta: ChunkedGraphMeta, cross_edges: Iterable): - """Computes the layer in which a cross chunk edge becomes relevant. + """Computes the layer in which an atomic cross chunk edge becomes relevant. I.e. if a cross chunk edge links two nodes in layer 4 this function returns 3. :param cross_edges: n x 2 array @@ -152,40 +161,7 @@ def get_cross_chunk_edges_layer(meta: ChunkedGraphMeta, cross_edges: Iterable): return cross_chunk_edge_layers -def filter_min_layer_cross_edges( - meta: ChunkedGraphMeta, cross_edges_d: Dict, node_layer: int = 2 -) -> Tuple[int, Iterable]: - """ - Given a dict of cross chunk edges {layer: edges} - Return the first layer with cross edges. - """ - for layer in range(node_layer, meta.layer_count): - edges_ = cross_edges_d.get(layer, empty_2d) - if edges_.size: - return (layer, edges_) - return (meta.layer_count, edges_) - - -def filter_min_layer_cross_edges_multiple( - meta: ChunkedGraphMeta, l2id_atomic_cross_edges_ds: Iterable, node_layer: int = 2 -) -> Tuple[int, Iterable]: - """ - Given a list of dicts of cross chunk edges [{layer: edges}] - Return the first layer with cross edges. - """ - min_layer = meta.layer_count - for edges_d in l2id_atomic_cross_edges_ds: - layer_, _ = filter_min_layer_cross_edges(meta, edges_d, node_layer=node_layer) - min_layer = min(min_layer, layer_) - edges = [empty_2d] - for edges_d in l2id_atomic_cross_edges_ds: - edges.append(edges_d.get(min_layer, empty_2d)) - return min_layer, np.concatenate(edges) - - def get_edges_status(cg, edges: Iterable, time_stamp: Optional[float] = None): - from ...utils.general import in2d - coords0 = chunk_utils.get_chunk_coordinates_multiple(cg.meta, edges[:, 0]) coords1 = chunk_utils.get_chunk_coordinates_multiple(cg.meta, edges[:, 1]) @@ -214,3 +190,20 @@ def get_edges_status(cg, edges: Iterable, time_stamp: Optional[float] = None): active_status.extend(mask) active_status = np.array(active_status, dtype=bool) return existence_status, active_status + + +def filter_inactive_cross_edges( + cg, all_chunk_edges: Edges, time_stamp: Optional[float] = None +): + result = [] + layers = cg.get_cross_chunk_edges_layer(all_chunk_edges.get_pairs()) + for layer in np.unique(layers): + layer_mask = layers == layer + parent_layer = layer + 1 + layer_edges = all_chunk_edges[layer_mask] + n1, n2 = layer_edges.node_ids1, layer_edges.node_ids2 + parents1 = cg.get_roots(n1, stop_layer=parent_layer, time_stamp=time_stamp) + parents2 = cg.get_roots(n2, stop_layer=parent_layer, time_stamp=time_stamp) + mask = parents1 == parents2 + result.append(layer_edges[mask]) + return reduce(lambda x, y: x + y, result, Edges([], [])) diff --git a/pychunkedgraph/graph/edits.py b/pychunkedgraph/graph/edits.py index be2eee1c6..a86037851 100644 --- a/pychunkedgraph/graph/edits.py +++ b/pychunkedgraph/graph/edits.py @@ -1,46 +1,71 @@ # pylint: disable=invalid-name, missing-docstring, too-many-locals, c-extension-no-member -import datetime +import datetime, random from typing import Dict from typing import List from typing import Tuple from typing import Iterable +from typing import Set from collections import defaultdict -import numpy as np import fastremap +import numpy as np + +from pychunkedgraph import get_logger +from pychunkedgraph.profiler import HierarchicalProfiler, get_profiler from . import types -from . import attributes +from pychunkedgraph.graph import attributes from . import cache as cache_utils +from .edges import get_latest_edges_wrapper, get_new_nodes from .edges.utils import concatenate_cross_edge_dicts from .edges.utils import merge_cross_edge_dicts -from .utils import basetypes +from pychunkedgraph.graph import basetypes from .utils import flatgraph -from .utils.serializers import serialize_uint64 -from ..logging.log_db import TimeIt +from pychunkedgraph.graph import serializers from ..utils.general import in2d +from ..debug.utils import sanity_check, sanity_check_single + +logger = get_logger(__name__) def _init_old_hierarchy(cg, l2ids: np.ndarray, parent_ts: datetime.datetime = None): - new_old_id_d = defaultdict(set) - old_new_id_d = defaultdict(set) + """ + Populates old hierarcy from child to root and also gets children of intermediate nodes. + These will be needed later and cached in cg.cache used during an edit. + """ + all_parents = [] old_hierarchy_d = {id_: {2: id_} for id_ in l2ids} + node_layer_parent_map = cg.get_all_parents_dict_multiple( + l2ids, time_stamp=parent_ts + ) for id_ in l2ids: - layer_parent_d = cg.get_all_parents_dict(id_, time_stamp=parent_ts) + layer_parent_d = node_layer_parent_map[id_] old_hierarchy_d[id_].update(layer_parent_d) for parent in layer_parent_d.values(): + all_parents.append(parent) old_hierarchy_d[parent] = old_hierarchy_d[id_] - return new_old_id_d, old_new_id_d, old_hierarchy_d + children = cg.get_children(all_parents, flatten=True) + _ = cg.get_parents(children, time_stamp=parent_ts) + return old_hierarchy_d + + +def flip_ids(id_map, node_ids): + """ + returns old or new ids according to the map + """ + ids = [np.asarray(list(id_map[id_]), dtype=basetypes.NODE_ID) for id_ in node_ids] + ids.append(types.empty_1d) # concatenate needs at least one array + return np.concatenate(ids).astype(basetypes.NODE_ID) def _analyze_affected_edges( cg, atomic_edges: Iterable[np.ndarray], parent_ts: datetime.datetime = None ) -> Tuple[Iterable, Dict]: """ - Determine if atomic edges are within the chunk. - If not, they are cross edges between two L2 IDs in adjacent chunks. - Returns edges between L2 IDs and atomic cross edges. + Returns l2 edges within chunk and self edges for nodes in cross chunk edges. + + Also returns new cross edges dicts for nodes crossing chunk boundary. """ supervoxels = np.unique(atomic_edges) parents = cg.get_parents(supervoxels, time_stamp=parent_ts) @@ -51,23 +76,29 @@ def _analyze_affected_edges( for edge_ in atomic_edges[edge_layers == 1] ] - # cross chunk edges - atomic_cross_edges_d = defaultdict(lambda: defaultdict(list)) + cross_edges_d = defaultdict(lambda: defaultdict(list)) for layer in range(2, cg.meta.layer_count): layer_edges = atomic_edges[edge_layers == layer] if not layer_edges.size: continue for edge in layer_edges: - parent_1 = sv_parent_d[edge[0]] - parent_2 = sv_parent_d[edge[1]] - atomic_cross_edges_d[parent_1][layer].append(edge) - atomic_cross_edges_d[parent_2][layer].append(edge[::-1]) - parent_edges.extend([[parent_1, parent_1], [parent_2, parent_2]]) - return (parent_edges, atomic_cross_edges_d) - - -def _get_relevant_components(edges: np.ndarray, supervoxels: np.ndarray) -> Tuple: - edges = np.concatenate([edges, np.vstack([supervoxels, supervoxels]).T]) + parent0 = sv_parent_d[edge[0]] + parent1 = sv_parent_d[edge[1]] + cross_edges_d[parent0][layer].append([parent0, parent1]) + cross_edges_d[parent1][layer].append([parent1, parent0]) + parent_edges.extend([[parent0, parent0], [parent1, parent1]]) + # Convert inner Python lists to typed numpy arrays to avoid + # dtype promotion issues when concatenated with uint64 arrays. + for node_id in cross_edges_d: + for layer in cross_edges_d[node_id]: + cross_edges_d[node_id][layer] = np.array( + cross_edges_d[node_id][layer], dtype=basetypes.NODE_ID + ).reshape(-1, 2) + return parent_edges, cross_edges_d + + +def _get_relevant_components(edges: np.ndarray, svs: np.ndarray) -> Tuple: + edges = np.concatenate([edges, np.vstack([svs, svs]).T]).astype(basetypes.NODE_ID) graph, _, _, graph_ids = flatgraph.build_gt_graph(edges, make_directed=True) ccs = flatgraph.connected_components(graph) relevant_ccs = [] @@ -75,7 +106,7 @@ def _get_relevant_components(edges: np.ndarray, supervoxels: np.ndarray) -> Tupl # when merging, there must be only two components for cc_idx in ccs: cc = graph_ids[cc_idx] - if np.any(np.in1d(supervoxels, cc)): + if np.any(np.isin(svs, cc)): relevant_ccs.append(cc) assert len(relevant_ccs) == 2, "must be 2 components" return relevant_ccs @@ -89,9 +120,7 @@ def merge_preprocess( parent_ts: datetime.datetime = None, ) -> np.ndarray: """ - Determine if a fake edge needs to be added. - Get subgraph within the bounding box - Add fake edge if there are no inactive edges between two components. + Check and return inactive edges in the subgraph. """ edge_layers = cg.get_cross_chunk_edges_layer(subgraph_edges) active_edges = [types.empty_2d] @@ -108,19 +137,20 @@ def merge_preprocess( active_edges.append(active) inactive_edges.append(inactive) - relevant_ccs = _get_relevant_components(np.concatenate(active_edges), supervoxels) - inactive = np.concatenate(inactive_edges) + active_edges = np.concatenate(active_edges).astype(basetypes.NODE_ID) + inactive_edges = np.concatenate(inactive_edges).astype(basetypes.NODE_ID) + relevant_ccs = _get_relevant_components(active_edges, supervoxels) _inactive = [types.empty_2d] # source to sink edges - source_mask = np.in1d(inactive[:, 0], relevant_ccs[0]) - sink_mask = np.in1d(inactive[:, 1], relevant_ccs[1]) - _inactive.append(inactive[source_mask & sink_mask]) + source_mask = np.isin(inactive_edges[:, 0], relevant_ccs[0]) + sink_mask = np.isin(inactive_edges[:, 1], relevant_ccs[1]) + _inactive.append(inactive_edges[source_mask & sink_mask]) # sink to source edges - sink_mask = np.in1d(inactive[:, 1], relevant_ccs[0]) - source_mask = np.in1d(inactive[:, 0], relevant_ccs[1]) - _inactive.append(inactive[source_mask & sink_mask]) - _inactive = np.concatenate(_inactive) + sink_mask = np.isin(inactive_edges[:, 1], relevant_ccs[0]) + source_mask = np.isin(inactive_edges[:, 0], relevant_ccs[1]) + _inactive.append(inactive_edges[source_mask & sink_mask]) + _inactive = np.concatenate(_inactive).astype(basetypes.NODE_ID) return np.unique(_inactive, axis=0) if _inactive.size else types.empty_2d @@ -141,12 +171,15 @@ def check_fake_edges( time_stamp=parent_ts, ) ) - assert len(roots) == 2, "edges must be from 2 roots" - print("found inactive", len(inactive_edges)) + assert len(roots) == 2, ( + f"edges must be from 2 roots; got {len(roots)} " + f"({roots.tolist()}); inactive_edge_count={len(inactive_edges)}" + ) return inactive_edges, [] rows = [] supervoxels = atomic_edges.ravel() + # fake edges are stored with l2 chunks chunk_ids = cg.get_chunk_ids_from_node_ids( cg.get_parents(supervoxels, time_stamp=parent_ts) ) @@ -157,7 +190,7 @@ def check_fake_edges( val_dict[attributes.Connectivity.FakeEdges] = np.array( [[edge]], dtype=basetypes.NODE_ID ) - id1 = serialize_uint64(id1, fake_edges=True) + id1 = serializers.serialize_uint64(id1, fake_edges=True) rows.append( cg.client.mutate_row( id1, @@ -169,7 +202,7 @@ def check_fake_edges( val_dict[attributes.Connectivity.FakeEdges] = np.array( [edge[::-1]], dtype=basetypes.NODE_ID ) - id2 = serialize_uint64(id2, fake_edges=True) + id2 = serializers.serialize_uint64(id2, fake_edges=True) rows.append( cg.client.mutate_row( id2, @@ -177,7 +210,6 @@ def check_fake_edges( time_stamp=time_stamp, ) ) - print("no inactive", len(atomic_edges)) return atomic_edges, rows @@ -189,90 +221,152 @@ def add_edges( time_stamp: datetime.datetime = None, parent_ts: datetime.datetime = None, allow_same_segment_merge=False, + stitch_mode: bool = False, + do_sanity_check: bool = True, ): - edges, l2_atomic_cross_edges_d = _analyze_affected_edges( + edges, l2_cross_edges_d = _analyze_affected_edges( cg, atomic_edges, parent_ts=parent_ts ) l2ids = np.unique(edges) - if not allow_same_segment_merge: - assert ( - np.unique(cg.get_roots(l2ids, assert_roots=True, time_stamp=parent_ts)).size - == 2 - ), "L2 IDs must belong to different roots." - new_old_id_d, old_new_id_d, old_hierarchy_d = _init_old_hierarchy( - cg, l2ids, parent_ts=parent_ts - ) + if not allow_same_segment_merge and not stitch_mode: + roots = cg.get_roots(l2ids, assert_roots=True, time_stamp=parent_ts) + assert np.unique(roots).size >= 2, ( + f"L2 IDs must belong to different roots; " + f"l2ids={l2ids.tolist()} all share root={np.unique(roots).tolist()}; " + f"parent_ts={parent_ts} op={operation_id}" + ) + + new_old_id_d = defaultdict(set) + old_new_id_d = defaultdict(set) + old_hierarchy_d = _init_old_hierarchy(cg, l2ids, parent_ts=parent_ts) atomic_children_d = cg.get_children(l2ids) - atomic_cross_edges_d = merge_cross_edge_dicts( - cg.get_atomic_cross_edges(l2ids), l2_atomic_cross_edges_d + cross_edges_d = merge_cross_edge_dicts( + cg.get_cross_chunk_edges(l2ids, time_stamp=parent_ts), l2_cross_edges_d ) - graph, _, _, graph_ids = flatgraph.build_gt_graph(edges, make_directed=True) components = flatgraph.connected_components(graph) + + chunk_count_map = defaultdict(int) + for cc_indices in components: + l2ids_ = graph_ids[cc_indices] + chunk = cg.get_chunk_id(l2ids_[0]) + chunk_count_map[chunk] += 1 + + chunk_ids = list(chunk_count_map.keys()) + random.shuffle(chunk_ids) + chunk_new_ids_map = {} + for chunk_id in chunk_ids: + new_ids = cg.id_client.create_node_ids(chunk_id, size=chunk_count_map[chunk_id]) + chunk_new_ids_map[chunk_id] = list(new_ids) + new_l2_ids = [] for cc_indices in components: l2ids_ = graph_ids[cc_indices] - new_id = cg.id_client.create_node_id(cg.get_chunk_id(l2ids_[0])) - cg.cache.children_cache[new_id] = np.concatenate( - [atomic_children_d[l2id] for l2id in l2ids_] - ) - cg.cache.atomic_cx_edges_cache[new_id] = concatenate_cross_edge_dicts( - [atomic_cross_edges_d[l2id] for l2id in l2ids_] - ) - cache_utils.update( - cg.cache.parents_cache, cg.cache.children_cache[new_id], new_id - ) + new_id = chunk_new_ids_map[cg.get_chunk_id(l2ids_[0])].pop() new_l2_ids.append(new_id) new_old_id_d[new_id].update(l2ids_) for id_ in l2ids_: old_new_id_d[id_].add(new_id) - create_parents = CreateParentNodes( - cg, - new_l2_ids=new_l2_ids, - old_hierarchy_d=old_hierarchy_d, - new_old_id_d=new_old_id_d, - old_new_id_d=old_new_id_d, - operation_id=operation_id, - time_stamp=time_stamp, - parent_ts=parent_ts, - ) + # update cache + # map parent to new merged children and vice versa + merged_children = [atomic_children_d[l2id] for l2id in l2ids_] + merged_children = np.concatenate(merged_children).astype(basetypes.NODE_ID) + cg.cache.children_cache[new_id] = merged_children + cache_utils.update(cg.cache.parents_cache, merged_children, new_id) - new_roots = create_parents.run() - new_entries = create_parents.create_new_entries() - return new_roots, new_l2_ids, new_entries + # update cross chunk edges by replacing old_ids with new + # this can be done only after all new IDs have been created + for new_id, cc_indices in zip(new_l2_ids, components): + l2ids_ = graph_ids[cc_indices] + new_cx_edges_d = {} + cx_edges = [cross_edges_d[l2id] for l2id in l2ids_] + cx_edges_d = concatenate_cross_edge_dicts(cx_edges, unique=True) + temp_map = {k: next(iter(v)) for k, v in old_new_id_d.items()} + for layer, edges in cx_edges_d.items(): + edges = fastremap.remap(edges, temp_map, preserve_missing_labels=True) + new_cx_edges_d[layer] = edges + assert np.all(edges[:, 0] == new_id), ( + f"layer {layer} cross-edges[:, 0] must equal new_id={new_id}; " + f"got unique values {np.unique(edges[:, 0]).tolist()}" + ) + cg.cache.cross_chunk_edges_cache[new_id] = new_cx_edges_d + + profiler = get_profiler() + profiler.reset() + with profiler.profile("run"): + create_parents = CreateParentNodes( + cg, + new_l2_ids=new_l2_ids, + old_hierarchy_d=old_hierarchy_d, + new_old_id_d=new_old_id_d, + old_new_id_d=old_new_id_d, + operation_id=operation_id, + time_stamp=time_stamp, + parent_ts=parent_ts, + stitch_mode=stitch_mode, + do_sanity_check=do_sanity_check, + profiler=profiler, + ) + new_roots = create_parents.run() + if do_sanity_check: + sanity_check(cg, new_roots, operation_id) + create_parents.create_new_entries() + profiler.print_report(operation_id) + return new_roots, new_l2_ids, create_parents.new_entries -def _process_l2_agglomeration( + +def _split_l2_agglomeration( + cg, + operation_id: int, agg: types.Agglomeration, removed_edges: np.ndarray, - atomic_cross_edges_d: Dict[int, np.ndarray], + parent_ts: datetime.datetime = None, ): """ - For a given L2 id, remove given edges - and calculate new connected components. + For a given L2 id, remove given edges; calculate new connected components. """ chunk_edges = agg.in_edges.get_pairs() - cross_edges = np.concatenate([types.empty_2d, *atomic_cross_edges_d.values()]) chunk_edges = chunk_edges[~in2d(chunk_edges, removed_edges)] - cross_edges = cross_edges[~in2d(cross_edges, removed_edges)] - isolated_ids = agg.supervoxels[~np.in1d(agg.supervoxels, chunk_edges)] - isolated_edges = np.column_stack((isolated_ids, isolated_ids)) - graph, _, _, graph_ids = flatgraph.build_gt_graph( - np.concatenate([chunk_edges, isolated_edges]), make_directed=True + cross_edges = agg.cross_edges.get_pairs() + # we must avoid the cache to read roots to get segment state before edit began + parents = cg.get_parents(cross_edges[:, 0], time_stamp=parent_ts, raw_only=True) + + # if there are cross edges, there must be a single parent. + # if there aren't any, there must be no parents. XOR these 2 conditions. + err = ( + f"got cross edges from more than one l2 node; op {operation_id}; " + f"unique_parents={np.unique(parents).tolist()} " + f"cross_edges_count={cross_edges.shape[0]}" ) + assert (np.unique(parents).size == 1) != (cross_edges.size == 0), err + + if cross_edges.size: + # inactive edges must be filtered out + root = cg.get_root(parents[0], time_stamp=parent_ts, raw_only=True) + neighbor_roots = cg.get_roots( + cross_edges[:, 1], raw_only=True, time_stamp=parent_ts + ) + active_mask = neighbor_roots == root + cross_edges = cross_edges[active_mask] + cross_edges = cross_edges[~in2d(cross_edges, removed_edges)] + isolated_ids = agg.supervoxels[~np.isin(agg.supervoxels, chunk_edges)] + isolated_edges = np.column_stack((isolated_ids, isolated_ids)) + _edges = np.concatenate([chunk_edges, isolated_edges]).astype(basetypes.NODE_ID) + graph, _, _, graph_ids = flatgraph.build_gt_graph(_edges, make_directed=True) return flatgraph.connected_components(graph), graph_ids, cross_edges def _filter_component_cross_edges( - cc_ids: np.ndarray, cross_edges: np.ndarray, cross_edge_layers: np.ndarray + component_ids: np.ndarray, cross_edges: np.ndarray, cross_edge_layers: np.ndarray ) -> Dict[int, np.ndarray]: """ Filters cross edges for a connected component `cc_ids` from `cross_edges` of the complete chunk. """ - mask = np.in1d(cross_edges[:, 0], cc_ids) + mask = np.isin(cross_edges[:, 0], component_ids) cross_edges_ = cross_edges[mask] cross_edge_layers_ = cross_edge_layers[mask] edges_d = {} @@ -288,45 +382,67 @@ def remove_edges( cg, *, atomic_edges: Iterable[np.ndarray], - l2id_agglomeration_d: Dict, - operation_id: basetypes.OPERATION_ID = None, + operation_id: basetypes.OPERATION_ID = None, # type: ignore time_stamp: datetime.datetime = None, parent_ts: datetime.datetime = None, + do_sanity_check: bool = True, ): edges, _ = _analyze_affected_edges(cg, atomic_edges, parent_ts=parent_ts) l2ids = np.unique(edges) - assert ( - np.unique(cg.get_roots(l2ids, assert_roots=True, time_stamp=parent_ts)).size - == 1 - ), "L2 IDs must belong to same root." - new_old_id_d, old_new_id_d, old_hierarchy_d = _init_old_hierarchy( - cg, l2ids, parent_ts=parent_ts + roots = cg.get_roots(l2ids, assert_roots=True, time_stamp=parent_ts) + unique_roots, counts = np.unique(roots, return_counts=True) + assert unique_roots.size == 1, ( + f"L2 IDs must belong to same root; got root→l2_count=" + f"{dict(zip(unique_roots.tolist(), counts.tolist()))}; " + f"parent_ts={parent_ts} op={operation_id}" + ) + + l2id_agglomeration_d, _ = cg.get_l2_agglomerations( + l2ids, active=True, time_stamp=parent_ts ) - l2id_chunk_id_d = dict(zip(l2ids.tolist(), cg.get_chunk_ids_from_node_ids(l2ids))) - atomic_cross_edges_d = cg.get_atomic_cross_edges(l2ids) + new_old_id_d = defaultdict(set) + old_new_id_d = defaultdict(set) + old_hierarchy_d = _init_old_hierarchy(cg, l2ids, parent_ts=parent_ts) + chunk_id_map = dict(zip(l2ids.tolist(), cg.get_chunk_ids_from_node_ids(l2ids))) - removed_edges = np.concatenate([atomic_edges, atomic_edges[:, ::-1]], axis=0) + removed_edges = [atomic_edges, atomic_edges[:, ::-1]] + removed_edges = np.concatenate(removed_edges, axis=0).astype(basetypes.NODE_ID) new_l2_ids = [] for id_ in l2ids: - l2_agg = l2id_agglomeration_d[id_] - ccs, graph_ids, cross_edges = _process_l2_agglomeration( - l2_agg, removed_edges, atomic_cross_edges_d[id_] + agg = l2id_agglomeration_d[id_] + ccs, graph_ids, cross_edges = _split_l2_agglomeration( + cg, operation_id, agg, removed_edges, parent_ts ) - # calculated here to avoid repeat computation in loop + new_parents = cg.id_client.create_node_ids(chunk_id_map[agg.node_id], len(ccs)) + cross_edge_layers = cg.get_cross_chunk_edges_layer(cross_edges) - new_parent_ids = cg.id_client.create_node_ids( - l2id_chunk_id_d[l2_agg.node_id], len(ccs) - ) for i_cc, cc in enumerate(ccs): - new_id = new_parent_ids[i_cc] - cg.cache.children_cache[new_id] = graph_ids[cc] - cg.cache.atomic_cx_edges_cache[new_id] = _filter_component_cross_edges( - graph_ids[cc], cross_edges, cross_edge_layers - ) - cache_utils.update(cg.cache.parents_cache, graph_ids[cc], new_id) + new_id = new_parents[i_cc] new_l2_ids.append(new_id) new_old_id_d[new_id].add(id_) old_new_id_d[id_].add(new_id) + cg.cache.children_cache[new_id] = graph_ids[cc] + cache_utils.update(cg.cache.parents_cache, graph_ids[cc], new_id) + cg.cache.cross_chunk_edges_cache[new_id] = _filter_component_cross_edges( + graph_ids[cc], cross_edges, cross_edge_layers + ) + + cx_edges_d = cg.get_cross_chunk_edges(new_l2_ids, time_stamp=parent_ts) + for new_id in new_l2_ids: + new_cx_edges_d = cx_edges_d.get(new_id, {}) + for layer, edges in new_cx_edges_d.items(): + svs = np.unique(edges) + parents = cg.get_parents(svs, time_stamp=parent_ts) + temp_map = dict(zip(svs, parents)) + + edges = fastremap.remap(edges, temp_map, preserve_missing_labels=True) + edges = np.unique(edges, axis=0) + new_cx_edges_d[layer] = edges + assert np.all(edges[:, 0] == new_id), ( + f"layer {layer} cross-edges[:, 0] must equal new_id={new_id}; " + f"got unique values {np.unique(edges[:, 0]).tolist()}" + ) + cg.cache.cross_chunk_edges_cache[new_id] = new_cx_edges_d create_parents = CreateParentNodes( cg, @@ -337,10 +453,157 @@ def remove_edges( operation_id=operation_id, time_stamp=time_stamp, parent_ts=parent_ts, + do_sanity_check=do_sanity_check, ) new_roots = create_parents.run() - new_entries = create_parents.create_new_entries() - return new_roots, new_l2_ids, new_entries + + if do_sanity_check: + sanity_check(cg, new_roots, operation_id) + create_parents.create_new_entries() + return new_roots, new_l2_ids, create_parents.new_entries + + +def _get_descendants_batch(cg, node_ids): + """Get all descendants at layers >= 2 for multiple node_ids. + Batches get_children calls by level to reduce IO. + Returns dict {node_id: np.ndarray of descendants}. + """ + if not node_ids: + return {} + results = {nid: [] for nid in node_ids} + # expand_map: {node_to_expand: root_node_id} + expand_map = {nid: nid for nid in node_ids} + + while expand_map: + next_expand = {} + children_d = cg.get_children(list(expand_map.keys())) + for parent, root in expand_map.items(): + children = children_d[parent] + layers = cg.get_chunk_layers(children) + mask = layers >= 2 + results[root].extend(children[mask]) + for c in children[layers > 2]: + next_expand[c] = root + expand_map = next_expand + return { + nid: np.array(desc, dtype=basetypes.NODE_ID) for nid, desc in results.items() + } + + +def _get_counterparts( + cg, node_id: int, cx_edges_d: dict +) -> Tuple[List[int], Dict[int, int]]: + """ + Extract counterparts and their corresponding layers from cross chunk edges. + Returns (counterparts list, counterpart_layers dict). + """ + node_layer = cg.get_chunk_layer(node_id) + counterparts = [] + counterpart_layers = {} + for layer in range(node_layer, cg.meta.layer_count): + layer_edges = cx_edges_d.get(layer, types.empty_2d) + if layer_edges.size == 0: + continue + counterparts.extend(layer_edges[:, 1]) + layers_d = dict(zip(layer_edges[:, 1], [layer] * len(layer_edges[:, 1]))) + counterpart_layers.update(layers_d) + return counterparts, counterpart_layers + + +def _update_neighbor_cx_edges_single( + cg, + new_id: int, + node_map: dict, + counterpart_layers: dict, + all_counterparts_cx_edges_d: dict, + descendants_d: dict, +) -> dict: + """ + For each new_id, update cross chunk edges of its counterparts. + Some of them maybe updated multiple times so we need to collect them first + and then write to storage to consolidate the mutations. + Returns updated counterparts. + """ + node_layer = cg.get_chunk_layer(new_id) + counterparts = list(counterpart_layers.keys()) + cp_cx_edges_d = {cp: all_counterparts_cx_edges_d.get(cp, {}) for cp in counterparts} + updated_counterparts = {} + for counterpart, edges_d in cp_cx_edges_d.items(): + val_dict = {} + counterpart_layer = counterpart_layers[counterpart] + for layer in range(node_layer, cg.meta.layer_count): + edges = edges_d.get(layer, types.empty_2d) + if edges.size == 0: + continue + assert np.all(edges[:, 0] == counterpart), ( + f"layer {layer} cross-edges[:, 0] must equal counterpart={counterpart}; " + f"got unique values {np.unique(edges[:, 0]).tolist()}" + ) + edges = fastremap.remap(edges, node_map, preserve_missing_labels=True) + if layer == counterpart_layer: + flip_edge = np.array([counterpart, new_id], dtype=basetypes.NODE_ID) + edges = np.concatenate([edges, [flip_edge]]).astype(basetypes.NODE_ID) + descendants = descendants_d[new_id] + mask = np.isin(edges[:, 1], descendants) + if np.any(mask): + masked_edges = edges[mask] + masked_edges[:, 1] = new_id + edges[mask] = masked_edges + edges = np.unique(edges, axis=0) + edges_d[layer] = edges + val_dict[attributes.Connectivity.CrossChunkEdge[layer]] = edges + if not val_dict: + continue + cg.cache.cross_chunk_edges_cache[counterpart] = edges_d + updated_counterparts[counterpart] = val_dict + return updated_counterparts + + +def _update_neighbor_cx_edges( + cg, + new_ids: List[int], + new_old_id: dict, + old_new_id, + *, + time_stamp, + parent_ts, +) -> List: + """ + For each new_id, get counterparts and update its cross chunk edges. + Some of them maybe updated multiple times so we need to collect them first + and then write to storage to consolidate the mutations. + Returns mutations to updated counterparts/partner nodes. + """ + updated_counterparts = {} + newid_cx_edges_d = cg.get_cross_chunk_edges(new_ids, time_stamp=parent_ts) + node_map = {} + for k, v in old_new_id.items(): + if len(v) == 1: + node_map[k] = next(iter(v)) + + all_cps = set() + newid_counterpart_info = {} + for _id in new_ids: + counterparts, cp_layers = _get_counterparts(cg, _id, newid_cx_edges_d[_id]) + all_cps.update(counterparts) + newid_counterpart_info[_id] = cp_layers + + all_cx_edges_d = cg.get_cross_chunk_edges(list(all_cps), time_stamp=parent_ts) + descendants_d = _get_descendants_batch(cg, new_ids) + for new_id in new_ids: + m = {old_id: new_id for old_id in flip_ids(new_old_id, [new_id])} + node_map.update(m) + cp_layers = newid_counterpart_info[new_id] + result = _update_neighbor_cx_edges_single( + cg, new_id, node_map, cp_layers, all_cx_edges_d, descendants_d + ) + updated_counterparts.update(result) + updated_entries = [] + for node, val_dict in updated_counterparts.items(): + rowkey = serializers.serialize_uint64(node) + row = cg.client.mutate_row(rowkey, val_dict, time_stamp=time_stamp) + updated_entries.append(row) + return updated_entries class CreateParentNodes: @@ -349,32 +612,39 @@ def __init__( cg, *, new_l2_ids: Iterable, - operation_id: basetypes.OPERATION_ID, + operation_id: basetypes.OPERATION_ID, # type: ignore time_stamp: datetime.datetime, - new_old_id_d: Dict[np.uint64, Iterable[np.uint64]] = None, - old_new_id_d: Dict[np.uint64, Iterable[np.uint64]] = None, + new_old_id_d: Dict[np.uint64, Set[np.uint64]] = None, + old_new_id_d: Dict[np.uint64, Set[np.uint64]] = None, old_hierarchy_d: Dict[np.uint64, Dict[int, np.uint64]] = None, parent_ts: datetime.datetime = None, + stitch_mode: bool = False, + do_sanity_check: bool = True, + profiler: HierarchicalProfiler = None, ): self.cg = cg + self.new_entries = [] self._new_l2_ids = new_l2_ids self._old_hierarchy_d = old_hierarchy_d self._new_old_id_d = new_old_id_d self._old_new_id_d = old_new_id_d - self._new_ids_d = defaultdict(list) # new IDs in each layer - self._cross_edges_d = {} - self._operation_id = operation_id + self._new_ids_d = defaultdict(list) + self._opid = operation_id self._time_stamp = time_stamp - self._last_successful_ts = parent_ts + self._last_ts = parent_ts + self.stitch_mode = stitch_mode + self.do_sanity_check = do_sanity_check + self._profiler = profiler if profiler else get_profiler() def _update_id_lineage( self, - parent: basetypes.NODE_ID, + parent: basetypes.NODE_ID, # type: ignore children: np.ndarray, layer: int, parent_layer: int, ): - mask = np.in1d(children, self._new_ids_d[layer]) + # update newly created children; mask others + mask = np.isin(children, self._new_ids_d[layer]) for child_id in children[mask]: child_old_ids = self._new_old_id_d[child_id] for id_ in child_old_ids: @@ -382,90 +652,148 @@ def _update_id_lineage( self._new_old_id_d[parent].add(old_id) self._old_new_id_d[old_id].add(parent) - def _get_old_ids(self, new_ids): - old_ids = [ - np.array(list(self._new_old_id_d[id_]), dtype=basetypes.NODE_ID) - for id_ in new_ids - ] - return np.concatenate(old_ids) - - def _map_sv_to_parent(self, node_ids, layer, node_map=None): - sv_parent_d = {} - sv_cross_edges = [types.empty_2d] - if node_map is None: - node_map = {} + def _get_connected_components(self, node_ids: np.ndarray, layer: int): + cross_edges_d = self.cg.get_cross_chunk_edges( + node_ids, time_stamp=self._last_ts + ) + cx_edges = [types.empty_2d] for id_ in node_ids: - id_eff = node_map.get(id_, id_) - edges_ = self._cross_edges_d[id_].get(layer, types.empty_2d) - sv_parent_d.update(dict(zip(edges_[:, 0], [id_eff] * len(edges_)))) - sv_cross_edges.append(edges_) - return sv_parent_d, np.concatenate(sv_cross_edges) - - def _get_connected_components( - self, node_ids: np.ndarray, layer: int, lower_layer_ids: np.ndarray - ): - _node_ids = np.concatenate([node_ids, lower_layer_ids]) - cached = np.fromiter(self._cross_edges_d.keys(), dtype=basetypes.NODE_ID) - not_cached = _node_ids[~np.in1d(_node_ids, cached)] - - with TimeIt( - f"get_cross_chunk_edges.{layer}", - self.cg.graph_id, - self._operation_id, - ): - self._cross_edges_d.update( - self.cg.get_cross_chunk_edges(not_cached, all_layers=True) - ) - - sv_parent_d, sv_cross_edges = self._map_sv_to_parent(node_ids, layer) - get_sv_parents = np.vectorize(sv_parent_d.get, otypes=[np.uint64]) - try: - cross_edges = get_sv_parents(sv_cross_edges) - except TypeError: # NoneType error - # if there is a missing parent, try including lower layer ids - # this can happen due to skip connections - - # we want to map all these lower IDs to the current layer - lower_layer_to_layer = self.cg.get_roots( - lower_layer_ids, stop_layer=layer, ceil=False - ) - node_map = {k: v for k, v in zip(lower_layer_ids, lower_layer_to_layer)} - sv_parent_d, sv_cross_edges = self._map_sv_to_parent( - _node_ids, layer, node_map=node_map - ) - get_sv_parents = np.vectorize(sv_parent_d.get, otypes=[np.uint64]) - cross_edges = get_sv_parents(sv_cross_edges) + edges_ = cross_edges_d[id_].get(layer, types.empty_2d) + cx_edges.append(edges_) - cross_edges = np.concatenate([cross_edges, np.vstack([node_ids, node_ids]).T]) - graph, _, _, graph_ids = flatgraph.build_gt_graph( - cross_edges, make_directed=True - ) - return flatgraph.connected_components(graph), graph_ids + cx_edges = [*cx_edges, np.vstack([node_ids, node_ids]).T] + cx_edges = np.concatenate(cx_edges).astype(basetypes.NODE_ID) + graph, _, _, graph_ids = flatgraph.build_gt_graph(cx_edges, make_directed=True) + components = flatgraph.connected_components(graph) + return components, graph_ids def _get_layer_node_ids( self, new_ids: np.ndarray, layer: int ) -> Tuple[np.ndarray, np.ndarray]: # get old identities of new IDs - old_ids = self._get_old_ids(new_ids) + old_ids = flip_ids(self._new_old_id_d, new_ids) # get their parents, then children of those parents - node_ids = self.cg.get_children( - np.unique( - self.cg.get_parents(old_ids, time_stamp=self._last_successful_ts) - ), - flatten=True, - ) + old_parents = self.cg.get_parents(old_ids, time_stamp=self._last_ts) + siblings = self.cg.get_children(np.unique(old_parents), flatten=True) # replace old identities with new IDs - mask = np.in1d(node_ids, old_ids) - node_ids = np.concatenate( - [ - np.array(list(self._old_new_id_d[id_]), dtype=basetypes.NODE_ID) - for id_ in node_ids[mask] - ] - + [node_ids[~mask], new_ids] - ) + mask = np.isin(siblings, old_ids) + node_ids = [flip_ids(self._old_new_id_d, old_ids), siblings[~mask], new_ids] + node_ids = np.concatenate(node_ids).astype(basetypes.NODE_ID) node_ids = np.unique(node_ids) layer_mask = self.cg.get_chunk_layers(node_ids) == layer - return node_ids[layer_mask], node_ids[~layer_mask] + return node_ids[layer_mask] + + def _update_cross_edge_cache_batched(self, new_ids: list): + """ + Batch update cross chunk edges in cache for all new IDs at a layer. + """ + updated_entries = [] + if not new_ids: + return updated_entries + + parent_layer = self.cg.get_chunk_layer(new_ids[0]) + if parent_layer == 2: + # L2 cross edges have already been updated + return updated_entries + + all_children_d = self.cg.get_children(new_ids) + all_children = np.concatenate(list(all_children_d.values())) + all_cx_edges_raw = self.cg.get_cross_chunk_edges( + all_children, time_stamp=self._last_ts + ) + combined_cx_edges = concatenate_cross_edge_dicts(all_cx_edges_raw.values()) + with self._profiler.profile("latest"): + updated_cx_edges, edge_nodes = get_latest_edges_wrapper( + self.cg, combined_cx_edges, parent_ts=self._last_ts + ) + + # update cache with resolved stale edges + val_ds = defaultdict(dict) + children_cx_edges = defaultdict(dict) + for lyr in range(2, self.cg.meta.layer_count): + edges = updated_cx_edges.get(lyr, types.empty_2d) + if len(edges) == 0: + continue + children, inverse = np.unique(edges[:, 0], return_inverse=True) + masks = inverse == np.arange(len(children))[:, None] + for child, mask in zip(children, masks): + children_cx_edges[child][lyr] = edges[mask] + val_ds[child][attributes.Connectivity.CrossChunkEdge[lyr]] = edges[mask] + + for c, cx_edges_map in children_cx_edges.items(): + self.cg.cache.cross_chunk_edges_cache[c] = cx_edges_map + rowkey = serializers.serialize_uint64(c) + row = self.cg.client.mutate_row(rowkey, val_ds[c], time_stamp=self._last_ts) + updated_entries.append(row) + + # Distribute results back to each parent's cache + # Key insight: edges[:, 0] are children, map them to their parent + edge_parents = get_new_nodes(self.cg, edge_nodes, parent_layer, self._last_ts) + edge_parents_d = dict(zip(edge_nodes, edge_parents)) + for new_id in new_ids: + children_set = set(all_children_d[new_id]) + parent_cx_edges_d = {} + for layer in range(parent_layer, self.cg.meta.layer_count): + edges = updated_cx_edges.get(layer, types.empty_2d) + if len(edges) == 0: + continue + # Filter to edges whose source is one of this parent's children + mask = np.isin(edges[:, 0], list(children_set)) + if not np.any(mask): + continue + + pedges = edges[mask].copy() + pedges = fastremap.remap( + pedges, edge_parents_d, preserve_missing_labels=True + ) + parent_cx_edges_d[layer] = np.unique(pedges, axis=0) + assert np.all( + pedges[:, 0] == new_id + ), f"OP {self._opid}: mismatch {new_id} != {np.unique(pedges[:, 0])}" + self.cg.cache.cross_chunk_edges_cache[new_id] = parent_cx_edges_d + return updated_entries + + def _get_new_ids(self, chunk_id, count, is_root): + batch_size = count + new_ids = [] + while len(new_ids) < count: + candidate_ids = self.cg.id_client.create_node_ids( + chunk_id, batch_size, root_chunk=is_root + ) + existing = self.cg.client.read_nodes(node_ids=candidate_ids) + non_existing = set(candidate_ids) - existing.keys() + new_ids.extend(non_existing) + batch_size = min(batch_size * 2, 2**16) + return new_ids[:count] + + def _get_new_parents(self, layer, ccs, graph_ids) -> tuple[dict, dict]: + cc_layer_chunk_map = {} + size_map = defaultdict(int) + for i, cc_idx in enumerate(ccs): + parent_layer = layer + 1 # must be reset for each connected component + cc_ids = graph_ids[cc_idx] + if len(cc_ids) == 1: + # skip connection + parent_layer = self.cg.meta.layer_count + cx_edges_d = self.cg.get_cross_chunk_edges( + [cc_ids[0]], time_stamp=self._last_ts + ) + for l in range(layer + 1, self.cg.meta.layer_count): + if len(cx_edges_d[cc_ids[0]].get(l, types.empty_2d)) > 0: + parent_layer = l + break + chunk_id = self.cg.get_parent_chunk_id(cc_ids[0], parent_layer) + cc_layer_chunk_map[i] = (parent_layer, chunk_id) + size_map[chunk_id] += 1 + + chunk_ids = list(size_map.keys()) + random.shuffle(chunk_ids) + chunk_new_ids_map = {} + layers = self.cg.get_chunk_layers(chunk_ids) + for c, l in zip(chunk_ids, layers): + is_root = l == self.cg.meta.layer_count + chunk_new_ids_map[c] = self._get_new_ids(c, size_map[c], is_root) + return chunk_new_ids_map, cc_layer_chunk_map def _create_new_parents(self, layer: int): """ @@ -478,33 +806,35 @@ def _create_new_parents(self, layer: int): update parent old IDs """ new_ids = self._new_ids_d[layer] - layer_node_ids, lower_layer_ids = self._get_layer_node_ids(new_ids, layer) - components, graph_ids = self._get_connected_components( - layer_node_ids, layer, lower_layer_ids - ) - for cc_indices in components: - parent_layer = layer + 1 - cc_ids = graph_ids[cc_indices] - if len(cc_ids) == 1: - # skip connection - parent_layer = self.cg.meta.layer_count - for l in range(layer + 1, self.cg.meta.layer_count): - if len(self._cross_edges_d[cc_ids[0]].get(l, types.empty_2d)) > 0: - parent_layer = l - break + layer_node_ids = self._get_layer_node_ids(new_ids, layer) + ccs, _ids = self._get_connected_components(layer_node_ids, layer) + new_parents_map, cc_layer_chunk_map = self._get_new_parents(layer, ccs, _ids) + + for i, cc_indices in enumerate(ccs): + cc_ids = _ids[cc_indices] + parent_layer, chunk_id = cc_layer_chunk_map[i] + parent = new_parents_map[chunk_id].pop() + + self._new_ids_d[parent_layer].append(parent) + self._update_id_lineage(parent, cc_ids, layer, parent_layer) + self.cg.cache.children_cache[parent] = cc_ids + cache_utils.update(self.cg.cache.parents_cache, cc_ids, parent) + if not self.do_sanity_check: + continue - parent_id = self.cg.id_client.create_node_id( - self.cg.get_parent_chunk_id(cc_ids[0], parent_layer), - root_chunk=parent_layer == self.cg.meta.layer_count, - ) - self._new_ids_d[parent_layer].append(parent_id) - self.cg.cache.children_cache[parent_id] = cc_ids - cache_utils.update( - self.cg.cache.parents_cache, - cc_ids, - parent_id, - ) - self._update_id_lineage(parent_id, cc_ids, layer, parent_layer) + try: + sanity_check_single(self.cg, parent, self._opid) + except AssertionError: + pairs = [ + (a, b) for idx, a in enumerate(cc_ids) for b in cc_ids[idx + 1 :] + ] + for c1, c2 in pairs: + l2c1 = self.cg.get_l2children([c1]) + l2c2 = self.cg.get_l2children([c2]) + if np.intersect1d(l2c1, l2c2).size: + c = np.intersect1d(l2c1, l2c2) + msg = f"{self._opid}: {layer} {c1} {c2} common children {c}" + raise ValueError(msg) def run(self) -> Iterable: """ @@ -513,87 +843,123 @@ def run(self) -> Iterable: """ self._new_ids_d[2] = self._new_l2_ids for layer in range(2, self.cg.meta.layer_count): - if len(self._new_ids_d[layer]) == 0: + new_nodes = self._new_ids_d[layer] + if len(new_nodes) == 0: continue - with TimeIt( - f"create_new_parents_layer.{layer}", - self.cg.graph_id, - self._operation_id, - ): + self.cg.cache.new_ids.update(new_nodes) + # all new IDs in this layer have been created + # update their cross chunk edges and their neighbors' + with self._profiler.profile(f"l{layer}_update_cx_cache"): + entries = self._update_cross_edge_cache_batched(new_nodes) + self.new_entries.extend(entries) + + with self._profiler.profile(f"l{layer}_update_neighbor_cx"): + entries = _update_neighbor_cx_edges( + self.cg, + new_nodes, + self._new_old_id_d, + self._old_new_id_d, + time_stamp=self._time_stamp, + parent_ts=self._last_ts, + ) + self.new_entries.extend(entries) + with self._profiler.profile(f"l{layer}_create_new_parents"): self._create_new_parents(layer) return self._new_ids_d[self.cg.meta.layer_count] def _update_root_id_lineage(self): - new_root_ids = self._new_ids_d[self.cg.meta.layer_count] - former_root_ids = self._get_old_ids(new_root_ids) - former_root_ids = np.unique(former_root_ids) - assert ( - len(former_root_ids) < 2 or len(new_root_ids) < 2 - ), "Something went wrong." - rows = [] - for new_root_id in new_root_ids: + if self.stitch_mode: + return + new_roots = self._new_ids_d[self.cg.meta.layer_count] + former_roots = flip_ids(self._new_old_id_d, new_roots) + former_roots = np.unique(former_roots) + + err = f"new roots are inconsistent; op {self._opid}" + assert len(former_roots) < 2 or len(new_roots) < 2, err + for new_root_id in new_roots: val_dict = { - attributes.Hierarchy.FormerParent: np.array(former_root_ids), - attributes.OperationLogs.OperationID: self._operation_id, + attributes.Hierarchy.FormerParent: former_roots, + attributes.OperationLogs.OperationID: self._opid, } - rows.append( + self.new_entries.append( self.cg.client.mutate_row( - serialize_uint64(new_root_id), + serializers.serialize_uint64(new_root_id), val_dict, time_stamp=self._time_stamp, ) ) - for former_root_id in former_root_ids: + for former_root_id in former_roots: val_dict = { - attributes.Hierarchy.NewParent: np.array(new_root_ids), - attributes.OperationLogs.OperationID: self._operation_id, + attributes.Hierarchy.NewParent: np.array( + new_roots, dtype=basetypes.NODE_ID + ), + attributes.OperationLogs.OperationID: self._opid, } - rows.append( + self.new_entries.append( self.cg.client.mutate_row( - serialize_uint64(former_root_id), + serializers.serialize_uint64(former_root_id), val_dict, time_stamp=self._time_stamp, ) ) - return rows - def _get_atomic_cross_edges_val_dict(self): - new_ids = np.array(self._new_ids_d[2], dtype=basetypes.NODE_ID) + def _get_cross_edges_val_dicts(self): val_dicts = {} - atomic_cross_edges_d = self.cg.get_atomic_cross_edges(new_ids) - for id_ in new_ids: - val_dict = {} - for layer, edges in atomic_cross_edges_d[id_].items(): - val_dict[attributes.Connectivity.CrossChunkEdge[layer]] = edges - val_dicts[id_] = val_dict + for layer in range(2, self.cg.meta.layer_count): + new_ids = np.array(self._new_ids_d[layer], dtype=basetypes.NODE_ID) + cross_edges_d = self.cg.get_cross_chunk_edges( + new_ids, time_stamp=self._last_ts + ) + for id_ in new_ids: + val_dict = {} + for layer, edges in cross_edges_d[id_].items(): + val_dict[attributes.Connectivity.CrossChunkEdge[layer]] = edges + val_dicts[id_] = val_dict return val_dicts def create_new_entries(self) -> List: - rows = [] - val_dicts = self._get_atomic_cross_edges_val_dict() - for layer in range(2, self.cg.meta.layer_count + 1): + max_layer = self.cg.meta.layer_count + val_dicts = self._get_cross_edges_val_dicts() + for layer in range(2, max_layer + 1): new_ids = self._new_ids_d[layer] for id_ in new_ids: + if self.do_sanity_check: + root_layer = self.cg.get_chunk_layer(self.cg.get_root(id_)) + assert root_layer == max_layer, (id_, self.cg.get_root(id_)) + + if layer < max_layer: + try: + _parent = self.cg.get_parent(id_) + _children = self.cg.get_children(_parent) + assert id_ in _children, (layer, id_, _parent, _children) + except TypeError as e: + logger.error( + f"id={id_} parent={_parent} " + f"root={self.cg.get_root(id_)}" + ) + raise TypeError from e + val_dict = val_dicts.get(id_, {}) children = self.cg.get_children(id_) + err = f"parent layer less than children; op {self._opid}" assert np.max( self.cg.get_chunk_layers(children) - ) < self.cg.get_chunk_layer(id_), "Parent layer less than children." + ) < self.cg.get_chunk_layer(id_), err val_dict[attributes.Hierarchy.Child] = children - rows.append( + self.new_entries.append( self.cg.client.mutate_row( - serialize_uint64(id_), + serializers.serialize_uint64(id_), val_dict, time_stamp=self._time_stamp, ) ) for child_id in children: - rows.append( + self.new_entries.append( self.cg.client.mutate_row( - serialize_uint64(child_id), + serializers.serialize_uint64(child_id), {attributes.Hierarchy.Parent: id_}, time_stamp=self._time_stamp, ) ) - return rows + self._update_root_id_lineage() + self._update_root_id_lineage() diff --git a/pychunkedgraph/graph/err_dump.py b/pychunkedgraph/graph/err_dump.py new file mode 100644 index 000000000..03df8ff89 --- /dev/null +++ b/pychunkedgraph/graph/err_dump.py @@ -0,0 +1,113 @@ +"""Best-effort error artifact dumps for edit operations. + +Writes a JSON blob per failed operation to +`{WATERSHED}/graphene_errors/{cg.graph_id}/{op_id}.json` so logs stay +concise while the full payload (ids, edges, traceback) survives for +later inspection. `read_err_artifact(cg, op_id)` reads it back. +""" + +import traceback +from datetime import datetime + +import numpy as np + +from pychunkedgraph import get_logger + +logger = get_logger(__name__) + +_DUMP_ATTRS = ( + "source_ids", + "sink_ids", + "source_coords", + "sink_coords", + "added_edges", + "removed_edges", + "atomic_edges", + "affinities", + "bbox_offset", +) + + +def _json_safe(obj): + """Recursively coerce numpy + datetime into JSON-serializable types.""" + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, np.integer): + return int(obj) + if isinstance(obj, np.floating): + return float(obj) + if isinstance(obj, dict): + return {str(k): _json_safe(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_json_safe(v) for v in obj] + if isinstance(obj, datetime): + return obj.isoformat() + return obj + + +def _err_dir(cg): + """The `{WATERSHED}/graphene_errors/{cg.graph_id}` base; None if WATERSHED unset.""" + ws = getattr(cg.meta.data_source, "WATERSHED", None) + if not ws: + return None + return f"{ws.rstrip('/')}/graphene_errors/{cg.graph_id}" + + +def payload_summary(op) -> str: + """One-line operation source/sink summary for diagnostic log lines.""" + src = getattr(op, "source_ids", None) + snk = getattr(op, "sink_ids", None) + if src is None and snk is None: + return "" + return ( + f" source_ids={src.tolist() if src is not None else None}" + f" sink_ids={snk.tolist() if snk is not None else None}" + ) + + +def build_err_payload(op, op_id, err) -> dict: + """Structured operation snapshot capturing inputs + traceback for replay.""" + payload = { + "op_type": type(op).__name__, + "op_id": int(op_id), + "exception_class": type(err).__name__, + "exception_message": str(err), + "traceback": traceback.format_exc(), + "user_id": getattr(op, "user_id", None), + "parent_ts": getattr(op, "parent_ts", None), + } + for attr in _DUMP_ATTRS: + val = getattr(op, attr, None) + if val is not None: + payload[attr] = val + return payload + + +def dump_err_artifact(cg, op_id, payload): + """Write `{err_dir}/{op_id}.json`. Returns URL or None; never raises.""" + cf_dir = _err_dir(cg) + if cf_dir is None: + return None + try: + from cloudfiles import CloudFiles + + filename = f"{op_id}.json" + CloudFiles(cf_dir).put_json(filename, _json_safe(payload)) + return f"{cf_dir}/{filename}" + except Exception as e: + logger.warning(f"err_dump failed for op={op_id}: {e}") + return None + + +def read_err_artifact(cg, op_id): + """Return the artifact dict for this `op_id`, or None if missing.""" + cf_dir = _err_dir(cg) + if cf_dir is None: + return None + try: + from cloudfiles import CloudFiles + + return CloudFiles(cf_dir).get_json(f"{op_id}.json") + except Exception as e: + logger.warning(f"err_dump read failed for op={op_id}: {e}") + return None diff --git a/pychunkedgraph/graph/exceptions.py b/pychunkedgraph/graph/exceptions.py index 45aa57fc7..496f55e4f 100644 --- a/pychunkedgraph/graph/exceptions.py +++ b/pychunkedgraph/graph/exceptions.py @@ -1,23 +1,19 @@ from six.moves import http_client +from kvdbclient.exceptions import KVDBClientError +from kvdbclient.exceptions import LockingError +from kvdbclient.exceptions import PreconditionError -class ChunkedGraphError(Exception): - """Base class for all exceptions raised by the ChunkedGraph""" - pass - - -class LockingError(ChunkedGraphError): - """Raised when a Bigtable Lock could not be acquired""" - pass +class ChunkedGraphError(KVDBClientError): + """Base class for all exceptions raised by the ChunkedGraph""" -class PreconditionError(ChunkedGraphError): - """Raised when preconditions for Chunked Graph operations are not met""" pass class PostconditionError(ChunkedGraphError): """Raised when postconditions for Chunked Graph operations are not met""" + pass @@ -42,7 +38,7 @@ def __init__(self, message): self.message = message def __str__(self): - return f'[{self.status_code}]: {self.message}' + return f"[{self.status_code}]: {self.message}" class ClientError(ChunkedGraphAPIError): @@ -51,21 +47,25 @@ class ClientError(ChunkedGraphAPIError): class BadRequest(ClientError): """Exception mapping a ``400 Bad Request`` response.""" + status_code = http_client.BAD_REQUEST class Unauthorized(ClientError): """Exception mapping a ``401 Unauthorized`` response.""" + status_code = http_client.UNAUTHORIZED class Forbidden(ClientError): """Exception mapping a ``403 Forbidden`` response.""" + status_code = http_client.FORBIDDEN class Conflict(ClientError): """Exception mapping a ``409 Conflict`` response.""" + status_code = http_client.CONFLICT @@ -75,9 +75,29 @@ class ServerError(ChunkedGraphAPIError): class InternalServerError(ServerError): """Exception mapping a ``500 Internal Server Error`` response.""" + status_code = http_client.INTERNAL_SERVER_ERROR class GatewayTimeout(ServerError): """Exception mapping a ``504 Gateway Timeout`` response.""" + status_code = http_client.GATEWAY_TIMEOUT + + +class SupervoxelSplitRequiredError(ChunkedGraphError): + """ + Raised when supervoxel splitting is necessary. + Edit process should catch this error and retry after supervoxel has been split. + Saves remapping required for detecting which supervoxels need to be split. + """ + + def __init__( + self, + message: str, + sv_remapping: dict, + operation_id: int | None = None, + ): + super().__init__(message) + self.sv_remapping = sv_remapping + self.operation_id = operation_id diff --git a/pychunkedgraph/graph/lineage.py b/pychunkedgraph/graph/lineage.py index 6876ec563..957691fb9 100644 --- a/pychunkedgraph/graph/lineage.py +++ b/pychunkedgraph/graph/lineage.py @@ -1,26 +1,28 @@ """ Functions for tracking root ID changes over time. """ + +from __future__ import annotations + from typing import Union from typing import Optional from typing import Iterable -from datetime import datetime +from datetime import datetime, timezone from collections import defaultdict import numpy as np -from networkx import DiGraph -from . import attributes +from pychunkedgraph.graph import ( + attributes, + basetypes, + get_min_time, + get_max_time, + get_valid_timestamp, +) from .exceptions import ChunkedGraphError -from .attributes import Hierarchy -from .attributes import OperationLogs -from .utils.basetypes import NODE_ID -from .utils.generic import get_min_time -from .utils.generic import get_max_time -from .utils.generic import get_valid_timestamp -def get_latest_root_id(cg, root_id: NODE_ID.type) -> np.ndarray: +def get_latest_root_id(cg, root_id: basetypes.NODE_ID.type) -> np.ndarray: """Returns the latest root id associated with the provided root id""" id_working_set = [root_id] latest_root_ids = [] @@ -38,7 +40,7 @@ def get_latest_root_id(cg, root_id: NODE_ID.type) -> np.ndarray: def get_future_root_ids( cg, - root_id: NODE_ID, + root_id: basetypes.NODE_ID, time_stamp: Optional[datetime] = get_max_time(), ) -> np.ndarray: """ @@ -69,12 +71,12 @@ def get_future_root_ids( if next_id != root_id: id_history.append(next_id) next_ids = temp_next_ids - return np.unique(np.array(id_history, dtype=NODE_ID)) + return np.unique(np.array(id_history, dtype=basetypes.NODE_ID)) def get_past_root_ids( cg, - root_id: NODE_ID, + root_id: basetypes.NODE_ID, time_stamp: Optional[datetime] = get_min_time(), ) -> np.ndarray: """ @@ -108,12 +110,12 @@ def get_past_root_ids( if next_id != root_id: id_history.append(next_id) next_ids = temp_next_ids - return np.unique(np.array(id_history, dtype=NODE_ID)) + return np.unique(np.array(id_history, dtype=basetypes.NODE_ID)) def get_previous_root_ids( cg, - root_ids: Iterable[NODE_ID.type], + root_ids: Iterable[basetypes.NODE_ID.type], ) -> dict: """Returns immediate former root IDs (1 step history)""" nodes_d = cg.client.read_nodes( @@ -128,7 +130,7 @@ def get_previous_root_ids( def get_root_id_history( cg, - root_id: NODE_ID, + root_id: basetypes.NODE_ID, time_stamp_past: Optional[datetime] = get_min_time(), time_stamp_future: Optional[datetime] = get_max_time(), ) -> np.ndarray: @@ -140,18 +142,24 @@ def get_root_id_history( """ past_ids = get_past_root_ids(cg, root_id, time_stamp=time_stamp_past) future_ids = get_future_root_ids(cg, root_id, time_stamp=time_stamp_future) - return np.concatenate([past_ids, np.array([root_id], dtype=NODE_ID), future_ids]) + return np.concatenate( + [past_ids, np.array([root_id], dtype=basetypes.NODE_ID), future_ids] + ) def _get_node_properties(node_entry: dict) -> dict: node_d = {} - node_d["timestamp"] = node_entry[Hierarchy.Child][0].timestamp.timestamp() - if OperationLogs.OperationID in node_entry: - if len(node_entry[OperationLogs.OperationID]) == 2 or ( - len(node_entry[OperationLogs.OperationID]) == 1 - and Hierarchy.NewParent in node_entry + node_d["timestamp"] = node_entry[attributes.Hierarchy.Child][ + 0 + ].timestamp.timestamp() + if attributes.OperationLogs.OperationID in node_entry: + if len(node_entry[attributes.OperationLogs.OperationID]) == 2 or ( + len(node_entry[attributes.OperationLogs.OperationID]) == 1 + and attributes.Hierarchy.NewParent in node_entry ): - node_d["operation_id"] = node_entry[OperationLogs.OperationID][0].value + node_d["operation_id"] = node_entry[attributes.OperationLogs.OperationID][ + 0 + ].value return node_d @@ -166,15 +174,17 @@ def lineage_graph( going backwards in time until `timestamp_past` and in future until `timestamp_future` """ + from networkx import DiGraph + if not isinstance(node_ids, np.ndarray) and not isinstance(node_ids, list): node_ids = [node_ids] graph = DiGraph() - past_ids = np.array(node_ids, dtype=NODE_ID) - future_ids = np.array(node_ids, dtype=NODE_ID) + past_ids = np.array(node_ids, dtype=basetypes.NODE_ID) + future_ids = np.array(node_ids, dtype=basetypes.NODE_ID) timestamp_past = float(0) if timestamp_past is None else timestamp_past.timestamp() timestamp_future = ( - datetime.utcnow().timestamp() + datetime.now(timezone.utc).timestamp() if timestamp_future is None else timestamp_future.timestamp() ) @@ -190,10 +200,10 @@ def lineage_graph( graph.add_node(k, **node_d) if ( node_d["timestamp"] < timestamp_past - or not Hierarchy.FormerParent in val + or not attributes.Hierarchy.FormerParent in val ): continue - former_ids = val[Hierarchy.FormerParent][0].value + former_ids = val[attributes.Hierarchy.FormerParent][0].value next_past_ids.extend( [former_id for former_id in former_ids if not former_id in graph.nodes] ) @@ -206,7 +216,10 @@ def lineage_graph( val = nodes_raw[k] node_d = _get_node_properties(val) graph.add_node(k, **node_d) - if node_d["timestamp"] > timestamp_future or not Hierarchy.NewParent in val: + if ( + node_d["timestamp"] > timestamp_future + or not attributes.Hierarchy.NewParent in val + ): continue try: future_operation_id_dict[node_d["operation_id"]].append(k) @@ -215,13 +228,13 @@ def lineage_graph( logs_raw = cg.client.read_log_entries(list(future_operation_id_dict.keys())) for operation_id in future_operation_id_dict: - new_ids = logs_raw[operation_id][OperationLogs.RootID] + new_ids = logs_raw[operation_id][attributes.OperationLogs.RootID] next_future_ids.extend( [new_id for new_id in new_ids if not new_id in graph.nodes] ) for new_id in new_ids: for k in future_operation_id_dict[operation_id]: graph.add_edge(k, new_id) - past_ids = np.array(np.unique(next_past_ids), dtype=NODE_ID) - future_ids = np.array(np.unique(next_future_ids), dtype=NODE_ID) + past_ids = np.array(np.unique(next_past_ids), dtype=basetypes.NODE_ID) + future_ids = np.array(np.unique(next_future_ids), dtype=basetypes.NODE_ID) return graph diff --git a/pychunkedgraph/graph/locks.py b/pychunkedgraph/graph/locks.py index b3a3a0eb7..62df03e5e 100644 --- a/pychunkedgraph/graph/locks.py +++ b/pychunkedgraph/graph/locks.py @@ -1,12 +1,19 @@ -from typing import Union -from typing import Sequence +import hashlib +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Sequence, Union from collections import defaultdict import numpy as np -from . import exceptions +from pychunkedgraph import get_logger + +from . import attributes, exceptions, serializers from .types import empty_1d -from .lineage import get_future_root_ids +from .lineage import lineage_graph +from .dry_run import is_dry_run + +logger = get_logger(__name__) class RootLock: @@ -22,6 +29,7 @@ class RootLock: "lock_acquired", "operation_id", "privileged_mode", + "future_root_ids_d", ] # FIXME: `locked_root_ids` is only required and exposed because `cg.client.lock_roots` # currently might lock different (more recent) root IDs than requested. @@ -44,25 +52,35 @@ def __init__( # caused by failed writes. Must be used with `operation_id`, # meaning only existing failed operations can be run this way. self.privileged_mode = privileged_mode + self.future_root_ids_d = defaultdict(lambda: empty_1d) def __enter__(self): - if self.privileged_mode: - assert self.operation_id is not None, "Please provide operation ID." - from warnings import warn - - warn("Warning: Privileged mode without acquiring lock.") - return self if not self.operation_id: self.operation_id = self.cg.id_client.create_operation_id() - future_root_ids_d = defaultdict(lambda: empty_1d) + if is_dry_run(): + return self + + if self.privileged_mode: + return self + + import networkx as nx + + nodes_ts = self.cg.get_node_timestamps(self.root_ids, return_numpy=0) + min_ts = min(nodes_ts) + lgraph = lineage_graph(self.cg, self.root_ids, timestamp_past=min_ts) + self.future_root_ids_d = defaultdict(lambda: empty_1d) for id_ in self.root_ids: - future_root_ids_d[id_] = get_future_root_ids(self.cg, id_) + node_descendants = nx.descendants(lgraph, id_) + node_descendants = np.unique( + np.array(list(node_descendants), dtype=np.uint64) + ) + self.future_root_ids_d[id_] = node_descendants self.lock_acquired, self.locked_root_ids = self.cg.client.lock_roots( root_ids=self.root_ids, operation_id=self.operation_id, - future_root_ids_d=future_root_ids_d, + future_root_ids_d=self.future_root_ids_d, max_tries=7, ) if not self.lock_acquired: @@ -70,9 +88,22 @@ def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): + if is_dry_run(): + return if self.lock_acquired: - for locked_root_id in self.locked_root_ids: - self.cg.client.unlock_root(locked_root_id, self.operation_id) + max_workers = min(8, max(1, len(self.locked_root_ids))) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + unlock_futures = [ + executor.submit( + self.cg.client.unlock_root, root_id, self.operation_id + ) + for root_id in self.locked_root_ids + ] + for future in as_completed(unlock_futures): + try: + future.result() + except Exception as e: + logger.warning(f"Failed to unlock root: {e}") class IndefiniteRootLock: @@ -87,7 +118,14 @@ class IndefiniteRootLock: or when it has already been locked indefinitely. """ - __slots__ = ["cg", "root_ids", "acquired", "operation_id", "privileged_mode"] + __slots__ = [ + "cg", + "root_ids", + "acquired", + "operation_id", + "privileged_mode", + "future_root_ids_d", + ] def __init__( self, @@ -95,6 +133,7 @@ def __init__( operation_id: np.uint64, root_ids: Union[np.uint64, Sequence[np.uint64]], privileged_mode: bool = False, + future_root_ids_d=None, ) -> None: self.cg = cg self.operation_id = operation_id @@ -104,31 +143,450 @@ def __init__( # This is intended to be used in extremely rare cases to fix errors # caused by failed writes. self.privileged_mode = privileged_mode + self.future_root_ids_d = future_root_ids_d def __enter__(self): + if is_dry_run(): + return self if self.privileged_mode: - from warnings import warn - - warn("Warning: Privileged mode without acquiring indefinite lock.") return self if not self.cg.client.renew_locks(self.root_ids, self.operation_id): raise exceptions.LockingError("Could not renew locks before writing.") - future_root_ids_d = defaultdict(lambda: empty_1d) - for id_ in self.root_ids: - future_root_ids_d[id_] = get_future_root_ids(self.cg, id_) + if self.future_root_ids_d is None: + import networkx as nx + + nodes_ts = self.cg.get_node_timestamps(self.root_ids, return_numpy=0) + min_ts = min(nodes_ts) + lgraph = lineage_graph(self.cg, self.root_ids, timestamp_past=min_ts) + self.future_root_ids_d = defaultdict(lambda: empty_1d) + for id_ in self.root_ids: + node_descendants = nx.descendants(lgraph, id_) + node_descendants = np.unique( + np.array(list(node_descendants), dtype=np.uint64) + ) + self.future_root_ids_d[id_] = node_descendants + self.acquired, self.root_ids, failed = self.cg.client.lock_roots_indefinitely( root_ids=self.root_ids, operation_id=self.operation_id, - future_root_ids_d=future_root_ids_d, + future_root_ids_d=self.future_root_ids_d, ) if not self.acquired: - raise exceptions.LockingError(f"{failed} has been locked indefinitely.") + raise exceptions.LockingError(f"{failed} have been locked indefinitely.") return self def __exit__(self, exception_type, exception_value, traceback): + if is_dry_run(): + return + if exception_type is not None: + # Partial bigtable hierarchy writes may have landed before + # the exception propagated. Keep the indefinite cells held + # so subsequent ops on these roots refuse to acquire — + # forces operator recovery (`repair_operation(..., unlock= + # True)`) rather than letting a silent corruption slip into + # further edits. + return if self.acquired: - for locked_root_id in self.root_ids: - self.cg.client.unlock_indefinitely_locked_root( - locked_root_id, self.operation_id + max_workers = min(8, max(1, len(self.root_ids))) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + unlock_futures = [ + executor.submit( + self.cg.client.unlock_indefinitely_locked_root, + root_id, + self.operation_id, + ) + for root_id in self.root_ids + ] + for future in as_completed(unlock_futures): + try: + future.result() + except Exception as e: + logger.warning(f"Failed to unlock root: {e}") + + +def _downsample_block_lock_row_key(block_coord) -> bytes: + """Row key for one pyramid_block's downsample lock cell. + + Hash-prefixed so spatially-clustered block coords — common when a + team edits the same region — scatter across bigtable tablets instead + of piling up in one lexicographic range, which would hot-spot a + single tablet under concurrent load. + + 26 bytes total: + - 2-byte blake2b hash of the packed coord (tablet distribution). + - 24 bytes of packed coord (big-endian uint64 per axis). + uint64 per axis tracks the existing node-id width and puts no cap on + the block grid. The full coord in the key guarantees uniqueness even + if two coords share the 2-byte hash prefix. + """ + bx, by, bz = (int(c) for c in block_coord) + packed = ( + bx.to_bytes(8, "big", signed=False) + + by.to_bytes(8, "big", signed=False) + + bz.to_bytes(8, "big", signed=False) + ) + return hashlib.blake2b(packed, digest_size=2).digest() + packed + + +class DownsampleBlockLock: + """Lock a set of pyramid_blocks for the lifetime of a downsample task. + + The downsample worker holds one across read → tinybrain → write for + every block it touches. All-or-nothing: on partial acquisition we + release what we got and retry with backoff; on repeated failure we + raise so the pubsub message ends up un-acked and redelivered. + + Uses `cg.client.lock_by_row_key` with hash-prefixed row keys — the + generic row-key lock primitive in kvdbclient — so these rows never + collide with node-id-keyed root locks even though both use the same + `Concurrency.Lock` column. + """ + + __slots__ = ["cg", "block_coords", "operation_id", "acquired_keys"] + + # Retry budget for partial-acquire failures. Each attempt releases + # anything it got in the previous pass, then re-acquires from scratch. + _MAX_ACQUIRE_ATTEMPTS = 7 + _ACQUIRE_BACKOFF_BASE_SEC = 0.5 + + def __init__( + self, + cg, + block_coords: Sequence, + operation_id: np.uint64, + ) -> None: + self.cg = cg + # Sort so every `__enter__` uses a consistent acquisition order + # across workers — reduces contention between workers whose block + # sets overlap. Sort is on the coord tuple (not the hashed row + # key) so the order is stable and debuggable. + self.block_coords = sorted( + (int(bx), int(by), int(bz)) for bx, by, bz in block_coords + ) + self.operation_id = np.uint64(operation_id) + self.acquired_keys: list = [] + + def __enter__(self): + for attempt in range(self._MAX_ACQUIRE_ATTEMPTS): + self.acquired_keys = [] + all_ok = True + for coord in self.block_coords: + row_key = _downsample_block_lock_row_key(coord) + if self.cg.client.lock_by_row_key(row_key, self.operation_id): + self.acquired_keys.append(row_key) + else: + all_ok = False + break + if all_ok: + return self + self._release_acquired() + time.sleep(self._ACQUIRE_BACKOFF_BASE_SEC * (2**attempt)) + raise exceptions.LockingError( + f"Could not acquire downsample block locks for coords " + f"{self.block_coords} after {self._MAX_ACQUIRE_ATTEMPTS} attempts" + ) + + def __exit__(self, exception_type, exception_value, traceback): + self._release_acquired() + + def _release_acquired(self): + if not self.acquired_keys: + return + max_workers = min(8, max(1, len(self.acquired_keys))) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit( + self.cg.client.unlock_by_row_key, key, self.operation_id + ) + for key in self.acquired_keys + ] + for future in as_completed(futures): + try: + future.result() + except Exception as e: + logger.warning(f"Failed to unlock downsample block: {e}") + self.acquired_keys = [] + + def renew(self) -> bool: + """Extend expiry on every held lock. Returns False if any failed.""" + ok = True + for key in self.acquired_keys: + if not self.cg.client.renew_lock_by_row_key(key, self.operation_id): + logger.warning(f"Failed to renew downsample block lock {key!r}") + ok = False + return ok + + +def _l2_chunk_lock_row_key(chunk_id) -> bytes: + """Row key for one L2 chunk's spatial lock cell. + + Hash-prefixed so spatially-clustered chunk IDs scatter across + bigtable tablets instead of piling up in one lexicographic range, + which would hot-spot a single tablet under concurrent load. + + 10 bytes total: + - 2-byte blake2b hash of the chunk_id (tablet distribution). + - 8 bytes of big-endian uint64 chunk_id. + chunk_id already encodes layer+xyz in its bits, so the full key is + unique per L2 chunk. + """ + packed = int(chunk_id).to_bytes(8, "big", signed=False) + return hashlib.blake2b(packed, digest_size=2).digest() + packed + + +class L2ChunkLock: + """Lock a set of L2 chunks to serialize SV splits that touch them. + + Closes the cross-root spatial race: two SV splits on overlapping L2 + chunks but distinct roots acquire disjoint root-lock sets and would + otherwise race on seg state. This lock is held across the + `split_supervoxel` loop (seg write + SV-level hierarchy row write) + so the pair commits atomically. + + All-or-nothing: on partial acquisition we release what we got and + retry with backoff; on repeated failure we raise `LockingError`. + + Uses `cg.client.lock_by_row_key` — the generic row-key lock in + kvdbclient — with a row-key namespace distinct from root and + downsample block locks (all three share `attributes.Concurrency.Lock` + under the hood; the row key disambiguates). + """ + + __slots__ = [ + "cg", + "chunk_ids", + "operation_id", + "privileged_mode", + "acquired_keys", + ] + + # Retry budget for partial-acquire failures. Each attempt releases + # anything it got in the previous pass, then re-acquires from scratch. + _MAX_ACQUIRE_ATTEMPTS = 7 + _ACQUIRE_BACKOFF_BASE_SEC = 0.5 + + def __init__( + self, + cg, + chunk_ids: Sequence[int], + operation_id: np.uint64, + *, + privileged_mode: bool = False, + ) -> None: + self.cg = cg + # Sort so every `__enter__` uses a consistent acquisition order + # across workers — reduces contention when overlapping lock sets + # would otherwise race AB/BA. + self.chunk_ids = sorted(int(c) for c in chunk_ids) + self.operation_id = np.uint64(operation_id) + self.privileged_mode = privileged_mode + self.acquired_keys: list = [] + + def __enter__(self): + if is_dry_run(): + return self + if self.privileged_mode: + # Replay path: the crashed op's `IndefiniteL2ChunkLock` cells + # are still set on these chunks (that's what's blocking new + # ops), and `lock_by_row_key_with_indefinite` would refuse. + # Mirror `RootLock`/`IndefiniteRootLock`'s privileged escape + # hatch — skip temporal acquire, the indefinite cells are + # our de-facto lock and they'll be released by the inner + # `IndefiniteL2ChunkLock(privileged_mode=True)` on exit. + return self + for attempt in range(self._MAX_ACQUIRE_ATTEMPTS): + self.acquired_keys = [] + all_ok = True + for chunk_id in self.chunk_ids: + row_key = _l2_chunk_lock_row_key(chunk_id) + # `_with_indefinite`: the temporal acquire must also + # refuse if the indefinite column is set. Closes the + # crash-recovery race — a worker that died holding + # `IndefiniteL2ChunkLock` leaves the indefinite cell + # set, and the next op must see it rather than silently + # racing into partial state. + if self.cg.client.lock_by_row_key_with_indefinite( + row_key, self.operation_id + ): + self.acquired_keys.append(row_key) + else: + all_ok = False + break + if all_ok: + return self + self._release_acquired() + time.sleep(self._ACQUIRE_BACKOFF_BASE_SEC * (2**attempt)) + raise exceptions.LockingError( + f"Could not acquire L2 chunk locks for chunks {self.chunk_ids} " + f"after {self._MAX_ACQUIRE_ATTEMPTS} attempts" + ) + + def __exit__(self, exception_type, exception_value, traceback): + if is_dry_run(): + return + self._release_acquired() + + def _release_acquired(self): + if not self.acquired_keys: + return + max_workers = min(8, max(1, len(self.acquired_keys))) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit( + self.cg.client.unlock_by_row_key, key, self.operation_id + ) + for key in self.acquired_keys + ] + for future in as_completed(futures): + try: + future.result() + except Exception as e: + logger.warning(f"Failed to unlock L2 chunk: {e}") + self.acquired_keys = [] + + def renew(self) -> bool: + """Extend expiry on every held lock. Returns False if any failed.""" + ok = True + for key in self.acquired_keys: + if not self.cg.client.renew_lock_by_row_key(key, self.operation_id): + logger.warning(f"Failed to renew L2 chunk lock {key!r}") + ok = False + return ok + + +class IndefiniteL2ChunkLock: + """Upgrade held-temporal L2 chunk locks to indefinite. + + Structurally mirrors `IndefiniteRootLock`: acquired inside the + temporal lock (`L2ChunkLock`) context after preconditions are + established, and held across the write phase. Doesn't expire — the + cell persists on bigtable until explicitly released (or operator + recovery clears it), so a worker that dies with writes in flight + leaves the chunks marked indefinitely-held. + + The temporal `L2ChunkLock` must already be held by the same + `operation_id`; the acquire filter for temporal now rejects on + indefinite cells, so future temporal acquires on these chunks + refuse until this lock is released. + + Durable scope: `__enter__` writes `chunk_ids` to the op-log row's + `OperationLogs.L2ChunkLockScope` column. This persists through a + worker crash, giving `stuck_ops replay` the exact chunk set to + clean up without a bigtable-wide lock-row scan. + + `privileged_mode=True` is the operator recovery escape hatch: + skips the acquire step (the cells already exist, held by this same + op_id from the crashed attempt), pre-populates `acquired_keys` from + `chunk_ids` so `__exit__` still value-matches-releases those cells, + and does not re-write the op-log scope column. + """ + + __slots__ = ["cg", "chunk_ids", "operation_id", "privileged_mode", "acquired_keys"] + + def __init__( + self, + cg, + chunk_ids: Sequence[int], + operation_id: np.uint64, + *, + privileged_mode: bool = False, + ) -> None: + self.cg = cg + self.chunk_ids = sorted(int(c) for c in chunk_ids) + self.operation_id = np.uint64(operation_id) + self.privileged_mode = privileged_mode + self.acquired_keys: list = [] + + def __enter__(self): + if is_dry_run(): + return self + if self.privileged_mode: + # Recovery path: crashed op's indefinite cells already exist + # under this op_id. Populate acquired_keys so __exit__'s + # value-matched release deletes them after the replay writes + # succeed. + self.acquired_keys = [_l2_chunk_lock_row_key(c) for c in self.chunk_ids] + return self + for chunk_id in self.chunk_ids: + row_key = _l2_chunk_lock_row_key(chunk_id) + if not self.cg.client.lock_by_row_key_indefinitely( + row_key, self.operation_id + ): + # Partial acquire: release what we got and fail. No + # retry — an indefinite cell belongs to a currently- + # running or crashed op and won't clear on its own. + self._release_acquired() + raise exceptions.LockingError( + f"Could not upgrade L2 chunk {chunk_id} to indefinite lock " + f"(another op holds it)" + ) + self.acquired_keys.append(row_key) + self._write_scope_to_op_log() + return self + + def __exit__(self, exception_type, exception_value, traceback): + if is_dry_run(): + return + if exception_type is not None: + # Partial OCDBT seg / bigtable SV-hierarchy writes may have + # landed before the exception propagated. Leave the + # indefinite cells held and the op-log scope intact so + # subsequent ops refuse at `L2ChunkLock` acquire — forces + # operator recovery (`stuck_ops replay`) rather than + # leaking orphan SV IDs into downstream reads. + return + self._release_acquired() + self._clear_scope_on_op_log() + + def _write_scope_to_op_log(self): + """Record the chunk scope on the op-log row before seg/bigtable + writes begin. A worker crash after this point leaves both the + per-chunk indefinite cells AND this field set, so recovery can + locate the partial-write region without a bigtable scan. + """ + row_key = serializers.serialize_uint64(self.operation_id) + scope = np.asarray(self.chunk_ids, dtype=np.uint64) + entry = self.cg.client.mutate_row( + row_key, + {attributes.OperationLogs.L2ChunkLockScope: scope}, + ) + self.cg.client.write([entry]) + + def _clear_scope_on_op_log(self): + """Clear the scope record on normal exit — op completed or was + cleanly rolled back, so no partial state needs recovery. Overwrites + with an empty array; a subsequent `read_log_entries` returns an + empty scope (recovery skips). Best-effort; failures here are + logged but don't propagate. + """ + try: + row_key = serializers.serialize_uint64(self.operation_id) + empty = np.array([], dtype=np.uint64) + entry = self.cg.client.mutate_row( + row_key, + {attributes.OperationLogs.L2ChunkLockScope: empty}, + ) + self.cg.client.write([entry]) + except Exception as e: + logger.warning(f"Failed to clear L2ChunkLockScope on op-log row: {e}") + + def _release_acquired(self): + if not self.acquired_keys: + return + max_workers = min(8, max(1, len(self.acquired_keys))) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit( + self.cg.client.unlock_indefinitely_locked_by_row_key, + key, + self.operation_id, ) + for key in self.acquired_keys + ] + for future in as_completed(futures): + try: + future.result() + except Exception as e: + logger.warning(f"Failed to unlock indefinite L2 chunk: {e}") + self.acquired_keys = [] diff --git a/pychunkedgraph/graph/meta.py b/pychunkedgraph/graph/meta.py index 83d670ffe..9f93e2f4d 100644 --- a/pychunkedgraph/graph/meta.py +++ b/pychunkedgraph/graph/meta.py @@ -2,20 +2,26 @@ from datetime import timedelta from typing import Dict from typing import List -from typing import Tuple from typing import Sequence from collections import namedtuple import numpy as np -from cloudvolume import CloudVolume +import tensorstore as ts + +from pychunkedgraph.graph.ocdbt import ( + OcdbtConfig, + build_cg_ocdbt_spec, + ensure_fork_synced, + fork_base_manifest, + fork_exists, + get_seg_source_and_destination_ocdbt, + read_populate_meta, +) from .utils.generic import compute_bitmasks from .chunks.utils import get_chunks_boundary -from ..utils.redis import keys as r_keys -from ..utils.redis import get_rq_queue from ..utils.redis import get_redis_connection - _datasource_fields = ("EDGES", "COMPONENTS", "WATERSHED", "DATA_VERSION", "CV_MIP") _datasource_defaults = (None, None, None, None, 0) DataSource = namedtuple( @@ -52,6 +58,29 @@ ) +def _redis_cached_json(key: str, loader): + """Return JSON-decoded value at ``key`` in Redis, or call ``loader()`` and + write the result through. Spares distributed workers from re-fetching the + same GCS object on every CG instantiation. Silently bypasses Redis if it + is unreachable; returns ``loader()`` directly in that case. + """ + redis = None + try: + redis = get_redis_connection() + cached = redis.get(key) + if cached is not None: + return json.loads(cached) + except Exception: + redis = None + value = loader() + if value is not None and redis is not None: + try: + redis.set(key, json.dumps(value)) + except Exception: + ... + return value + + class ChunkedGraphMeta: def __init__( self, graph_config: GraphConfig, data_source: DataSource, custom_data: Dict = {} @@ -64,9 +93,23 @@ def __init__( self._custom_data = custom_data self._ws_cv = None + self._ws_ts_scales = {} + self._ws_info_d = None + # Multi-scale OCDBT handles + per-scale resolutions, populated lazily + # from source's info JSON. ws_ocdbt returns scale 0 for backward + # compatibility; ws_ocdbt_scales exposes the full pyramid. + self._ws_ocdbt_scales = None + self._ws_ocdbt_resolutions = None self._layer_bounds_d = None self._layer_count = None self._bitmasks = None + self._ocdbt_seg = None + self._ocdbt_config_cached = None + + @property + def graph_id(self): + assert self._graph_config.ID is not None, "graph_id required" + return self._graph_config.ID_PREFIX + self._graph_config.ID @property def graph_config(self): @@ -80,29 +123,163 @@ def data_source(self): def custom_data(self): return self._custom_data + def for_copied_graph(self, graph_id: str) -> "ChunkedGraphMeta": + """Rewrite this meta in place for a table copied/restored under ``graph_id`` — its + graph id and mesh dirs — so the copy's meshes never alias the source's; returns self + for a one-line ``update_meta`` call.""" + gc = self._graph_config._asdict() + gc["ID"] = graph_id + self._graph_config = GraphConfig(**gc) + mesh = self._custom_data.get("mesh") + if mesh and "dir" in mesh: + # Only an explicit graph-suffixed dynamic_mesh_dir shares initial meshes; a bare + # "dynamic" or an unset value defaults to a private per-graph top-level dir, so a + # copy can never alias the source. + rewrite = ( + shared_initial_mesh_dirs + if mesh.get("dynamic_mesh_dir") not in (None, "dynamic") + else private_mesh_dirs + ) + mesh["dir"], mesh["dynamic_mesh_dir"] = rewrite(mesh["dir"], graph_id) + return self + @property def ws_cv(self): + """Watershed CloudVolume — back-compat hatch (meshing / diagnostics).""" if self._ws_cv: return self._ws_cv + from cloudvolume import CloudVolume - cache_key = f"{self.graph_config.ID}:ws_cv_info_cached" - try: - # try reading a cached info file for distributed workers - # useful to avoid md5 errors on high gcs load - redis = get_redis_connection() - cached_info = json.loads(redis.get(cache_key)) - self._ws_cv = CloudVolume(self._data_source.WATERSHED, info=cached_info) - except Exception: - self._ws_cv = CloudVolume(self._data_source.WATERSHED) - try: - redis.set(cache_key, json.dumps(self._ws_cv.info)) - except Exception: - ... + ws = self._data_source.WATERSHED + info = _redis_cached_json( + f"ws_cv_info_cached:{ws}", + lambda: CloudVolume(ws, progress=False).info, + ) + self._ws_cv = CloudVolume(ws, info=info, progress=False) return self._ws_cv + def ws_ts_scale(self, mip: int): + """Watershed handle (tensorstore neuroglancer_precomputed) at scale ``mip``.""" + if mip not in self._ws_ts_scales: + ws = self._data_source.WATERSHED.rstrip("/") + self._ws_ts_scales[mip] = ts.open( + { + "driver": "neuroglancer_precomputed", + "kvstore": ws, + "scale_index": mip, + } + ).result() + return self._ws_ts_scales[mip] + + @property + def ws_ts(self): + """Watershed handle at base scale (mip 0).""" + return self.ws_ts_scale(0) + + @property + def _ws_info(self): + """Watershed precomputed ``info`` JSON, Redis-cached.""" + if self._ws_info_d is None: + # Base must not end in '/'; the leading '/' in '/info' supplies the + # separator — otherwise the GCS read returns empty. + ws = self._data_source.WATERSHED.rstrip("/") + self._ws_info_d = _redis_cached_json( + f"ws_info_cached:{ws}", + lambda: json.loads( + ts.KvStore.open(ws).result().read("/info").result().value + ), + ) + return self._ws_info_d + + @property + def ocdbt_config(self) -> OcdbtConfig: + """Per-CG OCDBT settings with precedence info-file > custom_data > defaults. + + The watershed's ``/ocdbt/.populated/meta.json`` is the authoritative + on-disk source for fields that affect the OCDBT format (compression, + max_inline_value_bytes, populate_layer). custom_data fills per-CG + fields (enabled, sv_split_threshold) and anything the info file + doesn't pin. Both layers fall through to dataclass defaults. + + The info-file fetch goes through a Redis cache (same pattern as + ``ws_cv``) so distributed workers don't re-read the same GCS + object on every CG instantiation. Result is also cached in + instance state after first access. Legacy ``custom_data["seg"]`` + shape is read when ``"ocdbt_config"`` is absent so pre-refactor + CGs still open. + """ + if self._ocdbt_config_cached is not None: + return self._ocdbt_config_cached + + meta_d = self._custom_data.get("ocdbt_config") + if meta_d is None: + seg = self._custom_data.get("seg", {}) + meta_d = { + "enabled": bool(seg.get("ocdbt", False)), + "sv_split_threshold": int(seg.get("sv_split_threshold", 10)), + } + + info_d = None + ws = self._data_source.WATERSHED + if ws: + info_d = _redis_cached_json( + f"ocdbt_info_cached:{ws}", + lambda: read_populate_meta(ws), + ) + + self._ocdbt_config_cached = OcdbtConfig.resolve(meta_d, info_d) + return self._ocdbt_config_cached + + @property + def ocdbt_seg(self) -> bool: + if self._ocdbt_seg is None: + self._ocdbt_seg = self.ocdbt_config.enabled + return self._ocdbt_seg + + @property + def ws_ocdbt(self): + """Base scale (MIP 0) handle. Backward-compatible single-handle access.""" + return self.ws_ocdbt_scales[0] + + @property + def ws_ocdbt_scales(self): + """List of TensorStore handles, one per MIP level. Lazily initialized. + + Opens the CG's delta OCDBT via the kvstack-layered fork spec — reads + merge the shared base + this CG's edits, writes go to the delta. + """ + assert self.ocdbt_seg, "make sure this pcg has segmentation in ocdbt format" + if self._ws_ocdbt_scales is None: + ws = self.data_source.WATERSHED + # Auto-create the fork on first open if missing — e.g. after a + # bigtable copy that gave us a new graph_id. Idempotent and + # race-safe: concurrent opens write identical base-manifest + # bytes to the same path. Can't race with an edit because an + # edit pre-supposes the fork exists. + if not fork_exists(ws, self.graph_id): + fork_base_manifest(ws, self.graph_id) + # Refresh the fork manifest from base if it's stale and edit-free. + # See ensure_fork_synced docstring; without this, post-fork-creation + # populate writes to base are invisible through the kvstack view + # and reads return zeros. + ensure_fork_synced(ws, self.graph_id) + _, self._ws_ocdbt_scales, self._ws_ocdbt_resolutions = ( + get_seg_source_and_destination_ocdbt( + ws, self.graph_id, self.ocdbt_config + ) + ) + return self._ws_ocdbt_scales + + @property + def ws_ocdbt_resolutions(self): + """Per-scale [x,y,z] resolutions (used to derive downsample factors).""" + # Trigger lazy init via ws_ocdbt_scales — both are populated together. + _ = self.ws_ocdbt_scales + return self._ws_ocdbt_resolutions + @property def resolution(self): - return self.ws_cv.resolution # pylint: disable=no-member + return np.array(self._ws_info["scales"][0]["resolution"]) @property def layer_count(self) -> int: @@ -110,8 +287,6 @@ def layer_count(self) -> int: if self._layer_count: return self._layer_count - bbox = np.array(self.ws_cv.bounds.to_list()) # pylint: disable=no-member - bbox = bbox.reshape(2, 3) n_chunks = get_chunks_boundary( self.voxel_counts, np.array(self._graph_config.CHUNK_SIZE, dtype=int) ) @@ -145,18 +320,14 @@ def bitmasks(self): @property def voxel_bounds(self): - bounds = np.array(self.ws_cv.bounds.to_list()) # pylint: disable=no-member - return bounds.reshape(2, -1).T + s0 = self._ws_info["scales"][0] + vo = np.array(s0["voxel_offset"]) + return np.array([vo, vo + np.array(s0["size"])]).T @property def voxel_counts(self) -> Sequence[int]: """returns number of voxels in each dimension""" - cv_bounds = np.array(self.ws_cv.bounds.to_list()) # pylint: disable=no-member - cv_bounds = cv_bounds.reshape(2, -1).T - voxel_counts = cv_bounds.copy() - voxel_counts -= cv_bounds[:, 0:1] # pylint: disable=unsubscriptable-object - voxel_counts = voxel_counts[:, 1] - return voxel_counts + return np.array(self._ws_info["scales"][0]["size"]) @property def layer_chunk_bounds(self) -> Dict: @@ -225,6 +396,10 @@ def edge_dtype(self): def READ_ONLY(self): return self.custom_data.get("READ_ONLY", False) + @property + def sv_split_threshold(self) -> int: + return self.ocdbt_config.sv_split_threshold + @property def split_bounding_offset(self): return self.custom_data.get( @@ -234,8 +409,7 @@ def split_bounding_offset(self): @property def dataset_info(self) -> Dict: - info = self.ws_cv.info # pylint: disable=no-member - + info = dict(self._ws_info) info.update( { "chunks_start_at_voxel_offset": True, @@ -247,6 +421,24 @@ def dataset_info(self) -> Dict: "cv_mip": self.data_source.CV_MIP, "n_layers": self.layer_count, "spatial_bit_masks": self.bitmasks, + "ocdbt_seg": self.ocdbt_seg, + # Full kvstore spec a reader hands to tensorstore's + # `neuroglancer_precomputed` driver. Server owns the + # contract — paths, data prefixes, and OCDBT config + # (e.g. `max_inline_value_bytes`) are all resolved + # here, so readers don't duplicate configuration and + # future schema changes are picked up on re-fetch. + # Readers pass this verbatim as `kvstore`; add a + # `version` field for time-travel reads. + "ocdbt_kvstore_spec": ( + build_cg_ocdbt_spec( + self._data_source.WATERSHED, + self.graph_id, + self.ocdbt_config, + ) + if self.ocdbt_seg and self._graph_config.ID + else None + ), }, } ) @@ -288,3 +480,17 @@ def is_out_of_bounds(self, chunk_coordinate): return np.any(chunk_coordinate < 0) or np.any( chunk_coordinate > 2 ** self.bitmasks[1] ) + + +def private_mesh_dirs(mesh_dir: str, graph_id: str) -> tuple[str, str]: + """(dir, dynamic_mesh_dir) for a copied table whose meshes are all its own (source + dynamic dir unset or the bare "dynamic"). Suffix the top-level dir per graph so even + initial meshes stay private; the dynamic subdir keeps the bare "dynamic" inside it.""" + return f"{mesh_dir}_{graph_id}", "dynamic" + + +def shared_initial_mesh_dirs(mesh_dir: str, graph_id: str) -> tuple[str, str]: + """(dir, dynamic_mesh_dir) for a copied table that shares initial meshes with siblings + from the same backup — the source dynamic dir was graph-suffixed. Keep the top-level dir + shared; re-derive only the dynamic subdir per graph.""" + return mesh_dir, f"dynamic_{graph_id}" diff --git a/pychunkedgraph/graph/misc.py b/pychunkedgraph/graph/misc.py index b33e8a6fd..a9d1fbcac 100644 --- a/pychunkedgraph/graph/misc.py +++ b/pychunkedgraph/graph/misc.py @@ -8,10 +8,9 @@ import fastremap import numpy as np -from multiwrapper import multiprocessing_utils as mu from . import ChunkedGraph -from . import attributes +from pychunkedgraph.graph import attributes from .edges import Edges from .utils import flatgraph from .types import Agglomeration @@ -51,22 +50,6 @@ def _read_delta_root_rows( return new_root_ids, expired_root_ids -def _read_root_rows_thread(args) -> list: - start_seg_id, end_seg_id, serialized_cg_info, time_stamp = args - cg = ChunkedGraph(**serialized_cg_info) - start_id = cg.get_node_id(segment_id=start_seg_id, chunk_id=cg.root_chunk_id) - end_id = cg.get_node_id(segment_id=end_seg_id, chunk_id=cg.root_chunk_id) - rows = cg.client.read_nodes( - start_id=start_id, - end_id=end_id, - end_id_inclusive=False, - end_time=time_stamp, - end_time_inclusive=True, - ) - root_ids = [k for (k, v) in rows.items() if attributes.Hierarchy.NewParent not in v] - return root_ids - - def get_proofread_root_ids( cg: ChunkedGraph, start_time: Optional[datetime.datetime] = None, @@ -94,43 +77,12 @@ def get_proofread_root_ids( def get_latest_roots( - cg, time_stamp: Optional[datetime.datetime] = None, n_threads: int = 1 + cg: ChunkedGraph, time_stamp: Optional[datetime.datetime] = None, n_threads: int = 1 ) -> Sequence[np.uint64]: - # Create filters: time and id range - max_seg_id = cg.get_max_seg_id(cg.root_chunk_id) + 1 - n_blocks = 1 if n_threads == 1 else int(np.min([n_threads * 3 + 1, max_seg_id])) - seg_id_blocks = np.linspace(1, max_seg_id, n_blocks + 1, dtype=np.uint64) - cg_serialized_info = cg.get_serialized_info() - if n_threads > 1: - del cg_serialized_info["credentials"] - - multi_args = [] - for i_id_block in range(0, len(seg_id_blocks) - 1): - multi_args.append( - [ - seg_id_blocks[i_id_block], - seg_id_blocks[i_id_block + 1], - cg_serialized_info, - time_stamp, - ] - ) - - if n_threads == 1: - results = mu.multiprocess_func( - _read_root_rows_thread, - multi_args, - n_threads=n_threads, - verbose=False, - debug=n_threads == 1, - ) - else: - results = mu.multisubprocess_func( - _read_root_rows_thread, multi_args, n_threads=n_threads - ) - root_ids = [] - for result in results: - root_ids.extend(result) - return np.array(root_ids, dtype=np.uint64) + root_chunk = cg.get_chunk_id(layer=cg.meta.layer_count, x=0, y=0, z=0) + rr = cg.range_read_chunk(root_chunk, time_stamp=time_stamp) + roots = [k for k, v in rr.items() if attributes.Hierarchy.NewParent not in v] + return np.array(roots, dtype=np.uint64) def get_delta_roots( @@ -190,7 +142,7 @@ def get_contact_sites( ) # Build area lookup dictionary - cs_svs = edges[~np.in1d(edges, sv_ids).reshape(-1, 2)] + cs_svs = edges[~np.isin(edges, sv_ids)] area_dict = collections.defaultdict(int) for area, sv_id in zip(areas, cs_svs): @@ -202,7 +154,6 @@ def get_contact_sites( # Load edges of these cs_svs edges_cs_svs_rows = cg.client.read_nodes( node_ids=u_cs_svs, - # columns=[attributes.Connectivity.Partner, attributes.Connectivity.Connected], ) pre_cs_edges = [] for ri in edges_cs_svs_rows.items(): @@ -214,7 +165,7 @@ def get_contact_sites( cs_dict = collections.defaultdict(list) for cc in ccs: cc_sv_ids = unique_ids[cc] - cc_sv_ids = cc_sv_ids[np.in1d(cc_sv_ids, u_cs_svs)] + cc_sv_ids = cc_sv_ids[np.isin(cc_sv_ids, u_cs_svs)] cs_areas = area_dict_vec(cc_sv_ids) partner_root_id = ( int(cg.get_root(cc_sv_ids[0], time_stamp=time_stamp)) diff --git a/pychunkedgraph/graph/ocdbt/TENSORSTORE_REFERENCE.md b/pychunkedgraph/graph/ocdbt/TENSORSTORE_REFERENCE.md new file mode 100644 index 000000000..7faf4f140 --- /dev/null +++ b/pychunkedgraph/graph/ocdbt/TENSORSTORE_REFERENCE.md @@ -0,0 +1,134 @@ +# tensorstore OCDBT reference + +Every entry below was verified by probing tensorstore directly (intentional-bad-value + spec round-trip) against the binary in this workspace's venv. Re-verify if the tensorstore version changes. + +## OCDBT kvstore spec — top-level fields + +Sibling of `driver: "ocdbt"`: + +| Field | Type | Default | Notes | +|---|---|---|---| +| `base` | kvstore spec or URL | — | underlying kvstore (gcs/file/s3/…) | +| `manifest` | kvstore spec or URL | (under `base`) | the manifest *can* live in a separate kvstore from data | +| `config` | object | `{}` | see Config sub-fields below | +| `assume_config` | bool | `false` | skip reading the existing config from the manifest (use with care) | +| `coordinator` | ocdbt_coordinator resource | named ref `"ocdbt_coordinator"` | enables distributed mode when set | +| `cache_pool` | cache_pool resource | named ref `"cache_pool"` | | +| `data_copy_concurrency` | data_copy_concurrency resource | named ref | | +| `target_data_file_size` | uint64 | driver default | when a single commit's d/ writes exceed this, the writer rolls a new d/ file | +| `experimental_read_coalescing_threshold_bytes` | uint64 | — | | +| `experimental_read_coalescing_merged_bytes` | uint64 | — | | +| `experimental_read_coalescing_interval` | uint64 | — | | +| `btree_node_data_prefix` | string | `"d/"` | path prefix for btree-node files | +| `value_data_prefix` | string | `"d/"` | path prefix for value files | +| `version_tree_node_data_prefix` | string | `"d/"` | path prefix for version-tree files | +| `path` | string | `""` | sub-prefix in the kvstore | + +**Not fields**: `data_file_prefixes`, `version_spec`, `recheck_cached*`, `transaction`, `btree_writer_concurrency`, `manifest_kind` (lives under `config`). + +## OCDBT `config` sub-fields + +| Field | Type | tensorstore default | Notes | +|---|---|---|---| +| `compression` | object | `{}` (none) | `{"id": "zstd", "level": N}` — zstd level 1–22 | +| `max_inline_value_bytes` | uint64 | `100` | values ≤ this size live inline in the btree leaf bytes; larger values get written to a d/ file and the mutation carries only an `IndirectDataReference`. In distributed mode this **directly bounds cooperator-forwarded RPC size**: inline values are carried inside the `WriteRequest.mutations` field, so a leaf's batch blows past the 4 MiB gRPC max-receive whenever multiple inline values pile up on one node. Source: `distributed/btree_writer.cc` `StagePending`. Setting low (≤ a few KB) pushes chunk values out-of-line → small mutations → small RPCs. | +| `max_decoded_node_bytes` | uint64 | `8388608` (8 MiB) | btree node split threshold. Larger nodes → shallower tree → fewer per-commit node touches. Setting this *smaller* than the default INCREASES per-commit forwarded bytes — empirically went from ~8 MiB to ~23 MiB RPCs when set to 1 MiB. | +| `version_tree_arity_log2` | int | — | controls version tree branching; rarely tuned | +| `manifest_kind` | enum | `"single"` | `"single"` or `"numbered"` (manifest history retained — needed for time-travel reads) | +| `uuid` | string | (auto) | 32-hex per-base UUID assigned at create time | + +**Not fields**: `data_file_prefixes`, `data_file_prefix`, `btree_node_arity_log2`, `version_tree_node_arity`. + +## `ocdbt_coordinator` context resource + +| Field | Type | Default | Notes | +|---|---|---|---| +| `address` | string | — | `"host:port"` of the DistributedCoordinatorServer | +| `lease_duration` | duration string (`"1s"`, `"500ms"`, etc.) | — | how long a lease holder owns a btree node | +| `security` | object | `{method: "insecure"}` | requires `method` key. This build has **no** security methods registered (build flag) — all calls cleartext. | + +## `DistributedCoordinatorServer({...})` + +| Field | Type | Default | Notes | +|---|---|---|---| +| `bind_addresses` | list[string] | one ephemeral port | gRPC server bind address(es). `.port` after construction gives the ephemeral port. | +| `security` | object | insecure | same shape as the resource's security | + +**There is NO Python knob for the gRPC server's max-receive message size.** The 4 MiB default is set inside tensorstore's gRPC server builder. Confirmed by strings on the binary: no `TENSORSTORE_*` env var, no spec/resource field, no Context resource that maps to `grpc.max_receive_message_length`. + +## Distributed vs non-distributed write paths + +The OCDBT driver picks one of two compiled implementations at open time: + +- **non-distributed** (`btree_writer.cc`): coordinator absent. Each commit writes the manifest itself. Concurrent writers race the manifest CAS; losers retry; their pre-commit d/ writes become orphans. +- **distributed** (`distributed/btree_writer.cc`, `cooperator_*.cc`): coordinator present. One lease holder per btree node serializes commits. Other cooperators **forward their mutations over gRPC** to the lease holder. + +### Constraints unique to distributed mode + +1. **`ts.Transaction(atomic=True)` is incompatible.** "Cannot read/write … as single atomic transaction" — verified on (info + chunk) and on (cross-key). A plain `ts.Transaction()` still batches all writes into one OCDBT commit; only the *atomicity* across keys is lost. +2. **Cooperator-forwarded RPC ≤ ~4 MiB.** Carries (btree node delta) + (value bytes for keys committed into that node). +3. **Disjoint user-key writes still trigger forwarding.** Leases are per btree node, not per user-key range. Two workers writing distinct keys into the same node → one forwards to the other. + +## Cooperator batching + +`cooperator_submit_mutation_batch.cc` `SendToPeer` is the gRPC sender. The `WriteRequest` proto has `repeated bytes mutations` — each entry is one encoded `BtreeNodeWriteMutation` destined for the same leaf. The encoded mutation embeds the value_reference inline if it's an `absl::Cord`, or carries just an `IndirectDataReference` (small struct) otherwise. So **what's actually on the wire per RPC = (small request header) + Σ encoded mutations**, and each encoded mutation's size is dominated by its value bytes IF the value is inline. + +Threshold for inline-vs-ref is `max_inline_value_bytes` (see config table). That's the real lever for RPC size. + +What changes RPC size (verified by production dumps): +- `max_inline_value_bytes=1 MiB`, default node bytes → RPCs 5–8 MiB (inline chunks pile up in the batch) +- `max_inline_value_bytes=1 MiB` + `max_decoded_node_bytes=1 MiB` → RPCs up to 23 MiB (smaller nodes ≠ smaller RPCs) +- `max_inline_value_bytes=1 MiB` + dst `chunk_size` halved → RPCs grew to 12 MiB (more mutations per node → bigger batches) +- `max_inline_value_bytes=4 KiB` (chunks go out-of-line) → mutations carry only refs; RPC = small header + N×(key + ref + generation) → fits 4 MiB regardless of value sizes (this is the path our code takes) + +## Defaults visible from spec round-trip + +```json +{ + "assume_config": false, + "btree_node_data_prefix": "d/", + "config": {}, + "coordinator": "ocdbt_coordinator", + "cache_pool": "cache_pool", + "data_copy_concurrency": "data_copy_concurrency", + "value_data_prefix": "d/", + "version_tree_node_data_prefix": "d/" +} +``` + +## Env vars + +- `OCDBT_COORDINATOR_HOST`, `OCDBT_COORDINATOR_PORT`: **NO EFFECT**. Not referenced anywhere in the binary. Address must go in spec's `coordinator.address`. +- `TENSORSTORE_VERBOSE_LOGGING`: comma-separated tag list to stderr. Tags include `ocdbt`, `coordinator`. + +Other `TENSORSTORE_*` vars exist (CA paths, S3/GCS concurrency, etc.) — grep the binary. + +## On-disk layout + +- `manifest.ocdbt` at the base — root btree node + current data file refs. +- `d/` — directory of "data files" each holding concatenated values + (optionally) btree node bytes + version-tree node bytes. +- Each commit creates **at least one** d/ file holding all values + nodes for that commit, then a CAS-update of `manifest.ocdbt`. +- `target_data_file_size` controls when a single commit splits its d/ writes across files. + +## How this maps onto pychunkedgraph + +- `OcdbtConfig` (`pychunkedgraph/graph/ocdbt/meta.py`) → `compression: zstd 12`, `max_inline_value_bytes = 4 KiB`. The 4 KiB threshold keeps small metadata (info JSON, populate markers) inline while forcing every chunk value out-of-line into d/ files — this is what keeps cooperator RPCs under the 4 MiB gRPC ceiling. +- `create_base_ocdbt` is the **only** path that embeds `config.ts_config()` in its kvstore spec — that write persists the values into `manifest.ocdbt`. Every open path (`open_base_ocdbt`, `build_cg_ocdbt_spec`) omits the `config` block: tensorstore would otherwise assert our in-code defaults against the on-disk manifest and raise `FAILED_PRECONDITION` on any drift, bricking every existing base. On-disk wins. +- `populate_chunk` (`pychunkedgraph/ingest/ocdbt.py`) opens the base with `coordinator_address` (distributed mode). +- `copy_ws_bbox_multiscale` uses **non-atomic** `ts.Transaction()` because of the distributed-mode constraint above. +- `_dump_failure_to_gcs` writes JSON failure forensics when `ERROR_DUMP` env is set. + +## Empirically tried and ruled out + +- `OCDBT_COORDINATOR_HOST/PORT` env vars — no effect. +- Bumping gRPC max-receive via env / channel arg / spec field — no such knob. +- Smaller `dst chunk_size` alone — RPC size grew (more mutations per node). +- Smaller `max_decoded_node_bytes` alone — RPC size grew (more per-commit node touches). +- `--ocdbt-edges` legacy path — decommissioned, removed. +- `ts.Transaction(atomic=True)` with distributed coordinator — incompatible. + +## Open observations (not verified at production scale) + +- `lease_duration` may reduce cross-cooperator forwarding if held long enough that a worker's whole task lands on its own nodes. +- `target_data_file_size` may affect manifest growth but not RPC size. +- Switching dst encoding from `compressed_segmentation` to `raw` would make per-value size predictable (`chunk_volume × bytes_per_voxel`), bypassing the dense-region pathological CS encoding (one observed key encoded to 23 MiB at 256×256×64). diff --git a/pychunkedgraph/graph/ocdbt/__init__.py b/pychunkedgraph/graph/ocdbt/__init__.py new file mode 100644 index 000000000..633964e1c --- /dev/null +++ b/pychunkedgraph/graph/ocdbt/__init__.py @@ -0,0 +1,57 @@ +"""Public API for the OCDBT-backed segmentation store. + +See ``main.py`` for the architectural notes. This module just re-exports +the names that external callers (ingest, edits, runtime, tests) reach for. +""" + +from .meta import OcdbtConfig +from .utils import ( + _layer_bbox, + _read_source_scales, + base_exists, + fork_exists, + is_chunk_populated, + mark_chunk_populated, + read_populate_meta, + write_populate_meta, +) +from .main import ( + _mode_downsample, + build_cg_ocdbt_spec, + copy_ws_bbox_multiscale, + copy_ws_chunk, + copy_ws_chunk_multiscale, + create_base_ocdbt, + ensure_fork_synced, + fork_base_manifest, + get_seg_source_and_destination_ocdbt, + open_base_ocdbt, + propagate_to_coarser_scales, + wipe_base_ocdbt, + write_seg_chunks, +) + +__all__ = [ + "OcdbtConfig", + "_layer_bbox", + "_mode_downsample", + "_read_source_scales", + "base_exists", + "build_cg_ocdbt_spec", + "copy_ws_bbox_multiscale", + "copy_ws_chunk", + "copy_ws_chunk_multiscale", + "create_base_ocdbt", + "ensure_fork_synced", + "fork_base_manifest", + "fork_exists", + "get_seg_source_and_destination_ocdbt", + "is_chunk_populated", + "mark_chunk_populated", + "open_base_ocdbt", + "propagate_to_coarser_scales", + "read_populate_meta", + "wipe_base_ocdbt", + "write_populate_meta", + "write_seg_chunks", +] diff --git a/pychunkedgraph/graph/ocdbt/debug.py b/pychunkedgraph/graph/ocdbt/debug.py new file mode 100644 index 000000000..b64376695 --- /dev/null +++ b/pychunkedgraph/graph/ocdbt/debug.py @@ -0,0 +1,143 @@ +"""Diagnostic plumbing for OCDBT failures. + +Humanize-count for log lines, generic failure envelope (host/pod/versions/ +traceback/timestamp), bbox-failure payload builder, and a GCS dump helper +that writes per-task forensic JSON under ``$ERROR_DUMP/__.json``. +Kept out of ``main.py`` and ``utils.py`` so the core OCDBT code stays +free of import bloat that's only used on failure paths. +""" + +import json +import logging +import os +import socket +import sys +import traceback +from datetime import datetime, timezone +from os import environ +from typing import Optional + +import tensorstore as ts + +_logger = logging.getLogger(__name__) + + +def humanize_count(n: int) -> str: + """Compact count for log lines: 1234567 → '1.2M', 950 → '950'.""" + for unit, scale in (("G", 1_000_000_000), ("M", 1_000_000), ("K", 1_000)): + if n >= scale: + return f"{n / scale:.1f}{unit}" + return str(n) + + +def failure_envelope(exc: BaseException, dump_tag: Optional[str]) -> dict: + """Generic metadata for any failure dump — host, pod, versions, + timestamp, traceback, coordinator env. Caller merges with the + failure-specific fields to build the final payload. + """ + return { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "dump_tag": dump_tag, + "host": { + "hostname": socket.gethostname(), + "pid": os.getpid(), + "pod_name": environ.get("MY_POD_NAME"), + "pod_ip": environ.get("MY_POD_IP"), + "node_name": environ.get("MY_NODE_NAME"), + }, + "versions": { + "tensorstore": getattr(ts, "__version__", None), + "python": sys.version, + }, + "ocdbt_coordinator_env": { + "OCDBT_COORDINATOR_HOST": environ.get("OCDBT_COORDINATOR_HOST"), + "OCDBT_COORDINATOR_PORT": environ.get("OCDBT_COORDINATOR_PORT"), + }, + "exception": { + "type": type(exc).__name__, + "module": type(exc).__module__, + "message": str(exc), + "traceback": traceback.format_exc(), + }, + } + + +def bbox_failure_payload( + exc: BaseException, + dump_tag: Optional[str], + bbox_lo, + bbox_hi, + resolutions, + per_scale, + dst_handle, + src_handle, +) -> dict: + """Build the full diagnostic dict for a ``copy_ws_bbox_multiscale`` + commit failure. + + Merges generic ``failure_envelope`` metadata with bbox-specific + fields (per-scale shape / chunk / key-count / raw-bytes, src+dst + kvstore specs). Spec dumps are wrapped in try/except so a malformed + handle doesn't shadow the original exception. + """ + try: + dst_spec = dst_handle.kvstore.spec().to_json() + except Exception as e: + dst_spec = f"" + try: + src_spec = src_handle.kvstore.spec().to_json() + except Exception as e: + src_spec = f"" + total_voxels = sum(p[2] for p in per_scale) + total_raw = sum(p[3] for p in per_scale) + total_keys = sum(p[5] for p in per_scale) + return { + **failure_envelope(exc, dump_tag), + "bbox_lo": [int(c) for c in bbox_lo], + "bbox_hi": [int(c) for c in bbox_hi], + "resolutions": [list(map(int, r)) for r in resolutions], + "n_scales": len(per_scale), + "total_voxels": total_voxels, + "total_raw_bytes": total_raw, + "total_keys": total_keys, + "per_scale": [ + { + "scale_index": i, + "dims": list(dims), + "voxels": nvox, + "raw_bytes": raw_bytes, + "chunk_shape": list(chunk_shape), + "n_keys": n_keys, + "max_raw_per_key_bytes": max_per_key, + } + for i, dims, nvox, raw_bytes, chunk_shape, n_keys, max_per_key in per_scale + ], + "dst_kvstore_spec": dst_spec, + "src_kvstore_spec": src_spec, + } + + +def dump_failure_to_gcs(payload: dict, dump_tag: str) -> Optional[str]: + """Write a per-task failure report to ``$ERROR_DUMP/__.json``. + + Returns the full path or None (env unset, dump_tag empty, or write + error). ``dump_tag`` carries the calling-context identifier (graph + id, layer, coords, …) so multiple experiments can share one + ``ERROR_DUMP`` bucket without collisions. + """ + root = environ.get("ERROR_DUMP", "").strip() + if not root or not dump_tag: + return None + if not root.endswith("/"): + root += "/" + utc = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + rel = f"{dump_tag}__{utc}.json" + full = root + rel + try: + ts.KvStore.open(root).result().write( + rel, json.dumps(payload, indent=2).encode("utf-8") + ).result() + return full + except Exception as e: + _logger.warning("failed to write ERROR_DUMP at %s: %r", full, e) + return None diff --git a/pychunkedgraph/graph/ocdbt/main.py b/pychunkedgraph/graph/ocdbt/main.py new file mode 100644 index 000000000..f24ad6a74 --- /dev/null +++ b/pychunkedgraph/graph/ocdbt/main.py @@ -0,0 +1,583 @@ +"""OCDBT-backed neuroglancer_precomputed segmentation store — public API. + +Architecture: one immutable base OCDBT per watershed + one delta OCDBT per +ChunkedGraph. Reads merge base + delta via tensorstore's kvstack driver. +Writes land in the delta via OCDBT's ``*_data_prefix`` options. + +Multi-scale (MIP pyramid) is supported: the source watershed's info JSON +drives the scale layout. All scales share one OCDBT kvstore; the precomputed +driver prefixes keys by scale key automatically. + +Versioned reads +--------------- +Every OCDBT commit gets a monotonically-increasing ``generation_number`` and +an ``absl::Now()``-stamped ``commit_time`` (nanoseconds since epoch). The +tensorstore OCDBT driver lets callers pin a read-only open to a prior version +via the ``version`` spec field; accepts either an integer generation number +or an ISO-8601 UTC timestamp string. The timestamp form requires a ``Z`` +suffix (not ``+00:00``) and is interpreted as ``commit_time <= T`` — the open +returns the latest version at or before the pinned time. + +The commit_time itself cannot be overridden by the caller: OCDBT stamps each +commit from the writer's local clock (``absl::Now()`` in +``btree_writer_commit_operation.cc``). This means we can't make OCDBT commits +align exactly with a caller-provided operation timestamp. What the L2 chunk +lock guarantees instead: no other writer can commit to our chunks while we +hold the lock, so any timestamp captured under the lock before our first +commit is a valid pin for "pre-op state of our chunks." + +Retention: the OCDBT spec exposes no pruning fields. All versions are +retained by default. +""" + +from os import environ + +import numpy as np +import tensorstore as ts + +from pychunkedgraph import get_logger + +from .debug import bbox_failure_payload, dump_failure_to_gcs +from .meta import OcdbtConfig +from ..dry_run import is_dry_run +from .utils import ( + _base_ocdbt_path, + _ensure_trailing_slash, + _open_precomputed_scale, + _read_source_scales, + _schema_from_src, + base_exists, + fork_exists, +) + +logger = get_logger(__name__) + + +def create_base_ocdbt(ws_path: str, config: OcdbtConfig): + """One-time bootstrap: create the shared base OCDBT at ``/ocdbt/base/``. + + Wipes any existing base first, then opens each scale with create=True + so the info JSON is built from the source. Populating the base with + actual chunk data happens separately via ``copy_ws_chunk_multiscale`` + or ``copy_ws_bbox_multiscale`` during the per-chunk ingest tasks. + + Returns (src_list, dst_list, resolutions) for the caller to use with + the copy helpers. + """ + base = _base_ocdbt_path(ws_path) + # Wipe via the underlying GCS/file driver, NOT through the ocdbt + # driver. Opening as ocdbt on an empty dir creates a default-config + # `manifest.ocdbt` stub (max_inline_value_bytes=100); on a dir with + # an existing manifest it only clears the B+tree, leaving the + # manifest's config in place. Either way the subsequent open with a + # different config mismatches. + try: + kvs = ts.KvStore.open(base).result() + kvs.delete_range(ts.KvStore.KeyRange()).result() + except Exception: + pass + + scales = _read_source_scales(ws_path) + resolutions = [s["resolution"] for s in scales] + base_kvstore = {"driver": "ocdbt", "base": base, "config": config.ts_config()} + + src_list, dst_list = [], [] + for i in range(len(scales)): + src_i = ts.open( + {"driver": "neuroglancer_precomputed", "kvstore": ws_path, "scale_index": i} + ).result() + dst_i = _open_precomputed_scale( + base_kvstore, i, create=True, **_schema_from_src(src_i) + ) + src_list.append(src_i) + dst_list.append(dst_i) + return src_list, dst_list, resolutions + + +def wipe_base_ocdbt(ws_path: str): + """Wipe the base OCDBT entirely (for --reset-ocdbt).""" + base = _base_ocdbt_path(ws_path) + # Wipe via the underlying GCS/file driver so the manifest file is + # deleted too. Opening as ocdbt only clears the B+tree. + try: + kvs = ts.KvStore.open(base).result() + kvs.delete_range(ts.KvStore.KeyRange()).result() + except Exception: + pass + + +def open_base_ocdbt( + ws_path: str, config: OcdbtConfig, coordinator_address: str | None = None +): + """Open the existing base OCDBT (read/write) for populating during ingest. + + Used by per-chunk ingest tasks that copy precomputed data into the shared + base. NOT used at runtime — runtime always goes through the per-CG fork + spec via ``get_seg_source_and_destination_ocdbt``. + + ``coordinator_address`` (``"host:port"``) routes every OCDBT commit + through a ``DistributedCoordinatorServer`` so parallel workers don't + race the shared manifest's CAS — the only thing that prevents the + orphan ``d/`` file explosion. Required for any concurrent writer; the + arg is optional so single-process callers (e.g. tests, notebooks) can + skip it. + + Returns (src_list, dst_list, resolutions). + """ + base = _base_ocdbt_path(ws_path) + scales = _read_source_scales(ws_path) + resolutions = [s["resolution"] for s in scales] + # No `config` block: the base already exists (created by + # `create_base_ocdbt`), so its on-disk manifest is authoritative. + # Embedding `config.ts_config()` here would assert our in-code + # defaults against whatever is persisted and raise + # FAILED_PRECONDITION on any drift. + base_kvstore = {"driver": "ocdbt", "base": base} + if coordinator_address: + base_kvstore["coordinator"] = {"address": coordinator_address} + + src_list, dst_list = [], [] + for i in range(len(scales)): + src_i = ts.open( + {"driver": "neuroglancer_precomputed", "kvstore": ws_path, "scale_index": i} + ).result() + dst_i = _open_precomputed_scale(base_kvstore, i, **_schema_from_src(src_i)) + src_list.append(src_i) + dst_list.append(dst_i) + return src_list, dst_list, resolutions + + +def build_cg_ocdbt_spec( + ws_path: str, + graph_id: str, + config: OcdbtConfig, + *, + pinned_at: "int | str | None" = None, +) -> dict: + """Open-time kvstore spec for a CG's OCDBT, backed by a shared immutable base. + + This function is a pure spec-constructor — it doesn't materialize + the fork. The fork's ``manifest.ocdbt`` must exist before ``ts.open`` + on this spec will succeed; it's created by ``fork_base_manifest`` + (invoked from the ingest CLI's OCDBT path or the ``seg_ocdbt`` + notebook). ``ChunkedGraphMeta.ws_ocdbt_scales`` asserts presence via + ``fork_exists`` so callers get a clear error instead of a tensorstore + internal failure. + + All three kvstack layers below AND all three ``*_data_prefix`` options + are load-bearing; removing any of them causes fork writes to leak + into the immutable base (verified empirically). + + When ``pinned_at`` is set, the opened kvstore is read-only and returns + state as of the specified version. Accepts an integer generation + number (exact) or an ISO-8601 UTC timestamp string with ``Z`` suffix + (interpreted as ``commit_time <= T``). + """ + base = _base_ocdbt_path(ws_path) + fork_dir = _ensure_trailing_slash(f"{ws_path.rstrip('/')}/ocdbt/{graph_id}") + data_prefix = f"{graph_id}_d/" + + # Catch-all. Lets the fork READ base's B+tree (manifest + d/ + # data files) via fall-through. Must be first so later layers can + # override sub-ranges. + base_layer = {"base": base} + + # Single-key override. Routes the fork's manifest file so new + # commits by this CG are visible only to this CG. Without this layer + # manifest writes silently clobber base's manifest. + fork_manifest_layer = { + "exact": "manifest.ocdbt", + "base": fork_dir + "manifest.ocdbt", + } + + # Catches OCDBT's new data-file writes for the fork. Pairs with the + # *_data_prefix options: OCDBT would otherwise write under the + # default `d/` prefix — no later layer claims `d/`, so kvstack + # falls through to the base catch-all and the writes corrupt base. + fork_data_layer = { + "prefix": data_prefix, + "base": _ensure_trailing_slash(fork_dir + data_prefix), + } + + # No `config` block: this spec opens an existing OCDBT (the shared + # base + this fork's manifest+data layers). Tensorstore validates + # every field of `config` against the on-disk manifest and raises + # FAILED_PRECONDITION on mismatch, so embedding our in-code defaults + # here would break any base that was created with different values + # (e.g. an older default for `max_inline_value_bytes`). On-disk wins. + spec = { + "driver": "ocdbt", + "base": { + "driver": "kvstack", + "layers": [base_layer, fork_manifest_layer, fork_data_layer], + }, + # Steer every kind of OCDBT write under `_d/` so the + # fork_data_layer catches them. + "value_data_prefix": data_prefix, + "btree_node_data_prefix": data_prefix, + "version_tree_node_data_prefix": data_prefix, + } + if pinned_at is not None: + spec["version"] = pinned_at + return spec + + +def fork_base_manifest(ws_path: str, graph_id: str, wipe_existing: bool = False): + """Initialize a CG's delta directory by copying the base manifest. + + If wipe_existing=True, deletes the existing fork directory first (for + --retry when a prior ingest failed and left partial delta state). + """ + assert base_exists(ws_path), "base OCDBT must exist before forking" + base = _base_ocdbt_path(ws_path) + fork_dir = _ensure_trailing_slash(f"{ws_path.rstrip('/')}/ocdbt/{graph_id}") + + if wipe_existing: + try: + kvs = ts.KvStore.open(fork_dir).result() + kvs.delete_range(ts.KvStore.KeyRange()).result() + except Exception: + pass + + base_kvs = ts.KvStore.open(base).result() + fork_kvs = ts.KvStore.open(fork_dir).result() + manifest = base_kvs.read("manifest.ocdbt").result().value + fork_kvs.write("manifest.ocdbt", manifest).result() + + +def ensure_fork_synced(ws_path: str, graph_id: str) -> bool: + """Sync fork manifest to base — but only before the fork's first edit. + + Invariant we enforce: a fresh, edit-free fork must reflect base's + *current* manifest at open time. ``setup_base`` calls + ``fork_base_manifest`` once at graph creation, possibly before + populate has committed most of its writes; any subsequent populate + commit to base would otherwise be invisible through the fork + (symptom: meshing reads return zeros). We close that window by + re-snapshotting on the first runtime open before any edit lands. + + Once the fork has any edit (anything under ``_d/``), the + function has no work to do: base is immutable post-setup, so the + fork manifest cannot fall behind in any way that matters — its + divergence from base is just the fork's own forward progress. + Edit files are stable, so listing the prefix is a sufficient + short-circuit and skips reading both manifests on every runtime + open. + + Returns True iff the fork manifest was refreshed. + """ + if is_dry_run(): + return False + if not fork_exists(ws_path, graph_id): + return False + fork_dir = _ensure_trailing_slash(f"{ws_path.rstrip('/')}/ocdbt/{graph_id}") + fork_kvs = ts.KvStore.open(fork_dir).result() + data_prefix = f"{graph_id}_d/" + edit_files = fork_kvs.list( + ts.KvStore.KeyRange(data_prefix, data_prefix[:-1] + chr(ord("/") + 1)) + ).result() + if len(edit_files) > 0: + # Steady state — fork has progressed forward by design. + return False + base = _base_ocdbt_path(ws_path) + base_kvs = ts.KvStore.open(base).result() + base_manifest = base_kvs.read("manifest.ocdbt").result().value + fork_manifest = fork_kvs.read("manifest.ocdbt").result().value + if base_manifest == fork_manifest: + return False + fork_kvs.write("manifest.ocdbt", base_manifest).result() + logger.note(f"refreshed fork manifest at {fork_dir} from base (no edits)") + return True + + +def get_seg_source_and_destination_ocdbt( + ws_path: str, + graph_id: str, + config: OcdbtConfig, + *, + pinned_at: "int | str | None" = None, +) -> tuple: + """Open source watershed + CG's delta OCDBT destination (all scales). + + Always uses the fork-based kvstack spec. Requires the base to exist and + the fork's manifest to be present (set up at ingest time). + + When ``pinned_at`` is set, the destination OCDBT handles are opened + read-only at that version — used by the recovery path to read + pre-op seg values via ``ChunkedGraphMeta.pinned_seg_reads``. + + Returns: + (src_list, dst_list, resolutions): per-scale TensorStore handles + and [x,y,z] resolutions. + """ + scales = _read_source_scales(ws_path) + resolutions = [s["resolution"] for s in scales] + cg_kvstore = build_cg_ocdbt_spec(ws_path, graph_id, config, pinned_at=pinned_at) + + src_list, dst_list = [], [] + for i in range(len(scales)): + src_i = ts.open( + {"driver": "neuroglancer_precomputed", "kvstore": ws_path, "scale_index": i} + ).result() + dst_i = _open_precomputed_scale(cg_kvstore, i, **_schema_from_src(src_i)) + src_list.append(src_i) + dst_list.append(dst_i) + return src_list, dst_list, resolutions + + +def copy_ws_chunk( + source, + destination, + chunk_size: tuple, + coords: list, + voxel_bounds: np.ndarray, +): + """Copy one chunk from source watershed to OCDBT destination at the same scale. + + Coordinates are interpreted at the source/destination's native scale — + callers must pre-scale them when copying coarser MIP levels. + """ + coords = np.array(coords, dtype=int) + chunk_size = np.array(chunk_size, dtype=int) + vx_start = coords * chunk_size + voxel_bounds[:, 0] + vx_end = vx_start + chunk_size + xE, yE, zE = voxel_bounds[:, 1] + + x0, y0, z0 = vx_start + x1, y1, z1 = vx_end + x1 = min(x1, xE) + y1 = min(y1, yE) + z1 = min(z1, zE) + + data = source[x0:x1, y0:y1, z0:z1].read().result() + destination[x0:x1, y0:y1, z0:z1].write(data).result() + + +def copy_ws_chunk_multiscale( + src_list, + dst_list, + resolutions, + chunk_size: tuple, + coords: list, + voxel_bounds: np.ndarray, +): + """Copy a base-resolution chunk's physical region across all MIP scales. + + The graph's chunk grid is defined at base resolution. For each coarser + scale we copy the SAME physical region — voxel coordinates are divided + by the cumulative downsample factor (derived from resolution ratios). + Source already has correct data at every scale, so this is a pure copy + with no recomputation. + """ + assert len(src_list) == len(dst_list) == len(resolutions) + coords = np.array(coords, dtype=int) + chunk_size_arr = np.array(chunk_size, dtype=int) + base_res = np.array(resolutions[0]) + + # Physical region at base resolution. + vx_start_base = coords * chunk_size_arr + voxel_bounds[:, 0] + vx_end_base = np.minimum(vx_start_base + chunk_size_arr, voxel_bounds[:, 1]) + + for i, (src, dst) in enumerate(zip(src_list, dst_list)): + # Cumulative factor from base to this scale (e.g. [2,2,1] per level). + factor = (np.array(resolutions[i]) / base_res).astype(int) + x0, y0, z0 = vx_start_base // factor + x1, y1, z1 = vx_end_base // factor + if x1 <= x0 or y1 <= y0 or z1 <= z0: + logger.debug(f"skipping empty region at scale {i}") + continue + data = src[x0:x1, y0:y1, z0:z1].read().result() + dst[x0:x1, y0:y1, z0:z1].write(data).result() + + +def copy_ws_bbox_multiscale( + src_list, + dst_list, + resolutions, + bbox_lo: np.ndarray, + bbox_hi: np.ndarray, + dump_tag: str | None = None, +): + """Copy a base-resolution voxel bbox across all MIP scales under one + transaction so the whole multi-scale write lands as a single OCDBT commit. + + The transaction (not ``atomic=True``) is what's load-bearing: it batches + every per-chunk underlying-kvstore write across every scale into one + commit, so the d/ file count for one call is constant in bbox size and + grows only with scale count. ``atomic=True`` would add cross-key + isolation but is rejected by tensorstore's distributed-OCDBT path — + when the kvstore is opened with a ``coordinator``, atomic transactions + cannot span multiple keys (verified empirically). Non-atomic still + batches; the coordinator handles concurrency by serializing the commit + on the wire. + + Passing the source TensorStore directly into ``write(...)`` lets + tensorstore stream the copy without materializing an intermediate + numpy array in Python — peak RSS drops by roughly one scale's + worth versus the read-into-numpy-then-write pattern. + """ + assert len(src_list) == len(dst_list) == len(resolutions) + dump_enabled = bool(environ.get("ERROR_DUMP")) + base_res = np.array(resolutions[0]) + txn = ts.Transaction() + # per_scale rows are only populated when dump_enabled, so the failure + # path has enough context for the structured GCS dump without paying any + # bookkeeping cost on the happy path. + per_scale: list = [] + for i, (src, dst) in enumerate(zip(src_list, dst_list)): + factor = (np.array(resolutions[i]) / base_res).astype(int) + x0, y0, z0 = bbox_lo // factor + x1, y1, z1 = bbox_hi // factor + if x1 <= x0 or y1 <= y0 or z1 <= z0: + continue + if dump_enabled: + dims = (int(x1 - x0), int(y1 - y0), int(z1 - z0)) + nvox = dims[0] * dims[1] * dims[2] + bpv = int(np.dtype(dst.dtype.numpy_dtype).itemsize) + # The precomputed driver's read_chunk shape includes a channel + # axis; the spatial chunk shape is the first three dims. + chunk_shape = tuple(int(s) for s in dst.chunk_layout.read_chunk.shape[:3]) + n_keys = int( + np.prod( + [int(np.ceil(d / c)) if c else 0 for d, c in zip(dims, chunk_shape)] + ) + ) + max_raw_per_key = int(np.prod(chunk_shape)) * bpv + per_scale.append( + (i, dims, nvox, nvox * bpv, chunk_shape, n_keys, max_raw_per_key) + ) + dst.with_transaction(txn)[x0:x1, y0:y1, z0:z1].write( + src[x0:x1, y0:y1, z0:z1] + ).result() + try: + txn.commit_async().result() + except Exception as exc: + if dump_enabled: + payload = bbox_failure_payload( + exc, + dump_tag, + bbox_lo, + bbox_hi, + resolutions, + per_scale, + dst_list[0], + src_list[0], + ) + path = dump_failure_to_gcs(payload, dump_tag) + if path: + logger.note(f"OCDBT commit failure dump → {path}") + raise + + +def _mode_downsample(data: np.ndarray, factors: tuple) -> np.ndarray: + """Mode downsample 4D segmentation array [X,Y,Z,C] by per-axis factors. + + Mode (most-frequent label) is the correct downsampling for segmentation: + it preserves exact label IDs (no interpolation) and biases toward the + dominant label in each block. + + Fast path for 2x2x1: uses a vectorized 4-element pairwise comparison. + Among 4 voxels {a,b,c,d}, if any value appears at least twice it is the + mode. Order of comparisons biases ties toward the top-left corner, which + is the standard convention for segmentation downsampling. + """ + fx, fy, fz = factors + X, Y, Z, C = data.shape + + # Pad with edge values so dimensions are divisible by the factor. + # Using 'edge' (not zeros) avoids introducing a phantom background label. + pad = [(0, (-X % fx) % fx), (0, (-Y % fy) % fy), (0, (-Z % fz) % fz), (0, 0)] + if any(p[1] > 0 for p in pad): + data = np.pad(data, pad, mode="edge") + X, Y, Z, C = data.shape + + if fx == 2 and fy == 2 and fz == 1: + # Fast vectorized path for the common 2x2x1 case. + reshaped = data.reshape(X // 2, 2, Y // 2, 2, Z, C) + a = reshaped[:, 0, :, 0] + b = reshaped[:, 0, :, 1] + c = reshaped[:, 1, :, 0] + d = reshaped[:, 1, :, 1] + return np.where( + (a == b) | (a == c) | (a == d), + a, + np.where((b == c) | (b == d), b, np.where(c == d, c, a)), + ) + + if fx == 2 and fy == 2 and fz == 2: + # 2x2x2 (8-element mode) — strided subsample is fast and label-safe + # for typical segmentation where adjacent voxels share labels. + return data[::2, ::2, ::2] + + # Generic factor: reshape into blocks, take strided first element. + # This is label-safe but loses the mode property; downsampling factor + # ratios in production are 2x2x1 or 2x2x2 so the fast paths cover them. + reshaped = data.reshape(X // fx, fx, Y // fy, fy, Z // fz, fz, C) + return reshaped[:, 0, :, 0, :, 0] + + +def propagate_to_coarser_scales(dst_scales, resolutions, base_slices): + """Cascade-downsample data from base scale through all coarser scales. + + Called after writing to the base scale (e.g. after an SV split). Each + coarser scale reads from the level below it (not from base directly), + so total downsample cost shrinks geometrically — each level processes + 1/N the data of the previous one. + + Args: + dst_scales: TensorStore handles, one per MIP level. + resolutions: [x,y,z] resolution arrays per scale, used to derive + per-axis downsample factors from consecutive resolution ratios. + base_slices: tuple of 3 slices (x, y, z) covering the region written + at base resolution. + """ + prev_slices = base_slices + for i in range(1, len(dst_scales)): + # Per-axis downsample factor from actual resolution ratio. + # Never hardcoded — different datasets may have different ratios. + factor = (np.array(resolutions[i]) / np.array(resolutions[i - 1])).astype(int) + + # Map prev-level slices to this level's coordinates. + # Ceil division on stop ensures we cover any partial block. + target_slices = tuple( + slice(s.start // f, -(-s.stop // f)) for s, f in zip(prev_slices, factor) + ) + + data = dst_scales[i - 1][prev_slices + (slice(None),)].read().result() + downsampled = _mode_downsample(data, tuple(int(f) for f in factor)) + dst_scales[i][target_slices + (slice(None),)].write(downsampled).result() + + prev_slices = target_slices + + +def write_seg_chunks(meta, seg_writes): + """Write a flat batch of pre-sliced L2 chunks to OCDBT in parallel. + + ``seg_writes`` is the aggregated output of ``sv_split.edits.split_supervoxels`` + across every rep in an operation — each pair is one L2 chunk's worth + of ``(voxel_slices, data)``. Flattening across reps matters: one + ``write_seg_chunks`` call fires every chunk write in one parallel + tensorstore batch instead of serializing rep-by-rep. + + Only chunks that actually received new SV IDs appear here; gap + chunks between cross-chunk-connected pieces and neighbor chunks the + overlap read touched are skipped by the split planner. + + Coarser MIP levels stay the downsample worker's job — it picks up + the pubsub message ``publish_edit`` sends after this returns. + + Args: + meta: ChunkedGraphMeta with ``ws_ocdbt`` (base-scale handle). + seg_writes: iterable of ``(voxel_slices, data)`` pairs, where + ``voxel_slices`` is a 3-tuple of ``slice`` objects covering one + L2 chunk's x/y/z extent and ``data`` is the 3D label block + (shape matches the slice extents). + """ + if is_dry_run(): + return + futures = [ + meta.ws_ocdbt[voxel_slices + (slice(None),)].write(data[..., np.newaxis]) + for voxel_slices, data in seg_writes + ] + for f in futures: + f.result() diff --git a/pychunkedgraph/graph/ocdbt/meta.py b/pychunkedgraph/graph/ocdbt/meta.py new file mode 100644 index 000000000..df18b4dca --- /dev/null +++ b/pychunkedgraph/graph/ocdbt/meta.py @@ -0,0 +1,78 @@ +"""OcdbtConfig dataclass — single source of truth for per-CG OCDBT settings.""" + +from dataclasses import asdict, dataclass, field +from typing import Dict, Optional + + +@dataclass +class OcdbtConfig: + """Per-CG OCDBT settings, persisted in ``ChunkedGraphMeta.custom_data["ocdbt_config"]``. + + Carries both ingest-time choices (populate base? at which layer?) and + tensorstore kvstore options (compression, inline byte cap) that must + stay consistent for the lifetime of the OCDBT base. Built once from + the dataset yaml's ``ocdbt_config:`` section and stored alongside the + CG so every code path that opens an OCDBT store reads back the same + values. + """ + + enabled: bool = False + populate_base: bool = False + populate_layer: int = 3 + sv_split_threshold: int = 10 + compression: Dict = field(default_factory=lambda: {"id": "zstd", "level": 12}) + # Inline-vs-out-of-line threshold. Values ≤ this size live in the btree + # leaf bytes; larger values get written to a d/ file and the mutation + # carries only an IndirectDataReference. This directly determines + # cooperator-forwarded RPC size in distributed mode: inline values are + # carried inside the gRPC WriteRequest's `mutations` field, so a leaf's + # batch can blow past tensorstore's hardcoded 4 MiB gRPC max-receive + # whenever multiple inline values pile up on the same node. Verified + # by reading btree_writer.cc StagePending in v0.1.81. + # + # 4 KiB keeps small metadata (info JSON ~1.5 KB, populate-marker files) + # inline while forcing every segmentation chunk value out-of-line — + # chunks compress to 100s of KB even for the smallest scales. With + # chunk bytes out-of-line the WriteRequest stays tiny regardless of + # how many keys a worker commits at once. Tradeoff vs the previous + # 1 MiB cap: each chunk now has its own zstd-framed d/ blob instead of + # sharing a leaf's compression context, which can cost a few percent + # of compression ratio (much less than the originally-feared "7× + # bloat", which only applied at the 100-byte default). + max_inline_value_bytes: int = 4096 + + @classmethod + def from_dict(cls, d: Optional[Dict]) -> "OcdbtConfig": + """Build from a dict. Unknown keys are ignored so older configs don't + break newer code, and newer fields default in when older configs are + loaded. + """ + if not d: + return cls() + known = {f.name for f in cls.__dataclass_fields__.values()} + return cls(**{k: v for k, v in d.items() if k in known}) + + @classmethod + def resolve(cls, *dicts: Optional[Dict]) -> "OcdbtConfig": + """Layered merge: later dicts override earlier ones, all over defaults. + + Use to express precedence — e.g. ``resolve(yaml_dict, info_file_dict)`` + gives info-file values priority over yaml-supplied ones, with + dataclass defaults filling anything neither side specifies. + ``None`` and empty dicts are no-ops. + """ + merged: Dict = {} + for d in dicts: + if d: + merged.update(d) + return cls.from_dict(merged) + + def to_dict(self) -> Dict: + return asdict(self) + + def ts_config(self) -> Dict: + """The subset that belongs inside a tensorstore OCDBT kvstore ``config``.""" + return { + "compression": dict(self.compression), + "max_inline_value_bytes": self.max_inline_value_bytes, + } diff --git a/pychunkedgraph/graph/ocdbt/utils.py b/pychunkedgraph/graph/ocdbt/utils.py new file mode 100644 index 000000000..4c04e5799 --- /dev/null +++ b/pychunkedgraph/graph/ocdbt/utils.py @@ -0,0 +1,156 @@ +"""Internal helpers for the OCDBT package. + +Path builders, schema extraction, populate-marker IO, layer-bbox math. +Not part of the public API except for the marker IO and ``_layer_bbox`` +which the ingest worker uses across the package boundary. +""" + +import json +from typing import Optional + +import numpy as np +import tensorstore as ts +from tenacity import ( + retry, + retry_if_exception_message, + stop_after_attempt, + wait_exponential, +) + +# tensorstore raises ValueError with an absl/grpc status-code prefix. Retry +# only the transient classes — DNS hiccups, deadline-blown reads, server +# 5xx — so a single flaky GCS call doesn't kill the populate task. Persistent +# errors (NOT_FOUND, INVALID_ARGUMENT, RESOURCE_EXHAUSTED, …) propagate. +_transient = retry( + retry=retry_if_exception_message( + match=r"^(UNAVAILABLE|DEADLINE_EXCEEDED|ABORTED|INTERNAL):" + ), + stop=stop_after_attempt(5), + wait=wait_exponential(multiplier=0.5, min=0.5, max=8), + reraise=True, +) + + +def _ensure_trailing_slash(path: str) -> str: + """Ensure kvstore paths end with / so they're treated as directories.""" + return path if path.endswith("/") else path + "/" + + +def _base_ocdbt_path(ws_path: str) -> str: + return _ensure_trailing_slash(f"{ws_path.rstrip('/')}/ocdbt/base") + + +def _populate_markers_path(ws_path: str) -> str: + return _ensure_trailing_slash(f"{ws_path.rstrip('/')}/ocdbt/.populated") + + +def _marker_key(layer: int, coords) -> str: + return f"l{int(layer)}_{int(coords[0])}_{int(coords[1])}_{int(coords[2])}" + + +def _read_source_scales(ws_path: str): + """Read the source precomputed ``info`` JSON to get scale count and resolutions. + + The leading '/' in '/info' is required for GCS — without it the read + returns empty. + """ + kvs = ts.KvStore.open(ws_path).result() + info = json.loads(kvs.read("/info").result().value) + return info["scales"] + + +def _open_precomputed_scale( + kvstore, scale_index: int, create: bool = False, **schema_kw +): + """Open one neuroglancer_precomputed scale on top of a kvstore spec.""" + spec = { + "driver": "neuroglancer_precomputed", + "kvstore": kvstore, + "scale_index": scale_index, + } + return ts.open(spec, create=create, **schema_kw).result() + + +def _schema_from_src(src_handle) -> dict: + """Extract the schema kwargs needed to open a matching destination. + + ``domain`` already carries both extent and origin (voxel_offset). Passing + ``shape`` alongside conflicts with non-zero-origin sources because shape + implies origin=0 — tensorstore refuses to merge ``[0, N)`` with + ``[offset, offset+N)``. + """ + s = src_handle.schema + return dict( + rank=s.rank, + dtype=s.dtype, + codec=s.codec, + domain=s.domain, + chunk_layout=s.chunk_layout, + dimension_units=s.dimension_units, + ) + + +@_transient +def is_chunk_populated(ws_path: str, layer: int, coords) -> bool: + """Check whether this chunk's precomputed→OCDBT copy has already completed. + + Markers live outside the OCDBT keyspace at + ``/ocdbt/.populated/l___`` so retried ingest tasks + don't re-copy chunks and bloat the database with redundant versioned + writes. + """ + kvs = ts.KvStore.open(_populate_markers_path(ws_path)).result() + result = kvs.read(_marker_key(layer, coords)).result() + return result.value is not None and len(result.value) > 0 + + +@_transient +def mark_chunk_populated(ws_path: str, layer: int, coords) -> None: + """Record that this chunk's precomputed→OCDBT copy completed.""" + kvs = ts.KvStore.open(_populate_markers_path(ws_path)).result() + kvs.write(_marker_key(layer, coords), b"1").result() + + +@_transient +def read_populate_meta(ws_path: str) -> Optional[dict]: + """Return the per-base populate config dict, or None if not yet written.""" + kvs = ts.KvStore.open(_populate_markers_path(ws_path)).result() + r = kvs.read("meta.json").result() + if r.value is None or len(r.value) == 0: + return None + return json.loads(r.value) + + +@_transient +def write_populate_meta(ws_path: str, meta: dict) -> None: + """Persist the per-base populate config (layer, etc.) alongside markers.""" + kvs = ts.KvStore.open(_populate_markers_path(ws_path)).result() + kvs.write("meta.json", json.dumps(meta).encode()).result() + + +def base_exists(ws_path: str) -> bool: + """Check if the base OCDBT has already been created for this watershed.""" + base = _base_ocdbt_path(ws_path) + kvs = ts.KvStore.open(base).result() + result = kvs.read("manifest.ocdbt").result() + return result.value is not None and len(result.value) > 0 + + +def fork_exists(ws_path: str, graph_id: str) -> bool: + """Check if this ChunkedGraph's fork has been initialized.""" + fork_dir = _ensure_trailing_slash(f"{ws_path.rstrip('/')}/ocdbt/{graph_id}") + kvs = ts.KvStore.open(fork_dir).result() + result = kvs.read("manifest.ocdbt").result() + return result.value is not None and len(result.value) > 0 + + +def _layer_bbox(meta, layer: int, coords) -> tuple: + """Base-resolution voxel bbox of a chunk at this layer.""" + chunk_size = np.array(meta.graph_config.CHUNK_SIZE, dtype=int) + layer_chunk_size = chunk_size * (1 << (layer - 2)) + coords = np.array(coords, dtype=int) + vol_start = meta.voxel_bounds[:, 0] + vol_end = meta.voxel_bounds[:, 1] + lo = coords * layer_chunk_size + vol_start + hi = np.minimum(lo + layer_chunk_size, vol_end) + return lo, hi diff --git a/pychunkedgraph/graph/operation.py b/pychunkedgraph/graph/operation.py index 68abc17bc..d226b2eef 100644 --- a/pychunkedgraph/graph/operation.py +++ b/pychunkedgraph/graph/operation.py @@ -1,5 +1,6 @@ -# pylint: disable=invalid-name, missing-docstring, too-many-lines, protected-access +# pylint: disable=invalid-name, missing-docstring, too-many-lines, protected-access, broad-exception-raised +import time from abc import ABC, abstractmethod from collections import namedtuple from datetime import datetime @@ -9,32 +10,47 @@ from typing import Type from typing import Tuple from typing import Union +from typing import Any from typing import Optional from typing import Sequence from functools import reduce import numpy as np -from google.cloud import bigtable +from pychunkedgraph import get_logger +from . import err_dump from . import locks from . import edits +from . import sv_split from . import types -from . import attributes +from .ocdbt import write_seg_chunks +from .dry_run import is_dry_run +from pychunkedgraph.graph import attributes from .edges import Edges from .edges.utils import get_edges_status -from .utils import basetypes -from .utils import serializers +from pychunkedgraph.graph import basetypes +from pychunkedgraph.graph import serializers from .cache import CacheService -from .cutting import run_multicut +from .cutting import Cut, SvSplitRequired, run_multicut from .exceptions import PreconditionError from .exceptions import PostconditionError -from .utils.generic import get_bounding_box as get_bbox -from ..logging.log_db import TimeIt - +from .utils.generic import get_bounding_box as get_bbox, assert_same_root +from pychunkedgraph.graph import get_valid_timestamp if TYPE_CHECKING: from .chunkedgraph import ChunkedGraph +logger = get_logger(__name__) + + +def _log_edit_done(result, op_type, elapsed): + new_roots = list(map(int, np.asarray(result.new_root_ids).tolist())) + old_roots = list(map(int, np.asarray(result.old_root_ids).tolist())) + logger.note( + f"<{result.operation_id}> {op_type} done " + f"new_roots={new_roots} old_roots={old_roots} elapsed={elapsed:.2f}s" + ) + class GraphEditOperation(ABC): __slots__ = [ @@ -44,9 +60,12 @@ class GraphEditOperation(ABC): "sink_coords", "parent_ts", "privileged_mode", + "do_sanity_check", ] Result = namedtuple( - "Result", ["operation_id", "new_root_ids", "new_lvl2_ids", "old_root_ids"] + "Result", + ["operation_id", "new_root_ids", "new_lvl2_ids", "old_root_ids", "seg_bbox"], + defaults=(None,), ) def __init__( @@ -362,10 +381,10 @@ def _update_root_ids(self) -> np.ndarray: @abstractmethod def _apply( self, *, operation_id, timestamp - ) -> Tuple[np.ndarray, np.ndarray, List["bigtable.row.Row"]]: + ) -> Tuple[np.ndarray, np.ndarray, List[Any]]: """Initiates the graph operation calculation. :return: New root IDs, new Lvl2 node IDs, and affected records - :rtype: Tuple[np.ndarray, np.ndarray, List["bigtable.row.Row"]] + :rtype: Tuple[np.ndarray, np.ndarray, List[Any]] """ @abstractmethod @@ -378,11 +397,11 @@ def _create_log_record( new_root_ids, status=1, exception="", - ) -> "bigtable.row.Row": + ) -> Any: """Creates a log record with all necessary information to replay the current GraphEditOperation :return: Bigtable row containing the log record - :rtype: bigtable.row.Row + :rtype: row mutation object """ @abstractmethod @@ -417,6 +436,7 @@ def execute( is_merge = isinstance(self, MergeOperation) op_type = "merge" if is_merge else "split" self.parent_ts = parent_ts + t_edit_start = time.time() root_ids = self._update_root_ids() with locks.RootLock( self.cg, @@ -430,6 +450,8 @@ def execute( lock.locked_root_ids, np.array([lock.operation_id] * len(lock.locked_root_ids)), ) + if timestamp is None: + timestamp = get_valid_timestamp(timestamp) log_record_before_edit = self._create_log_record( operation_id=lock.operation_id, @@ -438,15 +460,14 @@ def execute( operation_ts=override_ts if override_ts else timestamp, status=attributes.OperationLogs.StatusCodes.CREATED.value, ) - self.cg.client.write([log_record_before_edit]) + self._persist_rows([log_record_before_edit]) try: - with TimeIt(f"{op_type}.apply", self.cg.graph_id, lock.operation_id): - new_root_ids, new_lvl2_ids, affected_records = self._apply( - operation_id=lock.operation_id, - timestamp=override_ts if override_ts else timestamp, - ) - if self.cg.meta.READ_ONLY: + new_root_ids, new_lvl2_ids, affected_records = self._apply( + operation_id=lock.operation_id, + timestamp=override_ts if override_ts else timestamp, + ) + if is_dry_run(): # return without persisting changes return GraphEditOperation.Result( operation_id=lock.operation_id, @@ -460,9 +481,32 @@ def execute( except PostconditionError as err: self.cg.cache = None raise PostconditionError(err) from err + except (AssertionError, RuntimeError) as err: + self.cg.cache = None + dump_url = err_dump.dump_err_artifact( + self.cg, + lock.operation_id, + err_dump.build_err_payload(self, lock.operation_id, err), + ) + logger.error( + f"<{lock.operation_id}> {type(self).__name__} failed: " + f"{type(err).__name__}: {err}" + f"{err_dump.payload_summary(self)} dump={dump_url}" + ) + raise RuntimeError(err) from err except Exception as err: # unknown exception, update log record with error self.cg.cache = None + dump_url = err_dump.dump_err_artifact( + self.cg, + lock.operation_id, + err_dump.build_err_payload(self, lock.operation_id, err), + ) + logger.error( + f"<{lock.operation_id}> {type(self).__name__} failed: " + f"{type(err).__name__}: {err}" + f"{err_dump.payload_summary(self)} dump={dump_url}" + ) log_record_error = self._create_log_record( operation_id=lock.operation_id, new_root_ids=types.empty_1d, @@ -471,19 +515,19 @@ def execute( status=attributes.OperationLogs.StatusCodes.EXCEPTION.value, exception=repr(err), ) - self.cg.client.write([log_record_error]) - raise Exception(err) - - with TimeIt(f"{op_type}.write", self.cg.graph_id, lock.operation_id): - result = self._write( - lock, - override_ts if override_ts else timestamp, - new_root_ids, - new_lvl2_ids, - affected_records, - root_ids, - ) - return result + self._persist_rows([log_record_error]) + raise Exception(err) from err + + result = self._write( + lock, + override_ts if override_ts else timestamp, + new_root_ids, + new_lvl2_ids, + affected_records, + root_ids, + ) + _log_edit_done(result, op_type, time.time() - t_edit_start) + return result def _write( self, @@ -512,6 +556,7 @@ def _write( lock.operation_id, lock.locked_root_ids, privileged_mode=lock.privileged_mode, + future_root_ids_d=lock.future_root_ids_d, ): # indefinite lock for writing, if a node instance or pod dies during this # the roots must stay locked indefinitely to prevent further corruption. @@ -535,8 +580,18 @@ def _write( new_root_ids=new_root_ids, new_lvl2_ids=new_lvl2_ids, old_root_ids=old_root_ids, + # Only set when the operation actually ran SV splits (MulticutOperation + # populates this; other operations leave the attr absent and it defaults + # to None via the Result namedtuple's default). + seg_bbox=getattr(self, "seg_bboxes", None) or None, ) + def _persist_rows(self, rows): + """Persist BT mutation rows; no-op under ``PCG_DRY_RUN=1``.""" + if is_dry_run(): + return + self.cg.client.write(rows) + class MergeOperation(GraphEditOperation): """Merge Operation: Connect *known* pairs of supervoxels by adding a (weighted) edge. @@ -565,6 +620,8 @@ class MergeOperation(GraphEditOperation): "affinities", "bbox_offset", "allow_same_segment_merge", + "do_sanity_check", + "stitch_mode", ] def __init__( @@ -578,6 +635,8 @@ def __init__( bbox_offset: Tuple[int, int, int] = (240, 240, 24), affinities: Optional[Sequence[np.float32]] = None, allow_same_segment_merge: Optional[bool] = False, + do_sanity_check: Optional[bool] = True, + stitch_mode: bool = False, ) -> None: super().__init__( cg, user_id=user_id, source_coords=source_coords, sink_coords=sink_coords @@ -585,6 +644,8 @@ def __init__( self.added_edges = np.atleast_2d(added_edges).astype(basetypes.NODE_ID) self.bbox_offset = np.atleast_1d(bbox_offset).astype(basetypes.COORDINATES) self.allow_same_segment_merge = allow_same_segment_merge + self.do_sanity_check = do_sanity_check + self.stitch_mode = stitch_mode self.affinities = None if affinities is not None: @@ -608,16 +669,20 @@ def _update_root_ids(self) -> np.ndarray: def _apply( self, *, operation_id, timestamp - ) -> Tuple[np.ndarray, np.ndarray, List["bigtable.row.Row"]]: - root_ids = set( - self.cg.get_roots( - self.added_edges.ravel(), assert_roots=True, time_stamp=self.parent_ts - ) - ) + ) -> Tuple[np.ndarray, np.ndarray, List[Any]]: + sv_ids = self.added_edges.ravel() + roots = self.cg.get_roots(sv_ids, assert_roots=True, time_stamp=self.parent_ts) + root_ids = set(roots) if len(root_ids) < 2 and not self.allow_same_segment_merge: - raise PreconditionError("Supervoxels must belong to different objects.") - bbox = get_bbox(self.source_coords, self.sink_coords, self.bbox_offset) - with TimeIt("subgraph", self.cg.graph_id, operation_id): + raise PreconditionError( + f"[MergeOperation._apply] Supervoxels must belong to different " + f"objects. sv_id->root: {dict(zip(sv_ids.tolist(), roots.tolist()))}" + ) + + atomic_edges = self.added_edges + fake_edge_rows = [] + if not self.stitch_mode: + bbox = get_bbox(self.source_coords, self.sink_coords, self.bbox_offset) edges = self.cg.get_subgraph( root_ids, bbox=bbox, @@ -625,29 +690,34 @@ def _apply( edges_only=True, ) - with TimeIt("preprocess", self.cg.graph_id, operation_id): - inactive_edges = edits.merge_preprocess( + if self.allow_same_segment_merge: + inactive_edges = types.empty_2d + else: + inactive_edges = edits.merge_preprocess( + self.cg, + subgraph_edges=edges, + supervoxels=self.added_edges.ravel(), + parent_ts=self.parent_ts, + ) + + atomic_edges, fake_edge_rows = edits.check_fake_edges( self.cg, - subgraph_edges=edges, - supervoxels=self.added_edges.ravel(), + atomic_edges=self.added_edges, + inactive_edges=inactive_edges, + time_stamp=timestamp, parent_ts=self.parent_ts, ) - atomic_edges, fake_edge_rows = edits.check_fake_edges( + new_roots, new_l2_ids, new_entries = edits.add_edges( self.cg, - atomic_edges=self.added_edges, - inactive_edges=inactive_edges, + atomic_edges=atomic_edges, + operation_id=operation_id, time_stamp=timestamp, parent_ts=self.parent_ts, + allow_same_segment_merge=self.allow_same_segment_merge, + do_sanity_check=self.do_sanity_check, + stitch_mode=self.stitch_mode, ) - with TimeIt("add_edges", self.cg.graph_id, operation_id): - new_roots, new_l2_ids, new_entries = edits.add_edges( - self.cg, - atomic_edges=atomic_edges, - operation_id=operation_id, - time_stamp=timestamp, - parent_ts=self.parent_ts, - ) return new_roots, new_l2_ids, fake_edge_rows + new_entries def _create_log_record( @@ -659,7 +729,7 @@ def _create_log_record( new_root_ids: Sequence[np.uint64], status: int = 1, exception: str = "", - ) -> "bigtable.row.Row": + ) -> Any: val_dict = { attributes.OperationLogs.UserID: self.user_id, attributes.OperationLogs.RootID: new_root_ids, @@ -705,7 +775,7 @@ class SplitOperation(GraphEditOperation): :type sink_coords: Optional[Sequence[Sequence[int]]], optional """ - __slots__ = ["removed_edges", "bbox_offset"] + __slots__ = ["removed_edges", "bbox_offset", "do_sanity_check"] def __init__( self, @@ -716,12 +786,14 @@ def __init__( source_coords: Optional[Sequence[Sequence[int]]] = None, sink_coords: Optional[Sequence[Sequence[int]]] = None, bbox_offset: Tuple[int] = (240, 240, 24), + do_sanity_check: Optional[bool] = True, ) -> None: super().__init__( cg, user_id=user_id, source_coords=source_coords, sink_coords=sink_coords ) self.removed_edges = np.atleast_2d(removed_edges).astype(basetypes.NODE_ID) self.bbox_offset = np.atleast_1d(bbox_offset).astype(basetypes.COORDINATES) + self.do_sanity_check = do_sanity_check if np.any(np.equal(self.removed_edges[:, 0], self.removed_edges[:, 1])): raise PreconditionError("Requested split contains at least 1 self-loop.") @@ -729,49 +801,25 @@ def __init__( assert np.sum(layers) == layers.size, "IDs must be supervoxels." def _update_root_ids(self) -> np.ndarray: - root_ids = np.unique( - self.cg.get_roots( - self.removed_edges.ravel(), - assert_roots=True, - time_stamp=self.parent_ts, - ) - ) - if len(root_ids) > 1: - raise PreconditionError("Supervoxels must belong to the same object.") - return root_ids + sv_ids = self.removed_edges.ravel() + roots = self.cg.get_roots(sv_ids, assert_roots=True, time_stamp=self.parent_ts) + return assert_same_root(sv_ids, roots, source="SplitOperation._update_root_ids") def _apply( self, *, operation_id, timestamp - ) -> Tuple[np.ndarray, np.ndarray, List["bigtable.row.Row"]]: - if ( - len( - set( - self.cg.get_roots( - self.removed_edges.ravel(), - assert_roots=True, - time_stamp=self.parent_ts, - ) - ) - ) - > 1 - ): - raise PreconditionError("Supervoxels must belong to the same object.") + ) -> Tuple[np.ndarray, np.ndarray, List[Any]]: + sv_ids = self.removed_edges.ravel() + roots = self.cg.get_roots(sv_ids, assert_roots=True, time_stamp=self.parent_ts) + assert_same_root(sv_ids, roots, source="SplitOperation._apply") - with TimeIt("subgraph", self.cg.graph_id, operation_id): - l2id_agglomeration_d, _ = self.cg.get_l2_agglomerations( - self.cg.get_parents( - self.removed_edges.ravel(), time_stamp=self.parent_ts - ), - ) - with TimeIt("remove_edges", self.cg.graph_id, operation_id): - return edits.remove_edges( - self.cg, - operation_id=operation_id, - atomic_edges=self.removed_edges, - l2id_agglomeration_d=l2id_agglomeration_d, - time_stamp=timestamp, - parent_ts=self.parent_ts, - ) + return edits.remove_edges( + self.cg, + operation_id=operation_id, + atomic_edges=self.removed_edges, + time_stamp=timestamp, + parent_ts=self.parent_ts, + do_sanity_check=self.do_sanity_check, + ) def _create_log_record( self, @@ -782,7 +830,7 @@ def _create_log_record( new_root_ids: Sequence[np.uint64], status: int = 1, exception: str = "", - ) -> "bigtable.row.Row": + ) -> Any: val_dict = { attributes.OperationLogs.UserID: self.user_id, attributes.OperationLogs.RootID: new_root_ids, @@ -839,6 +887,12 @@ class MulticutOperation(GraphEditOperation): "bbox_offset", "path_augment", "disallow_isolating_cut", + "do_sanity_check", + # Base-resolution bboxes of SV splits done as part of this op, one + # per rep. Populated only when the multicut hit SvSplitRequired and + # split_supervoxels actually ran. Surfaced on the Result so the + # downsample worker knows which regions to re-mip. + "seg_bboxes", ] def __init__( @@ -854,6 +908,7 @@ def __init__( removed_edges: Sequence[Sequence[np.uint64]] = types.empty_2d, path_augment: bool = True, disallow_isolating_cut: bool = True, + do_sanity_check: Optional[bool] = True, ) -> None: super().__init__( cg, user_id=user_id, source_coords=source_coords, sink_coords=sink_coords @@ -864,81 +919,148 @@ def __init__( self.bbox_offset = np.atleast_1d(bbox_offset).astype(basetypes.COORDINATES) self.path_augment = path_augment self.disallow_isolating_cut = disallow_isolating_cut - if np.any(np.in1d(self.sink_ids, self.source_ids)): - raise PreconditionError( - "Supervoxels exist in both sink and source, " - "try placing the points further apart." - ) + self.do_sanity_check = do_sanity_check + self.seg_bboxes = [] - ids = np.concatenate([self.source_ids, self.sink_ids]) + ids = np.concatenate([self.source_ids, self.sink_ids]).astype(basetypes.NODE_ID) layers = self.cg.get_chunk_layers(ids) assert np.sum(layers) == layers.size, "IDs must be supervoxels." def _update_root_ids(self) -> np.ndarray: - sink_and_source_ids = np.concatenate((self.source_ids, self.sink_ids)) - root_ids = np.unique( - self.cg.get_roots( - sink_and_source_ids, assert_roots=True, time_stamp=self.parent_ts - ) + sink_and_source_ids = np.concatenate((self.source_ids, self.sink_ids)).astype( + basetypes.NODE_ID + ) + roots = self.cg.get_roots( + sink_and_source_ids, assert_roots=True, time_stamp=self.parent_ts + ) + return assert_same_root( + sink_and_source_ids, + roots, + source="MulticutOperation._update_root_ids", ) - if len(root_ids) > 1: - raise PreconditionError("Supervoxels must belong to the same segment.") - return root_ids def _apply( self, *, operation_id, timestamp - ) -> Tuple[np.ndarray, np.ndarray, List["bigtable.row.Row"]]: - # Verify that sink and source are from the same root object - root_ids = set( - self.cg.get_roots( - np.concatenate([self.source_ids, self.sink_ids]), - assert_roots=True, - time_stamp=self.parent_ts, + ) -> Tuple[np.ndarray, np.ndarray, List[Any]]: + result = self._run_multicut(operation_id) + if isinstance(result, SvSplitRequired): + # Running under GraphEditOperation.execute's RootLock — no same-root + # edit can interleave between the SV split and the retry multicut. + # `plan_sv_splits` returns the chunk scope for both locks below, + # `split_supervoxels` is a pure planner that computes the full + # payload. Writes happen here inside nested L2 chunk locks: + # - `L2ChunkLock` (temporal) spans the seg reads (inside + # `split_supervoxels`) and the writes, so no concurrent + # op can mutate our chunks mid-compute. + # - `IndefiniteL2ChunkLock` is scoped tightly to the writes + # only. A worker death inside it leaves the indefinite + # cell set on every chunk row in scope, blocking future + # ops until operator replay clears them. + tasks, chunk_ids = sv_split.edits.plan_sv_splits( + self.cg, + sv_remapping=result.sv_remapping, + source_ids=self.source_ids, + sink_ids=self.sink_ids, + source_coords=self.source_coords, + sink_coords=self.sink_coords, ) + with locks.L2ChunkLock( + self.cg, + chunk_ids, + operation_id, + privileged_mode=self.privileged_mode, + ): + sv_result = sv_split.edits.split_supervoxels( + self.cg, + tasks=tasks, + sv_remapping=result.sv_remapping, + source_ids=self.source_ids, + sink_ids=self.sink_ids, + operation_id=operation_id, + timestamp=timestamp, + parent_ts=self.parent_ts, + ) + with locks.IndefiniteL2ChunkLock( + self.cg, + chunk_ids, + operation_id, + privileged_mode=self.privileged_mode, + ): + write_seg_chunks(self.cg.meta, sv_result.seg_writes) + self._persist_rows(sv_result.bigtable_rows) + self.seg_bboxes = sv_result.seg_bboxes + self.source_ids = sv_result.source_ids_fresh + self.sink_ids = sv_result.sink_ids_fresh + result = self._run_multicut(operation_id) + if isinstance(result, SvSplitRequired): + raise PreconditionError( + "Supervoxel split succeeded but source and sink remain " + "connected; place source and sink farther apart." + ) + + assert isinstance(result, Cut), f"unexpected multicut result: {result!r}" + self.removed_edges = result.atomic_edges + if not self.removed_edges.size: + raise PostconditionError("Mincut could not find any edges to remove.") + + return edits.remove_edges( + self.cg, + operation_id=operation_id, + atomic_edges=self.removed_edges, + time_stamp=timestamp, + parent_ts=self.parent_ts, + do_sanity_check=self.do_sanity_check, + ) + + def _run_multicut(self, operation_id): + """Build the local subgraph and run multicut; returns the tagged result. + + Factored so `_apply` can call it twice — once for initial detection + and again after an SV split to get fresh atomic_edges against the + post-split graph topology. + """ + sink_and_source_ids = np.concatenate([self.source_ids, self.sink_ids]).astype( + basetypes.NODE_ID + ) + roots = self.cg.get_roots( + sink_and_source_ids, + assert_roots=True, + time_stamp=self.parent_ts, + ) + root_ids = set( + assert_same_root( + sink_and_source_ids, + roots, + source="MulticutOperation._run_multicut", + ).tolist() ) - if len(root_ids) > 1: - raise PreconditionError("Supervoxels must belong to the same object.") bbox = get_bbox( self.source_coords, self.sink_coords, self.cg.meta.split_bounding_offset, ) - with TimeIt("get_subgraph", self.cg.graph_id, operation_id): - l2id_agglomeration_d, edges = self.cg.get_subgraph( - root_ids.pop(), bbox=bbox, bbox_is_coordinate=True - ) - - edges = reduce(lambda x, y: x + y, edges, Edges([], [])) - supervoxels = np.concatenate( - [agg.supervoxels for agg in l2id_agglomeration_d.values()] - ) - mask0 = np.in1d(edges.node_ids1, supervoxels) - mask1 = np.in1d(edges.node_ids2, supervoxels) - edges = edges[mask0 & mask1] + l2id_agglomeration_d, edges_tuple = self.cg.get_subgraph( + root_ids.pop(), bbox=bbox, bbox_is_coordinate=True + ) + edges = reduce(lambda x, y: x + y, edges_tuple, Edges([], [])) + supervoxels = np.concatenate( + [agg.supervoxels for agg in l2id_agglomeration_d.values()] + ).astype(basetypes.NODE_ID) + mask0 = np.isin(edges.node_ids1, supervoxels) + mask1 = np.isin(edges.node_ids2, supervoxels) + edges = edges[mask0 & mask1] if len(edges) == 0: raise PreconditionError("No local edges found.") - with TimeIt("multicut", self.cg.graph_id, operation_id): - self.removed_edges = run_multicut( - edges, - self.source_ids, - self.sink_ids, - path_augment=self.path_augment, - disallow_isolating_cut=self.disallow_isolating_cut, - ) - if not self.removed_edges.size: - raise PostconditionError("Mincut could not find any edges to remove.") - - with TimeIt("remove_edges", self.cg.graph_id, operation_id): - return edits.remove_edges( - self.cg, - operation_id=operation_id, - atomic_edges=self.removed_edges, - l2id_agglomeration_d=l2id_agglomeration_d, - time_stamp=timestamp, - parent_ts=self.parent_ts, - ) + return run_multicut( + edges, + self.source_ids, + self.sink_ids, + path_augment=self.path_augment, + disallow_isolating_cut=self.disallow_isolating_cut, + sv_split_supported=self.cg.meta.ocdbt_seg, + ) def _create_log_record( self, @@ -949,7 +1071,7 @@ def _create_log_record( new_root_ids: Sequence[np.uint64], status: int = 1, exception: str = "", - ) -> "bigtable.row.Row": + ) -> Any: val_dict = { attributes.OperationLogs.UserID: self.user_id, attributes.OperationLogs.RootID: new_root_ids, @@ -1045,7 +1167,7 @@ def _update_root_ids(self): def _apply( self, *, operation_id, timestamp - ) -> Tuple[np.ndarray, np.ndarray, List["bigtable.row.Row"]]: + ) -> Tuple[np.ndarray, np.ndarray, List[Any]]: return self.superseded_operation._apply( operation_id=operation_id, timestamp=timestamp ) @@ -1059,7 +1181,7 @@ def _create_log_record( new_root_ids: Sequence[np.uint64], status: int = 1, exception: str = "", - ) -> "bigtable.row.Row": + ) -> Any: val_dict = { attributes.OperationLogs.UserID: self.user_id, attributes.OperationLogs.RedoOperationID: self.superseded_operation_id, @@ -1178,7 +1300,7 @@ def _update_root_ids(self): def _apply( self, *, operation_id, timestamp - ) -> Tuple[np.ndarray, np.ndarray, List["bigtable.row.Row"]]: + ) -> Tuple[np.ndarray, np.ndarray, List[Any]]: if isinstance(self.inverse_superseded_operation, MergeOperation): return edits.add_edges( self.inverse_superseded_operation.cg, @@ -1201,7 +1323,7 @@ def _create_log_record( new_root_ids: Sequence[np.uint64], status: int = 1, exception: str = "", - ) -> "bigtable.row.Row": + ) -> Any: val_dict = { attributes.OperationLogs.UserID: self.user_id, attributes.OperationLogs.UndoOperationID: self.superseded_operation_id, diff --git a/pychunkedgraph/graph/segmenthistory.py b/pychunkedgraph/graph/segmenthistory.py index 30f42d15b..83dc8175a 100644 --- a/pychunkedgraph/graph/segmenthistory.py +++ b/pychunkedgraph/graph/segmenthistory.py @@ -1,13 +1,13 @@ import collections -from datetime import datetime +from datetime import datetime, timezone from typing import Iterable import numpy as np import fastremap from networkx.algorithms.dag import ancestors as nx_ancestors -from .attributes import OperationLogs -from .utils import basetypes +from pychunkedgraph.graph import attributes +from pychunkedgraph.graph import basetypes class SegmentHistory: @@ -31,7 +31,7 @@ def __init__( if timestamp_past is not None: self.timestamp_past = timestamp_past - self.timestamp_future = datetime.utcnow() + self.timestamp_future = datetime.now(timezone.utc) if timestamp_future is None: self.timestamp_future = timestamp_future @@ -78,7 +78,9 @@ def operation_id_root_id_dict(self): @property def operation_ids(self): - return np.array(list(self.operation_id_root_id_dict.keys())) + return np.array( + list(self.operation_id_root_id_dict.keys()), dtype=basetypes.OPERATION_ID + ) @property def _log_rows(self): @@ -328,7 +330,9 @@ def past_future_id_mapping(self, root_id=None): past_id_mapping = {} future_id_mapping = {} for root_id in root_ids: - ancestors = np.array(list(nx_ancestors(self.lineage_graph, root_id))) + ancestors = np.array( + list(nx_ancestors(self.lineage_graph, root_id)), dtype=np.uint64 + ) if len(ancestors) == 0: past_id_mapping[int(root_id)] = [root_id] else: @@ -375,11 +379,11 @@ def __init__(self, row, timestamp): @property def is_merge(self): - return OperationLogs.AddedEdge in self.row + return attributes.OperationLogs.AddedEdge in self.row @property def user_id(self): - return self.row[OperationLogs.UserID] + return self.row[attributes.OperationLogs.UserID] @property def log_type(self): @@ -387,7 +391,7 @@ def log_type(self): @property def root_ids(self): - return self.row[OperationLogs.RootID] + return self.row[attributes.OperationLogs.RootID] @property def edges_failsafe(self): @@ -403,27 +407,27 @@ def edges_failsafe(self): def sink_source_ids(self): return np.concatenate( [ - self.row[OperationLogs.SinkID], - self.row[OperationLogs.SourceID], + self.row[attributes.OperationLogs.SinkID], + self.row[attributes.OperationLogs.SourceID], ] ) @property def added_edges(self): assert self.is_merge, "Not a merge operation." - return self.row[OperationLogs.AddedEdge] + return self.row[attributes.OperationLogs.AddedEdge] @property def removed_edges(self): assert not self.is_merge, "Not a split operation." - return self.row[OperationLogs.RemovedEdge] + return self.row[attributes.OperationLogs.RemovedEdge] @property def coordinates(self): return np.array( [ - self.row[OperationLogs.SourceCoordinate], - self.row[OperationLogs.SinkCoordinate], + self.row[attributes.OperationLogs.SourceCoordinate], + self.row[attributes.OperationLogs.SinkCoordinate], ] ) diff --git a/pychunkedgraph/graph/subgraph.py b/pychunkedgraph/graph/subgraph.py index ab2593175..4f21f2489 100644 --- a/pychunkedgraph/graph/subgraph.py +++ b/pychunkedgraph/graph/subgraph.py @@ -1,3 +1,5 @@ +# pylint: disable=invalid-name, missing-docstring, import-outside-toplevel + from typing import List from typing import Dict from typing import Tuple @@ -30,9 +32,7 @@ def __init__(self, meta, node_ids, return_layers, serializable): # "Frontier" of nodes that cg.get_children will be called on self.cur_nodes = np.array(list(node_ids), dtype=np.uint64) # Mapping of current frontier to self.node_ids - self.cur_nodes_to_original_nodes = dict( - zip(self.cur_nodes, self.cur_nodes) - ) + self.cur_nodes_to_original_nodes = dict(zip(self.cur_nodes, self.cur_nodes)) self.stop_layer = max(1, min(return_layers)) self.create_initial_node_to_subgraph() @@ -107,13 +107,11 @@ def flatten_subgraph(self): for node_id in self.node_ids: for return_layer in self.return_layers: node_key = self.get_dict_key(node_id) - children_at_layer = self.node_to_subgraph[node_key][ - return_layer - ] + children_at_layer = self.node_to_subgraph[node_key][return_layer] if len(children_at_layer) > 0: - self.node_to_subgraph[node_key][ - return_layer - ] = np.concatenate(children_at_layer) + self.node_to_subgraph[node_key][return_layer] = np.concatenate( + children_at_layer + ) else: self.node_to_subgraph[node_key][return_layer] = empty_1d @@ -123,10 +121,12 @@ def get_subgraph_nodes( node_id_or_ids: Union[np.uint64, Iterable], bbox: Optional[Sequence[Sequence[int]]] = None, bbox_is_coordinate: bool = False, - return_layers: List = [2], + return_layers: List = None, serializable: bool = False, - return_flattened: bool = False + return_flattened: bool = False, ) -> Tuple[Dict, Dict, Edges]: + if return_layers is None: + return_layers = [2] single = False node_ids = node_id_or_ids bbox = normalize_bounding_box(cg.meta, bbox, bbox_is_coordinate) @@ -139,7 +139,7 @@ def get_subgraph_nodes( bounding_box=bbox, return_layers=return_layers, serializable=serializable, - return_flattened=return_flattened + return_flattened=return_flattened, ) if single: if serializable: @@ -155,7 +155,7 @@ def get_subgraph_edges_and_leaves( bbox_is_coordinate: bool = False, edges_only: bool = False, leaves_only: bool = False, -) -> Tuple[Dict, Dict, Edges]: +) -> Tuple[Dict, Tuple[Edges]]: """Get the edges and/or leaves of the specified node_ids within the specified bounding box.""" from .types import empty_1d @@ -183,11 +183,11 @@ def _get_subgraph_multiple_nodes( bounding_box: Optional[Sequence[Sequence[int]]], return_layers: Sequence[int], serializable: bool = False, - return_flattened: bool = False + return_flattened: bool = False, ): from collections import ChainMap - from multiwrapper.multiprocessing_utils import n_cpus - from multiwrapper.multiprocessing_utils import multithread_func + import os + from concurrent.futures import ThreadPoolExecutor from .utils.generic import mask_nodes_by_bounding_box @@ -224,23 +224,25 @@ def _get_subgraph_multiple_nodes_threaded( subgraph = SubgraphProgress(cg.meta, node_ids, return_layers, serializable) while not subgraph.done_processing(): this_n_threads = min( - [int(len(subgraph.cur_nodes) // 50000) + 1, n_cpus] - ) - cur_nodes_child_maps = multithread_func( - _get_subgraph_multiple_nodes_threaded, - np.array_split(subgraph.cur_nodes, this_n_threads), - n_threads=this_n_threads, - debug=this_n_threads == 1, + [int(len(subgraph.cur_nodes) // 50000) + 1, os.cpu_count()] ) + batches = np.array_split(subgraph.cur_nodes, this_n_threads) + if this_n_threads == 1: + cur_nodes_child_maps = [ + _get_subgraph_multiple_nodes_threaded(b) for b in batches + ] + else: + with ThreadPoolExecutor(max_workers=this_n_threads) as executor: + cur_nodes_child_maps = list( + executor.map(_get_subgraph_multiple_nodes_threaded, batches) + ) cur_nodes_children = dict(ChainMap(*cur_nodes_child_maps)) subgraph.process_batch_of_children(cur_nodes_children) if return_flattened and len(return_layers) == 1: for node_id in node_ids: - subgraph.node_to_subgraph[ - _get_dict_key(node_id) - ] = subgraph.node_to_subgraph[_get_dict_key(node_id)][ - return_layers[0] - ] + subgraph.node_to_subgraph[_get_dict_key(node_id)] = ( + subgraph.node_to_subgraph[_get_dict_key(node_id)][return_layers[0]] + ) - return subgraph.node_to_subgraph \ No newline at end of file + return subgraph.node_to_subgraph diff --git a/pychunkedgraph/graph/sv_lookup/README.md b/pychunkedgraph/graph/sv_lookup/README.md new file mode 100644 index 000000000..5c15880e9 --- /dev/null +++ b/pychunkedgraph/graph/sv_lookup/README.md @@ -0,0 +1,86 @@ +# sv_lookup + +Resolve voxel coordinates from interactive UI clicks (split, merge) to +the current supervoxel IDs that those coordinates physically belong to. +Every coord is answered from a **current segmentation read** — +segmentation-agnostic with respect to backend (OCDBT or precomputed +CloudVolume), dispatched by `graph.utils.generic.get_local_segmentation`. + +The caller-supplied `node_ids` are interpreted **only as a layer hint** +to distinguish a 2D slice click from a 3D mesh click. No SV id, root +id, or other value from the client is trusted as the answer. + +## Contract + +| Input layer | Click origin | Behaviour | +|---|---|---| +| `layer == 1` | 2D slice canvas (NG attaches the L1 SV from the slice view) | Return the literal seg SV at that voxel. No constraint, no search. Stale L1 ids from the UI are ignored because the seg is the source of truth. | +| `layer >= 2` | 3D mesh pick (NG attaches the root the user clicked) | Read the literal seg SV. If its current root equals the supplied root → return the literal. Otherwise run a parent-constrained nearest-SV search using the supplied root, at growing radii. | + +The cost ceiling, derived from the contract: + +| Case | seg reads | `get_roots` | `get_atomic_ids_from_coords` | +|---|---|---|---| +| All coords are 2D | 1 | 0 | 0 | +| All coords are 3D-interior (literal root already matches) | 1 | 1 | 0 | +| 3D coord(s) need a search, all share one root | 1 | 1 | ≤ `len(max_dist_steps)` | +| 3D coords need a search across `k` distinct roots | 1 | 1 | one growing-radius sequence per root | + +The growing-radius schedule defaults to +`np.array([4, 8, 14, 28]) * mean(meta.resolution)` nm — the loop breaks +as soon as one radius returns a result for a given root. + +## Layout + +``` +pychunkedgraph/graph/sv_lookup/ +├── __init__.py # re-exports the public API +├── main.py # resolve_supervoxels_at_coords (the orchestrator) +└── utils.py # lookup_svs_from_seg, get_atomic_id_from_coord, + # get_atomic_ids_from_coords (the low-level lookups) +``` + +- `main.resolve_supervoxels_at_coords(cg, coordinates, node_ids, + max_dist_steps=None)` — public entry point. Returns `(N,) uint64`. + Raises `cg_exceptions.BadRequest` on invalid input or unresolvable + coords. +- `utils.lookup_svs_from_seg(meta, coordinates)` — one batched seg read + over the coords' bbox; returns the literal SV per coord. +- `utils.get_atomic_ids_from_coords(meta, coordinates, parent_id, + parent_id_layer, parent_ts, get_roots, max_dist_nm)` — the + parent-constrained nearest-SV search. Reads one bbox-sized seg block + around the input coords, maps every voxel to its root via `get_roots`, + then picks the nm-closest voxel whose root matches `parent_id` for + each input coord. Returns `None` if no voxel within `max_dist_nm` matches. +- `utils.get_atomic_id_from_coord(...)` — single-coord variant retained + for the `cg.get_atomic_id_from_coord` method wrapper. + +## Wiring + +``` +app.app_utils.handle_supervoxel_id_lookup # thin Flask-layer wrapper + └── sv_lookup.resolve_supervoxels_at_coords # the contract above + ├── sv_lookup.utils.lookup_svs_from_seg + ├── cg.get_chunk_layers + ├── cg.get_roots + └── cg.get_atomic_ids_from_coords # for 3D-needs-search only + └── sv_lookup.utils.get_atomic_ids_from_coords +``` + +`cg.get_atomic_ids_from_coords` (defined on `ChunkedGraph`) is the +method wrapper around `utils.get_atomic_ids_from_coords`; it provides +`parent_ts` from the parent's node timestamps and short-circuits a +layer-1 parent to `[parent_id] * N`. The orchestrator only ever invokes +it with a root (layer ≥ 2), so the layer-1 short-circuit never fires +from this path. + +## Tests + +- `tests/graph/test_sv_lookup_main.py` — orchestrator behaviour using a + real bigtable-backed `gen_graph`, with `meta._ws_cv` swapped for a + small sliceable in-memory seg. Covers 2D-only, 3D-interior, + 3D-on-background, search-exhausted, growing-radius-breaks-on-success, + mixed-batch-with-multiple-roots, and the all-2D / all-3D-interior + cost-ceiling guarantees. +- `tests/graph/test_sv_lookup_utils.py` — the low-level + `get_atomic_id_from_coord` and `get_atomic_ids_from_coords` tests. diff --git a/pychunkedgraph/graph/sv_lookup/__init__.py b/pychunkedgraph/graph/sv_lookup/__init__.py new file mode 100644 index 000000000..cfdbb904e --- /dev/null +++ b/pychunkedgraph/graph/sv_lookup/__init__.py @@ -0,0 +1,13 @@ +from .main import resolve_supervoxels_at_coords +from .utils import ( + get_atomic_id_from_coord, + get_atomic_ids_from_coords, + lookup_svs_from_seg, +) + +__all__ = [ + "resolve_supervoxels_at_coords", + "get_atomic_id_from_coord", + "get_atomic_ids_from_coords", + "lookup_svs_from_seg", +] diff --git a/pychunkedgraph/graph/sv_lookup/main.py b/pychunkedgraph/graph/sv_lookup/main.py new file mode 100644 index 000000000..4ec306235 --- /dev/null +++ b/pychunkedgraph/graph/sv_lookup/main.py @@ -0,0 +1,101 @@ +"""Public SV-lookup orchestrator. + +See ``README.md`` for the segmentation-agnostic 2D/3D contract this +function enforces. +""" + +from typing import Optional +from typing import Sequence + +import numpy as np + +from .. import exceptions as cg_exceptions +from .utils import lookup_svs_from_seg + + +def _resolve_3d_with_root( + cg, + coords: np.ndarray, + root_id: np.uint64, + max_dist_steps: Sequence[float], +) -> np.ndarray: + """Per-root growing-radius parent-constrained search. + + Used for 3D click coords whose literal seg SV's root != ``root_id``. + One ``cg.get_atomic_ids_from_coords`` call per radius; breaks on + first success. Raises ``BadRequest`` if the largest radius fails. + """ + for max_dist_nm in max_dist_steps: + resolved = cg.get_atomic_ids_from_coords( + coords, parent_id=root_id, max_dist_nm=max_dist_nm + ) + if resolved is not None: + return np.asarray(resolved, dtype=np.uint64) + raise cg_exceptions.BadRequest( + f"Could not determine supervoxel ID for coordinates " + f"{coords.tolist()} - Lookup stage." + ) + + +def resolve_supervoxels_at_coords( + cg, + coordinates: Sequence[Sequence[int]], + node_ids: Sequence[np.uint64], + max_dist_steps: Optional[Sequence[float]] = None, +) -> np.ndarray: + """Resolve voxel coordinates to current supervoxel ids. + + Segmentation-agnostic (OCDBT or precomputed CV; dispatched by + ``get_local_segmentation``). The constraint never depends on + caller-supplied node ids beyond their layer: + + - layer 1 (2D slice click) -> accept the literal seg SV. + - layer >= 2 (3D mesh click, ``node_id`` interpreted as root) -> + accept the literal SV when its current root matches; otherwise + run a per-root parent-constrained nearest-SV search. + + Returns ``(N,)`` uint64. Raises ``cg_exceptions.BadRequest`` on + invalid input or unresolvable coords. + """ + coordinates = np.asarray(coordinates, dtype=int) + if coordinates.ndim != 2 or coordinates.shape[1] != 3: + raise cg_exceptions.BadRequest( + f"Could not determine supervoxel ID for coordinates " + f"{coordinates} - Validation stage." + ) + node_ids = np.asarray(node_ids, dtype=np.uint64) + if max_dist_steps is None: + max_dist_steps = np.array([4, 8, 14, 28], dtype=float) * np.mean( + cg.meta.resolution + ) + + lit_svs = lookup_svs_from_seg(cg.meta, coordinates) + layers = cg.get_chunk_layers(node_ids) + is_3d = layers >= 2 + out = lit_svs.astype(np.uint64, copy=True) + + if not is_3d.any(): + return out + + three_d_idx = np.where(is_3d)[0] + three_d_lit = lit_svs[three_d_idx] + nz_in_3d = three_d_lit != 0 + lit_roots_3d = np.zeros(len(three_d_idx), dtype=np.uint64) + if nz_in_3d.any(): + lit_roots_3d[nz_in_3d] = cg.get_roots(three_d_lit[nz_in_3d], fail_to_zero=True) + + needs_search = lit_roots_3d != node_ids[three_d_idx] + if not needs_search.any(): + return out + + search_idx_in_3d = np.where(needs_search)[0] + search_node_ids = node_ids[three_d_idx][search_idx_in_3d] + search_coords = coordinates[three_d_idx][search_idx_in_3d] + + for root_id in np.unique(search_node_ids): + root_mask = search_node_ids == root_id + sub_search_coords = search_coords[root_mask] + resolved = _resolve_3d_with_root(cg, sub_search_coords, root_id, max_dist_steps) + out_positions = three_d_idx[search_idx_in_3d[root_mask]] + out[out_positions] = resolved + return out diff --git a/pychunkedgraph/graph/sv_lookup/utils.py b/pychunkedgraph/graph/sv_lookup/utils.py new file mode 100644 index 000000000..9b97b0051 --- /dev/null +++ b/pychunkedgraph/graph/sv_lookup/utils.py @@ -0,0 +1,152 @@ +"""Low-level segmentation reads and parent-constrained nearest-SV search. + +See ``README.md`` for the overall contract. These primitives are +segmentation-agnostic: they go through ``get_local_segmentation`` which +dispatches to OCDBT or CloudVolume transparently. +""" + +from typing import Optional +from typing import Sequence +from typing import Callable +from datetime import datetime + +import numpy as np +import fastremap + +from ..meta import ChunkedGraphMeta +from ..utils.generic import get_local_segmentation + + +def lookup_svs_from_seg(meta: ChunkedGraphMeta, coordinates) -> np.ndarray: + """Read SV IDs at the given voxel coordinates. + + One batched seg read over the coords' bounding box; returns the + literal SV at each coord (0 for background). + """ + coordinates = np.asarray(coordinates, dtype=int) + bbox_start = coordinates.min(axis=0) + bbox_end = coordinates.max(axis=0) + 1 + seg = get_local_segmentation(meta, bbox_start, bbox_end)[..., 0] + local = coordinates - bbox_start + return seg[local[:, 0], local[:, 1], local[:, 2]].astype(np.uint64) + + +def get_atomic_id_from_coord( + meta: ChunkedGraphMeta, + get_root: Callable, + x: int, + y: int, + z: int, + parent_id: np.uint64, + n_tries: int = 5, + time_stamp: Optional[datetime] = None, +) -> np.uint64: + """Determines atomic id given a coordinate.""" + x = int(x / 2**meta.data_source.CV_MIP) + y = int(y / 2**meta.data_source.CV_MIP) + z = int(z) + xyz = np.array([x, y, z]) + + checked = [] + atomic_id = None + root_id = get_root(parent_id, time_stamp=time_stamp) + + for i_try in range(n_tries): + r = (i_try - 1) ** 2 + lo = np.maximum(xyz - r, 0) + hi = xyz + r + 1 + atomic_id_block = ( + meta.ws_ts[lo[0] : hi[0], lo[1] : hi[1], lo[2] : hi[2]].read().result() + ) + atomic_ids, atomic_id_count = np.unique(atomic_id_block, return_counts=True) + + sorted_atomic_ids = atomic_ids[np.argsort(atomic_id_count)] + sorted_atomic_ids = sorted_atomic_ids[~np.isin(sorted_atomic_ids, checked)] + + for candidate_atomic_id in sorted_atomic_ids: + if candidate_atomic_id != 0: + ass_root_id = get_root(candidate_atomic_id, time_stamp=time_stamp) + if ass_root_id == root_id: + atomic_id = candidate_atomic_id + break + else: + checked.append(candidate_atomic_id) + if atomic_id is not None: + break + return atomic_id + + +def get_atomic_ids_from_coords( + meta: ChunkedGraphMeta, + coordinates: Sequence[Sequence[int]], + parent_id: np.uint64, + parent_id_layer: int, + parent_ts: datetime, + get_roots: Callable, + max_dist_nm: int = 150, +) -> Sequence[np.uint64]: + """Parent-constrained nearest-SV search over multiple coords. + + Reads one bbox-sized seg block around the coords, masks every voxel + to its root via ``get_roots``, then picks the nm-closest voxel whose + root matches ``parent_id`` for each input coord. Returns ``None`` if + no voxel within ``max_dist_nm`` matches. + + :param coordinates: n x 3 np.ndarray of locations in voxel space + :param parent_id: parent id common to all coordinates at any layer + :param max_dist_nm: max distance explored + """ + if parent_id_layer == 1: + return np.array([parent_id] * len(coordinates), dtype=np.uint64) + + coordinates_nm = coordinates * np.array(meta.resolution) + max_dist_vx = np.ceil(max_dist_nm / meta.resolution).astype(dtype=np.int32) + bbox = np.array( + [ + np.min(coordinates, axis=0) - max_dist_vx, + np.max(coordinates, axis=0) + max_dist_vx + 1, + ] + ) + + local_sv_seg = get_local_segmentation(meta, bbox[0], bbox[1]).squeeze() + lower_bs = np.floor( + (np.array(coordinates_nm) - max_dist_nm) / np.array(meta.resolution) - bbox[0] + ).astype(np.int32) + upper_bs = np.ceil( + (np.array(coordinates_nm) + max_dist_nm) / np.array(meta.resolution) - bbox[0] + ).astype(np.int32) + local_sv_ids = [] + for lb, ub in zip(lower_bs, upper_bs): + local_sv_ids.extend( + fastremap.unique(local_sv_seg[lb[0] : ub[0], lb[1] : ub[1], lb[2] : ub[2]]) + ) + local_sv_ids = fastremap.unique(np.array(local_sv_ids, dtype=np.uint64)) + local_parent_ids = get_roots( + local_sv_ids, + time_stamp=parent_ts, + stop_layer=parent_id_layer, + fail_to_zero=True, + ) + + local_parent_seg = fastremap.remap( + local_sv_seg, + dict(zip(local_sv_ids, local_parent_ids)), + preserve_missing_labels=True, + ) + + parent_id_locs_vx = np.array(np.where(local_parent_seg == parent_id)).T + if len(parent_id_locs_vx) == 0: + return None + + parent_id_locs_nm = (parent_id_locs_vx + bbox[0]) * np.array(meta.resolution) + dist_mat = np.sqrt( + np.sum((parent_id_locs_nm[:, None] - coordinates_nm) ** 2, axis=-1) + ) + match_ids = np.argmin(dist_mat, axis=0) + matched_dists = np.array([dist_mat[idx, i] for i, idx in enumerate(match_ids)]) + if np.any(matched_dists > max_dist_nm): + return None + + local_coords = parent_id_locs_vx[match_ids] + matched_sv_ids = [local_sv_seg[tuple(c)] for c in local_coords] + return matched_sv_ids diff --git a/pychunkedgraph/graph/sv_split/NOTES.md b/pychunkedgraph/graph/sv_split/NOTES.md new file mode 100644 index 000000000..f9c852538 --- /dev/null +++ b/pychunkedgraph/graph/sv_split/NOTES.md @@ -0,0 +1,96 @@ +# sv_split — notes (known issues, test handles, future work) + +Secondary material that doesn't belong in `README.md` (the design reference). + +## Known issues + +### `{1, 2}` cross-side INF bridge through an unsplit partner + +Validation refuses the split when an L1 INF-affinity edge in an unsplit partner +SV connects a source-side fragment (label 1) to a sink-side fragment (label 2). +The carve inside the SV is correct in voxel space, but the cut surface +geometrically passes through the partner — INF means the partner is the same +segment as both fragments, and the chunked graph can't downgrade that affinity. +Mincut on the L2 graph would have an uncuttable source→partner→sink path. + +Current behavior: post-condition raises with the offending partner id. The user +re-issues with seeds that avoid the corridor through the partner, or includes +the partner in the seed set. A cascade-split — recursively splitting bridging +partners — is the long-term fix but expands the unit-of-edit and touches +lineage / op-log / multicut handshake. + +### Mincut precondition rejection after a successful SV-split + +On large objects the multicut may reject with `Sinks and sources are not +connected through the local graph` even though the SV-split, edge writeback, +and validation all pass. Open question: bbox window too tight to keep the +seeded L2 nodes in one component after the partner splits, or one of the +partner splits removed a bridge edge the mincut needed. + +## Future work + +### Dedup seg reads + subgraph fetches across reps in `split_supervoxels` + +When a multicut produces multiple cross-chunk reps that need splitting, +the orchestrator iterates `for task in tasks: split_supervoxel(task)` +and each rep independently (a) reads its bbox crop from OCDBT (~2–4 s +per call) and (b) fetches `cg.get_subgraph(root, bbox)` (~1.6 s per +call). When reps cluster spatially or share an L2 ancestor the same +OCDBT chunks and subgraph payload are fetched twice or more. + +Two surgical dedups, opt-in conditional on `len(tasks) > 1` so the +single-rep path is byte-equivalent to today: + +**Seg-read union per overlap cluster.** Cluster reps by overlap of +their padded voxel bboxes (`[bbs−1, bbe+1]`). Connected components over +that overlap relation = clusters. One `get_local_segmentation(union_bbs, +union_bbe)` per cluster; each rep takes a `.copy()` of its sub-slice at +the same point `_read_seg_and_ids` would have returned, so per-rep +mutation isolation is preserved (`_parse_results` and `mask_except` +both mutate seg in place). Saves the OCDBT I/O; per-rep copy cost +unchanged. + +**Subgraph union per shared root.** After `cg.get_roots([sv_id], parent_ts)` +resolves each rep's root, group reps by root. One `cg.get_subgraph(root, +union_bbox)` per shared root; each rep filters the returned edges by its +own `[bbs, bbe]`. `get_subgraph` already filters by bbox post-read +(subgraph.py:212–216 via `mask_nodes_by_bounding_box`), so a wider union +bbox returns a superset of any per-rep bbox's edges. Endpoints touching +multiple reps' bboxes get processed in each rep — matches today's +behavior since `_get_new_edges` is idempotent across duplicate edges. + +Correctness gates: a `TestMultiRepParity` covering both dedups against +the independent-per-rep baseline (byte-equal `seg_writes`, `bigtable_rows`, +`source_ids_fresh`, `sink_ids_fresh`, `old_new_map`, `new_id_label_map`) +plus the canonical pinky log lines remaining unchanged. Single-rep path +must take the original code path verbatim — the new helpers must not +fire when `n_tasks == 1`. + +Risks: union seg buffer holds longer than per-rep (mitigation: per-rep +copy at the same boundary today's path uses; union buffer freed after +the cluster's last rep finishes `_route_edges_and_rows`). Cluster +detection bug → larger union buffer, never wrong output. Cap union +volume to bound RSS on dense N-rep clusters. + +Not in scope: parallelizing reps, batching `id_client.create_node_ids`, +changing `_compute_split` / `_apply_and_capture` / `_parse_results` / +`_update_chunks`, modifying per-rep log lines. + +## Architecture + +The split algorithm lives in the external `supervoxel-splitter` package. +`splitter.get_splitter()` resolves an implementation class via the +`PCG_SV_SPLITTER` env var (dotted import path; default +`supervoxel_splitter.GeodesicSplitter`) and forwards `**kwargs` to its +constructor so call-site tuning propagates. `_coords.py` holds post-split +coord utilities consumed by `edges.py`. + +## Geodesic backend — `backend` kwarg on `GeodesicSplitter` + +`dj3d` (default) selects a faster geodesic kernel that has no anisotropy +parameter; the cost grid is pre-scaled by `mean(sampling_ds)` to approximate +per-axis anisotropy, and the cut surface diverges by a small amount on highly +anisotropic graphs where one axis is >5× the others. `mcp` selects the +anisotropy-correct kernel with per-axis sampling. Pass via +`get_splitter(backend="mcp")` from PCG to override (forwarded as a +`GeodesicSplitter` constructor kwarg). diff --git a/pychunkedgraph/graph/sv_split/README.md b/pychunkedgraph/graph/sv_split/README.md new file mode 100644 index 000000000..67b6be319 --- /dev/null +++ b/pychunkedgraph/graph/sv_split/README.md @@ -0,0 +1,706 @@ +# Supervoxel splitting + +Reference for the SV-split sub-flow of a multicut edit. Covers what the +operation does, where it sits in the edit pipeline, the voxel-level cut +algorithm, edge re-routing afterwards, the concurrency / storage contracts, +recovery, performance characteristics, and the failure modes. + +Companion: `NOTES.md` (sibling) holds known issues, open failure modes, and +future-work threads. Any code or design change that alters a contract / +invariant / failure mode covered here must update both files in the same +commit. + +--- + +## 1. Overview + +A *supervoxel split* bisects one physical supervoxel — a connected region in +the raw segmentation — along a user-seeded cut. The user supplies source and +sink seeds inside one SV; the system finds a cut surface, mints fresh L1 ids +for each side, rewrites the affected OCDBT segmentation chunks, and reroutes +every graph edge that used to reference the old SV. + +The SV-split runs only when the multicut detects that source and sink resolve +to the same physical SV across chunk boundaries — i.e. they sit in the same +INF-connected component of the chunk graph. A multicut cannot sever an +INF-affinity edge, so the only way to honour the user's cut is to actually +split the underlying voxels and produce new ids the graph mincut can then cut +between. + +Only graphs whose segmentation is OCDBT-backed can run SV-splits — the +operation needs a writable segmentation store. On read-only segmentation +backends the multicut surfaces a precondition error instead of attempting +the split. + +## 2. End-to-end flow + +``` +Split request (source coords, sink coords, source/sink ids) + │ + ▼ +Resolve coords → current L1 SV ids at those voxels + │ + ▼ +┌───────────────────────────────────────────────────────────────────────┐ +│ ROOT LOCK (held across the whole operation) │ +│ │ +│ MULTICUT #1: │ +│ build local subgraph around source / sink │ +│ merge cross-chunk INF edges into "rep" super-nodes │ +│ run mincut between source and sink │ +│ → Cut | SvSplitRequired(sv_remapping) | PreconditionError │ +│ │ +│ if SvSplitRequired: │ +│ ┌───────────────────────────────────────────────────────────────┐ │ +│ │ L2 CHUNK LOCK (temporal, per-chunk, 1-CG-chunk margin) │ │ +│ │ │ │ +│ │ PLAN: │ │ +│ │ enumerate one task per cross-chunk rep that bridges │ │ +│ │ src ↔ sink │ │ +│ │ bbox = full envelope of seed coords + 1 CG-chunk margin │ │ +│ │ │ │ +│ │ SPLIT (pure, no IO): │ │ +│ │ for each task: │ │ +│ │ read seg in [bbs − 1, bbe + 1] ← 1-voxel shell │ │ +│ │ geodesic cut → label map {0, 1, 2} │ │ +│ │ per chunk, ≥ 2 labels → mint fresh L1 ids │ │ +│ │ 1 label → keep original (no OCDBT) │ │ +│ │ route edges incident to split SVs (§4) │ │ +│ │ aggregate per-task results │ │ +│ │ │ │ +│ │ ┌─────────────────────────────────────────────────────┐ │ │ +│ │ │ INDEFINITE L2 CHUNK LOCK (durable scope record) │ │ │ +│ │ │ write seg chunks → OCDBT │ │ │ +│ │ │ persist BT rows → BT │ │ │ +│ │ └─────────────────────────────────────────────────────┘ │ │ +│ └───────────────────────────────────────────────────────────────┘ │ +│ │ +│ refresh src/sink ids to fresh fragments │ +│ │ +│ MULTICUT #2 (retry against post-split graph): │ +│ → Cut | SvSplitRequired │ +│ if SvSplitRequired: raise PreconditionError │ +│ │ +│ COMMIT CUT: │ +│ remove the cut atomic edges │ +│ mint new roots, write hierarchy + op-log │ +└───────────────────────────────────────────────────────────────────────┘ + │ + ▼ +Release root lock — edit is durable + │ + ▼ +Publish pubsub message; on SV-split, carries seg bboxes of rewritten regions + │ + ▼ +Async mesh / downsample workers consume bboxes under their own pyramid-block locks +``` + +Key timing rules: + +- `SvSplitRequired` is a **return value**, not an exception. The root lock + stays held across detect → split → retry → commit; nothing unwinds. +- SV-split writes (OCDBT seg + BT rows) land **before** the retry multicut. + If the retry rejects the cut, the SV writes stay durable (orphan + fragments) — the op aborts to the user but the segmentation has moved. +- SV-split is per-rep. Each cross-chunk-connected rep that bridges src↔sink + is handled by its own task under the same L2 lock set. + +## 3. Voxel-level cut + +The cut is a **geodesic region grow**, not a voxel graph mincut. Each voxel +of the SV being split is assigned to whichever seed it reaches first in an +anisotropy-aware geodesic cost field. + +### Inputs + +- A 3D boolean mask of the SV's voxels in the task bbox. +- Source and sink seed coords (mip0 voxel space). +- Anisotropic voxel sampling in nm. + +### Steps + +1. **Foreground bbox crop.** Tight bbox of the True voxels; the geodesic + runs on this sub-volume only. Geodesic-algorithm allocations scale with + foreground extent, not task bbox. +2. **Seed ridge bridge + snap.** Build a path from each seed to a + foreground ridge; snap each seed to the nearest in-mask voxel via + kdtree. Seeds need not land on a foreground voxel. +3. **EDT → speed → travel cost.** Anisotropic Euclidean distance transform + of the mask; speed = (distance / max distance)^γ_neck (γ ≈ 1.6), + clipped to a small floor; travel cost = 1 / speed. Interior is cheap, + neck is expensive — the cut tracks the thinnest section of the SV. +4. **Geodesic arrival.** Compute arrival times from each seed set on the + travel-cost field with the anisotropic sampling baked in; each voxel + goes to the side it reaches more cheaply. Optional axis-wise + downsample keeps the geodesic in physical-nm units regardless of stride. +5. **Narrow-band proximity boost.** Voxels within a relative threshold of + the opposing side get a cost boost so the boundary tracks the midline. +6. **Single-CC enforcement + stray resolution.** For each side keep the + 26-connected component(s) containing a seed; relabel orphan CCs to a + transient label-3. Then for every label-3 voxel, look up its 26 + neighbours' labels and per-component assign the side with the majority + of label-1 / label-2 neighbours; ties break on per-voxel anisotropic + distance to the nearest source vs sink seed. Work scales with the + label-3 voxel set, not the volume. See §12 for the full label-3 + lifecycle and why a second-pass enforce may leave some label-3 + fragments for the mincut to absorb. +7. **Embed back.** Result is a small-int label map over the original task + bbox; voxels outside the foreground crop are 0. + +### Single-CC invariant + +The chunk-update step mints **one fresh id per distinct label value per L1 +chunk** and does **no CC of its own**. If a side left two disconnected +pieces inside one chunk, both would collapse into one new id and the chunk +graph would gain a spuriously-connected fragment. Therefore every label +out of the geodesic must already be a single 26-connected component per +chunk — enforced at full resolution because upsampling under a foreground +mask can fragment. + +### Design rationale + +- **Geodesic region grow over voxel graph mincut.** A geodesic in an + anisotropy-aware cost field follows the SV's medial geometry and yields a + smooth midline cut without building / solving a per-voxel adjacency graph + on every split. +- **Snap seeds to the foreground.** Seeds from operator clicks or upstream + mincut output need not land on a true voxel; snapping keeps both arrival + fields rooted inside the SV. +- **Resolve label-3 strays, don't drop them.** Dropping voxels loses mass; + merging blindly can bridge sides. Border-count reassignment keeps every + voxel while respecting the cut. Label-3 fragments that survive after the + second-pass enforce propagate through to the mincut as their own SV ids + and the cut resolves them topology-first; see §12. +- **Per-voxel scan over full-volume dilation.** Stray border counts visit + only the label-3 voxel set and their 26 neighbours; work scales with + the size of that set, not the volume. +- **Reads pinned to the operation's parent timestamp.** Every graph read + during a split (parents, cross-chunk edges, SV-to-root map) is pinned so + a replay sees the same graph state and allocates the same fresh SV ids — + what makes interrupted splits safe to re-run under recovery. + +## 4. Edge re-routing + +After the chunk-write step produces an `old → fresh-fragments` map and a +`fragment id → cut-side label ∈ {1, 2, 3}` map (see §12 for label-3), +update every atomic edge that used to reference an old SV. + +### Sketch + +``` +update_edges: + 1. fetch atomic subgraph in [bbs, bbe] + 2. dedup, drop self-loops + 3. resolve partner roots via one batched call + 4. for each old SV being split: + inactive partner → broadcast edge to every fragment + active partner + INF + partner-split → match by cut-side label + INF + partner-unsplit → closest-fragment-only + finite → connect every fragment within threshold; + fall back to closest + 5. inter-fragment 0.001 bridges between every fragment pair + 6. validate (4 invariants — §4c) + 7. return (edges, affinities, areas) + +add_new_edges: + 1. duplicate bidirectional + 2. group by partner's L2 parent chunk + 3. per chunk: append to the append-only split-edge history, + rewrite the latest-only compacted snapshot (stale-filtered) +``` + +### Routing rules + +For each edge incident to a split SV, the partner's root determines the +path: + +- **Inactive partner** (different root). Broadcast: every fragment gets a + copy of the edge with affinity / area preserved. Costless if the roots + stay apart; collapses harmlessly to one root-level edge if they later + merge. +- **Active partner + INF + partner also split.** Connect each fragment to + the partner's fragment with the same cut-side label. Fallback (no + matching label in this iteration's new ids) writes closest-fragment INF + — a known class-C bridge risk (§9). +- **Active partner + INF + partner unsplit.** Write the edge to the + single closest fragment only. Broadcasting INF to both sides would form + an uncuttable bridge for the retry mincut. +- **Active partner + finite affinity.** Connect every fragment within the + threshold (configurable per-graph); fall back to closest. + +### Inter-fragment 0.001 bridges + +For each old SV with ≥ 2 fragments, every fragment pair gets a finite-0.001 +edge. These are **cuttable** by the retry multicut — they are the route the +mincut actually severs to separate the source side from the sink side. + +### Post-route validation + +| check | invariant | raises | +|---|---|---| +| A | No `{1, 2}` (source-side ↔ sink-side) INF bridge through any unsplit partner | `PostconditionError` | +| B | No self-loops | `PostconditionError` | +| C | Every old SV has at least one replacement edge | `PostconditionError` | +| D | Every fragment pair from same old SV has the 0.001 bridge | `PostconditionError` | + +Check A only fires on `{1, 2}` — bridges involving label-3 fragments are +valid (the unresolved fragment rides with whichever seeded side the +inf-cluster joins; see §12 deep dive). + +Failures abort before any write lands. + +### Distance computation + +- **Partner inside bbox.** Kdtree over the partner's voxels in the seg + crop; per fragment, query against it and return the minimum voxel + distance. +- **Partner outside bbox.** Fragment kdtree to the partner's chunk + boundary face. Over-estimate for non-boundary-aligned partners but the + only signal available without extra reads. + +### Persistence + +Edges write per L2 chunk into two parallel columns: + +- **Append-only split-edge history.** Time-travel reads at any timestamp + `T` walk cells with `ts ≤ T` and apply stale-edge resolution. +- **Latest-only compacted snapshot.** Single fresh cell per op. Previous + compacted rows are loaded, rows referencing any SV in the old→new map + are filtered out, new rows are appended, and the union is written + back. Current-time readers take this single cell directly. + +Edges write at the operation's logical timestamp so a parent-timestamp- +filtered reader sees atomic visibility. + +## 5. Concurrency + +Three lock layers, each scoped to the smallest window that closes its race +class: + +| lock | scope | duration | races closed | +|---|---|---|---| +| Root lock | the op's root id(s) | entire op | same-root edit interleaving | +| L2-chunk lock (temporal) | every L1 chunk the SV-split touches | SV-split + retry | cross-root spatial races during SV-split read / compute | +| L2-chunk lock (indefinite) | same chunks | OCDBT + BT writes | crash-mid-write isolation (§7) | + +### L2 chunk-set computation + +Walk each task's `[bbs − 1, bbe + 1]` bbox (the 1-voxel shell the edge +router needs to read), map it to overlapping L1 chunks, union across +tasks, and return a deterministically sorted list so workers with +overlapping sets never acquire in opposing orders. + +### Write-scope minimization + +Only chunks with ≥ 2 distinct labels (geodesic actually split the rep +there) get OCDBT writes. Annulus pieces — chunks where the rep is single- +label — keep their original L1 ids and skip the seg write entirely; the +segmentation backend is append-only so writing unchanged bytes would +inflate the on-disk delta for no value. + +### Post-split id refresh without an extra read + +After the SV-split lands, the original source/sink ids reference now- +superseded SVs. The retry multicut needs the *current* ids at the seed +voxels. The in-memory seg block produced during split is bitwise +identical to what just landed on OCDBT — the write is synchronous and +happens under the L2 lock, so nothing else mutated those voxels — and the +fresh source/sink ids fall out of the same in-memory lookup. No extra +round-trip. + +## 6. Storage substrate + +### Edges + +| source | role | +|---|---| +| Bucket (GCS protobuf), in-chunk channel | both endpoints in one L1 chunk | +| Bucket, cross-chunk channel | endpoints in different L1 chunks; INF-affinity (watershed's "same SV across this face" claim) | +| Bigtable, fake-edges column | operator-added merge edges | +| Bigtable, split-edge history columns | append-only SV-split history (edges + affinities + areas) | +| Bigtable, compacted split-edge columns | latest-only snapshot | + +The subgraph fetch merges bucket and bigtable edges, then drops edges +incident to SVs whose row carries a "new-identity" marker and are not in +the live agglomeration's parent map. On a clean table the filter is a +no-op (no bigtable L1 edges, no new-identity rows). + +### Hierarchy rows (per L1 SV) + +- Parent: single L1 parent id. +- Former-identity: array of node ids this row was created from. +- New-identity: array of node ids replacing this row (set on split). +- Per-layer cross-chunk edges. + +On split, the lineage step writes former-identity and the operation id on +each new fragment, copies the parent pointer (preserving the parent cell's +timestamp), updates the parent's child list to replace the old SV with +the new fragments, and writes new-identity on the old SV. + +### OCDBT segmentation + +Reads pull from the OCDBT-backed segmentation store when available, +falling back to CloudVolume otherwise. Writes aggregate all +`(voxel slices, data)` pairs across tasks into one flat list and issue +all tensorstore futures in parallel; per-task / per-rep loops would +serialize wall time. On-disk OCDBT config is authoritative — opening an +existing OCDBT must not pass a top-level `"config"` key; the manifest is +the source of truth. + +### ID allocation + +A per-chunk counter atomically increments, claims a contiguous range, and +OR-masks with the chunk id. Failed ops leak ids — the counter never rolls +back. + +### Locks + +Per-row columns for temporal and durable locks. Per-chunk rows are hash- +prefixed for write distribution. Acquire refuses if either column is set. +Value-matched release on `__exit__` prevents an op from clearing a lock +another op holds. The op log durably records the chunk set the indefinite +lock covers — recovery reads this to know which chunks need cleanup (§7). + +## 7. Recovery — worker crash mid-write + +Both writes inside the indefinite L2 chunk lock — OCDBT seg + BT rows — +must land for the op to be consistent. A worker death inside that block +leaves: + +- Indefinite lock cells set on the affected chunks. +- The op log's chunk-set record durably populated with the chunk ids. +- Possibly partial OCDBT writes and zero / partial BT rows. + +Future ops on any of those chunks refuse to start (the lock blocks them) +— the crashed state is isolated, not amplified. + +The authoritative signal that an op is stuck is the chunk-set record +non-empty past the clean exit point. A minimum-age threshold (≈ 10 min) +filters in-flight ops from definitively-dead ones. + +### Why a single pinned read is not enough + +The SV-split reads a 1-voxel shell around each chunk; the shell's +neighbouring chunks may have been mutated by other ops since the crash. +A single pinned read of the world at the op's parent timestamp would +return stale neighbour values; routing fresh edges to those stale ids +corrupts the graph. Recovery cannot rely on a single pinned view. + +### Cleanup-then-replay + +1. **Cleanup.** For each chunk in the crashed op's lock scope: read the + chunk's voxels at the op's parent timestamp (pinned handle), write + those values back to the latest (unpinned) handle. The crashed op's + chunks now show pre-op state at the latest manifest; neighbour chunks + and any concurrent ops' work are untouched. +2. **Replay.** Re-run the op under the privileged-repair path. Reads see + pre-op values on the op's chunks + current state on every other chunk + — a consistent world. Allocates fresh ids (the crashed op's ids leak), + writes new seg + hierarchy, lands the op-log row at success. The + indefinite lock's `__exit__` value-matched-releases the cells the + crashed op originally set (replay reuses the operation id), freeing + the chunks. + +### Orphan history + +OCDBT is append-only — the crashed op's partial writes still exist in +OCDBT commit history, just overshadowed at the latest manifest. Readers +pinning a historical version between crash and replay still see the +partial state; readers at latest never observe it. Orphan ids are never +referenced by any hierarchy row. + +## 8. Performance + +### Geodesic + +- **Foreground crop.** Reducing the geodesic's working volume to the + tight bbox of the SV's True voxels is the biggest single lever; big + win for thin reps; near no-op for dense reps that fill the bbox. +- **Downsample.** Axis-wise stride is configurable; the geodesic stays + in physical-nm units regardless of the stride choice. The principled + derivation is `resolution.max() // resolution`, which makes the + operation isotropic in nm space across any voxel anisotropy. +- **Narrow-band refinement.** Refines the cut surface near the boundary + at full resolution regardless of the global downsample. +- **Single-CC at full res.** The single-CC invariant must hold at chunk + granularity; running it on the DS grid would let upsampling fragment + a side into disconnected pieces. +- **Enforce on the foreground crop, not the read bbox.** Labels only + exist inside the foreground halo crop (the bbox of the union of cut + SVs, padded by one voxel). Single-CC enforcement and the label-3 + resolver pass a view of that crop to cc3d instead of the full read + bbox. cc3d's connected-component output scales with input volume, so + the savings grow with the gap between read-bbox and foreground — + largest on graphs with big `CHUNK_SIZE` where the one-chunk margin + inflates the read bbox far past the actual foreground. +- **Parallel arrival fields.** Source-side and sink-side arrival times + are computed concurrently in a two-worker fork pool. The underlying + geodesic kernel holds the GIL, so process-level parallelism is the + only lever; roughly halves the wall vs sequential. +- **Configurable backend.** `PYCG_GEODESIC_BACKEND=dj3d` (default) is the + faster kernel; it approximates anisotropy via a single mean-scale factor + on the cost grid. `=mcp` is anisotropy-correct via per-axis sampling; the + cut surface diverges by a small amount on highly-anisotropic graphs + depending on which backend is enabled. + +### Edge re-routing + +- One subgraph fetch per task at the task's `[bbs, bbe]`. +- Distances are kdtree queries on per-fragment voxel sets — no full- + volume scans. +- Edge writes per L2 chunk; bidirectional duplication; one parallel + batch. + +### Writes + +- OCDBT: only changed chunks. Annulus / shell-only chunks are untouched. +- Bigtable: lineage + new edges per affected chunk, one batched write. +- Seg-chunk writes aggregate all `(slices, data)` across reps into one + flat tensorstore future list — never per-rep loops. + +### Lock hold time + +- Temporal L2 lock covers the read + geodesic + edge-route compute — the + longest window in the SV-split. Released as soon as the compute + finishes. +- Indefinite L2 lock covers only the persist block — short, but durable + if the worker dies inside it. + +## 9. Failure modes + +| trigger | message / type | +|---|---| +| src and sink resolve to different roots | `PreconditionError("Supervoxels must belong to the same object ...")` | +| retry multicut still returns SvSplitRequired | `PreconditionError("Supervoxel split succeeded but source and sink remain connected; place source and sink farther apart.")` | +| mincut produced no removable edges | `PostconditionError("Mincut could not find any edges to remove.")` | +| no edges in retry's local subgraph | `PreconditionError("No local edges found.")` | +| src and sink in different CCs of retry's local subgraph | `PreconditionError("Sinks and sources are not connected through the local graph.")` | +| in-mask seed connection failed | `RuntimeError("In-mask connection failed for at least one team; skipping split.")` | +| new fragment landed in different chunk than its old SV | `PreconditionError("new supervoxel landed in a different chunk than the SV it split from")` | +| `{1, 2}` cross-side INF bridge via unsplit partner / self-loop / missing replacement / missing 0.001 bridge | `PostconditionError(...)` | + +### Diagnostic artifact on failure + +Every edit-flow `AssertionError` / `RuntimeError` / unknown `Exception` +writes a JSON snapshot to +`{cg.meta.data_source.WATERSHED}/graphene_errors/{cg.graph_id}/{op_id}.json` +before the exception propagates. The artifact contains op type, user, +source / sink ids + coords, removed / added edges, parent_ts, the +exception class + message, and the full traceback — enough to replay +the payload against current state. The error log line for the +operation includes the artifact URL. + +`PreconditionError` / `PostconditionError` (user-facing) are not +dumped. Read an artifact back with +`pychunkedgraph.graph.err_dump.read_err_artifact(cg, op_id)`. + +### Bridge classes (residual) + +After the retry mincut runs, the only structural paths from src-side to +sink-side are: + +- **0.001 inter-fragment bridges** (cuttable; intended). +- **Cross-chunk INF edges** routed by the by-label / closest-unsplit + rules (uncuttable; should be label-pure by construction). + +Known residual risks: + +- **C — by-label cross-label fallback.** When the iteration's new-ids set + contains fragments of only one label and the partner has the other + label, the fallback writes closest-fragment INF. Bridge survives. Fix + candidates: drop the edge or write finite cuttable affinity. +- **Annulus↔annulus bucket INF.** Two unsplit annulus pieces routed to + opposite labels by their respective closest-unsplit calls retain their + direct bucket INF edge (the routing loop only iterates over olds being + split — never sees unsplit↔unsplit pairs). Fix candidates: force-split + the bridge endpoints, or detect-and-abort with a dedicated error. + +## 10. Invariants + +- SV-split + retry multicut + commit are one atomic operation under a + single root lock. +- Within the SV-split step, concurrent splits on overlapping L2 chunks + serialize via the temporal L2 lock. +- OCDBT writes touch only chunks whose voxels actually changed. Annulus + and shell-only chunks are untouched. +- Every fresh L1 id minted in the SV-split has at least one voxel in + OCDBT carrying that id. OCDBT seg's per-voxel id ≡ the chunk graph's + child-derived L1 set at that voxel. +- Every label out of the geodesic is a single 26-connected component per + L1 chunk. +- After commit, readers at the op's timestamp see new SV ids in the cut + region and new roots reflecting the cut. +- Coarser MIP levels are eventually consistent with mip0, lagging at most + until the async downsample worker processes the pubsub message. + +## 11. Coordinate units + +- HTTP boundary: nm. +- After source/sink resolution: mip0 voxels. +- Bounding-box offset: mip0 voxels (default ≈ `(120, 120, 12)`). +- L1 chunk dimensions: mip0 voxels (graph config). +- Chunk grid origin: `voxel_bounds[:, 0]` — chunk `(0,0,0)` is at + `voxel_bounds[:, 0]`, not 0. + +## 12. Deep dive — the label-{1, 2, 3} lifecycle + +The geodesic cut produces a per-voxel label volume. Its valid values +evolve as the cut pipeline runs. Understanding why each step exists, and +what label-3 means at each stage, is the key to reading the routing and +the post-route validation. + +### Label legend (within the SV being split) + +| value | meaning at the end of a step | +|---|---| +| `0` | background — voxel outside the SV mask | +| `1` | source-side fragment voxel (source seed reaches it first) | +| `2` | sink-side fragment voxel (sink seed reaches it first) | +| `3` | **transient unresolved voxel** — a fragment the system has not yet decided belongs to source or sink | + +Label-3 is *internal* to the cut pipeline. After commit, no graph node +carries label 3 — every voxel ends up rooted under either the source-side +or sink-side new root after the retry mincut. + +### Step-by-step rationale + +**1) Geodesic cut → labels `{0, 1, 2}`.** `argmin(arrival_from_source, +arrival_from_sink)` is the geodesic Voronoi partition: each foreground +voxel is assigned to whichever seed set reaches it first along the cost- +weighted shortest path in the anisotropic field. This is the cut surface +as a voxel-level partition. + +**2-3) 1st-pass single-CC enforcement, each side, with stray-demotion +enabled.** The Voronoi partition can leave **islands** of a side stranded +— e.g. a thin pocket the geodesic happens to assign to side A but with no +internal path to any source seed without crossing label-2. The chunk- +update step commits one new SV per `(old_sv × connected_component × +label)`, so each `(old_sv, label)` must be a single CC inside one L1 +chunk. + +The rule: keep CCs that contain at least one **seed of that side** (those +represent the legitimate side territory). Demote every other CC. Demote +to **3** rather than back to the opposite side, because we don't actually +know if the island belongs to the opposite side — it might be +geographically inside one side's main mass with the geodesic having taken +a sub-optimal arc through it. Label 3 is the contract: *"unresolved — +let the smarter downstream resolver decide."* + +**4) Stray resolver.** For each connected component of label-3: + +- **Border vote (majority).** Count how many 26-neighbours of the CC's + voxels carry label 1 vs label 2. Majority side wins. Rationale: an + island's natural home is whatever it's geographically nestled inside; + local topology is the strongest signal for which side it *should* + belong to. +- **Seed-distance tiebreak.** If border counts are equal, fall back to + the per-voxel anisotropic distance to the nearest seed of each side + (kdtree on the seed point sets, scaled by sampling). Closer side + wins. Global fallback when local context is symmetric. + +After this step every voxel originally demoted in steps 2-3 has been +reassigned. **Invariant at exit of the resolver: labels ∈ `{0, 1, 2}`.** + +**5-6) 2nd-pass single-CC enforcement, each side, with stray-demotion +enabled.** Step 4 just moved potentially many voxels from label-3 to +label-1 / label-2. Those newly-reassigned voxels can form **new** +disconnected CCs of their new side — an island resolved to label-1 may +be sitting far from the main label-1 mass with no internal label-1 path +between them. The single-CC-per-chunk invariant must hold at the end of +the pipeline; this second pass re-checks. + +Same contract as steps 2-3: orphan CCs (no seed inside them) get +demoted to label 3. + +**7) End — label-3 fragments may survive.** Unlike the first round, no +third resolver pass runs. The label-3 voxels that the second pass +produced do not get reassigned. This is intentional — see step 8. + +**8) Per-chunk fresh-id minting.** The chunk-update step mints one fresh +L1 SV id per distinct label value per L1 chunk. Label 3 is treated like +any other label: every label-3 CC inside a chunk becomes its own new SV +id, and the routing step records the fragment-id → label-3 mapping so +edge routing knows how to handle it. + +**9) Edge routing treats label-3 fragments as first-class new SVs:** + +- Low-affinity inter-fragment edges (affinity `0.001`) are emitted for + *every pair* of fragments of the same old SV — including all pairings + with label-3 fragments. These edges are cuttable by the mincut. +- INF-affinity edges to **split partners** are routed by same-label + match — a label-3 fragment of this SV connects via INF to a label-3 + fragment of a split partner, if one exists. Two unresolved fragments + inf-joined is fine; both ride together. +- INF-affinity to **unsplit partners** uses closest-fragment-only. The + partner can land inf-edged to a label-1 fragment via one cross-chunk + face and inf-edged to a label-3 fragment via another, producing a + `{1, 3}` bridge through the partner. +- Inactive partners (different root) broadcast to every fragment, + including label-3. + +**10) Post-route validation accepts label-3 bridges.** The cross-label +check only raises on `{1, 2}` — bridges through an unsplit partner +involving only seeded-side labels. `{1, 3}`, `{2, 3}`, and `{3}` bridges +are valid because the unresolved fragment has no seed; the mincut places +the entire inf-connected cluster on whichever side carries the seed. + +**11) Retry multicut places label-3 fragments.** The mincut builds a +graph where: + +- Source nodes = the new SV ids at each source seed voxel — anchored by + *physical seed voxel location*, not by the label the fragment carries. +- Sink nodes = the new SV ids at each sink seed voxel. +- INF edges are uncuttable; the `0.001` inter-fragment edges are cheap + to cut. + +The cut algorithm finds the minimum-cost edge set to remove that +disconnects source nodes from sink nodes. Label-3 fragments end up on +whichever side has lower cut cost. For a partner-cluster `{label-1 frag, +label-3 frag, unsplit partner}` joined by INF, the cluster is one +indivisible unit; it joins the source side via the label-1 fragment's +source seed; the label-3 fragment rides along. Symmetrically for `{label-2 +frag, label-3 frag, partner}` → sink side. + +When a label-3 fragment has no INF anchor to either side, the cut runs +through its `0.001` neighbours: it ends up wherever the global cost +optimum places it. + +### Why this works without a third resolver pass + +Re-running the resolver after the second-pass enforce is *one option* — +it would produce labels `{0, 1, 2}` once again and the validate-routing +chain could keep the old "no cross-label" rule. The model we use instead +is: **the heuristic resolver and the global mincut decide different +questions, and the mincut is qualified to absorb whatever the heuristic +left undecided.** + +- The resolver's border vote answers *"based on the local 26- + neighbourhood, which side does this voxel belong to?"* — strong signal + when the neighbourhood is asymmetric, no signal when it is symmetric + or the fragment sits in a thin neck. +- The mincut answers *"given the global graph topology (INF clusters, + inter-fragment 0.001 edges, partner connectivity), what's the minimum- + cost cut that separates source seeds from sink seeds?"* — uses + information the resolver doesn't have (cross-chunk INF edges, partner + identity, multi-fragment routing). + +When the second-pass enforce produces label-3 voxels, the local signal +has *already failed once* (the resolver placed them, and the placement +created a new disconnected CC). Re-applying the same border vote is +unlikely to be more correct than letting the mincut weigh the global +topology. So we publish the label-3 fragment as a real SV, route it, and +let the multicut place it. + +### Mental model — labels are a routing hint, not a cut constraint + +The fragment-id → label-{1, 2, 3} mapping is consumed by *one* thing: +the edge router, to decide which fragment of *this* split SV gets +inf-edged to each fragment of a *split partner* SV via the by-label +match rule. After routing, the multicut treats every fragment as just a +node in a weighted graph — the label has no further role. Source and +sink anchors come from the *physical seed voxel position*, not the label +of the fragment that ended up at that voxel. This is what lets label-3 +fragments propagate through harmlessly: they are nodes with no seed +anchor, sitting wherever the cut places them. diff --git a/pychunkedgraph/graph/sv_split/__init__.py b/pychunkedgraph/graph/sv_split/__init__.py new file mode 100644 index 000000000..4c8c54c70 --- /dev/null +++ b/pychunkedgraph/graph/sv_split/__init__.py @@ -0,0 +1,5 @@ +""" +Supervoxel splitting. +""" + +from . import edges, edits diff --git a/pychunkedgraph/graph/sv_split/_coords.py b/pychunkedgraph/graph/sv_split/_coords.py new file mode 100644 index 000000000..9c8f0a0a6 --- /dev/null +++ b/pychunkedgraph/graph/sv_split/_coords.py @@ -0,0 +1,62 @@ +"""Per-label coordinate utilities used by edge routing after a SV split.""" + +from typing import Dict, Iterable, Optional + +import fastremap +import numpy as np + +from pychunkedgraph.profiler import get_profiler + +_prof = get_profiler() + + +def build_coords_by_label( + vol: np.ndarray, + *, + labels: Optional[Iterable[int]] = None, + background: int = 0, + min_points: int = 1, + dtype: np.dtype = np.float32, + boundary_only: bool = False, +) -> Dict[int, np.ndarray]: + """Group voxel coords by label via ``fastremap.point_cloud``. + + Returns ``{label: (M_label, 3) coords in (z, y, x)}`` cast to + ``dtype``. ``fastremap.point_cloud`` is a C++ single-pass + implementation that emits ``uint16`` coords grouped by label, + treating ``0`` as background. + + ``labels`` restricts the output dict; the underlying C++ scan + visits every voxel regardless (faster and lighter than a + label-filtered Python scan). ``min_points`` drops labels with + fewer than that many voxels. ``background != 0`` removes that + label from the result after the call. + + ``boundary_only=True`` emits only 6-conn boundary voxels per label + via fastremap's native ``shell=True`` path — no extra allocation, + no mutation of ``vol``. min-distance between any two labels' + boundary point sets equals min-distance between their interior + point sets, so this is correctness-preserving for nearest-neighbor + consumers. + """ + if vol.ndim != 3: + raise ValueError("`vol` must be a 3D array.") + with _prof.profile("point_cloud"): + raw = fastremap.point_cloud(vol, shell=boundary_only) + if background != 0: + raw.pop(background, None) + + if labels is not None: + wanted = {int(x) for x in labels} + if not wanted: + return {} + items = ((k, v) for k, v in raw.items() if int(k) in wanted) + else: + items = raw.items() + + result: Dict[int, np.ndarray] = {} + for k, coords in items: + if coords.shape[0] < min_points: + continue + result[int(k)] = coords.astype(dtype, copy=False) + return result diff --git a/pychunkedgraph/graph/sv_split/edges.py b/pychunkedgraph/graph/sv_split/edges.py new file mode 100644 index 000000000..e4d1a33b9 --- /dev/null +++ b/pychunkedgraph/graph/sv_split/edges.py @@ -0,0 +1,466 @@ +""" +Edge routing logic for supervoxel splits. + +When a supervoxel (SV) is split into multiple fragments, all edges that +connected the original SV to its neighbors must be reassigned to the +appropriate new fragment(s). This module handles that reassignment. + +Edge classification: + Active edges: partner SV shares the same root as the split SV. + These edges are routed based on affinity type: + - Inf-affinity (cross-chunk) to a split partner: matched by split label, + connecting fragments that received the same label during the split. + - Inf-affinity (cross-chunk) to an unsplit partner: assigned to the + closest fragment only. Broadcasting to all fragments would create an + uncuttable bridge between source/sink sides of the split. + - Finite-affinity: assigned to fragments within a distance threshold + of the partner, or the closest fragment if none are within threshold. + + Inactive edges: partner SV has a different root. + These are edges to neighboring objects. All fragments inherit the edge + since any fragment could border the neighbor. + +Distance computation: + For partners within the segmentation bbox, distances are precomputed via + kdtree pairwise distances. For active partners outside the bbox (e.g. + cross-chunk fragments not in the rep's CC member set), distances are + computed from each new fragment's kdtree to the partner's chunk boundary. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from datetime import datetime + +import fastremap +import numpy as np + +from pychunkedgraph import get_logger +from pychunkedgraph.profiler import get_profiler +from pychunkedgraph.graph import attributes, basetypes, serializers +from pychunkedgraph.graph.chunks import utils as chunk_utils +from pychunkedgraph.graph.exceptions import PostconditionError +from pykdtree.kdtree import KDTree as cKDTree + +if TYPE_CHECKING: + from pychunkedgraph.graph.chunkedgraph import ChunkedGraph + +logger = get_logger(__name__) + + +def _match_by_label(new_ids, partner, aff, area, new_id_label_map, distances_row): + """For inf-affinity (cross-chunk) edges: connect fragments with matching split label.""" + partner_label = new_id_label_map[partner] + matching = np.array( + [nid for nid in new_ids if new_id_label_map.get(nid) == partner_label], + dtype=basetypes.NODE_ID, + ) + if len(matching): + edges = np.column_stack( + [matching, np.full(len(matching), partner, dtype=np.uint64)] + ) + affs = np.full(len(matching), aff, dtype=basetypes.EDGE_AFFINITY) + areas = np.full(len(matching), area, dtype=basetypes.EDGE_AREA) + return edges, affs, areas + # fallback: closest fragment + close = new_ids[np.argmin(distances_row)] + return ( + np.array([[close, partner]], dtype=np.uint64), + np.array([aff], dtype=basetypes.EDGE_AFFINITY), + np.array([area], dtype=basetypes.EDGE_AREA), + ) + + +def _match_by_proximity(new_ids, partner, aff, area, distances_row, threshold): + """For regular edges: connect fragments within distance threshold.""" + close_mask = distances_row < threshold + nearby = new_ids[close_mask] + if len(nearby): + edges = np.column_stack( + [nearby, np.full(len(nearby), partner, dtype=np.uint64)] + ) + affs = np.full(len(nearby), aff, dtype=basetypes.EDGE_AFFINITY) + areas = np.full(len(nearby), area, dtype=basetypes.EDGE_AREA) + return edges, affs, areas + close = new_ids[np.argmin(distances_row)] + return ( + np.array([[close, partner]], dtype=np.uint64), + np.array([aff], dtype=basetypes.EDGE_AFFINITY), + np.array([area], dtype=basetypes.EDGE_AREA), + ) + + +def _match_inf_unsplit(new_ids, partner, aff, area, distances_row): + """Inf-affinity edge to an unsplit partner: assign to closest fragment only. + Connecting all fragments would create an uncuttable bridge between source/sink sides. + """ + closest = new_ids[np.argmin(distances_row)] + return ( + np.array([[closest, partner]], dtype=np.uint64), + np.array([aff], dtype=basetypes.EDGE_AFFINITY), + np.array([area], dtype=basetypes.EDGE_AREA), + ) + + +def _match_partner( + new_ids, partner, aff, area, distances_row, new_id_label_map, threshold +): + """Route a single old edge to the appropriate new fragment(s).""" + if np.isinf(aff): + if new_id_label_map and partner in new_id_label_map: + return _match_by_label( + new_ids, partner, aff, area, new_id_label_map, distances_row + ) + return _match_inf_unsplit(new_ids, partner, aff, area, distances_row) + return _match_by_proximity(new_ids, partner, aff, area, distances_row, threshold) + + +def _expand_partners(active_partners, active_affs, active_areas, old_new_map): + """If a partner was also split, expand it to its new fragment IDs.""" + remapped_lists = [ + np.asarray(list(old_new_map.get(p, {p})), dtype=np.uint64) + for p in active_partners + ] + if not remapped_lists: + return ( + [], + np.array([], dtype=basetypes.EDGE_AFFINITY), + np.array([], dtype=basetypes.EDGE_AREA), + ) + counts = np.array([len(r) for r in remapped_lists]) + partners = np.concatenate(remapped_lists) + affs = np.repeat(active_affs, counts) + areas = np.repeat(active_areas, counts) + return partners, affs, areas + + +def _compute_partner_distances(new_kdtrees, partner_coords, partner_tree=None): + """Min distance from each new fragment to a partner's voxel coords. + + `partner_tree` may be supplied pre-built; otherwise built from `partner_coords`. + """ + if partner_tree is None: + partner_tree = cKDTree(partner_coords) + distances = np.empty(len(new_kdtrees), dtype=float) + for i, kt in enumerate(new_kdtrees): + if kt.n <= partner_tree.n: + d, _ = partner_tree.query(kt.data.reshape(-1, 3), k=1) + else: + d, _ = kt.query(partner_coords, k=1) + distances[i] = float(np.min(d)) + return distances + + +def _compute_boundary_distances(new_kdtrees, partner_chunk, old_chunk, chunk_size): + """Compute distance from each new fragment to a partner's chunk boundary. + + Used for active partners outside the bbox that have no kdtree entry. + `partner_chunk`, `old_chunk`, `chunk_size` are precomputed by the caller. + """ + diff = partner_chunk.astype(int) - old_chunk.astype(int) + axis = np.argmax(np.abs(diff)) + if diff[axis] > 0: + boundary = (old_chunk[axis] + 1) * chunk_size[axis] + else: + boundary = old_chunk[axis] * chunk_size[axis] + return np.array([np.min(np.abs(kt.data[:, axis] - boundary)) for kt in new_kdtrees]) + + +def _get_new_edges( + edges_info: tuple, + old_new_map: dict, + coords_by_label: dict, + root_id: basetypes.NODE_ID, + sv_root_map: dict, + cg: "ChunkedGraph", + new_kdtrees: list, + new_ids_arr: np.ndarray, + new_id_label_map: dict = None, + threshold: int = 10, +): + edge_batches, aff_batches, area_batches = [], [], [] + edges, affinities, areas = edges_info + + # `new_kdtrees[id_to_idx[nid]]` retrieves the precomputed tree for any + # new fragment id — every per-old subset draws from this one pool. + id_to_idx = {int(nid): i for i, nid in enumerate(new_ids_arr)} + # Partner trees are pure functions of the partner's voxel coords; + # cache across the whole call so a partner shared by multiple olds + # builds its tree once. + partner_tree_cache: dict = {} + + for old, new in old_new_map.items(): + new_ids = np.array(list(new), dtype=basetypes.NODE_ID) + edges_m = (edges[:, 0] == old) | (edges[:, 1] == old) + selected_edges = edges[edges_m] + sel_m = selected_edges != old + bad_rows = np.sum(sel_m, axis=1) != 1 + assert not bad_rows.any(), ( + f"each selected edge must touch old={old} exactly once; " + f"bad_rows={selected_edges[bad_rows].tolist()}" + ) + + partners = selected_edges[sel_m] + edge_affs = affinities[edges_m] + edge_areas = areas[edges_m] + partner_roots = np.array( + [sv_root_map.get(p, 0) for p in partners], dtype=np.uint64 + ) + active_m = partner_roots == root_id + + # Inactive partners (different root): broadcast to all fragments + inactive_idx = np.where(~active_m)[0] + if len(inactive_idx) > 0: + inactive_partners = partners[inactive_idx] + n_frag = len(new_ids) + broadcast_edges = np.column_stack( + [ + np.repeat(new_ids, len(inactive_partners)), + np.tile(inactive_partners, n_frag), + ] + ) + edge_batches.append(broadcast_edges) + aff_batches.append(np.tile(edge_affs[inactive_idx], n_frag)) + area_batches.append(np.tile(edge_areas[inactive_idx], n_frag)) + + # Active partners (same root): route based on affinity type + active_partners, act_affs, act_areas = _expand_partners( + partners[active_m], edge_affs[active_m], edge_areas[active_m], old_new_map + ) + if len(active_partners) > 0: + frag_kdtrees = [new_kdtrees[id_to_idx[int(nid)]] for nid in new_ids] + old_chunk = cg.get_chunk_coordinates(new_ids[0]) if cg else None + chunk_size = cg.meta.graph_config.CHUNK_SIZE if cg else None + # Pre-resolve chunk coords for partners that will hit the + # boundary-distance fallback (no in-bbox voxel coords). One + # batched bit-shift beats N scalar calls and removes `cg` from + # `_compute_boundary_distances`. + boundary_partners = [ + int(p) for p in active_partners if coords_by_label.get(int(p)) is None + ] + if boundary_partners and cg is not None: + boundary_arr = np.array(boundary_partners, dtype=np.uint64) + boundary_coords = chunk_utils.get_chunk_coordinates_multiple( + cg.meta, boundary_arr + ) + partner_chunk_map = { + p: boundary_coords[i] for i, p in enumerate(boundary_partners) + } + else: + partner_chunk_map = {} + for k, partner in enumerate(active_partners): + partner_int = int(partner) + partner_coords = coords_by_label.get(partner_int) + if partner_coords is not None: + pt = partner_tree_cache.get(partner_int) + if pt is None: + pt = cKDTree(partner_coords) + partner_tree_cache[partner_int] = pt + act_dist_row = _compute_partner_distances( + frag_kdtrees, partner_coords, partner_tree=pt + ) + else: + act_dist_row = _compute_boundary_distances( + frag_kdtrees, + partner_chunk_map[partner_int], + old_chunk, + chunk_size, + ) + e, a, ar = _match_partner( + new_ids, + partner, + act_affs[k], + act_areas[k], + act_dist_row, + new_id_label_map, + threshold, + ) + edge_batches.append(e) + aff_batches.append(a) + area_batches.append(ar) + + # Low-affinity edges between split fragments (cuttable by mincut) + if len(new_ids) > 1: + i_idx, j_idx = np.triu_indices(len(new_ids), k=1) + pairs = np.column_stack([new_ids[i_idx], new_ids[j_idx]]) + edge_batches.append(pairs) + n_pairs = len(pairs) + aff_batches.append(np.full(n_pairs, 0.001, dtype=basetypes.EDGE_AFFINITY)) + area_batches.append(np.zeros(n_pairs, dtype=basetypes.EDGE_AREA)) + + if len(edge_batches) == 0: + return ( + np.array([], dtype=basetypes.NODE_ID).reshape(0, 2), + np.array([], dtype=basetypes.EDGE_AFFINITY), + np.array([], dtype=basetypes.EDGE_AREA), + ) + all_edges = np.concatenate(edge_batches) + all_affs = np.concatenate(aff_batches) + all_areas = np.concatenate(area_batches) + edges_ = np.sort(all_edges.astype(basetypes.NODE_ID), axis=1) + edges_, idx = np.unique(edges_, return_index=True, axis=0) + return edges_, all_affs[idx], all_areas[idx] + + +def validate_split_edges(edges, affinities, old_new_map, new_id_label_map=None): + """Validate edge routing results before writing to prevent graph corruption. + + Checks: + A. No cross-label inf bridges — if an unsplit partner connects via inf edges + to fragments with different labels (different sides of the split), that + creates an uncuttable bridge through mincut. + B. No self-loops. + C. All old SVs have replacement edges from their fragments. + D. Inter-fragment edges exist between all fragment pairs. + + Raises PostconditionError on any violation. + """ + if len(edges) == 0: + return + + all_new_ids_arr = np.array( + [nid for ids in old_new_map.values() for nid in ids], dtype=np.uint64 + ) + + # B. No self-loops (cheapest check first) + self_loops = edges[:, 0] == edges[:, 1] + if self_loops.any(): + raise PostconditionError(f"Self-loop edges detected: {edges[self_loops]}") + + # A. No cross-label inf bridges to unsplit partners + if new_id_label_map: + inf_mask = np.isinf(affinities) + if inf_mask.any(): + inf_edges = edges[inf_mask] + is_frag_0 = np.isin(inf_edges[:, 0], all_new_ids_arr) + is_frag_1 = np.isin(inf_edges[:, 1], all_new_ids_arr) + mixed_mask = is_frag_0 ^ is_frag_1 + if mixed_mask.any(): + mixed = inf_edges[mixed_mask] + mixed_frag0 = is_frag_0[mixed_mask] + partners = np.where(mixed_frag0, mixed[:, 1], mixed[:, 0]) + fragments = np.where(mixed_frag0, mixed[:, 0], mixed[:, 1]) + unsplit_mask = ~np.isin(partners, all_new_ids_arr) + if unsplit_mask.any(): + unsplit_partners = partners[unsplit_mask] + unsplit_fragments = fragments[unsplit_mask] + for p in np.unique(unsplit_partners): + p_frags = unsplit_fragments[unsplit_partners == p] + labels = { + new_id_label_map[int(f)] + for f in p_frags + if int(f) in new_id_label_map + } + # Only {1, 2} forces a source↔sink uncuttable path + # through the partner. Bridges that include label-3 + # (unresolved fragment, no seed) ride to whichever + # seeded side the inf-cluster ends up on — a valid + # cut. The label is a routing hint, not a cut + # constraint; the mincut decides side membership. + if {1, 2}.issubset(labels): + raise PostconditionError( + f"Inf-affinity edge to unsplit partner {p} bridges " + f"source-side and sink-side fragments {labels}. " + f"This creates an uncuttable bridge in mincut." + ) + + # C. All old SVs have replacement edges + edge_svs = np.unique(edges.ravel()) + for old_id, new_ids in old_new_map.items(): + new_arr = np.array(list(new_ids), dtype=np.uint64) + if not np.any(np.isin(new_arr, edge_svs)): + raise PostconditionError( + f"Old SV {old_id} has no replacement edges from fragments {new_ids}" + ) + + # D. Inter-fragment edges exist + # edges are already sorted (col0 < col1), so check sorted pairs directly + edge_set = set(map(tuple, edges.tolist())) + for new_ids in old_new_map.values(): + ids = sorted(new_ids) + for i in range(len(ids)): + for j in range(i + 1, len(ids)): + if (ids[i], ids[j]) not in edge_set: + raise PostconditionError( + f"Missing inter-fragment edge between {ids[i]} and {ids[j]}" + ) + + +def _edges_to_bidirectional(edges_, affinities_, areas_): + """Duplicate edges in both directions and map nodes to chunks.""" + return ( + np.r_[edges_, edges_[:, ::-1]], + np.r_[affinities_, affinities_], + np.r_[areas_, areas_], + ) + + +def _compact_chunk_edges(prev_data, new_edges, new_affs, new_areas, stale_svs): + """Merge new edges with existing compacted edges, filtering stale SVs.""" + prev_cells = prev_data.get(attributes.Connectivity.CompactedSplitEdges, []) + if prev_cells: + prev_e = prev_cells[-1].value + prev_a = prev_data[attributes.Connectivity.CompactedAffinity][-1].value + prev_ar = prev_data[attributes.Connectivity.CompactedArea][-1].value + keep = ~np.isin(prev_e[:, 0], stale_svs) & ~np.isin(prev_e[:, 1], stale_svs) + new_edges = np.concatenate([prev_e[keep], new_edges]) + new_affs = np.concatenate([prev_a[keep], new_affs]) + new_areas = np.concatenate([prev_ar[keep], new_areas]) + return { + attributes.Connectivity.CompactedSplitEdges: new_edges, + attributes.Connectivity.CompactedAffinity: new_affs, + attributes.Connectivity.CompactedArea: new_areas, + } + + +def add_new_edges( + cg: "ChunkedGraph", + edges_tuple: tuple, + old_new_map: dict, + time_stamp: datetime = None, +): + edges_, affinities_, areas_ = edges_tuple + nodes = fastremap.unique(edges_) + chunks = cg.get_chunk_ids_from_node_ids(cg.get_parents(nodes)) + node_chunks = dict(zip(nodes, chunks)) + + edges, affinities, areas = _edges_to_bidirectional(edges_, affinities_, areas_) + stale_svs = np.array(list(old_new_map.keys()), dtype=basetypes.NODE_ID) + unique_chunks = np.unique(chunks) + + existing = cg.client.read_nodes( + node_ids=unique_chunks, + properties=[ + attributes.Connectivity.CompactedSplitEdges, + attributes.Connectivity.CompactedAffinity, + attributes.Connectivity.CompactedArea, + ], + fake_edges=True, + ) + + rows = [] + chunks_arr = fastremap.remap(edges, node_chunks) + for chunk_id in unique_chunks: + mask = chunks_arr[:, 0] == chunk_id + new_e, new_a, new_ar = edges[mask], affinities[mask], areas[mask] + row_key = serializers.serialize_uint64(chunk_id, fake_edges=True) + + # Append to SplitEdges (history, preserves all timestamps) + rows.append( + cg.client.mutate_row( + row_key, + { + attributes.Connectivity.SplitEdges: new_e, + attributes.Connectivity.Affinity: new_a, + attributes.Connectivity.Area: new_ar, + }, + time_stamp=time_stamp, + ) + ) + + # Write compacted edges (latest valid only) + compact_dict = _compact_chunk_edges( + existing.get(chunk_id, {}), new_e, new_a, new_ar, stale_svs + ) + rows.append(cg.client.mutate_row(row_key, compact_dict, time_stamp=time_stamp)) + return rows diff --git a/pychunkedgraph/graph/sv_split/edits.py b/pychunkedgraph/graph/sv_split/edits.py new file mode 100644 index 000000000..6110ab81a --- /dev/null +++ b/pychunkedgraph/graph/sv_split/edits.py @@ -0,0 +1,1004 @@ +""" +Manage new supervoxels after a supervoxel split. +""" + +import os +import time +from datetime import datetime +from collections import defaultdict +from functools import reduce +from typing import TYPE_CHECKING, List, Tuple + +import fastremap +import numpy as np +from pykdtree.kdtree import KDTree as cKDTree + +from pychunkedgraph import get_logger +from pychunkedgraph.profiler import get_profiler +from pychunkedgraph.graph import ( + attributes, + cache as cache_utils, + basetypes, + serializers, +) +from pychunkedgraph.graph.chunks.utils import chunks_overlapping_bbox +from pychunkedgraph.graph.edges import Edges +from pychunkedgraph.graph.exceptions import PostconditionError +from .splitter import get_splitter +from ._coords import build_coords_by_label +from .edges import _get_new_edges, add_new_edges, validate_split_edges +from .state import ( + ApplyResult, + SplitCtx, + SplitResult, + SvSplitOutcome, + SvSplitTask, +) +from pychunkedgraph.graph.utils import get_local_segmentation + +if TYPE_CHECKING: + from pychunkedgraph.graph.chunkedgraph import ChunkedGraph + +logger = get_logger(__name__) + + +def _coords_bbox( + cg: "ChunkedGraph", + src_coords_rep: np.ndarray, + sink_coords_rep: np.ndarray, +) -> tuple: + """Base-voxel bbox covering the user's source/sink seeds plus a margin. + + The cut surface lives between the user-placed source and sink + voxels; voxels of the rep that are far from those seeds never + contribute to the cut. So the read region is the seeds' envelope, + not the rep's full chunk envelope — for a physical SV cut into many + pieces across chunks, this can be orders of magnitude smaller. + + The margin is one CG chunk on each side. It matches the existing + L2 chunk lock margin and the 1-voxel shell read in + `split_supervoxel`, and gives `split_supervoxel_helper` headroom + around the seeds for the cut surface to travel along the SV. + + Pieces of the rep that fall outside the bbox keep their existing + IDs — they aren't read here and aren't rewritten. Cross-chunk-edge + routing for boundary-adjacent pieces is handled by the 1-voxel + shell at read time; cross-chunk edges entirely between unsplit + pieces don't change because their IDs don't change. + """ + coords = np.concatenate([src_coords_rep, sink_coords_rep], axis=0) + margin = np.array(cg.meta.graph_config.CHUNK_SIZE, dtype=int) + vol_start = cg.meta.voxel_bounds[:, 0] + vol_end = cg.meta.voxel_bounds[:, 1] + bbs = np.clip(coords.min(axis=0) - margin, vol_start, vol_end) + bbe = np.clip(coords.max(axis=0) + margin, vol_start, vol_end) + return bbs, bbe + + +def _l2_chunks_for_splits(cg: "ChunkedGraph", per_rep_bboxes: list) -> list[int]: + """Layer-2 chunk IDs every rep's split will read or write. + + Reads extend 1 voxel past `[bbs, bbe]` so `update_edges` has anchor + voxels for cross-chunk neighbors; the lock must cover those neighbor + chunks too, hence the `bbs - 1` / `bbe + 1` expansion. Clipped to + volume bounds so a bbox on the volume edge doesn't enumerate phantom + negative-index chunks. Sorted for deterministic lock-acquire order + (L2ChunkLock relies on sorted input for deadlock avoidance). + """ + vol_start = cg.meta.voxel_bounds[:, 0] + vol_end = cg.meta.voxel_bounds[:, 1] + chunk_size = cg.meta.graph_config.CHUNK_SIZE + chunk_coords = set() + for bbs, bbe in per_rep_bboxes: + read_lo = np.clip(bbs - 1, vol_start, vol_end) + read_hi = np.clip(bbe + 1, vol_start, vol_end) + chunk_coords.update( + chunks_overlapping_bbox( + read_lo, read_hi, chunk_size, origin=vol_start + ).keys() + ) + return sorted( + int(cg.get_chunk_id(layer=2, x=x, y=y, z=z)) for (x, y, z) in chunk_coords + ) + + +def _overlapping_reps( + *, + sv_remapping: dict, + source_ids: np.ndarray, + sink_ids: np.ndarray, + source_coords: np.ndarray, + sink_coords: np.ndarray, +): + """Yield per-rep data for every rep that links source and sink. + + A rep is a cross-chunk-representative SV shared by at least one + source and one sink in `sv_remapping`. These are the SVs that must + be split before the multicut can partition source from sink. + + Yields `(sv_id, src_coords_rep, sink_coords_rep, src_mask, sink_mask)`: + sv_id — one of the rep's source SV IDs, used as the + seed for `split_supervoxel`. + src_coords_rep — slice of source_coords whose SV maps to this rep. + sink_coords_rep — slice of sink_coords whose SV maps to this rep. + src_mask — positional boolean mask over source_ids; the + caller uses it to splice per-rep results back + into the full source arrays. + sink_mask — same, for sink_ids. + + Keyword-only signature — positional source/sink args of the same + shape are easy to swap without noticing. + """ + sources_remapped = fastremap.remap( + source_ids, sv_remapping, preserve_missing_labels=True, in_place=False + ) + sinks_remapped = fastremap.remap( + sink_ids, sv_remapping, preserve_missing_labels=True, in_place=False + ) + overlap_mask = np.isin(sources_remapped, sinks_remapped) + for rep in np.unique(sources_remapped[overlap_mask]): + src_mask = sources_remapped == rep + sink_mask = sinks_remapped == rep + yield ( + source_ids[src_mask][0], + source_coords[src_mask], + sink_coords[sink_mask], + src_mask, + sink_mask, + ) + + +def plan_sv_splits( + cg: "ChunkedGraph", + *, + sv_remapping: dict, + source_ids: np.ndarray, + sink_ids: np.ndarray, + source_coords: np.ndarray, + sink_coords: np.ndarray, +) -> Tuple[List[SvSplitTask], list]: + """Compute one `SvSplitTask` per rep and the L2 chunk set the splits + will touch. + + Pure function — no bigtable/OCDBT IO, no locks. Lets the caller + acquire the L2 chunk locks (both temporal and indefinite) around + `split_supervoxels` without recomputing the plan inside. + + Returns `(tasks, chunk_ids)` — `tasks` feeds `split_supervoxels`, + `chunk_ids` is the sorted union of read-expanded L2 chunks the full + operation touches. + """ + tasks: List[SvSplitTask] = [] + for ( + sv_id, + src_coords_rep, + sink_coords_rep, + src_mask, + sink_mask, + ) in _overlapping_reps( + sv_remapping=sv_remapping, + source_ids=source_ids, + sink_ids=sink_ids, + source_coords=source_coords, + sink_coords=sink_coords, + ): + bbs, bbe = _coords_bbox(cg, src_coords_rep, sink_coords_rep) + tasks.append( + SvSplitTask( + sv_id=sv_id, + src_coords=src_coords_rep, + sink_coords=sink_coords_rep, + src_mask=src_mask, + sink_mask=sink_mask, + bbs=bbs, + bbe=bbe, + ) + ) + chunk_ids = _l2_chunks_for_splits(cg, [(t.bbs, t.bbe) for t in tasks]) + return tasks, chunk_ids + + +def split_supervoxels( + cg: "ChunkedGraph", + *, + tasks: List[SvSplitTask], + sv_remapping: dict, + source_ids: np.ndarray, + sink_ids: np.ndarray, + operation_id: int, + timestamp: datetime = None, + parent_ts: datetime = None, +) -> SplitResult: + """Pure planner for the SV-split step. Returns a `SplitResult` with + all the data the caller needs to persist under locks. + + Does **not** write — the caller (`MulticutOperation._apply`) owns + the L2 chunk lock lifecycle and fires the OCDBT + bigtable writes + inside `IndefiniteL2ChunkLock`. + + Must be called inside the caller's `L2ChunkLock` for the + `plan.chunk_ids` set — the seg reads inside `split_supervoxel` need + to be consistent with concurrent writers. + + `timestamp` is the op's logical write time; threaded down to every + `mutate_row` in the persist block so all new-SV cells land at the + same logical time (atomic visibility for `parent_ts`-filtered + readers, and deterministic replay via `override_ts`). + + Fields on the returned `SplitResult`: + seg_bboxes: per-task base-resolution `(bbs, bbe)` — downsample + worker input. + source_ids_fresh / sink_ids_fresh: input `source_ids`/`sink_ids` + with positions touched by an overlap task replaced by the + new SV ID that now lives at that coord. Untouched positions + stay unchanged. Feeds the retry multicut. + seg_writes: flat list of `(voxel_slices, data)` pairs across all + tasks — one tensorstore write per pair, fired in parallel. + bigtable_rows: flattened rows from `copy_parents_and_add_lineage` + + `add_new_edges` across all tasks. + """ + source_ids_fresh = np.asarray(source_ids, dtype=basetypes.NODE_ID).copy() + sink_ids_fresh = np.asarray(sink_ids, dtype=basetypes.NODE_ID).copy() + + logger.note( + f"<{operation_id}> [sv_split:plan] {len(tasks)} supervoxels to split: " + f"{[int(t.sv_id) for t in tasks]}" + ) + + seg_bboxes = [] + seg_writes: List[Tuple[Tuple[slice, slice, slice], np.ndarray]] = [] + bigtable_rows: list = [] + for task in tasks: + out = split_supervoxel( + cg, + task, + operation_id, + sv_remapping=sv_remapping, + time_stamp=timestamp, + parent_ts=parent_ts, + ) + seg_bboxes.append(out.seg_bbox) + source_ids_fresh[task.src_mask] = out.src_new_ids + sink_ids_fresh[task.sink_mask] = out.sink_new_ids + seg_writes.extend(out.seg_write_pairs) + bigtable_rows.extend(out.bigtable_rows) + return SplitResult( + seg_bboxes=seg_bboxes, + source_ids_fresh=source_ids_fresh, + sink_ids_fresh=sink_ids_fresh, + seg_writes=seg_writes, + bigtable_rows=bigtable_rows, + ) + + +def _update_chunks(cg: "ChunkedGraph", chunks_bbox_map, seg, result_seg, bb_start): + """Process all chunks in a single pass: assign new SV IDs to split fragments. + + Returns `(results, change_chunks)`: + results: per-chunk (indices, old_values, new_values, label_id_map) + tuples; consumed by `_parse_results`. + change_chunks: `(chunk_coord, chunk_bbox)` for the chunks whose + voxels received new SV IDs. `write_seg_chunks` uses this to + rewrite only those chunks (skipping gap chunks that had no + split activity keeps the OCDBT delta proportional to actual + label changes). + """ + results = [] + change_chunks = [] + for chunk_coord, chunk_bbox in chunks_bbox_map.items(): + x, y, z = chunk_coord + chunk_id = cg.get_chunk_id(layer=1, x=x, y=y, z=z) + + _s, _e = chunk_bbox - bb_start + og_chunk_seg = seg[_s[0] : _e[0], _s[1] : _e[1], _s[2] : _e[2]] + chunk_seg = result_seg[_s[0] : _e[0], _s[1] : _e[1], _s[2] : _e[2]] + + labels = fastremap.unique(chunk_seg[chunk_seg != 0]) + if labels.size < 2: + continue + + new_ids = cg.id_client.create_node_ids(chunk_id, size=len(labels)) + _indices = [] + _old_values = [] + _new_values = [] + _label_id_map = {} + for _id, new_id in zip(labels, new_ids): + _mask = chunk_seg == _id + voxel_locs = np.where(_mask) + _og_value = og_chunk_seg[ + voxel_locs[0][0], voxel_locs[1][0], voxel_locs[2][0] + ] + _index = np.column_stack(voxel_locs) + n = len(_index) + _indices.append(_index) + _old_values.append(np.full(n, _og_value, dtype=basetypes.NODE_ID)) + _new_values.append(np.full(n, new_id, dtype=basetypes.NODE_ID)) + _label_id_map[int(_id)] = new_id + + _indices = np.concatenate(_indices) + (chunk_bbox[0] - bb_start) + _old_values = np.concatenate(_old_values) + _new_values = np.concatenate(_new_values) + results.append((_indices, _old_values, _new_values, _label_id_map)) + change_chunks.append((chunk_coord, chunk_bbox)) + return results, change_chunks + + +def _voxel_crop(bbs, bbe, bbs_, bbe_): + xS, yS, zS = bbs - bbs_ + xE, yE, zE = (None if i == 0 else -1 for i in bbe_ - bbe) + voxel_overlap_crop = np.s_[xS:xE, yS:yE, zS:zE] + return voxel_overlap_crop + + +def _assert_same_chunk(cg: "ChunkedGraph", old_new_map: dict) -> None: + """Every new SV must live in the same chunk as the SV it split from. + + PCG segment IDs are unique only within a chunk; a split fragment that + landed in a different chunk than its parent would break the hierarchy. + """ + olds = np.fromiter(old_new_map.keys(), dtype=basetypes.NODE_ID) + news = np.fromiter( + (n for ns in old_new_map.values() for n in ns), dtype=basetypes.NODE_ID + ) + expected = np.repeat( + cg.get_chunk_ids_from_node_ids(olds), [len(ns) for ns in old_new_map.values()] + ) + got = cg.get_chunk_ids_from_node_ids(news) + bad = np.flatnonzero(got != expected) + assert bad.size == 0, ( + "new SV landed in a different chunk than the SV it split from; " + f"(new_sv, got_chunk, expected_chunk): " + f"{[(int(news[i]), int(got[i]), int(expected[i])) for i in bad.tolist()]}" + ) + + +def _parse_results(results, seg, bbs, bbe): + """Merge per-chunk split results into a single segmentation volume. + + Applies new SV IDs from each chunk's split result to `seg` (in-place) + and builds the old→new mapping + label→new-id mapping. + + Returns (seg, old_new_map, new_id_label_map). + """ + old_new_map = defaultdict(set) + new_id_label_map = {} + for result in results: + if result: + indexer, old_values, new_values, label_id_map = result + seg[tuple(indexer.T)] = new_values + # old/new are per-voxel parallel arrays with only a handful + # of unique pairs per chunk; dedupe so the Python loop is + # over labels, not voxels. + unique_pairs = fastremap.unique( + np.column_stack([old_values, new_values]), axis=0 + ) + for old_sv, new_sv in unique_pairs: + old_new_map[old_sv].add(new_sv) + for label, new_id in label_id_map.items(): + new_id_label_map[new_id] = label + + assert np.all(seg.shape == bbe - bbs), f"{seg.shape} != {bbe - bbs}" + return seg, old_new_map, new_id_label_map + + +def _read_seg_and_ids(cg: "ChunkedGraph", bbs, bbe, *, sv_id=None, op_id=None): + """Read seg over [bbs-1, bbe+1] and return its distinct SV IDs. + + The 1-voxel shell gives update_edges anchor voxels from neighbouring + SVs. Returns (seg, sv_ids, bbs_, bbe_). + """ + vol_start = cg.meta.voxel_bounds[:, 0] + vol_end = cg.meta.voxel_bounds[:, 1] + bbs_ = np.clip(bbs - 1, vol_start, vol_end) + bbe_ = np.clip(bbe + 1, vol_start, vol_end) + _prof = get_profiler() + t0 = time.time() + with _prof.profile("seg_read"): + seg = get_local_segmentation(cg.meta, bbs_, bbe_).squeeze() + logger.note(f"<{op_id}> {sv_id}: read {seg.shape} ({time.time() - t0:.2f}s)") + + with _prof.profile("seg_unique"): + # Unique per chunk on the segment-id field only. Segment IDs are + # injective within a chunk and narrower than uint64, so the per- + # block sort is cheap; the chunk bits are OR'd back before the + # final union. The lattice is anchored at voxel_bounds[:, 0] so + # each block is exactly one chunk. Background 0 carries no chunk + # and is restored once at the end. + chunk_map = chunks_overlapping_bbox( + bbs_, bbe_, cg.meta.graph_config.CHUNK_SIZE, origin=vol_start + ) + parts = [] + has_zero = False + for (cx, cy, cz), cbbox in chunk_map.items(): + s, e = cbbox[0] - bbs_, cbbox[1] - bbs_ + sub = seg[s[0] : e[0], s[1] : e[1], s[2] : e[2]] + chunk_id = np.uint64(cg.get_chunk_id(layer=1, x=cx, y=cy, z=cz)) + limit = np.uint64(cg.get_segment_id_limit(chunk_id)) + narrow = np.min_scalar_type(int(limit)) + u = fastremap.unique((sub & limit).astype(narrow, copy=False)) + if u.size and u[0] == 0: + has_zero = True + u = u[1:] + parts.append(u.astype(np.uint64) | chunk_id) + sv_ids = ( + fastremap.unique(np.concatenate(parts)) + if parts + else np.array([], np.uint64) + ) + if has_zero: + sv_ids = np.concatenate([[np.uint64(0)], sv_ids]) + return seg, sv_ids, bbs_, bbe_ + + +def _select_cut_supervoxels(sv_id, sv_ids, rep_pieces, *, op_id=None): + """Narrow the rep to the pieces actually present in the bbox seg. + + Rep pieces whose voxels lie outside the seed-driven bbox don't appear + in seg and contribute nothing to the cut. Returns (cut_supervoxels, + supervoxel_ids). + """ + seg_ids = {int(x) for x in sv_ids if x != 0} + cut_supervoxels = rep_pieces & seg_ids + supervoxel_ids = np.array(list(cut_supervoxels), dtype=basetypes.NODE_ID) + logger.note( + f"<{op_id}> {sv_id}: whole_sv in_bbox={len(cut_supervoxels)} " + f"outside_bbox={len(rep_pieces) - len(cut_supervoxels)}" + ) + logger.verbose(f"<{op_id}> {sv_id}: pieces={supervoxel_ids.tolist()}") + return cut_supervoxels, supervoxel_ids + + +_SNAP_KWARGS = dict( + use_boundary=False, + downsample=False, + use_bbox=True, +) + + +def _log_split_result(op_id, sv_id, result): + """Emit the four OLD-format split log lines from a SplitResult.""" + d = result.diagnostics or {} + stage = d.get("stage_elapsed_s") or {} + pre = d.get("pre_resolve_label_counts") or d.get("label_counts") or {} + final = d.get("label_counts") or {} + logger.note( + f"<{op_id}> {sv_id}: connect_seeds ({stage.get('seed_prep', 0.0):.2f}s)" + ) + logger.note( + f"<{op_id}> {sv_id}: geodesic " + f"backend={d.get('backend')} ds={d.get('downsample_zyx')} " + f"{stage.get('arrival', 0.0):.3f}s" + ) + logger.note( + f"<{op_id}> {sv_id}: resolve3 " + f"label-1 {pre.get(1, 0)} label-2 {pre.get(2, 0)} label-3 stray {pre.get(3, 0)}" + ) + logger.note( + f"<{op_id}> {sv_id}: final " + f"label-1 {final.get(1, 0)} label-2 {final.get(2, 0)} label-3 unresolved {final.get(3, 0)}" + ) + + +def split_supervoxel_helper(ctx: SplitCtx, binary_seg: np.ndarray): + """Run the configured SV cut for one task; returns the label ndarray.""" + voxel_size = np.array(ctx.cg.meta.resolution) + downsample = voxel_size.max() // voxel_size # xyz order + # Per-axis clamp: + # - max_axis_ds bounds the per-axis stride so the cut surface + # precision stays within a small multiple of the finest voxel + # dimension (3× ≈ 24 nm on pinky). + # - min_grid_per_axis ensures ≥ N cells per axis post-downsample so + # narrow_band_rel has room to refine and small SVs fall back to + # full-res rather than collapsing the geodesic grid. + # binary_seg.shape is xyz (seg_overlap is xyz from ctx.seg); zip downsample + # and shape both in xyz, then reverse the final tuple to zyx for the + # geodesic call's axis convention. + max_axis_ds = 3 + min_grid_per_axis = 16 + ds_xyz = tuple( + max(1, min(int(s), max_axis_ds, dim // min_grid_per_axis)) + for s, dim in zip(downsample, binary_seg.shape) + ) + ds_zyx = ds_xyz[::-1] + src = ctx.source_coords - ctx.bbs + sink = ctx.sink_coords - ctx.bbs + backend_kwargs = {} + backend_env = os.environ.get("PYCG_GEODESIC_BACKEND") + if backend_env: + backend_kwargs["backend"] = backend_env + splitter = get_splitter( + downsample_geodesic=ds_zyx, + seed_prep_downsample=tuple(int(d) for d in downsample), + snap_kwargs=dict(_SNAP_KWARGS), + raise_if_multi_cc=True, + profiler=get_profiler(), + **backend_kwargs, + ) + with get_profiler().profile("split"): + result = splitter.split( + binary_seg, + src, + sink, + voxel_size=voxel_size, + vol_order="xyz", + vox_order="xyz", + seed_order="xyz", + ) + _log_split_result(ctx.operation_id, ctx.sv_id, result) + return result.labels + + +def _compute_split(ctx: SplitCtx, supervoxel_ids): + """Build the binary mask over the overlap crop and run the cut. + + Returns (split_result, voxel_overlap_crop). + """ + _prof = get_profiler() + with _prof.profile("binary_seg"): + # Chunked per-SV OR over the overlap crop. The plain loop would + # peak at 2× the bool output (binary_seg + one per-iter transient). + # Slabbing in z caps the transient at the byte budget below; np.isin + # is worse on memory here because the SV-id value range is too + # wide for `kind='table'` and `kind='sort'` allocates an int64 + # permutation buffer ≈ 8× the input. + voxel_overlap_crop = _voxel_crop(ctx.bbs, ctx.bbe, ctx.bbs_, ctx.bbe_) + seg_overlap = ctx.seg[voxel_overlap_crop] + binary_seg = np.empty(seg_overlap.shape, dtype=bool) + yx = int(seg_overlap.shape[1]) * int(seg_overlap.shape[2]) + slab_bytes = 128 * 1024 * 1024 + slab_z = max(1, slab_bytes // max(yx, 1)) + for z0 in range(0, seg_overlap.shape[0], slab_z): + z1 = min(z0 + slab_z, seg_overlap.shape[0]) + slab = seg_overlap[z0:z1] + binary_seg[z0:z1] = slab == supervoxel_ids[0] + for sv in supervoxel_ids[1:]: + binary_seg[z0:z1] |= slab == sv + t0 = time.time() + logger.note( + f"<{ctx.operation_id}> {ctx.sv_id}: split computation starting shape={binary_seg.shape}" + ) + split_result = split_supervoxel_helper(ctx, binary_seg) + logger.note( + f"<{ctx.operation_id}> {ctx.sv_id}: split computation done " + f"shape={split_result.shape} ({time.time() - t0:.2f}s)" + ) + return split_result, voxel_overlap_crop + + +def _pick_fresh_source_sink_ids( + old_new_map: dict, + new_id_label_map: dict, + n_source: int, + n_sink: int, + *, + sv_id, +) -> Tuple[np.ndarray, np.ndarray]: + """Pick a label-1 and a label-2 fragment from the same old SV (same chunk) + and broadcast to ``n_source`` / ``n_sink`` length arrays. + + Same-chunk fragments are directly connected by a 0.001 inter-fragment + bridge in ``add_new_edges``, giving the retry multicut a guaranteed + one-hop path between sources and sinks irrespective of what extends + beyond the local-subgraph bbox. Falls back to any label-1 / label-2 + only if no old SV produced both sides (degenerate cut). + """ + src_frag = sink_frag = None + for old_sv in sorted(old_new_map): + new_ids = sorted(old_new_map[old_sv]) + l1 = [n for n in new_ids if new_id_label_map.get(n) == 1] + l2 = [n for n in new_ids if new_id_label_map.get(n) == 2] + if l1 and l2: + src_frag, sink_frag = l1[0], l2[0] + break + if src_frag is None: + src_frag = next( + (n for n in sorted(new_id_label_map) if new_id_label_map[n] == 1), None + ) + if sink_frag is None: + sink_frag = next( + (n for n in sorted(new_id_label_map) if new_id_label_map[n] == 2), None + ) + if src_frag is None or sink_frag is None: + raise PostconditionError( + f"cut for sv {sv_id} produced no fragments on " + f"{'source' if src_frag is None else 'sink'} side" + ) + return ( + np.full(n_source, src_frag, dtype=basetypes.NODE_ID), + np.full(n_sink, sink_frag, dtype=basetypes.NODE_ID), + ) + + +def _apply_and_capture( + ctx: SplitCtx, voxel_overlap_crop, split_result, cut_supervoxels +): + """Apply fresh IDs to seg's crop and capture the write/lookup outputs. + + Writes fresh SV IDs into seg's overlap crop in place (a view, no + full-crop copy; _parse_results only writes, never reads crop values), + then captures the OCDBT write payloads and the src/sink id lookups + while the crop still holds unmasked neighbour IDs. Everything here + runs before the root mask, which would otherwise zero the neighbour + IDs the write must preserve. Returns an `ApplyResult`. + """ + cg, seg, bbs, bbe = ctx.cg, ctx.seg, ctx.bbs, ctx.bbe + _prof = get_profiler() + chunks_bbox_map = chunks_overlapping_bbox( + bbs, bbe, cg.meta.graph_config.CHUNK_SIZE, origin=cg.meta.voxel_bounds[:, 0] + ) + t0 = time.time() + results, change_chunks = _update_chunks( + cg, chunks_bbox_map, seg[voxel_overlap_crop], split_result, bbs + ) + logger.note( + f"<{ctx.operation_id}> {ctx.sv_id}: chunk updates {len(chunks_bbox_map)} chunks, " + f"{len(change_chunks)} with splits ({time.time() - t0:.2f}s)" + ) + + with _prof.profile("parse_results"): + new_seg = seg[voxel_overlap_crop] + new_seg, old_new_map, new_id_label_map = _parse_results( + results, new_seg, bbs, bbe + ) + del results + _assert_same_chunk(cg, old_new_map) + unsplit = cut_supervoxels - set(old_new_map.keys()) + logger.note( + f"<{ctx.operation_id}> {ctx.sv_id}: split_svs={len(old_new_map)} " + f"unsplit_kept={len(unsplit)}" + ) + if unsplit: + logger.verbose(f"<{ctx.operation_id}> {ctx.sv_id}: unsplit kept IDs: {unsplit}") + + # .copy() per changed chunk detaches each payload from seg before the + # mask / update_edges mutate it; changed chunks only, so the copies + # stay proportional to the edit. The caller batches them into one + # parallel tensorstore write. + seg_write_pairs: List[Tuple[Tuple[slice, slice, slice], np.ndarray]] = [] + for _, chunk_bbox in change_chunks: + lo, hi = chunk_bbox[0], chunk_bbox[1] + local_lo = lo - bbs + local_hi = hi - bbs + data = new_seg[ + local_lo[0] : local_hi[0], + local_lo[1] : local_hi[1], + local_lo[2] : local_hi[2], + ].copy() + voxel_slices = tuple(slice(int(s), int(e)) for s, e in zip(lo, hi)) + seg_write_pairs.append((voxel_slices, data)) + + src_new_ids, sink_new_ids = _pick_fresh_source_sink_ids( + old_new_map, + new_id_label_map, + len(ctx.source_coords), + len(ctx.sink_coords), + sv_id=ctx.sv_id, + ) + return ApplyResult( + old_new_map=old_new_map, + new_id_label_map=new_id_label_map, + seg_write_pairs=seg_write_pairs, + src_new_ids=src_new_ids, + sink_new_ids=sink_new_ids, + ) + + +def _fetch_subgraph_for_edges(cg, root, bbox, parent_ts, _prof): + """Fetch subgraph edges, dedup, then resolve roots for all endpoints. + Returns the trimmed edge tuple, the SV→root map, and per-stage metrics.""" + t0 = time.time() + with _prof.profile("subgraph"): + _, sg_edges_iter = cg.get_subgraph(root, bbox, bbox_is_coordinate=True) + edges_ = reduce(lambda x, y: x + y, sg_edges_iter, Edges([], [])) + n_subgraph = len(edges_.get_pairs()) + t_subgraph = time.time() - t0 + + edges = edges_.get_pairs() + affinities = edges_.affinities + areas = edges_.areas + + edges = np.sort(edges, axis=1) + _, edges_idx = np.unique(edges, axis=0, return_index=True) + edges_idx = edges_idx[edges[edges_idx, 0] != edges[edges_idx, 1]] + edges = edges[edges_idx] + affinities = affinities[edges_idx] + areas = areas[edges_idx] + + t0 = time.time() + with _prof.profile("roots"): + all_edge_svs = np.unique(edges) + all_roots = cg.get_roots(all_edge_svs, time_stamp=parent_ts) + sv_root_map = dict(zip(all_edge_svs, all_roots)) + n_roots = len(all_edge_svs) + t_roots = time.time() - t0 + + return ( + edges, + affinities, + areas, + all_edge_svs, + sv_root_map, + n_subgraph, + t_subgraph, + n_roots, + t_roots, + ) + + +def _compute_route_edges( + coords_by_label, + edges, + affinities, + areas, + sv_root_map, + old_new_map, + new_id_label_map, + cg, + root, + new_ids, + _prof, +): + """Build per-fragment KDTrees and route subgraph edges to new SV IDs. + Returns (edges_tuple, t_new) after validate_split_edges passes.""" + with _prof.profile("kdtrees"): + new_kdtrees = [cKDTree(coords_by_label[int(k)]) for k in new_ids] + + t0 = time.time() + with _prof.profile("get_new_edges"): + edges_tuple = _get_new_edges( + (edges, affinities, areas), + old_new_map, + coords_by_label, + root, + sv_root_map, + cg, + new_kdtrees, + new_ids, + new_id_label_map, + threshold=cg.meta.sv_split_threshold, + ) + t_new = time.time() - t0 + del new_kdtrees + with _prof.profile("validate_edges"): + validate_split_edges( + edges_tuple[0], edges_tuple[1], old_new_map, new_id_label_map + ) + return edges_tuple, t_new + + +def _route_edges_and_rows(ctx: SplitCtx, old_new_map, new_id_label_map): + """Resolve the split's root, route edges, build bigtable rows. + + Returns the flat list of bigtable rows. + """ + cg = ctx.cg + seg = ctx.seg + ctx.seg = None + _prof = get_profiler() + with _prof.profile("get_roots"): + roots = cg.get_roots(ctx.sv_ids, time_stamp=ctx.parent_ts) + root = roots[np.flatnonzero(ctx.sv_ids == ctx.sv_id)[0]] + + bbox = np.array([ctx.bbs, ctx.bbe]) + op_id = ctx.operation_id + sv_id = ctx.sv_id + parent_ts = ctx.parent_ts + old_new_map = dict(old_new_map) + new_ids = np.array(list(set.union(*old_new_map.values())), dtype=basetypes.NODE_ID) + + t_outer = time.time() + with _prof.profile("update_edges"): + ( + edges, + affinities, + areas, + all_edge_svs, + sv_root_map, + n_subgraph, + t_subgraph, + n_roots, + t_roots, + ) = _fetch_subgraph_for_edges(cg, root, bbox, parent_ts, _prof) + + # Inline: seg lives & dies in this frame. renumber's in-place rebind + # drops the only ref to the uint64 buffer; splitting this section + # across helpers would leave a pinning frame-local alive. + t0 = time.time() + with _prof.profile("build_coords"): + wanted_labels = np.union1d(new_ids, all_edge_svs) + with _prof.profile("mask_except"): + fastremap.mask_except(seg, list(wanted_labels), in_place=True) + if len(wanted_labels) <= np.iinfo(np.uint32).max: + with _prof.profile("renumber"): + seg, remap = fastremap.renumber(seg, in_place=True) + coords_small = build_coords_by_label(seg, boundary_only=True) + inv = {v: k for k, v in remap.items()} + coords_by_label = {int(inv[k]): v for k, v in coords_small.items()} + else: + coords_by_label = build_coords_by_label(seg, boundary_only=True) + del seg + n_labels = len(coords_by_label) + t_coords = time.time() - t0 + + edges_tuple, t_new = _compute_route_edges( + coords_by_label, + edges, + affinities, + areas, + sv_root_map, + old_new_map, + new_id_label_map, + cg, + root, + new_ids, + _prof, + ) + + logger.note( + f"<{op_id}> {sv_id} update_edges: subgraph={n_subgraph}/{t_subgraph:.2f}s " + f"roots={n_roots}/{t_roots:.2f}s coords={n_labels}/{t_coords:.2f}s " + f"_get_new_edges/{t_new:.2f}s" + ) + logger.note( + f"<{op_id}> {sv_id} -> {root} new_edges {edges_tuple[0].shape} " + f"({time.time() - t_outer:.2f}s)" + ) + + rows0 = copy_parents_and_add_lineage( + cg, ctx.operation_id, old_new_map, time_stamp=ctx.time_stamp + ) + rows1 = add_new_edges(cg, edges_tuple, old_new_map, time_stamp=ctx.time_stamp) + return rows0 + rows1 + + +def split_supervoxel( + cg: "ChunkedGraph", + task: SvSplitTask, + operation_id: int, + *, + sv_remapping: dict, + time_stamp: datetime = None, + parent_ts: datetime = None, +) -> SvSplitOutcome: + """Split one cross-chunk-connected SV into connected components. + + `task.bbs` / `task.bbe` are the base-voxel bbox covering the user's + source and sink seeds plus a one-chunk margin — `plan_sv_splits` + pre-computed this via `_coords_bbox`. The bbox is driven by where + the user wants the cut, not by the rep's full chunk envelope; rep + pieces outside the bbox aren't read and keep their existing IDs. + + `time_stamp` is the op's logical write time; threaded through to + `copy_parents_and_add_lineage` + `add_new_edges` so every new-SV + mutation lands at the same timestamp. + """ + sv_id = task.sv_id + bbs = task.bbs + bbe = task.bbe + + op_id = operation_id + t_start = time.time() + logger.note(f"<{op_id}> [sv_split:start] {sv_id} bbox=({bbs}, {bbe})") + + rep = sv_remapping.get(sv_id, sv_id) + rep_pieces = {int(sv) for sv, r in sv_remapping.items() if r == rep} + + seg, sv_ids, bbs_, bbe_ = _read_seg_and_ids(cg, bbs, bbe, sv_id=sv_id, op_id=op_id) + ctx = SplitCtx( + cg=cg, + seg=seg, + bbs=bbs, + bbe=bbe, + bbs_=bbs_, + bbe_=bbe_, + sv_id=sv_id, + sv_ids=sv_ids, + source_coords=task.src_coords, + sink_coords=task.sink_coords, + operation_id=operation_id, + time_stamp=time_stamp, + parent_ts=parent_ts, + ) + del seg + cut_supervoxels, supervoxel_ids = _select_cut_supervoxels( + sv_id, sv_ids, rep_pieces, op_id=op_id + ) + split_result, voxel_overlap_crop = _compute_split(ctx, supervoxel_ids) + applied = _apply_and_capture(ctx, voxel_overlap_crop, split_result, cut_supervoxels) + del split_result, voxel_overlap_crop + rows = _route_edges_and_rows(ctx, applied.old_new_map, applied.new_id_label_map) + + logger.note( + f"<{op_id}> [sv_split:end] {sv_id} elapsed={time.time() - t_start:.2f}s" + ) + return SvSplitOutcome( + seg_bbox=(bbs, bbe), + src_new_ids=applied.src_new_ids, + sink_new_ids=applied.sink_new_ids, + seg_write_pairs=applied.seg_write_pairs, + bigtable_rows=rows, + ) + + +def copy_parents_and_add_lineage( + cg: "ChunkedGraph", + operation_id: int, + old_new_map: dict, + *, + time_stamp: datetime = None, +) -> list: + """Copy parent pointers from old SVs onto their new-ID fragments + and write the lineage (FormerIdentity / NewIdentity) + L2 Child + list updates. + + `time_stamp` is the op's logical write time — used for every new-SV + cell this function writes so a `parent_ts`-filtered reader sees the + op atomically. The Parent-copy and Child-list writes deliberately + preserve the old cell's timestamp (so pre-op readers still see the + old hierarchy via the old timestamp). + + Returns a list of mutations to be persisted. + """ + result = [] + parents = set() + old_new_map = {k: list(v) for k, v in old_new_map.items()} + parent_cells_map = cg.client.read_nodes( + node_ids=list(old_new_map.keys()), properties=attributes.Hierarchy.Parent + ) + for old_id, new_ids in old_new_map.items(): + for new_id in new_ids: + val_dict = { + attributes.Hierarchy.FormerIdentity: np.array( + [old_id], dtype=basetypes.NODE_ID + ), + attributes.OperationLogs.OperationID: operation_id, + } + result.append( + cg.client.mutate_row( + serializers.serialize_uint64(new_id), + val_dict, + time_stamp=time_stamp, + ) + ) + for cell in parent_cells_map[old_id]: + cache_utils.update(cg.cache.parents_cache, [new_id], cell.value) + parents.add(cell.value) + result.append( + cg.client.mutate_row( + serializers.serialize_uint64(new_id), + {attributes.Hierarchy.Parent: cell.value}, + time_stamp=cell.timestamp, + ) + ) + val_dict = { + attributes.Hierarchy.NewIdentity: np.array(new_ids, dtype=basetypes.NODE_ID) + } + result.append( + cg.client.mutate_row( + serializers.serialize_uint64(old_id), + val_dict, + time_stamp=time_stamp, + ) + ) + + children_cells_map = cg.client.read_nodes( + node_ids=list(parents), properties=attributes.Hierarchy.Child + ) + for parent, children_cells in children_cells_map.items(): + assert len(children_cells) == 1, ( + f"expected 1 Child cell for parent={parent}; " + f"got {len(children_cells)}: {children_cells}" + ) + for cell in children_cells: + mask = np.isin(cell.value, list(old_new_map.keys())) + replace = np.concatenate([old_new_map[x] for x in cell.value[mask]]) + children = np.concatenate([cell.value[~mask], replace]) + cg.cache.children_cache[parent] = children + result.append( + cg.client.mutate_row( + serializers.serialize_uint64(parent), + {attributes.Hierarchy.Child: children}, + time_stamp=cell.timestamp, + ) + ) + return result diff --git a/pychunkedgraph/graph/sv_split/profile.py b/pychunkedgraph/graph/sv_split/profile.py new file mode 100644 index 000000000..8f8cfc592 --- /dev/null +++ b/pychunkedgraph/graph/sv_split/profile.py @@ -0,0 +1,376 @@ +"""Re-runnable dry-run profile harness for SV splits. + +Drives an SV-split operation end-to-end under ``PCG_DRY_RUN=1`` so no +BT or OCDBT state is mutated, captures per-stage timing + memory + IO +metrics into a ``HierarchicalProfiler`` (one ``BlockMetrics`` row per +stage), and snapshots each stage's intermediate result into a +``RunRecord`` dataclass that's persisted alongside the profiler. + +The persisted run lets the user iterate on a single heavy stage in +isolation (e.g. profile just ``split_supervoxels`` after editing it) +without re-running the prior stages. +""" + +import hashlib +import json +import pickle +import shutil +import sys +import tempfile +import traceback +from contextlib import contextmanager, nullcontext, redirect_stdout +from dataclasses import dataclass +from io import StringIO +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +from pychunkedgraph import NOTICE, VERBOSE, configure_logging +from pychunkedgraph.app.segmentation.common import _get_sources_and_sinks +from pychunkedgraph.profiler import HierarchicalProfiler, get_profiler +from . import edits +from pychunkedgraph.graph import utils as _utils_pkg +from pychunkedgraph.graph.dry_run import dry_run_scope +from pychunkedgraph.graph.operation import Cut, MulticutOperation, SvSplitRequired +from pychunkedgraph.graph.utils import generic as _utils_generic +from pychunkedgraph.graph.sv_lookup import utils as _sv_lookup_utils + +_CACHE_ROOT = Path(tempfile.gettempdir()) / "pcg_split_profile" + + +@dataclass +class RunRecord: + """Per-stage record/outputs captured during a ``run_split_profile`` run. + + Persisted to disk alongside the profiler so single-stage replays + can reuse the record without re-running prior stages. Every field + matches the exact value at the corresponding call site in + ``MulticutOperation._apply``. + """ + + operation_id: Optional[int] = None + timestamp: Any = None + source_ids_pre: Optional[np.ndarray] = None + sink_ids_pre: Optional[np.ndarray] = None + source_coords: Optional[np.ndarray] = None + sink_coords: Optional[np.ndarray] = None + sv_remapping: Optional[dict] = None + plan_tasks: Any = None + plan_chunk_ids: Any = None + sv_result: Any = None + cut: Any = None + result: Any = None + + +def _payload_canonical(payload: dict) -> str: + """Canonical JSON encoding used for hashing and collision detection.""" + return json.dumps(payload, sort_keys=True) + + +def _payload_sha(payload: dict) -> str: + """First 8 hex chars of sha256 over the canonical-JSON payload.""" + return hashlib.sha256(_payload_canonical(payload).encode()).hexdigest()[:8] + + +def run_dir(cg, payload: dict) -> Path: + """Cache directory for ``(cg.graph_id, payload)`` under system tmp.""" + return _CACHE_ROOT / cg.graph_id / _payload_sha(payload) + + +@contextmanager +def count_io(cg): + """Count BT row reads + OCDBT bytes read for the wrapped block. + + Wraps ``cg.client.read_nodes`` and ``cg.client.read_log_entries`` + (BT reads), plus ``get_local_segmentation`` at every top-level + binding site reached from the SV-split flow (OCDBT reads). All + originals are restored on exit. + """ + counters: Dict[str, int] = { + "bt_row_reads": 0, + "bt_log_reads": 0, + "ocdbt_reads": 0, + "ocdbt_bytes": 0, + } + + orig_read_nodes = cg.client.read_nodes + + def wrap_read_nodes(*a, **k): + result = orig_read_nodes(*a, **k) + counters["bt_row_reads"] += len(result) if result is not None else 0 + return result + + orig_read_log = cg.client.read_log_entries + + def wrap_read_log(*a, **k): + result = orig_read_log(*a, **k) + counters["bt_log_reads"] += len(result) if result is not None else 0 + return result + + cg.client.read_nodes = wrap_read_nodes + cg.client.read_log_entries = wrap_read_log + + # Patch every binding of get_local_segmentation reached from the + # SV-split flow. The source module is _utils_generic; the others + # imported it by name at module load time, so they hold separate + # references that need their own swap. + seg_modules = [_utils_generic, _utils_pkg, edits, _sv_lookup_utils] + orig_seg_fns = {m: m.get_local_segmentation for m in seg_modules} + + def wrap_get_local_seg(meta, bbox_start, bbox_end, mip=0): + # Always call the source function so we don't double-count if + # one wrapped binding calls another. + arr = orig_seg_fns[_utils_generic](meta, bbox_start, bbox_end, mip) + counters["ocdbt_bytes"] += int(arr.nbytes) + counters["ocdbt_reads"] += 1 + return arr + + for m in seg_modules: + m.get_local_segmentation = wrap_get_local_seg + + try: + yield counters + finally: + cg.client.read_nodes = orig_read_nodes + cg.client.read_log_entries = orig_read_log + for m, fn in orig_seg_fns.items(): + m.get_local_segmentation = fn + + +def profile_call(cg, name, fn, *args, **kwargs): + """Profile a single callable under dry-run with IO counters. + + Standalone replay helper for per-stage profiling (e.g. after + editing a single function's source). Opens ``dry_run_scope`` + + ``count_io``, runs ``profiler.profile(name, with_memory=True, + with_rss=True, counters=counters)`` around ``fn(*args, **kwargs)``, + returns ``(profiler, result)``. The profiler has exactly one block. + """ + profiler = HierarchicalProfiler(enabled=True) + with dry_run_scope(), count_io(cg) as counters: + with profiler.profile(name, counters=counters): + result = fn(*args, **kwargs) + return profiler, result + + +def build_op( + cg, + payload: dict, + *, + user_id: str = "dry_run_profile", + bbox_offset: Tuple[int, int, int] = (240, 240, 24), +) -> MulticutOperation: + """Decode a /split JSON payload and instantiate ``MulticutOperation``. + + Mirrors ``ChunkedGraph.remove_edges`` direct instantiation pattern. + The caller drives ``op.execute()``. + """ + source_ids, sink_ids, source_coords, sink_coords = _get_sources_and_sinks( + cg, payload + ) + op = MulticutOperation( + cg, + user_id=user_id, + source_ids=source_ids, + sink_ids=sink_ids, + source_coords=source_coords, + sink_coords=sink_coords, + bbox_offset=bbox_offset, + ) + return op + + +def annotate_chunks(cg, chunk_ids) -> List[str]: + """Annotate each chunk id with its NGL-navigable center voxel.""" + out: List[str] = [] + for cid in chunk_ids: + coord = cg.get_chunk_center_voxel(int(cid)).tolist() + out.append(f"{int(cid):#x} -> voxel {coord}") + return out + + +def _save_run( + cg, + payload: dict, + profiler: HierarchicalProfiler, + record: RunRecord, +) -> Path: + """Write run artifacts under ``run_dir``; raise on payload-hash collision.""" + target = run_dir(cg, payload) + payload_path = target / "payload.json" + incoming = _payload_canonical(payload) + if payload_path.exists(): + existing = payload_path.read_text() + if existing != incoming: + raise RuntimeError( + f"payload-hash collision at {target}: " + "existing payload != incoming payload" + ) + target.mkdir(parents=True, exist_ok=True) + with open(target / "record.pkl", "wb") as f: + pickle.dump(record, f) + with open(target / "profiler.pkl", "wb") as f: + pickle.dump(profiler, f) + buf = StringIO() + with redirect_stdout(buf): + profiler.metrics_report() + (target / "metrics.txt").write_text(buf.getvalue()) + payload_path.write_text(incoming) + return target + + +def load_run(cg, payload: dict) -> Tuple[HierarchicalProfiler, RunRecord]: + """Restore a prior ``run_split_profile`` result from disk.""" + target = run_dir(cg, payload) + with open(target / "profiler.pkl", "rb") as f: + profiler = pickle.load(f) + with open(target / "record.pkl", "rb") as f: + record = pickle.load(f) + return profiler, record + + +def _clean_traceback(tb_text: str) -> str: + """Strip caret-pointer lines (e.g. `` ^^^^``) and blank lines.""" + lines = [] + for line in tb_text.splitlines(): + stripped = line.strip() + if not stripped: + continue + if set(stripped) <= {"^"}: + continue + lines.append(line) + return "\n".join(lines) + + +def load_traceback(cg, payload: dict) -> Optional[str]: + """Return the saved traceback for a run (cleaned), or ``None`` if it succeeded. + + ``run_split_profile`` writes ``traceback.txt`` only when + ``op.execute()`` raised; its absence means the run completed. + """ + tb_path = run_dir(cg, payload) / "traceback.txt" + return _clean_traceback(tb_path.read_text()) if tb_path.exists() else None + + +def run_split_profile( + cg, + payload: dict, + *, + overwrite: bool = False, + dry_run: bool = True, + verbose: bool = True, +) -> Tuple[HierarchicalProfiler, RunRecord]: + """Drive an SV split with per-stage metrics captured. + + Returns ``(profiler, record)``. Uses the global profiler so inline + ``get_profiler().profile()`` blocks inside the SV-split call path + are captured automatically. ``record`` holds each stage's + intermediate values for standalone replay. + + ``dry_run=True`` (default) wraps the op in ``dry_run_scope`` + + ``edge_writeback_overlay`` so persistence is intercepted while the + retry multicut still sees post-split topology in memory. + ``dry_run=False`` lets every write land normally and skips the + in-memory overlay. + + ``overwrite=True`` wipes any existing cached run for this payload + before starting. + + Always writes a cache (profiler + record + metrics.txt) to + ``run_dir(cg, payload)`` on completion — even when ``op.execute()`` + raises — and prints the cache path. + """ + configure_logging(level=VERBOSE if verbose else NOTICE) + + target_dir = run_dir(cg, payload) + if target_dir.exists(): + if overwrite: + shutil.rmtree(target_dir) + else: + raise FileExistsError( + f"cached run already exists at {target_dir}; " + "pass overwrite=True to wipe and re-run, or " + "load_run(cg, payload) to read it" + ) + + profiler = get_profiler() + profiler.reset() + profiler.enabled = True + record = RunRecord() + + op = build_op(cg, payload) + record.source_ids_pre = op.source_ids.copy() + record.sink_ids_pre = op.sink_ids.copy() + record.source_coords = op.source_coords + record.sink_coords = op.sink_coords + + dry_ctx = dry_run_scope() if dry_run else nullcontext() + with dry_ctx, count_io(cg) as counters: + profiler.default_counters = counters + + # Capture-only wrappers for RunRecord replay — no profile() + # blocks. The real per-step metrics come from inline profile() + # blocks inside the called functions. + orig_run_multicut = MulticutOperation._run_multicut + orig_plan_sv_splits = edits.plan_sv_splits + orig_split_supervoxels = edits.split_supervoxels + + mincut_call_count = [0] + + def wrap_run_multicut(self_op, operation_id): + result = orig_run_multicut(self_op, operation_id) + mincut_call_count[0] += 1 + if mincut_call_count[0] == 1 and isinstance(result, SvSplitRequired): + record.sv_remapping = result.sv_remapping + elif isinstance(result, Cut): + record.cut = result + return result + + def wrap_plan_sv_splits(*a, **k): + result = orig_plan_sv_splits(*a, **k) + record.plan_tasks, record.plan_chunk_ids = result + return result + + def wrap_split_supervoxels(*a, **k): + if "operation_id" in k: + record.operation_id = k["operation_id"] + if "timestamp" in k: + record.timestamp = k["timestamp"] + result = orig_split_supervoxels(*a, **k) + record.sv_result = result + return result + + MulticutOperation._run_multicut = wrap_run_multicut + edits.plan_sv_splits = wrap_plan_sv_splits + edits.split_supervoxels = wrap_split_supervoxels + + tb_text = None + try: + record.result = op.execute() + except Exception: + tb_text = traceback.format_exc() + finally: + MulticutOperation._run_multicut = orig_run_multicut + edits.plan_sv_splits = orig_plan_sv_splits + edits.split_supervoxels = orig_split_supervoxels + profiler.default_counters = None + + try: + target = _save_run(cg, payload, profiler, record) + if tb_text is not None: + (target / "traceback.txt").write_text(tb_text) + artifact = target / ("traceback.txt" if tb_text else "metrics.txt") + print(f"[split_profile] run cached at {artifact}") + except Exception as save_err: + print(f"[split_profile] cache save failed: {save_err}", file=sys.stderr) + + if tb_text: + print("result:", _clean_traceback(tb_text).splitlines()[-1]) + else: + print("result:", record.result) + + # Disable so the global profiler is a no-op for callers outside + # this harness (production code paths included). + profiler.enabled = False + return profiler, record diff --git a/pychunkedgraph/graph/sv_split/splitter.py b/pychunkedgraph/graph/sv_split/splitter.py new file mode 100644 index 000000000..21ee698c7 --- /dev/null +++ b/pychunkedgraph/graph/sv_split/splitter.py @@ -0,0 +1,23 @@ +"""Resolve the configured SV splitter implementation. + +`PCG_SV_SPLITTER` env var holds the dotted import path of the Splitter +class. To use a different splitter, install its package and set the +env var — no PCG code change required. +""" + +import importlib +import os + +DEFAULT_SPLITTER = "supervoxel_splitter.GeodesicSplitter" + + +def get_splitter(**kwargs): + """Import the configured Splitter class and instantiate. + + `**kwargs` forward to the class constructor so caller-side tuning + propagates without dispatch. + """ + path = os.environ.get("PCG_SV_SPLITTER", DEFAULT_SPLITTER) + module_path, _, class_name = path.rpartition(".") + cls = getattr(importlib.import_module(module_path), class_name) + return cls(**kwargs) diff --git a/pychunkedgraph/graph/sv_split/state.py b/pychunkedgraph/graph/sv_split/state.py new file mode 100644 index 000000000..17a87cd5d --- /dev/null +++ b/pychunkedgraph/graph/sv_split/state.py @@ -0,0 +1,117 @@ +"""Per-stage data containers shared across the SV-split modules. + +Pulled into a leaf module so that ``edits``, ``profile``, ``inspect``, +``bridge_check``, and any future helper can all type-reference the same +dataclasses without import cycles. +""" + +from dataclasses import dataclass +from datetime import datetime +from typing import TYPE_CHECKING, List, Optional, Tuple + +import numpy as np + +if TYPE_CHECKING: + from pychunkedgraph.graph.chunkedgraph import ChunkedGraph + + +@dataclass +class SvSplitTask: + """One SV-split task per cross-chunk rep. + + Produced by ``plan_sv_splits`` (pure, no IO), consumed by + ``split_supervoxel``. ``src_mask`` / ``sink_mask`` are positional masks + back into the caller's ``source_ids`` / ``sink_ids`` arrays so the + aggregator can splice the per-task fresh IDs in at the right + positions. + """ + + sv_id: int + src_coords: np.ndarray + sink_coords: np.ndarray + src_mask: np.ndarray + sink_mask: np.ndarray + bbs: np.ndarray + bbe: np.ndarray + + +@dataclass +class SplitCtx: + """Per-task context shared across the split stage helpers. + + Holds the inputs every stage threads through unchanged. ``seg`` is + mutated in place across stages (fresh IDs written, then root mask); + the reference is stable, so storing it here is sound. + """ + + cg: "ChunkedGraph" + seg: np.ndarray + bbs: np.ndarray + bbe: np.ndarray + bbs_: np.ndarray + bbe_: np.ndarray + sv_id: int + sv_ids: np.ndarray + source_coords: np.ndarray + sink_coords: np.ndarray + operation_id: int + time_stamp: datetime + parent_ts: datetime + + +@dataclass +class ApplyResult: + """Outputs of ``_apply_and_capture`` consumed by the orchestrator. + + ``full_shape`` / ``fg_shape`` are stamped on by ``split_supervoxel`` + after ``_compute_split`` so the run record exposes the bbox the + geodesic actually ran on (the foreground crop) vs the bbox derived + from seed coords. Diff between them measures how much the rep is + concentrated near the seeds. + """ + + old_new_map: dict + new_id_label_map: dict + seg_write_pairs: List[Tuple[Tuple[slice, slice, slice], np.ndarray]] + src_new_ids: np.ndarray + sink_new_ids: np.ndarray + new_edges_tuple: Optional[tuple] = None + full_shape: Optional[Tuple[int, int, int]] = None + fg_shape: Optional[Tuple[int, int, int]] = None + + +@dataclass +class SvSplitOutcome: + """Output of ``split_supervoxel`` for one task. Aggregated into + ``SplitResult`` by ``split_supervoxels``.""" + + seg_bbox: Tuple[np.ndarray, np.ndarray] + src_new_ids: np.ndarray + sink_new_ids: np.ndarray + # Per-chunk OCDBT write payloads for this task. + seg_write_pairs: List[Tuple[Tuple[slice, slice, slice], np.ndarray]] + bigtable_rows: list + applied: Optional[ApplyResult] = None + + +@dataclass +class SplitResult: + """Pure planner output of ``split_supervoxels``. + + The caller (``MulticutOperation._apply``) performs the actual writes + under the L2 chunk locks: + - ``seg_writes`` is fed to ``write_seg_chunks`` as one flat parallel batch. + - ``bigtable_rows`` is written via ``cg.client.write`` in one batch. + """ + + seg_bboxes: List[Tuple[np.ndarray, np.ndarray]] + source_ids_fresh: np.ndarray + sink_ids_fresh: np.ndarray + # Flat list across all tasks: (voxel_slices, data_block) per OCDBT + # chunk write. ``voxel_slices`` is a 3-tuple of ``slice`` objects; the + # caller appends the channel slice and writes to ``meta.ws_ocdbt``. + seg_writes: List[Tuple[Tuple[slice, slice, slice], np.ndarray]] + bigtable_rows: list + # Per-task outcomes — carries each task's ApplyResult for post-split + # graph inspection (defence-in-depth bridge_check). + task_outcomes: Optional[List[SvSplitOutcome]] = None diff --git a/pychunkedgraph/graph/types.py b/pychunkedgraph/graph/types.py index 9a551f35c..fb7789cf1 100644 --- a/pychunkedgraph/graph/types.py +++ b/pychunkedgraph/graph/types.py @@ -1,14 +1,14 @@ -from typing import Dict -from typing import Iterable +# pylint: disable=invalid-name, missing-docstring from collections import namedtuple import numpy as np -from .utils import basetypes +from pychunkedgraph.graph import basetypes empty_1d = np.empty(0, dtype=basetypes.NODE_ID) empty_2d = np.empty((0, 2), dtype=basetypes.NODE_ID) - +empty_affinities = np.empty(0, dtype=basetypes.EDGE_AFFINITY) +empty_areas = np.empty(0, dtype=basetypes.EDGE_AREA) """ An Agglomeration is syntactic sugar for representing diff --git a/pychunkedgraph/graph/utils/__init__.py b/pychunkedgraph/graph/utils/__init__.py index e69de29bb..c1d56e0fe 100644 --- a/pychunkedgraph/graph/utils/__init__.py +++ b/pychunkedgraph/graph/utils/__init__.py @@ -0,0 +1 @@ +from .generic import get_local_segmentation \ No newline at end of file diff --git a/pychunkedgraph/graph/utils/_graph_tool.py b/pychunkedgraph/graph/utils/_graph_tool.py new file mode 100644 index 000000000..7b12ce462 --- /dev/null +++ b/pychunkedgraph/graph/utils/_graph_tool.py @@ -0,0 +1,8 @@ +"""Lazy graph_tool surface: import this module (not graph_tool) so the heavy +import and its scipy pull are deferred to first use, not package-import time. +""" + +import graph_tool +import graph_tool.flow as flow +import graph_tool.topology as topology +from graph_tool import Graph, GraphView diff --git a/pychunkedgraph/graph/utils/basetypes.py b/pychunkedgraph/graph/utils/basetypes.py deleted file mode 100644 index e55324e6a..000000000 --- a/pychunkedgraph/graph/utils/basetypes.py +++ /dev/null @@ -1,16 +0,0 @@ -import numpy as np - - -CHUNK_ID = SEGMENT_ID = NODE_ID = OPERATION_ID = np.dtype('uint64').newbyteorder('L') -EDGE_AFFINITY = np.dtype('float32').newbyteorder('L') -EDGE_AREA = np.dtype('uint64').newbyteorder('L') - -COUNTER = np.dtype('int64').newbyteorder('B') - -COORDINATES = np.dtype('int64').newbyteorder('L') -CHUNKSIZE = np.dtype('uint64').newbyteorder('L') -FANOUT = np.dtype('uint64').newbyteorder('L') -LAYERCOUNT = np.dtype('uint64').newbyteorder('L') -SPATIALBITS = np.dtype('uint64').newbyteorder('L') -ROOTCOUNTERBITS = np.dtype('uint64').newbyteorder('L') -SKIPCONNECTIONS = np.dtype('uint64').newbyteorder('L') \ No newline at end of file diff --git a/pychunkedgraph/graph/utils/flatgraph.py b/pychunkedgraph/graph/utils/flatgraph.py index df469d728..9bd700c4e 100644 --- a/pychunkedgraph/graph/utils/flatgraph.py +++ b/pychunkedgraph/graph/utils/flatgraph.py @@ -1,8 +1,9 @@ +# pylint: disable=invalid-name, missing-docstring, c-extension-no-member + +from itertools import combinations, chain + import fastremap import numpy as np -from itertools import combinations, chain -from graph_tool import Graph, GraphView -from graph_tool import topology, search def build_gt_graph( @@ -16,6 +17,8 @@ def build_gt_graph( :param hashed: bool :return: graph, capacities """ + from ._graph_tool import Graph + edges = np.array(edges, np.uint64) if weights is not None: assert len(weights) == len(edges) @@ -52,6 +55,8 @@ def connected_components(graph): :param graph: graph_tool.Graph :return: np.array of len == number of nodes """ + from ._graph_tool import Graph, topology + assert isinstance(graph, Graph) cc_labels = topology.label_components(graph)[0].a @@ -66,6 +71,8 @@ def connected_components(graph): def team_paths_all_to_all(graph, capacity, team_vertex_ids): + from ._graph_tool import topology + dprop = capacity.copy() # Use inverse affinity as the distance between vertices. dprop.a = 1 / (dprop.a + np.finfo(np.float64).eps) @@ -88,7 +95,10 @@ def team_paths_all_to_all(graph, capacity, team_vertex_ids): def neighboring_edges(graph, vertex_id): - """Returns vertex and edge lists of a seed vertex, in the same format as team_paths_all_to_all.""" + """ + Returns vertex and edge lists of a seed vertex, + in the same format as team_paths_all_to_all. + """ add_v = [] add_e = [] v0 = graph.vertex(vertex_id) @@ -106,7 +116,7 @@ def intersect_nodes(paths_v_s, paths_v_y): def harmonic_mean_paths(x): - return np.power(np.product(x), 1 / len(x)) + return np.power(np.prod(x), 1 / len(x)) def compute_filtered_paths( @@ -116,6 +126,8 @@ def compute_filtered_paths( intersect_vertices, ): """Make a filtered GraphView that excludes intersect vertices and recompute shortest paths""" + from ._graph_tool import GraphView + intersection_filter = np.full(graph.num_vertices(), True) intersection_filter[intersect_vertices] = False vfilt = graph.new_vertex_property("bool", vals=intersection_filter) @@ -124,7 +136,8 @@ def compute_filtered_paths( gfilt, capacity, team_vertex_ids ) - # graph-tool will invalidate the vertex and edge properties if I don't rebase them on the main graph + # graph-tool will invalidate the vertex and + # edge properties if I don't rebase them on the main graph # before tearing down the GraphView new_paths_e = [] for pth in paths_e: @@ -169,6 +182,8 @@ def remove_overlapping_edges(paths_v_s, paths_e_s, paths_v_y, paths_e_y): def check_connectedness(vertices, edges, expected_number=1): """Returns True if the augmenting edges still form a single connected component""" + from ._graph_tool import Graph, topology + paths_inds = np.unique([int(v) for v in chain.from_iterable(vertices)]) edge_list_inds = np.array( [[int(e.source()), int(e.target())] for e in chain.from_iterable(edges)] diff --git a/pychunkedgraph/graph/utils/generic.py b/pychunkedgraph/graph/utils/generic.py index 9a2b6f979..e5c9e2883 100644 --- a/pychunkedgraph/graph/utils/generic.py +++ b/pychunkedgraph/graph/utils/generic.py @@ -3,7 +3,9 @@ TODO categorize properly """ +from __future__ import annotations +import bisect import datetime from typing import Dict from typing import Iterable @@ -14,10 +16,33 @@ from collections import defaultdict import numpy as np -import pandas as pd -import pytz from ..chunks import utils as chunk_utils +from ..exceptions import PreconditionError + + +def assert_same_root( + sv_ids: np.ndarray, roots: np.ndarray, *, source: str +) -> np.ndarray: + """Raise PreconditionError if `roots` spans more than one root. + + Required root = the most common root. Offenders are the supervoxels + whose root differs from it. `source` tags the call site so the same + error wording fired from different places is greppable. + """ + root_ids, root_counts = np.unique(roots, return_counts=True) + if len(root_ids) > 1: + required_root = int(root_ids[np.argmax(root_counts)]) + offenders_by_root: dict = {} + for sv, r in zip(sv_ids.tolist(), roots.tolist()): + if int(r) != required_root: + offenders_by_root.setdefault(int(r), []).append(int(sv)) + raise PreconditionError( + f"[{source}] Supervoxels must belong to the same object " + f"(required root {required_root}). " + f"offenders by root: {offenders_by_root}" + ) + return root_ids def compute_indices_pandas(data) -> pd.Series: @@ -27,6 +52,8 @@ def compute_indices_pandas(data) -> pd.Series: :param data: np.ndarray :return: pandas dataframe """ + import pandas as pd + d = data.ravel() f = lambda x: np.unravel_index(x.index, data.shape) return pd.Series(d).groupby(d).apply(f) @@ -75,36 +102,6 @@ def compute_bitmasks(n_layers: int, s_bits_atomic_layer: int = 8) -> Dict[int, i return bitmask_dict -def get_max_time(): - """Returns the (almost) max time in datetime.datetime - :return: datetime.datetime - """ - return datetime.datetime(9999, 12, 31, 23, 59, 59, 0) - - -def get_min_time(): - """Returns the min time in datetime.datetime - :return: datetime.datetime - """ - return datetime.datetime.strptime("01/01/00 00:00", "%d/%m/%y %H:%M") - - -def time_min(): - """Returns a minimal time stamp that still works with google - :return: datetime.datetime - """ - return datetime.datetime.strptime("01/01/00 00:00", "%d/%m/%y %H:%M") - - -def get_valid_timestamp(timestamp): - if timestamp is None: - timestamp = datetime.datetime.utcnow() - if timestamp.tzinfo is None: - timestamp = pytz.UTC.localize(timestamp) - # Comply to resolution of BigTables TimeRange - return _get_google_compatible_time_stamp(timestamp, round_up=False) - - def get_bounding_box( source_coords: Sequence[Sequence[int]], sink_coords: Sequence[Sequence[int]], @@ -137,27 +134,6 @@ def filter_failed_node_ids(row_ids, segment_ids, max_children_ids): return row_ids[max_child_ids_occ_so_far == 0] -def _get_google_compatible_time_stamp( - time_stamp: datetime.datetime, round_up: bool = False -) -> datetime.datetime: - """Makes a datetime.datetime time stamp compatible with googles' services. - Google restricts the accuracy of time stamps to milliseconds. Hence, the - microseconds are cut of. By default, time stamps are rounded to the lower - number. - :param time_stamp: datetime.datetime - :param round_up: bool - :return: datetime.datetime - """ - micro_s_gap = datetime.timedelta(microseconds=time_stamp.microsecond % 1000) - if micro_s_gap == 0: - return time_stamp - if round_up: - time_stamp += datetime.timedelta(microseconds=1000) - micro_s_gap - else: - time_stamp -= micro_s_gap - return time_stamp - - def mask_nodes_by_bounding_box( meta, nodes: Union[Iterable[np.uint64], np.uint64], @@ -173,9 +149,7 @@ def mask_nodes_by_bounding_box( adapt_layers = layers - 2 adapt_layers[adapt_layers < 0] = 0 fanout = meta.graph_config.FANOUT - bounding_box_layer = ( - bounding_box[None] / (fanout ** adapt_layers)[:, None, None] - ) + bounding_box_layer = bounding_box[None] / (fanout**adapt_layers)[:, None, None] bound_check = np.array( [ np.all(chunk_coordinates < bounding_box_layer[:, 1], axis=1), @@ -183,4 +157,43 @@ def mask_nodes_by_bounding_box( ] ).T - return np.all(bound_check, axis=1) \ No newline at end of file + return np.all(bound_check, axis=1) + + +def get_parents_at_timestamp(nodes, parents_ts_map, time_stamp, unique: bool = False): + """ + Search for the first parent with ts <= `time_stamp`. + `parents_ts_map[node]` is a map of ts:parent with sorted timestamps (desc). + """ + skipped_nodes = [] + parents = set() if unique else [] + for node in nodes: + try: + ts_parent_map = parents_ts_map[node] + ts_list = list(ts_parent_map.keys()) + asc_ts_list = ts_list[::-1] + idx = bisect.bisect_right(asc_ts_list, time_stamp) + ts = asc_ts_list[idx - 1] + parent = ts_parent_map[ts] + parents.add(parent) if unique else parents.append(parent) + except KeyError: + skipped_nodes.append(node) + return list(parents), skipped_nodes + + +def get_local_segmentation(meta, bbox_start, bbox_end, mip: int = 0) -> np.ndarray: + """Read a segmentation region from OCDBT (or the watershed via tensorstore). + + `bbox_start` and `bbox_end` must already be in the requested MIP level's + coordinate space — this function does not rescale them. Meshing computes + chunk bounds at the target MIP and passes them through directly; SV split + and coordinate lookup always use base resolution (mip=0). + """ + xL, yL, zL = bbox_start + xH, yH, zH = bbox_end + if meta.ocdbt_seg: + # mip > 0 reads from a coarser scale; saves bandwidth and is what + # meshing wants when it operates at a non-base MIP. + store = meta.ws_ocdbt if mip == 0 else meta.ws_ocdbt_scales[mip] + return store[xL:xH, yL:yH, zL:zH].read().result() + return meta.ws_ts_scale(mip)[xL:xH, yL:yH, zL:zH].read().result() diff --git a/pychunkedgraph/graph/utils/id_helpers.py b/pychunkedgraph/graph/utils/id_helpers.py index aa486ac84..54792759a 100644 --- a/pychunkedgraph/graph/utils/id_helpers.py +++ b/pychunkedgraph/graph/utils/id_helpers.py @@ -2,14 +2,9 @@ Utils functions for node and segment IDs. """ -from typing import Optional -from typing import Sequence -from typing import Callable -from datetime import datetime - import numpy as np -from . import basetypes +from pychunkedgraph.graph import basetypes from ..meta import ChunkedGraphMeta from ..chunks import utils as chunk_utils @@ -20,7 +15,7 @@ def get_segment_id_limit( """Get maximum possible Segment ID for given Node ID or Chunk ID.""" layer = chunk_utils.get_chunk_layer(meta, node_or_chunk_id) chunk_offset = 64 - meta.graph_config.LAYER_ID_BITS - 3 * meta.bitmasks[layer] - return np.uint64(2 ** chunk_offset - 1) + return np.uint64(2**chunk_offset - 1) def get_segment_id( @@ -47,143 +42,3 @@ def get_node_id( return chunk_id | segment_id else: return chunk_utils.get_chunk_id(meta, layer=layer, x=x, y=y, z=z) | segment_id - - -def get_atomic_id_from_coord( - meta: ChunkedGraphMeta, - get_root: callable, - x: int, - y: int, - z: int, - parent_id: np.uint64, - n_tries: int = 5, - time_stamp: Optional[datetime] = None, -) -> np.uint64: - """Determines atomic id given a coordinate.""" - x = int(x / 2 ** meta.data_source.CV_MIP) - y = int(y / 2 ** meta.data_source.CV_MIP) - z = int(z) - - checked = [] - atomic_id = None - root_id = get_root(parent_id, time_stamp=time_stamp) - - for i_try in range(n_tries): - # Define block size -- increase by one each try - x_l = x - (i_try - 1) ** 2 - y_l = y - (i_try - 1) ** 2 - z_l = z - (i_try - 1) ** 2 - - x_h = x + 1 + (i_try - 1) ** 2 - y_h = y + 1 + (i_try - 1) ** 2 - z_h = z + 1 + (i_try - 1) ** 2 - - x_l = 0 if x_l < 0 else x_l - y_l = 0 if y_l < 0 else y_l - z_l = 0 if z_l < 0 else z_l - - # Get atomic ids from cloudvolume - atomic_id_block = meta.cv[x_l:x_h, y_l:y_h, z_l:z_h] - atomic_ids, atomic_id_count = np.unique(atomic_id_block, return_counts=True) - - # sort by frequency and discard those ids that have been checked - # previously - sorted_atomic_ids = atomic_ids[np.argsort(atomic_id_count)] - sorted_atomic_ids = sorted_atomic_ids[~np.in1d(sorted_atomic_ids, checked)] - - # For each candidate id check whether its root id corresponds to the - # given root id - for candidate_atomic_id in sorted_atomic_ids: - if candidate_atomic_id != 0: - ass_root_id = get_root(candidate_atomic_id, time_stamp=time_stamp) - if ass_root_id == root_id: - # atomic_id is not None will be our indicator that the - # search was successful - atomic_id = candidate_atomic_id - break - else: - checked.append(candidate_atomic_id) - if atomic_id is not None: - break - # Returns None if unsuccessful - return atomic_id - - -def get_atomic_ids_from_coords( - meta: ChunkedGraphMeta, - coordinates: Sequence[Sequence[int]], - parent_id: np.uint64, - parent_id_layer: int, - parent_ts: datetime, - get_roots: Callable, - max_dist_nm: int = 150, -) -> Sequence[np.uint64]: - """Retrieves supervoxel ids for multiple coords. - - :param coordinates: n x 3 np.ndarray of locations in voxel space - :param parent_id: parent id common to all coordinates at any layer - :param max_dist_nm: max distance explored - :return: supervoxel ids; returns None if no solution was found - """ - import fastremap - - if parent_id_layer == 1: - return np.array([parent_id] * len(coordinates), dtype=np.uint64) - - coordinates_nm = coordinates * np.array(meta.resolution) - # Define bounding box to be explored - max_dist_vx = np.ceil(max_dist_nm / meta.resolution).astype(dtype=np.int32) - bbox = np.array( - [ - np.min(coordinates, axis=0) - max_dist_vx, - np.max(coordinates, axis=0) + max_dist_vx + 1, - ] - ) - - local_sv_seg = meta.cv[ - bbox[0, 0] : bbox[1, 0], bbox[0, 1] : bbox[1, 1], bbox[0, 2] : bbox[1, 2] - ].squeeze() - - # limit get_roots calls to the relevant areas of the data - lower_bs = np.floor( - (np.array(coordinates_nm) - max_dist_nm) / np.array(meta.resolution) - bbox[0] - ).astype(np.int32) - upper_bs = np.ceil( - (np.array(coordinates_nm) + max_dist_nm) / np.array(meta.resolution) - bbox[0] - ).astype(np.int32) - local_sv_ids = [] - for lb, ub in zip(lower_bs, upper_bs): - local_sv_ids.extend( - fastremap.unique(local_sv_seg[lb[0] : ub[0], lb[1] : ub[1], lb[2] : ub[2]]) - ) - local_sv_ids = fastremap.unique(np.array(local_sv_ids, dtype=np.uint64)) - local_parent_ids = get_roots( - local_sv_ids, - time_stamp=parent_ts, - stop_layer=parent_id_layer, - fail_to_zero=True - ) - - local_parent_seg = fastremap.remap( - local_sv_seg, - dict(zip(local_sv_ids, local_parent_ids)), - preserve_missing_labels=True, - ) - - parent_id_locs_vx = np.array(np.where(local_parent_seg == parent_id)).T - if len(parent_id_locs_vx) == 0: - return None - - parent_id_locs_nm = (parent_id_locs_vx + bbox[0]) * np.array(meta.resolution) - # find closest supervoxel ids and check that they are closer than the limit - dist_mat = np.sqrt( - np.sum((parent_id_locs_nm[:, None] - coordinates_nm) ** 2, axis=-1) - ) - match_ids = np.argmin(dist_mat, axis=0) - matched_dists = np.array([dist_mat[idx, i] for i, idx in enumerate(match_ids)]) - if np.any(matched_dists > max_dist_nm): - return None - - local_coords = parent_id_locs_vx[match_ids] - matched_sv_ids = [local_sv_seg[tuple(c)] for c in local_coords] - return matched_sv_ids diff --git a/pychunkedgraph/graph/utils/serializers.py b/pychunkedgraph/graph/utils/serializers.py deleted file mode 100644 index 09c0f63b0..000000000 --- a/pychunkedgraph/graph/utils/serializers.py +++ /dev/null @@ -1,156 +0,0 @@ -from typing import Any, Iterable -import json -import pickle - -import numpy as np -import zstandard as zstd - - -class _Serializer: - def __init__(self, serializer, deserializer, basetype=Any, compression_level=None): - self._serializer = serializer - self._deserializer = deserializer - self._basetype = basetype - self._compression_level = compression_level - - def serialize(self, obj): - content = self._serializer(obj) - if self._compression_level: - return zstd.ZstdCompressor(level=self._compression_level).compress(content) - return content - - def deserialize(self, obj): - if self._compression_level: - obj = zstd.ZstdDecompressor().decompressobj().decompress(obj) - return self._deserializer(obj) - - @property - def basetype(self): - return self._basetype - - -class NumPyArray(_Serializer): - @staticmethod - def _deserialize(val, dtype, shape=None, order=None): - data = np.frombuffer(val, dtype=dtype) - if shape is not None: - return data.reshape(shape, order=order) - if order is not None: - return data.reshape(data.shape, order=order) - return data - - def __init__(self, dtype, shape=None, order=None, compression_level=None): - super().__init__( - serializer=lambda x: x.newbyteorder(dtype.byteorder).tobytes(), - deserializer=lambda x: NumPyArray._deserialize( - x, dtype, shape=shape, order=order - ), - basetype=dtype.type, - compression_level=compression_level, - ) - - -class NumPyValue(_Serializer): - def __init__(self, dtype): - super().__init__( - serializer=lambda x: x.newbyteorder(dtype.byteorder).tobytes(), - deserializer=lambda x: np.frombuffer(x, dtype=dtype)[0], - basetype=dtype.type, - ) - - -class String(_Serializer): - def __init__(self, encoding="utf-8"): - super().__init__( - serializer=lambda x: x.encode(encoding), - deserializer=lambda x: x.decode(), - basetype=str, - ) - - -class JSON(_Serializer): - def __init__(self): - super().__init__( - serializer=lambda x: json.dumps(x).encode("utf-8"), - deserializer=lambda x: json.loads(x.decode()), - basetype=str, - ) - - -class Pickle(_Serializer): - def __init__(self): - super().__init__( - serializer=lambda x: pickle.dumps(x), - deserializer=lambda x: pickle.loads(x), - basetype=str, - ) - - -class UInt64String(_Serializer): - def __init__(self): - super().__init__( - serializer=serialize_uint64, - deserializer=deserialize_uint64, - basetype=np.uint64, - ) - - -def pad_node_id(node_id: np.uint64) -> str: - """ Pad node id to 20 digits - - :param node_id: int - :return: str - """ - return "%.20d" % node_id - - -def serialize_uint64(node_id: np.uint64, counter=False, fake_edges=False) -> bytes: - """ Serializes an id to be ingested by a bigtable table row - - :param node_id: int - :return: str - """ - if counter: - return serialize_key("i%s" % pad_node_id(node_id)) # type: ignore - if fake_edges: - return serialize_key("f%s" % pad_node_id(node_id)) # type: ignore - return serialize_key(pad_node_id(node_id)) # type: ignore - - -def serialize_uint64s_to_regex(node_ids: Iterable[np.uint64]) -> bytes: - """ Serializes an id to be ingested by a bigtable table row - - :param node_id: int - :return: str - """ - node_id_str = "".join(["%s|" % pad_node_id(node_id) for node_id in node_ids])[:-1] - return serialize_key(node_id_str) # type: ignore - - -def deserialize_uint64(node_id: bytes, fake_edges=False) -> np.uint64: - """ De-serializes a node id from a BigTable row - - :param node_id: bytes - :return: np.uint64 - """ - if fake_edges: - return np.uint64(node_id[1:].decode()) # type: ignore - return np.uint64(node_id.decode()) # type: ignore - - -def serialize_key(key: str) -> bytes: - """ Serializes a key to be ingested by a bigtable table row - - :param key: str - :return: bytes - """ - return key.encode("utf-8") - - -def deserialize_key(key: bytes) -> str: - """ Deserializes a row key - - :param key: bytes - :return: str - """ - return key.decode() diff --git a/pychunkedgraph/ingest/__init__.py b/pychunkedgraph/ingest/__init__.py index b3d832d5e..482dfbb5f 100644 --- a/pychunkedgraph/ingest/__init__.py +++ b/pychunkedgraph/ingest/__init__.py @@ -1,32 +1,17 @@ from collections import namedtuple +from pychunkedgraph import configure_logging, NOTICE -_cluster_ingest_config_fields = ( - "ATOMIC_Q_NAME", - "ATOMIC_Q_LIMIT", - "ATOMIC_Q_INTERVAL", -) -_cluster_ingest_defaults = ( - "l2", - 100000, - 120, -) -ClusterIngestConfig = namedtuple( - "ClusterIngestConfig", - _cluster_ingest_config_fields, - defaults=_cluster_ingest_defaults, -) - +configure_logging(level=NOTICE) _ingestconfig_fields = ( - "CLUSTER", # cluster config "AGGLOMERATION", "WATERSHED", "USE_RAW_EDGES", "USE_RAW_COMPONENTS", "TEST_RUN", ) -_ingestconfig_defaults = (None, None, None, False, False, False) +_ingestconfig_defaults = (None, None, False, False, False) IngestConfig = namedtuple( "IngestConfig", _ingestconfig_fields, defaults=_ingestconfig_defaults ) diff --git a/pychunkedgraph/ingest/cli.py b/pychunkedgraph/ingest/cli.py index 7668e8f24..d28fbbce1 100644 --- a/pychunkedgraph/ingest/cli.py +++ b/pychunkedgraph/ingest/cli.py @@ -1,24 +1,39 @@ +# pylint: disable=invalid-name, missing-function-docstring, unspecified-encoding + """ cli for running ingest """ -from os import environ +import os +from functools import partial from time import sleep +from pychunkedgraph import configure_logging, DEBUG + import click import yaml from flask.cli import AppGroup -from rq import Queue +from .cluster import create_atomic_chunk, create_parent_chunk, enqueue_l2_tasks from .manager import IngestionManager -from .utils import bootstrap -from .cluster import randomize_grid_points +from .ocdbt import coordinator, setup_base +from .utils import ( + bootstrap, + job_type_guard, + print_completion_rate, + print_status, + purge_layer_state, + queue_layer_helper, + requeue_chunk, +) +from .simple_tests import run_all +from .create.parent_layer import add_parent_chunk from ..graph.chunkedgraph import ChunkedGraph -from ..utils.redis import get_redis_connection -from ..utils.redis import keys as r_keys -from ..utils.general import chunked +from ..graph.ocdbt import OcdbtConfig +from ..utils.redis import get_redis_connection, keys as r_keys -ingest_cli = AppGroup("ingest") +group_name = "ingest" +ingest_cli = AppGroup(group_name) def init_ingest_cmds(app): @@ -26,6 +41,8 @@ def init_ingest_cmds(app): @ingest_cli.command("flush_redis") +@click.confirmation_option(prompt="Are you sure you want to flush redis?") +@job_type_guard(group_name) def flush_redis(): """FLush redis db.""" redis = get_redis_connection() @@ -34,38 +51,111 @@ def flush_redis(): @ingest_cli.command("graph") @click.argument("graph_id", type=str) -@click.argument("dataset", type=click.Path(exists=True)) -@click.option("--raw", is_flag=True) -@click.option("--test", is_flag=True) -@click.option("--retry", is_flag=True) +@click.argument("dataset", type=click.Path(exists=True), required=False) +@click.option("--raw", is_flag=True, help="Read edges from agglomeration output.") +@click.option( + "--retry", + "-r", + is_flag=True, + help="Re-run setup against the existing table (no cg.create()).", +) +@click.option( + "--skip-queue", + "-s", + is_flag=True, + help="Set up everything but don't enqueue L2 tasks.", +) +@click.option( + "--test", + "-t", + is_flag=True, + help="Test 8 chunks at the center of dataset.", +) +@job_type_guard(group_name) def ingest_graph( - graph_id: str, dataset: click.Path, raw: bool, test: bool, retry: bool + graph_id: str, + dataset: click.Path, + raw: bool, + retry: bool, + skip_queue: bool, + test: bool, ): + """Main ingest command. Takes config from yaml, queues atomic tasks. + + Purely about the bigtable graph: creates the table and enqueues L2 + tasks. OCDBT base + fork creation happens in ``ingest layer N`` when + N matches ``ocdbt_populate_layer``; that's the single owner of the + OCDBT lifecycle. + + ``--retry`` reuses the existing IngestionManager from redis and skips + ``cg.create()``. Pair with ``--skip-queue`` to skip L2 enqueue too. """ - Main ingest command. - Takes ingest config from a yaml file and queues atomic tasks. + redis = get_redis_connection() + if test: + configure_logging(level=DEBUG) + + if retry: + imanager_pickle = redis.get(r_keys.INGESTION_MANAGER) + if imanager_pickle is None: + raise click.ClickException( + f"--retry requires an existing `{group_name}` job in redis. " + f"Run without --retry to start a new job." + ) + imanager = IngestionManager.from_pickle(imanager_pickle) + else: + if dataset is None: + raise click.ClickException("dataset is required unless --retry is passed.") + redis.set(r_keys.JOB_TYPE, group_name) + with open(dataset, "r") as stream: + config = yaml.safe_load(stream) + meta, ingest_config, client_info, ocdbt_config_dict = bootstrap( + graph_id, config, raw, test + ) + cg = ChunkedGraph(meta=meta, client_info=client_info) + cg.create() + imanager = IngestionManager( + ingest_config, + meta, + ocdbt_config=ocdbt_config_dict, + ) + + if not skip_queue: + enqueue_l2_tasks(imanager, create_atomic_chunk) + os._exit(0) + + +@ingest_cli.command("mesh_meta") +@click.argument("graph_id", type=str) +@click.argument("dataset", type=click.Path(exists=True)) +@job_type_guard(group_name) +def mesh_meta(graph_id: str, dataset: click.Path): + """Set up every mesh.* metadata field for GRAPH_ID from DATASET yaml. + + Reads ``mesh_config:`` from the yaml, applies it to the graph. Run + once per new/copied graph, after the operator has verified initial + ingest (including the root layer) is complete — no automatic gate. """ - from .cluster import enqueue_atomic_tasks + # nested: pulls meshing/cloudvolume, only needed at call time + from ..meshing.meta import MeshConfig + from ..meshing.setup import setup_mesh_meta with open(dataset, "r") as stream: config = yaml.safe_load(stream) - - meta, ingest_config, client_info = bootstrap( - graph_id, - config=config, - raw=raw, - test_run=test, - ) - cg = ChunkedGraph(meta=meta, client_info=client_info) - if not retry: - cg.create() - enqueue_atomic_tasks(IngestionManager(ingest_config, meta)) + if "mesh_config" not in config: + raise click.ClickException( + f"{dataset} has no `mesh_config:` block — required for mesh_meta." + ) + mesh_cfg = MeshConfig.from_dict(config["mesh_config"]) + cg = ChunkedGraph(graph_id=graph_id) + result = setup_mesh_meta(cg, mesh_cfg) + click.echo(f"mesh meta written for {graph_id}: {result}") @ingest_cli.command("imanager") @click.argument("graph_id", type=str) @click.argument("dataset", type=click.Path(exists=True)) @click.option("--raw", is_flag=True) +@job_type_guard(group_name) def pickle_imanager(graph_id: str, dataset: click.Path, raw: bool): """ Load ingest config into redis server. @@ -77,118 +167,157 @@ def pickle_imanager(graph_id: str, dataset: click.Path, raw: bool): except yaml.YAMLError as exc: print(exc) - meta, ingest_config, _ = bootstrap(graph_id, config=config, raw=raw) - imanager = IngestionManager(ingest_config, meta) - imanager.redis + meta, ingest_config, _, ocdbt_config_dict = bootstrap( + graph_id, config=config, raw=raw + ) + imanager = IngestionManager(ingest_config, meta, ocdbt_config=ocdbt_config_dict) + imanager.redis.set(r_keys.JOB_TYPE, group_name) @ingest_cli.command("layer") @click.argument("parent_layer", type=int) -def queue_layer(parent_layer): +@click.option( + "--queue-only", + "-q", + is_flag=True, + help="Only enqueue tasks; do not start the OCDBT coordinator. " + "Use when a coordinator is already running in another process.", +) +@click.option( + "--ocdbt-only", + "-o", + is_flag=True, + help="Workers run only OCDBT populate (skip add_parent_chunk). " + "Requires the OCDBT populate layer.", +) +@click.option( + "--ingest-only", + "-i", + is_flag=True, + help="Workers run only add_parent_chunk (skip OCDBT populate). " + "Use when the OCDBT base is already populated for this layer.", +) +@job_type_guard(group_name) +def queue_layer(parent_layer, queue_only, ocdbt_only, ingest_only): """ Queue all chunk tasks at a given layer. Must be used when all the chunks at `parent_layer - 1` have completed. - """ - from itertools import product - import numpy as np - from .cluster import create_parent_chunk - from .utils import chunk_id_str + When this layer is the OCDBT populate layer, this command also owns the + OCDBT lifecycle: idempotently creates the base + fork via ``setup_base`` + and starts a ``DistributedCoordinatorServer`` so every worker's commit + routes through one process (eliminates manifest-CAS races and orphan + ``d/`` files). Stays in the foreground until killed. + + Flags: + ``--queue-only`` skips the coordinator (one is assumed running elsewhere). + ``--ocdbt-only`` task body = OCDBT populate only. + ``--ingest-only`` task body = add_parent_chunk only. + """ assert parent_layer > 2, "This command is for layers 3 and above." + if ocdbt_only and ingest_only: + raise click.ClickException( + "--ocdbt-only and --ingest-only are mutually exclusive." + ) redis = get_redis_connection() imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) - if parent_layer == imanager.cg_meta.layer_count: - chunk_coords = [(0, 0, 0)] - else: - bounds = imanager.cg_meta.layer_chunk_bounds[parent_layer] - chunk_coords = randomize_grid_points(*bounds) + is_populate_layer = imanager.is_ocdbt_populate_layer(parent_layer) + if ocdbt_only and not is_populate_layer: + raise click.ClickException( + "--ocdbt-only requires running at the OCDBT populate layer." + ) - def get_chunks_not_done(coords: list) -> list: - """check for set membership in redis in batches""" - coords_strs = ["_".join(map(str, coord)) for coord in coords] - try: - completed = imanager.redis.smismember(f"{parent_layer}c", coords_strs) - except Exception: - return coords - return [coord for coord, c in zip(coords, completed) if not c] - - batch_size = int(environ.get("JOB_BATCH_SIZE", 10000)) - batches = chunked(chunk_coords, batch_size) - q = imanager.get_task_queue(f"l{parent_layer}") - - for batch in batches: - _coords = get_chunks_not_done(batch) - # buffer for optimal use of redis memory - if len(q) > int(environ.get("QUEUE_SIZE", 100000)): - interval = int(environ.get("QUEUE_INTERVAL", 300)) - sleep(interval) - - job_datas = [] - for chunk_coord in _coords: - job_datas.append( - Queue.prepare_data( - create_parent_chunk, - args=(parent_layer, chunk_coord), - result_ttl=0, - job_id=chunk_id_str(parent_layer, chunk_coord), - timeout=f"{int(parent_layer * parent_layer)}m", - ) - ) - q.enqueue_many(job_datas) + if is_populate_layer: + # Single owner of the OCDBT lifecycle: create base + fork if + # missing, reconcile config with on-disk meta, then re-pickle + # imanager so queued workers read the resolved config. + resolved = setup_base(imanager.cg, OcdbtConfig.from_dict(imanager.ocdbt_config)) + imanager.ocdbt_config = resolved.to_dict() + imanager.redis.set(r_keys.INGESTION_MANAGER, imanager.serialized(pickled=True)) + + mode = "ocdbt" if ocdbt_only else ("ingest" if ingest_only else "full") + task_fn = ( + partial(create_parent_chunk, mode=mode) + if mode != "full" + else create_parent_chunk + ) + + # Coordinator only matters when OCDBT populate will actually run. + needs_coordinator = ( + is_populate_layer and mode in ("full", "ocdbt") and not queue_only + ) + if needs_coordinator: + with coordinator(imanager.redis): + queue_layer_helper(parent_layer, imanager, task_fn) + while True: + sleep(60) + else: + queue_layer_helper(parent_layer, imanager, task_fn) @ingest_cli.command("status") -def ingest_status(): +@click.option("--refresh", type=int, default=5, help="Seconds between redis polls.") +@job_type_guard(group_name) +def ingest_status(refresh: int): """Print ingest status to console by layer.""" redis = get_redis_connection() - imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) - layers = range(2, imanager.cg_meta.layer_count + 1) - for layer, layer_count in zip(layers, imanager.cg_meta.layer_chunk_counts): - completed = redis.scard(f"{layer}c") - print(f"{layer}\t: {completed} / {layer_count}") + try: + imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) + print_status(imanager, redis, refresh_seconds=refresh) + except TypeError as err: + print(f"\nNo current `{group_name}` job found in redis: {err}") @ingest_cli.command("chunk") @click.argument("queue", type=str) @click.argument("chunk_info", nargs=4, type=int) +@job_type_guard(group_name) def ingest_chunk(queue: str, chunk_info): """Manually queue chunk when a job is stuck for whatever reason.""" - from .cluster import _create_atomic_chunk - from .cluster import create_parent_chunk - from .utils import chunk_id_str - - redis = get_redis_connection() - imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) - layer = chunk_info[0] - coords = chunk_info[1:] - queue = imanager.get_task_queue(queue) - if layer == 2: - func = _create_atomic_chunk - args = (coords,) - else: - func = create_parent_chunk - args = (layer, coords) - queue.enqueue( - func, - job_id=chunk_id_str(layer, coords), - job_timeout=f"{int(layer * layer)}m", - result_ttl=0, - args=args, - ) + requeue_chunk(queue, chunk_info, create_atomic_chunk, create_parent_chunk) @ingest_cli.command("chunk_local") @click.argument("graph_id", type=str) @click.argument("chunk_info", nargs=4, type=int) -@click.option("--n_threads", type=int, default=1) -def ingest_chunk_local(graph_id: str, chunk_info, n_threads: int): +@click.option("--n_processes", type=int, default=1) +@job_type_guard(group_name) +def ingest_chunk_local(graph_id: str, chunk_info, n_processes: int): """Manually ingest a chunk on a local machine.""" - from .create.abstract_layers import add_layer - from .cluster import _create_atomic_chunk - - if chunk_info[0] == 2: - _create_atomic_chunk(chunk_info[1:]) + layer, coords = chunk_info[0], chunk_info[1:] + if layer == 2: + create_atomic_chunk(coords) else: cg = ChunkedGraph(graph_id=graph_id) - add_layer(cg, chunk_info[0], chunk_info[1:], n_threads=n_threads) + add_parent_chunk(cg, layer, coords, n_processes=n_processes) + cg = ChunkedGraph(graph_id=graph_id) + add_parent_chunk(cg, layer, coords, n_processes=n_processes) + + +@ingest_cli.command("rate") +@click.argument("layer", type=int) +@click.option("--span", default=10, help="Time span to calculate rate.") +@job_type_guard(group_name) +def rate(layer: int, span: int): + redis = get_redis_connection() + imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) + print_completion_rate(imanager, layer, span=span) + + +@ingest_cli.command("run_tests") +@click.argument("graph_id", type=str) +@job_type_guard(group_name) +def run_tests(graph_id): + run_all(ChunkedGraph(graph_id=graph_id)) + + +@ingest_cli.command("purge_layer") +@click.argument("layer", type=int) +@click.confirmation_option(prompt="Purge ALL redis state for this layer?") +@job_type_guard(group_name) +def purge_layer(layer: int): + """Drop the per-layer RQ queue + registries + completion set so the + layer can be re-run from a previous layer's backup.""" + purge_layer_state(get_redis_connection(), layer) + click.echo(f"purged redis state for layer {layer}") diff --git a/pychunkedgraph/ingest/cli_upgrade.py b/pychunkedgraph/ingest/cli_upgrade.py new file mode 100644 index 000000000..83e9e53c8 --- /dev/null +++ b/pychunkedgraph/ingest/cli_upgrade.py @@ -0,0 +1,142 @@ +# pylint: disable=invalid-name, missing-function-docstring, unspecified-encoding + +""" +cli for running upgrade +""" + +import click +from flask.cli import AppGroup + +from pychunkedgraph import __version__, get_logger +from pychunkedgraph.graph.meta import GraphConfig + +from . import IngestConfig +from .cluster import enqueue_l2_tasks, upgrade_atomic_chunk, upgrade_parent_chunk +from .manager import IngestionManager +from .ocdbt import setup_base +from .utils import ( + job_type_guard, + print_completion_rate, + print_status, + queue_layer_helper, + requeue_chunk, +) +from ..graph.chunkedgraph import ChunkedGraph, ChunkedGraphMeta +from ..graph.ocdbt import OcdbtConfig +from ..utils.redis import get_redis_connection +from ..utils.redis import keys as r_keys + +logger = get_logger(__name__) + +group_name = "upgrade" +upgrade_cli = AppGroup(group_name) + + +def init_upgrade_cmds(app): + app.cli.add_command(upgrade_cli) + + +@upgrade_cli.command("flush_redis") +@click.confirmation_option(prompt="Are you sure you want to flush redis?") +@job_type_guard(group_name) +def flush_redis(): + """FLush redis db.""" + redis = get_redis_connection() + redis.flushdb() + + +@upgrade_cli.command("graph") +@click.argument("graph_id", type=str) +@click.option("--test", is_flag=True, help="Test 8 chunks at the center of dataset.") +@click.option("--ocdbt", is_flag=True, help="Enable ocdbt seg (SV splitting support).") +@click.option( + "--sv-split-threshold", + type=int, + default=10, + help="Distance threshold for SV split edge matching.", +) +@job_type_guard(group_name) +def upgrade_graph( + graph_id: str, + test: bool, + ocdbt: bool, + sv_split_threshold: int, +): + """ + Main upgrade command. Queues atomic tasks. + """ + redis = get_redis_connection() + redis.set(r_keys.JOB_TYPE, group_name) + ingest_config = IngestConfig(TEST_RUN=test) + cg = ChunkedGraph(graph_id=graph_id) + cg.client.add_table_version(__version__, overwrite=True) + + if graph_id != cg.graph_id: + gc = cg.meta.graph_config._asdict() + gc["ID"] = graph_id + new_meta = ChunkedGraphMeta( + GraphConfig(**gc), cg.meta.data_source, cg.meta.custom_data + ) + cg.update_meta(new_meta, overwrite=True) + cg = ChunkedGraph(graph_id=graph_id) + + if ocdbt: + ocdbt_cfg = OcdbtConfig.from_dict(cg.meta.custom_data.get("ocdbt_config")) + ocdbt_cfg.enabled = True + ocdbt_cfg.sv_split_threshold = sv_split_threshold + setup_base(cg, ocdbt_cfg) + logger.note(f"enabled ocdbt seg with sv_split_threshold={sv_split_threshold}") + try: + cg.client.create_column_family("4") + except Exception: + ... + + imanager = IngestionManager(ingest_config, cg.meta) + enqueue_l2_tasks(imanager, upgrade_atomic_chunk) + + +@upgrade_cli.command("layer") +@click.argument("parent_layer", type=int) +@click.option("--splits", default=0, help="Split chunks into multiple tasks.") +@job_type_guard(group_name) +def queue_layer(parent_layer: int, splits: int = 0): + """ + Queue all chunk tasks at a given layer. + Must be used when all the chunks at `parent_layer - 1` have completed. + """ + assert parent_layer > 2, "This command is for layers 3 and above." + redis = get_redis_connection() + imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) + queue_layer_helper(parent_layer, imanager, upgrade_parent_chunk, splits=splits) + + +@upgrade_cli.command("status") +@click.option("--refresh", type=int, default=5, help="Seconds between redis polls.") +@job_type_guard(group_name) +def upgrade_status(refresh: int): + """Print upgrade status to console.""" + redis = get_redis_connection() + try: + imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) + print_status(imanager, redis, upgrade=True, refresh_seconds=refresh) + except TypeError as err: + print(f"\nNo current `{group_name}` job found in redis: {err}") + + +@upgrade_cli.command("chunk") +@click.argument("queue", type=str) +@click.argument("chunk_info", nargs=4, type=int) +@job_type_guard(group_name) +def upgrade_chunk(queue: str, chunk_info): + """Manually queue chunk when a job is stuck for whatever reason.""" + requeue_chunk(queue, chunk_info, upgrade_atomic_chunk, upgrade_parent_chunk) + + +@upgrade_cli.command("rate") +@click.argument("layer", type=int) +@click.option("--span", default=10, help="Time span to calculate rate.") +@job_type_guard(group_name) +def rate(layer: int, span: int): + redis = get_redis_connection() + imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) + print_completion_rate(imanager, layer, span=span) diff --git a/pychunkedgraph/ingest/cluster.py b/pychunkedgraph/ingest/cluster.py index cf9417024..36d111f1a 100644 --- a/pychunkedgraph/ingest/cluster.py +++ b/pychunkedgraph/ingest/cluster.py @@ -1,195 +1,241 @@ +# pylint: disable=invalid-name, missing-function-docstring, import-outside-toplevel + """ -Ingest / create chunkedgraph with workers. +Ingest / create chunkedgraph with workers on a cluster. """ -from typing import Sequence, Tuple +from os import environ +from time import sleep +from typing import Callable, Dict, Iterable, Tuple, Sequence import numpy as np +from rq import Queue as RQueue, Retry + +from pychunkedgraph import get_logger + +logger = get_logger(__name__) + -from .utils import chunk_id_str +from .utils import chunk_id_str, get_chunks_not_done, randomize_grid_points from .manager import IngestionManager -from .common import get_atomic_chunk_data -from .ran_agglomeration import get_active_edges -from .create.atomic_layer import add_atomic_edges -from .create.abstract_layers import add_layer -from ..graph.meta import ChunkedGraphMeta +from .ocdbt import get_coordinator_address, populate_chunk +from .ran_agglomeration import ( + get_active_edges, + read_raw_edge_data, + read_raw_agglomeration_data, +) +from .create.atomic_layer import add_atomic_chunk +from .create.parent_layer import add_parent_chunk +from .upgrade.atomic_layer import update_chunk as update_atomic_chunk +from .upgrade.parent_layer import update_chunk as update_parent_chunk +from ..graph.edges import EDGE_TYPES +from ..graph import ChunkedGraph, ChunkedGraphMeta +from ..graph.ocdbt import is_chunk_populated from ..graph.chunks.hierarchy import get_children_chunk_coords -from ..utils.redis import keys as r_keys -from ..utils.redis import get_redis_connection +from ..io.edges import get_chunk_edges +from ..io.components import get_chunk_components +from ..utils.redis import keys as r_keys, get_redis_connection +from ..utils.general import chunked +_CACHED_IMANAGER = None -def _post_task_completion(imanager: IngestionManager, layer: int, coords: np.ndarray): - from os import environ +def _get_imanager(): + """Cache IngestionManager per worker process to avoid repeated Redis GETs + deserializations.""" + global _CACHED_IMANAGER + if _CACHED_IMANAGER is not None: + return _CACHED_IMANAGER + redis = get_redis_connection() + _CACHED_IMANAGER = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) + return _CACHED_IMANAGER + + +def _post_task_completion( + imanager: IngestionManager, layer: int, coords: np.ndarray, split: int = None +): chunk_str = "_".join(map(str, coords)) + if split is not None: + chunk_str += f"_{split}" # mark chunk as completed - "c" imanager.redis.sadd(f"{layer}c", chunk_str) - - if environ.get("DO_NOT_AUTOQUEUE_PARENT_CHUNKS", None) is not None: - return - - parent_layer = layer + 1 - if parent_layer > imanager.cg_meta.layer_count: - return - - parent_coords = np.array(coords, int) // imanager.cg_meta.graph_config.FANOUT - parent_id_str = chunk_id_str(parent_layer, parent_coords) - imanager.redis.sadd(parent_id_str, chunk_str) - - parent_chunk_str = "_".join(map(str, parent_coords)) - if not imanager.redis.hget(parent_layer, parent_chunk_str): - # cache children chunk count - # checked by tracker worker to enqueue parent chunk - children_count = len( - get_children_chunk_coords(imanager.cg_meta, parent_layer, parent_coords) - ) - imanager.redis.hset(parent_layer, parent_chunk_str, children_count) - - tracker_queue = imanager.get_task_queue(f"t{layer}") - tracker_queue.enqueue( - enqueue_parent_task, - job_id=f"t{layer}_{chunk_str}", - job_timeout=f"30s", - result_ttl=0, - args=( - parent_layer, - parent_coords, - ), - ) + logger.note(f"{chunk_str} marked as complete") -def enqueue_parent_task( +def create_parent_chunk( parent_layer: int, parent_coords: Sequence[int], -): - redis = get_redis_connection() - imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) - parent_id_str = chunk_id_str(parent_layer, parent_coords) - parent_chunk_str = "_".join(map(str, parent_coords)) - - children_done = redis.scard(parent_id_str) - # if zero then this key was deleted and parent already queued. - if children_done == 0: - print("parent already queued.") - return - - # if the previous layer is complete - # no need to check children progress for each parent chunk - child_layer = parent_layer - 1 - child_layer_done = redis.scard(f"{child_layer}c") - child_layer_count = imanager.cg_meta.layer_chunk_counts[child_layer - 2] - child_layer_finished = child_layer_done == child_layer_count - - if not child_layer_finished: - children_count = int(redis.hget(parent_layer, parent_chunk_str).decode("utf-8")) - if children_done != children_count: - print("children not done.") - return - - queue = imanager.get_task_queue(f"l{parent_layer}") - queue.enqueue( - create_parent_chunk, - job_id=parent_id_str, - job_timeout=f"{int(parent_layer * parent_layer)}m", - result_ttl=0, - args=( + mode: str = "full", +) -> None: + """One parent-chunk task. ``mode`` (bound at queue time via partial) + selects which halves run: + ``full`` : OCDBT populate (if eligible) + add_parent_chunk + ``ocdbt`` : only OCDBT populate (skip add_parent_chunk) + ``ingest`` : only add_parent_chunk (skip OCDBT populate) + + ``_post_task_completion`` always runs so the layer's progress tracking + in redis stays consistent. + + OCDBT populate runs FIRST so any failure aborts the task BEFORE graph + mutation; otherwise a half-built graph would force corrupt-state retries. + """ + imanager = _get_imanager() + + do_ocdbt = mode in ("full", "ocdbt") and imanager.is_ocdbt_populate_layer( + parent_layer + ) + do_ingest = mode in ("full", "ingest") + + if do_ocdbt: + ws = imanager.cg.meta.data_source.WATERSHED + if not is_chunk_populated(ws, parent_layer, parent_coords): + address = get_coordinator_address(imanager.redis) + populate_chunk( + imanager, ws, parent_layer, parent_coords, coordinator_address=address + ) + + if do_ingest: + add_parent_chunk( + imanager.cg, parent_layer, parent_coords, - ), - ) - redis.hdel(parent_layer, parent_chunk_str) - redis.delete(parent_id_str) + get_children_chunk_coords( + imanager.cg_meta, + parent_layer, + parent_coords, + ), + ) + + _post_task_completion(imanager, parent_layer, parent_coords) -def create_parent_chunk( +def upgrade_parent_chunk( parent_layer: int, parent_coords: Sequence[int], + split: int = None, + splits: int = None, ) -> None: - redis = get_redis_connection() - imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) - add_layer( - imanager.cg, - parent_layer, - parent_coords, - get_children_chunk_coords( - imanager.cg_meta, - parent_layer, - parent_coords, - ), + imanager = _get_imanager() + update_parent_chunk( + imanager.cg, parent_coords, layer=parent_layer, split=split, splits=splits + ) + _post_task_completion(imanager, parent_layer, parent_coords, split=split) + + +def _get_atomic_chunk_data( + imanager: IngestionManager, coord: Sequence[int] +) -> Tuple[Dict, Dict]: + """ + Helper to read either raw data or processed data + If reading from raw data, save it as processed data + """ + chunk_edges = ( + read_raw_edge_data(imanager, coord) + if imanager.config.USE_RAW_EDGES + else get_chunk_edges(imanager.cg_meta.data_source.EDGES, [coord]) ) - _post_task_completion(imanager, parent_layer, parent_coords) + _check_edges_direction(chunk_edges, imanager.cg, coord) -def randomize_grid_points(X: int, Y: int, Z: int) -> Tuple[int, int, int]: - indices = np.arange(X * Y * Z) - np.random.shuffle(indices) - for index in indices: - yield np.unravel_index(index, (X, Y, Z)) + mapping = ( + read_raw_agglomeration_data(imanager, coord) + if imanager.config.USE_RAW_COMPONENTS + else get_chunk_components(imanager.cg_meta.data_source.COMPONENTS, coord) + ) + return chunk_edges, mapping -def enqueue_atomic_tasks(imanager: IngestionManager): - from os import environ - from time import sleep - from rq import Queue as RQueue +def _check_edges_direction( + chunk_edges: dict, cg: ChunkedGraph, coord: Sequence[int] +) -> None: + """ + For between and cross chunk edges: + Checks and flips edges such that nodes1 are always within a chunk and nodes2 outside the chunk. + Where nodes1 = edges[:,0] and nodes2 = edges[:,1]. + """ + x, y, z = coord + chunk_id = cg.get_chunk_id(layer=1, x=x, y=y, z=z) + for edge_type in [EDGE_TYPES.between_chunk, EDGE_TYPES.cross_chunk]: + edges = chunk_edges[edge_type] + chunk_ids = cg.get_chunk_ids_from_node_ids(edges.node_ids1) + mask = chunk_ids == chunk_id + assert np.all(mask), "all IDs must belong to same chunk" + + +def create_atomic_chunk(coords: Sequence[int]): + """Creates single atomic chunk""" + imanager = _get_imanager() + coords = np.array(list(coords), dtype=int) - chunk_coords = _get_test_chunks(imanager.cg.meta) - chunk_count = len(chunk_coords) - if not imanager.config.TEST_RUN: - atomic_chunk_bounds = imanager.cg_meta.layer_chunk_bounds[2] - chunk_coords = randomize_grid_points(*atomic_chunk_bounds) - chunk_count = imanager.cg_meta.layer_chunk_counts[0] + chunk_edges_all, mapping = _get_atomic_chunk_data(imanager, coords) + chunk_edges_active, isolated_ids = get_active_edges(chunk_edges_all, mapping) + add_atomic_chunk(imanager.cg, coords, chunk_edges_active, isolated=isolated_ids) - print(f"total chunk count: {chunk_count}, queuing...") - batch_size = int(environ.get("L2JOB_BATCH_SIZE", 1000)) + for k, v in chunk_edges_all.items(): + logger.debug(f"{k}: {len(v)}") + for k, v in chunk_edges_active.items(): + logger.debug(f"active_{k}: {len(v)}") - job_datas = [] - for chunk_coord in chunk_coords: - q = imanager.get_task_queue(imanager.config.CLUSTER.ATOMIC_Q_NAME) - # buffer for optimal use of redis memory - if len(q) > imanager.config.CLUSTER.ATOMIC_Q_LIMIT: - print(f"Sleeping {imanager.config.CLUSTER.ATOMIC_Q_INTERVAL}s...") - sleep(imanager.config.CLUSTER.ATOMIC_Q_INTERVAL) - - x, y, z = chunk_coord - chunk_str = f"{x}_{y}_{z}" - if imanager.redis.sismember("2c", chunk_str): - # already done, skip - continue - job_datas.append( - RQueue.prepare_data( - _create_atomic_chunk, - args=(chunk_coord,), - timeout=environ.get("L2JOB_TIMEOUT", "3m"), - result_ttl=0, - job_id=chunk_id_str(2, chunk_coord), - ) - ) - if len(job_datas) % batch_size == 0: - q.enqueue_many(job_datas) - job_datas = [] - q.enqueue_many(job_datas) + _post_task_completion(imanager, 2, coords) -def _create_atomic_chunk(coords: Sequence[int]): - """Creates single atomic chunk""" - redis = get_redis_connection() - imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) +def upgrade_atomic_chunk(coords: Sequence[int]): + """Upgrades single atomic chunk""" + imanager = _get_imanager() coords = np.array(list(coords), dtype=int) - chunk_edges_all, mapping = get_atomic_chunk_data(imanager, coords) - chunk_edges_active, isolated_ids = get_active_edges(chunk_edges_all, mapping) - add_atomic_edges(imanager.cg, coords, chunk_edges_active, isolated=isolated_ids) - if imanager.config.TEST_RUN: - # print for debugging - for k, v in chunk_edges_all.items(): - print(k, len(v)) - for k, v in chunk_edges_active.items(): - print(f"active_{k}", len(v)) + update_atomic_chunk(imanager.cg, coords) _post_task_completion(imanager, 2, coords) def _get_test_chunks(meta: ChunkedGraphMeta): - """Chunks at center of the dataset most likely not to be empty""" + """Chunks at the center most likely not to be empty""" parent_coords = np.array(meta.layer_chunk_bounds[3]) // 2 return get_children_chunk_coords(meta, 3, parent_coords) - # f = lambda r1, r2, r3: np.array(np.meshgrid(r1, r2, r3), dtype=int).T.reshape(-1, 3) - # return f((x, x + 1), (y, y + 1), (z, z + 1)) + + +def _queue_tasks(imanager: IngestionManager, chunk_fn: Callable, coords: Iterable): + queue_name = "l2" + q = imanager.get_task_queue(queue_name) + batch_size = int(environ.get("JOB_BATCH_SIZE", 10000)) + batches = chunked(coords, batch_size) + retry = int(environ.get("RETRY_COUNT", 0)) + failure_ttl = int(environ.get("FAILURE_TTL", 300)) + max_queue_size = int(environ.get("QUEUE_SIZE", 1000000)) + for batch in batches: + _coords = get_chunks_not_done(imanager, 2, batch) + # buffer for optimal use of redis memory + while len(q) > max_queue_size: + logger.note( + f"Queue has {len(q)} items (limit {max_queue_size}), waiting..." + ) + sleep(10) + + job_datas = [] + for chunk_coord in _coords: + job_datas.append( + RQueue.prepare_data( + chunk_fn, + args=(chunk_coord,), + timeout=environ.get("L2JOB_TIMEOUT", "3m"), + result_ttl=0, + job_id=chunk_id_str(2, chunk_coord), + retry=Retry(retry) if retry > 1 else None, + description="", + failure_ttl=failure_ttl, + ) + ) + q.enqueue_many(job_datas) + logger.note(f"Queued {len(job_datas)} chunks.") + + +def enqueue_l2_tasks(imanager: IngestionManager, chunk_fn: Callable): + """ + `chunk_fn`: function to process a given layer 2 chunk. + """ + chunk_coords = _get_test_chunks(imanager.cg.meta) + chunk_count = len(chunk_coords) + if not imanager.config.TEST_RUN: + atomic_chunk_bounds = imanager.cg_meta.layer_chunk_bounds[2] + chunk_coords = randomize_grid_points(*atomic_chunk_bounds) + chunk_count = imanager.cg_meta.layer_chunk_counts[0] + logger.note(f"Chunk count: {chunk_count}, queuing...") + _queue_tasks(imanager, chunk_fn, chunk_coords) diff --git a/pychunkedgraph/ingest/common.py b/pychunkedgraph/ingest/common.py deleted file mode 100644 index dccf58602..000000000 --- a/pychunkedgraph/ingest/common.py +++ /dev/null @@ -1,61 +0,0 @@ -from typing import Dict -from typing import Tuple -from typing import Sequence - -from .manager import IngestionManager -from .ran_agglomeration import read_raw_edge_data -from .ran_agglomeration import read_raw_agglomeration_data -from ..graph import ChunkedGraph -from ..io.edges import get_chunk_edges -from ..io.components import get_chunk_components - - -def get_atomic_chunk_data( - imanager: IngestionManager, coord: Sequence[int] -) -> Tuple[Dict, Dict]: - """ - Helper to read either raw data or processed data - If reading from raw data, save it as processed data - """ - chunk_edges = ( - read_raw_edge_data(imanager, coord) - if imanager.config.USE_RAW_EDGES - else get_chunk_edges(imanager.cg_meta.data_source.EDGES, [coord]) - ) - - _check_edges_direction(chunk_edges, imanager.cg, coord) - - mapping = ( - read_raw_agglomeration_data(imanager, coord) - if imanager.config.USE_RAW_COMPONENTS - else get_chunk_components(imanager.cg_meta.data_source.COMPONENTS, coord) - ) - return chunk_edges, mapping - - -def _check_edges_direction( - chunk_edges: dict, cg: ChunkedGraph, coord: Sequence[int] -) -> None: - """ - For between and cross chunk edges: - Checks and flips edges such that nodes1 are always within a chunk and nodes2 outside the chunk. - Where nodes1 = edges[:,0] and nodes2 = edges[:,1]. - """ - import numpy as np - from ..graph.edges import Edges - from ..graph.edges import EDGE_TYPES - - x, y, z = coord - chunk_id = cg.get_chunk_id(layer=1, x=x, y=y, z=z) - for edge_type in [EDGE_TYPES.between_chunk, EDGE_TYPES.cross_chunk]: - edges = chunk_edges[edge_type] - e1 = edges.node_ids1 - e2 = edges.node_ids2 - - e2_chunk_ids = cg.get_chunk_ids_from_node_ids(e2) - mask = e2_chunk_ids == chunk_id - e1[mask], e2[mask] = e2[mask], e1[mask] - - e1_chunk_ids = cg.get_chunk_ids_from_node_ids(e1) - mask = e1_chunk_ids == chunk_id - assert np.all(mask), "all IDs must belong to same chunk" diff --git a/pychunkedgraph/ingest/create/abstract_layers.py b/pychunkedgraph/ingest/create/abstract_layers.py deleted file mode 100644 index 529a6846f..000000000 --- a/pychunkedgraph/ingest/create/abstract_layers.py +++ /dev/null @@ -1,247 +0,0 @@ -""" -Functions for creating parents in level 3 and above -""" - -import time -import math -import datetime -import multiprocessing as mp -from collections import defaultdict -from typing import Optional -from typing import Sequence -from typing import List - -import numpy as np -from multiwrapper import multiprocessing_utils as mu - -from ...graph import types -from ...graph import attributes -from ...utils.general import chunked -from ...graph.utils import flatgraph -from ...graph.utils import basetypes -from ...graph.utils import serializers -from ...graph.chunkedgraph import ChunkedGraph -from ...graph.utils.generic import get_valid_timestamp -from ...graph.utils.generic import filter_failed_node_ids -from ...graph.chunks.hierarchy import get_children_chunk_coords -from ...graph.connectivity.cross_edges import get_children_chunk_cross_edges -from ...graph.connectivity.cross_edges import get_chunk_nodes_cross_edge_layer - - -def add_layer( - cg: ChunkedGraph, - layer_id: int, - parent_coords: Sequence[int], - children_coords: Sequence[Sequence[int]] = np.array([]), - *, - time_stamp: Optional[datetime.datetime] = None, - n_threads: int = 4, -) -> None: - if not children_coords.size: - children_coords = get_children_chunk_coords(cg.meta, layer_id, parent_coords) - children_ids = _read_children_chunks(cg, layer_id, children_coords, n_threads > 1) - edge_ids = get_children_chunk_cross_edges( - cg, layer_id, parent_coords, use_threads=n_threads > 1 - ) - - print("children_coords", children_coords.size, layer_id, parent_coords) - print( - "n e", len(children_ids), len(edge_ids), layer_id, parent_coords, - ) - - node_layers = cg.get_chunk_layers(children_ids) - edge_layers = cg.get_chunk_layers(np.unique(edge_ids)) - assert np.all(node_layers < layer_id), "invalid node layers" - assert np.all(edge_layers < layer_id), "invalid edge layers" - # Extract connected components - # isolated_node_mask = ~np.in1d(children_ids, np.unique(edge_ids)) - # add_node_ids = children_ids[isolated_node_mask].squeeze() - add_edge_ids = np.vstack([children_ids, children_ids]).T - - edge_ids = list(edge_ids) - edge_ids.extend(add_edge_ids) - graph, _, _, graph_ids = flatgraph.build_gt_graph(edge_ids, make_directed=True) - ccs = flatgraph.connected_components(graph) - print("ccs", len(ccs)) - _write_connected_components( - cg, - layer_id, - parent_coords, - ccs, - graph_ids, - get_valid_timestamp(time_stamp), - n_threads > 1, - ) - return f"{layer_id}_{'_'.join(map(str, parent_coords))}" - - -def _read_children_chunks( - cg: ChunkedGraph, layer_id, children_coords, use_threads=True -): - if not use_threads: - children_ids = [types.empty_1d] - for child_coord in children_coords: - children_ids.append(_read_chunk([], cg, layer_id - 1, child_coord)) - return np.concatenate(children_ids) - - print("_read_children_chunks") - with mp.Manager() as manager: - children_ids_shared = manager.list() - multi_args = [] - for child_coord in children_coords: - multi_args.append( - ( - children_ids_shared, - cg.get_serialized_info(), - layer_id - 1, - child_coord, - ) - ) - mu.multiprocess_func( - _read_chunk_helper, - multi_args, - n_threads=min(len(multi_args), mp.cpu_count()), - ) - print("_read_children_chunks done") - return np.concatenate(children_ids_shared) - - -def _read_chunk_helper(args): - children_ids_shared, cg_info, layer_id, chunk_coord = args - cg = ChunkedGraph(**cg_info) - _read_chunk(children_ids_shared, cg, layer_id, chunk_coord) - - -def _read_chunk(children_ids_shared, cg: ChunkedGraph, layer_id: int, chunk_coord): - print(f"_read_chunk {layer_id}, {chunk_coord}") - x, y, z = chunk_coord - range_read = cg.range_read_chunk( - cg.get_chunk_id(layer=layer_id, x=x, y=y, z=z), - properties=attributes.Hierarchy.Child, - ) - row_ids = [] - max_children_ids = [] - for row_id, row_data in range_read.items(): - row_ids.append(row_id) - max_children_ids.append(np.max(row_data[0].value)) - row_ids = np.array(row_ids, dtype=basetypes.NODE_ID) - segment_ids = np.array([cg.get_segment_id(r_id) for r_id in row_ids]) - - row_ids = filter_failed_node_ids(row_ids, segment_ids, max_children_ids) - children_ids_shared.append(row_ids) - print(f"_read_chunk {layer_id}, {chunk_coord} done {len(row_ids)}") - return row_ids - - -def _write_connected_components( - cg: ChunkedGraph, - layer_id: int, - parent_coords, - ccs, - graph_ids, - time_stamp, - use_threads=True, -) -> None: - if not ccs: - return - - node_layer_d_shared = {} - if layer_id < cg.meta.layer_count: - print("getting node_layer_d_shared") - node_layer_d_shared = get_chunk_nodes_cross_edge_layer( - cg, layer_id, parent_coords, use_threads=use_threads - ) - - print("node_layer_d_shared", len(node_layer_d_shared)) - - ccs_with_node_ids = [] - for cc in ccs: - ccs_with_node_ids.append(graph_ids[cc]) - - if not use_threads: - _write( - cg, - layer_id, - parent_coords, - ccs_with_node_ids, - node_layer_d_shared, - time_stamp, - use_threads=use_threads, - ) - return - - task_size = int(math.ceil(len(ccs_with_node_ids) / mp.cpu_count() / 10)) - chunked_ccs = chunked(ccs_with_node_ids, task_size) - cg_info = cg.get_serialized_info() - multi_args = [] - for ccs in chunked_ccs: - multi_args.append( - (cg_info, layer_id, parent_coords, ccs, node_layer_d_shared, time_stamp) - ) - mu.multiprocess_func( - _write_components_helper, - multi_args, - n_threads=min(len(multi_args), mp.cpu_count()), - ) - - -def _write_components_helper(args): - print("running _write_components_helper") - cg_info, layer_id, parent_coords, ccs, node_layer_d_shared, time_stamp = args - cg = ChunkedGraph(**cg_info) - _write(cg, layer_id, parent_coords, ccs, node_layer_d_shared, time_stamp) - - -def _write( - cg, layer_id, parent_coords, ccs, node_layer_d_shared, time_stamp, use_threads=True -): - parent_layer_ids = range(layer_id, cg.meta.layer_count + 1) - cc_connections = {l: [] for l in parent_layer_ids} - for node_ids in ccs: - layer = layer_id - if len(node_ids) == 1: - layer = node_layer_d_shared.get(node_ids[0], cg.meta.layer_count) - cc_connections[layer].append(node_ids) - - rows = [] - x, y, z = parent_coords - parent_chunk_id = cg.get_chunk_id(layer=layer_id, x=x, y=y, z=z) - parent_chunk_id_dict = cg.get_parent_chunk_id_dict(parent_chunk_id) - - # Iterate through layers - for parent_layer_id in parent_layer_ids: - if len(cc_connections[parent_layer_id]) == 0: - continue - - parent_chunk_id = parent_chunk_id_dict[parent_layer_id] - reserved_parent_ids = cg.id_client.create_node_ids( - parent_chunk_id, - size=len(cc_connections[parent_layer_id]), - root_chunk=parent_layer_id == cg.meta.layer_count and use_threads, - ) - - for i_cc, node_ids in enumerate(cc_connections[parent_layer_id]): - parent_id = reserved_parent_ids[i_cc] - for node_id in node_ids: - rows.append( - cg.client.mutate_row( - serializers.serialize_uint64(node_id), - {attributes.Hierarchy.Parent: parent_id}, - time_stamp=time_stamp, - ) - ) - - rows.append( - cg.client.mutate_row( - serializers.serialize_uint64(parent_id), - {attributes.Hierarchy.Child: node_ids}, - time_stamp=time_stamp, - ) - ) - - if len(rows) > 100000: - cg.client.write(rows) - print("wrote rows", len(rows), layer_id, parent_coords) - rows = [] - cg.client.write(rows) - print("wrote rows", len(rows), layer_id, parent_coords) diff --git a/pychunkedgraph/ingest/create/atomic_layer.py b/pychunkedgraph/ingest/create/atomic_layer.py index 4fa1f1688..bdce44283 100644 --- a/pychunkedgraph/ingest/create/atomic_layer.py +++ b/pychunkedgraph/ingest/create/atomic_layer.py @@ -1,48 +1,54 @@ +# pylint: disable=invalid-name, missing-function-docstring, import-outside-toplevel + """ Functions for creating atomic nodes and their level 2 abstract parents """ import datetime from typing import Dict -from typing import List from typing import Optional from typing import Sequence -import pytz import numpy as np -from ...graph import attributes +from pychunkedgraph import get_logger + +from ...graph import attributes, basetypes, serializers, get_valid_timestamp from ...graph.chunkedgraph import ChunkedGraph -from ...graph.utils import basetypes -from ...graph.utils import serializers from ...graph.edges import Edges from ...graph.edges import EDGE_TYPES from ...graph.utils.generic import compute_indices_pandas -from ...graph.utils.generic import get_valid_timestamp from ...graph.utils.flatgraph import build_gt_graph from ...graph.utils.flatgraph import connected_components +logger = get_logger(__name__) -def add_atomic_edges( + +def add_atomic_chunk( cg: ChunkedGraph, - chunk_coord: np.ndarray, + coords: Sequence[int], chunk_edges_d: Dict[str, Edges], isolated: Sequence[int], time_stamp: Optional[datetime.datetime] = None, ): chunk_node_ids, chunk_edge_ids = _get_chunk_nodes_and_edges(chunk_edges_d, isolated) + logger.note( + f"L2 chunk {tuple(map(int, coords))}: nodes={len(chunk_node_ids):,} " + f"edges={len(chunk_edge_ids):,}" + ) if not chunk_node_ids.size: return chunk_ids = cg.get_chunk_ids_from_node_ids(chunk_node_ids) - assert len(np.unique(chunk_ids)) == 1 + assert len(np.unique(chunk_ids)) == 1, np.unique(chunk_ids) + + max_node_id = np.max(chunk_node_ids) + cg.id_client.set_max_node_id(chunk_ids[0], max_node_id) graph, _, _, unique_ids = build_gt_graph(chunk_edge_ids, make_directed=True) ccs = connected_components(graph) - parent_chunk_id = cg.get_chunk_id( - layer=2, x=chunk_coord[0], y=chunk_coord[1], z=chunk_coord[2] - ) + parent_chunk_id = cg.get_chunk_id(layer=2, x=coords[0], y=coords[1], z=coords[2]) parent_ids = cg.id_client.create_node_ids(parent_chunk_id, size=len(ccs)) sparse_indices, remapping = _get_remapping(chunk_edges_d) @@ -70,8 +76,11 @@ def _get_chunk_nodes_and_edges(chunk_edges_d: dict, isolated_ids: Sequence[int]) in-chunk edges and nodes_ids """ isolated_nodes_self_edges = np.vstack([isolated_ids, isolated_ids]).T - node_ids = [isolated_ids] - edge_ids = [isolated_nodes_self_edges] + + node_ids = [isolated_ids] if len(isolated_ids) != 0 else [] + edge_ids = ( + [isolated_nodes_self_edges] if len(isolated_nodes_self_edges) != 0 else [] + ) for edge_type in EDGE_TYPES: edges = chunk_edges_d[edge_type] node_ids.append(edges.node_ids1) @@ -79,9 +88,9 @@ def _get_chunk_nodes_and_edges(chunk_edges_d: dict, isolated_ids: Sequence[int]) node_ids.append(edges.node_ids2) edge_ids.append(edges.get_pairs()) - chunk_node_ids = np.unique(np.concatenate(node_ids)) + chunk_node_ids = np.unique(np.concatenate(node_ids).astype(basetypes.NODE_ID)) edge_ids.append(np.vstack([chunk_node_ids, chunk_node_ids]).T) - return (chunk_node_ids, np.concatenate(edge_ids)) + return (chunk_node_ids, np.concatenate(edge_ids).astype(basetypes.NODE_ID)) def _get_remapping(chunk_edges_d: dict): @@ -101,7 +110,13 @@ def _get_remapping(chunk_edges_d: dict): def _process_component( - cg, chunk_edges_d, parent_id, node_ids, sparse_indices, remapping, time_stamp, + cg, + chunk_edges_d, + parent_id, + node_ids, + sparse_indices, + remapping, + time_stamp, ): nodes = [] chunk_out_edges = [] # out = between + cross @@ -112,7 +127,7 @@ def _process_component( r_key = serializers.serialize_uint64(node_id) nodes.append(cg.client.mutate_row(r_key, val_dict, time_stamp=time_stamp)) - chunk_out_edges = np.concatenate(chunk_out_edges) + chunk_out_edges = np.concatenate(chunk_out_edges).astype(basetypes.NODE_ID) cce_layers = cg.get_cross_chunk_edges_layer(chunk_out_edges) u_cce_layers = np.unique(cce_layers) @@ -120,7 +135,7 @@ def _process_component( for cc_layer in u_cce_layers: layer_out_edges = chunk_out_edges[cce_layers == cc_layer] if layer_out_edges.size: - col = attributes.Connectivity.CrossChunkEdge[cc_layer] + col = attributes.Connectivity.AtomicCrossChunkEdge[cc_layer] val_dict[col] = layer_out_edges r_key = serializers.serialize_uint64(parent_id) @@ -143,5 +158,7 @@ def _get_outgoing_edges(node_id, chunk_edges_d, sparse_indices, remapping): ] row_ids = row_ids[column_ids == 0] # edges that this node is part of - chunk_out_edges = np.concatenate([chunk_out_edges, edges[row_ids]]) + chunk_out_edges = np.concatenate([chunk_out_edges, edges[row_ids]]).astype( + basetypes.NODE_ID + ) return chunk_out_edges diff --git a/pychunkedgraph/ingest/create/cross_edges.py b/pychunkedgraph/ingest/create/cross_edges.py new file mode 100644 index 000000000..873941d7e --- /dev/null +++ b/pychunkedgraph/ingest/create/cross_edges.py @@ -0,0 +1,254 @@ +# pylint: disable=invalid-name, missing-docstring + +import math +import multiprocessing as mp +from collections import defaultdict +from typing import Sequence +from typing import Dict + +import numpy as np +from ...graph import attributes, basetypes +from ...graph.types import empty_2d +from ...graph.chunkedgraph import ChunkedGraph +from ...graph.utils.generic import filter_failed_node_ids +from ...graph.chunks.atomic import get_touching_atomic_chunks +from ...graph.chunks.atomic import get_bounding_atomic_chunks +from ...utils.general import chunked + + +def get_children_chunk_cross_edges( + cg: ChunkedGraph, layer, chunk_coord, *, n_processes: int = 1 +) -> np.ndarray: + """ + Cross edges that connect children chunks. + The edges are between node IDs in the given layer. + """ + atomic_chunks = get_touching_atomic_chunks(cg.meta, layer, chunk_coord) + if len(atomic_chunks) == 0: + return [] + + if n_processes <= 1: + return _get_children_chunk_cross_edges(cg, atomic_chunks, layer - 1) + + with mp.Manager() as manager: + edge_ids_shared = manager.list() + edge_ids_shared.append(empty_2d) + + task_size = int(math.ceil(len(atomic_chunks) / n_processes / 10)) + chunked_l2chunk_list = chunked(atomic_chunks, task_size) + multi_args = [] + for atomic_chunks in chunked_l2chunk_list: + multi_args.append( + (edge_ids_shared, cg.get_serialized_info(), atomic_chunks, layer - 1) + ) + + with mp.Pool(processes=min(len(multi_args), n_processes)) as pool: + pool.map(_get_children_chunk_cross_edges_helper, multi_args) + + cross_edges = np.concatenate(edge_ids_shared) + if cross_edges.size: + return np.unique(cross_edges, axis=0) + return cross_edges + + +def _get_children_chunk_cross_edges_helper(args) -> None: + edge_ids_shared, cg_info, atomic_chunks, layer = args + # Re-raise as a bare RuntimeError: the original may hold an unpicklable client + # handle, which the pool would surface as MaybeEncodingError, losing the cause. + try: + cg = ChunkedGraph(**cg_info) + edge_ids_shared.append( + _get_children_chunk_cross_edges(cg, atomic_chunks, layer) + ) + except Exception as exc: + raise RuntimeError( + f"_get_children_chunk_cross_edges failed at layer {layer}: {exc!r}" + ) from None + + +def _get_children_chunk_cross_edges( + cg: ChunkedGraph, atomic_chunks, layer +) -> np.ndarray: + """ + Non parallelized version + Cross edges that connect children chunks. + The edges are between node IDs in the given layer (not atomic). + """ + cross_edges = [empty_2d] + for layer2_chunk in atomic_chunks: + edges = _read_atomic_chunk_cross_edges(cg, layer2_chunk, layer) + cross_edges.append(edges) + + cross_edges = np.concatenate(cross_edges) + if not cross_edges.size: + return empty_2d + + cross_edges[:, 0] = cg.get_roots(cross_edges[:, 0], stop_layer=layer, ceil=False) + cross_edges[:, 1] = cg.get_roots(cross_edges[:, 1], stop_layer=layer, ceil=False) + result = np.unique(cross_edges, axis=0) if cross_edges.size else empty_2d + return result + + +def _read_atomic_chunk_cross_edges( + cg: ChunkedGraph, chunk_coord: Sequence[int], cross_edge_layer: int +) -> np.ndarray: + """ + Returns cross edges between l2 nodes in current chunk and + l1 supervoxels from neighbor chunks. + """ + cross_edge_col = attributes.Connectivity.AtomicCrossChunkEdge[cross_edge_layer] + range_read, l2ids = _read_atomic_chunk(cg, chunk_coord, [cross_edge_layer]) + + parent_neighboring_chunk_supervoxels_d = defaultdict(list) + for l2id in l2ids: + if not cross_edge_col in range_read[l2id]: + continue + edges = range_read[l2id][cross_edge_col][0].value + parent_neighboring_chunk_supervoxels_d[l2id] = edges[:, 1] + + cross_edges = [empty_2d] + for l2id, nebor_svs in parent_neighboring_chunk_supervoxels_d.items(): + chunk_parent_ids = np.array([l2id] * len(nebor_svs), dtype=basetypes.NODE_ID) + cross_edges.append(np.vstack([chunk_parent_ids, nebor_svs]).T) + cross_edges = np.concatenate(cross_edges) + return cross_edges + + +def get_chunk_nodes_cross_edge_layer( + cg: ChunkedGraph, layer: int, chunk_coord: Sequence[int], n_processes: int = 1 +) -> Dict: + """ + gets nodes in a chunk that are part of cross chunk edges + return_type dict {node_id: layer} + the lowest layer (>= current layer) at which a node_id is part of a cross edge + """ + atomic_chunks = get_bounding_atomic_chunks(cg.meta, layer, chunk_coord) + if len(atomic_chunks) == 0: + return {} + + if n_processes <= 1: + return _get_chunk_nodes_cross_edge_layer(cg, atomic_chunks, layer) + + cg_info = cg.get_serialized_info() + manager = mp.Manager() + node_ids_shared = manager.list() + node_layers_shared = manager.list() + task_size = int(math.ceil(len(atomic_chunks) / n_processes / 10)) + chunked_l2chunk_list = chunked(atomic_chunks, task_size) + multi_args = [] + for atomic_chunks in chunked_l2chunk_list: + multi_args.append( + (node_ids_shared, node_layers_shared, cg_info, atomic_chunks, layer) + ) + + with mp.Pool(processes=min(len(multi_args), n_processes)) as pool: + pool.map(_get_chunk_nodes_cross_edge_layer_helper, multi_args) + + node_layer_d_shared = manager.dict() + _find_min_layer(node_layer_d_shared, node_ids_shared, node_layers_shared) + return node_layer_d_shared + + +def _get_chunk_nodes_cross_edge_layer_helper(args): + node_ids_shared, node_layers_shared, cg_info, atomic_chunks, layer = args + # See _get_children_chunk_cross_edges_helper: keep the failure picklable. + try: + cg = ChunkedGraph(**cg_info) + node_layer_d = _get_chunk_nodes_cross_edge_layer(cg, atomic_chunks, layer) + node_ids_shared.append( + np.fromiter(node_layer_d.keys(), dtype=basetypes.NODE_ID) + ) + node_layers_shared.append(np.fromiter(node_layer_d.values(), dtype=np.uint8)) + except Exception as exc: + raise RuntimeError( + f"_get_chunk_nodes_cross_edge_layer failed at layer {layer}: {exc!r}" + ) from None + + +def _get_chunk_nodes_cross_edge_layer(cg: ChunkedGraph, atomic_chunks, layer): + """ + Non parallelized version + gets nodes in a chunk that are part of cross chunk edges + return_type dict {node_id: layer} + the lowest layer (>= current layer) at which a node_id is part of a cross edge + """ + atomic_node_layer_d = {} + for atomic_chunk in atomic_chunks: + chunk_node_layer_d = _read_atomic_chunk_cross_edge_nodes( + cg, atomic_chunk, layer + ) + atomic_node_layer_d.update(chunk_node_layer_d) + + l2ids = np.fromiter(atomic_node_layer_d.keys(), dtype=basetypes.NODE_ID) + parents = cg.get_roots(l2ids, stop_layer=layer - 1, ceil=False) + layers = np.fromiter(atomic_node_layer_d.values(), dtype=int) + + node_layer_d = defaultdict(lambda: cg.meta.layer_count) + for i, parent in enumerate(parents): + node_layer_d[parent] = min(node_layer_d[parent], layers[i]) + return node_layer_d + + +def _read_atomic_chunk_cross_edge_nodes(cg: ChunkedGraph, chunk_coord, layer): + """ + the lowest layer at which an l2 node is part of a cross edge + """ + node_layer_d = {} + relevant_layers = range(layer, cg.meta.layer_count) + range_read, l2ids = _read_atomic_chunk(cg, chunk_coord, relevant_layers) + for l2id in l2ids: + for layer in relevant_layers: + if attributes.Connectivity.AtomicCrossChunkEdge[layer] in range_read[l2id]: + node_layer_d[l2id] = layer + break + return node_layer_d + + +def _find_min_layer(node_layer_d_shared, node_ids_shared, node_layers_shared): + """ + `node_layer_d_shared`: DictProxy + + `node_ids_shared`: ListProxy + + `node_layers_shared`: ListProxy + + Due to parallelization, there will be multiple values for min_layer of a node. + We need to find the global min_layer after all multiprocesses return. + For eg: + At some indices p and q, there will be a node_id x + i.e. `node_ids_shared[p] == node_ids_shared[q]` + + and node_layers_shared[p] != node_layers_shared[q] + so we need: + `node_layer_d_shared[x] = min(node_layers_shared[p], node_layers_shared[q])` + """ + node_ids = np.concatenate(node_ids_shared) + layers = np.concatenate(node_layers_shared) + for i, node_id in enumerate(node_ids): + layer = node_layer_d_shared.get(node_id, layers[i]) + node_layer_d_shared[node_id] = min(layer, layers[i]) + + +def _read_atomic_chunk(cg: ChunkedGraph, chunk_coord, layers): + """ + read entire atomic chunk; all nodes and their relevant cross edges + filter out invalid nodes generated by failed tasks + """ + x, y, z = chunk_coord + child_col = attributes.Hierarchy.Child + range_read = cg.range_read_chunk( + cg.get_chunk_id(layer=2, x=x, y=y, z=z), + properties=[child_col] + + [attributes.Connectivity.AtomicCrossChunkEdge[l] for l in layers], + ) + + row_ids = [] + max_children_ids = [] + for row_id, row_data in range_read.items(): + row_ids.append(row_id) + max_children_ids.append(np.max(row_data[child_col][0].value)) + + row_ids = np.array(row_ids, dtype=basetypes.NODE_ID) + segment_ids = np.array([cg.get_segment_id(r_id) for r_id in row_ids]) + l2ids = filter_failed_node_ids(row_ids, segment_ids, max_children_ids) + return range_read, l2ids diff --git a/pychunkedgraph/ingest/create/parent_layer.py b/pychunkedgraph/ingest/create/parent_layer.py new file mode 100644 index 000000000..a3eb31d39 --- /dev/null +++ b/pychunkedgraph/ingest/create/parent_layer.py @@ -0,0 +1,266 @@ +# pylint: disable=invalid-name, missing-docstring, import-outside-toplevel, c-extension-no-member + +""" +Functions for creating parents in level 3 and above +""" + +import math +import datetime +import multiprocessing as mp +from typing import Optional +from typing import Sequence + +import fastremap +import numpy as np + +from pychunkedgraph import get_logger + +from ...graph import types, attributes, basetypes, serializers, get_valid_timestamp +from ...utils.general import chunked +from ...graph.utils import flatgraph +from ...graph.chunkedgraph import ChunkedGraph +from ...graph.edges.utils import concatenate_cross_edge_dicts +from ...graph.utils.generic import filter_failed_node_ids +from ...graph.chunks.hierarchy import get_children_chunk_coords +from .cross_edges import get_children_chunk_cross_edges +from .cross_edges import get_chunk_nodes_cross_edge_layer + +logger = get_logger(__name__) + + +def add_parent_chunk( + cg: ChunkedGraph, + layer_id: int, + coords: Sequence[int], + children_coords: Sequence[Sequence[int]] = np.array([]), + *, + time_stamp: Optional[datetime.datetime] = None, + n_processes: int = 1, +) -> None: + """``n_processes`` bounds every worker pool below; it is the pod's CPU + allocation (PCG_N_PROCESSES), never the node's core count.""" + if not children_coords.size: + children_coords = get_children_chunk_coords(cg.meta, layer_id, coords) + children_ids = _read_children_chunks(cg, layer_id, children_coords, n_processes) + cx_edges = get_children_chunk_cross_edges( + cg, layer_id, coords, n_processes=n_processes + ) + + node_layers = cg.get_chunk_layers(children_ids) + edge_layers = cg.get_chunk_layers(np.unique(cx_edges)) + assert np.all(node_layers < layer_id), "invalid node layers" + assert np.all(edge_layers < layer_id), "invalid edge layers" + + cx_edges = list(cx_edges) + cx_edges.extend(np.vstack([children_ids, children_ids]).T) # add self-edges + graph, _, _, graph_ids = flatgraph.build_gt_graph(cx_edges, make_directed=True) + raw_ccs = flatgraph.connected_components(graph) # connected components with indices + connected_components = [graph_ids[cc] for cc in raw_ccs] + + logger.note( + f"L{layer_id} chunk {tuple(map(int, coords))}: nodes={len(connected_components):,} " + f"cx_edges={len(cx_edges):,}" + ) + + ts = get_valid_timestamp(time_stamp) + _write_connected_components( + cg, layer_id, coords, connected_components, ts, n_processes + ) + + # Stamp the post-ingest boundary meshing reads to split initial from edited roots. + # ts is the explicit cell timestamp shared by every root just written; +500ms (the + # same guard get_earliest_timestamp puts below the first op) lifts the boundary + # strictly above them. + if layer_id == cg.meta.layer_count: + boundary = ts + datetime.timedelta(milliseconds=500) + cg.meta.custom_data["earliest_ts"] = boundary.isoformat() + cg.update_meta(cg.meta, overwrite=True) + + +def _read_children_chunks( + cg: ChunkedGraph, layer_id, children_coords, n_processes: int = 1 +): + if n_processes <= 1: + children_ids = [types.empty_1d] + for child_coord in children_coords: + children_ids.append(_read_chunk([], cg, layer_id - 1, child_coord)) + return np.concatenate(children_ids).astype(basetypes.NODE_ID) + + with mp.Manager() as manager: + children_ids_shared = manager.list() + multi_args = [] + for child_coord in children_coords: + multi_args.append( + ( + children_ids_shared, + cg.get_serialized_info(), + layer_id - 1, + child_coord, + ) + ) + with mp.Pool(processes=min(len(multi_args), n_processes)) as pool: + pool.map(_read_chunk_helper, multi_args) + return np.concatenate(children_ids_shared).astype(basetypes.NODE_ID) + + +def _read_chunk_helper(args): + children_ids_shared, cg_info, layer_id, chunk_coord = args + # Re-raise as a bare RuntimeError: the original may hold an unpicklable client + # handle, which the pool would surface as MaybeEncodingError, losing the cause. + try: + cg = ChunkedGraph(**cg_info) + _read_chunk(children_ids_shared, cg, layer_id, chunk_coord) + except Exception as exc: + raise RuntimeError( + f"_read_chunk failed at layer {layer_id} chunk " + f"{tuple(map(int, chunk_coord))}: {exc!r}" + ) from None + + +def _read_chunk(children_ids_shared, cg: ChunkedGraph, layer_id: int, chunk_coord): + x, y, z = chunk_coord + range_read = cg.range_read_chunk( + cg.get_chunk_id(layer=layer_id, x=x, y=y, z=z), + properties=attributes.Hierarchy.Child, + ) + row_ids = [] + max_children_ids = [] + for row_id, row_data in range_read.items(): + row_ids.append(row_id) + max_children_ids.append(np.max(row_data[0].value)) + row_ids = np.array(row_ids, dtype=basetypes.NODE_ID) + segment_ids = np.array([cg.get_segment_id(r_id) for r_id in row_ids]) + + row_ids = filter_failed_node_ids(row_ids, segment_ids, max_children_ids) + children_ids_shared.append(row_ids) + return row_ids + + +def _write_connected_components( + cg, layer, pcoords, components, time_stamp, n_processes: int = 1 +): + if len(components) == 0: + return + + node_layer_d = {} + if layer < cg.meta.layer_count: + node_layer_d = get_chunk_nodes_cross_edge_layer(cg, layer, pcoords, n_processes) + + if n_processes <= 1: + _write(cg, layer, pcoords, components, node_layer_d, time_stamp, False) + return + + task_size = int(math.ceil(len(components) / n_processes / 10)) + chunked_ccs = chunked(components, task_size) + cg_info = cg.get_serialized_info() + multi_args = [] + for ccs in chunked_ccs: + args = (cg_info, layer, pcoords, ccs, node_layer_d, time_stamp) + multi_args.append(args) + with mp.Pool(processes=min(len(multi_args), n_processes)) as pool: + pool.map(_write_components_helper, multi_args) + + +def _write_components_helper(args): + cg_info, layer, pcoords, ccs, node_layer_d, time_stamp = args + # See _read_chunk_helper: keep the failure picklable. + try: + cg = ChunkedGraph(**cg_info) + _write(cg, layer, pcoords, ccs, node_layer_d, time_stamp) + except Exception as exc: + raise RuntimeError( + f"_write failed at layer {layer} chunk {tuple(map(int, pcoords))}: {exc!r}" + ) from None + + +def _children_rows( + cg: ChunkedGraph, parent_id, children: Sequence, cx_edges_d: dict, time_stamp +): + """ + Update children rows to point to the parent_id, collect cached children + cross chunk edges to lift and update parent cross chunk edges. + Returns list of mutations to children and list of children cross edges. + """ + rows = [] + children_cx_edges = [] + children_layers = cg.get_chunk_layers(children) + for child, node_layer in zip(children, children_layers): + node_layer = cg.get_chunk_layer(child) + row_id = serializers.serialize_uint64(child) + val_dict = {attributes.Hierarchy.Parent: parent_id} + node_cx_edges_d = cx_edges_d.get(child, {}) + if not node_cx_edges_d: + rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp)) + continue + for layer in range(node_layer, cg.meta.layer_count): + if not layer in node_cx_edges_d: + continue + layer_edges = node_cx_edges_d[layer] + nodes = np.unique(layer_edges) + parents = cg.get_roots(nodes, stop_layer=node_layer, ceil=False) + edge_parents_d = dict(zip(nodes, parents)) + layer_edges = fastremap.remap( + layer_edges, edge_parents_d, preserve_missing_labels=True + ) + layer_edges = np.unique(layer_edges, axis=0) + col = attributes.Connectivity.CrossChunkEdge[layer] + val_dict[col] = layer_edges + node_cx_edges_d[layer] = layer_edges + children_cx_edges.append(node_cx_edges_d) + rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp)) + return rows, children_cx_edges + + +def _write( + cg: ChunkedGraph, + layer_id, + parent_coords, + components, + node_layer_d, + ts, + use_threads=True, +): + parent_layers = range(layer_id, cg.meta.layer_count + 1) + cc_connections = {l: [] for l in parent_layers} + for node_ids in components: + layer = layer_id + if len(node_ids) == 1: + layer = node_layer_d.get(node_ids[0], cg.meta.layer_count) + cc_connections[layer].append(node_ids) + + rows = [] + x, y, z = parent_coords + parent_chunk_id = cg.get_chunk_id(layer=layer_id, x=x, y=y, z=z) + parent_chunk_id_dict = cg.get_parent_chunk_id_dict(parent_chunk_id) + for parent_layer in parent_layers: + if len(cc_connections[parent_layer]) == 0: + continue + parent_chunk_id = parent_chunk_id_dict[parent_layer] + reserved_parent_ids = cg.id_client.create_node_ids( + parent_chunk_id, + size=len(cc_connections[parent_layer]), + root_chunk=parent_layer == cg.meta.layer_count and use_threads, + ) + for i_cc, children in enumerate(cc_connections[parent_layer]): + parent = reserved_parent_ids[i_cc] + if layer_id == 3: + # when layer 3 is being processed, children chunks are at layer 2 + # layer 2 chunks at this time will only have atomic cross edges + cx_edges_d = cg.get_atomic_cross_edges(children) + else: + cx_edges_d = cg.get_cross_chunk_edges(children, raw_only=True) + _rows, cx_edges = _children_rows(cg, parent, children, cx_edges_d, ts) + rows.extend(_rows) + row_id = serializers.serialize_uint64(parent) + val_dict = {attributes.Hierarchy.Child: children} + parent_cx_edges_d = concatenate_cross_edge_dicts(cx_edges, unique=True) + for layer in range(parent_layer, cg.meta.layer_count): + if not layer in parent_cx_edges_d: + continue + col = attributes.Connectivity.CrossChunkEdge[layer] + val_dict[col] = parent_cx_edges_d[layer] + rows.append(cg.client.mutate_row(row_id, val_dict, ts)) + if len(rows) > 100000: + cg.client.write(rows) + rows = [] + cg.client.write(rows) diff --git a/pychunkedgraph/ingest/manager.py b/pychunkedgraph/ingest/manager.py index f5f870810..566558e05 100644 --- a/pychunkedgraph/ingest/manager.py +++ b/pychunkedgraph/ingest/manager.py @@ -1,3 +1,5 @@ +# pylint: disable=invalid-name, missing-docstring + import pickle from . import IngestConfig @@ -9,13 +11,24 @@ class IngestionManager: - def __init__(self, config: IngestConfig, chunkedgraph_meta: ChunkedGraphMeta): + def __init__( + self, + config: IngestConfig, + chunkedgraph_meta: ChunkedGraphMeta, + ocdbt_config: dict = None, + _from_pickle: bool = False, + ): self._config = config self._chunkedgraph_meta = chunkedgraph_meta self._cg = None self._redis = None self._task_queues = {} - self.redis # initiate and cache info + self._from_pickle = _from_pickle + self.ocdbt_config = ocdbt_config or {} + + if not _from_pickle: + # initiate redis and store serialized state + self.redis # pylint: disable=pointless-statement @property def config(self): @@ -36,18 +49,46 @@ def redis(self): if self._redis is not None: return self._redis self._redis = get_redis_connection() - self._redis.set(r_keys.INGESTION_MANAGER, self.serialized(pickled=True)) + if not self._from_pickle: + self._redis.set(r_keys.INGESTION_MANAGER, self.serialized(pickled=True)) return self._redis + @property + def ocdbt_seg(self) -> bool: + return bool(self.ocdbt_config.get("enabled")) + + @property + def ocdbt_populate_base(self) -> bool: + return bool(self.ocdbt_config.get("populate_base")) + + @property + def ocdbt_populate_layer(self) -> int: + return int(self.ocdbt_config.get("populate_layer", 3)) + + def is_ocdbt_populate_layer(self, layer: int) -> bool: + """True iff OCDBT is enabled, base-populate is on, AND the given + layer matches the configured populate layer. Single guard for any + code that branches on 'should this layer touch OCDBT?'. + """ + return ( + self.ocdbt_seg + and self.ocdbt_populate_base + and layer == self.ocdbt_populate_layer + ) + def serialized(self, pickled=False): - params = {"config": self._config, "chunkedgraph_meta": self._chunkedgraph_meta} + params = { + "config": self._config, + "chunkedgraph_meta": self._chunkedgraph_meta, + "ocdbt_config": self.ocdbt_config, + } if pickled: return pickle.dumps(params) return params @classmethod def from_pickle(cls, serialized_info): - return cls(**pickle.loads(serialized_info)) + return cls(**pickle.loads(serialized_info), _from_pickle=True) def get_task_queue(self, q_name): if q_name in self._task_queues: diff --git a/pychunkedgraph/ingest/ocdbt.py b/pychunkedgraph/ingest/ocdbt.py new file mode 100644 index 000000000..a8b10666f --- /dev/null +++ b/pychunkedgraph/ingest/ocdbt.py @@ -0,0 +1,128 @@ +"""OCDBT-specific ingest helpers. + +Single home for everything OCDBT-related at the ingest layer: + * coordinator-server lifecycle (`coordinator`) + * per-chunk populate task (`populate_chunk`), used from `create_parent_chunk` + * shared base setup (`setup_base`), used by both ingest and upgrade CLIs +""" + +from contextlib import contextmanager +from os import environ + +import tensorstore as ts + +from pychunkedgraph import get_logger + +from ..graph.ocdbt import ( + OcdbtConfig, + _layer_bbox, + base_exists, + copy_ws_bbox_multiscale, + create_base_ocdbt, + fork_base_manifest, + mark_chunk_populated, + open_base_ocdbt, + read_populate_meta, + write_populate_meta, +) + +logger = get_logger(__name__) + +_COORD_HOST_KEY = "OCDBT_COORDINATOR_HOST" +_COORD_PORT_KEY = "OCDBT_COORDINATOR_PORT" + + +@contextmanager +def coordinator(redis): + """Start a ``DistributedCoordinatorServer`` and advertise its address in + Redis so parallel populate workers route every OCDBT commit through this + one server — no manifest-CAS races, no orphan ``d/`` files. + + The server lives as long as the ``with`` block does; on exit the Redis + advertisement is cleared so a stale address can't outlive the server. + Caller blocks inside the ``with`` body (e.g. ``while True: sleep(60)``) + to keep the server reference alive across the populate phase. + """ + server = ts.ocdbt.DistributedCoordinatorServer() + host = environ.get("MY_POD_IP", "localhost") + redis.set(_COORD_HOST_KEY, host) + redis.set(_COORD_PORT_KEY, str(server.port)) + logger.note(f"OCDBT Coordinator listening at {host}:{server.port}") + try: + yield server + finally: + redis.delete(_COORD_HOST_KEY, _COORD_PORT_KEY) + logger.note("OCDBT Coordinator advertisement cleared.") + + +def get_coordinator_address(redis) -> str: + """Return the advertised ``"host:port"`` for the OCDBT coordinator. + + The address goes into the OCDBT kvstore spec's ``coordinator`` field — + the only routing knob tensorstore actually honors (verified against + the tensorstore binary; ``OCDBT_COORDINATOR_HOST/PORT`` env vars are + not consulted). + + Distributed callers MUST go through this getter so the populate fails + loudly when the coordinator isn't advertised — uncoordinated parallel + commits race the shared manifest and leak orphan ``d/`` files, the + exact bug this code exists to prevent. + """ + host = redis.get(_COORD_HOST_KEY) + port = redis.get(_COORD_PORT_KEY) + if not host or not port: + raise RuntimeError( + "OCDBT coordinator address not advertised in Redis " + f"({_COORD_HOST_KEY}/{_COORD_PORT_KEY} unset). " + "Run `flask ingest layer N` (with N == ocdbt_populate_layer) to " + "start the coordinator before queuing populate workers." + ) + return f"{host.decode()}:{port.decode()}" + + +def populate_chunk( + imanager, ws: str, layer: int, coords, coordinator_address: str | None = None +) -> None: + """One LN parent-layer task's OCDBT populate. + + When ``coordinator_address`` is set, every commit routes through that + server (mandatory for distributed workers — see ``get_coordinator_address``). + Single-process callers (notebooks, local one-off runs) can omit it and + write directly; safe as long as no other writer is committing concurrently. + + Copies the base-resolution bbox at every scale under one atomic + transaction and records the per-chunk completion marker. + """ + cfg = OcdbtConfig.from_dict(imanager.ocdbt_config) + src_list, dst_list, resolutions = open_base_ocdbt( + ws, cfg, coordinator_address=coordinator_address + ) + lo, hi = _layer_bbox(imanager.cg.meta, layer, coords) + coord_str = "_".join(str(int(c)) for c in coords) + dump_tag = f"{imanager.cg.meta.graph_id}/L{layer}/{coord_str}" + logger.note(f"L{layer} OCDBT populate {tuple(int(c) for c in coords)}") + copy_ws_bbox_multiscale(src_list, dst_list, resolutions, lo, hi, dump_tag=dump_tag) + mark_chunk_populated(ws, layer, coords) + + +def setup_base(cg, ocdbt_cfg: OcdbtConfig) -> OcdbtConfig: + """Idempotent OCDBT base + fork setup, shared by ingest and upgrade. + + Creates the base if missing; reconciles the yaml/CLI-supplied config + with the on-disk populate_meta (info-file wins per + ``OcdbtConfig.resolve``); persists the resolved config to + ``cg.meta.custom_data["ocdbt_config"]``; forks the manifest for this + CG. Returns the resolved OcdbtConfig. To wipe and start over, use + ``gcloud storage rm -r gs:///ocdbt/`` before invoking. + """ + ws = cg.meta.data_source.WATERSHED + if not base_exists(ws): + create_base_ocdbt(ws, ocdbt_cfg) + info = read_populate_meta(ws) + resolved = OcdbtConfig.resolve(ocdbt_cfg.to_dict(), info) + if resolved.populate_base: + write_populate_meta(ws, resolved.to_dict()) + cg.meta.custom_data["ocdbt_config"] = resolved.to_dict() + cg.update_meta(cg.meta, overwrite=True) + fork_base_manifest(ws, cg.meta.graph_id) + return resolved diff --git a/pychunkedgraph/ingest/ran_agglomeration.py b/pychunkedgraph/ingest/ran_agglomeration.py index 7c4af51f7..c386f88e0 100644 --- a/pychunkedgraph/ingest/ran_agglomeration.py +++ b/pychunkedgraph/ingest/ran_agglomeration.py @@ -5,10 +5,7 @@ from collections import defaultdict from itertools import product -from typing import Dict -from typing import Iterable -from typing import Tuple -from typing import Union +from typing import Dict, Iterable, Tuple, Union from binascii import crc32 @@ -22,9 +19,8 @@ from .utils import postprocess_edge_data from ..io.edges import put_chunk_edges from ..io.components import put_chunk_components -from ..graph.utils import basetypes -from ..graph.edges import Edges -from ..graph.edges import EDGE_TYPES +from ..graph import basetypes +from ..graph.edges import EDGE_TYPES, Edges from ..graph.types import empty_2d from ..graph.chunks.utils import get_chunk_id @@ -318,7 +314,9 @@ def get_active_edges(edges_d, mapping): if edge_type == EDGE_TYPES.in_chunk: pseudo_isolated_ids.append(edges.node_ids2) - return chunk_edges_active, np.unique(np.concatenate(pseudo_isolated_ids)) + return chunk_edges_active, np.unique( + np.concatenate(pseudo_isolated_ids).astype(basetypes.NODE_ID) + ) def define_active_edges(edge_dict, mapping) -> Union[Dict, np.ndarray]: @@ -384,7 +382,7 @@ def read_raw_agglomeration_data(imanager: IngestionManager, chunk_coord: np.ndar edges_list = _read_agg_files(filenames, chunk_ids, path) G = nx.Graph() - G.add_edges_from(np.concatenate(edges_list)) + G.add_edges_from(np.concatenate(edges_list).astype(basetypes.NODE_ID)) mapping = {} components = list(nx.connected_components(G)) for i_cc, cc in enumerate(components): diff --git a/pychunkedgraph/ingest/rq_cli.py b/pychunkedgraph/ingest/rq_cli.py index 27b9c865d..62367860b 100644 --- a/pychunkedgraph/ingest/rq_cli.py +++ b/pychunkedgraph/ingest/rq_cli.py @@ -1,47 +1,25 @@ +# pylint: disable=invalid-name, missing-function-docstring + """ cli for redis jobs """ -import os + import sys import click -from redis import Redis from rq import Queue -from rq import Worker -from rq.worker import WorkerStatus from rq.job import Job from rq.exceptions import InvalidJobOperationError from rq.exceptions import NoSuchJobError from rq.registry import StartedJobRegistry from rq.registry import FailedJobRegistry -from flask import current_app from flask.cli import AppGroup -from ..utils.redis import REDIS_HOST -from ..utils.redis import REDIS_PORT -from ..utils.redis import REDIS_PASSWORD - +from ..utils.redis import get_redis_connection # rq extended rq_cli = AppGroup("rq") -connection = Redis(host=REDIS_HOST, port=REDIS_PORT, db=0, password=REDIS_PASSWORD) - - -@rq_cli.command("status") -@click.argument("queues", nargs=-1, type=str) -@click.option("--show-busy", is_flag=True) -def get_status(queues, show_busy): - print("NOTE: Use --show-busy to display count of non idle workers\n") - for queue in queues: - q = Queue(queue, connection=connection) - print(f"Queue name \t: {queue}") - print(f"Jobs queued \t: {len(q)}") - print(f"Workers total \t: {Worker.count(queue=q)}") - if show_busy: - workers = Worker.all(queue=q) - count = sum([worker.get_state() == WorkerStatus.BUSY for worker in workers]) - print(f"Workers busy \t: {count}") - print(f"Jobs failed \t: {q.failed_job_registry.count}\n") +connection = get_redis_connection() @rq_cli.command("failed") @@ -129,9 +107,14 @@ def clean_start_registry(queue): def clear_failed_registry(queue): failed_job_registry = FailedJobRegistry(queue, connection=connection) job_ids = failed_job_registry.get_job_ids() + count = 0 for job_id in job_ids: - failed_job_registry.remove(job_id, delete_job=True) - print(f"Deleted {len(job_ids)} jobs from the failed job registry.") + try: + failed_job_registry.remove(job_id, delete_job=True) + count += 1 + except Exception: + ... + print(f"Deleted {count} jobs from the failed job registry.") def init_rq_cmds(app): diff --git a/pychunkedgraph/ingest/simple_tests.py b/pychunkedgraph/ingest/simple_tests.py new file mode 100644 index 000000000..61292eb85 --- /dev/null +++ b/pychunkedgraph/ingest/simple_tests.py @@ -0,0 +1,180 @@ +# pylint: disable=invalid-name, missing-function-docstring, broad-exception-caught + +""" +Some sanity tests to ensure chunkedgraph was created properly. +""" + +from datetime import datetime, timezone +import numpy as np + +from pychunkedgraph.graph import ChunkedGraph +from pychunkedgraph.graph import attributes + + +def family(cg: ChunkedGraph): + np.random.seed(42) + n_chunks = 100 + n_segments_per_chunk = 200 + timestamp = datetime.now(timezone.utc) + + node_ids = [] + for layer in range(2, cg.meta.layer_count - 1): + for _ in range(n_chunks): + c_x = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][0]) + c_y = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][1]) + c_z = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][2]) + chunk_id = cg.get_chunk_id(layer=layer, x=c_x, y=c_y, z=c_z) + max_segment_id = cg.get_segment_id(cg.id_client.get_max_node_id(chunk_id)) + if max_segment_id < 10: + continue + + segment_ids = np.random.randint(1, max_segment_id, n_segments_per_chunk) + for segment_id in segment_ids: + node_ids.append( + cg.get_node_id(np.uint64(segment_id), np.uint64(chunk_id)) + ) + + rows = cg.client.read_nodes( + node_ids=node_ids, end_time=timestamp, properties=attributes.Hierarchy.Parent + ) + valid_node_ids = [] + non_valid_node_ids = [] + for k in rows.keys(): + if len(rows[k]) > 0: + valid_node_ids.append(k) + else: + non_valid_node_ids.append(k) + + parents = cg.get_parents(valid_node_ids, time_stamp=timestamp) + children_dict = cg.get_children(parents) + for child, parent in zip(valid_node_ids, parents): + assert child in children_dict[parent] + print("success") + + +def existence(cg: ChunkedGraph): + np.random.seed(42) + layer = 2 + n_chunks = 100 + n_segments_per_chunk = 200 + timestamp = datetime.now(timezone.utc) + node_ids = [] + for _ in range(n_chunks): + c_x = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][0]) + c_y = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][1]) + c_z = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][2]) + chunk_id = cg.get_chunk_id(layer=layer, x=c_x, y=c_y, z=c_z) + max_segment_id = cg.get_segment_id(cg.id_client.get_max_node_id(chunk_id)) + if max_segment_id < 10: + continue + + segment_ids = np.random.randint(1, max_segment_id, n_segments_per_chunk) + for segment_id in segment_ids: + node_ids.append(cg.get_node_id(np.uint64(segment_id), np.uint64(chunk_id))) + + rows = cg.client.read_nodes( + node_ids=node_ids, end_time=timestamp, properties=attributes.Hierarchy.Parent + ) + valid_node_ids = [] + non_valid_node_ids = [] + for k in rows.keys(): + if len(rows[k]) > 0: + valid_node_ids.append(k) + else: + non_valid_node_ids.append(k) + + roots = [] + try: + roots = cg.get_roots(valid_node_ids) + assert len(roots) == len(valid_node_ids) + print("success") + except Exception as e: + print(f"Something went wrong: {e}") + print("At least one node failed. Checking nodes one by one:") + + if len(roots) != len(valid_node_ids): + log_dict = {} + success_dict = {} + for node_id in valid_node_ids: + try: + _ = cg.get_root(node_id, time_stamp=timestamp) + print(f"Success: {node_id} from chunk {cg.get_chunk_id(node_id)}") + success_dict[node_id] = True + except Exception as e: + print(f"{node_id} - chunk {cg.get_chunk_id(node_id)} failed: {e}") + success_dict[node_id] = False + t_id = node_id + while t_id is not None: + last_working_chunk = cg.get_chunk_id(t_id) + t_id = cg.get_parent(t_id) + + layer = cg.get_chunk_layer(last_working_chunk) + print(f"Failed on layer {layer} in chunk {last_working_chunk}") + log_dict[node_id] = last_working_chunk + if log_dict: # diagnostics above are informational; the suite must still fail + raise AssertionError(f"{len(log_dict)} nodes failed the existence check") + + +def cross_edges(cg: ChunkedGraph): + np.random.seed(42) + layer = 2 + n_chunks = 10 + n_segments_per_chunk = 200 + timestamp = datetime.now(timezone.utc) + node_ids = [] + for _ in range(n_chunks): + c_x = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][0]) + c_y = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][1]) + c_z = np.random.randint(0, cg.meta.layer_chunk_bounds[layer][2]) + chunk_id = cg.get_chunk_id(layer=layer, x=c_x, y=c_y, z=c_z) + max_segment_id = cg.get_segment_id(cg.id_client.get_max_node_id(chunk_id)) + if max_segment_id < 10: + continue + + segment_ids = np.random.randint(1, max_segment_id, n_segments_per_chunk) + for segment_id in segment_ids: + node_ids.append(cg.get_node_id(np.uint64(segment_id), np.uint64(chunk_id))) + + rows = cg.client.read_nodes( + node_ids=node_ids, end_time=timestamp, properties=attributes.Hierarchy.Parent + ) + valid_node_ids = [] + non_valid_node_ids = [] + for k in rows.keys(): + if len(rows[k]) > 0: + valid_node_ids.append(k) + else: + non_valid_node_ids.append(k) + + cc_edges = cg.get_atomic_cross_edges(valid_node_ids) + cc_ids = np.unique( + np.concatenate( + [ + np.concatenate(list(v.values())) + for v in list(cc_edges.values()) + if len(v.values()) + ] + ) + ) + + roots = cg.get_roots(cc_ids) + root_dict = dict(zip(cc_ids, roots)) + root_dict_vec = np.vectorize(root_dict.get) + + for k in cc_edges: + if len(cc_edges[k]) == 0: + continue + local_ids = np.unique(np.concatenate(list(cc_edges[k].values()))) + assert len(np.unique(root_dict_vec(local_ids))) + print("success") + + +def run_all(cg: ChunkedGraph): + print("Running family tests:") + family(cg) + + print("\nRunning existence tests:") + existence(cg) + + print("\nRunning cross_edges tests:") + cross_edges(cg) diff --git a/pychunkedgraph/ingest/upgrade/__init__.py b/pychunkedgraph/ingest/upgrade/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pychunkedgraph/ingest/upgrade/atomic_layer.py b/pychunkedgraph/ingest/upgrade/atomic_layer.py new file mode 100644 index 000000000..9cc34aefa --- /dev/null +++ b/pychunkedgraph/ingest/upgrade/atomic_layer.py @@ -0,0 +1,162 @@ +# pylint: disable=invalid-name, missing-docstring, c-extension-no-member + +from collections import defaultdict +from datetime import datetime, timedelta, timezone +import time + +from pychunkedgraph import get_logger + +logger = get_logger(__name__) +from copy import copy + +import fastremap +import numpy as np +from pychunkedgraph.graph import ChunkedGraph, types +from pychunkedgraph.graph import attributes, serializers +from pychunkedgraph.graph.utils.generic import get_parents_at_timestamp + +from .utils import fix_corrupt_nodes, get_end_timestamps, get_parent_timestamps + +CHILDREN = {} + + +def update_cross_edges( + cg: ChunkedGraph, + node, + cx_edges_d: dict, + node_ts, + node_end_ts, + timestamps_map: defaultdict[int, set], + parents_ts_map: defaultdict[int, dict], +) -> list: + """ + Helper function to update a single L2 ID. + Returns a list of mutations with given timestamps. + """ + rows = [] + edges = np.concatenate(list(cx_edges_d.values())) + partners = np.unique(edges[:, 1]) + + timestamps = copy(timestamps_map[node]) + for partner in partners: + timestamps.update(timestamps_map[partner]) + + node_end_ts = node_end_ts or datetime.now(timezone.utc) + for ts in sorted(timestamps): + if ts < node_ts: + continue + if ts > node_end_ts: + break + + val_dict = {} + parents, _ = get_parents_at_timestamp(partners, parents_ts_map, ts) + edge_parents_d = dict(zip(partners, parents)) + for layer, layer_edges in cx_edges_d.items(): + layer_edges = fastremap.remap( + layer_edges, edge_parents_d, preserve_missing_labels=True + ) + layer_edges[:, 0] = node + layer_edges = np.unique(layer_edges, axis=0) + col = attributes.Connectivity.CrossChunkEdge[layer] + val_dict[col] = layer_edges + row_id = serializers.serialize_uint64(node) + rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp=ts)) + return rows + + +def update_nodes(cg: ChunkedGraph, nodes, nodes_ts, children_map=None) -> list: + start = time.time() + if children_map is None: + children_map = CHILDREN + end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, children_map, layer=2) + + cx_edges_d = cg.get_atomic_cross_edges(nodes) + all_cx_edges = [types.empty_2d] + for _cx_edges_d in cx_edges_d.values(): + if _cx_edges_d: + all_cx_edges.append(np.concatenate(list(_cx_edges_d.values()))) + all_partners = np.unique(np.concatenate(all_cx_edges)[:, 1]) + timestamps_d = get_parent_timestamps(cg, np.concatenate([nodes, all_partners])) + + parents_ts_map = defaultdict(dict) + all_parents = cg.get_parents(all_partners, current=False) + for partner, parents in zip(all_partners, all_parents): + for parent, ts in parents: + parents_ts_map[partner][ts] = parent + logger.note(f"update_nodes init {len(nodes)}: {time.time() - start}") + + rows = [] + skipped = [] + for node, node_ts, end_ts in zip(nodes, nodes_ts, end_timestamps): + is_stale = end_ts is not None + _cx_edges_d = cx_edges_d.get(node, {}) + if is_stale: + end_ts -= timedelta(milliseconds=1) + row_id = serializers.serialize_uint64(node) + val_dict = {attributes.Hierarchy.StaleTimeStamp: 0} + rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp=end_ts)) + + if not _cx_edges_d: + skipped.append(node) + continue + + _rows = update_cross_edges( + cg, node, _cx_edges_d, node_ts, end_ts, timestamps_d, parents_ts_map + ) + rows.extend(_rows) + parents = cg.get_roots(skipped) + layers = cg.get_chunk_layers(parents) + assert np.all(layers == cg.meta.layer_count) + return rows + + +def update_chunk(cg: ChunkedGraph, chunk_coords: list[int], clean: bool = False): + """ + Iterate over all L2 IDs in a chunk and update their cross chunk edges, + within the periods they were valid/active. + """ + global CHILDREN + + start = time.time() + x, y, z = chunk_coords + chunk_id = cg.get_chunk_id(layer=2, x=x, y=y, z=z) + rr = cg.range_read_chunk(chunk_id) + + nodes = [] + nodes_ts = [] + earliest_ts = datetime.fromisoformat(cg.meta.custom_data["earliest_ts"]) + + corrupt_nodes = [] + for k, v in rr.items(): + try: + CHILDREN[k] = v[attributes.Hierarchy.Child][0].value + ts = v[attributes.Hierarchy.Child][0].timestamp + _ = v[attributes.Hierarchy.Parent] + nodes.append(k) + nodes_ts.append(earliest_ts if ts < earliest_ts else ts) + except KeyError: + # ignore invalid nodes from failed ingest tasks, w/o parent column entry + # retain invalid nodes from edits to fix the hierarchy + if ts > earliest_ts: + corrupt_nodes.append(k) + + if clean: + logger.note(f"found {len(corrupt_nodes)} corrupt nodes {corrupt_nodes[:3]}...") + fix_corrupt_nodes(cg, corrupt_nodes, CHILDREN) + return + + # Set max node ID for the L1 chunk (needed for SV splitting to create new IDs) + l1_chunk_id = cg.get_chunk_id(layer=1, x=x, y=y, z=z) + if CHILDREN: + all_svs = np.concatenate(list(CHILDREN.values())) + cg.id_client.set_max_node_id(l1_chunk_id, np.max(all_svs)) + + cg.copy_fake_edges(chunk_id) + if len(nodes) == 0: + return + + logger.note(f"processing {len(nodes)} nodes.") + assert len(CHILDREN) > 0, (nodes, CHILDREN) + rows = update_nodes(cg, nodes, nodes_ts) + cg.client.write(rows) + logger.note(f"mutations: {len(rows)}, time: {time.time() - start}") diff --git a/pychunkedgraph/ingest/upgrade/parent_layer.py b/pychunkedgraph/ingest/upgrade/parent_layer.py new file mode 100644 index 000000000..d2fef574e --- /dev/null +++ b/pychunkedgraph/ingest/upgrade/parent_layer.py @@ -0,0 +1,274 @@ +# pylint: disable=invalid-name, missing-docstring, c-extension-no-member + +from math import ceil +import bisect, random, time, os, gc + +from pychunkedgraph import get_logger + +logger = get_logger(__name__) +import multiprocessing as mp +from collections import defaultdict +from datetime import datetime, timezone + +import fastremap +import numpy as np +from tqdm import tqdm +from cachetools import LRUCache + +from pychunkedgraph.graph import ChunkedGraph +from pychunkedgraph.graph.edges import stale, get_latest_edges_wrapper +from pychunkedgraph.graph import attributes, serializers, basetypes +from pychunkedgraph.graph.types import empty_2d +from pychunkedgraph.utils.general import chunked + +from .utils import fix_corrupt_nodes, get_end_timestamps, get_parent_timestamps + +CHILDREN = {} +CX_EDGES = {} +CG: ChunkedGraph = None +PARENT_CACHE_LIMIT = int(os.environ.get("PARENT_CACHE_LIMIT", 256)) * 1024 + + +def _populate_nodes_and_children( + cg: ChunkedGraph, chunk_id: np.uint64, nodes: list = None +) -> dict: + global CHILDREN + if nodes: + children_map = cg.get_children(nodes) + for k, v in children_map.items(): + if len(v): + CHILDREN[k] = v + return + response = cg.range_read_chunk(chunk_id, properties=attributes.Hierarchy.Child) + for k, v in response.items(): + CHILDREN[k] = v[0].value + + +def _get_cx_edges_at_timestamp(node, response, ts): + result = defaultdict(list) + for child in CHILDREN[node]: + if child not in response: + continue + for key, cells in response[child].items(): + # cells are sorted in descending order of timestamps + asc_ts = [c.timestamp for c in reversed(cells)] + k = bisect.bisect_right(asc_ts, ts) - 1 + if k >= 0: + idx = len(cells) - 1 - k + try: + result[key.index].append(cells[idx].value) + except IndexError as e: + logger.error(f"{k}, {idx}, {len(cells)}, {asc_ts}") + raise IndexError from e + for layer, edges in result.items(): + result[layer] = np.concatenate(edges) + return result + + +def _populate_cx_edges_with_timestamps( + cg: ChunkedGraph, layer: int, nodes: list, nodes_ts: list, clean: bool = False +): + """ + Collect timestamps of edits from children, since we use the same timestamp + for all IDs involved in an edit, we can use the timestamps of + when cross edges of children were updated. + """ + # this data is not needed for clean tasks + if clean: + return + + start = time.time() + global CX_EDGES + attrs = [ + attributes.Connectivity.CrossChunkEdge[l] + for l in range(layer, cg.meta.layer_count) + ] + all_children = np.concatenate(list(CHILDREN.values())) + response = cg.client.read_nodes(node_ids=all_children, properties=attrs) + timestamps_d = get_parent_timestamps(cg, nodes) + end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, CHILDREN, layer=layer) + logger.note(f"_populate_cx_edges_with_timestamps init: {time.time() - start}") + + start = time.time() + partners_map = {} + for node, node_ts in zip(nodes, nodes_ts): + CX_EDGES[node] = {} + cx_edges_d_node_ts = _get_cx_edges_at_timestamp(node, response, node_ts) + edges = np.concatenate([empty_2d] + list(cx_edges_d_node_ts.values())) + partners_map[node] = edges[:, 1] + CX_EDGES[node][node_ts] = cx_edges_d_node_ts + + partners = np.unique(np.concatenate([*partners_map.values()])) + partner_parent_ts_d = get_parent_timestamps(cg, partners) + logger.note(f"get partners timestamps init: {time.time() - start}") + + rows = [] + for node, node_ts, node_end_ts in zip(nodes, nodes_ts, end_timestamps): + timestamps = timestamps_d[node] + for partner in partners_map[node]: + timestamps.update(partner_parent_ts_d[partner]) + + is_stale = node_end_ts is not None + node_end_ts = node_end_ts or datetime.now(timezone.utc) + for ts in sorted(timestamps): + if ts > node_end_ts: + break + CX_EDGES[node][ts] = _get_cx_edges_at_timestamp(node, response, ts) + + if is_stale: + row_id = serializers.serialize_uint64(node) + val_dict = {attributes.Hierarchy.StaleTimeStamp: 0} + rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp=node_end_ts)) + cg.client.write(rows) + + +def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts) -> list: + """ + Helper function to update a single ID. + Returns a list of mutations with timestamps. + """ + rows = [] + row_id = serializers.serialize_uint64(node) + for ts, edges_d in CX_EDGES[node].items(): + if ts < node_ts: + continue + edges_d, _nodes = get_latest_edges_wrapper(cg, edges_d, parent_ts=ts) + if _nodes.size == 0: + continue + + parents = cg.get_roots(_nodes, time_stamp=ts, stop_layer=layer, ceil=False) + edge_parents_d = dict(zip(_nodes, parents)) + val_dict = {} + for _layer, layer_edges in edges_d.items(): + layer_edges = fastremap.remap( + layer_edges, edge_parents_d, preserve_missing_labels=True + ) + layer_edges[:, 0] = node + layer_edges = np.unique(layer_edges, axis=0) + col = attributes.Connectivity.CrossChunkEdge[_layer] + val_dict[col] = layer_edges + rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp=ts)) + return rows + + +def _update_cross_edges_helper(args): + global CG + stale.PARENTS_CACHE = LRUCache(PARENT_CACHE_LIMIT) + stale.CHILDREN_CACHE = LRUCache(1 * 1024) + cg_info, layer, nodes, nodes_ts, clean = args + + if CG is None: + CG = ChunkedGraph(**cg_info) + cg = CG + parents = cg.get_parents(nodes, fail_to_zero=True) + + tasks = [] + corrupt_nodes = [] + earliest_ts = None + if clean: + earliest_ts = datetime.fromisoformat(cg.meta.custom_data["earliest_ts"]) + + for node, parent, node_ts in zip(nodes, parents, nodes_ts): + if parent == 0: + # ignore invalid nodes from failed ingest tasks, w/o parent column entry + # retain invalid nodes from edits to fix the hierarchy + if clean and node_ts > earliest_ts: + corrupt_nodes.append(node) + else: + tasks.append((cg, layer, node, node_ts)) + + if clean: + logger.note(f"found {len(corrupt_nodes)} corrupt nodes {corrupt_nodes[:3]}...") + fix_corrupt_nodes(cg, corrupt_nodes, CHILDREN) + return + + rows = [] + for task in tasks: + rows.extend(update_cross_edges(*task)) + stale.PARENTS_CACHE.clear() + stale.CHILDREN_CACHE.clear() + cg.client.write(rows) + gc.collect() + + +def _get_split_nodes( + cg: ChunkedGraph, chunk_id: basetypes.CHUNK_ID, split: int, splits: int +): + max_id = cg.client.get_max_node_id(chunk_id) + total = max_id - chunk_id + split_size = int(ceil(total / splits)) + start = int(chunk_id + np.uint64(split * split_size)) + end = int(start + split_size) + return range(int(start), int(end)) + + +def update_chunk( + cg: ChunkedGraph, + chunk_coords: list[int], + layer: int, + nodes: list = None, + split: int = None, + splits: int = None, + clean: bool = False, + n_processes: int = 1, +): + """ + Iterate over all layer IDs in a chunk and update their cross chunk edges. + """ + debug = nodes is not None + start = time.time() + x, y, z = chunk_coords + chunk_id = cg.get_chunk_id(layer=layer, x=x, y=y, z=z) + + if splits is not None: + nodes = _get_split_nodes(cg, chunk_id, split, splits) + + _populate_nodes_and_children(cg, chunk_id, nodes=nodes) + logger.note(f"_populate_nodes_and_children: {time.time() - start}") + nodes = list(CHILDREN.keys()) + if len(nodes) == 0: + return + + logger.note(f"processing {len(nodes)} nodes.") + random.shuffle(nodes) + start = time.time() + nodes_ts = cg.get_node_timestamps(nodes, return_numpy=False, normalize=True) + logger.note(f"get_node_timestamps: {time.time() - start}") + + start = time.time() + _populate_cx_edges_with_timestamps(cg, layer, nodes, nodes_ts, clean) + logger.note(f"_populate_cx_edges_with_timestamps: {time.time() - start}") + + if debug: + rows = [] + stale.PARENTS_CACHE = LRUCache(PARENT_CACHE_LIMIT) + stale.CHILDREN_CACHE = LRUCache(1 * 1024) + logger.note(f"processing {len(nodes)} nodes with 1 worker.") + for node, node_ts in zip(nodes, nodes_ts): + rows.extend(update_cross_edges(cg, layer, node, node_ts)) + stale.PARENTS_CACHE.clear() + stale.CHILDREN_CACHE.clear() + logger.note(f"total elaspsed time: {time.time() - start}") + return + + task_size = int(os.environ.get("TASK_SIZE", 1)) + chunked_nodes = chunked(nodes, task_size) + chunked_nodes_ts = chunked(nodes_ts, task_size) + cg_info = cg.get_serialized_info() + + tasks = [] + for chunk, ts_chunk in zip(chunked_nodes, chunked_nodes_ts): + args = (cg_info, layer, chunk, ts_chunk, clean) + tasks.append(args) + + process_multiplier = int(os.environ.get("PROCESS_MULTIPLIER", 5)) + processes = min(n_processes * process_multiplier, len(tasks)) + logger.note(f"processing {len(nodes)} nodes with {processes} workers.") + with mp.Pool(processes) as pool: + _ = list( + tqdm( + pool.imap_unordered(_update_cross_edges_helper, tasks), + total=len(tasks), + ) + ) + logger.note(f"total elaspsed time: {time.time() - start}") diff --git a/pychunkedgraph/ingest/upgrade/utils.py b/pychunkedgraph/ingest/upgrade/utils.py new file mode 100644 index 000000000..caa1f067b --- /dev/null +++ b/pychunkedgraph/ingest/upgrade/utils.py @@ -0,0 +1,139 @@ +# pylint: disable=invalid-name, missing-docstring + +from collections import defaultdict +from datetime import datetime, timedelta + +import numpy as np +from pychunkedgraph.graph import ChunkedGraph +from pychunkedgraph.graph import attributes, serializers + + +def exists_as_parent(cg: ChunkedGraph, parent, nodes) -> bool: + """ + Check if a given l2 parent is in the history of given nodes. + """ + response = cg.client.read_nodes( + node_ids=nodes, properties=attributes.Hierarchy.Parent + ) + parents = set() + for cells in response.values(): + parents.update([cell.value for cell in cells]) + return parent in parents + + +def get_edit_timestamps(cg: ChunkedGraph, edges_d, start_ts, end_ts) -> list: + """ + Timestamps of when post-side nodes were involved in an edit. + Post-side - nodes in the neighbor chunk. + This is required because we need to update edges from both sides. + """ + cx_edges = np.concatenate(list(edges_d.values())) + timestamps = get_parent_timestamps( + cg, cx_edges[:, 1], start_time=start_ts, end_time=end_ts + ) + timestamps.add(start_ts) + return sorted(timestamps) + + +def _get_end_timestamps_helper(cg: ChunkedGraph, nodes: list) -> defaultdict[int, set]: + result = defaultdict(set) + response = cg.client.read_nodes( + node_ids=nodes, properties=attributes.Hierarchy.StaleTimeStamp + ) + for k, v in response.items(): + result[k].add(v[0].timestamp) + return result + + +def get_end_timestamps( + cg: ChunkedGraph, nodes: list, nodes_ts: datetime, children_map: dict, layer: int +): + """ + Gets the last timestamp for each node at which to update its cross edges. + For layer 2: + Get parent timestamps for all children of a node. + The first timestamp > node_timestamp among these is the last timestamp. + This is the timestamp at which one of node's children + got a new parent that superseded the current node. + These are cached in database. + For all nodes in each layer > 2: + Pick the earliest child node_end_ts > node_ts and cache in database. + """ + result = [] + children = np.concatenate([*children_map.values()]) + if layer == 2: + timestamps_d = get_parent_timestamps(cg, children) + else: + timestamps_d = _get_end_timestamps_helper(cg, children) + + for node, node_ts in zip(nodes, nodes_ts): + node_children = children_map[node] + _children_timestamps = [] + for k in node_children: + if k in timestamps_d: + _children_timestamps.append(timestamps_d[k]) + _timestamps = set().union(*_children_timestamps) + _timestamps.add(node_ts) + try: + _timestamps = sorted(_timestamps) + _index = np.searchsorted(_timestamps, node_ts) + end_ts = _timestamps[_index + 1] + except IndexError: + # this node has not been edited, but might have it edges updated + end_ts = None + result.append(end_ts) + return result + + +def get_parent_timestamps( + cg: ChunkedGraph, nodes, start_time=None, end_time=None +) -> defaultdict[int, set]: + """ + Timestamps of when the given nodes were edited. + """ + earliest_ts = cg.get_earliest_timestamp() + response = cg.client.read_nodes( + node_ids=nodes, + properties=[attributes.Hierarchy.Parent], + start_time=start_time, + end_time=end_time, + end_time_inclusive=False, + ) + + result = defaultdict(set) + for k, v in response.items(): + for cell in v[attributes.Hierarchy.Parent]: + ts = cell.timestamp + result[k].add(earliest_ts if ts < earliest_ts else ts) + return result + + +def fix_corrupt_nodes(cg: ChunkedGraph, nodes: list, children_d: dict): + """ + For each node: delete it from parent column of its children. + Then deletes the node itself, effectively erasing it from hierarchy. + """ + mutations = [] + row_keys_to_delete = [] + for node in nodes: + children = children_d[node] + _map = cg.client.read_nodes( + node_ids=children, properties=attributes.Hierarchy.Parent + ) + + for child, parent_cells in _map.items(): + timestamps_to_delete = [ + cell.timestamp for cell in parent_cells if cell.value == node + ] + if timestamps_to_delete: + mutations.append( + ( + serializers.serialize_uint64(child), + attributes.Hierarchy.Parent, + timestamps_to_delete, + ) + ) + row_keys_to_delete.append(serializers.serialize_uint64(node)) + + if mutations or row_keys_to_delete: + cg.client.delete_cells(mutations, row_keys_to_delete=row_keys_to_delete) diff --git a/pychunkedgraph/ingest/utils.py b/pychunkedgraph/ingest/utils.py index fa7ef7a3c..87b048be4 100644 --- a/pychunkedgraph/ingest/utils.py +++ b/pychunkedgraph/ingest/utils.py @@ -1,14 +1,44 @@ -from typing import Tuple +# pylint: disable=invalid-name, missing-docstring +import functools +import math +import sys +from os import environ +from time import sleep +from typing import Dict, Generator, Tuple + +import numpy as np +from kvdbclient import get_config_class +from rich import box +from rich.console import Group +from rich.live import Live +from rich.panel import Panel +from rich.rule import Rule +from rich.table import Table +from rich.text import Text +from rq import Queue, Retry +from rq.registry import ( + CanceledJobRegistry, + DeferredJobRegistry, + FailedJobRegistry, + FinishedJobRegistry, + ScheduledJobRegistry, + StartedJobRegistry, +) +from rq.worker_registration import WORKERS_BY_QUEUE_KEY + +from pychunkedgraph import get_logger -from . import ClusterIngestConfig from . import IngestConfig -from ..graph.meta import ChunkedGraphMeta -from ..graph.meta import DataSource -from ..graph.meta import GraphConfig +from .manager import IngestionManager +from ..graph import BackendClientInfo +from ..graph.meta import ChunkedGraphMeta, DataSource, GraphConfig +from ..graph.ocdbt import OcdbtConfig +from ..utils.general import chunked +from ..utils.redis import get_redis_connection +from ..utils.redis import keys as r_keys -from ..graph.client import BackendClientInfo -from ..graph.client.bigtable import BigTableConfig +logger = get_logger(__name__) chunk_id_str = lambda layer, coords: f"{layer}_{'_'.join(map(str, coords))}" @@ -16,30 +46,39 @@ def bootstrap( graph_id: str, config: dict, - overwrite: bool = False, raw: bool = False, test_run: bool = False, -) -> Tuple[ChunkedGraphMeta, IngestConfig, BackendClientInfo]: - """Parse config loaded from a yaml file.""" +) -> Tuple[ChunkedGraphMeta, IngestConfig, BackendClientInfo, Dict]: + """Parse config loaded from a yaml file. + + Returns ``(meta, ingest_config, client_info, ocdbt_config_dict)`` where the + ocdbt config dict is sanitized through ``OcdbtConfig.from_dict(...).to_dict()`` + so unknown yaml keys are dropped and missing fields take dataclass defaults. + """ ingest_config = IngestConfig( **config.get("ingest_config", {}), - CLUSTER=ClusterIngestConfig(), USE_RAW_EDGES=raw, USE_RAW_COMPONENTS=raw, TEST_RUN=test_run, ) - client_config = BigTableConfig(**config["backend_client"]["CONFIG"]) - client_info = BackendClientInfo(config["backend_client"]["TYPE"], client_config) + backend_type = config["backend_client"].get("TYPE", "bigtable") + client_config = get_config_class(backend_type)(**config["backend_client"]["CONFIG"]) + client_info = BackendClientInfo(backend_type, client_config) graph_config = GraphConfig( ID=f"{graph_id}", - OVERWRITE=overwrite, + OVERWRITE=False, **config["graph_config"], ) data_source = DataSource(**config["data_source"]) meta = ChunkedGraphMeta(graph_config, data_source) - return (meta, ingest_config, client_info) + ocdbt_config_dict = OcdbtConfig.from_dict(config.get("ocdbt_config")).to_dict() + return (meta, ingest_config, client_info, ocdbt_config_dict) + + +def move_up(lines: int = 1): + sys.stdout.write(f"\033[{lines}A") def postprocess_edge_data(im, edge_dict): @@ -72,4 +111,409 @@ def postprocess_edge_data(im, edge_dict): return new_edge_dict else: - raise Exception(f"Unknown data_version: {data_version}") + raise ValueError(f"Unknown data_version: {data_version}") + + +def randomize_grid_points(X: int, Y: int, Z: int) -> Generator[int, int, int]: + indices = np.arange(X * Y * Z) + np.random.shuffle(indices) + for index in indices: + yield np.unravel_index(index, (X, Y, Z)) + + +def get_chunks_not_done( + imanager: IngestionManager, layer: int, coords: list, splits: int = 0 +) -> list: + """check for set membership in redis in batches""" + coords_strs = [] + if splits > 0: + split_coords = [] + for coord in coords: + for split in range(splits): + jid = "_".join(map(str, coord)) + f"_{split}" + coords_strs.append(jid) + split_coords.append((coord, split)) + else: + coords_strs = ["_".join(map(str, coord)) for coord in coords] + try: + completed = imanager.redis.smismember(f"{layer}c", coords_strs) + except Exception: + return split_coords if splits > 0 else coords + + if splits > 0: + return [coord for coord, c in zip(split_coords, completed) if not c] + return [coord for coord, c in zip(coords, completed) if not c] + + +def print_completion_rate(imanager: IngestionManager, layer: int, span: int = 30): + rate = 0.0 + while True: + counts = [] + print(f"{rate} chunks per second.") + for _ in range(span + 1): + counts.append(imanager.redis.scard(f"{layer}c")) + sleep(1) + rate = np.diff(counts).sum() / span + move_up() + + +def _workers_busy_per_queue(redis, worker_keys_per_layer): + """For each layer's set of worker keys, return parallel (workers, busy) + string lists — "-" / "-" when no workers are registered for that layer. + + Two-round-trip approach: caller already fetched the SMEMBERS sets; this + function pipelines HGET state for every worker key and counts busy. + """ + state_pipe = redis.pipeline() + for keys in worker_keys_per_layer: + for wk in keys: + state_pipe.hget(wk, "state") + states = state_pipe.execute() if any(worker_keys_per_layer) else [] + + workers, busy = [], [] + idx = 0 + for keys in worker_keys_per_layer: + total = len(keys) + b = 0 + for _ in keys: + if states[idx] == b"busy": + b += 1 + idx += 1 + workers.append(f"{total}" if total else "-") + busy.append(f"{b}" if total else "-") + return workers, busy + + +def _layer_keys(layers) -> list: + """Stable per-layer redis keys (completed-set, queue list, failed zset, workers set). + + Returned once before the refresh loop so each refresh skips Queue / + FailedJobRegistry construction and the lazy rq.registry import. + """ + return [ + ( + f"{layer}c", + f"rq:queue:l{layer}", + f"rq:failed:l{layer}", + WORKERS_BY_QUEUE_KEY % f"l{layer}", + ) + for layer in layers + ] + + +def _layer_status(redis, layer_keys): + """Pipelined fetch of job_type + per-layer counts + busy-worker ratios.""" + pipeline = redis.pipeline() + pipeline.get(r_keys.JOB_TYPE) + for completed_key, queue_key, failed_key, workers_key in layer_keys: + pipeline.scard(completed_key) + pipeline.llen(queue_key) + pipeline.zcard(failed_key) + pipeline.smembers(workers_key) + results = pipeline.execute() + + job_type = results[0].decode() if results[0] else "not_available" + completed, queued, failed, worker_keys_per_layer = [], [], [], [] + for i in range(1, len(results), 4): + completed.append(results[i]) + queued.append(results[i + 1]) + failed.append(results[i + 2]) + worker_keys_per_layer.append(results[i + 3]) + + workers, busy = _workers_busy_per_queue(redis, worker_keys_per_layer) + return job_type, completed, queued, failed, workers, busy + + +def _sized_table(columns: list, rows: list, **table_kwargs) -> Table: + """Build a Rich Table whose column widths are sized to the actual data. + + `columns` is a list of (name, justify) tuples. + `rows` is a list of tuples of cell strings (one per column). + Each column gets width = max(len(name), max(len(cell)) over rows) so Rich + never wraps or crops because no column is implicitly squeezed. + """ + table = Table( + box=None, + pad_edge=False, + padding=(0, 2), + show_header=True, + header_style="bold", + **table_kwargs, + ) + for col_idx, (name, justify) in enumerate(columns): + width = max(len(name), max((len(row[col_idx]) for row in rows), default=0)) + # Header wrapped in Text so any brackets in `name` render literally + # rather than being parsed as Rich markup tags. + table.add_column( + Text(name, style="bold"), justify=justify, width=width, no_wrap=True + ) + for row in rows: + table.add_row(*row) + return table + + +def _aligned_kv_table(pairs: list, widths: list) -> Table: + """One-data-row mini-table with externally-provided per-column widths.""" + table = Table( + box=None, pad_edge=False, padding=(0, 1), show_header=True, header_style="bold" + ) + for (name, _), w in zip(pairs, widths): + table.add_column(name, justify="left", width=w, no_wrap=True) + table.add_row(*(v for _, v in pairs)) + return table + + +def _header_renderables(imanager: IngestionManager) -> list: + """Graph and ocdbt rows as mini-tables sharing column widths so columns line up.""" + graph_pairs = [ + ("version", str(imanager.cg.version)), + ("graph_id", imanager.cg.graph_id), + ("chunk_size", str(imanager.cg.meta.graph_config.CHUNK_SIZE)), + ] + ocdbt_pairs = [] + if imanager.ocdbt_seg: + ocdbt_pairs = [ + ("ocdbt", str(imanager.ocdbt_seg)), + ("populate_base", str(imanager.ocdbt_populate_base)), + ("populate_layer", str(imanager.ocdbt_populate_layer)), + ] + + # Per-column width = max length seen in EITHER row's header or value at that index. + n = max(len(graph_pairs), len(ocdbt_pairs)) + widths = [] + for i in range(n): + sizes = [] + if i < len(graph_pairs): + sizes.append(len(graph_pairs[i][0])) + sizes.append(len(graph_pairs[i][1])) + if i < len(ocdbt_pairs): + sizes.append(len(ocdbt_pairs[i][0])) + sizes.append(len(ocdbt_pairs[i][1])) + widths.append(max(sizes)) + + out = [_aligned_kv_table(graph_pairs, widths)] + if ocdbt_pairs: + out.append(Rule(style="dim")) + out.append(_aligned_kv_table(ocdbt_pairs, widths)) + return out + + +def _status_table( + layers, layer_counts, completed, queued, failed, workers, busy +) -> Table: + """One row per layer with progress, queue, and worker stats.""" + columns = [ + ("layer", "center"), + ("queued", "right"), + ("completed", "right"), + ("total", "right"), + ("progress", "right"), + ("failed", "right"), + ("workers", "right"), + ("busy", "right"), + ] + rows = [] + for layer, done, count, q, f, w, b in zip( + layers, completed, layer_counts, queued, failed, workers, busy + ): + pct = math.floor((done / count) * 100) if count else 0 + rows.append( + ( + str(layer), + f"{q:,}", + f"{done:,}", + f"{count:,}", + f"{pct}%", + f"{f:,}", + str(w), + str(b), + ) + ) + return _sized_table(columns, rows) + + +def _status_renderable( + imanager, + layers, + layer_counts, + job_type, + completed, + queued, + failed, + workers, + busy, +): + """Combine header rows + per-layer table inside one Panel; job_type goes in the title.""" + body = Group( + *_header_renderables(imanager), + Rule(style="dim"), + _status_table(layers, layer_counts, completed, queued, failed, workers, busy), + ) + return Panel( + body, + title=job_type, + title_align="left", + box=box.ROUNDED, + padding=(0, 1), + expand=False, + ) + + +def print_status( + imanager: IngestionManager, + redis, + upgrade: bool = False, + refresh_seconds: int = 5, +): + """ + Print status to console. + If `upgrade=True`, status does not include the root layer, + since there is no need to update cross edges for root ids. + `refresh_seconds` is how often redis is re-polled between redraws. + """ + layers = range(2, imanager.cg_meta.layer_count + 1) + if upgrade: + layers = range(2, imanager.cg_meta.layer_count) + layer_counts = imanager.cg_meta.layer_chunk_counts + layer_keys = _layer_keys(layers) + + def render(): + return _status_renderable( + imanager, layers, layer_counts, *_layer_status(redis, layer_keys) + ) + + # Start Live with a placeholder so the panel paints instantly; the first + # real fetch (which includes redis connection setup) replaces it. + with Live(Text("loading…"), screen=False) as live: + while True: + live.update(render()) + sleep(refresh_seconds) + + +def queue_layer_helper( + parent_layer: int, imanager: IngestionManager, fn, splits: int = 0 +): + if parent_layer == imanager.cg_meta.layer_count: + chunk_coords = [(0, 0, 0)] + else: + bounds = imanager.cg_meta.layer_chunk_bounds[parent_layer] + chunk_coords = randomize_grid_points(*bounds) + + q = imanager.get_task_queue(f"l{parent_layer}") + batch_size = int(environ.get("JOB_BATCH_SIZE", 10000)) + timeout_scale = int(environ.get("TIMEOUT_SCALE_FACTOR", 1)) + batches = chunked(chunk_coords, batch_size) + failure_ttl = int(environ.get("FAILURE_TTL", 300)) + retry = int(environ.get("RETRY_COUNT", 0)) + max_queue_size = int(environ.get("QUEUE_SIZE", 100000)) + for batch in batches: + _coords = get_chunks_not_done(imanager, parent_layer, batch, splits=splits) + # buffer for optimal use of redis memory + while len(q) > max_queue_size: + logger.note( + f"Queue has {len(q)} items (limit {max_queue_size}), waiting..." + ) + sleep(10) + + job_datas = [] + for chunk_coord in _coords: + if splits > 0: + coord, split = chunk_coord + jid = chunk_id_str(parent_layer, coord) + f"_{split}" + job_datas.append( + Queue.prepare_data( + fn, + args=(parent_layer, coord, split, splits), + result_ttl=0, + job_id=jid, + timeout=f"{timeout_scale * int(parent_layer * parent_layer)}m", + retry=Retry(retry) if retry > 1 else None, + description="", + failure_ttl=failure_ttl, + ) + ) + else: + job_datas.append( + Queue.prepare_data( + fn, + args=(parent_layer, chunk_coord), + result_ttl=0, + job_id=chunk_id_str(parent_layer, chunk_coord), + timeout=f"{timeout_scale * int(parent_layer * parent_layer)}m", + retry=Retry(retry) if retry > 1 else None, + description="", + failure_ttl=failure_ttl, + ) + ) + q.enqueue_many(job_datas) + logger.note(f"Queued {len(job_datas)} chunks.") + + +_RQ_REGISTRY_CLASSES = ( + FailedJobRegistry, + StartedJobRegistry, + DeferredJobRegistry, + ScheduledJobRegistry, + FinishedJobRegistry, + CanceledJobRegistry, +) + + +def purge_layer_state(redis, layer: int) -> None: + """Reset per-layer state so a layer can be re-run from a previous + layer's backup: drop the RQ queue (deletes jobs too), wipe each RQ + registry by its own ``.key`` attribute (so we don't hardcode RQ's + internal key naming), and clear the pychunkedgraph completion set + ``f"{layer}c"``. + """ + name = f"l{layer}" + Queue(name=name, connection=redis).delete(delete_jobs=True) + for cls in _RQ_REGISTRY_CLASSES: + redis.delete(cls(name=name, connection=redis).key) + redis.delete(f"{layer}c") + + +def requeue_chunk(queue_name: str, chunk_info, atomic_fn, parent_fn): + """Body of the ``chunk`` CLI command (shared by ingest and upgrade). + + Loads the manager from Redis, dispatches ``atomic_fn`` for L2 or + ``parent_fn`` for L3+, and enqueues a single task with the standard + job_id / timeout convention. + """ + redis = get_redis_connection() + imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER)) + layer, coords = chunk_info[0], chunk_info[1:] + if layer == 2: + fn, args = atomic_fn, (coords,) + else: + fn, args = parent_fn, (layer, coords) + queue = imanager.get_task_queue(queue_name) + queue.enqueue( + fn, + job_id=chunk_id_str(layer, coords), + job_timeout=f"{int(layer * layer)}m", + result_ttl=0, + args=args, + ) + + +def job_type_guard(job_type: str): + def decorator_job_type_guard(func): + @functools.wraps(func) + def wrapper_job_type_guard(*args, **kwargs): + redis = get_redis_connection() + current_type = redis.get(r_keys.JOB_TYPE) + if current_type is not None: + current_type = current_type.decode() + msg = ( + f"Currently running `{current_type}`. You're attempting to run `{job_type}`." + f"\nRun `[flask] {current_type} flush_redis` to clear the current job and restart." + ) + if current_type != job_type: + print(f"\n*WARNING*\n{msg}") + exit(1) + return func(*args, **kwargs) + + return wrapper_job_type_guard + + return decorator_job_type_guard diff --git a/pychunkedgraph/io/components.py b/pychunkedgraph/io/components.py index a6301c7d2..6d554c7e5 100644 --- a/pychunkedgraph/io/components.py +++ b/pychunkedgraph/io/components.py @@ -4,7 +4,7 @@ from cloudfiles import CloudFiles from .protobuf.chunkComponents_pb2 import ChunkComponentsMsg -from ..graph.utils import basetypes +from ..graph import basetypes def serialize(connected_components: Iterable) -> ChunkComponentsMsg: diff --git a/pychunkedgraph/io/edges.py b/pychunkedgraph/io/edges.py index 82595e139..a9fa76aa6 100644 --- a/pychunkedgraph/io/edges.py +++ b/pychunkedgraph/io/edges.py @@ -2,6 +2,7 @@ """ Functions for reading and writing edges from cloud storage. """ + import os from typing import Dict from typing import List @@ -15,7 +16,7 @@ from .protobuf.chunkEdges_pb2 import ChunkEdgesMsg from ..graph.edges import Edges from ..graph.edges import EDGE_TYPES -from ..graph.utils import basetypes +from ..graph import basetypes from ..graph.edges.utils import concatenate_chunk_edges @@ -38,7 +39,7 @@ def deserialize(edges_message: EdgesMsg) -> Tuple[np.ndarray, np.ndarray, np.nda def _parse_edges(compressed: List[bytes]) -> List[Dict]: result = [] - if(len(compressed) == 0): + if len(compressed) == 0: return result zdc = zstd.ZstdDecompressor() try: diff --git a/pychunkedgraph/logging/log_db.py b/pychunkedgraph/logging/log_db.py index 89680500a..4a4244022 100644 --- a/pychunkedgraph/logging/log_db.py +++ b/pychunkedgraph/logging/log_db.py @@ -4,7 +4,7 @@ import threading import time import queue -from datetime import datetime +from datetime import datetime, timezone from google.api_core.exceptions import GoogleAPIError from datastoreflex import DatastoreFlex @@ -109,7 +109,7 @@ def __init__(self, name: str, graph_id: str, operation_id=-1, **kwargs): self.names.append(name) self._start = None self._graph_id = graph_id - self._ts = datetime.utcnow() + self._ts = datetime.now(timezone.utc) self._kwargs = kwargs if operation_id != -1: self.operation_id = operation_id diff --git a/pychunkedgraph/meshing/manifest/cache.py b/pychunkedgraph/meshing/manifest/cache.py index f38a830c2..0decc3a65 100644 --- a/pychunkedgraph/meshing/manifest/cache.py +++ b/pychunkedgraph/meshing/manifest/cache.py @@ -12,6 +12,8 @@ DOES_NOT_EXIST = "X" INITIAL_PATH_PREFIX = "initial_path_prefix" +MANIFEST_TTL_SECONDS = 3 * 24 * 3600 + REDIS_HOST = os.environ.get("MANIFEST_CACHE_REDIS_HOST", "localhost") REDIS_PORT = os.environ.get("MANIFEST_CACHE_REDIS_PORT", "6379") REDIS_PASSWORD = os.environ.get("MANIFEST_CACHE_REDIS_PASSWORD", "") @@ -74,6 +76,27 @@ def clear_fragments(self, node_ids) -> None: keys = [f"{self.namespace}:{n}" for n in node_ids] REDIS.delete(*keys) + def clear_namespace(self, batch_size: int = 1000) -> int: + """Delete every key under this graph_id's namespace. + + SCAN-based pattern delete (non-blocking on the redis side). + Returns the number of keys deleted. + """ + if REDIS is None: + return 0 + + pattern = f"{self.namespace}:*" + deleted = 0 + batch = [] + for key in REDIS.scan_iter(match=pattern, count=batch_size): + batch.append(key) + if len(batch) >= batch_size: + deleted += REDIS.delete(*batch) + batch.clear() + if batch: + deleted += REDIS.delete(*batch) + return deleted + def _get_cached_initial_fragments(self, node_ids: List[np.uint64]): if REDIS is None: return {}, node_ids, [] @@ -140,10 +163,16 @@ def _set_cached_initial_fragments( for node_id, fragment_info in fragments_d.items(): path, offset, size = fragment_info key = f"{self.namespace}:{node_id}" - pipeline.set(key, f"{path[prefix_idx:]}:{offset}:{size}") + pipeline.set( + key, f"{path[prefix_idx:]}:{offset}:{size}", ex=MANIFEST_TTL_SECONDS + ) for node_id in not_existing: - pipeline.set(f"{self.namespace}:{node_id}", DOES_NOT_EXIST) + pipeline.set( + f"{self.namespace}:{node_id}", + DOES_NOT_EXIST, + ex=MANIFEST_TTL_SECONDS, + ) pipeline.execute() @@ -155,9 +184,15 @@ def _set_cached_dynamic_fragments( pipeline = REDIS.pipeline() for node_id, fragment in fragments_d.items(): - pipeline.set(f"{self.namespace}:{node_id}", fragment) + pipeline.set( + f"{self.namespace}:{node_id}", fragment, ex=MANIFEST_TTL_SECONDS + ) for node_id in not_existing: - pipeline.set(f"{self.namespace}:{node_id}", DOES_NOT_EXIST) + pipeline.set( + f"{self.namespace}:{node_id}", + DOES_NOT_EXIST, + ex=MANIFEST_TTL_SECONDS, + ) pipeline.execute() diff --git a/pychunkedgraph/meshing/manifest/sharded.py b/pychunkedgraph/meshing/manifest/sharded.py index 2576fcb2f..8b122b235 100644 --- a/pychunkedgraph/meshing/manifest/sharded.py +++ b/pychunkedgraph/meshing/manifest/sharded.py @@ -8,7 +8,7 @@ from .utils import get_children_before_start_layer from ...graph import ChunkedGraph from ...graph.types import empty_1d -from ...graph.utils.basetypes import NODE_ID +from ...graph.basetypes import NODE_ID from ...graph.chunks import utils as chunk_utils diff --git a/pychunkedgraph/meshing/manifest/utils.py b/pychunkedgraph/meshing/manifest/utils.py index 67e600653..3c4538797 100644 --- a/pychunkedgraph/meshing/manifest/utils.py +++ b/pychunkedgraph/meshing/manifest/utils.py @@ -16,7 +16,7 @@ from ..meshgen_utils import get_json_info from ...graph import ChunkedGraph from ...graph.types import empty_1d -from ...graph.utils.basetypes import NODE_ID +from ...graph.basetypes import NODE_ID from ...graph.utils import generic as misc_utils @@ -40,7 +40,7 @@ def _get_children(cg, node_ids: Sequence[np.uint64], children_cache: Dict): if len(node_ids) == 0: return empty_1d.copy() node_ids = np.array(node_ids, dtype=NODE_ID) - mask = np.in1d(node_ids, np.fromiter(children_cache.keys(), dtype=NODE_ID)) + mask = np.isin(node_ids, np.fromiter(children_cache.keys(), dtype=NODE_ID)) children_d = cg.get_children(node_ids[~mask]) children_cache.update(children_d) @@ -105,8 +105,10 @@ def _get_dynamic_meshes(cg, node_ids: Sequence[np.uint64]) -> Tuple[Dict, List]: if len(node_ids) == 0: return result, not_existing - mesh_dir = cg.meta.custom_data.get("mesh", {}).get("dir", "graphene_meshes") - mesh_path = f"{cg.meta.data_source.WATERSHED}/{mesh_dir}/dynamic" + mesh_meta = cg.meta.custom_data.get("mesh", {}) + mesh_dir = mesh_meta.get("dir", "graphene_meshes") + dynamic_dir = mesh_meta.get("dynamic_mesh_dir", "dynamic") + mesh_path = f"{cg.meta.data_source.WATERSHED}/{mesh_dir}/{dynamic_dir}" cf = CloudFiles(mesh_path) manifest_cache = ManifestCache(cg.graph_id, initial=False) diff --git a/pychunkedgraph/meshing/mesh_analysis.py b/pychunkedgraph/meshing/mesh_analysis.py index 97bb28f5b..abdf95957 100644 --- a/pychunkedgraph/meshing/mesh_analysis.py +++ b/pychunkedgraph/meshing/mesh_analysis.py @@ -16,10 +16,10 @@ def compute_centroid_with_chunk_boundary(cg, vertices, l2_id, last_l2_id): a path, return the center point of the mesh on the chunk boundary separating the two ids, and the center point of the entire mesh. :param cg: ChunkedGraph object - :param vertices: [[np.float]] + :param vertices: [[np.float64]] :param l2_id: np.uint64 :param last_l2_id: np.uint64 or None - :return: [np.float] + :return: [np.float64] """ centroid_by_range = compute_centroid_by_range(vertices) if last_l2_id is None: diff --git a/pychunkedgraph/meshing/mesh_io.py b/pychunkedgraph/meshing/mesh_io.py index 40c02bba0..4a6eac7c4 100644 --- a/pychunkedgraph/meshing/mesh_io.py +++ b/pychunkedgraph/meshing/mesh_io.py @@ -7,19 +7,23 @@ import networkx as nx import cloudvolume -from multiwrapper import multiprocessing_utils as mu +from concurrent.futures import ProcessPoolExecutor + def read_mesh_h5(): pass + def write_mesh_h5(): pass + def read_obj(path): return Mesh(path) + def _download_meshes_thread(args): - """ Downloads meshes into target directory + """Downloads meshes into target directory :param args: list """ @@ -33,7 +37,7 @@ def _download_meshes_thread(args): def download_meshes(seg_ids, target_dir, cv_path, n_threads=1): - """ Downloads meshes in target directory (parallel) + """Downloads meshes in target directory (parallel) :param seg_ids: list of ints :param target_dir: str @@ -52,12 +56,11 @@ def download_meshes(seg_ids, target_dir, cv_path, n_threads=1): multi_args.append([seg_id_block, cv_path, target_dir]) if n_jobs == 1: - mu.multiprocess_func(_download_meshes_thread, - multi_args, debug=True, - verbose=True, n_threads=1) + for args in multi_args: + _download_meshes_thread(args) else: - mu.multisubprocess_func(_download_meshes_thread, - multi_args, n_threads=n_threads) + with ProcessPoolExecutor(max_workers=n_threads) as executor: + list(executor.map(_download_meshes_thread, multi_args)) def refine_mesh(): @@ -77,6 +80,7 @@ def mesh(self, filename): return self.filename_dict[filename] + class Mesh(object): def __init__(self, filename): self._vertices = [] @@ -117,8 +121,9 @@ def normals(self): @property def edges(self): if self._edges is None: - self._edges = np.concatenate([self.faces[:, :2], - self.faces[:, 1:3]], axis=0) + self._edges = np.concatenate( + [self.faces[:, :2], self.faces[:, 1:3]], axis=0 + ) return self._edges @property @@ -141,21 +146,23 @@ def load_obj(self): normals = [] for line in open(self.filename, "r"): - if line.startswith('#'): continue + if line.startswith("#"): + continue values = line.split() - if not values: continue - if values[0] == 'v': + if not values: + continue + if values[0] == "v": v = values[1:4] vertices.append(v) - elif values[0] == 'vn': + elif values[0] == "vn": v = map(float, values[1:4]) normals.append(v) - elif values[0] == 'f': + elif values[0] == "f": face = [] texcoords = [] norms = [] for v in values[1:]: - w = v.split('/') + w = v.split("/") face.append(int(w[0])) if len(w) >= 2 and len(w[1]) > 0: texcoords.append(int(w[1])) @@ -168,8 +175,8 @@ def load_obj(self): faces.append(face) self._faces = np.array(faces, dtype=int) - 1 - self._vertices = np.array(vertices, dtype=np.float) - self._normals = np.array(normals, dtype=np.float) + self._vertices = np.array(vertices, dtype=np.float64) + self._normals = np.array(normals, dtype=np.float64) def load_h5(self): with h5py.File(self.filename, "r") as f: @@ -191,7 +198,8 @@ def write_vertices_ply(self, out_fname, coords=None): tweaked_array = np.array( list(zip(coords[:, 0], coords[:, 1], coords[:, 2])), - dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) + dtype=[("x", "f4"), ("y", "f4"), ("z", "f4")], + ) vertex_element = plyfile.PlyElement.describe(tweaked_array, "vertex") @@ -200,8 +208,15 @@ def write_vertices_ply(self, out_fname, coords=None): plyfile.PlyData([vertex_element]).write(out_fname) - def get_local_view(self, n_points, pc_align=False, center_node_id=None, - center_coord=None, method="kdtree", verbose=False): + def get_local_view( + self, + n_points, + pc_align=False, + center_node_id=None, + center_coord=None, + method="kdtree", + verbose=False, + ): if center_node_id is None and center_coord is None: center_node_id = np.random.randint(len(self.vertices)) @@ -215,11 +230,11 @@ def get_local_view(self, n_points, pc_align=False, center_node_id=None, if verbose: print(np.mean(dists), np.max(dists), np.min(dists)) elif method == "graph": - dist_dict = nx.single_source_dijkstra_path_length(self.graph, - center_node_id, - weight="weight") - sorting = np.argsort(np.array(list(dist_dict.values()))) - node_ids = np.array(list(dist_dict.keys()))[sorting[:n_points]] + dist_dict = nx.single_source_dijkstra_path_length( + self.graph, center_node_id, weight="weight" + ) + sorting = np.argsort(np.array(list(dist_dict.values()))) + node_ids = np.array(list(dist_dict.keys()))[sorting[:n_points]] else: raise Exception("unknow method") @@ -236,7 +251,9 @@ def calc_pc_align(self, vertices): return pca.transform(vertices) def create_nx_graph(self): - weights = np.linalg.norm(self.vertices[self.edges[:, 0]] - self.vertices[self.edges[:, 1]], axis=1) + weights = np.linalg.norm( + self.vertices[self.edges[:, 0]] - self.vertices[self.edges[:, 1]], axis=1 + ) print(weights.shape) @@ -244,8 +261,6 @@ def create_nx_graph(self): weighted_graph.add_edges_from(self.edges) for i_edge, edge in enumerate(self.edges): - weighted_graph[edge[0]][edge[1]]['weight'] = weights[i_edge] + weighted_graph[edge[0]][edge[1]]["weight"] = weights[i_edge] return weighted_graph - - diff --git a/pychunkedgraph/meshing/meshengine.py b/pychunkedgraph/meshing/meshengine.py index 615e6cdb6..3f86fd7b3 100644 --- a/pychunkedgraph/meshing/meshengine.py +++ b/pychunkedgraph/meshing/meshengine.py @@ -3,19 +3,21 @@ import itertools import random +from concurrent.futures import ProcessPoolExecutor from pychunkedgraph.graph import chunkedgraph -from multiwrapper import multiprocessing_utils as mu from . import meshgen class MeshEngine(object): - def __init__(self, - table_id: str, - instance_id: str = "pychunkedgraph", - project_id: str = "neuromancer-seung-import", - mesh_mip: int = 3, - highest_mesh_layer: int = 5): + def __init__( + self, + table_id: str, + instance_id: str = "pychunkedgraph", + project_id: str = "neuromancer-seung-import", + mesh_mip: int = 3, + highest_mesh_layer: int = 5, + ): self._table_id = table_id self._instance_id = instance_id @@ -62,7 +64,8 @@ def cg(self): self._cg = chunkedgraph.ChunkedGraph( table_id=self.table_id, instance_id=self.instance_id, - project_id=self.project_id) + project_id=self.project_id, + ) return self._cg @property @@ -80,8 +83,9 @@ def cv(self): self._cv.info["mesh"] = self.cv_mesh_dir return self._cv - def mesh_multiple_layers(self, layers=None, bounding_box=None, - block_factor=2, n_threads=128): + def mesh_multiple_layers( + self, layers=None, bounding_box=None, block_factor=2, n_threads=128 + ): if layers is None: layers = range(1, int(self.cg.n_layers + 1)) @@ -94,28 +98,30 @@ def mesh_multiple_layers(self, layers=None, bounding_box=None, for layer in layers: print("Now: layer %d" % layer) - self.mesh_single_layer(layer, bounding_box=bounding_box, - block_factor=block_factor, - n_threads=n_threads) - - def mesh_single_layer(self, layer, bounding_box=None, block_factor=2, - n_threads=128): + self.mesh_single_layer( + layer, + bounding_box=bounding_box, + block_factor=block_factor, + n_threads=n_threads, + ) + + def mesh_single_layer( + self, layer, bounding_box=None, block_factor=2, n_threads=128 + ): assert layer <= self.highest_mesh_layer dataset_bounding_box = np.array(self.cv.bounds.to_list()) - block_bounding_box_cg = \ - [np.floor(dataset_bounding_box[:3] / - self.cg.chunk_size).astype(int), - np.ceil(dataset_bounding_box[3:] / - self.cg.chunk_size).astype(int)] + block_bounding_box_cg = [ + np.floor(dataset_bounding_box[:3] / self.cg.chunk_size).astype(int), + np.ceil(dataset_bounding_box[3:] / self.cg.chunk_size).astype(int), + ] if bounding_box is not None: - bounding_box_cg = \ - [np.floor(bounding_box[0] / - self.cg.chunk_size).astype(int), - np.ceil(bounding_box[1] / - self.cg.chunk_size).astype(int)] + bounding_box_cg = [ + np.floor(bounding_box[0] / self.cg.chunk_size).astype(int), + np.ceil(bounding_box[1] / self.cg.chunk_size).astype(int), + ] m = block_bounding_box_cg[0] < bounding_box_cg[0] block_bounding_box_cg[0][m] = bounding_box_cg[0][m] @@ -126,31 +132,37 @@ def mesh_single_layer(self, layer, bounding_box=None, block_factor=2, block_bounding_box_cg /= 2 ** np.max([0, layer - 2]) block_bounding_box_cg = np.ceil(block_bounding_box_cg) - n_jobs = np.product(block_bounding_box_cg[1] - - block_bounding_box_cg[0]) / \ - block_factor ** 2 < n_threads + n_jobs = ( + np.prod(block_bounding_box_cg[1] - block_bounding_box_cg[0]) + / block_factor**2 + < n_threads + ) while n_jobs < n_threads and block_factor > 1: block_factor -= 1 - n_jobs = np.product(block_bounding_box_cg[1] - - block_bounding_box_cg[0]) / \ - block_factor ** 2 < n_threads - - block_iter = itertools.product(np.arange(block_bounding_box_cg[0][0], - block_bounding_box_cg[1][0], - block_factor), - np.arange(block_bounding_box_cg[0][1], - block_bounding_box_cg[1][1], - block_factor), - np.arange(block_bounding_box_cg[0][2], - block_bounding_box_cg[1][2], - block_factor)) + n_jobs = ( + np.prod(block_bounding_box_cg[1] - block_bounding_box_cg[0]) + / block_factor**2 + < n_threads + ) + + block_iter = itertools.product( + np.arange( + block_bounding_box_cg[0][0], block_bounding_box_cg[1][0], block_factor + ), + np.arange( + block_bounding_box_cg[0][1], block_bounding_box_cg[1][1], block_factor + ), + np.arange( + block_bounding_box_cg[0][2], block_bounding_box_cg[1][2], block_factor + ), + ) blocks = np.array(list(block_iter), dtype=int) cg_info = self.cg.get_serialized_info() - del (cg_info['credentials']) + del cg_info["credentials"] multi_args = [] for start_block in blocks: @@ -158,44 +170,57 @@ def mesh_single_layer(self, layer, bounding_box=None, block_factor=2, m = end_block > block_bounding_box_cg[1] end_block[m] = block_bounding_box_cg[1][m] - multi_args.append([cg_info, start_block, end_block, self.cg._cv_path, - self.cv_mesh_dir, self.mesh_mip, layer]) + multi_args.append( + [ + cg_info, + start_block, + end_block, + self.cg._cv_path, + self.cv_mesh_dir, + self.mesh_mip, + layer, + ] + ) random.shuffle(multi_args) random.shuffle(multi_args) # Run parallelizing if n_threads == 1: - mu.multiprocess_func(meshgen._mesh_layer_thread, multi_args, - n_threads=n_threads, verbose=True, - debug=n_threads == 1) + for args in multi_args: + meshgen._mesh_layer_thread(args) else: - mu.multisubprocess_func(meshgen._mesh_layer_thread, multi_args, - n_threads=n_threads, - suffix="%s_%d" % (self.table_id, layer)) + with ProcessPoolExecutor(max_workers=n_threads) as executor: + list(executor.map(meshgen._mesh_layer_thread, multi_args)) def create_manifests_for_higher_layers(self, n_threads=1): root_id_max = self.cg.get_max_node_id( - self.cg.get_chunk_id(layer=int(self.cg.n_layers), - x=int(0), y=int(0), - z=int(0))) + self.cg.get_chunk_id( + layer=int(self.cg.n_layers), x=int(0), y=int(0), z=int(0) + ) + ) - root_id_blocks = np.linspace(1, root_id_max, n_threads*3).astype(int) + root_id_blocks = np.linspace(1, root_id_max, n_threads * 3).astype(int) cg_info = self.cg.get_serialized_info() - del (cg_info['credentials']) + del cg_info["credentials"] multi_args = [] for i_block in range(len(root_id_blocks) - 1): - multi_args.append([cg_info, self.cv_path, self.cv_mesh_dir, - root_id_blocks[i_block], - root_id_blocks[i_block + 1], - self.highest_mesh_layer]) + multi_args.append( + [ + cg_info, + self.cv_path, + self.cv_mesh_dir, + root_id_blocks[i_block], + root_id_blocks[i_block + 1], + self.highest_mesh_layer, + ] + ) # Run parallelizing if n_threads == 1: - mu.multiprocess_func(meshgen._create_manifest_files_thread, - multi_args, n_threads=n_threads, verbose=True, - debug=n_threads == 1) + for args in multi_args: + meshgen._create_manifest_files_thread(args) else: - mu.multisubprocess_func(meshgen._create_manifest_files_thread, - multi_args, n_threads=n_threads) + with ProcessPoolExecutor(max_workers=n_threads) as executor: + list(executor.map(meshgen._create_manifest_files_thread, multi_args)) diff --git a/pychunkedgraph/meshing/meshgen.py b/pychunkedgraph/meshing/meshgen.py index a8da89b1f..42975c896 100644 --- a/pychunkedgraph/meshing/meshgen.py +++ b/pychunkedgraph/meshing/meshgen.py @@ -10,7 +10,7 @@ import pytz from scipy import ndimage -from multiwrapper import multiprocessing_utils as mu +from concurrent.futures import ThreadPoolExecutor from cloudfiles import CloudFiles from cloudvolume import CloudVolume from cloudvolume.datasource.precomputed.sharding import ShardingSpecification @@ -23,7 +23,6 @@ from pychunkedgraph.meshing import meshgen_utils # noqa from pychunkedgraph.meshing.manifest.cache import ManifestCache - UTC = pytz.UTC # Change below to true if debugging and want to see results in stdout @@ -40,8 +39,13 @@ def decode_draco_mesh_buffer(fragment): try: mesh_object = DracoPy.decode_buffer_to_mesh(fragment) - vertices = np.array(mesh_object.points) - faces = np.array(mesh_object.faces) + # asarray, not array: points/faces are already ndarrays, so this is + # zero-copy and aliases mesh_object's buffers. Callers mutate + # "vertices" in place (transform_draco_vertices) but never read + # mesh_object.points again, and each decode allocates its own buffer, + # so the alias is safe. + vertices = np.asarray(mesh_object.points) + faces = np.asarray(mesh_object.faces) except ValueError as exc: raise ValueError("Not a valid draco mesh") from exc @@ -75,7 +79,7 @@ def remap_seg_using_unsafe_dict(seg, unsafe_dict): overlaps.extend(np.unique(seg[:, :, -2][bin_cc_seg[:, :, -1]])) overlaps = np.unique(overlaps) - linked_l2_ids = overlaps[np.in1d(overlaps, unsafe_dict[unsafe_root_id])] + linked_l2_ids = overlaps[np.isin(overlaps, unsafe_dict[unsafe_root_id])] if len(linked_l2_ids) == 0: seg[bin_cc_seg] = 0 @@ -263,7 +267,12 @@ def _get_root_ids(args): multi_args.append([start_ids[i_block], start_ids[i_block + 1]]) if n_jobs > 0: - mu.multithread_func(_get_root_ids, multi_args, n_threads=n_threads) + if n_threads == 1: + for args in multi_args: + _get_root_ids(args) + else: + with ThreadPoolExecutor(max_workers=n_threads) as executor: + list(executor.map(_get_root_ids, multi_args)) return lx_ids, np.array(root_ids), lx_id_remap @@ -298,10 +307,13 @@ def calculate_stop_layer(cg, chunk_id): # Find lowest common chunk neigh_parent_chunk_ids = np.array(neigh_parent_chunk_ids) - layer_agreement = np.all( - (neigh_parent_chunk_ids - neigh_parent_chunk_ids[0]) == 0, axis=0 - ) - stop_layer = np.where(layer_agreement)[0][0] + chunk_layer + if chunk_layer + 1 == cg.meta.layer_count: + stop_layer = cg.meta.layer_count + else: + layer_agreement = np.all( + (neigh_parent_chunk_ids - neigh_parent_chunk_ids[0]) == 0, axis=0 + ) + stop_layer = np.where(layer_agreement)[0][0] + chunk_layer return stop_layer, neigh_chunk_ids @@ -317,12 +329,13 @@ def get_lx_overlapping_remappings(cg, chunk_id, time_stamp=None, n_threads=1): :return: multiples """ if time_stamp is None: - time_stamp = datetime.datetime.utcnow() + time_stamp = datetime.datetime.now(datetime.timezone.utc) if time_stamp.tzinfo is None: time_stamp = UTC.localize(time_stamp) stop_layer, neigh_chunk_ids = calculate_stop_layer(cg, chunk_id) - print(f"Stop layer: {stop_layer}") + if PRINT_FOR_DEBUGGING: + print(f"Stop layer: {stop_layer}") # Find the parent in the lowest common chunk for each l2 id. These parent # ids are referred to as root ids even though they are not necessarily the @@ -337,7 +350,8 @@ def get_lx_overlapping_remappings(cg, chunk_id, time_stamp=None, n_threads=1): # This loop is the main bottleneck for neigh_chunk_id in neigh_chunk_ids: - print(f"Neigh: {neigh_chunk_id} --------------") + if PRINT_FOR_DEBUGGING: + print(f"Neigh: {neigh_chunk_id} --------------") lx_ids, root_ids, lx_id_remap = get_root_lx_remapping( cg, neigh_chunk_id, stop_layer, time_stamp=time_stamp, n_threads=n_threads @@ -357,7 +371,7 @@ def get_lx_overlapping_remappings(cg, chunk_id, time_stamp=None, n_threads=1): ) safe_lx_ids = lx_ids[u_idx[c_root_ids == 1]] - unsafe_lx_ids = lx_ids[~np.in1d(lx_ids, safe_lx_ids)] + unsafe_lx_ids = lx_ids[~np.isin(lx_ids, safe_lx_ids)] unsafe_root_ids = np.unique(root_ids[u_idx[c_root_ids != 1]]) lx_root_dict = dict(zip(neigh_lx_ids, neigh_root_ids)) @@ -387,7 +401,7 @@ def get_lx_overlapping_remappings(cg, chunk_id, time_stamp=None, n_threads=1): unsafe_dict = collections.defaultdict(list) for root_id in unsafe_root_ids: - if np.sum(~np.in1d(root_lx_dict[root_id], unsafe_lx_ids)) == 0: + if np.sum(~np.isin(root_lx_dict[root_id], unsafe_lx_ids)) == 0: continue for neigh_lx_id in root_lx_dict[root_id]: @@ -443,7 +457,12 @@ def _get_root_ids(args): multi_args.append([start_ids[i_block], start_ids[i_block + 1]]) if n_jobs > 0: - mu.multithread_func(_get_root_ids, multi_args, n_threads=n_threads) + if n_threads == 1: + for args in multi_args: + _get_root_ids(args) + else: + with ThreadPoolExecutor(max_workers=n_threads) as executor: + list(executor.map(_get_root_ids, multi_args)) sv_ids_index = len(node_ids) chunk_ids_index = len(node_ids) + len(sv_ids) @@ -475,12 +494,13 @@ def get_lx_overlapping_remappings_for_nodes_and_svs( :return: multiples """ if time_stamp is None: - time_stamp = datetime.datetime.utcnow() + time_stamp = datetime.datetime.now(datetime.timezone.utc) if time_stamp.tzinfo is None: time_stamp = UTC.localize(time_stamp) stop_layer, _ = calculate_stop_layer(cg, chunk_id) - print(f"Stop layer: {stop_layer}") + if PRINT_FOR_DEBUGGING: + print(f"Stop layer: {stop_layer}") # Find the parent in the lowest common chunk for each node id and sv id. These parent # ids are referred to as root ids even though they are not necessarily the @@ -946,10 +966,11 @@ def chunk_initial_mesh_task( mesh_dst = cv_unsharded_mesh_path result.append((chunk_id, layer, cx, cy, cz)) - print( - "Retrieving remap table for chunk %s -- (%s, %s, %s, %s)" - % (chunk_id, layer, cx, cy, cz) - ) + if PRINT_FOR_DEBUGGING: + print( + "Retrieving remap table for chunk %s -- (%s, %s, %s, %s)" + % (chunk_id, layer, cx, cy, cz) + ) mesher = zmesh.Mesher(cg.meta.cv.mip_resolution(mip)) draco_encoding_settings = get_draco_encoding_settings_for_chunk( cg, chunk_id, mip, high_padding @@ -1040,7 +1061,8 @@ def get_multi_child_nodes(cg, chunk_id, node_id_subset=None, chunk_bbox_string=F fragment.value for child_fragments_for_node in node_rows for fragment in child_fragments_for_node - ], dtype=object + ], + dtype=object, ) # Filter out node ids that do not have roots (caused by failed ingest tasks) root_ids = cg.get_roots(node_ids, fail_to_zero=True) @@ -1107,16 +1129,19 @@ def chunk_stitch_remeshing_task( assert layer > 2 - print( - "Retrieving children for chunk %s -- (%s, %s, %s, %s)" - % (chunk_id, layer, cx, cy, cz) - ) + if PRINT_FOR_DEBUGGING: + print( + "Retrieving children for chunk %s -- (%s, %s, %s, %s)" + % (chunk_id, layer, cx, cy, cz) + ) multi_child_nodes, _ = get_multi_child_nodes(cg, chunk_id, node_id_subset, False) - print(f"{len(multi_child_nodes)} nodes with more than one child") + if PRINT_FOR_DEBUGGING: + print(f"{len(multi_child_nodes)} nodes with more than one child") result.append((chunk_id, len(multi_child_nodes))) if not multi_child_nodes: - print("Nothing to do", cx, cy, cz) + if PRINT_FOR_DEBUGGING: + print("Nothing to do", cx, cy, cz) return ", ".join(str(x) for x in result) cv = CloudVolume( @@ -1146,7 +1171,7 @@ def chunk_stitch_remeshing_task( fragments_d = {} for new_fragment_id, fragment_ids_to_fetch in multi_child_nodes.items(): i += 1 - if i % max(1, len(multi_child_nodes) // 10) == 0: + if PRINT_FOR_DEBUGGING and i % max(1, len(multi_child_nodes) // 10) == 0: print(f"{i}/{len(multi_child_nodes)}") old_fragments = [] @@ -1240,6 +1265,10 @@ def chunk_stitch_remeshing_task( def chunk_initial_sharded_stitching_task( cg_name, chunk_id, mip, cg=None, high_padding=1, cache=True ): + """DEPRECATED: single-threaded sharded stitch. ``meshing.meshing_sqs.MeshTask`` + now dispatches ``meshing.stitch.chunk_initial_sharded_stitching_task_mp`` + (parallel, mesh-equivalent output) instead. Kept as the reference + implementation the parallel path is gated against.""" start_existence_check_time = time.time() if cg is None: cg = ChunkedGraph(graph_id=cg_name) diff --git a/pychunkedgraph/meshing/meshgen_utils.py b/pychunkedgraph/meshing/meshgen_utils.py index 711c09322..7b3471be7 100644 --- a/pychunkedgraph/meshing/meshgen_utils.py +++ b/pychunkedgraph/meshing/meshgen_utils.py @@ -1,19 +1,13 @@ import re -import multiprocessing as mp -from time import time -from typing import List -from typing import Dict -from typing import Tuple from typing import Sequence from functools import lru_cache import numpy as np -from cloudvolume import CloudVolume from cloudvolume.lib import Vec -from multiwrapper import multiprocessing_utils as mu -from pychunkedgraph.graph.utils.basetypes import NODE_ID # noqa +from pychunkedgraph.graph.basetypes import NODE_ID # noqa from ..graph.types import empty_1d +from pychunkedgraph.graph.utils import get_local_segmentation def str_to_slice(slice_str: str): @@ -129,7 +123,13 @@ def recursive_helper(cur_node_ids): only_child_mask = np.array( [len(children_for_node) == 1 for children_for_node in children_array] ) - only_children = children_array[only_child_mask].astype(np.uint64).ravel() + # Extract children from object array - each filtered element is a 1-element array + filtered_children = children_array[only_child_mask] + only_children = ( + np.concatenate(filtered_children).astype(np.uint64) + if filtered_children.size + else np.array([], dtype=np.uint64) + ) if np.any(only_child_mask): temp_array = cur_node_ids[stop_layer_mask] temp_array[only_child_mask] = recursive_helper(only_children) @@ -145,19 +145,29 @@ def get_json_info(cg): dataset_info = cg.meta.dataset_info dummy_app_info = {"app": {"supported_api_versions": [0, 1]}} info = {**dataset_info, **dummy_app_info} - info["mesh"] = cg.meta.custom_data.get("mesh", {}).get("dir", "graphene_meshes") + mesh_meta = cg.meta.custom_data.get("mesh", {}) + info["mesh"] = mesh_meta.get("dir", "graphene_meshes") + # `dynamic_mesh_dir` lets a dataset name the unsharded dynamic-mesh + # subdir explicitly. Default `"dynamic"` matches the mesh worker's + # fallback and NG's current hardcoded subdir name — see the + # spelunker-ocdbt graphene backend (looks up + # `dynamic/`). NG must be patched to read + # this info field before non-default values route correctly. + dynamic_dir = mesh_meta.get("dynamic_mesh_dir", "dynamic") + info["dynamic_mesh_dir"] = dynamic_dir + # cloud-volume reads the dynamic dir from mesh_metadata.unsharded_mesh_dir, not + # dynamic_mesh_dir; mirror it so an unpatched client fetches dynamic meshes from + # the right dir. Copy the dict so cg.meta.dataset_info is untouched. + mesh_metadata = dict(info.get("mesh_metadata", {})) + mesh_metadata["unsharded_mesh_dir"] = dynamic_dir + info["mesh_metadata"] = mesh_metadata info_str = dumps(info) return loads(info_str) def get_ws_seg_for_chunk(cg, chunk_id, mip, overlap_vx=1): - cv = CloudVolume(cg.meta.cv.cloudpath, mip=mip, fill_missing=True) - mip_diff = mip - cg.meta.cv.mip - - mip_chunk_size = np.array(cg.meta.graph_config.CHUNK_SIZE, dtype=int) / np.array( - [2 ** mip_diff, 2 ** mip_diff, 1] - ) - mip_chunk_size = mip_chunk_size.astype(int) + layer = cg.get_chunk_layer(chunk_id) + mip_chunk_size = get_mesh_block_shape_for_mip(cg, layer, mip) chunk_start = ( cg.meta.cv.mip_voxel_offset(mip) @@ -169,11 +179,6 @@ def get_ws_seg_for_chunk(cg, chunk_id, mip, overlap_vx=1): cg.meta.cv.mip_voxel_offset(mip), cg.meta.cv.mip_voxel_offset(mip) + cg.meta.cv.mip_volume_size(mip), ) - - ws_seg = cv[ - chunk_start[0] : chunk_end[0], - chunk_start[1] : chunk_end[1], - chunk_start[2] : chunk_end[2], - ].squeeze() - + # Coordinates are at the target MIP; get_local_segmentation reads that scale. + ws_seg = get_local_segmentation(cg.meta, chunk_start, chunk_end, mip=mip).squeeze() return ws_seg diff --git a/pychunkedgraph/meshing/meshing_batch.py b/pychunkedgraph/meshing/meshing_batch.py index 6f40fb0a0..926783114 100644 --- a/pychunkedgraph/meshing/meshing_batch.py +++ b/pychunkedgraph/meshing/meshing_batch.py @@ -4,20 +4,22 @@ from cloudfiles import CloudFiles from taskqueue import TaskQueue, LocalTaskQueue -from pychunkedgraph.graph.chunkedgraph import ChunkedGraph # noqa +from pychunkedgraph.graph.chunkedgraph import ChunkedGraph # noqa from pychunkedgraph.meshing.meshing_sqs import MeshTask from pychunkedgraph.meshing import meshgen_utils # noqa if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--queue_name', type=str, default=None) - parser.add_argument('--chunk_start', nargs=3, type=int) - parser.add_argument('--chunk_end', nargs=3, type=int) - parser.add_argument('--cg_name', type=str) - parser.add_argument('--layer', type=int) - parser.add_argument('--mip', type=int) - parser.add_argument('--skip_cache', action='store_true') - parser.add_argument('--overwrite', type=bool, default=False) + parser.add_argument("--queue_name", type=str, default=None) + parser.add_argument("--cg_name", type=str) + parser.add_argument("--layer", type=int) + parser.add_argument("--mip", type=int) + parser.add_argument("--skip_cache", action="store_true") + parser.add_argument( + "--skip", + action="store_true", + help="do not queue a chunk whose shard already exists", + ) args = parser.parse_args() cache = not args.skip_cache @@ -27,31 +29,23 @@ f"graphene://https://localhost/segmentation/table/dummy", info=meshgen_utils.get_json_info(cg), ) - dst = os.path.join( - cv.cloudpath, cv.mesh.meta.mesh_path, "initial", str(args.layer) - ) + dst = os.path.join(cv.cloudpath, cv.mesh.meta.mesh_path, "initial", str(args.layer)) cf = CloudFiles(dst) - if len(list(cf.list())) > 0 and not args.overwrite: - raise ValueError(f"Destination {dst} is not empty. Use `--overwrite true` to proceed anyway.") - - chunks_arr = [] - for x in range(args.chunk_start[0],args.chunk_end[0]): - for y in range(args.chunk_start[1], args.chunk_end[1]): - for z in range(args.chunk_start[2], args.chunk_end[2]): - chunks_arr.append((x, y, z)) + bounds = cg.meta.layer_chunk_bounds[args.layer] + chunks_arr = np.indices(tuple(int(b) for b in bounds)).reshape(3, -1).T np.random.shuffle(chunks_arr) class MeshTaskIterator(object): def __init__(self, chunks): self.chunks = chunks + def __iter__(self): - if args.overwrite: - meshed = set() - else: - meshed = set(cf.list()) - for chunk in self.chunks: - chunk_id = cg.get_chunk_id(layer=args.layer, x=chunk[0], y=chunk[1], z=chunk[2]) + meshed = set(cf.list()) if args.skip else set() + for x, y, z in self.chunks: + chunk_id = cg.get_chunk_id( + layer=args.layer, x=int(x), y=int(y), z=int(z) + ) shard_filename = cv.mesh.readers[args.layer].get_filename(chunk_id) if shard_filename in meshed: continue @@ -62,4 +56,4 @@ def __iter__(self): tq.insert_all(MeshTaskIterator(chunks_arr)) else: tq = LocalTaskQueue(parallel=1) - tq.insert_all(MeshTaskIterator(chunks_arr)) \ No newline at end of file + tq.insert_all(MeshTaskIterator(chunks_arr)) diff --git a/pychunkedgraph/meshing/meshing_sqs.py b/pychunkedgraph/meshing/meshing_sqs.py index b302a1744..59ce840ee 100644 --- a/pychunkedgraph/meshing/meshing_sqs.py +++ b/pychunkedgraph/meshing/meshing_sqs.py @@ -1,6 +1,37 @@ +import traceback +import multiprocessing + +import numpy as np from taskqueue import RegisteredTask + from pychunkedgraph.meshing import meshgen -import numpy as np +from pychunkedgraph.meshing.stitch import chunk_initial_sharded_stitching_task_mp + + +def _mesh_chunk(cg_name, layer, chunk_id, mip, cache): + """Mesh one chunk. Top-level (not a method/closure) so it is picklable as the + forked-child target. Resets the inherited cloudfiles connection pool so the + child's reads use fresh sockets, not the parent's.""" + from cloudfiles import reset_connection_pools + + reset_connection_pools() + chunk_id = np.uint64(chunk_id) + if layer == 2: + return meshgen.chunk_initial_mesh_task( + cg_name, chunk_id, None, mip=mip, sharded=True, cache=cache + ) + return chunk_initial_sharded_stitching_task_mp(cg_name, chunk_id, mip, cache=cache) + + +def _run_in_child(conn, fn, args): + """Forked-child entry: run ``fn``, send the outcome back, exit. Module-level so + the forked child can import it as the Process target.""" + try: + conn.send(("ok", fn(*args))) + except BaseException: # pylint: disable=broad-except + conn.send(("error", traceback.format_exc())) + finally: + conn.close() class MeshTask(RegisteredTask): @@ -8,22 +39,29 @@ def __init__(self, cg_name, layer, chunk_id, mip, cache=True): super().__init__(cg_name, layer, chunk_id, mip, cache) def execute(self): - cg_name = self.cg_name - chunk_id = np.uint64(self.chunk_id) - mip = self.mip - layer = self.layer - if layer == 2: - result = meshgen.chunk_initial_mesh_task( - cg_name, - chunk_id, - None, - mip=mip, - sharded=True, - cache=self.cache + """Mesh the chunk in a fresh forked child that exits when done. A long-lived + poll process accumulates dirty network state (grpc/s2n/SSL/tensorstore + + cloudfiles sockets) and unreleased heap across chunks; a worker pool forked + from that dirty parent inherits both, which corrupts sharded reads and bloats + memory. The child's exit returns all of it to the OS, so the next chunk forks + from a clean parent. Re-raises the child's failure so the poll loop fails the + task loudly (and restarts the pod).""" + ctx = multiprocessing.get_context("fork") + parent_conn, child_conn = ctx.Pipe(duplex=False) + args = (self.cg_name, self.layer, int(self.chunk_id), self.mip, self.cache) + proc = ctx.Process(target=_run_in_child, args=(child_conn, _mesh_chunk, args)) + proc.start() + child_conn.close() + try: + status, payload = parent_conn.recv() + except EOFError: + proc.join() + raise RuntimeError( + f"mesh child died without a result (exitcode {proc.exitcode})" ) - else: - result = meshgen.chunk_initial_sharded_stitching_task( - cg_name, chunk_id, mip, cache=self.cache - ) - print(result) - + finally: + parent_conn.close() + proc.join() + if status == "error": + raise RuntimeError(f"mesh child raised:\n{payload}") + return payload diff --git a/pychunkedgraph/meshing/meshlabserver.py b/pychunkedgraph/meshing/meshlabserver.py index 3065d6707..65c7439db 100644 --- a/pychunkedgraph/meshing/meshlabserver.py +++ b/pychunkedgraph/meshing/meshlabserver.py @@ -3,7 +3,7 @@ import glob import numpy as np -from multiwrapper import multiprocessing_utils as mu +from concurrent.futures import ProcessPoolExecutor HOME = os.path.expanduser("~") @@ -12,17 +12,19 @@ def run_meshlab_script(script_name, arg_dict): - """ Runs meshlabserver script --headless + """Runs meshlabserver script --headless No X-Server required :param script_name: str :param arg_dict: dict [str: str] """ - arg_string = "".join(["-{0} {1} ".format(k, arg_dict[k]) - for k in arg_dict.keys()]) - command = "xvfb-run --auto-servernum --server-num=1 meshlabserver -s {0}/{1} {2}".\ - format(path_to_scripts, script_name, arg_string) + arg_string = "".join(["-{0} {1} ".format(k, arg_dict[k]) for k in arg_dict.keys()]) + command = ( + "xvfb-run --auto-servernum --server-num=1 meshlabserver -s {0}/{1} {2}".format( + path_to_scripts, script_name, arg_string + ) + ) p = subprocess.Popen(command, shell=True, stderr=subprocess.PIPE) p.wait() @@ -31,8 +33,9 @@ def _run_meshlab_script_on_dir_thread(args): script_name, path_block, out_dir, suffix, arg_dict = args for path in path_block: - out_path = "{}/{}{}.obj".format(out_dir, - "".join(os.path.basename(path).split(".")[:-1]), suffix) + out_path = "{}/{}{}.obj".format( + out_dir, "".join(os.path.basename(path).split(".")[:-1]), suffix + ) this_arg_dict = {"i": path, "o": out_path} this_arg_dict.update(arg_dict) @@ -40,8 +43,9 @@ def _run_meshlab_script_on_dir_thread(args): run_meshlab_script(script_name, this_arg_dict) -def run_meshlab_script_on_dir(script_name, in_dir, out_dir, suffix, arg_dict={}, - n_threads=1): +def run_meshlab_script_on_dir( + script_name, in_dir, out_dir, suffix, arg_dict={}, n_threads=1 +): paths = glob.glob(in_dir + "/*.obj") print(len(paths)) @@ -60,10 +64,8 @@ def run_meshlab_script_on_dir(script_name, in_dir, out_dir, suffix, arg_dict={}, multi_args.append([script_name, path_block, out_dir, suffix, arg_dict]) if n_threads == 1: - mu.multiprocess_func(_run_meshlab_script_on_dir_thread, - multi_args, debug=True, - verbose=True, n_threads=1) + for args in multi_args: + _run_meshlab_script_on_dir_thread(args) else: - mu.multisubprocess_func(_run_meshlab_script_on_dir_thread, - multi_args, n_threads=n_threads) - + with ProcessPoolExecutor(max_workers=n_threads) as executor: + list(executor.map(_run_meshlab_script_on_dir_thread, multi_args)) diff --git a/pychunkedgraph/meshing/meta.py b/pychunkedgraph/meshing/meta.py new file mode 100644 index 000000000..fe9052a5e --- /dev/null +++ b/pychunkedgraph/meshing/meta.py @@ -0,0 +1,65 @@ +"""MeshConfig dataclass — single source of truth for per-CG mesh setup values. + +Read from the dataset yaml under a ``mesh_config:`` block, exactly like +``OcdbtConfig`` is read from ``ocdbt_config:``. Every static field is +required — the helper that applies it (``setup_mesh_meta``) does not +substitute defaults for missing yaml entries. The only optional field +is :attr:`dynamic_mesh_dir`, which is graph-id-derived and filled in by +:meth:`with_graph_id` when omitted from the yaml. The mesh chunk_size is +derived (CG CHUNK_SIZE / per-axis downsample at ``mip``), not configured. + +Example yaml block:: + + mesh_config: + dir: graphene_meshes + mip: 0 + max_layer: 6 + max_error: 40 + minishard_bits: {2: 1, 3: 3, 4: 6, 5: 9, 6: 12} + # dynamic_mesh_dir: my_custom_dir # optional; default "dynamic_" +""" + +from dataclasses import asdict, dataclass, replace +from typing import Dict, Optional + + +@dataclass +class MeshConfig: + """Per-CG mesh setup config. See module docstring for yaml schema.""" + + dir: str + mip: int + max_layer: int + max_error: int + minishard_bits: Dict[int, int] + dynamic_mesh_dir: Optional[str] = None + + @classmethod + def from_dict(cls, d: Dict) -> "MeshConfig": + """Build from a yaml-parsed dict. + + Unknown keys are dropped (so older yamls don't break newer code). + ``minishard_bits`` keys are coerced to ``int`` so the yaml is + tolerant of bare-int vs quoted-string keys. + """ + if not d: + raise ValueError( + "MeshConfig.from_dict: empty config — yaml `mesh_config:` " + "block is required" + ) + known = {f for f in cls.__dataclass_fields__} + kwargs = {k: v for k, v in d.items() if k in known} + if "minishard_bits" in kwargs: + kwargs["minishard_bits"] = { + int(k): int(v) for k, v in kwargs["minishard_bits"].items() + } + return cls(**kwargs) + + def with_graph_id(self, graph_id: str) -> "MeshConfig": + """Return a copy with ``dynamic_mesh_dir`` filled in if unset.""" + if self.dynamic_mesh_dir is not None: + return self + return replace(self, dynamic_mesh_dir=f"dynamic_{graph_id}") + + def to_dict(self) -> Dict: + return asdict(self) diff --git a/pychunkedgraph/meshing/setup.py b/pychunkedgraph/meshing/setup.py new file mode 100644 index 000000000..c126bba1c --- /dev/null +++ b/pychunkedgraph/meshing/setup.py @@ -0,0 +1,127 @@ +"""One-shot mesh metadata setup for a CG. + +Writes every mesh-related field a new or freshly-copied graph needs +before any mesh fragment can be served: + + 1. ``cg.meta.ws_cv.info["mesh"]`` (mesh dir in the watershed cv info.json) + 2. ``cg.meta.ws_cv.mesh.meta.info`` (per-layer sharded mesh spec) + 3. ``cg.meta.ws_cv.info["mesh_metadata"]`` (uniform draco grid + dynamic dir) + 4. ``cg.meta.custom_data["mesh"]`` (CG bigtable meta block) + +Idempotent: re-running overwrites the same fields with the current +inputs, with one exception — ``initial_ts`` is preserved if already +set, because changing it after the fact would reclassify every node +id and silently break served manifests. +""" + +import logging +from datetime import datetime, timezone + +from ..graph.chunkedgraph import ChunkedGraph +from .meshgen import get_draco_encoding_settings_for_chunk +from .meshgen_utils import get_mesh_block_shape_for_mip +from .meta import MeshConfig + +logger = logging.getLogger(__name__) + + +def derive_initial_ts(cg: ChunkedGraph) -> int: + """Unix-seconds boundary for ``mesh.initial_ts`` (see ``segregate_node_ids``). + + ``get_earliest_timestamp`` returns the first edit, or — pre-edit — the + ingest-completion boundary stamped during the root-layer build. ``+1`` makes the + second-granularity threshold strictly above the last initial root (the check + is ``<`` and ``int()`` truncates). + """ + earliest = cg.get_earliest_timestamp() + if earliest <= datetime.fromtimestamp(0, tz=timezone.utc): + raise RuntimeError( + "derive_initial_ts: no operations and no ingest earliest_ts stamped" + ) + return int(earliest.timestamp()) + 1 + + +def setup_mesh_meta( + cg: ChunkedGraph, + mesh_config: MeshConfig, +) -> dict: + """Write every mesh.* metadata field this graph needs to serve meshes. + + Writes go to two places: the watershed CloudVolume (steps 1-3, via + ``info.json`` / ``mesh/info`` on GCS) and the CG's bigtable meta + block (step 4). + + ``initial_ts`` is set once and never overwritten — if the existing + bigtable mesh meta already has one, it is reused as-is. Otherwise + it is derived via :func:`derive_initial_ts` and persisted. + + Returns the mesh meta dict persisted into bigtable. + """ + cfg = mesh_config.with_graph_id(cg.graph_id) + n_scales = len(cg.meta.ws_cv.info["scales"]) + if not 0 <= cfg.mip < n_scales: + raise ValueError( + f"mesh_config.mip {cfg.mip} exceeds watershed scales (available 0..{n_scales - 1})" + ) + existing_mesh = cg.meta.custom_data.get("mesh", {}) + existing_ts = existing_mesh.get("initial_ts") + initial_ts = int(existing_ts) if existing_ts is not None else derive_initial_ts(cg) + if existing_ts is not None: + logger.info("preserving existing initial_ts=%d", initial_ts) + + # 1. watershed CV info — mesh dir. + cg.meta.ws_cv.info["mesh"] = cfg.dir + cg.meta.ws_cv.commit_info() + logger.info("wrote ws_cv.info['mesh']=%r", cfg.dir) + + # 2. sharded mesh spec — same template per layer, layer-specific bits. + layer_shard_spec = { + "@type": "neuroglancer_uint64_sharded_v1", + "preshift_bits": 0, + "hash": "murmurhash3_x86_128", + "shard_bits": 0, + "minishard_index_encoding": "gzip", + "data_encoding": "raw", + } + sharding = { + str(layer): {**layer_shard_spec, "minishard_bits": int(bits)} + for layer, bits in cfg.minishard_bits.items() + if layer <= cfg.max_layer + } + mesh_chunk_size = get_mesh_block_shape_for_mip(cg, 2, cfg.mip) + mesh_spec = { + "@type": "neuroglancer_legacy_mesh", + "spatial_index": None, + "mip": int(cfg.mip), + "chunk_size": [int(x) for x in mesh_chunk_size], + "sharding": sharding, + } + cg.meta.ws_cv.mesh.meta.info = mesh_spec + cg.meta.ws_cv.mesh.meta.commit_info() + logger.info("wrote sharded mesh spec for layers %s", sorted(sharding.keys())) + + # 3. uniform draco grid size, derived from layer-2 draco settings. + draco = get_draco_encoding_settings_for_chunk( + cg, cg.get_chunk_id(layer=2, x=0, y=0, z=0), mip=cfg.mip + ) + grid_size = draco["quantization_range"] / (2 ** draco["quantization_bits"] - 1) + cg.meta.ws_cv.info["mesh_metadata"] = { + "uniform_draco_grid_size": grid_size, + "unsharded_mesh_dir": "dynamic", + } + cg.meta.ws_cv.commit_info() + logger.info("wrote mesh_metadata uniform_draco_grid_size=%s", grid_size) + + # 4. CG-side bigtable meta block. + mesh_meta = { + "max_layer": int(cfg.max_layer), + "dynamic_mesh_dir": cfg.dynamic_mesh_dir, + "mip": int(cfg.mip), + "max_error": int(cfg.max_error), + "dir": cfg.dir, + "initial_ts": int(initial_ts), + } + cg.meta.custom_data["mesh"] = mesh_meta + cg.update_meta(cg.meta, overwrite=True) + logger.info("wrote cg.meta.custom_data['mesh']=%r", mesh_meta) + return mesh_meta diff --git a/pychunkedgraph/meshing/stitch/NOTES.md b/pychunkedgraph/meshing/stitch/NOTES.md new file mode 100644 index 000000000..0ddd8c4dd --- /dev/null +++ b/pychunkedgraph/meshing/stitch/NOTES.md @@ -0,0 +1,54 @@ +# stitch — notes (known issues, test handles, future work) + +Secondary material that doesn't belong in `README.md` (the design reference). + +## Large-chunk test handle + +`experimental_pinky_lgsv_split_v0` layer-6 chunk `451485862643892224` — ~28k +parents, ~5.5 GB shard, an ~83M-vertex giant fragment. One of the larger shards in +that dataset; a good memory/perf optimization target. Reproduce with +`debug.repro_chunk(cg, 451485862643892224)`. + +## Known issue — peak memory OOM on small nodes + +Heavy layer-6 chunks drive whole-cgroup peak RSS above a small node's limit +(~40 GB peak at 24 workers vs a 28 GB node), OOM-ing a worker mid-run. Reproduces +only under that memory pressure — the chunk above completes cleanly with more +headroom / fewer workers. + +Measured split (the chunk above, 16 workers): per-worker peak RSS p50 ≈ 0.75 GB, +max ≈ 2.2 GB (the giant's worker). The peak is dominated by the **parent synthesize +transient** — `acc.merged_meshes` accumulates every encoded mesh, then +`synthesize_shard` holds ~2-3× the ~5.5 GB shard — **plus** the broad sum of +concurrent per-worker working sets (~0.75 GB × N). No single giant worker dominates; +both terms stack to ~40 GB at 24 workers. + +Eager `del` of merge/decoded-fragment intermediates (in `worker._stitch_one` and +`utils.merge_draco_meshes_across_boundaries_pure`) trims only a few GB — not enough +to close the gap on its own. + +## Fork-per-chunk invariant — keep the consumer single-threaded + +`MeshTask.execute` forks each chunk into a child that exits (clean grpc/s2n/heap +per chunk). The fork is only deadlock-safe while the consumer process is +single-threaded at fork time — the prod consumer (`meshing/mesh_worker.py`, +`TaskQueue(..., n_threads=0)`) is. Do not give the consumer worker threads +(`n_threads>0`) or otherwise spawn a background thread before `execute`: forking a +multi-threaded process can leave the child holding a lock no thread will release. +The inner stitch `mp.Pool` must keep forking before any cloud I/O in the child +(`task.py` comment at the pool construction) for the same s2n-atfork reason. + +## Future levers (neither done; parent term is the bigger one) + +- **Stream the shard via a tmp dir.** Workers write each encoded mesh to a + `mesh_path//