Pcgv2 ingest

.dockerignore

@@ -0,0 +1,124 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+.idea/*
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# Visual Code
+.vscode/
+
+# terraform
+.terraform/
+*.lock.hcl
+*.tfstate
+*.tfstate.*
+
+
+# local dev stuff
+.devcontainer/
+*.ipynb
+*.rdb
+/protobuf*
+
+# Git
+.git/

Dockerfile

@@ -1,9 +1,8 @@
-FROM tiangolo/uwsgi-nginx-flask:python3.7
+ARG PYTHON_VERSION=3.11
+FROM tiangolo/uwsgi-nginx-flask:python${PYTHON_VERSION}
 
 ENV GIT_SSL_NO_VERIFY=1
 RUN mkdir -p /home/nginx/.cloudvolume/secrets && chown -R nginx /home/nginx && usermod -d /home/nginx -s /bin/bash nginx
 
 COPY . /app
-RUN pip install pip==20.0.1 \
-    && pip install --no-cache-dir --upgrade -r requirements.txt \
-    && pip install --upgrade git+https://github.com/seung-lab/KVDbClient.git@main
+RUN pip install --no-cache-dir --upgrade -r requirements.txt

cloudbuild.v1.ingest.yaml

@@ -7,23 +7,23 @@ steps:
     timeout: 600s
     args:
       - build
-      - "--tag=gcr.io/neuromancer-seung-import/pcgl2cache:ingest_v1_$TAG_NAME"
-      - "--file=./ingest.Dockerfile"
+      - "--tag=gcr.io/$PROJECT_ID/pcgl2cache:ingest_v1"
+      - "--file=./ingest.v1.Dockerfile"
       - .
   - name: "gcr.io/cloud-builders/docker"
     entrypoint: "bash"
     args:
       [
         "-c",
-        "docker tag gcr.io/$PROJECT_ID/pcgl2cache:ingest_v1_$TAG_NAME $$USERNAME/pcgl2cache:ingest_v1_$TAG_NAME",
+        "docker tag gcr.io/$PROJECT_ID/pcgl2cache:ingest_v1 $$USERNAME/pcgl2cache:ingest_v1",
       ]
     secretEnv: ["USERNAME"]
   - name: "gcr.io/cloud-builders/docker"
     entrypoint: "bash"
-    args: ["-c", "docker push $$USERNAME/pcgl2cache:ingest_v1_$TAG_NAME"]
+    args: ["-c", "docker push $$USERNAME/pcgl2cache:ingest_v1"]
     secretEnv: ["USERNAME"]
 images:
-  - "gcr.io/neuromancer-seung-import/pcgl2cache:ingest_v1_$TAG_NAME"
+  - "gcr.io/$PROJECT_ID/pcgl2cache:ingest_v1"
 availableSecrets:
   secretManager:
     - versionName: projects/$PROJECT_ID/secrets/docker-password/versions/1

cloudbuild.v1.worker.yaml

@@ -7,8 +7,8 @@ steps:
     timeout: 600s
     args:
       - build
-      - "--tag=gcr.io/neuromancer-seung-import/pcgl2cache:worker_pcgv1_$TAG_NAME"
-      - "--file=./ingest.Dockerfile"
+      - "--tag=gcr.io/$PROJECT_ID/pcgl2cache:worker_pcgv1_$TAG_NAME"
+      - "--file=./ingest.v1.Dockerfile"
       - .
   - name: "gcr.io/cloud-builders/docker"
     entrypoint: "bash"
@@ -23,7 +23,7 @@ steps:
     args: ["-c", "docker push $$USERNAME/pcgl2cache:worker_pcgv1_$TAG_NAME"]
     secretEnv: ["USERNAME"]
 images:
-  - "gcr.io/neuromancer-seung-import/pcgl2cache:worker_pcgv1_$TAG_NAME"
+  - "gcr.io/$PROJECT_ID/pcgl2cache:worker_pcgv1_$TAG_NAME"
 availableSecrets:
   secretManager:
     - versionName: projects/$PROJECT_ID/secrets/docker-password/versions/1

cloudbuild.v2.ingest.yaml

@@ -0,0 +1,32 @@
+steps:
+  - name: "gcr.io/cloud-builders/docker"
+    entrypoint: "bash"
+    args: ["-c", "docker login --username=$$USERNAME --password=$$PASSWORD"]
+    secretEnv: ["USERNAME", "PASSWORD"]
+  - name: "gcr.io/cloud-builders/docker"
+    timeout: 600s
+    args:
+      - build
+      - "--tag=gcr.io/$PROJECT_ID/pcgl2cache:ingest_v2"
+      - "--file=./ingest.v2.Dockerfile"
+      - .
+  - name: "gcr.io/cloud-builders/docker"
+    entrypoint: "bash"
+    args:
+      [
+        "-c",
+        "docker tag gcr.io/$PROJECT_ID/pcgl2cache:ingest_v2 $$USERNAME/pcgl2cache:ingest_v2",
+      ]
+    secretEnv: ["USERNAME"]
+  - name: "gcr.io/cloud-builders/docker"
+    entrypoint: "bash"
+    args: ["-c", "docker push $$USERNAME/pcgl2cache:ingest_v2"]
+    secretEnv: ["USERNAME"]
+images:
+  - "gcr.io/$PROJECT_ID/pcgl2cache:ingest_v2"
+availableSecrets:
+  secretManager:
+    - versionName: projects/$PROJECT_ID/secrets/docker-password/versions/1
+      env: "PASSWORD"
+    - versionName: projects/$PROJECT_ID/secrets/docker-username/versions/1
+      env: "USERNAME"

cloudbuild.v2.worker.yaml

@@ -0,0 +1,32 @@
+steps:
+  - name: "gcr.io/cloud-builders/docker"
+    entrypoint: "bash"
+    args: ["-c", "docker login --username=$$USERNAME --password=$$PASSWORD"]
+    secretEnv: ["USERNAME", "PASSWORD"]
+  - name: "gcr.io/cloud-builders/docker"
+    timeout: 600s
+    args:
+      - build
+      - "--tag=gcr.io/$PROJECT_ID/pcgl2cache:worker_pcgv2_$TAG_NAME"
+      - "--file=./ingest.v2.Dockerfile"
+      - .
+  - name: "gcr.io/cloud-builders/docker"
+    entrypoint: "bash"
+    args:
+      [
+        "-c",
+        "docker tag gcr.io/$PROJECT_ID/pcgl2cache:worker_pcgv2_$TAG_NAME $$USERNAME/pcgl2cache:worker_pcgv2_$TAG_NAME",
+      ]
+    secretEnv: ["USERNAME"]
+  - name: "gcr.io/cloud-builders/docker"
+    entrypoint: "bash"
+    args: ["-c", "docker push $$USERNAME/pcgl2cache:worker_pcgv2_$TAG_NAME"]
+    secretEnv: ["USERNAME"]
+images:
+  - "gcr.io/$PROJECT_ID/pcgl2cache:worker_pcgv2_$TAG_NAME"
+availableSecrets:
+  secretManager:
+    - versionName: projects/$PROJECT_ID/secrets/docker-password/versions/1
+      env: "PASSWORD"
+    - versionName: projects/$PROJECT_ID/secrets/docker-username/versions/1
+      env: "USERNAME"

ingest.v1.Dockerfile

@@ -0,0 +1,8 @@
+FROM caveconnectome/pychunkedgraph:base_042124
+
+ENV GIT_SSL_NO_VERIFY=1
+RUN mkdir -p /home/nginx/.cloudvolume/secrets && chown -R nginx /home/nginx && usermod -d /home/nginx -s /bin/bash nginx
+
+COPY . /app
+RUN pip install --no-cache-dir --upgrade -r requirements.txt \
+    && pip install --upgrade git+https://github.com/CAVEconnectome/PyChunkedGraph.git@pcgv1

ingest.v2.Dockerfile

@@ -0,0 +1,13 @@
+FROM caveconnectome/pychunkedgraph:base_042124
+ENV VIRTUAL_ENV=/app/venv
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+ENV GIT_SSL_NO_VERIFY=1
+ENV CHUNKEDGRAPH_VERSION=2
+
+RUN mkdir -p /home/nginx/.cloudvolume/secrets && chown -R nginx /home/nginx && usermod -d /home/nginx -s /bin/bash nginx
+
+COPY requirements.txt .
+RUN pip install --upgrade --no-cache-dir -r requirements.txt \
+    && pip install --upgrade git+https://github.com/CAVEconnectome/PyChunkedGraph.git@main
+
+COPY . /app

pcgl2cache/core/attributes.py

@@ -1,7 +1,7 @@
 import numpy as np
 from kvdbclient.serializers import NumPyArray
 from kvdbclient.serializers import NumPyValue
-from kvdbclient.bigtable.attributes import Attribute
+from kvdbclient.attributes import _Attribute as Attribute
 
 UINT64 = np.dtype("uint64").newbyteorder("L")
 UINT32 = np.dtype("uint32").newbyteorder("L")

pcgl2cache/core/features.py

@@ -6,6 +6,7 @@
 from edt import edt
 from sklearn import decomposition
 from kvdbclient import BigTableClient
+from kvdbclient.base import serialize_uint64
 from cloudvolume import CloudVolume
 
 from . import attributes
@@ -103,7 +104,9 @@ def get_remapped_segmentation(self, l2id=None):
 
 def _get_l2_ids(l2vol: L2ChunkVolume, svids: np.array) -> np.array:
     if l2vol.cg:
-        l2ids = l2vol.cg.get_roots(svids, stop_layer=2, time_stamp=l2vol.timestamp)
+        l2ids = l2vol.cg.get_roots(
+            svids, stop_layer=2, fail_to_zero=True, time_stamp=l2vol.timestamp
+        )
         layers = l2vol.cg.get_chunk_layers(l2ids)
         sv_mask = layers == 1
         l2ids[sv_mask] = 0
@@ -312,10 +315,7 @@ def run_l2cache(
 
 
 def write_to_db(client: BigTableClient, result_d: dict) -> None:
-    from kvdbclient.base import Entry
-    from kvdbclient.base import EntryKey
-
-    entries = []
+    rows = []
     for tup in zip(*result_d.values()):
         (
             l2id,
@@ -338,5 +338,5 @@ def write_to_db(client: BigTableClient, result_d: dict) -> None:
             attributes.PCA: pca_comp,
             attributes.PCA_VAL: pca_vals,
         }
-        entries.append(Entry(EntryKey(l2id), val_d))
-    client.write_entries(entries)
+        rows.append(client.mutate_row(serialize_uint64(l2id), val_d))
+    client.write(rows)

pcgl2cache/ingest/__init__.py

@@ -1,3 +1,4 @@
+from os import environ
 from collections import namedtuple
 from os import environ
 
@@ -6,16 +7,14 @@
 
 _cluster_ingest_config_fields = (
     "REDIS_URL",
-    "BATCH_SIZE",
-    "L2CACHE_Q_NAME",
-    "L2CACHE_Q_LIMIT",  # these limits ensure the queue won't use too much memory
-    "L2CACHE_Q_INTERVAL",  # sleep interval before queuing the next job when limit is reached
+    "QUEUE_NAME",
+    "QUEUE_SIZE",  # these limits ensure the queue won't use too much memory
+    "QUEUE_INTERVAL",  # sleep interval before queuing the next job when limit is reached
 )
 _cluster_ingest_defaults = (
     REDIS_URL,
-    10,
-    "atomic",
-    int(environ.get("L2CACHE_Q_LIMIT", 500000)),
+    "l2",
+    int(environ.get("QUEUE_SIZE", 1000000)),
     60,
 )
 ClusterIngestConfig = namedtuple(