Merge branch 'master' of github.com:CAVEconnectome/MaterializationEngine

Author: fcollman

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 5.18.0
+current_version = 5.20.1
 commit = True
 tag = True
 

Dockerfile

@@ -33,7 +33,7 @@ COPY uv.lock pyproject.toml ./
 
 ENV UV_PROJECT_ENVIRONMENT="/usr/local/"
 RUN --mount=type=cache,target=/root/.cache/uv \
-  UV_VENV_ARGS="--system-site-packages" uv sync --frozen --no-install-project --no-default-groups
+  UV_VENV_ARGS="--system-site-packages" uv sync --frozen --no-install-project --no-default-groups --group deploy
 
 # COPY . ./
 # RUN --mount=type=cache,target=/root/.cache/uv \

docs/conf.py

@@ -23,7 +23,7 @@
 author = "Derrick Brittain,  Forrest Collman, Sven Dorkenwald"
 
 # The full version, including alpha/beta/rc tags
-release = "5.18.0"
+release = "5.20.1"
 
 
 

materializationengine/__init__.py

@@ -1 +1 @@
-__version__ = "5.18.0"
+__version__ = "5.20.1"

materializationengine/blueprints/client/api.py

@@ -32,7 +32,7 @@
 from materializationengine.models import MaterializedMetadata
 from materializationengine.schemas import AnalysisTableSchema, AnalysisVersionSchema
 
