Spatial aligned svid lookup

.github/workflows/materialization_ci.yml

@@ -21,6 +21,10 @@ jobs:
 
     name: Test against different Python versions
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10"]
+
     services:
       postgres:
         image: postgis/postgis:13-master
@@ -32,11 +36,11 @@ jobs:
           - 5432:5432
         options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
     steps:
-      - uses: actions/checkout@v4
-      - name: Set up Python 3.9
-        uses: actions/setup-python@v5
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
         with:
-          python-version: 3.9
+          python-version: ${{ matrix.python-version }}
       - uses: actions/cache@v2
         with:
           path: ~/.cache/pip
@@ -47,7 +51,7 @@ jobs:
         with:
           auto-update-conda: true
           auto-activate-base: true
-          python-version: 3.9
+          python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         shell: bash -l {0}
         run: |

Dockerfile

@@ -1,4 +1,4 @@
-FROM tiangolo/uwsgi-nginx-flask:python3.9
+FROM tiangolo/uwsgi-nginx-flask:python3.10
 
 ENV UWSGI_INI /app/uwsgi.ini
 RUN mkdir -p /home/nginx/.cloudvolume/secrets \

materializationengine/blueprints/client/query.py

@@ -155,7 +155,7 @@ def _fix_decimal_column(df_col):
     if np.all(is_integer_col(df_col)):
         return df_col.apply(int)
     else:
-        return df_col.apply(np.float)
+        return df_col.apply(float)
 
 
 def get_column(model, column):

materializationengine/blueprints/materialize/api.py

@@ -1,8 +1,8 @@
 import datetime
 import logging
-
+import redis
 from dynamicannotationdb.models import AnalysisTable, Base
-from flask import abort, current_app, request
+from flask import abort, current_app, request, jsonify
 from flask_accepts import accepts
 from flask_restx import Namespace, Resource, inputs, reqparse
 from materializationengine.blueprints.reset_auth import reset_auth
@@ -61,6 +61,13 @@
 materialize_parser.add_argument("days_to_expire", required=True, default=None, type=int)
 materialize_parser.add_argument("merge_tables", required=True, type=inputs.boolean)
 
+
+spatial_svid_parser = reqparse.RequestParser()
+spatial_svid_parser.add_argument("chunk_scale_factor", default=12, type=int)
+spatial_svid_parser.add_argument("get_root_ids", default=True, type=inputs.boolean)
+spatial_svid_parser.add_argument("upload_to_database", default=True, type=inputs.boolean)
+
+
 authorizations = {
     "apikey": {"type": "apiKey", "in": "query", "name": "middle_auth_token"}
 }
@@ -233,11 +240,52 @@ def post(self, datastack_name: str, table_name: str):
         ).apply_async()
         return 200
 
-
+@mat_bp.expect(spatial_svid_parser)
 @mat_bp.route(
-    "/materialize/run/dense_lookup_root_ids/datastack/<string:datastack_name>"
+    "/materialize/run/spatial_lookup/datastack/<string:datastack_name>/<string:table_name>"
 )
-class LookupDenseMissingRootIdsResource(Resource):
+class SpatialSVIDLookupTableResource(Resource):
+    @reset_auth
+    @auth_requires_permission("edit", table_arg="datastack_name")
+    @mat_bp.doc("Lookup spatially chunked svid workflow", security="apikey")
+    def post(self, datastack_name: str, table_name: str):
+        """Process newly added annotations and lookup segmentation data using
+        a spatially chunked svid lookup strategy. Optionally also lookups root ids.
+
+        Args:
+            datastack_name (str): name of datastack from infoservice
+            table_name (str): name of table
+        """
+        from materializationengine.workflows.spatial_lookup import (
+            run_spatial_lookup_workflow,
+        )
+        args = spatial_svid_parser.parse_args()
+
+
+        if datastack_name not in current_app.config["DATASTACKS"]:
+            abort(404, f"datastack {datastack_name} not configured for materialization")
+
+        datastack_info = get_datastack_info(datastack_name)
+
+        chunk_scale_factor = args["chunk_scale_factor"]
+        get_root_ids = args["get_root_ids"]
+        upload_to_database = args["upload_to_database"]
+        try:
+            run_spatial_lookup_workflow.si(
+                datastack_info,
+                table_name=table_name,
+                chunk_scale_factor=chunk_scale_factor,
+                get_root_ids=get_root_ids,
+                upload_to_database=upload_to_database,
+            ).apply_async()
+        except Exception as e:
+            logging.error(e)
+            return abort(400, f"Error running spatial lookup workflow: {e}")
+        return 200
+
+
+@mat_bp.route("/materialize/run/lookup_root_ids/datastack/<string:datastack_name>")
+class LookupMissingRootIdsResource(Resource):
     @reset_auth
     @auth_requires_admin
     @mat_bp.doc("Find all null root ids and lookup new roots", security="apikey")

