From 4be8a2a06601662a60e5d377c22fd93446b15592 Mon Sep 17 00:00:00 2001 From: Jordan Matelsky Date: Wed, 29 Oct 2025 10:33:03 -0400 Subject: [PATCH] Checkpoint from VS Code for coding agent session --- cloudome.py | 201 ++++++++---- cloudome_core/__init__.py | 1 + cloudome_core/processing.py | 173 ++++++++++ cloudome_core/ptq.py | 86 +++++ cloudome_core/storage.py | 143 +++++++++ manage.py | 411 ++++++++++++++++-------- src/cloudome/config.py | 213 ++++++++++++ database.py => src/cloudome/database.py | 0 src/cloudome/storage.py | 143 +++++++++ src/cloudome_core/__init__.py | 1 + src/cloudome_core/storage.py | 143 +++++++++ 11 files changed, 1317 insertions(+), 198 deletions(-) create mode 100644 cloudome_core/__init__.py create mode 100644 cloudome_core/processing.py create mode 100644 cloudome_core/ptq.py create mode 100644 cloudome_core/storage.py create mode 100644 src/cloudome/config.py rename database.py => src/cloudome/database.py (100%) create mode 100644 src/cloudome/storage.py create mode 100644 src/cloudome_core/__init__.py create mode 100644 src/cloudome_core/storage.py diff --git a/cloudome.py b/cloudome.py index ee42c5e..399f133 100644 --- a/cloudome.py +++ b/cloudome.py @@ -2,7 +2,7 @@ import json import numpy as np from flask import Flask -from database import ( +from cloudome.database import ( SynapseEdgeTask, SynapseEdgeTaskPayload, ContactomeEdgeTaskPayload, @@ -15,45 +15,76 @@ import math from collections import Counter -os.environ['CLOUD_VOLUME_DIR'] = '/tmp/cloudvolume' -os.makedirs('/tmp/cloudvolume', exist_ok=True) +os.environ["CLOUD_VOLUME_DIR"] = "/tmp/cloudvolume" +os.makedirs("/tmp/cloudvolume", exist_ok=True) from cloudvolume import CloudVolume SegmentID = int app = Flask(__name__) + @app.route("/") def _(): - return "Cloudome 2025-04-03" + return "Cloudome v2025-10-29" RADIUS = 10 PRESYNAPTIC = 2 POSTSYNAPTIC = 1 + def return_seg_edge(task: SynapseEdgeTaskPayload) -> tuple[SegmentID, SegmentID]: - xyz_center = task['centroid_xyz'] + xyz_center = task["centroid_xyz"] try: # Get the CloudVolume dimensions - synapse_volume = CloudVolume(task['synapse_channel'], use_https=True, cache=False, secrets="", mip=task['mip']) - segmentation_volume = CloudVolume(task['segmentation_channel'], use_https=True, cache=False, secrets="", mip=task['mip']) + synapse_volume = CloudVolume( + task["synapse_channel"], + use_https=True, + cache=False, + secrets="", + mip=task["mip"], + ) + segmentation_volume = CloudVolume( + task["segmentation_channel"], + use_https=True, + cache=False, + secrets="", + mip=task["mip"], + ) # Calculate bounding box bounds = synapse_volume.shape - x_min, x_max = max(0, xyz_center[0] - RADIUS), min(bounds[0], xyz_center[0] + RADIUS) - y_min, y_max = max(0, xyz_center[1] - RADIUS), min(bounds[1], xyz_center[1] + RADIUS) - z_min, z_max = max(0, xyz_center[2] - RADIUS), min(bounds[2], xyz_center[2] + RADIUS) + x_min, x_max = ( + max(0, xyz_center[0] - RADIUS), + min(bounds[0], xyz_center[0] + RADIUS), + ) + y_min, y_max = ( + max(0, xyz_center[1] - RADIUS), + min(bounds[1], xyz_center[1] + RADIUS), + ) + z_min, z_max = ( + max(0, xyz_center[2] - RADIUS), + min(bounds[2], xyz_center[2] + RADIUS), + ) if x_min >= x_max or y_min >= y_max or z_min >= z_max: - raise ValueError("Slicing range is invalid due to out-of-bounds coordinates.") + raise ValueError( + "Slicing range is invalid due to out-of-bounds coordinates." + ) # Pull volumes inside bounding box for both synapse and segmentation paint - prepost_mask = synapse_volume[x_min:x_max, y_min:y_max, z_min:z_max, 0].squeeze() - seg_mask = segmentation_volume[x_min:x_max, y_min:y_max, z_min:z_max, 0].squeeze() + prepost_mask = synapse_volume[ + x_min:x_max, y_min:y_max, z_min:z_max, 0 + ].squeeze() + seg_mask = segmentation_volume[ + x_min:x_max, y_min:y_max, z_min:z_max, 0 + ].squeeze() # Count seg voxels per id in pre, get ID with most common count => pre_id - vals, counts = np.unique(seg_mask[prepost_mask == PRESYNAPTIC], return_counts=True) + vals, counts = np.unique( + seg_mask[prepost_mask == PRESYNAPTIC], return_counts=True + ) unique_counts = zip(counts, vals) unique_counts = sorted(unique_counts, reverse=True) counts, vals = zip(*unique_counts) @@ -66,13 +97,13 @@ def return_seg_edge(task: SynapseEdgeTaskPayload) -> tuple[SegmentID, SegmentID] pre_max_id = vals[1] elif pre_max_id == 0: raise ValueError("The only presynaptic ID returned is 0.") - + # Subsample PSD voxels to only the one we care about labels_out, N = cc3d.connected_components(prepost_mask, return_N=True) stats = cc3d.statistics(labels_out) distance = math.inf label = -1 - for i, syn_centroid in enumerate(stats['centroids']): + for i, syn_centroid in enumerate(stats["centroids"]): temp_distance = math.dist(syn_centroid, [RADIUS, RADIUS, RADIUS]) if temp_distance < distance: distance = temp_distance @@ -87,7 +118,9 @@ def return_seg_edge(task: SynapseEdgeTaskPayload) -> tuple[SegmentID, SegmentID] counts, vals = zip(*unique_counts) # Error handling for 0 case if len(vals) == 0: - raise ValueError("No postsynaptic ID pixels found at {}.".format(xyz_center)) + raise ValueError( + "No postsynaptic ID pixels found at {}.".format(xyz_center) + ) else: post_max_id = vals[0] # Throw out id zero and presyn id if they are in indices 0 and/or 1 @@ -96,16 +129,20 @@ def return_seg_edge(task: SynapseEdgeTaskPayload) -> tuple[SegmentID, SegmentID] if (post_max_id == 0 or post_max_id == pre_max_id) and len(vals) > 2: post_max_id = vals[2] # If no postsynaptic ID is found, add in contact voxels - if (post_max_id == 0 or post_max_id == pre_max_id): + if post_max_id == 0 or post_max_id == pre_max_id: label_mask_encoding = 1 - label_mask = (labels_out == label) - masked_synapse_seg_vol = seg_mask + label_mask = labels_out == label + masked_synapse_seg_vol = seg_mask masked_synapse_seg_vol[label_mask] = label_mask_encoding contacts = cc3d.contacts(masked_synapse_seg_vol, connectivity=26) max_contact = 0 max_contact_id = -1 for contact in contacts: - if (label_mask_encoding in contact) and (pre_max_id not in contact) and (contacts[contact] > max_contact): + if ( + (label_mask_encoding in contact) + and (pre_max_id not in contact) + and (contacts[contact] > max_contact) + ): max_contact = contacts[contact] max_contact_id = contact[1] post_max_id = max_contact_id @@ -116,23 +153,24 @@ def return_seg_edge(task: SynapseEdgeTaskPayload) -> tuple[SegmentID, SegmentID] print(f"[ERROR]\t{e}") return -1, -1 + def count_contact_voxels(segmentation, resolution): - contacts = cc3d.contacts(segmentation, - connectivity=6, - anisotropy=tuple(resolution), - surface_area=True + contacts = cc3d.contacts( + segmentation, connectivity=6, anisotropy=tuple(resolution), surface_area=True ) return contacts + def remove_contact_overlap( contact_counts: dict[tuple[SegmentID], int], - counts_to_remove: list[dict[tuple[SegmentID], int]] + counts_to_remove: list[dict[tuple[SegmentID], int]], ): final_contacts = contact_counts for count_to_remove in counts_to_remove: final_contacts = dict(Counter(final_contacts) - Counter(count_to_remove)) return final_contacts + def count_volume_voxels(segmentation) -> dict[SegmentID, int]: """ Count the number of voxels for each segment ID in the segmentation volume. @@ -145,54 +183,100 @@ def count_volume_voxels(segmentation) -> dict[SegmentID, int]: return volume_counts + def return_ctc_edges(task: ContactomeEdgeTaskPayload): - xyz_start = task['cuboid_start'] + xyz_start = task["cuboid_start"] print(xyz_start) - xyz_radius = task['cuboid_radius'] + xyz_radius = task["cuboid_radius"] try: # Get the CloudVolume dimensions - segmentation_volume = CloudVolume(task['segmentation_channel'], use_https=True, parallel=False, cache=False, secrets="", mip=task['mip']) + segmentation_volume = CloudVolume( + task["segmentation_channel"], + use_https=True, + parallel=False, + cache=False, + secrets="", + mip=task["mip"], + ) bounds = segmentation_volume.shape - + # Add +1 to each leading edge coord so that contacts with adjacent cuboids are properly recorded - x_min, x_max = max(0, xyz_start[0]), min(bounds[0], xyz_start[0] + xyz_radius[0] + 1) - y_min, y_max = max(0, xyz_start[1]), min(bounds[1], xyz_start[1] + xyz_radius[1] + 1) - z_min, z_max = max(0, xyz_start[2]), min(bounds[2], xyz_start[2] + xyz_radius[2] + 1) + x_min, x_max = ( + max(0, xyz_start[0]), + min(bounds[0], xyz_start[0] + xyz_radius[0] + 1), + ) + y_min, y_max = ( + max(0, xyz_start[1]), + min(bounds[1], xyz_start[1] + xyz_radius[1] + 1), + ) + z_min, z_max = ( + max(0, xyz_start[2]), + min(bounds[2], xyz_start[2] + xyz_radius[2] + 1), + ) if x_min >= x_max or y_min >= y_max or z_min >= z_max: - raise ValueError("Slicing range is invalid due to out-of-bounds coordinates.") - seg_mask = segmentation_volume[x_min:x_max, y_min:y_max, z_min:z_max, 0].squeeze() + raise ValueError( + "Slicing range is invalid due to out-of-bounds coordinates." + ) + seg_mask = segmentation_volume[ + x_min:x_max, y_min:y_max, z_min:z_max, 0 + ].squeeze() # Generate contacts for whole volume - initial_contact_counts = count_contact_voxels(seg_mask, segmentation_volume.resolution) + initial_contact_counts = count_contact_voxels( + seg_mask, segmentation_volume.resolution + ) # Remove doubly-counted contacts at the edges o_x = count_contact_voxels(seg_mask[-1:, :, :], segmentation_volume.resolution) o_y = count_contact_voxels(seg_mask[:, -1:, :], segmentation_volume.resolution) o_z = count_contact_voxels(seg_mask[:, :, -1:], segmentation_volume.resolution) - final_contact_counts = remove_contact_overlap(initial_contact_counts, [o_x, o_y, o_z]) - + final_contact_counts = remove_contact_overlap( + initial_contact_counts, [o_x, o_y, o_z] + ) + return final_contact_counts except Exception as e: print(f"[ERROR]\t[ctc] {e}") return [] + def return_volume_counts(task: VolumeTaskPayload): - xyz_start = task['cuboid_start'] - xyz_radius = task['cuboid_radius'] + xyz_start = task["cuboid_start"] + xyz_radius = task["cuboid_radius"] try: # Get the CloudVolume dimensions - segmentation_volume = CloudVolume(task['segmentation_channel'], use_https=True, parallel=False, cache=False, secrets="", mip=task['mip']) + segmentation_volume = CloudVolume( + task["segmentation_channel"], + use_https=True, + parallel=False, + cache=False, + secrets="", + mip=task["mip"], + ) bounds = segmentation_volume.shape - x_min, x_max = max(0, xyz_start[0]), min(bounds[0], xyz_start[0] + xyz_radius[0]) - y_min, y_max = max(0, xyz_start[1]), min(bounds[1], xyz_start[1] + xyz_radius[1]) - z_min, z_max = max(0, xyz_start[2]), min(bounds[2], xyz_start[2] + xyz_radius[2]) + x_min, x_max = ( + max(0, xyz_start[0]), + min(bounds[0], xyz_start[0] + xyz_radius[0]), + ) + y_min, y_max = ( + max(0, xyz_start[1]), + min(bounds[1], xyz_start[1] + xyz_radius[1]), + ) + z_min, z_max = ( + max(0, xyz_start[2]), + min(bounds[2], xyz_start[2] + xyz_radius[2]), + ) if x_min >= x_max or y_min >= y_max or z_min >= z_max: - raise ValueError("Slicing range is invalid due to out-of-bounds coordinates.") + raise ValueError( + "Slicing range is invalid due to out-of-bounds coordinates." + ) - seg_mask = segmentation_volume[x_min:x_max, y_min:y_max, z_min:z_max, 0].squeeze() + seg_mask = segmentation_volume[ + x_min:x_max, y_min:y_max, z_min:z_max, 0 + ].squeeze() volume_counts = count_volume_voxels(seg_mask) return volume_counts @@ -200,12 +284,16 @@ def return_volume_counts(task: VolumeTaskPayload): print(f"[ERROR]\t[volume] {e}") return {} + def process_queue_job(event, context): # event_records = json.loads(event['Records'][0]['body']) - for record in event['Records']: - payload = json.loads(record['body']) + for record in event["Records"]: + payload = json.loads(record["body"]) # Uses "contactome" as the default task_type for back-compat. - if "cuboid_start" in payload and payload.get("task_type", "contactome") == "contactome": + if ( + "cuboid_start" in payload + and payload.get("task_type", "contactome") == "contactome" + ): # Contactome edge payload = ContactomeEdgeTaskPayload(**payload) graph_id = payload.pop("graph_id") @@ -217,18 +305,21 @@ def process_queue_job(event, context): graph_id=graph_id, # XYZ goes first so that it can still serve as a useful key to retrieve # a specific centroid from the listing: - synapse_id=f"ctc_x{payload['cuboid_start'][0]}_y{payload['cuboid_start'][1]}_z{payload['cuboid_start'][2]}_pre{ids[0]}_post{ids[1]}_w{edges[ids]}" + synapse_id=f"ctc_x{payload['cuboid_start'][0]}_y{payload['cuboid_start'][1]}_z{payload['cuboid_start'][2]}_pre{ids[0]}_post{ids[1]}_w{edges[ids]}", ).save() - elif "cuboid_start" in payload and payload.get("task_type", "contactome") == "volume": + elif ( + "cuboid_start" in payload + and payload.get("task_type", "contactome") == "volume" + ): # Volume task payload = VolumeTaskPayload(**payload) graph_id = payload.pop("graph_id") volume_counts = return_volume_counts(payload) - for (seg_id, count) in volume_counts.items(): + for seg_id, count in volume_counts.items(): VolumeCountResultsModel( graph_id=graph_id, - synapse_id=f"vol_x{payload['cuboid_start'][0]}_y{payload['cuboid_start'][1]}_z{payload['cuboid_start'][2]}_seg{seg_id}_v{count}" + synapse_id=f"vol_x{payload['cuboid_start'][0]}_y{payload['cuboid_start'][1]}_z{payload['cuboid_start'][2]}_seg{seg_id}_v{count}", ).save() else: # Synapse edge @@ -243,7 +334,5 @@ def process_queue_job(event, context): graph_id=graph_id, # XYZ goes first so that it can still serve as a useful key to retrieve # a specific centroid from the listing: - synapse_id=f"syn_x{payload['centroid_xyz'][0]}_y{payload['centroid_xyz'][1]}_z{payload['centroid_xyz'][2]}_pre{u}_post{v}" + synapse_id=f"syn_x{payload['centroid_xyz'][0]}_y{payload['centroid_xyz'][1]}_z{payload['centroid_xyz'][2]}_pre{u}_post{v}", ).save() - - diff --git a/cloudome_core/__init__.py b/cloudome_core/__init__.py new file mode 100644 index 0000000..40036c8 --- /dev/null +++ b/cloudome_core/__init__.py @@ -0,0 +1 @@ +"""Core Cloudome utilities shared across taskqueue workers and CLIs.""" diff --git a/cloudome_core/processing.py b/cloudome_core/processing.py new file mode 100644 index 0000000..d1c6804 --- /dev/null +++ b/cloudome_core/processing.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import math +from collections import Counter +from typing import Any, Iterable, cast + +import cc3d +import numpy as np +from cloudvolume import CloudVolume + +from database import ( # type: ignore[import-not-found] + ContactomeEdgeTaskPayload, + SynapseEdgeTaskPayload, + VolumeTaskPayload, +) + +RADIUS = 10 +PRESYNAPTIC = 2 +POSTSYNAPTIC = 1 + + +def _load_volume(path: str, mip: list[int] | int, *, parallel: bool = False) -> Any: + """Helper to construct CloudVolume objects while appeasing static type checking.""" + + kwargs: dict[str, Any] = { + "use_https": True, + "cache": False, + "secrets": "", + "mip": mip, + } + if parallel: + kwargs["parallel"] = True + return cast(Any, CloudVolume(path, **kwargs)) + + +def return_seg_edge(task: SynapseEdgeTaskPayload) -> tuple[int, int]: + xyz_center = task["centroid_xyz"] + try: + synapse_volume = _load_volume(task["synapse_channel"], task["mip"]) + segmentation_volume = _load_volume(task["segmentation_channel"], task["mip"]) + + bounds = synapse_volume.shape + x_min, x_max = max(0, xyz_center[0] - RADIUS), min(bounds[0], xyz_center[0] + RADIUS) + y_min, y_max = max(0, xyz_center[1] - RADIUS), min(bounds[1], xyz_center[1] + RADIUS) + z_min, z_max = max(0, xyz_center[2] - RADIUS), min(bounds[2], xyz_center[2] + RADIUS) + + if x_min >= x_max or y_min >= y_max or z_min >= z_max: + raise ValueError("Slicing range is invalid due to out-of-bounds coordinates.") + + prepost_mask = synapse_volume[x_min:x_max, y_min:y_max, z_min:z_max, 0].squeeze() + seg_mask = segmentation_volume[x_min:x_max, y_min:y_max, z_min:z_max, 0].squeeze() + + values, counts = np.unique(seg_mask[prepost_mask == PRESYNAPTIC], return_counts=True) + counts_with_ids = sorted(zip(counts, values), reverse=True) + if not counts_with_ids: + raise ValueError(f"No presynaptic ID pixels found at {xyz_center}.") + pre_max_id = counts_with_ids[0][1] + if pre_max_id == 0 and len(counts_with_ids) > 1: + pre_max_id = counts_with_ids[1][1] + elif pre_max_id == 0: + raise ValueError("The only presynaptic ID returned is 0.") + + labels_out, _ = cc3d.connected_components(prepost_mask, return_N=True) + stats = cc3d.statistics(labels_out) + distance = math.inf + label = -1 + for i, syn_centroid in enumerate(stats["centroids"]): + temp_distance = math.dist(syn_centroid, [RADIUS, RADIUS, RADIUS]) + if temp_distance < distance: + distance = temp_distance + label = i + if label == -1: + raise ValueError("No synapse centroids found in given subvolume.") + + values, counts = np.unique(seg_mask[labels_out == label], return_counts=True) + counts_with_ids = sorted(zip(counts, values), reverse=True) + if not counts_with_ids: + raise ValueError(f"No postsynaptic ID pixels found at {xyz_center}.") + post_max_id = counts_with_ids[0][1] + if (post_max_id == 0 or post_max_id == pre_max_id) and len(counts_with_ids) > 1: + post_max_id = counts_with_ids[1][1] + if (post_max_id == 0 or post_max_id == pre_max_id) and len(counts_with_ids) > 2: + post_max_id = counts_with_ids[2][1] + if post_max_id == 0 or post_max_id == pre_max_id: + label_mask_encoding = 1 + label_mask = labels_out == label + masked_synapse_seg_vol = seg_mask.copy() + masked_synapse_seg_vol[label_mask] = label_mask_encoding + contacts = cc3d.contacts(masked_synapse_seg_vol, connectivity=26) + max_contact = 0 + max_contact_id = -1 + for contact in contacts: + if ( + label_mask_encoding in contact + and pre_max_id not in contact + and contacts[contact] > max_contact + ): + max_contact = contacts[contact] + max_contact_id = contact[1] + post_max_id = max_contact_id + + return int(pre_max_id), int(post_max_id) + + except Exception as exc: # pragma: no cover - defensive logging + print(f"[ERROR]\t{exc}") + return -1, -1 + + +def count_contact_voxels(segmentation: np.ndarray, resolution: Iterable[float]): + return cc3d.contacts(segmentation, connectivity=6, anisotropy=tuple(resolution), surface_area=True) + + +def remove_contact_overlap( + contact_counts: dict[tuple[int, int], int], + counts_to_remove: list[dict[tuple[int, int], int]], +): + final_contacts = contact_counts + for count_to_remove in counts_to_remove: + final_contacts = dict(Counter(final_contacts) - Counter(count_to_remove)) + return final_contacts + + +def count_volume_voxels(segmentation: np.ndarray) -> dict[int, int]: + segment_ids = np.unique(segmentation) + volume_counts = {int(i): int(np.sum(segmentation == i)) for i in segment_ids} + return volume_counts + + +def return_ctc_edges(task: ContactomeEdgeTaskPayload): + xyz_start = task["cuboid_start"] + try: + segmentation_volume = _load_volume(task["segmentation_channel"], task["mip"], parallel=False) + + bounds = segmentation_volume.shape + x_min, x_max = max(0, xyz_start[0]), min(bounds[0], xyz_start[0] + task["cuboid_radius"][0] + 1) + y_min, y_max = max(0, xyz_start[1]), min(bounds[1], xyz_start[1] + task["cuboid_radius"][1] + 1) + z_min, z_max = max(0, xyz_start[2]), min(bounds[2], xyz_start[2] + task["cuboid_radius"][2] + 1) + if x_min >= x_max or y_min >= y_max or z_min >= z_max: + raise ValueError("Slicing range is invalid due to out-of-bounds coordinates.") + seg_mask = segmentation_volume[x_min:x_max, y_min:y_max, z_min:z_max, 0].squeeze() + + initial_contact_counts = count_contact_voxels(seg_mask, segmentation_volume.resolution) + + o_x = count_contact_voxels(seg_mask[-1:, :, :], segmentation_volume.resolution) + o_y = count_contact_voxels(seg_mask[:, -1:, :], segmentation_volume.resolution) + o_z = count_contact_voxels(seg_mask[:, :, -1:], segmentation_volume.resolution) + final_contact_counts = remove_contact_overlap(initial_contact_counts, [o_x, o_y, o_z]) + + return final_contact_counts + except Exception as exc: # pragma: no cover - defensive logging + print(f"[ERROR]\t[ctc] {exc}") + return {} + + +def return_volume_counts(task: VolumeTaskPayload): + xyz_start = task["cuboid_start"] + try: + segmentation_volume = _load_volume(task["segmentation_channel"], task["mip"], parallel=False) + + bounds = segmentation_volume.shape + x_min, x_max = max(0, xyz_start[0]), min(bounds[0], xyz_start[0] + task["cuboid_radius"][0]) + y_min, y_max = max(0, xyz_start[1]), min(bounds[1], xyz_start[1] + task["cuboid_radius"][1]) + z_min, z_max = max(0, xyz_start[2]), min(bounds[2], xyz_start[2] + task["cuboid_radius"][2]) + + if x_min >= x_max or y_min >= y_max or z_min >= z_max: + raise ValueError("Slicing range is invalid due to out-of-bounds coordinates.") + + seg_mask = segmentation_volume[x_min:x_max, y_min:y_max, z_min:z_max, 0].squeeze() + volume_counts = count_volume_voxels(seg_mask) + return volume_counts + except Exception as exc: # pragma: no cover - defensive logging + print(f"[ERROR]\t[volume] {exc}") + return {} diff --git a/cloudome_core/ptq.py b/cloudome_core/ptq.py new file mode 100644 index 0000000..8dfd764 --- /dev/null +++ b/cloudome_core/ptq.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from functools import partial +from typing import Iterable + +from taskqueue import queueable + +from .processing import return_ctc_edges, return_seg_edge, return_volume_counts +from .storage import ResultRecord, get_result_store + + +@queueable +def process_task(payload: dict, store_backend: str = "dynamodb", store_kwargs: dict | None = None) -> None: + """Queueable function executed by ptq workers.""" + + store = get_result_store(store_backend, **(store_kwargs or {})) + task_type = payload.get("task_type", "synapse") + graph_id = payload.get("graph_id") + + records: list[ResultRecord] = [] + + if task_type == "contactome": + edges = return_ctc_edges(payload) + if graph_id is None: + raise ValueError("contactome payload missing graph_id") + for (pre_id, post_id), weight in edges.items(): + records.append( + ResultRecord( + result_type="contactome", + graph_id=graph_id, + payload=( + f"ctc_x{payload['cuboid_start'][0]}_y{payload['cuboid_start'][1]}_" + f"z{payload['cuboid_start'][2]}_pre{pre_id}_post{post_id}_w{int(weight)}" + ), + ) + ) + elif task_type == "volume": + counts = return_volume_counts(payload) + if graph_id is None: + raise ValueError("volume payload missing graph_id") + for seg_id, count in counts.items(): + records.append( + ResultRecord( + result_type="volume", + graph_id=graph_id, + payload=( + f"vol_x{payload['cuboid_start'][0]}_y{payload['cuboid_start'][1]}_" + f"z{payload['cuboid_start'][2]}_seg{seg_id}_v{int(count)}" + ), + ) + ) + else: + # default to synapse edge processing + if graph_id is None: + raise ValueError("synapse payload missing graph_id") + pre_id, post_id = return_seg_edge(payload) + records.append( + ResultRecord( + result_type="synapse", + graph_id=graph_id, + payload=( + f"syn_x{payload['centroid_xyz'][0]}_y{payload['centroid_xyz'][1]}_" + f"z{payload['centroid_xyz'][2]}_pre{pre_id}_post{post_id}" + ), + ) + ) + + if records: + store.save_records(records) + + +def make_queueable_tasks( + payloads: Iterable[dict], + *, + store_backend: str, + store_kwargs: dict | None = None, +) -> Iterable[partial]: + """Convert JSON payloads into queueable partials ready for insertion.""" + + for payload in payloads: + yield partial( + process_task, + payload=payload, + store_backend=store_backend, + store_kwargs=store_kwargs or {}, + ) diff --git a/cloudome_core/storage.py b/cloudome_core/storage.py new file mode 100644 index 0000000..7de5a10 --- /dev/null +++ b/cloudome_core/storage.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import os +import sqlite3 +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Optional + +from database import ( + ContactEdgeResultsModel, + SynapseEdgeResultsModel, + VolumeCountResultsModel, +) + + +@dataclass(slots=True) +class ResultRecord: + result_type: str + graph_id: str + payload: str + + +class ResultStore: + """Abstract interface for persisting task results.""" + + def save_records(self, records: Iterable[ResultRecord]) -> None: + raise NotImplementedError + + +@dataclass(slots=True) +class DynamoDBOptions: + table_name: str = "CloudomeResults" + region_name: str = "us-east-1" + endpoint_url: Optional[str] = None + profile_name: Optional[str] = None + + +@dataclass(slots=True) +class SQLiteOptions: + path: Path | str = Path("./cloudome.db") + pragmas: dict[str, Any] | None = None + + def __post_init__(self) -> None: + if isinstance(self.path, str): + self.path = Path(self.path) + + +class DynamoResultStore(ResultStore): + def __init__(self, options: DynamoDBOptions | None = None): + self.options = options or DynamoDBOptions() + self._configure_models(self.options) + if self.options.profile_name: + os.environ.setdefault("AWS_PROFILE", self.options.profile_name) + + @staticmethod + def _configure_models(config: DynamoDBOptions) -> None: + models = ( + SynapseEdgeResultsModel, + ContactEdgeResultsModel, + VolumeCountResultsModel, + ) + for model in models: + model.Meta.table_name = config.table_name + model.Meta.region = config.region_name + if config.endpoint_url: + model.Meta.host = config.endpoint_url + else: + if hasattr(model.Meta, "host"): + setattr(model.Meta, "host", None) + + def save_records(self, records: Iterable[ResultRecord]) -> None: + for record in records: + if record.result_type == "synapse": + SynapseEdgeResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + elif record.result_type == "contactome": + ContactEdgeResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + elif record.result_type == "volume": + VolumeCountResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + else: + raise ValueError(f"Unknown result_type {record.result_type}") + + +class SQLiteResultStore(ResultStore): + def __init__(self, options: SQLiteOptions | None = None): + options = options or SQLiteOptions() + self.path = Path(options.path) + self.pragmas = options.pragmas or {} + self.path.parent.mkdir(parents=True, exist_ok=True) + self._initialize() + + def _initialize(self) -> None: + with self._connect() as conn: + conn.execute( + """ + CREATE TABLE IF NOT EXISTS results ( + result_type TEXT NOT NULL, + graph_id TEXT NOT NULL, + payload TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (result_type, graph_id, payload) + ) + """ + ) + + def _connect(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.path) + for pragma, value in self.pragmas.items(): + conn.execute(f"PRAGMA {pragma} = {value}") + return conn + + def save_records(self, records: Iterable[ResultRecord]) -> None: + with self._connect() as conn: + conn.executemany( + """ + INSERT OR REPLACE INTO results (result_type, graph_id, payload) + VALUES (?, ?, ?) + """, + ((r.result_type, r.graph_id, r.payload) for r in records), + ) + conn.commit() + + +def get_result_store( + backend: str = "dynamodb", + **options: Any, +) -> ResultStore: + """Instantiate a result store directly from CLI-provided options.""" + + backend = backend.lower() + if backend == "dynamodb": + return DynamoResultStore(DynamoDBOptions(**options)) + if backend == "sqlite": + return SQLiteResultStore(SQLiteOptions(**options)) + raise ValueError(f"Unsupported backend '{backend}'") diff --git a/manage.py b/manage.py index 771f7ef..03c275a 100644 --- a/manage.py +++ b/manage.py @@ -16,22 +16,30 @@ from tqdm import tqdm import re -from database import SynapseEdgeResultsModel, ContactomeEdgeTaskPayload, ContactEdgeResultsModel, TaskType, VolumeTaskPayload +from cloudome.database import ( + SynapseEdgeResultsModel, + ContactomeEdgeTaskPayload, + ContactEdgeResultsModel, + TaskType, + VolumeTaskPayload, +) -sqs = boto3.client('sqs', region_name='us-east-1') +sqs = boto3.client("sqs", region_name="us-east-1") -def get_centroids_for_syn_mask(synapse_channel: str, output_file: str, mip: list|int): +def get_centroids_for_syn_mask(synapse_channel: str, output_file: str, mip: list | int): # Use post synaptic densities as centroids. One synapse per PSD - binary_syn_mask = (CloudVolume(synapse_channel, mip=mip, cache=True)[..., 0].squeeze() == 1) + binary_syn_mask = ( + CloudVolume(synapse_channel, mip=mip, cache=True)[..., 0].squeeze() == 1 + ) labels_out, N = cc3d.connected_components(binary_syn_mask, return_N=True) dust_threshold = 1 stats = cc3d.statistics(labels_out) - with open(output_file, 'w') as fh: - for i, syn_centroid in tqdm(enumerate(stats['centroids'])): + with open(output_file, "w") as fh: + for i, syn_centroid in tqdm(enumerate(stats["centroids"])): if np.any(np.isnan(syn_centroid)): continue size = stats["voxel_counts"][i] @@ -39,43 +47,47 @@ def get_centroids_for_syn_mask(synapse_channel: str, output_file: str, mip: list fh.write(",".join(map(str, map(int, syn_centroid))) + "\n") - -def enqueue_centroids_from_file(sqs_url: str, graph_id: str, filename: str, synapse_channel: str, segmentation_channel: str, mip: list|int, enqueue_limit: int = None): +def enqueue_centroids_from_file( + sqs_url: str, + graph_id: str, + filename: str, + synapse_channel: str, + segmentation_channel: str, + mip: list | int, + enqueue_limit: int = None, +): """ Read centroids from a file and enqueue them to SQS for processing. """ - with open(filename, 'r') as fh: + with open(filename, "r") as fh: for i, line in enumerate(tqdm(fh)): if enqueue_limit is not None and i >= enqueue_limit: break - centroid_xyz = tuple(map(int, line.strip().split(','))) + centroid_xyz = tuple(map(int, line.strip().split(","))) payload: SynapseEdgeTaskPayload = { "graph_id": graph_id, "centroid_xyz": centroid_xyz, "synapse_channel": synapse_channel, "segmentation_channel": segmentation_channel, - "mip": mip + "mip": mip, } - sqs.send_message( - QueueUrl=sqs_url, - MessageBody=json.dumps(payload) - ) + sqs.send_message(QueueUrl=sqs_url, MessageBody=json.dumps(payload)) def generate_cuboidwise_tasks_for_contactome_or_volume( - sqs_url: str, - graph_id: str, - task_type: TaskType, - segmentation_channel: str, - mip: list|int, - block_size: tuple = (64, 64, 64), - z_start: int = None, - z_end: int = None, - enqueue_limit: int = None, - ): + sqs_url: str, + graph_id: str, + task_type: TaskType, + segmentation_channel: str, + mip: list | int, + block_size: tuple = (64, 64, 64), + z_start: int = None, + z_end: int = None, + enqueue_limit: int = None, +): # Create a file with each line being a cuboid start and radius seg_data = CloudVolume(segmentation_channel, mip=mip, cache=True) - + if z_end: z_end = z_end if z_end < int(seg_data.shape[2]) else int(seg_data.shape[2]) else: @@ -96,7 +108,9 @@ def generate_cuboidwise_tasks_for_contactome_or_volume( else: print(f"Queueing {len(blocks)} blocks") - for i, ((x_start, x_stop), (y_start, y_stop), (z_start, z_stop)) in tqdm(enumerate(blocks)): + for i, ((x_start, x_stop), (y_start, y_stop), (z_start, z_stop)) in tqdm( + enumerate(blocks) + ): if enqueue_limit is not None and i >= enqueue_limit: break @@ -108,26 +122,18 @@ def generate_cuboidwise_tasks_for_contactome_or_volume( "cuboid_radius": (x_stop - x_start, y_stop - y_start, z_stop - z_start), "segmentation_channel": segmentation_channel, "mip": mip, - } # ContactomeEdgeTaskPayload | VolumeTaskPayload - sqs.send_message( - QueueUrl=sqs_url, - MessageBody=json.dumps(payload) - ) + } # ContactomeEdgeTaskPayload | VolumeTaskPayload + sqs.send_message(QueueUrl=sqs_url, MessageBody=json.dumps(payload)) def local_dequeue(sqs_url: str): import cloudome - response = sqs.receive_message( - QueueUrl=sqs_url, - MaxNumberOfMessages=1 - ) + + response = sqs.receive_message(QueueUrl=sqs_url, MaxNumberOfMessages=1) if "Messages" in response: for message in response["Messages"]: cloudome.process_queue_job({"Records": [message]}, None) - sqs.delete_message( - QueueUrl=sqs_url, - ReceiptHandle=message["ReceiptHandle"] - ) + sqs.delete_message(QueueUrl=sqs_url, ReceiptHandle=message["ReceiptHandle"]) def initialize_resources(): @@ -142,8 +148,8 @@ def export_dynamodb_results_to_csv(graph_id: str, output_file: str): Export results for a given graph_id to a CSV file. This function streams the results to handle large datasets efficiently. """ - with open(output_file, 'w', newline='') as csvfile: - fieldnames = ['graph_id', 'synapse_id'] + with open(output_file, "w", newline="") as csvfile: + fieldnames = ["graph_id", "synapse_id"] writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writeheader() @@ -154,7 +160,10 @@ def export_dynamodb_results_to_csv(graph_id: str, output_file: str): # Stream results from ContactEdgeResultsModel for result in ContactEdgeResultsModel.query(graph_id): - writer.writerow({'graph_id': result.graph_id, 'synapse_id': result.synapse_id}) + writer.writerow( + {"graph_id": result.graph_id, "synapse_id": result.synapse_id} + ) + def simplify_contactome_data(instream: TextIOWrapper, outstream: TextIOWrapper): """ @@ -205,7 +214,9 @@ def simplify_volume_data(instream: TextIOWrapper, outstream: TextIOWrapper): writer.writerow([seg_id, total_count]) -def simplify_synapse_data(raw_file: TextIOWrapper, output_file: str, invalid_nodes: list, simple: bool = False): +def simplify_synapse_data( + raw_file: TextIOWrapper, output_file: str, invalid_nodes: list, simple: bool = False +): """ Simplify synapse data by processing raw export and generating an edgelist CSV. """ @@ -223,8 +234,8 @@ def simplify_synapse_data(raw_file: TextIOWrapper, output_file: str, invalid_nod # Parse synapse data (e.g., syn_x1000_y1068_z444_pre-1_post-1) _, x, y, z, pre, post = edge_raw.split("_") x, y, z = int(x[1:]), int(y[1:]), int(z[1:]) - pre = pre[len("pre"):] - post = post[len("post"):] + pre = pre[len("pre") :] + post = post[len("post") :] g.add_edge(pre, post, pos=(x, y, z)) # Remove invalid nodes (-1 and 0) @@ -242,115 +253,229 @@ def parse_arguments(): parser = argparse.ArgumentParser(description="Cloudome Command Line Interface") # Global arguments for multiple commands - parser.add_argument("--mip", type=str, default="72,72,84", - help="MIP value as either a single int or comma-separated values (e.g., 72,72,84)") - parser.add_argument("--sqs-url", type=str, default="https://sqs.us-east-1.amazonaws.com/407510763690/CloudomeJobs", - help="SQS URL for job queue") + parser.add_argument( + "--mip", + type=str, + default="72,72,84", + help="MIP value as either a single int or comma-separated values (e.g., 72,72,84)", + ) + parser.add_argument( + "--sqs-url", + type=str, + default="https://sqs.us-east-1.amazonaws.com/407510763690/CloudomeJobs", + help="SQS URL for job queue", + ) subparsers = parser.add_subparsers(dest="namespace", required=True) # Namespace: synapses - synapses_parser = subparsers.add_parser("synapses", help="Commands related to synapses") + synapses_parser = subparsers.add_parser( + "synapses", help="Commands related to synapses" + ) synapses_subparsers = synapses_parser.add_subparsers(dest="command", required=True) # Subcommand: generate (synapses) - syn_generate_parser = synapses_subparsers.add_parser("generate", help="Generate centroids for synapse mask") - syn_generate_parser.add_argument("--synapse-channel", type=str, required=True, - help="S3 path to synapse channel data") - syn_generate_parser.add_argument("--output-file", type=str, default="centroids.csv", - help="Output file path for centroids") + syn_generate_parser = synapses_subparsers.add_parser( + "generate", help="Generate centroids for synapse mask" + ) + syn_generate_parser.add_argument( + "--synapse-channel", + type=str, + required=True, + help="S3 path to synapse channel data", + ) + syn_generate_parser.add_argument( + "--output-file", + type=str, + default="centroids.csv", + help="Output file path for centroids", + ) # Subcommand: enqueue (synapses) - enqueue_parser = synapses_subparsers.add_parser("enqueue", help="Enqueue centroids from file") - enqueue_parser.add_argument("--graph-id", type=str, required=True, - help="Graph ID for processing") - enqueue_parser.add_argument("--centroids-file", type=str, required=True, - help="Path to centroids file") - enqueue_parser.add_argument("--synapse-channel", type=str, required=True, - help="S3 path to synapse channel data") - enqueue_parser.add_argument("--segmentation-channel", type=str, required=True, - help="S3 path to segmentation channel data") - enqueue_parser.add_argument("--enqueue-limit", type=int, default=None, - help="Limit the number of centroids to enqueue") + enqueue_parser = synapses_subparsers.add_parser( + "enqueue", help="Enqueue centroids from file" + ) + enqueue_parser.add_argument( + "--graph-id", type=str, required=True, help="Graph ID for processing" + ) + enqueue_parser.add_argument( + "--centroids-file", type=str, required=True, help="Path to centroids file" + ) + enqueue_parser.add_argument( + "--synapse-channel", + type=str, + required=True, + help="S3 path to synapse channel data", + ) + enqueue_parser.add_argument( + "--segmentation-channel", + type=str, + required=True, + help="S3 path to segmentation channel data", + ) + enqueue_parser.add_argument( + "--enqueue-limit", + type=int, + default=None, + help="Limit the number of centroids to enqueue", + ) # Subcommand: simplify (synapses) - syn_simplify_parser = synapses_subparsers.add_parser("simplify", help="Simplify raw synapse data to a CSV file") - syn_simplify_parser.add_argument("--raw-file", type=str, required=True, - help="Path to CSV file with raw exported synapse data (from `export` command)") - syn_simplify_parser.add_argument("--output-file", type=str, required=True, - help="Output file path for simplified synapse data") - syn_simplify_parser.add_argument("--invalid-nodes", type=str, nargs='*', default=["-1", "0"], - help="List of invalid nodes to remove from the graph") - syn_simplify_parser.add_argument("--simple", action='store_true', - help="If set, downcast to a simple graph") + syn_simplify_parser = synapses_subparsers.add_parser( + "simplify", help="Simplify raw synapse data to a CSV file" + ) + syn_simplify_parser.add_argument( + "--raw-file", + type=str, + required=True, + help="Path to CSV file with raw exported synapse data (from `export` command)", + ) + syn_simplify_parser.add_argument( + "--output-file", + type=str, + required=True, + help="Output file path for simplified synapse data", + ) + syn_simplify_parser.add_argument( + "--invalid-nodes", + type=str, + nargs="*", + default=["-1", "0"], + help="List of invalid nodes to remove from the graph", + ) + syn_simplify_parser.add_argument( + "--simple", action="store_true", help="If set, downcast to a simple graph" + ) # Namespace: contactome - contactome_parser = subparsers.add_parser("contactome", help="Commands related to contactome") - contactome_subparsers = contactome_parser.add_subparsers(dest="command", required=True) + contactome_parser = subparsers.add_parser( + "contactome", help="Commands related to contactome" + ) + contactome_subparsers = contactome_parser.add_subparsers( + dest="command", required=True + ) # Subcommand: generate (contactome) - contactome_generate_parser = contactome_subparsers.add_parser("generate", help="Generate cuboidwise tasks for contactome") - contactome_generate_parser.add_argument("--graph-id", type=str, required=True, - help="Graph ID for processing") - contactome_generate_parser.add_argument("--segmentation-channel", type=str, required=True, - help="S3 path to segmentation channel data") - contactome_generate_parser.add_argument("--block-size-x", type=int, default=64, - help="Block size for X dimension") - contactome_generate_parser.add_argument("--block-size-y", type=int, default=64, - help="Block size for Y dimension") - contactome_generate_parser.add_argument("--block-size-z", type=int, default=32, - help="Block size for Z dimension") - contactome_generate_parser.add_argument("--z-start", type=int, default=None, - help="Starting Z slice") - contactome_generate_parser.add_argument("--z-end", type=int, default=None, - help="Ending Z slice") - contactome_generate_parser.add_argument("--enqueue-limit", type=int, default=None, - help="Limit the number of tasks to enqueue") + contactome_generate_parser = contactome_subparsers.add_parser( + "generate", help="Generate cuboidwise tasks for contactome" + ) + contactome_generate_parser.add_argument( + "--graph-id", type=str, required=True, help="Graph ID for processing" + ) + contactome_generate_parser.add_argument( + "--segmentation-channel", + type=str, + required=True, + help="S3 path to segmentation channel data", + ) + contactome_generate_parser.add_argument( + "--block-size-x", type=int, default=64, help="Block size for X dimension" + ) + contactome_generate_parser.add_argument( + "--block-size-y", type=int, default=64, help="Block size for Y dimension" + ) + contactome_generate_parser.add_argument( + "--block-size-z", type=int, default=32, help="Block size for Z dimension" + ) + contactome_generate_parser.add_argument( + "--z-start", type=int, default=None, help="Starting Z slice" + ) + contactome_generate_parser.add_argument( + "--z-end", type=int, default=None, help="Ending Z slice" + ) + contactome_generate_parser.add_argument( + "--enqueue-limit", + type=int, + default=None, + help="Limit the number of tasks to enqueue", + ) # Subcommand: simplify (contactome) - contactome_simplify_parser = contactome_subparsers.add_parser("simplify", help="Simplify contactome raw export to edgelist CSV") - contactome_simplify_parser.add_argument("--raw-file", type=str, required=True, - help="Path to CSV file with raw exported contactome data (from `export` command)") - contactome_simplify_parser.add_argument("--output-file", type=str, required=True, - help="Output file path for simplified contactome data") - + contactome_simplify_parser = contactome_subparsers.add_parser( + "simplify", help="Simplify contactome raw export to edgelist CSV" + ) + contactome_simplify_parser.add_argument( + "--raw-file", + type=str, + required=True, + help="Path to CSV file with raw exported contactome data (from `export` command)", + ) + contactome_simplify_parser.add_argument( + "--output-file", + type=str, + required=True, + help="Output file path for simplified contactome data", + ) + # Namespace: volume - volume_parser = subparsers.add_parser("volume", help="Commands related to volume computation") + volume_parser = subparsers.add_parser( + "volume", help="Commands related to volume computation" + ) volume_subparsers = volume_parser.add_subparsers(dest="command", required=True) # Subcommand: generate (volume) - volume_generate_parser = volume_subparsers.add_parser("generate", help="Generate cuboidwise tasks for volume") - volume_generate_parser.add_argument("--graph-id", type=str, required=True, - help="Graph ID for processing") - volume_generate_parser.add_argument("--segmentation-channel", type=str, required=True, - help="S3 path to segmentation channel data") - volume_generate_parser.add_argument("--block-size-x", type=int, default=64, - help="Block size for X dimension") - volume_generate_parser.add_argument("--block-size-y", type=int, default=64, - help="Block size for Y dimension") - volume_generate_parser.add_argument("--block-size-z", type=int, default=32, - help="Block size for Z dimension") - volume_generate_parser.add_argument("--z-start", type=int, default=None, - help="Starting Z slice") - volume_generate_parser.add_argument("--z-end", type=int, default=None, - help="Ending Z slice") - volume_generate_parser.add_argument("--enqueue-limit", type=int, default=None, - help="Limit the number of tasks to enqueue") + volume_generate_parser = volume_subparsers.add_parser( + "generate", help="Generate cuboidwise tasks for volume" + ) + volume_generate_parser.add_argument( + "--graph-id", type=str, required=True, help="Graph ID for processing" + ) + volume_generate_parser.add_argument( + "--segmentation-channel", + type=str, + required=True, + help="S3 path to segmentation channel data", + ) + volume_generate_parser.add_argument( + "--block-size-x", type=int, default=64, help="Block size for X dimension" + ) + volume_generate_parser.add_argument( + "--block-size-y", type=int, default=64, help="Block size for Y dimension" + ) + volume_generate_parser.add_argument( + "--block-size-z", type=int, default=32, help="Block size for Z dimension" + ) + volume_generate_parser.add_argument( + "--z-start", type=int, default=None, help="Starting Z slice" + ) + volume_generate_parser.add_argument( + "--z-end", type=int, default=None, help="Ending Z slice" + ) + volume_generate_parser.add_argument( + "--enqueue-limit", + type=int, + default=None, + help="Limit the number of tasks to enqueue", + ) # Subcommand: simplify (volume) - volume_simplify_parser = volume_subparsers.add_parser("simplify", help="Simplify volume raw export to edgelist CSV") - volume_simplify_parser.add_argument("--raw-file", type=str, required=True, - help="Path to CSV file with raw exported volume data (from `export` command)") - volume_simplify_parser.add_argument("--output-file", type=str, required=True, - help="Output file path for simplified volume data") + volume_simplify_parser = volume_subparsers.add_parser( + "simplify", help="Simplify volume raw export to edgelist CSV" + ) + volume_simplify_parser.add_argument( + "--raw-file", + type=str, + required=True, + help="Path to CSV file with raw exported volume data (from `export` command)", + ) + volume_simplify_parser.add_argument( + "--output-file", + type=str, + required=True, + help="Output file path for simplified volume data", + ) # Namespace: export export_parser = subparsers.add_parser("export", help="Export results to CSV") - export_parser.add_argument("graph_id", type=str, help="Graph ID to export results for") + export_parser.add_argument( + "graph_id", type=str, help="Graph ID to export results for" + ) export_parser.add_argument("output_file", type=str, help="Output CSV file path") # Namespace: dequeue - dequeue_parser = subparsers.add_parser("dequeue", help="Process a single job from the queue") + dequeue_parser = subparsers.add_parser( + "dequeue", help="Process a single job from the queue" + ) return parser.parse_args() @@ -367,7 +492,9 @@ def main(): single_mip = int(args.mip) mip = [single_mip, single_mip, single_mip] except ValueError: - print(f"Error: MIP value '{args.mip}' is not valid. Use a single integer or comma-separated integers.") + print( + f"Error: MIP value '{args.mip}' is not valid. Use a single integer or comma-separated integers." + ) exit(1) if args.namespace == "synapses": @@ -385,15 +512,15 @@ def main(): synapse_channel=args.synapse_channel, segmentation_channel=args.segmentation_channel, mip=mip, - enqueue_limit=args.enqueue_limit + enqueue_limit=args.enqueue_limit, ) elif args.command == "simplify": - with open(args.raw_file, 'r') as infile: + with open(args.raw_file, "r") as infile: simplify_synapse_data( infile, output_file=args.output_file, invalid_nodes=args.invalid_nodes, - simple=args.simple + simple=args.simple, ) elif args.namespace == "contactome": if args.command == "generate": @@ -407,14 +534,14 @@ def main(): block_size=block_size, z_start=args.z_start, z_end=args.z_end, - enqueue_limit=args.enqueue_limit + enqueue_limit=args.enqueue_limit, ) elif args.command == "simplify": - with open(args.raw_file, 'r') as infile, open(args.output_file, 'w') as outfile: - simplify_contactome_data( - instream=infile, - outstream=outfile - ) + with ( + open(args.raw_file, "r") as infile, + open(args.output_file, "w") as outfile, + ): + simplify_contactome_data(instream=infile, outstream=outfile) elif args.namespace == "volume": if args.command == "generate": block_size = (args.block_size_x, args.block_size_y, args.block_size_z) @@ -427,14 +554,14 @@ def main(): block_size=block_size, z_start=args.z_start, z_end=args.z_end, - enqueue_limit=args.enqueue_limit + enqueue_limit=args.enqueue_limit, ) elif args.command == "simplify": - with open(args.raw_file, 'r') as infile, open(args.output_file, 'w') as outfile: - simplify_volume_data( - instream=infile, - outstream=outfile - ) + with ( + open(args.raw_file, "r") as infile, + open(args.output_file, "w") as outfile, + ): + simplify_volume_data(instream=infile, outstream=outfile) elif args.namespace == "export": export_dynamodb_results_to_csv(args.graph_id, args.output_file) elif args.namespace == "dequeue": diff --git a/src/cloudome/config.py b/src/cloudome/config.py new file mode 100644 index 0000000..a32f8ff --- /dev/null +++ b/src/cloudome/config.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, Optional +import os +import tomllib + +QueueBackend = Literal["sqs", "file"] +StorageBackend = Literal["dynamodb", "sqlite"] + + +@dataclass(slots=True) +class QueuePollSettings: + """Common polling options shared across ptq backends.""" + + lease_seconds: int = 300 + poll_interval: float = 1.0 + max_messages: int = 10 + green: bool = False + parallel: int = 1 + + +@dataclass(slots=True) +class SQSQueueSettings: + queue_url: str + region_name: Optional[str] = None + endpoint_url: Optional[str] = None + visibility_timeout: Optional[int] = None + profile_name: Optional[str] = None + + +@dataclass(slots=True) +class FileQueueSettings: + directory: Path + tally: bool = True + + +@dataclass(slots=True) +class QueueSettings: + backend: QueueBackend + poll: QueuePollSettings + sqs: Optional[SQSQueueSettings] = None + file: Optional[FileQueueSettings] = None + + @property + def is_sqs(self) -> bool: + return self.backend == "sqs" + + @property + def is_file(self) -> bool: + return self.backend == "file" + + +@dataclass(slots=True) +class DynamoDBSettings: + table_name: str = "CloudomeResults" + region_name: str = "us-east-1" + endpoint_url: Optional[str] = None + profile_name: Optional[str] = None + + +@dataclass(slots=True) +class SQLiteSettings: + path: Path = Path("./cloudome.db") + pragmas: dict[str, Any] | None = None + + +@dataclass(slots=True) +class StorageSettings: + backend: StorageBackend + dynamodb: Optional[DynamoDBSettings] = None + sqlite: Optional[SQLiteSettings] = None + + @property + def is_dynamodb(self) -> bool: + return self.backend == "dynamodb" + + @property + def is_sqlite(self) -> bool: + return self.backend == "sqlite" + + +@dataclass(slots=True) +class WorkerSettings: + module: str = "cloudome" + log_level: str = "INFO" + import_paths: tuple[str, ...] = () + + +@dataclass(slots=True) +class AppSettings: + queue: QueueSettings + storage: StorageSettings + worker: WorkerSettings = WorkerSettings() + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "AppSettings": + queue_cfg = cls._parse_queue(data.get("queue", {})) + storage_cfg = cls._parse_storage(data.get("storage", {})) + worker_cfg = cls._parse_worker(data.get("worker", {})) + return cls(queue=queue_cfg, storage=storage_cfg, worker=worker_cfg) + + @staticmethod + def _parse_queue(raw: dict[str, Any]) -> QueueSettings: + backend: QueueBackend = raw.get("backend", "file") + poll_raw = raw.get("poll", {}) + poll = QueuePollSettings( + lease_seconds=int(poll_raw.get("lease_seconds", 300)), + poll_interval=float(poll_raw.get("poll_interval", 1.0)), + max_messages=int(poll_raw.get("max_messages", 10)), + green=bool(poll_raw.get("green", False)), + parallel=int(poll_raw.get("parallel", 1)), + ) + + sqs_settings: Optional[SQSQueueSettings] = None + file_settings: Optional[FileQueueSettings] = None + + if backend == "sqs": + sqs_raw = raw.get("sqs", {}) + queue_url = sqs_raw.get("queue_url") + if not queue_url: + raise ValueError("SQS queue backend selected but queue_url is missing") + sqs_settings = SQSQueueSettings( + queue_url=queue_url, + region_name=sqs_raw.get("region_name"), + endpoint_url=sqs_raw.get("endpoint_url"), + visibility_timeout=sqs_raw.get("visibility_timeout"), + profile_name=sqs_raw.get("profile_name"), + ) + else: + file_raw = raw.get("file", {}) + directory = file_raw.get("directory", "./queues/default") + file_settings = FileQueueSettings( + directory=Path(directory).expanduser().resolve(), + tally=bool(file_raw.get("tally", True)), + ) + + return QueueSettings( + backend=backend, poll=poll, sqs=sqs_settings, file=file_settings + ) + + @staticmethod + def _parse_storage(raw: dict[str, Any]) -> StorageSettings: + backend: StorageBackend = raw.get("backend", "dynamodb") + dynamo_cfg: Optional[DynamoDBSettings] = None + sqlite_cfg: Optional[SQLiteSettings] = None + + if backend == "dynamodb": + dynamo_raw = raw.get("dynamodb", {}) + dynamo_cfg = DynamoDBSettings( + table_name=dynamo_raw.get("table_name", "CloudomeResults"), + region_name=dynamo_raw.get("region_name", "us-east-1"), + endpoint_url=dynamo_raw.get("endpoint_url"), + profile_name=dynamo_raw.get("profile_name"), + ) + else: + sqlite_raw = raw.get("sqlite", {}) + path = sqlite_raw.get("path", "./cloudome.db") + pragmas = sqlite_raw.get("pragmas") + sqlite_cfg = SQLiteSettings( + path=Path(path).expanduser().resolve(), pragmas=pragmas + ) + + return StorageSettings(backend=backend, dynamodb=dynamo_cfg, sqlite=sqlite_cfg) + + @staticmethod + def _parse_worker(raw: dict[str, Any]) -> WorkerSettings: + import_paths = raw.get("import_paths", []) + if isinstance(import_paths, str): + import_paths = [import_paths] + return WorkerSettings( + module=raw.get("module", "cloudome"), + log_level=raw.get("log_level", "INFO"), + import_paths=tuple(import_paths), + ) + + +def locate_config_path(path: Optional[str | Path] = None) -> Path: + """Resolve a configuration file path using defaults and environment overrides.""" + + if path: + candidate = Path(path) + if candidate.is_file(): + return candidate + raise FileNotFoundError(candidate) + + env_path = os.environ.get("CLOUDOME_CONFIG") + if env_path: + candidate = Path(env_path) + if candidate.is_file(): + return candidate + raise FileNotFoundError(candidate) + + default_candidates = [ + Path.cwd() / "cloudome.toml", + Path(__file__).resolve().parents[2] / "cloudome.toml", + ] + + for candidate in default_candidates: + if candidate.is_file(): + return candidate + + raise FileNotFoundError( + "No configuration file found. Set CLOUDOME_CONFIG or create cloudome.toml." + ) + + +def load_settings(path: Optional[str | Path] = None) -> AppSettings: + config_path = locate_config_path(path) + with open(config_path, "rb") as fp: + payload = tomllib.load(fp) + return AppSettings.from_dict(payload) diff --git a/database.py b/src/cloudome/database.py similarity index 100% rename from database.py rename to src/cloudome/database.py diff --git a/src/cloudome/storage.py b/src/cloudome/storage.py new file mode 100644 index 0000000..04f1183 --- /dev/null +++ b/src/cloudome/storage.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import os +import sqlite3 +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Optional + +from .database import ( + ContactEdgeResultsModel, + SynapseEdgeResultsModel, + VolumeCountResultsModel, +) + + +@dataclass(slots=True) +class ResultRecord: + result_type: str + graph_id: str + payload: str + + +class ResultStore: + """Abstract interface for persisting task results.""" + + def save_records(self, records: Iterable[ResultRecord]) -> None: + raise NotImplementedError + + +@dataclass(slots=True) +class DynamoDBOptions: + table_name: str = "CloudomeResults" + region_name: str = "us-east-1" + endpoint_url: Optional[str] = None + profile_name: Optional[str] = None + + +@dataclass(slots=True) +class SQLiteOptions: + path: Path | str = Path("./cloudome.db") + pragmas: dict[str, Any] | None = None + + def __post_init__(self) -> None: + if isinstance(self.path, str): + self.path = Path(self.path) + + +class DynamoResultStore(ResultStore): + def __init__(self, options: DynamoDBOptions | None = None): + self.options = options or DynamoDBOptions() + self._configure_models(self.options) + if self.options.profile_name: + os.environ.setdefault("AWS_PROFILE", self.options.profile_name) + + @staticmethod + def _configure_models(config: DynamoDBOptions) -> None: + models = ( + SynapseEdgeResultsModel, + ContactEdgeResultsModel, + VolumeCountResultsModel, + ) + for model in models: + model.Meta.table_name = config.table_name + model.Meta.region = config.region_name + if config.endpoint_url: + model.Meta.host = config.endpoint_url + else: + if hasattr(model.Meta, "host"): + setattr(model.Meta, "host", None) + + def save_records(self, records: Iterable[ResultRecord]) -> None: + for record in records: + if record.result_type == "synapse": + SynapseEdgeResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + elif record.result_type == "contactome": + ContactEdgeResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + elif record.result_type == "volume": + VolumeCountResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + else: + raise ValueError(f"Unknown result_type {record.result_type}") + + +class SQLiteResultStore(ResultStore): + def __init__(self, options: SQLiteOptions | None = None): + options = options or SQLiteOptions() + self.path = Path(options.path) + self.pragmas = options.pragmas or {} + self.path.parent.mkdir(parents=True, exist_ok=True) + self._initialize() + + def _initialize(self) -> None: + with self._connect() as conn: + conn.execute( + """ + CREATE TABLE IF NOT EXISTS results ( + result_type TEXT NOT NULL, + graph_id TEXT NOT NULL, + payload TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (result_type, graph_id, payload) + ) + """ + ) + + def _connect(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.path) + for pragma, value in self.pragmas.items(): + conn.execute(f"PRAGMA {pragma} = {value}") + return conn + + def save_records(self, records: Iterable[ResultRecord]) -> None: + with self._connect() as conn: + conn.executemany( + """ + INSERT OR REPLACE INTO results (result_type, graph_id, payload) + VALUES (?, ?, ?) + """, + ((r.result_type, r.graph_id, r.payload) for r in records), + ) + conn.commit() + + +def get_result_store( + backend: str = "dynamodb", + **options: Any, +) -> ResultStore: + """Instantiate a result store directly from CLI-provided options.""" + + backend = backend.lower() + if backend == "dynamodb": + return DynamoResultStore(DynamoDBOptions(**options)) + if backend == "sqlite": + return SQLiteResultStore(SQLiteOptions(**options)) + raise ValueError(f"Unsupported backend '{backend}'") diff --git a/src/cloudome_core/__init__.py b/src/cloudome_core/__init__.py new file mode 100644 index 0000000..40036c8 --- /dev/null +++ b/src/cloudome_core/__init__.py @@ -0,0 +1 @@ +"""Core Cloudome utilities shared across taskqueue workers and CLIs.""" diff --git a/src/cloudome_core/storage.py b/src/cloudome_core/storage.py new file mode 100644 index 0000000..7de5a10 --- /dev/null +++ b/src/cloudome_core/storage.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import os +import sqlite3 +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Optional + +from database import ( + ContactEdgeResultsModel, + SynapseEdgeResultsModel, + VolumeCountResultsModel, +) + + +@dataclass(slots=True) +class ResultRecord: + result_type: str + graph_id: str + payload: str + + +class ResultStore: + """Abstract interface for persisting task results.""" + + def save_records(self, records: Iterable[ResultRecord]) -> None: + raise NotImplementedError + + +@dataclass(slots=True) +class DynamoDBOptions: + table_name: str = "CloudomeResults" + region_name: str = "us-east-1" + endpoint_url: Optional[str] = None + profile_name: Optional[str] = None + + +@dataclass(slots=True) +class SQLiteOptions: + path: Path | str = Path("./cloudome.db") + pragmas: dict[str, Any] | None = None + + def __post_init__(self) -> None: + if isinstance(self.path, str): + self.path = Path(self.path) + + +class DynamoResultStore(ResultStore): + def __init__(self, options: DynamoDBOptions | None = None): + self.options = options or DynamoDBOptions() + self._configure_models(self.options) + if self.options.profile_name: + os.environ.setdefault("AWS_PROFILE", self.options.profile_name) + + @staticmethod + def _configure_models(config: DynamoDBOptions) -> None: + models = ( + SynapseEdgeResultsModel, + ContactEdgeResultsModel, + VolumeCountResultsModel, + ) + for model in models: + model.Meta.table_name = config.table_name + model.Meta.region = config.region_name + if config.endpoint_url: + model.Meta.host = config.endpoint_url + else: + if hasattr(model.Meta, "host"): + setattr(model.Meta, "host", None) + + def save_records(self, records: Iterable[ResultRecord]) -> None: + for record in records: + if record.result_type == "synapse": + SynapseEdgeResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + elif record.result_type == "contactome": + ContactEdgeResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + elif record.result_type == "volume": + VolumeCountResultsModel( + graph_id=record.graph_id, + synapse_id=record.payload, + ).save() + else: + raise ValueError(f"Unknown result_type {record.result_type}") + + +class SQLiteResultStore(ResultStore): + def __init__(self, options: SQLiteOptions | None = None): + options = options or SQLiteOptions() + self.path = Path(options.path) + self.pragmas = options.pragmas or {} + self.path.parent.mkdir(parents=True, exist_ok=True) + self._initialize() + + def _initialize(self) -> None: + with self._connect() as conn: + conn.execute( + """ + CREATE TABLE IF NOT EXISTS results ( + result_type TEXT NOT NULL, + graph_id TEXT NOT NULL, + payload TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (result_type, graph_id, payload) + ) + """ + ) + + def _connect(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.path) + for pragma, value in self.pragmas.items(): + conn.execute(f"PRAGMA {pragma} = {value}") + return conn + + def save_records(self, records: Iterable[ResultRecord]) -> None: + with self._connect() as conn: + conn.executemany( + """ + INSERT OR REPLACE INTO results (result_type, graph_id, payload) + VALUES (?, ?, ?) + """, + ((r.result_type, r.graph_id, r.payload) for r in records), + ) + conn.commit() + + +def get_result_store( + backend: str = "dynamodb", + **options: Any, +) -> ResultStore: + """Instantiate a result store directly from CLI-provided options.""" + + backend = backend.lower() + if backend == "dynamodb": + return DynamoResultStore(DynamoDBOptions(**options)) + if backend == "sqlite": + return SQLiteResultStore(SQLiteOptions(**options)) + raise ValueError(f"Unsupported backend '{backend}'")