-__version__ = "5.18.0"
+__version__ = "5.20.1"
 
 
 authorizations = {

materializationengine/blueprints/client/api2.py

@@ -71,7 +71,7 @@
 from materializationengine.schemas import AnalysisTableSchema, AnalysisVersionSchema
 from materializationengine.utils import check_read_permission
 
-__version__ = "5.18.0"
+__version__ = "5.20.1"
 
 
 authorizations = {

materializationengine/blueprints/materialize/api.py

@@ -1,41 +1,43 @@
 import datetime
 import logging
+import os
+import subprocess
+
+import cloudfiles
 import redis
-from dynamicannotationdb.models import AnalysisTable, Base
-from flask import abort, current_app, request, jsonify
+from dynamicannotationdb.models import AnalysisTable, AnalysisVersion, Base
+from flask import abort, current_app, jsonify, request
 from flask_accepts import accepts
-from flask_restx import Namespace, Resource, inputs, reqparse, fields
+from flask_restx import Namespace, Resource, fields, inputs, reqparse
+from middle_auth_client import (
+    auth_requires_admin,
+    auth_requires_dataset_admin,
+    auth_requires_permission,
+)
+from sqlalchemy import MetaData, Table
+from sqlalchemy.engine.url import make_url
+from sqlalchemy.exc import NoSuchTableError
+
 from materializationengine.blueprints.client.utils import get_latest_version
+from materializationengine.blueprints.materialize.schemas import (
+    AnnotationIDListSchema,
+    BadRootsSchema,
+    VirtualVersionSchema,
+)
 from materializationengine.blueprints.reset_auth import reset_auth
 from materializationengine.database import (
-    dynamic_annotation_cache,
     db_manager,
+    dynamic_annotation_cache,
 )
 from materializationengine.info_client import (
     get_aligned_volumes,
     get_datastack_info,
     get_relevant_datastack_info,
 )
-from dynamicannotationdb.models import AnalysisVersion
 from materializationengine.schemas import AnalysisTableSchema, AnalysisVersionSchema
-from materializationengine.blueprints.materialize.schemas import BadRootsSchema
-from middle_auth_client import auth_requires_admin, auth_requires_permission, auth_requires_dataset_admin
-from sqlalchemy import MetaData, Table
-from sqlalchemy.engine.url import make_url
-from sqlalchemy.exc import NoSuchTableError
 from materializationengine.utils import check_write_permission
-import os
-import subprocess
-import cloudfiles
-
-
-from materializationengine.blueprints.materialize.schemas import (
-    VirtualVersionSchema,
-    AnnotationIDListSchema,
-)
-
 
-__version__ = "5.18.0"
+__version__ = "5.20.1"
 
 
 bulk_upload_parser = reqparse.RequestParser()
@@ -715,7 +717,7 @@ class TableResource(Resource):
     @mat_bp.doc("get_all_tables", security="apikey")
     def get(self, aligned_volume_name, version):
         check_aligned_volume(aligned_volume_name)
-        
+
         with db_manager.session_scope(aligned_volume_name) as session:
             response = (
                 session.query(AnalysisTable)
@@ -742,11 +744,11 @@ class AnnotationResource(Resource):
     @mat_bp.doc("get_top_materialized_annotations", security="apikey")
     def get(self, aligned_volume_name: str, version: int, tablename: str):
         check_aligned_volume(aligned_volume_name)
-        
+
         try:
             with db_manager.session_scope(aligned_volume_name) as session:
                 engine = db_manager.get_engine(aligned_volume_name)
-                
+
                 metadata = MetaData()
                 try:
                     annotation_table = Table(
@@ -755,23 +757,24 @@ def get(self, aligned_volume_name: str, version: int, tablename: str):
                 except NoSuchTableError as e:
                     logging.error(f"No table exists {e}")
                     return abort(404)
-                
+
                 response = session.query(annotation_table).limit(10).all()
                 annotations = [r._asdict() for r in response]
-                
+
                 return (annotations, 200) if annotations else abort(404)
-                
+
         except Exception as e:
             logging.error(f"Error querying annotations: {e}")
             return abort(500)
 
+
 @mat_bp.route("/materialize/run/create_virtual/datastack/<string:datastack_name>")
 class CreateVirtualPublicVersionResource(Resource):
     @reset_auth
     @auth_requires_dataset_admin(table_arg="datastack_name")
     @mat_bp.doc("create virtual materialization", security="apikey")
     @accepts("VirtualVersionSchema", schema=VirtualVersionSchema, api=mat_bp)
-    def post(self, datastack_name:str):
+    def post(self, datastack_name: str):
         """Create a virtual version from an existing frozen version.
 
         Args:
@@ -790,8 +793,7 @@ def post(self, datastack_name:str):
         if not tables_to_include:
             return abort(400, "No tables included")
 
-        with db_manager.session_scope(aligned_volume) as session:        
-
+        with db_manager.session_scope(aligned_volume) as session:
             analysis_version = (
                 session.query(AnalysisVersion)
                 .filter(AnalysisVersion.version == target_version)
@@ -851,3 +853,80 @@ def post(self, datastack_name:str):
             analysis_version.expires_on = expiration_timestamp
 
         return f"{virtual_datastack_name} created", 200
+
+
+@mat_bp.route(
+    "/materialize/run/write_deltalake/datastack/<string:datastack_name>/version/<int(signed=True):version>/table_name/<string:table_name>/"
+)
+class WriteDeltalakeResource(Resource):
+    @reset_auth
+    @auth_requires_dataset_admin(table_arg="datastack_name")
+    @mat_bp.doc("Export a table to Delta Lake", security="apikey")
+    def post(self, datastack_name: str, version: int, table_name: str):
+        """Export a frozen materialization table to partitioned Delta Lake.
+
+        Args:
+            datastack_name (str): name of datastack from infoservice
+            version (int): materialization version (-1 for latest)
+            table_name (str): annotation table name to export
+        """
+        from materializationengine.workflows.deltalake_export import (
+            write_deltalake_table,
+        )
+
+        if version == -1:
+            version = get_latest_version(datastack_name)
+
+        datastack_info = get_datastack_info(datastack_name)
+
+        # Accept optional output_specs from JSON body.
+        output_specs = None
+        if request.is_json and request.json:
+            output_specs = request.json.get("output_specs", None)
+
+        if output_specs is not None:
+            from materializationengine.workflows.deltalake_export import (
+                DeltaLakeOutputSpec,
+            )
+
+            if not isinstance(output_specs, list):
+                return abort(400, "output_specs must be a list")
+            for item in output_specs:
+                if not isinstance(item, dict):
+                    return abort(400, "each entry in output_specs must be an object")
+                try:
+                    DeltaLakeOutputSpec(**item)
+                except TypeError as exc:
+                    return abort(400, f"invalid output_specs entry: {exc}")
+
+        write_deltalake_table.s(
+            datastack_info, version, table_name, output_specs=output_specs
+        ).apply_async()
+
+        return {
+            "message": f"Delta Lake export enqueued for {table_name} v{version}"
+        }, 200
+
+    @reset_auth
+    @auth_requires_dataset_admin(table_arg="datastack_name")
+    @mat_bp.doc("Get Delta Lake export progress", security="apikey")
+    def get(self, datastack_name: str, version: int, table_name: str):
+        """Get progress of a Delta Lake export for a table.
+
+        Returns JSON with ``status``, ``rows_processed``, ``total_rows``,
+        and ``percent_complete``.  Returns 404 if no export is tracked.
+        """
+        from materializationengine.workflows.deltalake_export import (
+            get_deltalake_export_progress,
+        )
+
+        if version == -1:
+            version = get_latest_version(datastack_name)
+
+        progress = get_deltalake_export_progress(datastack_name, version, table_name)
+        if progress is None:
+            return {
+                "message": f"No export progress found for {table_name} v{version}"
+            }, 404
+
+        return progress, 200

materializationengine/celery_init.py

@@ -1,6 +1,5 @@
 from celery import Celery
 
-
 celery = Celery(
     include=[
         "materializationengine.workflows.ingest_new_annotations",
@@ -12,6 +11,7 @@
         "materializationengine.workflows.periodic_database_removal",
         "materializationengine.workflows.periodic_materialization",
         "materializationengine.workflows.spatial_lookup",
+        "materializationengine.workflows.deltalake_export",
         "materializationengine.shared_tasks",
         "materializationengine.views",
         "materializationengine.monitor",

materializationengine/celery_worker.py

@@ -71,7 +71,7 @@ def create_celery(app=None):
             "result_expires": 86400,  # results expire in broker after 1 day
             "redis_socket_connect_timeout": 10,
             "broker_transport_options": {
-                "visibility_timeout": 8000,
+                "visibility_timeout": 21600,
                 "socket_timeout": 20,
                 "socket_connect_timeout": 20,
             },  # timeout (s) for tasks to be sent back to broker queue

materializationengine/config.py

@@ -48,9 +48,9 @@ class BaseConfig:
     MERGE_TABLES = True
     AUTH_SERVICE_NAMESPACE = "datastack"
 
-    REDIS_HOST="localhost"
-    REDIS_PORT=6379
-    REDIS_PASSWORD=""
+    REDIS_HOST = "localhost"
+    REDIS_PORT = 6379
+    REDIS_PASSWORD = ""
     SESSION_TYPE = "redis"
     PERMANENT_SESSION_LIFETIME = timedelta(hours=24)
     SESSION_PREFIX = "annotation_upload_"
@@ -62,6 +62,31 @@ class BaseConfig:
     STAGING_DATABASE_NAME = "staging"
     MATERIALIZATION_UPLOAD_BUCKET_PATH = "test_annotation_csv_upload"
 
+    # Delta Lake export settings
+    DELTALAKE_OUTPUT_BUCKET = os.environ.get("DELTALAKE_OUTPUT_BUCKET", "")
+    DELTALAKE_FLUSH_THRESHOLD_BYTES = int(
+        os.environ.get("DELTALAKE_FLUSH_THRESHOLD_BYTES", 2 * 1024 * 1024 * 1024)
+    )
+    DELTALAKE_TARGET_PARTITION_SIZE_MB = int(
+        os.environ.get("DELTALAKE_TARGET_PARTITION_SIZE_MB", 256)
+    )
+    DELTALAKE_OPTIMIZE_MAX_CONCURRENT_TASKS = int(
+        os.environ.get("DELTALAKE_OPTIMIZE_MAX_CONCURRENT_TASKS", 1)
+    )
+    DELTALAKE_OPTIMIZE_TARGET_SIZE_BYTES = (
+        int(os.environ["DELTALAKE_OPTIMIZE_TARGET_SIZE_BYTES"])
+        if "DELTALAKE_OPTIMIZE_TARGET_SIZE_BYTES" in os.environ
+        else None
+    )
+    # this one should help with memory during optimize if it is still a problem,
+    # but has not been tested on the mesh worker nodes. im not sure how spilling to 
+    # disk on those will work out of the box
+    DELTALAKE_OPTIMIZE_MAX_SPILL_SIZE_BYTES = (
+        int(os.environ["DELTALAKE_OPTIMIZE_MAX_SPILL_SIZE_BYTES"])
+        if "DELTALAKE_OPTIMIZE_MAX_SPILL_SIZE_BYTES" in os.environ
+        else None
+    )
+
     if os.environ.get("DAF_CREDENTIALS", None) is not None:
         with open(os.environ.get("DAF_CREDENTIALS"), "r") as f:
             AUTH_TOKEN = json.load(f)["token"]
@@ -150,7 +175,6 @@ class DevConfig(BaseConfig):
     CELERY_BROKER_URL = REDIS_URL
     CELERY_RESULT_BACKEND = REDIS_URL
     USE_SENTINEL = os.environ.get("USE_SENTINEL", False)
-    
 
 
 class TestConfig(BaseConfig):
@@ -193,18 +217,20 @@ def configure_app(app: Flask) -> Flask:
     beat_schedules_before = app.config.get("BEAT_SCHEDULES", "NOT_SET")
     app.config.from_pyfile("config.cfg", silent=True)
     beat_schedules_after = app.config.get("BEAT_SCHEDULES", "NOT_SET")
-    
+
     handler = logging.StreamHandler(sys.stdout)
     handler.setLevel(app.config["LOGGING_LEVEL"])
     app.logger.removeHandler(default_handler)
     app.logger.addHandler(handler)
     app.logger.setLevel(app.config["LOGGING_LEVEL"])
     app.logger.propagate = False
-    
+
     # Log BEAT_SCHEDULES loading status (debug level)
     app.logger.debug(f"BEAT_SCHEDULES before config.cfg: {beat_schedules_before}")
     app.logger.debug(f"BEAT_SCHEDULES after config.cfg: {beat_schedules_after}")
-    app.logger.debug(f"BEAT_SCHEDULES type: {type(beat_schedules_after)}, length: {len(beat_schedules_after) if isinstance(beat_schedules_after, (list, dict)) else 'N/A'}")
+    app.logger.debug(
+        f"BEAT_SCHEDULES type: {type(beat_schedules_after)}, length: {len(beat_schedules_after) if isinstance(beat_schedules_after, (list, dict)) else 'N/A'}"
+    )
     app.logger.debug(app.config)
     app.app_context().push()
     return app