materializationengine/celery_init.py

@@ -11,6 +11,7 @@
         "materializationengine.workflows.dummy_workflow",
         "materializationengine.workflows.periodic_database_removal",
         "materializationengine.workflows.periodic_materialization",
+        "materializationengine.workflows.spatial_lookup",
         "materializationengine.shared_tasks",
         "materializationengine.views",
         "materializationengine.monitor",

materializationengine/cloudvolume_gateway.py

@@ -0,0 +1,61 @@
+import cloudvolume
+import os
+
+
+class CloudVolumeGateway:
+    """A class to manage cloudvolume clients and cache them for reuse."""
+
+    def __init__(self, lru_bytes: int = 0):
+        self._cv_clients = {}
+        self._lru_bytes = lru_bytes
+
+    def get_cv(self, seg_source: str, mip_level: int = 0) -> cloudvolume.CloudVolume:
+        """A function to get a cloudvolume client for a given source.
+
+        Args:
+            seg_source (str): The cloudvolume source string.
+            mip_level (int, optional): The MIP level to use. Defaults to 0.
+
+        Returns:
+            cloudvolume.CloudVolume: The cloudvolume client.
+        """
+        seg_source_key = seg_source.split("/")[-1]
+        return (
+            self._get_cv_client(seg_source, seg_source_key, mip_level)
+            if seg_source_key not in self._cv_clients
+            else self._cv_clients[seg_source_key]
+        )
+
+    def _get_cv_client(
+        self, seg_source: str, seg_source_key: str, mip_level: int = 0
+    ) -> cloudvolume.CloudVolume:
+        """A helper function to create a cloudvolume client and cache it for reuse.
+
+        Args:
+            seg_source (str): The cloudvolume source string.
+            seg_source_key (str): The cloudvolume source key name to use for caching.
+            mip_level (int, optional): The MIP level to use. Defaults to 0.
+
+        Returns:
+            cloudvolume.CloudVolume: _description_
+        """
+        cv_client = cloudvolume.CloudVolume(
+            seg_source,
+            mip=mip_level,
+            use_https=True,
+            bounded=False,
+            fill_missing=True,
+            lru_bytes=self._lru_bytes,
+        )
+
+        self._cv_clients[seg_source_key] = cv_client
+        return self._cv_clients[seg_source_key]
+
+    def invalidate_cache(self):
+        """Clear the cache of cloudvolume clients."""
+        self._cv_clients = {}
+
+
+cloudvolume_cache = CloudVolumeGateway(
+    lru_bytes=int(os.environ.get("CELERY_CLOUDVOLUME_CACHE_BYTES", 0))
+)

materializationengine/workflows/ingest_new_annotations.py

@@ -655,8 +655,8 @@ def ingest_new_annotations_workflow(mat_metadata: dict):
         return fin.si()
 
 
-@celery.task(name="workflow:create_missing_segmentation_table")
-def create_missing_segmentation_table(mat_metadata: dict) -> dict:
+@celery.task(name="workflow:create_missing_segmentation_table", acks_late=True)
+def create_missing_segmentation_table(mat_metadata: dict) -> bool:
     """Create missing segmentation tables associated with an annotation table if it
     does not already exist.
 
@@ -666,7 +666,7 @@ def create_missing_segmentation_table(mat_metadata: dict) -> dict:
         Materialization metadata
 
     Returns:
-        dict: Materialization metadata
+        bool: if segmentation table was created or already exists
     """
     segmentation_table_name = mat_metadata.get("segmentation_table_name")
     aligned_volume = mat_metadata.get("aligned_volume")
@@ -698,6 +698,8 @@ def create_missing_segmentation_table(mat_metadata: dict) -> dict:
         except Exception as e:
             celery_logger.error(f"SQL ERROR: {e}")
             session.rollback()
+        finally:
+            session.close()
     else:
         session.close()
     return True
@@ -993,8 +995,10 @@ def get_new_root_ids(materialization_data: dict, mat_metadata: dict) -> dict:
 
     cg_client = chunkedgraph_cache.init_pcg(pcg_table_name)
 
-    # filter missing root_ids and lookup root_ids if missing
-    mask = np.logical_and.reduce([root_ids_df[col].isna() for col in cols])
+    # filter missing root_ids and lookup root_ids if missing or zero
+    mask = np.logical_and.reduce(
+        [(root_ids_df[col].isna() | (root_ids_df[col] == 0)) for col in cols]
+    )
     missing_root_rows = root_ids_df.loc[mask]
     if not missing_root_rows.empty:
         supervoxel_data = missing_root_rows.loc[:, supervoxel_col_names]