pcgv2 docker and ingest upgrade

.github/workflows/main.yml

@@ -13,10 +13,20 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check out code
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
 
-      - name: Build image and run tests
-        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"
+      - 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
+        run: |
+          docker run --rm seunglab/pychunkedgraph:${{ github.sha }} /bin/sh -c "pytest --cov-config .coveragerc --cov=pychunkedgraph ./pychunkedgraph/tests && codecov"

Dockerfile

@@ -1,11 +1,86 @@
-FROM caveconnectome/pychunkedgraph:base_042124
+# syntax=docker/dockerfile:1
+ARG PYTHON_VERSION=3.11
+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 ./
+
+RUN --mount=type=cache,target=/opt/conda/pkgs \
+    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 on 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.21.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 conda venv; setup.py
+# reads __version__ from pychunkedgraph/__init__.py.
+RUN pip install --no-deps -e .
+
+CMD ["/usr/bin/supervisord", "-n", "-c", "/etc/supervisor/supervisord.conf"]

cloudbuild.yaml

@@ -5,19 +5,24 @@ steps:
     args: ["-c", "docker login --username=$$USERNAME --password=$$PASSWORD"]
     secretEnv: ["USERNAME", "PASSWORD"]
 
+  # Create a buildx builder; the docker-container driver is required to export a
+  # registry-backed layer cache (the default docker driver cannot).
   - name: "gcr.io/cloud-builders/docker"
     entrypoint: "bash"
-    args:
-      - "-c"
-      - |
-        docker build -t $$USERNAME/pychunkedgraph:$TAG_NAME .
-    timeout: 600s
-    secretEnv: ["USERNAME"]
+    args: ["-c", "docker buildx create --name pcgbuilder --driver docker-container --use"]
 
-  # Push the final image to Dockerhub
+  # Build with a registry layer cache and push. mode=max caches intermediate
+  # stages (conda env, cbtemulator) so tagged rebuilds reuse them.
   - name: "gcr.io/cloud-builders/docker"
     entrypoint: "bash"
-    args: ["-c", "docker push $$USERNAME/pychunkedgraph:$TAG_NAME"]
+    args:
+      - "-c"
+      - |
+        docker buildx build \
+          --cache-from type=registry,ref=$$USERNAME/pychunkedgraph:buildcache \
+          --cache-to type=registry,ref=$$USERNAME/pychunkedgraph:buildcache,mode=max \
+          --tag $$USERNAME/pychunkedgraph:$TAG_NAME \
+          --push .
     secretEnv: ["USERNAME"]
 
 availableSecrets:
@@ -26,3 +31,7 @@ availableSecrets:
       env: "PASSWORD"
     - versionName: projects/$PROJECT_ID/secrets/docker-username/versions/1
       env: "USERNAME"
+
+timeout: 3600s
+options:
+  machineType: "E2_HIGHCPU_8"

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;
+        }
+    }
+}

pychunkedgraph/pipeline/ingest/setup.py

@@ -1,13 +1,13 @@
 """One-shot ingest setup — create the Bigtable table and write graph meta.
 
 Run once per graph, inside the image, before any layer Jobs:
-    python -m pychunkedgraph.pipeline.ingest.setup <graph_id> [--raw]
+    python -m pychunkedgraph.pipeline.ingest.setup <graph_id> [--raw] [--exist-ok]
 
 Reads the dataset yaml from its mounted location (no path passed); this is the
 only step that touches the yaml. Folds the agglomeration source into
 ``meta.custom_data["agg"] = {"path": str, "raw": bool}`` so workers read
 everything from Bigtable at run time — no yaml, no Redis. Errors out if the
-table already exists (the operator runs setup explicitly once).
+table already exists, unless ``--exist-ok`` (idempotent re-run for converge/resume).
 """
 
 import argparse
@@ -25,7 +25,12 @@
 DATASET_PATH = environ.get("PCG_DATASET", "/app/datasets/dataset.yml")
 
 
-def setup(graph_id: str, raw: bool = False, dataset_path: str = DATASET_PATH) -> None:
+def setup(
+    graph_id: str,
+    raw: bool = False,
+    exist_ok: bool = False,
+    dataset_path: str = DATASET_PATH,
+) -> None:
     with open(dataset_path) as stream:
         config = yaml.safe_load(stream)
     client_config = BigTableConfig(**config["backend_client"]["CONFIG"])
@@ -35,6 +40,8 @@ def setup(graph_id: str, raw: bool = False, dataset_path: str = DATASET_PATH) ->
     agg = {"path": config.get("ingest_config", {}).get("AGGLOMERATION"), "raw": raw}
     meta = ChunkedGraphMeta(graph_config, data_source, custom_data={"agg": agg})
     cg = ChunkedGraph(meta=meta, client_info=client_info)
+    if exist_ok and cg.client._table.exists():
+        return  # idempotent re-run: the graph is already set up
     cg.create()
 
 
@@ -46,8 +53,13 @@ def main() -> None:
         action="store_true",
         help="raw agglomeration input; the L2 workers convert it to processed data",
     )
+    parser.add_argument(
+        "--exist-ok",
+        action="store_true",
+        help="idempotent re-run: skip creation if the graph table already exists",
+    )
     args = parser.parse_args()
-    setup(args.graph_id, raw=args.raw)
+    setup(args.graph_id, raw=args.raw, exist_ok=args.exist_ok)
 
 
 if __name__ == "__main__":