migrate to the new kvdbclient

Author: akhileshh

Dockerfile

@@ -5,5 +5,4 @@ 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/seung-lab/KVDbClient.git@main
+RUN pip install --no-cache-dir --upgrade -r requirements.txt

ingest.v1.Dockerfile

@@ -5,5 +5,4 @@ RUN mkdir -p /home/nginx/.cloudvolume/secrets && chown -R nginx /home/nginx && u
 
 COPY . /app
 RUN pip install --no-cache-dir --upgrade -r requirements.txt \
-    && pip install --upgrade git+https://github.com/CAVEconnectome/PyChunkedGraph.git@pcgv1 \
-    && pip install --upgrade git+https://github.com/seung-lab/KVDbClient.git@main
+    && pip install --upgrade git+https://github.com/CAVEconnectome/PyChunkedGraph.git@pcgv1

ingest.v2.Dockerfile

@@ -8,7 +8,6 @@ RUN mkdir -p /home/nginx/.cloudvolume/secrets && chown -R nginx /home/nginx && u
 
 COPY requirements.txt .
 RUN pip install --upgrade --no-cache-dir -r requirements.txt \
-    && pip install --upgrade git+https://github.com/CAVEconnectome/PyChunkedGraph.git@main \
-    && pip install --upgrade git+https://github.com/seung-lab/KVDbClient.git@main
+    && 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
@@ -314,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,
@@ -340,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/cli.py

@@ -5,6 +5,7 @@
 import click
 from flask.cli import AppGroup
 
+import pcgl2cache
 from .manager import IngestionManager
 from .redis import get_redis_connection
 
@@ -43,8 +44,9 @@ def ingest_cache(
         from kvdbclient import BigTableClient
         from kvdbclient import get_default_client_info
 
-        client = BigTableClient(config=get_default_client_info().CONFIG)
-        client.create_table(cache_id)
+        meta = {"graph_id": graph_id, "cv_path": cv_path, "timestamp": timestamp}
+        client = BigTableClient(cache_id, config=get_default_client_info().CONFIG)
+        client.create_graph(meta=meta, version=pcgl2cache.__version__)
 
     # example format 2018-06-29 08:15:27
     timestamp = datetime.strptime(timestamp, "%Y-%m-%d %H:%M:%S").replace(
@@ -86,8 +88,9 @@ def ingest_cache_v2(
         from kvdbclient import BigTableClient
         from kvdbclient import get_default_client_info
 
-        client = BigTableClient(config=get_default_client_info().CONFIG)
-        client.create_table(cache_id)
+        meta = {"graph_id": graph_id, "cv_path": cv_path, "timestamp": timestamp}
+        client = BigTableClient(cache_id, config=get_default_client_info().CONFIG)
+        client.create_graph(meta=meta, version=pcgl2cache.__version__)
 
     # example format 2022-06-29 08:15:27
     timestamp = datetime.strptime(timestamp, "%Y-%m-%d %H:%M:%S").replace(

requirements.txt

@@ -17,6 +17,7 @@ multiwrapper
 middle-auth-client>=3.11.0
 python-json-logger
 zstandard
+kvdbclient>=0.1.1
 edt
 pyyaml
 cachetools

workers/common.py

@@ -3,8 +3,7 @@ def calculate_features(l2ids, l2cache_id, cv_path):
     import numpy as np
     from cloudvolume import CloudVolume
     from kvdbclient import BigTableClient, get_default_client_info
-    from kvdbclient.base import Entry
-    from kvdbclient.base import EntryKey
+    from kvdbclient.base import serialize_uint64
     from pcgl2cache.core.attributes import SIZE_NM3
     from pcgl2cache.core.features import run_l2cache
     from pcgl2cache.core.features import write_to_db
@@ -17,7 +16,7 @@ def calculate_features(l2ids, l2cache_id, cv_path):
             continue
         result = run_l2cache(cv, l2id=_id)
         if not result:
-            entry = Entry(EntryKey(_id), {SIZE_NM3: np.uint64(0)})
-            client.write_entries([entry])
+            row = client.mutate_row(serialize_uint64(_id), {SIZE_NM3: np.uint64(0)})
+            client.write([row])
         write_to_db(client, result)
         gc.collect()