From 7f713e9f12ce56e406ee61c75746a36b615f6626 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 2 Jan 2026 01:21:22 -0800 Subject: [PATCH 1/9] feat: BBoxStrider rewrite for superchunking --- tests/unit/geometry/test_bbox_strider.py | 124 ++++++++++++++++++++++- zetta_utils/geometry/bbox_strider.py | 112 +++++++++++--------- 2 files changed, 188 insertions(+), 48 deletions(-) diff --git a/tests/unit/geometry/test_bbox_strider.py b/tests/unit/geometry/test_bbox_strider.py index aa58457a9..8fba1e7ac 100644 --- a/tests/unit/geometry/test_bbox_strider.py +++ b/tests/unit/geometry/test_bbox_strider.py @@ -274,6 +274,30 @@ def test_bbox_strider_len( 0, BBox3D.from_slices((slice(0, 1), slice(-1, 0), slice(0, 2))), ], + [ + Vec3D(0, -1, 0), + Vec3D(1, 2, 3), + Vec3D(1, 1, 1), + IntVec3D(2, 2, 2), + IntVec3D(2, 2, 2), + IntVec3D(0, 0, 0), + "exact", + IntVec3D(4, 4, 4), + 0, + BBox3D.from_slices((slice(0, 1), slice(-1, 2), slice(0, 3))), + ], + [ + Vec3D(0, -1, 0), + Vec3D(1, 2, 3), + Vec3D(1, 1, 1), + IntVec3D(2, 2, 2), + IntVec3D(2, 2, 2), + IntVec3D(0, 0, 0), + "exact", + IntVec3D(4, 2, 4), + 0, + BBox3D.from_slices((slice(0, 1), slice(-1, 0), slice(0, 3))), + ], [ Vec3D(0, -1, 0), Vec3D(1, 2, 3), @@ -296,7 +320,7 @@ def test_bbox_strider_len( "exact", IntVec3D(4, 4, 4), 0, - BBox3D.from_slices((slice(0, 4), slice(-1, 0), slice(0, 4))), + BBox3D.from_slices((slice(0, 4), slice(-1, 2), slice(0, 4))), ], [ Vec3D(0, 0, 0), @@ -411,6 +435,104 @@ def test_bbox_strider_get_nth_res( assert bbox.start.allclose(expected.start) and bbox.end.allclose(expected.end) +@pytest.mark.parametrize( + "start_coord, end_coord, resolution, chunk_size, stride, stride_start_offset, mode, max_superchunk_size, expected_num_chunks", + [ + # bbox smaller than max_superchunk_size in all dimensions + [ + Vec3D(0, 0, 0), + Vec3D(100, 100, 100), + Vec3D(1, 1, 1), + IntVec3D(50, 50, 50), + IntVec3D(50, 50, 50), + None, + "expand", + IntVec3D(200, 200, 200), + 1, + ], + # bbox smaller in some dimensions only (anisotropic) + [ + Vec3D(0, 0, 0), + Vec3D(100, 500, 100), + Vec3D(1, 1, 1), + IntVec3D(50, 50, 50), + IntVec3D(50, 50, 50), + None, + "expand", + IntVec3D(200, 200, 200), + 3, + ], + # small bbox with weird alignment + [ + Vec3D(15, 23, 7), + Vec3D(115, 123, 107), + Vec3D(1, 1, 1), + IntVec3D(50, 50, 50), + IntVec3D(50, 50, 50), + IntVec3D(5, 3, 7), + "expand", + IntVec3D(300, 300, 300), + 1, + ], + # very small bbox (smaller than chunk_size) + [ + Vec3D(0, 0, 0), + Vec3D(30, 30, 30), + Vec3D(1, 1, 1), + IntVec3D(50, 50, 50), + IntVec3D(50, 50, 50), + None, + "expand", + IntVec3D(200, 200, 200), + 1, + ], + # edge case: bbox exactly equals max_superchunk_size + [ + Vec3D(0, 0, 0), + Vec3D(200, 200, 200), + Vec3D(1, 1, 1), + IntVec3D(50, 50, 50), + IntVec3D(50, 50, 50), + None, + "expand", + IntVec3D(200, 200, 200), + 1, + ], + ], +) +def test_bbox_strider_superchunking_small_bbox( + start_coord, + end_coord, + resolution, + chunk_size, + stride, + stride_start_offset, + mode, + max_superchunk_size, + expected_num_chunks, +): + strider = BBoxStrider( + bbox=BBox3D.from_coords( + start_coord=start_coord, end_coord=end_coord, resolution=resolution + ), + chunk_size=chunk_size, + resolution=resolution, + stride=stride, + stride_start_offset=stride_start_offset, + mode=mode, + max_superchunk_size=max_superchunk_size, + ) + assert strider.num_chunks == expected_num_chunks + + # Verify no overlapping chunks (for non-overlapping stride) + all_bboxes = strider.get_all_chunk_bboxes() + for i, bbox1 in enumerate(all_bboxes): + for bbox2 in all_bboxes[i + 1 :]: + assert not bbox1.intersects( + bbox2 + ), f"Chunks {i} overlap: {bbox1.bounds} and {bbox2.bounds}" + + def test_bbox_strider_exc(mocker): bbox = BBox3D.from_coords( start_coord=Vec3D(0, 0, 0), end_coord=Vec3D(1, 1, 2), resolution=Vec3D(1, 1, 1) diff --git a/zetta_utils/geometry/bbox_strider.py b/zetta_utils/geometry/bbox_strider.py index 42a761738..68288f9a1 100644 --- a/zetta_utils/geometry/bbox_strider.py +++ b/zetta_utils/geometry/bbox_strider.py @@ -52,6 +52,8 @@ class BBoxStrider: stride_start_offset_in_unit: Vec3D = attrs.field(init=False) bbox_snapped: BBox3D = attrs.field(init=False) step_limits: Vec3D[int] = attrs.field(init=False) + atomic_step_limits: Vec3D[int] = attrs.field(init=False) + super_factors: Vec3D[int] = attrs.field(init=False) step_start_partial: Tuple[bool, bool, bool] = attrs.field(init=False) step_end_partial: Tuple[bool, bool, bool] = attrs.field(init=False) mode: Optional[Literal["shrink", "expand", "exact"]] = "expand" @@ -67,28 +69,25 @@ def __attrs_post_init__(self) -> None: if self.mode == "exact": self._attrs_post_init_exact() - # recursively call __attrs_post_init__ if superchunking is set - if self.max_superchunk_size is None: - return - else: - if self.mode in ("expand", "shrink"): - object.__setattr__(self, "bbox", self.bbox_snapped) + if self.max_superchunk_size is not None: if not self.max_superchunk_size >= self.chunk_size: raise ValueError("`max_superchunk_size` must be at least as large as `chunk_size`") if self.chunk_size != self.stride: raise NotImplementedError( "superchunking is only supported when the `chunk_size` and `stride` are equal" ) - superchunk_size = self.chunk_size * (self.max_superchunk_size // self.chunk_size) - object.__setattr__(self, "chunk_size", superchunk_size) - object.__setattr__(self, "stride", superchunk_size) - object.__setattr__(self, "mode", "exact") - object.__setattr__(self, "max_superchunk_size", None) - object.__setattr__( - self, "stride_start_offset", self.bbox_snapped.start / self.resolution + super_factors = self.max_superchunk_size // self.chunk_size + object.__setattr__(self, "super_factors", super_factors) + super_step_limits = Vec3D[int]( + *( + (self.atomic_step_limits[i] + super_factors[i] - 1) // super_factors[i] + for i in range(3) + ) ) - object.__setattr__(self, "stride_start_offset_in_unit", self.bbox_snapped.start) - self.__attrs_post_init__() + object.__setattr__(self, "step_limits", super_step_limits) + else: + object.__setattr__(self, "super_factors", Vec3D[int](1, 1, 1)) + object.__setattr__(self, "step_limits", self.atomic_step_limits) def _attrs_post_init_exact(self) -> None: if self.chunk_size != self.stride: @@ -115,21 +114,20 @@ def _attrs_post_init_exact(self) -> None: ) ) ) - step_limits = floor(round(step_limits_snapped, VEC3D_PRECISION)) + atomic_step_limits = floor(round(step_limits_snapped, VEC3D_PRECISION)) bbox_start_diff = bbox_snapped.start - self.bbox.start bbox_end_diff = self.bbox.end - bbox_snapped.end step_start_partial = tuple(round(e, VEC3D_PRECISION) > 0 for e in bbox_start_diff) step_end_partial = tuple(round(e, VEC3D_PRECISION) > 0 for e in bbox_end_diff) - step_limits += Vec3D[int](*(int(e) for e in step_start_partial)) - step_limits += Vec3D[int](*(int(e) for e in step_end_partial)) + atomic_step_limits += Vec3D[int](*(int(e) for e in step_start_partial)) + atomic_step_limits += Vec3D[int](*(int(e) for e in step_end_partial)) logger.debug( f"Exact bbox requested; out of {self.bbox.bounds}," f" full cubes are in {bbox_snapped.bounds}, given offset" f" {stride_start_offset_in_unit}{self.bbox.unit} with chunk size" f" {self.chunk_size_in_unit}{self.bbox.unit}." ) - # Use `__setattr__` to keep the object frozen. - object.__setattr__(self, "step_limits", step_limits) + object.__setattr__(self, "atomic_step_limits", atomic_step_limits) object.__setattr__(self, "bbox_snapped", bbox_snapped) object.__setattr__(self, "step_start_partial", step_start_partial) object.__setattr__(self, "step_end_partial", step_end_partial) @@ -139,7 +137,6 @@ def _attrs_post_init_nonexact(self) -> None: step_start_partial = (False, False, False) step_end_partial = (False, False, False) if self.stride_start_offset is not None: - # align stride_start_offset to just larger than the start of the bbox stride_start_offset_in_unit = Vec3D[float](*self.stride_start_offset) * self.resolution stride_start_offset_in_unit += ( (self.bbox.start - stride_start_offset_in_unit) @@ -154,7 +151,7 @@ def _attrs_post_init_nonexact(self) -> None: .snapped( grid_offset=stride_start_offset_in_unit, grid_size=self.stride_in_unit, - mode=self.mode, # type: ignore #mypy doesn't realise that mode has been checked + mode=self.mode, # type: ignore ) .translated_end(self.chunk_size_in_unit, resolution=Vec3D(1, 1, 1)) ) @@ -181,16 +178,17 @@ def _attrs_post_init_nonexact(self) -> None: ) ) ) + atomic_step_limits = Vec3D[int](0, 0, 0) if self.mode == "shrink": - step_limits = floor(round(step_limits_snapped, VEC3D_PRECISION)) - if not step_limits_raw.allclose(step_limits): + atomic_step_limits = floor(round(step_limits_snapped, VEC3D_PRECISION)) + if not step_limits_raw.allclose(atomic_step_limits): rounded_bbox_bounds = tuple( ( bbox_snapped.bounds[i][0], ( bbox_snapped.bounds[i][0] + self.chunk_size_in_unit[i] - + (step_limits[i] - 1) * self.stride_in_unit[i] + + (atomic_step_limits[i] - 1) * self.stride_in_unit[i] ), ) for i in range(3) @@ -202,15 +200,15 @@ def _attrs_post_init_nonexact(self) -> None: f" {self.chunk_size_in_unit}{self.bbox.unit}." ) if self.mode == "expand": - step_limits = ceil(round(step_limits_snapped, VEC3D_PRECISION)) - if not step_limits_raw.allclose(step_limits): + atomic_step_limits = ceil(round(step_limits_snapped, VEC3D_PRECISION)) + if not step_limits_raw.allclose(atomic_step_limits): rounded_bbox_bounds = tuple( ( bbox_snapped.bounds[i][0], ( bbox_snapped.bounds[i][0] + self.chunk_size_in_unit[i] - + step_limits[i] * self.stride_in_unit[i] + + atomic_step_limits[i] * self.stride_in_unit[i] ), ) for i in range(3) @@ -221,8 +219,7 @@ def _attrs_post_init_nonexact(self) -> None: f" {self.stride_in_unit}{self.bbox.unit} with chunk size" f" {self.chunk_size_in_unit}{self.bbox.unit}." ) - # Use `__setattr__` to keep the object frozen. - object.__setattr__(self, "step_limits", step_limits) + object.__setattr__(self, "atomic_step_limits", atomic_step_limits) object.__setattr__(self, "bbox_snapped", bbox_snapped) object.__setattr__(self, "step_start_partial", step_start_partial) object.__setattr__(self, "step_end_partial", step_end_partial) @@ -252,20 +249,7 @@ def get_all_chunk_bboxes(self) -> List[BBox3D]: ] # TODO: generator? return result - def get_nth_chunk_bbox(self, n: int) -> BBox3D: - """Get nth chunk bbox, in order. - - :param n: Integer chunk index. - :return: Volumetric index for the chunk, including - ``self.desired_resolution`` and the slice representation of the region - at ``self.index_resolution``. - - """ - steps_along_dim = Vec3D[int]( - n % self.step_limits[0], - (n // self.step_limits[0]) % self.step_limits[1], - (n // (self.step_limits[0] * self.step_limits[1])) % self.step_limits[2], - ) + def _get_atomic_bbox(self, steps_along_dim: Vec3D[int]) -> BBox3D: if self.mode in ("shrink", "expand"): chunk_origin_in_unit = [ self.bbox_snapped.bounds[i][0] + self.stride_in_unit[i] * steps_along_dim[i] @@ -289,7 +273,10 @@ def get_nth_chunk_bbox(self, n: int) -> BBox3D: if steps_along_dim[i] == 0 and self.step_start_partial[i]: chunk_origin_in_unit[i] = self.bbox.start[i] chunk_end_in_unit[i] = self.bbox_snapped.start[i] - if steps_along_dim[i] == self.step_limits[i] - 1 and self.step_end_partial[i]: + if ( + steps_along_dim[i] == self.atomic_step_limits[i] - 1 + and self.step_end_partial[i] + ): chunk_origin_in_unit[i] = self.bbox_snapped.end[i] chunk_end_in_unit[i] = self.bbox.end[i] slices = ( @@ -298,5 +285,36 @@ def get_nth_chunk_bbox(self, n: int) -> BBox3D: slice(chunk_origin_in_unit[2], chunk_end_in_unit[2]), ) - result = BBox3D.from_slices(slices) - return result + return BBox3D.from_slices(slices) + + def get_nth_chunk_bbox(self, n: int) -> BBox3D: + """Get nth chunk bbox, in order. + + :param n: Integer chunk index. + :return: Volumetric index for the chunk, including + ``self.desired_resolution`` and the slice representation of the region + at ``self.index_resolution``. + + """ + super_steps_along_dim = Vec3D[int]( + n % self.step_limits[0], + (n // self.step_limits[0]) % self.step_limits[1], + (n // (self.step_limits[0] * self.step_limits[1])) % self.step_limits[2], + ) + + atomic_start_steps = super_steps_along_dim * self.super_factors + atomic_end_steps = Vec3D[int]( + *( + min( + (super_steps_along_dim[i] + 1) * self.super_factors[i], + self.atomic_step_limits[i], + ) + for i in range(3) + ) + ) + + bbox_start = self._get_atomic_bbox(atomic_start_steps) + if not self.max_superchunk_size: + return bbox_start + bbox_end = self._get_atomic_bbox(atomic_end_steps - Vec3D[int](1, 1, 1)) + return bbox_start.supremum(bbox_end) From 6f9f964f9ced9542524c24a4290fd1c0a9ca4252 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 2 Jan 2026 01:54:21 -0800 Subject: [PATCH 2/9] chore: update subchunkable tests to nondeprecated --- ...est_uint8_exc_auto_divisibility_and_expand_bbox_backend.cue | 3 +++ ...uint8_exc_auto_divisibility_and_shrink_processing_chunk.cue | 3 +++ ...nt8_exc_auto_divisibility_but_no_expand_bbox_processing.cue | 3 +++ ...t8_exc_bbox_non_integral_without_expand_bbox_resolution.cue | 3 +++ ...thout_expand_bbox_resolution_but_expand_bbox_processing.cue | 3 +++ ...hout_expand_bbox_resolution_but_shrink_processing_chunk.cue | 3 +++ .../subchunkable/specs/exc/test_uint8_exc_blend_too_large.cue | 3 +++ .../specs/exc/test_uint8_exc_both_bbox_and_auto_bbox.cue | 3 +++ .../specs/exc/test_uint8_exc_both_bbox_and_coords.cue | 3 +++ .../specs/exc/test_uint8_exc_both_coords_and_auto_bbox.cue | 3 +++ .../subchunkable/specs/exc/test_uint8_exc_both_fn_and_op.cue | 3 +++ .../specs/exc/test_uint8_exc_both_fn_semaphores_and_op.cue | 3 +++ .../specs/exc/test_uint8_exc_defer_but_skip_intermediaries.cue | 3 +++ .../specs/exc/test_uint8_exc_defer_on_not_toplevel.cue | 3 +++ .../specs/exc/test_uint8_exc_dst_tighten_bounds.cue | 3 +++ .../test_uint8_exc_generate_ng_link_but_not_print_summary.cue | 3 +++ .../exc/test_uint8_exc_level_intermediaries_dirs_not_equal.cue | 3 +++ ...xc_max_reduction_chunk_size_too_small_for_backend_chunk.cue | 3 +++ ...max_reduction_chunk_size_too_small_for_processing_chunk.cue | 3 +++ .../exc/test_uint8_exc_no_bbox_or_coords_or_auto_bbox.cue | 3 +++ .../subchunkable/specs/exc/test_uint8_exc_no_dst_but_defer.cue | 1 + .../exc/test_uint8_exc_no_dst_but_expand_bbox_backend.cue | 1 + .../exc/test_uint8_exc_no_dst_but_max_reduction_chunk_size.cue | 1 + .../exc/test_uint8_exc_no_dst_but_not_skip_intermediaries.cue | 1 + .../subchunkable/specs/exc/test_uint8_exc_no_fn_or_op.cue | 3 +++ .../exc/test_uint8_exc_nondivisible_and_not_recommendable.cue | 3 +++ .../exc/test_uint8_exc_nondivisible_but_recommendable.cue | 3 +++ ...ot_skip_intermediaries_but_no_level_intermediaries_dirs.cue | 3 +++ .../test_uint8_exc_processing_gap_but_auto_divisibility.cue | 3 +++ .../test_uint8_exc_processing_gap_but_blend_pad_toplevel.cue | 3 +++ .../test_uint8_exc_processing_gap_but_expand_bbox_backend.cue | 3 +++ ...est_uint8_exc_processing_gap_but_expand_bbox_resolution.cue | 3 +++ ...st_uint8_exc_processing_gap_but_shrink_processing_chunk.cue | 3 +++ .../specs/exc/test_uint8_exc_processing_gap_but_uneven.cue | 3 +++ .../specs/exc/test_uint8_exc_seq_of_seq_not_equal.cue | 3 +++ ..._exc_shrink_processing_chunk_and_expand_bbox_processing.cue | 3 +++ .../exc/test_uint8_exc_skip_intermediaries_but_blend_pad.cue | 3 +++ .../exc/test_uint8_exc_skip_intermediaries_but_crop_pad.cue | 3 +++ .../specs/exc/test_uint8_exc_skip_intermediaries_but_defer.cue | 3 +++ ...8_exc_skip_intermediaries_but_level_intermediaries_dirs.cue | 3 +++ ...uint8_exc_skip_intermediaries_but_reduction_worker_type.cue | 3 +++ tests/integration/subchunkable/specs/test_float32_copy.cue | 3 +++ .../integration/subchunkable/specs/test_float32_copy_blend.cue | 3 +++ .../integration/subchunkable/specs/test_float32_copy_crop.cue | 3 +++ .../specs/test_float32_copy_multilevel_checkerboard.cue | 3 +++ .../specs/test_float32_copy_multilevel_no_checkerboard.cue | 3 +++ .../subchunkable/specs/test_float32_copy_writeproc.cue | 3 +++ .../test_float32_copy_writeproc_multilevel_checkerboard.cue | 3 +++ ...float32_copy_writeproc_multilevel_checkerboard_parallel.cue | 3 +++ .../test_float32_copy_writeproc_multilevel_no_checkerboard.cue | 3 +++ .../subchunkable/specs/test_uint8_copy_auto_bbox.cue | 3 +++ .../subchunkable/specs/test_uint8_copy_auto_divisibility.cue | 3 +++ tests/integration/subchunkable/specs/test_uint8_copy_bbox.cue | 3 +++ tests/integration/subchunkable/specs/test_uint8_copy_blend.cue | 3 +++ .../integration/subchunkable/specs/test_uint8_copy_coords.cue | 3 +++ tests/integration/subchunkable/specs/test_uint8_copy_crop.cue | 3 +++ tests/integration/subchunkable/specs/test_uint8_copy_defer.cue | 3 +++ .../specs/test_uint8_copy_dont_skip_intermediaries.cue | 3 +++ .../subchunkable/specs/test_uint8_copy_expand_bbox_backend.cue | 3 +++ .../specs/test_uint8_copy_expand_bbox_processing.cue | 3 +++ .../specs/test_uint8_copy_expand_bbox_resolution.cue | 3 +++ ...py_expand_bbox_resolution_backend_processing_do_nothing.cue | 3 +++ .../subchunkable/specs/test_uint8_copy_fn_semaphores.cue | 3 +++ .../specs/test_uint8_copy_multilevel_checkerboard.cue | 3 +++ .../test_uint8_copy_multilevel_checkerboard_cache_up_to_l0.cue | 3 +++ .../specs/test_uint8_copy_multilevel_no_checkerboard.cue | 3 +++ .../subchunkable/specs/test_uint8_copy_no_op_kwargs.cue | 3 +++ tests/integration/subchunkable/specs/test_uint8_copy_op.cue | 3 +++ .../integration/subchunkable/specs/test_uint8_copy_postpad.cue | 3 +++ .../subchunkable/specs/test_uint8_copy_processing_gap.cue | 3 +++ .../specs/test_uint8_copy_shrink_processing_chunk.cue | 3 +++ .../subchunkable/specs/test_uint8_copy_skip_intermediaries.cue | 3 +++ .../specs/test_uint8_copy_top_level_checkerboard.cue | 3 +++ .../subchunkable/specs/test_uint8_copy_writeproc.cue | 3 +++ .../specs/test_uint8_copy_writeproc_multilevel.cue | 3 +++ .../test_uint8_copy_writeproc_multilevel_checkerboard.cue | 3 +++ ...t8_copy_writeproc_multilevel_checkerboard_cache_up_to_0.cue | 3 +++ .../test_uint8_copy_writeproc_multilevel_no_checkerboard.cue | 3 +++ tests/integration/subchunkable/specs/test_uint8_no_dst.cue | 1 + 79 files changed, 227 insertions(+) diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_and_expand_bbox_backend.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_and_expand_bbox_backend.cue index d1a1c9d95..4b6d7ab5f 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_and_expand_bbox_backend.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_and_expand_bbox_backend.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_auto_divisibility_and_expand_bbox_backend" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_and_shrink_processing_chunk.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_and_shrink_processing_chunk.cue index da49f344a..74c2ed441 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_and_shrink_processing_chunk.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_and_shrink_processing_chunk.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_auto_divisibility_and_shrink_processing_chunk" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_but_no_expand_bbox_processing.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_but_no_expand_bbox_processing.cue index 46eb8ba9b..e83356c35 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_but_no_expand_bbox_processing.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_auto_divisibility_but_no_expand_bbox_processing.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_auto_divisibility_but_no_expand_bbox_processing" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution.cue index f21bb2123..4aabfc890 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution" #BBOX: { @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_expand_bbox_processing.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_expand_bbox_processing.cue index a4c267063..4a8893b0c 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_expand_bbox_processing.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_expand_bbox_processing.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_expand_bbox_processing" #BBOX: { @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_shrink_processing_chunk.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_shrink_processing_chunk.cue index ffa80a6f3..7d3187f76 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_shrink_processing_chunk.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_shrink_processing_chunk.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_bbox_non_integral_without_expand_bbox_resolution_but_shrink_processing_chunk" #BBOX: { @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_blend_too_large.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_blend_too_large.cue index cb8cc20d2..0bdf87b5a 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_blend_too_large.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_blend_too_large.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_blend_too_large" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_bbox_and_auto_bbox.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_bbox_and_auto_bbox.cue index 561584078..6109023e9 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_bbox_and_auto_bbox.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_bbox_and_auto_bbox.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_both_bbox_and_auto_bbox" #BBOX: { @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_bbox_and_coords.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_bbox_and_coords.cue index 233e2935e..b9eec08ff 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_bbox_and_coords.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_bbox_and_coords.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_both_bbox_and_coords" #BBOX: { @@ -29,6 +30,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } start_coord: [64 * 1024, 64 * 1024, 2000] end_coord: [96 * 1024, 96 * 1024, 2005] diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_coords_and_auto_bbox.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_coords_and_auto_bbox.cue index abd732ae7..5356f6d89 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_coords_and_auto_bbox.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_coords_and_auto_bbox.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_both_coords_and_auto_bbox" @@ -22,6 +23,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } start_coord: [64 * 1024, 64 * 1024, 2000] end_coord: [96 * 1024, 96 * 1024, 2005] diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_fn_and_op.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_fn_and_op.cue index a1a2d8ccd..b79747be3 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_fn_and_op.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_fn_and_op.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_both_fn_and_op" #BBOX: { @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_fn_semaphores_and_op.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_fn_semaphores_and_op.cue index 5536d1410..04b3d6de7 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_fn_semaphores_and_op.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_both_fn_semaphores_and_op.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_both_fn_semaphores_and_op" #BBOX: { @@ -29,6 +30,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_defer_but_skip_intermediaries.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_defer_but_skip_intermediaries.cue index aca5b4953..3068541df 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_defer_but_skip_intermediaries.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_defer_but_skip_intermediaries.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_defer_but_skip_intermediaries" @@ -33,6 +34,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_defer_on_not_toplevel.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_defer_on_not_toplevel.cue index 44e6b0948..3f16522db 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_defer_on_not_toplevel.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_defer_on_not_toplevel.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_defer_on_not_toplevel" @@ -33,6 +34,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_dst_tighten_bounds.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_dst_tighten_bounds.cue index fb03ae3b5..cf0584fa1 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_dst_tighten_bounds.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_dst_tighten_bounds.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_dst_tighten_bounds" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } dst_tighten_bounds: true } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_generate_ng_link_but_not_print_summary.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_generate_ng_link_but_not_print_summary.cue index dc5356a8b..f7c0d8b1f 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_generate_ng_link_but_not_print_summary.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_generate_ng_link_but_not_print_summary.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_generate_ng_link_but_not_print_summary" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } generate_ng_link: true print_summary: false diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_level_intermediaries_dirs_not_equal.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_level_intermediaries_dirs_not_equal.cue index a6484629d..8a39bdd50 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_level_intermediaries_dirs_not_equal.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_level_intermediaries_dirs_not_equal.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_shrink_processing_and_expand_bbox_processing" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_max_reduction_chunk_size_too_small_for_backend_chunk.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_max_reduction_chunk_size_too_small_for_backend_chunk.cue index 2537e176c..189061ecf 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_max_reduction_chunk_size_too_small_for_backend_chunk.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_max_reduction_chunk_size_too_small_for_backend_chunk.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_max_reduction_chunk_size_too_small_for_backend_chunk" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_max_reduction_chunk_size_too_small_for_processing_chunk.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_max_reduction_chunk_size_too_small_for_processing_chunk.cue index 018d5098a..e38c8d618 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_max_reduction_chunk_size_too_small_for_processing_chunk.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_max_reduction_chunk_size_too_small_for_processing_chunk.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_max_reduction_chunk_size_too_small_for_processing_chunk" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_bbox_or_coords_or_auto_bbox.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_bbox_or_coords_or_auto_bbox.cue index 8ed5adec0..bf366056f 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_bbox_or_coords_or_auto_bbox.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_bbox_or_coords_or_auto_bbox.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_no_bbox_or_coords_or_auto_bbox" @@ -22,6 +23,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_defer.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_defer.cue index 32ffb5cd5..0703db841 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_defer.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_defer.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_no_dst_but_defer" diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_expand_bbox_backend.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_expand_bbox_backend.cue index 162d1a690..7a4a9da36 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_expand_bbox_backend.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_expand_bbox_backend.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_no_dst_but_expand_bbox_backend" diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_max_reduction_chunk_size.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_max_reduction_chunk_size.cue index 2302db36e..3ad47289f 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_max_reduction_chunk_size.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_max_reduction_chunk_size.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_no_dst_but_max_reduction_chunk_size" diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_not_skip_intermediaries.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_not_skip_intermediaries.cue index ec9e7e11f..7ad5a7c44 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_not_skip_intermediaries.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_dst_but_not_skip_intermediaries.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_no_dst_but_not_skip_intermediaries" diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_fn_or_op.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_fn_or_op.cue index a6f000540..d8a2b7988 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_fn_or_op.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_no_fn_or_op.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_no_fn_or_op" #BBOX: { @@ -26,6 +27,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_nondivisible_and_not_recommendable.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_nondivisible_and_not_recommendable.cue index 2e342e20f..b8bcf4c62 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_nondivisible_and_not_recommendable.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_nondivisible_and_not_recommendable.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_nondivisible_and_not_recommendable" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_nondivisible_but_recommendable.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_nondivisible_but_recommendable.cue index 24cd02e9b..1f11b3d8d 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_nondivisible_but_recommendable.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_nondivisible_but_recommendable.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_nondivisible_but_recommendable" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_not_skip_intermediaries_but_no_level_intermediaries_dirs.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_not_skip_intermediaries_but_no_level_intermediaries_dirs.cue index d2b48cd68..9ae376e77 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_not_skip_intermediaries_but_no_level_intermediaries_dirs.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_not_skip_intermediaries_but_no_level_intermediaries_dirs.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_not_skip_intermediaries_but_no_level_intermediaries_dirs" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_auto_divisibility.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_auto_divisibility.cue index 87f172561..94bb60126 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_auto_divisibility.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_auto_divisibility.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_processing_gap_but_auto_divisibility" @@ -37,6 +38,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_blend_pad_toplevel.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_blend_pad_toplevel.cue index 25901bae5..a6ae58972 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_blend_pad_toplevel.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_blend_pad_toplevel.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_processing_gap_but_blend_toplevel" @@ -37,6 +38,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_expand_bbox_backend.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_expand_bbox_backend.cue index 57229f57b..bb70a4f3b 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_expand_bbox_backend.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_expand_bbox_backend.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_processing_gap_but_expand_bbox_backend" @@ -38,6 +39,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_expand_bbox_resolution.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_expand_bbox_resolution.cue index 66e6c3d48..a0ce3bde4 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_expand_bbox_resolution.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_expand_bbox_resolution.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_processing_gap_but_expand_bbox_resolution" @@ -38,6 +39,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_shrink_processing_chunk.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_shrink_processing_chunk.cue index 77a7720e5..48045a9b4 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_shrink_processing_chunk.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_shrink_processing_chunk.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_processing_gap_but_shrink_processing_chunk" @@ -38,6 +39,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_uneven.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_uneven.cue index 8b85d7f15..d60f1d5d1 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_uneven.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_processing_gap_but_uneven.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_processing_gap_but_uneven" @@ -37,6 +38,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_seq_of_seq_not_equal.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_seq_of_seq_not_equal.cue index 23b20217b..3022f8567 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_seq_of_seq_not_equal.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_seq_of_seq_not_equal.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_seq_of_seq_not_equal" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_shrink_processing_chunk_and_expand_bbox_processing.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_shrink_processing_chunk_and_expand_bbox_processing.cue index 987cf67de..156624b60 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_shrink_processing_chunk_and_expand_bbox_processing.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_shrink_processing_chunk_and_expand_bbox_processing.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_shrink_processing_chunk_and_expand_bbox_processing" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_blend_pad.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_blend_pad.cue index 0977a529b..4e91ef56c 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_blend_pad.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_blend_pad.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_skip_intermediaries_but_blend_pad" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_crop_pad.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_crop_pad.cue index 7575b0755..7708b6b73 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_crop_pad.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_crop_pad.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_skip_intermediaries_but_crop_pad" @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_defer.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_defer.cue index 4fdb9f437..d8746349e 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_defer.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_defer.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_skip_intermediaries_but_defer" @@ -33,6 +34,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_level_intermediaries_dirs.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_level_intermediaries_dirs.cue index 0f0cc59a3..ce9694dc0 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_level_intermediaries_dirs.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_level_intermediaries_dirs.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_skip_intermediaries_but_level_intermediaries_dirs" @@ -33,6 +34,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_reduction_worker_type.cue b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_reduction_worker_type.cue index 23fa003e3..b11ed5ee8 100644 --- a/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_reduction_worker_type.cue +++ b/tests/integration/subchunkable/specs/exc/test_uint8_exc_skip_intermediaries_but_reduction_worker_type.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_exc_skip_intermediaries_but_reduction_worker_type" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } reduction_worker_type: "some_worker_type" } diff --git a/tests/integration/subchunkable/specs/test_float32_copy.cue b/tests/integration/subchunkable/specs/test_float32_copy.cue index 369a992a2..c4fe5c0f6 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_float32_copy_blend.cue b/tests/integration/subchunkable/specs/test_float32_copy_blend.cue index 6c68b8b61..49f13d921 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy_blend.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy_blend.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy_blend" @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_float32_copy_crop.cue b/tests/integration/subchunkable/specs/test_float32_copy_crop.cue index d8ba0c1e7..539d7aca5 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy_crop.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy_crop.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy_crop" @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_float32_copy_multilevel_checkerboard.cue b/tests/integration/subchunkable/specs/test_float32_copy_multilevel_checkerboard.cue index 57ba9392d..cf68322c5 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy_multilevel_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy_multilevel_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy_multilevel_checkerboard" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_float32_copy_multilevel_no_checkerboard.cue b/tests/integration/subchunkable/specs/test_float32_copy_multilevel_no_checkerboard.cue index 38721108c..c659ea35d 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy_multilevel_no_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy_multilevel_no_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy_multilevel_no_checkerboard" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_float32_copy_writeproc.cue b/tests/integration/subchunkable/specs/test_float32_copy_writeproc.cue index e9410c4e8..9a1e604b2 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy_writeproc.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy_writeproc.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy_writeproc" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_checkerboard.cue b/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_checkerboard.cue index 12f84319c..301104490 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy_writeproc_multilevel_checkerboard" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_checkerboard_parallel.cue b/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_checkerboard_parallel.cue index 54ea4be52..88c54539a 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_checkerboard_parallel.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_checkerboard_parallel.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy_writeproc_multilevel_checkerboard_parallel" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_no_checkerboard.cue b/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_no_checkerboard.cue index efb153f1d..d017f9195 100644 --- a/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_no_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_float32_copy_writeproc_multilevel_no_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_float32" #DST_PATH: "assets/outputs/test_float32_copy_writeproc_multilevel_no_checkerboard" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_auto_bbox.cue b/tests/integration/subchunkable/specs/test_uint8_copy_auto_bbox.cue index aafc29f15..be6d5014c 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_auto_bbox.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_auto_bbox.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_auto_bbox" @@ -23,6 +24,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_auto_divisibility.cue b/tests/integration/subchunkable/specs/test_uint8_copy_auto_divisibility.cue index de4d6c116..fcfd41e35 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_auto_divisibility.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_auto_divisibility.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_auto_divisibility" @@ -33,6 +34,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_bbox.cue b/tests/integration/subchunkable/specs/test_uint8_copy_bbox.cue index 4204cebeb..e22c32bb2 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_bbox.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_bbox.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_bbox" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_blend.cue b/tests/integration/subchunkable/specs/test_uint8_copy_blend.cue index ade5393f8..0a1b32d42 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_blend.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_blend.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_blend" @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_coords.cue b/tests/integration/subchunkable/specs/test_uint8_copy_coords.cue index 8db6540fa..cf347fcf1 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_coords.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_coords.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_coords" @@ -25,6 +26,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_crop.cue b/tests/integration/subchunkable/specs/test_uint8_copy_crop.cue index b93dc73c5..39e8f283e 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_crop.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_crop.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_crop" @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_defer.cue b/tests/integration/subchunkable/specs/test_uint8_copy_defer.cue index fcf8c32ca..2a483a26f 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_defer.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_defer.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_defer" @@ -35,6 +36,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_dont_skip_intermediaries.cue b/tests/integration/subchunkable/specs/test_uint8_copy_dont_skip_intermediaries.cue index 44d6d2f3f..f9a5056fa 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_dont_skip_intermediaries.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_dont_skip_intermediaries.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_dont_skip_intermediaries" @@ -28,6 +29,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_backend.cue b/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_backend.cue index d77a9064c..2be8de4db 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_backend.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_backend.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_expand_bbox_backend" @@ -33,6 +34,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_processing.cue b/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_processing.cue index a7716d571..e4c45c79e 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_processing.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_processing.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_expand_bbox_processing" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_resolution.cue b/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_resolution.cue index a0091f6f0..6290958f0 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_resolution.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_resolution.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_expand_bbox_resolution" @@ -33,6 +34,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_resolution_backend_processing_do_nothing.cue b/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_resolution_backend_processing_do_nothing.cue index fc70964a3..9dac3d43f 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_resolution_backend_processing_do_nothing.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_expand_bbox_resolution_backend_processing_do_nothing.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_expand_bbox_resolution_backend_processing_do_nothing" @@ -35,6 +36,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_fn_semaphores.cue b/tests/integration/subchunkable/specs/test_uint8_copy_fn_semaphores.cue index 10b37df31..24bb7f6d9 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_fn_semaphores.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_fn_semaphores.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_fn_semaphores" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_checkerboard.cue b/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_checkerboard.cue index ec978fc90..1d8267599 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_multilevel_checkerboard" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_checkerboard_cache_up_to_l0.cue b/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_checkerboard_cache_up_to_l0.cue index 849683c67..28f77cd4b 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_checkerboard_cache_up_to_l0.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_checkerboard_cache_up_to_l0.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_multilevel_checkerboard_cache_up_to_l0" @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_no_checkerboard.cue b/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_no_checkerboard.cue index 84ba386ba..140e4475b 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_no_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_multilevel_no_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_multilevel_no_checkerboard" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_no_op_kwargs.cue b/tests/integration/subchunkable/specs/test_uint8_copy_no_op_kwargs.cue index 644766671..76a3fbe94 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_no_op_kwargs.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_no_op_kwargs.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_no_op_kwargs" @@ -24,6 +25,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_op.cue b/tests/integration/subchunkable/specs/test_uint8_copy_op.cue index b7cda3678..161fb970c 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_op.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_op.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_op" @@ -33,6 +34,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_postpad.cue b/tests/integration/subchunkable/specs/test_uint8_copy_postpad.cue index 72a2e35de..2564c2cd1 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_postpad.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_postpad.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_postpad" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_processing_gap.cue b/tests/integration/subchunkable/specs/test_uint8_copy_processing_gap.cue index e0aa0a0e4..c4bf89363 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_processing_gap.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_processing_gap.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_processing_gap" @@ -37,6 +38,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_shrink_processing_chunk.cue b/tests/integration/subchunkable/specs/test_uint8_copy_shrink_processing_chunk.cue index f561e351b..c17ed5ecd 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_shrink_processing_chunk.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_shrink_processing_chunk.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_shrink_processing_chunk" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_skip_intermediaries.cue b/tests/integration/subchunkable/specs/test_uint8_copy_skip_intermediaries.cue index 0d5d10599..494c148e3 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_skip_intermediaries.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_skip_intermediaries.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_skip_intermediaries" @@ -28,6 +29,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_top_level_checkerboard.cue b/tests/integration/subchunkable/specs/test_uint8_copy_top_level_checkerboard.cue index caa35471b..ae952d444 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_top_level_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_top_level_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_top_level_checkerboard" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true } } diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc.cue b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc.cue index bedf0acaf..9ef59162b 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_writeproc" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel.cue b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel.cue index aeef9868b..f9d9aca02 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_writeproc_multilevel" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_checkerboard.cue b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_checkerboard.cue index 77ead8449..0ca12965e 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_writeproc_multilevel_checkerboard" @@ -31,6 +32,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_checkerboard_cache_up_to_0.cue b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_checkerboard_cache_up_to_0.cue index b7262f222..9eac24a18 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_checkerboard_cache_up_to_0.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_checkerboard_cache_up_to_0.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_writeproc_multilevel_checkerboard_cache_up_to_0" @@ -32,6 +33,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_no_checkerboard.cue b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_no_checkerboard.cue index 988cd3bde..77a880a92 100644 --- a/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_no_checkerboard.cue +++ b/tests/integration/subchunkable/specs/test_uint8_copy_writeproc_multilevel_no_checkerboard.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_copy_writeproc_multilevel_no_checkerboard" @@ -30,6 +31,8 @@ "@type": "build_cv_layer" path: #DST_PATH info_reference_path: #SRC_PATH + info_scales: [[128, 128, 40]] + info_inherit_all_params: true write_procs: [ { "@type": "lambda" diff --git a/tests/integration/subchunkable/specs/test_uint8_no_dst.cue b/tests/integration/subchunkable/specs/test_uint8_no_dst.cue index e03ad60d7..db74558eb 100644 --- a/tests/integration/subchunkable/specs/test_uint8_no_dst.cue +++ b/tests/integration/subchunkable/specs/test_uint8_no_dst.cue @@ -1,3 +1,4 @@ +"@version": "0.4" #SRC_PATH: "assets/inputs/fafb_v15_img_128_128_40-2048-3072_2000-2050_uint8" #DST_PATH: "assets/outputs/test_uint8_no_dst" From 5c8066efaf125a303938d0cf3a7784ceb8173a15 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Mon, 5 Jan 2026 21:19:16 -0800 Subject: [PATCH 3/9] feat: separate process handles SQS upkeep --- zetta_utils/mazepa/upkeep_handlers.py | 360 ++++++++++++++++++ zetta_utils/mazepa/worker.py | 211 +++++++--- .../configurations/worker_pool.py | 51 +-- zetta_utils/message_queues/sqs/queue.py | 2 +- zetta_utils/message_queues/sqs/utils.py | 56 ++- 5 files changed, 579 insertions(+), 101 deletions(-) create mode 100644 zetta_utils/mazepa/upkeep_handlers.py diff --git a/zetta_utils/mazepa/upkeep_handlers.py b/zetta_utils/mazepa/upkeep_handlers.py new file mode 100644 index 000000000..a3f96cfcb --- /dev/null +++ b/zetta_utils/mazepa/upkeep_handlers.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +import logging +import multiprocessing +import os +import threading +import time +from dataclasses import dataclass +from typing import Callable + +import tenacity + +from zetta_utils import log +from zetta_utils.common.partial import ComparablePartial + +logger = log.get_logger("mazepa") + + +def perform_direct_upkeep( + extend_lease_fn: Callable, + extend_duration: int, + task_start_time: float, +) -> None: + """ + Perform upkeep by directly calling extend_lease_fn. + + Used as a fallback for non-SQS queues where the process-based handler + cannot be used. + """ + current_time = time.time() + elapsed_since_start = current_time - task_start_time + logger.debug( + f"UPKEEP: [T+{elapsed_since_start:.1f}s] Timer fired, calling extend_lease_fn directly" + ) + try: + start_time = time.time() + extend_lease_fn(extend_duration) + api_duration = time.time() - start_time + logger.debug( + f"UPKEEP: [T+{elapsed_since_start:.1f}s] Successfully extended lease by " + f"{extend_duration}s (API call took {api_duration:.1f}s)" + ) + except tenacity.RetryError as e: # pragma: no cover + logger.error(f"UPKEEP: Failed to extend lease after retries: {e}") + except Exception as e: # pragma: no cover # pylint: disable=broad-except + logger.error(f"UPKEEP: Unexpected error: {type(e).__name__}: {e}") + + +def extract_sqs_metadata(extend_lease_fn: Callable) -> dict | None: + """ + Extract SQS metadata from extend_lease_fn if it's a ComparablePartial wrapping + an SQS queue's _extend_msg_lease method. + + Returns a dict with queue_name, region_name, endpoint_url, receipt_handle if found, + or None if this is not an SQS-based extend function. + """ + if not isinstance(extend_lease_fn, ComparablePartial): + return None + + msg = extend_lease_fn.kwargs.get("msg") + if msg is None: + return None + + # Check if msg has the SQS-specific attributes + if not all(hasattr(msg, attr) for attr in ("receipt_handle", "queue_name", "region_name")): + return None + + return { + "receipt_handle": msg.receipt_handle, + "queue_name": msg.queue_name, + "region_name": msg.region_name, + "endpoint_url": getattr(msg, "endpoint_url", None), + } + + +@dataclass +class UpkeepCommand: + """Command sent to the SQS upkeep handler process.""" + + action: str # "start_upkeep", "stop_upkeep", or "shutdown" + # Required for start_upkeep + task_id: str | None = None + receipt_handle: str | None = None + visibility_timeout: int | None = None + interval_sec: float | None = None + queue_name: str | None = None + region_name: str | None = None + endpoint_url: str | None = None + + +def _is_parent_alive(parent_pid: int) -> bool: # pragma: no cover + """Check if the parent process is still alive.""" + try: + os.kill(parent_pid, 0) # Signal 0 just checks if process exists + return True + except OSError: + return False + + +def run_sqs_upkeep_handler( # pylint: disable=too-many-statements + command_queue: multiprocessing.Queue, + log_level: str = "INFO", + parent_pid: int | None = None, + interval_sec: float = 5.0, +) -> None: + """ + Main loop for the SQS upkeep handler process. + + Runs in a separate process to handle SQS visibility extensions. This isolates + SQS operations from the main worker process's GIL, ensuring that heavy CPU work + in the main process doesn't delay upkeep operations. + + The handler manages its own timer, so timing is not affected by the main + process's CPU usage. + + If parent_pid is provided, the handler will exit if the parent process dies. + This prevents orphaned handler processes when workers are force-killed. + """ + # pylint: disable=import-outside-toplevel + from zetta_utils.mazepa.worker import worker_init + from zetta_utils.message_queues.sqs import utils + + # Initialize the process (logging, signal handlers, etc.) + # Don't set start method or load train/inference for the upkeep handler + worker_init(log_level=log_level) + + logger.info( + "SQS_HANDLER: Upkeep handler process started (PID: %d, parent: %s)", + multiprocessing.current_process().pid, + parent_pid, + ) + + # Track active upkeep tasks: task_id -> (stop_event, thread) + active_upkeeps: dict[str, tuple[threading.Event, threading.Thread]] = {} + + def _run_upkeep_loop( + task_id: str, + stop_event: threading.Event, + receipt_handle: str, + visibility_timeout: int, + interval_sec: float, + queue_name: str, + region_name: str, + endpoint_url: str | None, + ): + """Timer loop that extends visibility at regular intervals.""" + task_start_time = time.time() + logger.info( + f"SQS_HANDLER: [{task_id}] Starting upkeep loop: interval={interval_sec}s, " + f"extend_by={visibility_timeout}s" + ) + + while not stop_event.wait(timeout=interval_sec): + elapsed = time.time() - task_start_time + logger.debug( + f"SQS_HANDLER: [{task_id}] [T+{elapsed:.1f}s] Extending visibility to " + f"{visibility_timeout}s for queue '{queue_name}'" + ) + try: + api_start = time.time() + utils.change_message_visibility( + receipt_handle=receipt_handle, + visibility_timeout=visibility_timeout, + queue_name=queue_name, + region_name=region_name, + endpoint_url=endpoint_url, + ) + api_duration = time.time() - api_start + logger.debug( + f"SQS_HANDLER: [{task_id}] [T+{elapsed:.1f}s] Successfully extended " + f"(API took {api_duration:.1f}s)" + ) + except Exception as e: # pylint: disable=broad-except + logger.error( + f"SQS_HANDLER: [{task_id}] Failed to extend visibility: " + f"{type(e).__name__}: {e}" + ) + + elapsed = time.time() - task_start_time + logger.info(f"SQS_HANDLER: [{task_id}] Upkeep loop stopped after {elapsed:.1f}s") + + while True: + try: + # Use timeout so we can periodically check if parent is alive + try: + cmd = command_queue.get(timeout=interval_sec) + except: # pylint: disable=bare-except # pragma: no cover + # Check if parent is still alive + if parent_pid is not None and not _is_parent_alive(parent_pid): + logger.warning( + "SQS_HANDLER: Parent process %d died, exiting handler", parent_pid + ) + # Stop all active upkeeps before exiting + for task_id, (stop_event, thread) in active_upkeeps.items(): + stop_event.set() + thread.join(timeout=1.0) + break + continue + + if cmd.action == "shutdown": + logger.info("SQS_HANDLER: Received shutdown command") + # Stop all active upkeeps + for task_id, (stop_event, thread) in active_upkeeps.items(): + logger.info(f"SQS_HANDLER: Stopping upkeep for task {task_id}") + stop_event.set() + thread.join(timeout=2.0) + break + + if cmd.action == "start_upkeep": + if cmd.task_id in active_upkeeps: + logger.warning( + f"SQS_HANDLER: Upkeep already active for task {cmd.task_id}, ignoring" + ) + continue + + stop_event = threading.Event() + thread = threading.Thread( + target=_run_upkeep_loop, + args=( + cmd.task_id, + stop_event, + cmd.receipt_handle, + cmd.visibility_timeout, + cmd.interval_sec, + cmd.queue_name, + cmd.region_name, + cmd.endpoint_url, + ), + daemon=True, + name=f"upkeep-{cmd.task_id[:8]}", + ) + thread.start() + active_upkeeps[cmd.task_id] = (stop_event, thread) + logger.info(f"SQS_HANDLER: Started upkeep for task {cmd.task_id}") + + elif cmd.action == "stop_upkeep": + if cmd.task_id not in active_upkeeps: + logger.warning( + f"SQS_HANDLER: No active upkeep for task {cmd.task_id}, ignoring" + ) + continue + + stop_event, thread = active_upkeeps.pop(cmd.task_id) + stop_event.set() + thread.join(timeout=2.0) + logger.info(f"SQS_HANDLER: Stopped upkeep for task {cmd.task_id}") + + else: + logger.warning(f"SQS_HANDLER: Unknown action: {cmd.action}") + + except Exception as e: # pylint: disable=broad-except + logger.error(f"SQS_HANDLER: Error processing command: {type(e).__name__}: {e}") + + logger.info("SQS_HANDLER: Handler process exiting") + + +class SQSUpkeepHandlerManager: + """ + Manages the lifecycle of an SQS upkeep handler process. + + The handler process manages its own timers for visibility extensions, + completely isolated from the main process's GIL. + + Usage: + manager = SQSUpkeepHandlerManager() + manager.start() + try: + manager.start_upkeep(task_id, ...) # Handler starts its own timer + # ... task runs ... + manager.stop_upkeep(task_id) # Handler stops the timer + finally: + manager.shutdown() + """ + + def __init__(self): + self._command_queue: multiprocessing.Queue | None = None + self._handler_process: multiprocessing.Process | None = None + + def start(self) -> None: + """Start the handler process.""" + if self._command_queue is not None: + return # Already running + + # Get current log level to pass to handler process + current_log_level = logging.getLevelName(logging.getLogger("mazepa").getEffectiveLevel()) + + # Pass parent PID so handler can detect if parent dies + parent_pid = os.getpid() + + self._command_queue = multiprocessing.Queue() + self._handler_process = multiprocessing.Process( + target=run_sqs_upkeep_handler, + args=(self._command_queue, current_log_level, parent_pid), + daemon=True, + name="sqs-upkeep-handler", + ) + self._handler_process.start() + logger.info(f"Started SQS upkeep handler process (PID: {self._handler_process.pid})") + + def shutdown(self, timeout: float = 10.0) -> None: + """Shutdown the handler process gracefully.""" + if self._command_queue is None: + return # Not running + + logger.info("Shutting down SQS upkeep handler process...") + self._command_queue.put(UpkeepCommand(action="shutdown")) + + if self._handler_process is not None: + self._handler_process.join(timeout=timeout) + if self._handler_process.is_alive(): + logger.warning("Handler process did not stop gracefully, terminating...") + self._handler_process.terminate() + self._handler_process.join(timeout=1.0) + + self._command_queue = None + self._handler_process = None + logger.info("SQS upkeep handler process stopped.") + + def start_upkeep( + self, + task_id: str, + receipt_handle: str, + visibility_timeout: int, + interval_sec: float, + queue_name: str, + region_name: str, + endpoint_url: str | None = None, + ) -> None: + """ + Start upkeep for a task. The handler process will manage its own timer + and extend visibility at regular intervals. + """ + if self._command_queue is None: + logger.warning("SQS_HANDLER: Handler not running, start_upkeep ignored") + return + + cmd = UpkeepCommand( + action="start_upkeep", + task_id=task_id, + receipt_handle=receipt_handle, + visibility_timeout=visibility_timeout, + interval_sec=interval_sec, + queue_name=queue_name, + region_name=region_name, + endpoint_url=endpoint_url, + ) + self._command_queue.put_nowait(cmd) + + def stop_upkeep(self, task_id: str) -> None: + """Stop upkeep for a task.""" + if self._command_queue is None: + logger.warning("SQS_HANDLER: Handler not running, stop_upkeep ignored") + return + + self._command_queue.put_nowait(UpkeepCommand(action="stop_upkeep", task_id=task_id)) + + @property + def is_running(self) -> bool: + """Check if the handler process is running.""" + return self._handler_process is not None and self._handler_process.is_alive() diff --git a/zetta_utils/mazepa/worker.py b/zetta_utils/mazepa/worker.py index 81a0b6361..67611a7fd 100644 --- a/zetta_utils/mazepa/worker.py +++ b/zetta_utils/mazepa/worker.py @@ -1,15 +1,15 @@ from __future__ import annotations import math +import multiprocessing +import os import sys import time import traceback from typing import Any, Callable, Optional -import tenacity - -from zetta_utils import builder, log -from zetta_utils.common import RepeatTimer, monitor_resources +from zetta_utils import builder, log, try_load_train_inference +from zetta_utils.common import RepeatTimer, monitor_resources, reset_signal_handlers from zetta_utils.mazepa import constants, exceptions from zetta_utils.mazepa.exceptions import MazepaCancel, MazepaTimeoutError from zetta_utils.mazepa.pool_activity import PoolActivityTracker @@ -18,11 +18,67 @@ MAX_TRANSIENT_RETRIES, TRANSIENT_ERROR_CONDITIONS, ) +from zetta_utils.mazepa.upkeep_handlers import ( + SQSUpkeepHandlerManager, + extract_sqs_metadata, + perform_direct_upkeep, +) from zetta_utils.message_queues.base import MessageQueue, ReceivedMessage from . import Task +class DummyBuffer: + def read(self, data): + pass + + def write(self, data): + pass + + def flush(self): + pass + + +def redirect_buffers() -> None: + sys.stdin = DummyBuffer() # type: ignore + sys.stdout = DummyBuffer() # type: ignore + sys.stderr = DummyBuffer() # type: ignore + + +def worker_init( + log_level: str, + suppress_logs: bool = False, + set_start_method: bool = False, + multiprocessing_start_method: str = "spawn", + load_train_inference: bool = False, +) -> None: + """ + Initialize a worker process with proper logging and signal handling. + + Args: + log_level: Log level string (e.g., "INFO", "DEBUG") + suppress_logs: If True, redirect stdout/stderr to dummy buffers + set_start_method: If True, set multiprocessing start method (for worker pools) + multiprocessing_start_method: The start method to use if set_start_method is True + load_train_inference: If True, try to load train/inference modules (for worker pools) + """ + # Reset signal handlers inherited from parent to default behavior + reset_signal_handlers() + # For Kubernetes compatibility, ensure unbuffered output + os.environ["PYTHONUNBUFFERED"] = "1" + + if suppress_logs: + redirect_buffers() + else: + log.configure_logger(level=log_level, force=True) + + if set_start_method: + multiprocessing.set_start_method(multiprocessing_start_method, force=True) + + if load_train_inference: + try_load_train_inference() + + class AcceptAllTasks: def __call__(self, task: Task): return True @@ -90,6 +146,7 @@ def _process_task_batch( outcome_queue: MessageQueue[OutcomeReport], activity_tracker: PoolActivityTracker | None, debug: bool, + upkeep_handler: SQSUpkeepHandlerManager | None = None, ) -> None: logger.info("STARTING: task batch execution.") @@ -104,15 +161,15 @@ def _process_task_batch( with log.logging_tag_ctx("task_id", task.id_): with log.logging_tag_ctx("execution_id", task.execution_id): if task_filter_fn(task): - ack_task, outcome = process_task_message(msg=msg, debug=debug) + ack_task, outcome = process_task_message( + msg=msg, debug=debug, upkeep_handler=upkeep_handler + ) else: ack_task = True outcome = TaskOutcome(exception=MazepaCancel()) if ack_task: - outcome_report = OutcomeReport( - task_id=msg.payload.id_, outcome=outcome - ) + outcome_report = OutcomeReport(task_id=msg.payload.id_, outcome=outcome) outcome_queue.push([outcome_report]) msg.acknowledge_fn() @@ -154,38 +211,57 @@ def run_worker( idle_timeout: float | None = None, pool_name: str | None = None, ) -> str: - with monitor_resources(resource_monitor_interval): - start_time = time.time() - activity_tracker = PoolActivityTracker(pool_name) if pool_name else None + # Start SQS upkeep handler process for handling visibility extensions + upkeep_handler = SQSUpkeepHandlerManager() + upkeep_handler.start() - while True: - task_msgs = _pull_tasks_with_error_handling( - task_queue, outcome_queue, max_pull_num - ) + try: + with monitor_resources(resource_monitor_interval): + start_time = time.time() + activity_tracker = PoolActivityTracker(pool_name) if pool_name else None - if len(task_msgs) == 0: - _handle_idle_state(sleep_sec, idle_timeout, activity_tracker) - else: - _process_task_batch( - task_msgs, task_filter_fn, outcome_queue, activity_tracker, debug + while True: + task_msgs = _pull_tasks_with_error_handling( + task_queue, outcome_queue, max_pull_num ) - should_exit, reason = _check_exit_conditions( - start_time, max_runtime, idle_timeout, activity_tracker - ) - if should_exit: - assert reason is not None - logger.info(f"Worker exiting: {reason}") - return reason + if len(task_msgs) == 0: + _handle_idle_state(sleep_sec, idle_timeout, activity_tracker) + else: + _process_task_batch( + task_msgs, + task_filter_fn, + outcome_queue, + activity_tracker, + debug, + upkeep_handler, + ) + + should_exit, reason = _check_exit_conditions( + start_time, max_runtime, idle_timeout, activity_tracker + ) + if should_exit: + assert reason is not None + logger.info(f"Worker exiting: {reason}") + return reason + finally: + upkeep_handler.shutdown() def process_task_message( - msg: ReceivedMessage[Task], debug: bool, handle_exceptions: bool = True + msg: ReceivedMessage[Task], + debug: bool, + handle_exceptions: bool = True, + upkeep_handler: SQSUpkeepHandlerManager | None = None, ) -> tuple[bool, TaskOutcome]: task = msg.payload if task.upkeep_settings.perform_upkeep: outcome = _run_task_with_upkeep( - task, msg.extend_lease_fn, debug=debug, handle_exceptions=handle_exceptions + task, + msg.extend_lease_fn, + debug=debug, + handle_exceptions=handle_exceptions, + upkeep_handler=upkeep_handler, ) else: outcome = task(debug=debug, handle_exceptions=handle_exceptions) @@ -209,23 +285,70 @@ def process_task_message( def _run_task_with_upkeep( - task: Task, extend_lease_fn: Callable, debug: bool, handle_exceptions: bool + task: Task, + extend_lease_fn: Callable, + debug: bool, + handle_exceptions: bool, + upkeep_handler: SQSUpkeepHandlerManager | None = None, ) -> TaskOutcome: - def _perform_upkeep_callbacks(): - assert task.upkeep_settings.interval_sec is not None + task_start_time = time.time() + assert task.upkeep_settings.interval_sec is not None + extend_duration = math.ceil(task.upkeep_settings.interval_sec * 10) + + # Try to extract SQS metadata and use process-based handler + sqs_metadata = extract_sqs_metadata(extend_lease_fn) + use_process_handler = sqs_metadata is not None and upkeep_handler is not None + + if use_process_handler: + # Handler process manages its own timer - completely isolated from main process GIL + logger.debug( + f"UPKEEP: Starting upkeep via handler process: " + f"interval={task.upkeep_settings.interval_sec}s, extend_by={extend_duration}s" + ) + assert sqs_metadata is not None + assert upkeep_handler is not None + upkeep_handler.start_upkeep( + task_id=task.id_, + receipt_handle=sqs_metadata["receipt_handle"], + visibility_timeout=extend_duration, + interval_sec=task.upkeep_settings.interval_sec, + queue_name=sqs_metadata["queue_name"], + region_name=sqs_metadata["region_name"], + endpoint_url=sqs_metadata["endpoint_url"], + ) try: - extend_lease_fn(math.ceil(task.upkeep_settings.interval_sec * 5)) - except tenacity.RetryError as e: # pragma: no cover - logger.info(f"Couldn't perform upkeep: {e}") + logger.info("Task execution starting") + result = task(debug=debug, handle_exceptions=handle_exceptions) + elapsed = time.time() - task_start_time + logger.info(f"Task execution completed successfully after {elapsed:.2f}s") + except Exception as e: # pragma: no cover # pylint: disable=broad-except + elapsed = time.time() - task_start_time + logger.error( + f"Task execution failed with {type(e).__name__}: {e} after {elapsed:.2f}s" + ) + raise e + finally: + logger.debug("UPKEEP: Stopping upkeep via handler process") + upkeep_handler.stop_upkeep(task.id_) + else: + # Fallback: use RepeatTimer in main process for non-SQS queues + def upkeep_callback(): + perform_direct_upkeep(extend_lease_fn, extend_duration, task_start_time) - assert task.upkeep_settings.interval_sec is not None - upkeep = RepeatTimer(task.upkeep_settings.interval_sec, _perform_upkeep_callbacks) - upkeep.start() - try: - result = task(debug=debug, handle_exceptions=handle_exceptions) - except Exception as e: # pragma: no cover # pylint: disable=broad-except - raise e - finally: - upkeep.cancel() + upkeep = RepeatTimer(task.upkeep_settings.interval_sec, upkeep_callback) + upkeep.start() + try: + logger.info("Task execution starting") + result = task(debug=debug, handle_exceptions=handle_exceptions) + elapsed = time.time() - task_start_time + logger.info(f"Task execution completed successfully after {elapsed:.2f}s") + except Exception as e: # pragma: no cover # pylint: disable=broad-except + elapsed = time.time() - task_start_time + logger.error( + f"Task execution failed with {type(e).__name__}: {e} after {elapsed:.2f}s" + ) + raise e + finally: + upkeep.cancel() return result diff --git a/zetta_utils/mazepa_addons/configurations/worker_pool.py b/zetta_utils/mazepa_addons/configurations/worker_pool.py index 78eccf7d6..9a4394b34 100644 --- a/zetta_utils/mazepa_addons/configurations/worker_pool.py +++ b/zetta_utils/mazepa_addons/configurations/worker_pool.py @@ -3,59 +3,22 @@ import contextlib import logging import multiprocessing -import os -import sys import time from contextlib import ExitStack import pebble -from zetta_utils import builder, log, try_load_train_inference -from zetta_utils.common import monitor_resources, reset_signal_handlers +from zetta_utils import builder, log +from zetta_utils.common import monitor_resources from zetta_utils.mazepa import SemaphoreType, Task, configure_semaphores, run_worker from zetta_utils.mazepa.pool_activity import PoolActivityTracker from zetta_utils.mazepa.task_outcome import OutcomeReport +from zetta_utils.mazepa.worker import worker_init from zetta_utils.message_queues import FileQueue, SQSQueue logger = log.get_logger("mazepa") -class DummyBuffer: - def read(self, data): - pass - - def write(self, data): - pass - - def flush(self): - pass - - -def redirect_buffers() -> None: # Do not need to implement 14 passes for typing.FileIO - sys.stdin = DummyBuffer() # type: ignore - sys.stdout = DummyBuffer() # type: ignore - sys.stderr = DummyBuffer() # type: ignore - - -def worker_init( - suppress_worker_logs: bool, multiprocessing_start_method: str, log_level: str -) -> None: - # Reset signal handlers inherited from parent to default behavior - # This prevents parent's signal handlers from interfering with worker cleanup - reset_signal_handlers() - # For Kubernetes compatibility, ensure unbuffered output - os.environ["PYTHONUNBUFFERED"] = "1" - if suppress_worker_logs: - redirect_buffers() - else: - # Reconfigure logging in worker process with parent's log level - log.configure_logger(level=log_level, force=True) - - # Inherit the start method from the calling process - multiprocessing.set_start_method(multiprocessing_start_method, force=True) - try_load_train_inference() - - def run_local_worker( task_queue_name: str, outcome_queue_name: str, @@ -115,9 +78,11 @@ def setup_local_worker_pool( ), # 'fork' has issues with CV sharded reads initializer=worker_init, initargs=[ - suppress_worker_logs, - multiprocessing.get_start_method(), - current_log_level, + current_log_level, # log_level + suppress_worker_logs, # suppress_logs + True, # set_start_method + multiprocessing.get_start_method(), # multiprocessing_start_method + True, # load_train_inference ], ) try: diff --git a/zetta_utils/message_queues/sqs/queue.py b/zetta_utils/message_queues/sqs/queue.py index d4aebb84e..f033a6244 100644 --- a/zetta_utils/message_queues/sqs/queue.py +++ b/zetta_utils/message_queues/sqs/queue.py @@ -42,7 +42,7 @@ class SQSQueue(MessageQueue[T]): insertion_threads: int = 5 _queue: Any = attrs.field(init=False, default=None) pull_wait_sec: int = 0 - pull_lease_sec: int = 10 # TODO: get a better value + pull_lease_sec: int = 30 def _get_tq_queue(self) -> Any: if self._queue is None: diff --git a/zetta_utils/message_queues/sqs/utils.py b/zetta_utils/message_queues/sqs/utils.py index d13b3f5ce..5806865a1 100644 --- a/zetta_utils/message_queues/sqs/utils.py +++ b/zetta_utils/message_queues/sqs/utils.py @@ -48,6 +48,10 @@ def receive_msgs( msg_batch_size: int = 10, visibility_timeout: int = 60, ) -> list[SQSReceivedMsg]: + logger.debug( + f"RECEIVE: Attempting to receive messages from queue '{queue_name}' " + f"with visibility_timeout={visibility_timeout}s" + ) result = [] # type: list[SQSReceivedMsg] start_ts = time.time() while True: @@ -73,6 +77,12 @@ def receive_msgs( ) for message in resp["Messages"] ] + for msg in message_batch: + logger.debug( + f"RECEIVE: Message from queue '{queue_name}' with handle " + f"{msg.receipt_handle[:30]}... has initial visibility timeout of " + f"{visibility_timeout}s (expires at relative T+{visibility_timeout}s)" + ) result += message_batch if len(result) >= max_msg_num: @@ -91,14 +101,22 @@ def delete_msg_by_receipt_handle( region_name: str, endpoint_url: Optional[str] = None, ): + queue_url = get_queue_url(queue_name, region_name, endpoint_url=endpoint_url) logger.debug( - f"Deleting message with handle '{receipt_handle}' from queue '{queue_name}'" - f"in region '{region_name}'" - ) - get_sqs_client(region_name, endpoint_url=endpoint_url).delete_message( - QueueUrl=get_queue_url(queue_name, region_name, endpoint_url=endpoint_url), - ReceiptHandle=receipt_handle, + f"DELETE: Deleting message from queue '{queue_name}' (URL: {queue_url}), " + f"handle: {receipt_handle[:30]}..." ) + try: + get_sqs_client(region_name, endpoint_url=endpoint_url).delete_message( + QueueUrl=queue_url, + ReceiptHandle=receipt_handle, + ) + logger.debug(f"DELETE: Successfully deleted message from queue '{queue_name}'") + except Exception as e: + logger.error( + f"DELETE: Failed to delete message from queue '{queue_name}': {type(e).__name__}: {e}" + ) + raise @retry(stop=stop_after_attempt(10), wait=wait_random(min=0.1, max=5)) @@ -109,15 +127,27 @@ def change_message_visibility( region_name: str, endpoint_url: Optional[str] = None, ): + queue_url = get_queue_url(queue_name, region_name, endpoint_url=endpoint_url) logger.debug( - f"Changing visibility of the message with handle '{receipt_handle}' " - f"from queue '{queue_name}' in region '{region_name}' to {visibility_timeout}." - ) - get_sqs_client(region_name, endpoint_url=endpoint_url).change_message_visibility( - QueueUrl=get_queue_url(queue_name, region_name, endpoint_url=endpoint_url), - ReceiptHandle=receipt_handle, - VisibilityTimeout=visibility_timeout, + f"VISIBILITY: Changing visibility to {visibility_timeout}s for queue '{queue_name}' " + f"(URL: {queue_url}), handle: {receipt_handle[:30]}..." ) + try: + get_sqs_client(region_name, endpoint_url=endpoint_url).change_message_visibility( + QueueUrl=queue_url, + ReceiptHandle=receipt_handle, + VisibilityTimeout=visibility_timeout, + ) + logger.debug( + f"VISIBILITY: Successfully changed visibility to {visibility_timeout}s " + f"for queue '{queue_name}'" + ) + except Exception as e: + logger.error( + f"VISIBILITY: Failed to change visibility for queue '{queue_name}' " + f"to {visibility_timeout}s: {type(e).__name__}: {e}" + ) + raise # To be revived if we need batch deletes: From 674c8ea5f3e5c84b43ae8f0970ff3e227c4dc996 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 2 Jan 2026 01:20:01 -0800 Subject: [PATCH 4/9] feat: StackedVolumetricOperations --- pyproject.toml | 5 +- requirements.all.txt | 166 ++--- requirements.modules.txt | 161 ++--- zetta_utils/common/path.py | 2 +- zetta_utils/layer/backend_base.py | 4 + .../layer/db_layer/datastore/backend.py | 3 + .../layer/db_layer/firestore/backend.py | 3 + zetta_utils/layer/layer_base.py | 3 + zetta_utils/layer/layer_set/backend.py | 4 + zetta_utils/layer/protocols.py | 6 + .../layer/volumetric/annotation/backend.py | 4 + zetta_utils/layer/volumetric/backend.py | 6 + .../layer/volumetric/cloudvol/backend.py | 39 +- .../layer/volumetric/cloudvol/build.py | 2 +- .../volumetric/cloudvol/deprecated/backend.py | 13 + .../layer/volumetric/constant/backend.py | 14 + zetta_utils/layer/volumetric/layer.py | 3 + .../layer/volumetric/layer_set/backend.py | 23 + zetta_utils/layer/volumetric/protocols.py | 3 + .../layer/volumetric/seg_contact/backend.py | 3 + .../layer/volumetric/tensorstore/backend.py | 36 +- .../tensorstore/deprecated/backend.py | 13 + .../mazepa_layer_processing/__init__.py | 1 + .../common/__init__.py | 1 + .../common/stacked_volumetric_operations.py | 141 ++++ .../common/subchunkable_apply_flow.py | 22 +- .../common/volumetric_apply_flow.py | 654 +++++++++++------- .../common/volumetric_callable_operation.py | 73 +- .../operation_protocols.py | 45 ++ 29 files changed, 1001 insertions(+), 452 deletions(-) create mode 100644 zetta_utils/mazepa_layer_processing/common/stacked_volumetric_operations.py diff --git a/pyproject.toml b/pyproject.toml index 1d558a6e6..28f87feac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,7 +55,10 @@ augmentations = [ "torchvision", ] cli = ["click >= 8.0.1"] -cloudvol = ["cloud-volume[all_codecs] >= 12.6.2", "zetta_utils[tensor_ops]"] +cloudvol = [ + "cloud-volume[all_codecs] == 12.9.1", + "zetta_utils[tensor_ops]" +] tensorstore = [ "boto3 >= 1.38.15", "tensorstore >= 0.1.73, < 0.1.80", diff --git a/requirements.all.txt b/requirements.all.txt index db588cc2b..e598ac8ba 100644 --- a/requirements.all.txt +++ b/requirements.all.txt @@ -6,7 +6,7 @@ affine==2.4.0 # via zetta-utils (pyproject.toml) aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.13.2 +aiohttp==3.13.3 # via # zetta-utils (pyproject.toml) # fsspec @@ -17,13 +17,13 @@ alabaster==0.7.16 # via sphinx annotated-types==0.7.0 # via pydantic -anyio==4.12.0 +anyio==4.12.1 # via httpx apeye==1.4.1 # via sphinx-toolbox apeye-core==1.1.5 # via apeye -astroid==4.0.2 +astroid==4.0.3 # via pylint asttokens==3.0.1 # via stack-data @@ -38,7 +38,7 @@ attrs==25.4.0 # referencing autodocsumm==0.2.14 # via sphinx-toolbox -awscli==1.43.15 +awscli==1.44.20 # via zetta-utils (pyproject.toml) babel==2.17.0 # via sphinx @@ -46,21 +46,21 @@ beautifulsoup4==4.14.3 # via sphinx-toolbox black==21.9b0 # via zetta-utils (pyproject.toml) -blosc==1.11.3 +blosc==1.11.4 # via # cloud-volume # intern # meshparty blosc2==3.12.2 # via tables -boto3==1.42.9 +boto3==1.42.30 # via # zetta-utils (pyproject.toml) # cloud-files # cloud-volume # moto # task-queue -botocore==1.42.9 +botocore==1.42.30 # via # awscli # boto3 @@ -73,15 +73,14 @@ brotli==1.2.0 # urllib3 cachecontrol==0.14.4 # via sphinx-toolbox -cachetools==6.2.3 +cachetools==6.2.4 # via # zetta-utils (pyproject.toml) # caveclient - # google-auth # middle-auth-client caveclient==8.0.1 # via pcg-skel -certifi==2025.11.12 +certifi==2026.1.4 # via # httpcore # httpx @@ -120,7 +119,7 @@ cloud-files==5.8.2 # zetta-utils (pyproject.toml) # cloud-volume # datastoreflex -cloud-volume==12.8.0 +cloud-volume==12.9.1 # via # zetta-utils (pyproject.toml) # meshparty @@ -144,7 +143,7 @@ contourpy==1.3.3 # via matplotlib coolname==2.2.0 # via zetta-utils (pyproject.toml) -coverage==7.13.0 +coverage==7.13.1 # via pytest-cov crackle-codec==0.36.0 # via cloud-volume @@ -158,7 +157,7 @@ cssutils==2.11.1 # via dict2css cycler==0.12.1 # via matplotlib -cython==3.2.3 +cython==3.2.4 # via # abiss # lsds @@ -177,7 +176,7 @@ dict2css==0.3.0.post1 # via sphinx-toolbox dijkstra3d==1.15.2 # via kimimaro -dill==0.4.0 +dill==0.4.1 # via # zetta-utils (pyproject.toml) # multiprocess @@ -234,7 +233,7 @@ fastremap==1.17.7 # kimimaro # osteoid # pcg-skel -filelock==3.20.0 +filelock==3.20.3 # via # cachecontrol # sphinx-toolbox @@ -259,7 +258,7 @@ frozenlist==1.8.0 # via # aiohttp # aiosignal -fsspec==2025.12.0 +fsspec==2026.1.0 # via # zetta-utils (pyproject.toml) # gcsfs @@ -268,7 +267,7 @@ fsspec==2025.12.0 # torch furl==2.1.4 # via middle-auth-client -gcsfs==2025.12.0 +gcsfs==2026.1.0 # via zetta-utils (pyproject.toml) gevent==25.9.1 # via @@ -277,9 +276,9 @@ gevent==25.9.1 # task-queue gitdb==4.0.12 # via gitpython -gitpython==3.1.45 +gitpython==3.1.46 # via wandb -google-api-core==2.28.1 +google-api-core==2.29.0 # via # google-api-python-client # google-cloud-bigtable @@ -292,11 +291,11 @@ google-api-core==2.28.1 # google-cloud-pubsub # google-cloud-storage # google-cloud-storage-control -google-api-python-client==2.187.0 +google-api-python-client==2.188.0 # via zetta-utils (pyproject.toml) google-apitools==0.5.35 # via neuroglancer -google-auth==2.43.0 +google-auth==2.47.0 # via # zetta-utils (pyproject.toml) # cloud-files @@ -316,20 +315,19 @@ google-auth==2.43.0 # google-cloud-pubsub # google-cloud-storage # google-cloud-storage-control - # kubernetes # neuroglancer # task-queue -google-auth-httplib2==0.2.1 +google-auth-httplib2==0.3.0 # via google-api-python-client -google-auth-oauthlib==1.2.2 +google-auth-oauthlib==1.2.4 # via gcsfs -google-cloud-bigtable==2.34.0 +google-cloud-bigtable==2.35.0 # via zetta-utils (pyproject.toml) -google-cloud-billing==1.17.0 +google-cloud-billing==1.18.0 # via zetta-utils (pyproject.toml) -google-cloud-compute==1.40.0 +google-cloud-compute==1.42.0 # via zetta-utils (pyproject.toml) -google-cloud-container==2.61.0 +google-cloud-container==2.62.0 # via zetta-utils (pyproject.toml) google-cloud-core==2.5.0 # via @@ -340,22 +338,22 @@ google-cloud-core==2.5.0 # google-cloud-firestore # google-cloud-storage # task-queue -google-cloud-datastore==2.21.0 +google-cloud-datastore==2.23.0 # via # zetta-utils (pyproject.toml) # datastoreflex -google-cloud-firestore==2.21.0 +google-cloud-firestore==2.23.0 # via zetta-utils (pyproject.toml) -google-cloud-pubsub==2.33.0 +google-cloud-pubsub==2.34.0 # via messagingclient -google-cloud-storage==3.7.0 +google-cloud-storage==3.8.0 # via # cloud-files # cloud-volume # gcsfs -google-cloud-storage-control==1.8.0 +google-cloud-storage-control==1.9.0 # via gcsfs -google-crc32c==1.7.1 +google-crc32c==1.8.0 # via # cloud-files # crackle-codec @@ -386,6 +384,7 @@ grpcio==1.76.0 # google-cloud-billing # google-cloud-compute # google-cloud-container + # google-cloud-datastore # google-cloud-pubsub # google-cloud-storage-control # googleapis-common-protos @@ -403,7 +402,7 @@ html5lib==1.1 # via sphinx-toolbox httpcore==1.0.9 # via httpx -httplib2==0.31.0 +httplib2==0.31.1 # via # google-api-python-client # google-apitools @@ -411,7 +410,7 @@ httplib2==0.31.0 # oauth2client httpx==0.28.1 # via trimesh -identify==2.6.15 +identify==2.6.16 # via pre-commit idna==3.11 # via @@ -421,7 +420,7 @@ idna==3.11 # requests # sphinx-prompt # yarl -imagecodecs==2025.11.11 +imagecodecs==2026.1.14 # via cloud-volume imagecorruptions-imaug==1.1.3 # via @@ -435,7 +434,7 @@ imagesize==1.4.1 # via sphinx imaug==0.4.2 # via zetta-utils (pyproject.toml) -importlib-metadata==8.7.0 +importlib-metadata==8.7.1 # via opentelemetry-api inflection==0.5.1 # via python-jsonschema-objects @@ -443,9 +442,9 @@ iniconfig==2.3.0 # via pytest intern==1.4.2 # via cloud-volume -intervaltree==3.1.0 +intervaltree==3.2.1 # via cloud-files -ipython==9.8.0 +ipython==9.9.0 # via # caveclient # ipywidgets @@ -473,13 +472,13 @@ jmespath==1.0.1 # via # boto3 # botocore -joblib==1.5.2 +joblib==1.5.3 # via # intern # scikit-learn -json5==0.12.1 +json5==0.13.0 # via cloud-volume -jsonschema==4.25.1 +jsonschema==4.26.0 # via # caveclient # cloud-volume @@ -489,7 +488,7 @@ jsonschema-specifications==2025.9.1 # via jsonschema jupyterlab-widgets==3.0.16 # via ipywidgets -kimimaro==5.8.0 +kimimaro==5.8.1 # via zetta-utils (pyproject.toml) kiwisolver==1.4.9 # via matplotlib @@ -497,7 +496,7 @@ kornia==0.8.2 # via zetta-utils (pyproject.toml) kornia-rs==0.1.10 # via kornia -kubernetes==34.1.0 +kubernetes==35.0.0 # via zetta-utils (pyproject.toml) lazy-loader==0.4 # via scikit-image @@ -508,7 +507,7 @@ lightning-utilities==0.15.2 # lightning # pytorch-lightning # torchmetrics -llvmlite==0.46.0 +llvmlite==0.36.0 # via numba lsds @ git+https://github.com/ZettaAI/lsd.git@cebe976133efb991ebc40171eeed7a22c57ba50a # via zetta-utils (pyproject.toml) @@ -547,7 +546,7 @@ messagingclient==0.3.0 # via zetta-utils (pyproject.toml) microviewer==1.20.0 # via cloud-volume -middle-auth-client==3.19.0 +middle-auth-client==3.19.2 # via zetta-utils (pyproject.toml) ml-dtypes==0.5.4 # via @@ -555,7 +554,7 @@ ml-dtypes==0.5.4 # tensorstore more-itertools==10.8.0 # via cssutils -moto==5.1.18 +moto==5.1.20 # via zetta-utils (pyproject.toml) mpmath==1.3.0 # via sympy @@ -567,7 +566,7 @@ multidict==6.7.0 # via # aiohttp # yarl -multiprocess==0.70.18 +multiprocess==0.70.19 # via pathos multiwrapper==0.1.1 # via @@ -597,17 +596,17 @@ networkx==3.6.1 # trimesh neuroglancer==2.41.2 # via zetta-utils (pyproject.toml) -nodeenv==1.9.1 +nodeenv==1.10.0 # via pre-commit nose2==0.15.1 # via intern -numba==0.63.1 +numba==0.53.1 # via imagecorruptions-imaug numexpr==2.14.1 # via # blosc2 # tables -numpy==2.3.5 +numpy==2.4.1 # via # zetta-utils (pyproject.toml) # abiss @@ -719,17 +718,17 @@ oauth2client==4.1.3 # via google-apitools oauthlib==3.3.1 # via requests-oauthlib -onnx==1.20.0 +onnx==1.20.1 # via # zetta-utils (pyproject.toml) # onnx2torch onnx2torch==1.5.15 # via zetta-utils (pyproject.toml) -opencv-python==4.11.0.86 +opencv-python==4.13.0.90 # via # imagecorruptions-imaug # imaug -opencv-python-headless==4.11.0.86 +opencv-python-headless==4.13.0.90 # via # zetta-utils (pyproject.toml) # imaug @@ -778,7 +777,7 @@ pandas==2.3.3 # pcg-skel parso==0.8.5 # via jedi -pathos==0.3.4 +pathos==0.3.5 # via # cloud-files # cloud-volume @@ -790,7 +789,7 @@ pbr==7.0.3 # via task-queue pcg-skel==1.3.1 # via zetta-utils (pyproject.toml) -pdbp==1.8.1 +pdbp==1.8.2 # via zetta-utils (pyproject.toml) pebble==5.1.3 # via zetta-utils (pyproject.toml) @@ -798,7 +797,7 @@ pexpect==4.9.0 # via ipython piccolo-theme==0.24.0 # via zetta-utils (pyproject.toml) -pillow==12.0.0 +pillow==12.1.0 # via # imagecorruptions-imaug # imageio @@ -823,9 +822,9 @@ pluggy==1.6.0 # pytest-cov posix-ipc==1.3.2 # via cloud-volume -pox==0.3.6 +pox==0.3.7 # via pathos -ppft==1.7.7 +ppft==1.7.8 # via pathos pre-commit==2.19.0 # via zetta-utils (pyproject.toml) @@ -835,7 +834,7 @@ propcache==0.4.1 # via # aiohttp # yarl -proto-plus==1.26.1 +proto-plus==1.27.0 # via # google-api-core # google-cloud-bigtable @@ -846,7 +845,7 @@ proto-plus==1.26.1 # google-cloud-firestore # google-cloud-pubsub # google-cloud-storage-control -protobuf==6.33.2 +protobuf==6.33.4 # via # zetta-utils (pyproject.toml) # cloud-files @@ -866,7 +865,7 @@ protobuf==6.33.2 # onnx # proto-plus # wandb -psutil==7.1.3 +psutil==7.2.1 # via # zetta-utils (pyproject.toml) # cloud-volume @@ -881,9 +880,9 @@ py-cpuinfo==9.0.0 # via # blosc2 # tables -pyarrow==22.0.0 +pyarrow==23.0.0 # via caveclient -pyasn1==0.6.1 +pyasn1==0.6.2 # via # oauth2client # pyasn1-modules @@ -920,7 +919,7 @@ pygments==2.19.2 # sphinx-tabs pylint==4.0.4 # via zetta-utils (pyproject.toml) -pyparsing==3.2.5 +pyparsing==3.3.1 # via # httplib2 # matplotlib @@ -978,7 +977,7 @@ referencing==0.37.0 # via # jsonschema # jsonschema-specifications -regex==2025.11.3 +regex==2026.1.15 # via black requests==2.32.5 # via @@ -1040,7 +1039,7 @@ s3transfer==0.16.0 # via # awscli # boto3 -scikit-image==0.25.2 +scikit-image==0.26.0 # via # zetta-utils (pyproject.toml) # imagecorruptions-imaug @@ -1049,7 +1048,7 @@ scikit-learn==1.8.0 # via # zetta-utils (pyproject.toml) # meshparty -scipy==1.16.3 +scipy==1.17.0 # via # imagecorruptions-imaug # imaug @@ -1060,12 +1059,13 @@ scipy==1.16.3 # scikit-image # scikit-learn # trimesh -sentry-sdk==2.47.0 +sentry-sdk==2.49.0 # via wandb setuptools==80.9.0 # via # imagecorruptions-imaug # lightning-utilities + # numba # pbr # torch # torchfields @@ -1101,7 +1101,7 @@ snowballstemmer==3.0.1 # sphinx sortedcontainers==2.4.0 # via intervaltree -soupsieve==2.8 +soupsieve==2.8.2 # via beautifulsoup4 sphinx==6.2.1 # via @@ -1132,7 +1132,7 @@ sphinx-tabs==3.4.5 # via # zetta-utils (pyproject.toml) # sphinx-toolbox -sphinx-toolbox==4.1.0 +sphinx-toolbox==4.1.2 # via zetta-utils (pyproject.toml) sphinxcontrib-applehelp==2.0.0 # via @@ -1184,11 +1184,11 @@ tenacity==9.1.2 # task-queue tensorstore==0.1.79 # via zetta-utils (pyproject.toml) -testcontainers==4.13.3 +testcontainers==4.14.0 # via zetta-utils (pyproject.toml) threadpoolctl==3.6.0 # via scikit-learn -tifffile==2025.12.12 +tifffile==2026.1.14 # via scikit-image tinybrain==1.7.0 # via zetta-utils (pyproject.toml) @@ -1196,7 +1196,7 @@ toml==0.10.2 # via pre-commit tomli==1.2.3 # via black -tomlkit==0.13.3 +tomlkit==0.14.0 # via pylint torch==2.9.1 # via @@ -1219,7 +1219,7 @@ torchvision==0.24.1 # via # zetta-utils (pyproject.toml) # onnx2torch -tornado==6.5.3 +tornado==6.5.4 # via neuroglancer tqdm==4.67.1 # via @@ -1237,7 +1237,7 @@ traitlets==5.14.3 # ipython # ipywidgets # matplotlib-inline -trimesh==4.10.1 +trimesh==4.11.1 # via # zetta-utils (pyproject.toml) # meshparty @@ -1254,7 +1254,7 @@ types-pillow==10.2.0.20240822 # via zetta-utils (pyproject.toml) types-pyyaml==6.0.12.20250915 # via zetta-utils (pyproject.toml) -types-requests==2.32.4.20250913 +types-requests==2.32.4.20260107 # via zetta-utils (pyproject.toml) types-tabulate==0.9.0.20241207 # via zetta-utils (pyproject.toml) @@ -1291,7 +1291,7 @@ tzdata==2025.3 # via pandas uritemplate==4.2.0 # via google-api-python-client -urllib3==2.3.0 +urllib3==2.6.3 # via # botocore # caveclient @@ -1307,9 +1307,9 @@ urllib3==2.3.0 # types-requests vhacdx==0.0.10 # via trimesh -virtualenv==20.35.4 +virtualenv==20.36.1 # via pre-commit -wandb==0.23.1 +wandb==0.24.0 # via zetta-utils (pyproject.toml) waterz @ git+https://github.com/ZettaAI/waterz.git@0bac4be4c864b293fa49f11cf0abf52bd1ec81df # via zetta-utils (pyproject.toml) @@ -1319,7 +1319,7 @@ webencodings==0.5.1 # via html5lib websocket-client==1.9.0 # via kubernetes -werkzeug==3.1.4 +werkzeug==3.1.5 # via # zetta-utils (pyproject.toml) # flask @@ -1349,7 +1349,7 @@ zmesh==1.9.0 # via zetta-utils (pyproject.toml) zope-event==6.1 # via gevent -zope-interface==8.1.1 +zope-interface==8.2 # via gevent zstandard==0.25.0 # via diff --git a/requirements.modules.txt b/requirements.modules.txt index bdaaada90..4a5a6d92b 100644 --- a/requirements.modules.txt +++ b/requirements.modules.txt @@ -6,7 +6,7 @@ affine==2.4.0 # via zetta-utils (pyproject.toml) aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.13.2 +aiohttp==3.13.3 # via # zetta-utils (pyproject.toml) # fsspec @@ -15,7 +15,7 @@ aiosignal==1.4.0 # via aiohttp annotated-types==0.7.0 # via pydantic -anyio==4.12.0 +anyio==4.12.1 # via httpx asttokens==3.0.1 # via stack-data @@ -28,24 +28,24 @@ attrs==25.4.0 # caveclient # jsonschema # referencing -awscli==1.43.15 +awscli==1.44.20 # via zetta-utils (pyproject.toml) blinker==1.9.0 # via flask -blosc==1.11.3 +blosc==1.11.4 # via # cloud-volume # intern # meshparty blosc2==3.12.2 # via tables -boto3==1.42.9 +boto3==1.42.30 # via # zetta-utils (pyproject.toml) # cloud-files # cloud-volume # task-queue -botocore==1.42.9 +botocore==1.42.30 # via # awscli # boto3 @@ -55,15 +55,14 @@ brotli==1.2.0 # cloud-files # mapbuffer # urllib3 -cachetools==6.2.3 +cachetools==6.2.4 # via # zetta-utils (pyproject.toml) # caveclient - # google-auth # middle-auth-client caveclient==8.0.1 # via pcg-skel -certifi==2025.11.12 +certifi==2026.1.4 # via # httpcore # httpx @@ -96,15 +95,13 @@ cloud-files==5.8.2 # zetta-utils (pyproject.toml) # cloud-volume # datastoreflex -cloud-volume==12.8.0 +cloud-volume==12.9.1 # via # zetta-utils (pyproject.toml) # meshparty # pcg-skel colorama==0.4.6 - # via - # awscli - # taichi + # via awscli colorlog==6.10.1 # via trimesh comm==0.2.3 @@ -130,7 +127,7 @@ crc32c==2.8 # mapbuffer cycler==0.12.1 # via matplotlib -cython==3.2.3 +cython==3.2.4 # via # abiss # lsds @@ -147,12 +144,11 @@ deflate==0.8.1 # mapbuffer dijkstra3d==1.15.2 # via kimimaro -dill==0.4.0 +dill==0.4.1 # via # zetta-utils (pyproject.toml) # multiprocess # pathos - # taichi docutils==0.19 # via awscli dracopy==1.7.0 @@ -166,9 +162,7 @@ edt==3.1.0 # fastmorph # kimimaro einops==0.8.1 - # via - # zetta-utils (pyproject.toml) - # pointnet + # via zetta-utils (pyproject.toml) embreex==2.17.7.post7 # via trimesh executing==2.2.1 @@ -189,7 +183,7 @@ fastremap==1.17.7 # kimimaro # osteoid # pcg-skel -filelock==3.20.0 +filelock==3.20.3 # via torch fill-voids==2.1.1 # via @@ -210,7 +204,7 @@ frozenlist==1.8.0 # via # aiohttp # aiosignal -fsspec==2025.12.0 +fsspec==2026.1.0 # via # zetta-utils (pyproject.toml) # gcsfs @@ -219,7 +213,7 @@ fsspec==2025.12.0 # torch furl==2.1.4 # via middle-auth-client -gcsfs==2025.12.0 +gcsfs==2026.1.0 # via zetta-utils (pyproject.toml) gevent==25.9.1 # via @@ -228,9 +222,9 @@ gevent==25.9.1 # task-queue gitdb==4.0.12 # via gitpython -gitpython==3.1.45 +gitpython==3.1.46 # via wandb -google-api-core==2.28.1 +google-api-core==2.29.0 # via # google-api-python-client # google-cloud-bigtable @@ -243,11 +237,11 @@ google-api-core==2.28.1 # google-cloud-pubsub # google-cloud-storage # google-cloud-storage-control -google-api-python-client==2.187.0 +google-api-python-client==2.188.0 # via zetta-utils (pyproject.toml) google-apitools==0.5.35 # via neuroglancer -google-auth==2.43.0 +google-auth==2.47.0 # via # zetta-utils (pyproject.toml) # cloud-files @@ -267,20 +261,19 @@ google-auth==2.43.0 # google-cloud-pubsub # google-cloud-storage # google-cloud-storage-control - # kubernetes # neuroglancer # task-queue -google-auth-httplib2==0.2.1 +google-auth-httplib2==0.3.0 # via google-api-python-client -google-auth-oauthlib==1.2.2 +google-auth-oauthlib==1.2.4 # via gcsfs -google-cloud-bigtable==2.34.0 +google-cloud-bigtable==2.35.0 # via zetta-utils (pyproject.toml) -google-cloud-billing==1.17.0 +google-cloud-billing==1.18.0 # via zetta-utils (pyproject.toml) -google-cloud-compute==1.40.0 +google-cloud-compute==1.42.0 # via zetta-utils (pyproject.toml) -google-cloud-container==2.61.0 +google-cloud-container==2.62.0 # via zetta-utils (pyproject.toml) google-cloud-core==2.5.0 # via @@ -291,22 +284,22 @@ google-cloud-core==2.5.0 # google-cloud-firestore # google-cloud-storage # task-queue -google-cloud-datastore==2.21.0 +google-cloud-datastore==2.23.0 # via # zetta-utils (pyproject.toml) # datastoreflex -google-cloud-firestore==2.21.0 +google-cloud-firestore==2.23.0 # via zetta-utils (pyproject.toml) -google-cloud-pubsub==2.33.0 +google-cloud-pubsub==2.34.0 # via messagingclient -google-cloud-storage==3.7.0 +google-cloud-storage==3.8.0 # via # cloud-files # cloud-volume # gcsfs -google-cloud-storage-control==1.8.0 +google-cloud-storage-control==1.9.0 # via gcsfs -google-crc32c==1.7.1 +google-crc32c==1.8.0 # via # cloud-files # crackle-codec @@ -337,6 +330,7 @@ grpcio==1.76.0 # google-cloud-billing # google-cloud-compute # google-cloud-container + # google-cloud-datastore # google-cloud-pubsub # google-cloud-storage-control # googleapis-common-protos @@ -352,7 +346,7 @@ h5py==3.15.1 # via meshparty httpcore==1.0.9 # via httpx -httplib2==0.31.0 +httplib2==0.31.1 # via # google-api-python-client # google-apitools @@ -366,7 +360,7 @@ idna==3.11 # httpx # requests # yarl -imagecodecs==2025.11.11 +imagecodecs==2026.1.14 # via cloud-volume imagecorruptions-imaug==1.1.3 # via @@ -378,7 +372,7 @@ imageio==2.37.2 # scikit-image imaug==0.4.2 # via zetta-utils (pyproject.toml) -importlib-metadata==8.7.0 +importlib-metadata==8.7.1 # via opentelemetry-api inflection==0.5.1 # via python-jsonschema-objects @@ -386,9 +380,9 @@ iniconfig==2.3.0 # via pytest intern==1.4.2 # via cloud-volume -intervaltree==3.1.0 +intervaltree==3.2.1 # via cloud-files -ipython==9.8.0 +ipython==9.9.0 # via # caveclient # ipywidgets @@ -408,13 +402,13 @@ jmespath==1.0.1 # via # boto3 # botocore -joblib==1.5.2 +joblib==1.5.3 # via # intern # scikit-learn -json5==0.12.1 +json5==0.13.0 # via cloud-volume -jsonschema==4.25.1 +jsonschema==4.26.0 # via # caveclient # cloud-volume @@ -424,7 +418,7 @@ jsonschema-specifications==2025.9.1 # via jsonschema jupyterlab-widgets==3.0.16 # via ipywidgets -kimimaro==5.8.0 +kimimaro==5.8.1 # via zetta-utils (pyproject.toml) kiwisolver==1.4.9 # via matplotlib @@ -432,7 +426,7 @@ kornia==0.8.2 # via zetta-utils (pyproject.toml) kornia-rs==0.1.10 # via kornia -kubernetes==34.1.0 +kubernetes==35.0.0 # via zetta-utils (pyproject.toml) lazy-loader==0.4 # via scikit-image @@ -443,7 +437,7 @@ lightning-utilities==0.15.2 # lightning # pytorch-lightning # torchmetrics -llvmlite==0.46.0 +llvmlite==0.36.0 # via numba lsds @ git+https://github.com/ZettaAI/lsd.git@cebe976133efb991ebc40171eeed7a22c57ba50a # via zetta-utils (pyproject.toml) @@ -480,7 +474,7 @@ messagingclient==0.3.0 # via zetta-utils (pyproject.toml) microviewer==1.20.0 # via cloud-volume -middle-auth-client==3.19.0 +middle-auth-client==3.19.2 # via zetta-utils (pyproject.toml) ml-dtypes==0.5.4 # via @@ -494,7 +488,7 @@ multidict==6.7.0 # via # aiohttp # yarl -multiprocess==0.70.18 +multiprocess==0.70.19 # via pathos multiwrapper==0.1.1 # via @@ -517,13 +511,13 @@ neuroglancer==2.41.2 # via zetta-utils (pyproject.toml) nose2==0.15.1 # via intern -numba==0.63.1 +numba==0.53.1 # via imagecorruptions-imaug numexpr==2.14.1 # via # blosc2 # tables -numpy==2.2.6 +numpy==2.4.1 # via # zetta-utils (pyproject.toml) # abiss @@ -579,7 +573,6 @@ numpy==2.2.6 # shapely # simplejpeg # tables - # taichi # task-queue # tensorstore # tifffile @@ -636,17 +629,17 @@ oauth2client==4.1.3 # via google-apitools oauthlib==3.3.1 # via requests-oauthlib -onnx==1.20.0 +onnx==1.20.1 # via # zetta-utils (pyproject.toml) # onnx2torch onnx2torch==1.5.15 # via zetta-utils (pyproject.toml) -opencv-python==4.12.0.88 +opencv-python==4.13.0.90 # via # imagecorruptions-imaug # imaug -opencv-python-headless==4.12.0.88 +opencv-python-headless==4.13.0.90 # via # zetta-utils (pyproject.toml) # imaug @@ -694,7 +687,7 @@ pandas==2.3.3 # pcg-skel parso==0.8.5 # via jedi -pathos==0.3.4 +pathos==0.3.5 # via # cloud-files # cloud-volume @@ -704,13 +697,13 @@ pbr==7.0.3 # via task-queue pcg-skel==1.3.1 # via zetta-utils (pyproject.toml) -pdbp==1.8.1 +pdbp==1.8.2 # via zetta-utils (pyproject.toml) pebble==5.1.3 # via zetta-utils (pyproject.toml) pexpect==4.9.0 # via ipython -pillow==12.0.0 +pillow==12.1.0 # via # imagecorruptions-imaug # imageio @@ -727,13 +720,11 @@ platformdirs==4.5.1 # wandb pluggy==1.6.0 # via pytest -pointnet==0.1.1 - # via zetta-utils (pyproject.toml) posix-ipc==1.3.2 # via cloud-volume -pox==0.3.6 +pox==0.3.7 # via pathos -ppft==1.7.7 +ppft==1.7.8 # via pathos prompt-toolkit==3.0.52 # via ipython @@ -741,7 +732,7 @@ propcache==0.4.1 # via # aiohttp # yarl -proto-plus==1.26.1 +proto-plus==1.27.0 # via # google-api-core # google-cloud-bigtable @@ -752,7 +743,7 @@ proto-plus==1.26.1 # google-cloud-firestore # google-cloud-pubsub # google-cloud-storage-control -protobuf==6.33.2 +protobuf==6.33.4 # via # zetta-utils (pyproject.toml) # cloud-files @@ -772,7 +763,7 @@ protobuf==6.33.2 # onnx # proto-plus # wandb -psutil==7.1.3 +psutil==7.2.1 # via # zetta-utils (pyproject.toml) # cloud-volume @@ -787,9 +778,9 @@ py-cpuinfo==9.0.0 # via # blosc2 # tables -pyarrow==22.0.0 +pyarrow==23.0.0 # via caveclient -pyasn1==0.6.1 +pyasn1==0.6.2 # via # oauth2client # pyasn1-modules @@ -817,7 +808,7 @@ pygments==2.19.2 # pdbp # pytest # rich -pyparsing==3.2.5 +pyparsing==3.3.1 # via # httplib2 # matplotlib @@ -889,9 +880,7 @@ requests-oauthlib==1.3.1 rfc3339==6.2 # via python-logging-loki rich==14.2.0 - # via - # zetta-utils (pyproject.toml) - # taichi + # via zetta-utils (pyproject.toml) rpds-py==0.30.0 # via # jsonschema @@ -910,7 +899,7 @@ s3transfer==0.16.0 # via # awscli # boto3 -scikit-image==0.25.2 +scikit-image==0.26.0 # via # zetta-utils (pyproject.toml) # imagecorruptions-imaug @@ -919,7 +908,7 @@ scikit-learn==1.8.0 # via # zetta-utils (pyproject.toml) # meshparty -scipy==1.16.3 +scipy==1.17.0 # via # imagecorruptions-imaug # imaug @@ -930,12 +919,13 @@ scipy==1.16.3 # scikit-image # scikit-learn # trimesh -sentry-sdk==2.47.0 +sentry-sdk==2.49.0 # via wandb setuptools==80.9.0 # via # imagecorruptions-imaug # lightning-utilities + # numba # pbr # torch # torchfields @@ -982,8 +972,6 @@ tables==3.10.2 # via meshparty tabulate==0.9.0 # via zetta-utils (pyproject.toml) -taichi==1.7.4 - # via pointnet task-queue==2.14.3 # via zetta-utils (pyproject.toml) tenacity==9.1.2 @@ -996,7 +984,7 @@ tensorstore==0.1.79 # via zetta-utils (pyproject.toml) threadpoolctl==3.6.0 # via scikit-learn -tifffile==2025.12.12 +tifffile==2026.1.14 # via scikit-image tinybrain==1.7.0 # via zetta-utils (pyproject.toml) @@ -1006,7 +994,6 @@ torch==2.9.1 # kornia # lightning # onnx2torch - # pointnet # pytorch-lightning # torchfields # torchmetrics @@ -1022,7 +1009,7 @@ torchvision==0.24.1 # via # zetta-utils (pyproject.toml) # onnx2torch -tornado==6.5.3 +tornado==6.5.4 # via neuroglancer tqdm==4.67.1 # via @@ -1040,7 +1027,7 @@ traitlets==5.14.3 # ipython # ipywidgets # matplotlib-inline -trimesh==4.10.1 +trimesh==4.11.1 # via # zetta-utils (pyproject.toml) # meshparty @@ -1080,7 +1067,7 @@ tzdata==2025.3 # via pandas uritemplate==4.2.0 # via google-api-python-client -urllib3==2.3.0 +urllib3==2.6.3 # via # botocore # caveclient @@ -1091,7 +1078,7 @@ urllib3==2.3.0 # sentry-sdk vhacdx==0.0.10 # via trimesh -wandb==0.23.1 +wandb==0.24.0 # via zetta-utils (pyproject.toml) waterz @ git+https://github.com/ZettaAI/waterz.git@0bac4be4c864b293fa49f11cf0abf52bd1ec81df # via zetta-utils (pyproject.toml) @@ -1099,7 +1086,7 @@ wcwidth==0.2.14 # via prompt-toolkit websocket-client==1.9.0 # via kubernetes -werkzeug==3.1.4 +werkzeug==3.1.5 # via # zetta-utils (pyproject.toml) # flask @@ -1124,7 +1111,7 @@ zmesh==1.9.0 # via zetta-utils (pyproject.toml) zope-event==6.1 # via gevent -zope-interface==8.1.1 +zope-interface==8.2 # via gevent zstandard==0.25.0 # via diff --git a/zetta_utils/common/path.py b/zetta_utils/common/path.py index 4afe29215..9119fb992 100644 --- a/zetta_utils/common/path.py +++ b/zetta_utils/common/path.py @@ -21,5 +21,5 @@ def strip_prefix(path: str) -> str: # pragma: no cover def is_local(path: str) -> bool: # pragma: no cover - local_prefixes = ["file://", "fq://"] + local_prefixes = ["file://", "fq://", "mem://"] return any(abspath(path).startswith(local_prefix) for local_prefix in local_prefixes) diff --git a/zetta_utils/layer/backend_base.py b/zetta_utils/layer/backend_base.py index 1421fe700..893b29ce5 100644 --- a/zetta_utils/layer/backend_base.py +++ b/zetta_utils/layer/backend_base.py @@ -33,3 +33,7 @@ def with_changes(self, **kwargs) -> Backend[IndexT, DataT, DataWriteT]: # pragm `P.kwargs`.""" # return attrs.evolve(self, **kwargs) # has to be implemented by the # child class, as it's not necesserily an `attrs` class + + @abstractmethod + def delete(self): + """Deletes all data associated with this backend""" diff --git a/zetta_utils/layer/db_layer/datastore/backend.py b/zetta_utils/layer/db_layer/datastore/backend.py index 835a21b38..4650194af 100644 --- a/zetta_utils/layer/db_layer/datastore/backend.py +++ b/zetta_utils/layer/db_layer/datastore/backend.py @@ -313,6 +313,9 @@ def with_changes(self, **kwargs) -> DatastoreBackend: deepcopy(self), namespace=kwargs["namespace"], project=kwargs.get("project") ) + def delete(self): # pragma: no cover + raise NotImplementedError("delete() not implemented for DatastoreBackend") + def _get_data_from_entities(row_keys: list[str], entities: list[Entity]) -> DBDataT: row_entities = defaultdict(list) diff --git a/zetta_utils/layer/db_layer/firestore/backend.py b/zetta_utils/layer/db_layer/firestore/backend.py index e2bff3784..fe01244ce 100644 --- a/zetta_utils/layer/db_layer/firestore/backend.py +++ b/zetta_utils/layer/db_layer/firestore/backend.py @@ -263,3 +263,6 @@ def with_changes(self, **kwargs) -> FirestoreBackend: database=kwargs.get("database"), project=kwargs.get("project"), ) + + def delete(self): # pragma: no cover + raise NotImplementedError("delete() not implemented for FirestoreBackend") diff --git a/zetta_utils/layer/layer_base.py b/zetta_utils/layer/layer_base.py index 7f169189b..6c0ec2b38 100644 --- a/zetta_utils/layer/layer_base.py +++ b/zetta_utils/layer/layer_base.py @@ -127,3 +127,6 @@ def with_procs( proc_mods["write_procs"] = tuple(write_procs) return attrs.evolve(self, **proc_mods) + + def delete(self): + self.backend.delete() diff --git a/zetta_utils/layer/layer_set/backend.py b/zetta_utils/layer/layer_set/backend.py index 82fd6b60b..b597bec35 100644 --- a/zetta_utils/layer/layer_set/backend.py +++ b/zetta_utils/layer/layer_set/backend.py @@ -31,3 +31,7 @@ def name(self) -> str: def with_changes(self, **kwargs) -> LayerSetBackend[IndexT, DataT, DataWriteT]: return attrs.evolve(self, **kwargs) # pragma: no cover + + def delete(self): # pragma: no cover + for layer in self.layers.values(): + layer.delete() diff --git a/zetta_utils/layer/protocols.py b/zetta_utils/layer/protocols.py index 720a7b062..861cd728e 100644 --- a/zetta_utils/layer/protocols.py +++ b/zetta_utils/layer/protocols.py @@ -20,6 +20,9 @@ def __getitem__(self, idx: IndexT_contra) -> Any: def __setitem__(self, idx: IndexT_contra, data: Any): ... + def delete(self): + ... + @runtime_checkable class LayerWithIndexDataT(Protocol[IndexT_contra, DataT]): @@ -34,3 +37,6 @@ def __getitem__(self, idx: IndexT_contra) -> DataT: def __setitem__(self, idx: IndexT_contra, data: DataT): ... + + def delete(self): + ... diff --git a/zetta_utils/layer/volumetric/annotation/backend.py b/zetta_utils/layer/volumetric/annotation/backend.py index 2d7652df6..5fad4f178 100644 --- a/zetta_utils/layer/volumetric/annotation/backend.py +++ b/zetta_utils/layer/volumetric/annotation/backend.py @@ -937,6 +937,10 @@ def allow_cache(self) -> bool: def enforce_chunk_aligned_writes(self) -> bool: return False # pragma: no cover + @property + def overwrite_partial_chunks(self) -> bool: + return False # pragma: no cover + @property def use_compression(self) -> bool: return False # pragma: no cover diff --git a/zetta_utils/layer/volumetric/backend.py b/zetta_utils/layer/volumetric/backend.py index 7a3bec460..852e12095 100644 --- a/zetta_utils/layer/volumetric/backend.py +++ b/zetta_utils/layer/volumetric/backend.py @@ -45,6 +45,11 @@ def allow_cache(self) -> bool: def enforce_chunk_aligned_writes(self) -> bool: ... + @property + @abstractmethod + def overwrite_partial_chunks(self) -> bool: + ... + @property @abstractmethod def use_compression(self) -> bool: @@ -77,6 +82,7 @@ def get_bounds(self, resolution: Vec3D) -> VolumetricIndex: "allow_cache" = value: Union[bool, str] "use_compression" = value: str "enforce_chunk_aligned_writes" = value: bool + "overwrite_partial_chunks" = value: bool "voxel_offset_res" = (voxel_offset, resolution): Tuple[Vec3D[int], Vec3D] "chunk_size_res" = (chunk_size, resolution): Tuple[Vec3D[int], Vec3D] "dataset_size_res" = (dataset_size, resolution): Tuple[Vec3D[int], Vec3D] diff --git a/zetta_utils/layer/volumetric/cloudvol/backend.py b/zetta_utils/layer/volumetric/cloudvol/backend.py index f37af0be1..eda05a046 100644 --- a/zetta_utils/layer/volumetric/cloudvol/backend.py +++ b/zetta_utils/layer/volumetric/cloudvol/backend.py @@ -1,12 +1,14 @@ # pylint: disable=missing-docstring from __future__ import annotations +import gc import json from copy import deepcopy from typing import Any, Dict, Optional, Union import attrs import cachetools +import cloudfiles import cloudvolume as cv import numpy as np from cloudvolume import CloudVolume @@ -22,7 +24,7 @@ _cv_cache: cachetools.LRUCache = cachetools.LRUCache(maxsize=2048) _cv_cached: Dict[str, set] = {} -IN_MEM_CACHE_NUM_BYTES_PER_CV = 128 * 1024**2 +IN_MEM_CACHE_NUM_BYTES_PER_CV = 128 * 1024 ** 2 def _serialize_kwargs(kwargs: Dict[str, Any]) -> str: @@ -64,7 +66,7 @@ def _get_cv_cached( info=info, provenance={}, mip=tuple(resolution), - lru_bytes=cache_bytes_limit, + lru_bytes=0 if path_.startswith("mem://") else cache_bytes_limit, **kwargs, ) except ScaleUnavailableError as e: @@ -155,6 +157,7 @@ def _set_cv_defaults(self): self.cv_kwargs.setdefault("delete_black_uploads", True) self.cv_kwargs.setdefault("agglomerate", True) self.cv_kwargs.setdefault("lru_encoding", "raw") + self.cv_kwargs.setdefault("overwrite_partial_chunks", False) @property def name(self) -> str: # pragma: no cover @@ -219,6 +222,17 @@ def use_compression(self, value: bool) -> None: # pragma: no cover " use `backend.with_changes(use_compression=value:bool)` instead." ) + @property + def overwrite_partial_chunks(self) -> bool: # pragma: no cover + return self.cv_kwargs.get("overwrite_partial_chunks", False) + + @overwrite_partial_chunks.setter + def overwrite_partial_chunks(self, value: bool) -> None: # pragma: no cover + raise NotImplementedError( + "cannot set `overwrite_partial_chunks` for CVBackend directly;" + " use `backend.with_changes(overwrite_partial_chunks=value:bool)` instead." + ) + def clear_disk_cache(self) -> None: # pragma: no cover info = get_info(self.path) for scale in info["scales"]: @@ -231,6 +245,12 @@ def clear_disk_cache(self) -> None: # pragma: no cover ).cache.flush() def clear_cache(self) -> None: # pragma: no cover + path_ = abspath(self.path) + resolution_kwargs_pairs = _cv_cached.get(path_, set()) + for resolution, kwargs_key in resolution_kwargs_pairs: + cache_key = (path_, resolution, kwargs_key) + if cache_key in _cv_cache: + _cv_cache[cache_key].image.lru.clear() _clear_cv_cache(self.path) def read(self, idx: VolumetricIndex) -> npt.NDArray: @@ -270,6 +290,7 @@ def write(self, idx: VolumetricIndex, data: npt.NDArray): if data_final.min() < np.int64(0): raise ValueError("Attempting to write negative values to a uint64 CloudVolume") data_final = data_final.astype(np.uint64) + cvol[slices] = data_final cvol.autocrop = False @@ -278,6 +299,7 @@ def with_changes(self, **kwargs) -> CVBackend: "name" = value: str "use_compression" = value: str "enforce_chunk_aligned_writes" = value: bool + "overwrite_partial_chunks" = value: bool "voxel_offset_res" = (voxel_offset, resolution): Tuple[Vec3D[int], Vec3D] "chunk_size_res" = (chunk_size, resolution): Tuple[Vec3D[int], Vec3D] "dataset_size_res" = (dataset_size, resolution): Tuple[Vec3D[int], Vec3D] @@ -293,6 +315,7 @@ def with_changes(self, **kwargs) -> CVBackend: "allow_cache", "use_compression", "enforce_chunk_aligned_writes", + "overwrite_partial_chunks", "voxel_offset_res", "chunk_size_res", "dataset_size_res", @@ -306,6 +329,7 @@ def with_changes(self, **kwargs) -> CVBackend: keys_to_cv_kwargs = { "use_compression": "compress", "enforce_chunk_aligned_writes": "non_aligned_writes", + "overwrite_partial_chunks": "overwrite_partial_chunks", } keys_to_reverse = ["enforce_chunk_aligned_writes"] evolve_kwargs = {} @@ -369,3 +393,14 @@ def get_bounds(self, resolution: Vec3D) -> VolumetricIndex: # pragma: no cover def pformat(self) -> str: # pragma: no cover return self.name + + def delete(self): + self.clear_cache() + gc.collect() + + path = abspath(self.path) + + cf = cloudfiles.CloudFiles(path) + files_to_delete = list(cf.list()) + if files_to_delete: + cf.delete(files_to_delete) diff --git a/zetta_utils/layer/volumetric/cloudvol/build.py b/zetta_utils/layer/volumetric/cloudvol/build.py index 8dd2e0e4e..d69622032 100644 --- a/zetta_utils/layer/volumetric/cloudvol/build.py +++ b/zetta_utils/layer/volumetric/cloudvol/build.py @@ -18,7 +18,7 @@ # from typeguard import typechecked -# @typechecked # ypeError: isinstance() arg 2 must be a type or tuple of types on p3.9 +# @typechecked # TypeError: isinstance() arg 2 must be a type or tuple of types on p3.9 @builder.register("build_cv_layer", versions=">=0.4") def build_cv_layer( # pylint: disable=too-many-locals path: str, diff --git a/zetta_utils/layer/volumetric/cloudvol/deprecated/backend.py b/zetta_utils/layer/volumetric/cloudvol/deprecated/backend.py index d3fed7e6d..7a7985de2 100644 --- a/zetta_utils/layer/volumetric/cloudvol/deprecated/backend.py +++ b/zetta_utils/layer/volumetric/cloudvol/deprecated/backend.py @@ -349,3 +349,16 @@ def get_bounds(self, resolution: Vec3D) -> VolumetricIndex: # pragma: no cover def pformat(self) -> str: # pragma: no cover return self.name + + @property + def overwrite_partial_chunks(self) -> bool: # pragma: no cover + return False + + @overwrite_partial_chunks.setter + def overwrite_partial_chunks(self, value: bool) -> None: # pragma: no cover + raise NotImplementedError( + "cannot set `overwrite_partial_chunks` for deprecated CVBackend directly" + ) + + def delete(self): # pragma: no cover + raise NotImplementedError("delete() not implemented for deprecated CVBackend") diff --git a/zetta_utils/layer/volumetric/constant/backend.py b/zetta_utils/layer/volumetric/constant/backend.py index 45ec5538c..142223ed0 100644 --- a/zetta_utils/layer/volumetric/constant/backend.py +++ b/zetta_utils/layer/volumetric/constant/backend.py @@ -64,6 +64,16 @@ def allow_cache(self, value: Union[bool, str]) -> None: # pragma: no cover "cannot set `allow_cache` for ConstantVolumetricBackend directly;" ) + @property + def overwrite_partial_chunks(self) -> bool: # pragma: no cover + return False + + @overwrite_partial_chunks.setter + def overwrite_partial_chunks(self, value: bool) -> None: # pragma: no cover + raise NotImplementedError( + "cannot set `overwrite_partial_chunks` for ConstantVolumetricBackend directly;" + ) + @property def use_compression(self) -> bool: # pragma: no cover return False @@ -126,3 +136,7 @@ def assert_idx_is_chunk_aligned(self, idx: VolumetricIndex) -> None: # pragma: def pformat(self) -> str: # pragma: no cover return self.name + + def delete(self): # pragma: no cover + # No-op for constant backend (no data to delete) + pass diff --git a/zetta_utils/layer/volumetric/layer.py b/zetta_utils/layer/volumetric/layer.py index 43e4f298e..778dca8c9 100644 --- a/zetta_utils/layer/volumetric/layer.py +++ b/zetta_utils/layer/volumetric/layer.py @@ -62,3 +62,6 @@ def with_changes( **kwargs, ): return attrs.evolve(self, **kwargs) # pragma: no cover + + def delete(self): + self.backend.delete() diff --git a/zetta_utils/layer/volumetric/layer_set/backend.py b/zetta_utils/layer/volumetric/layer_set/backend.py index 83fa95077..ae28a9199 100644 --- a/zetta_utils/layer/volumetric/layer_set/backend.py +++ b/zetta_utils/layer/volumetric/layer_set/backend.py @@ -79,6 +79,25 @@ def enforce_chunk_aligned_writes(self, value: bool) -> None: # pragma: no cover " use `backend.with_changes(non_aligned_writes=value:bool)` instead." ) + @property + def overwrite_partial_chunks(self) -> bool: # pragma: no cover + overwrite_partial_chunks = { + k: v.backend.overwrite_partial_chunks for k, v in self.layers.items() + } + if not len(set(overwrite_partial_chunks.values())) == 1: + raise ValueError( + "Cannot determine consistent `overwrite_partial_chunks` for the " + f"volumetric layer set backend. Got: {overwrite_partial_chunks}" + ) + return list(overwrite_partial_chunks.values())[0] + + @overwrite_partial_chunks.setter + def overwrite_partial_chunks(self, value: bool) -> None: # pragma: no cover + raise NotImplementedError( + "cannot set `overwrite_partial_chunks` for VolumetricSetBackend directly;" + " use `backend.with_changes(overwrite_partial_chunks=value:bool)` instead." + ) + @property def allow_cache(self) -> bool: # pragma: no cover allow_caches = {k: v.backend.allow_cache for k, v in self.layers.items()} @@ -197,3 +216,7 @@ def with_changes(self, **kwargs) -> VolumetricSetBackend: # pragma: no cover def pformat(self) -> str: # pragma: no cover return "\n".join([f"{k}: {v.pformat()}" for (k, v) in self.layers.items()]) + + def delete(self): + for layer in self.layers.values(): + layer.delete() diff --git a/zetta_utils/layer/volumetric/protocols.py b/zetta_utils/layer/volumetric/protocols.py index f27cc8db1..f0ff10e65 100644 --- a/zetta_utils/layer/volumetric/protocols.py +++ b/zetta_utils/layer/volumetric/protocols.py @@ -61,3 +61,6 @@ def with_changes( **kwargs, ) -> VolumetricBasedLayerProtocol[DataT, IndexT]: ... + + def delete(self): + ... diff --git a/zetta_utils/layer/volumetric/seg_contact/backend.py b/zetta_utils/layer/volumetric/seg_contact/backend.py index 15190eac8..2bc1f3d97 100644 --- a/zetta_utils/layer/volumetric/seg_contact/backend.py +++ b/zetta_utils/layer/volumetric/seg_contact/backend.py @@ -37,6 +37,9 @@ def name(self) -> str: def with_changes(self, **kwargs) -> SegContactLayerBackend: return attrs.evolve(self, **kwargs) + def delete(self): # pragma: no cover + raise NotImplementedError("delete() not implemented for SegContactLayerBackend") + @classmethod def from_path(cls, path: str) -> SegContactLayerBackend: """Load backend from existing info file.""" diff --git a/zetta_utils/layer/volumetric/tensorstore/backend.py b/zetta_utils/layer/volumetric/tensorstore/backend.py index 26ff5e3a4..6007ebaa1 100644 --- a/zetta_utils/layer/volumetric/tensorstore/backend.py +++ b/zetta_utils/layer/volumetric/tensorstore/backend.py @@ -7,6 +7,7 @@ import attrs import cachetools +import cloudfiles import numpy as np import tensorstore import torch @@ -76,6 +77,8 @@ class TSBackend(VolumetricBackend): # pylint: disable=too-few-public-methods :param on_info_exists: Behavior mode for when both `info_spec` is given and the layer info already exists. :param enforce_chunk_aligned_writes: Whether to allow non-chunk-aligned writes. + :param overwrite_partial_chunks: Whether to overwrite partial chunks. Not supported + for TSBackend - must be False or NotImplementedError will be raised. """ @@ -85,8 +88,13 @@ class TSBackend(VolumetricBackend): # pylint: disable=too-few-public-methods info_keep_existing_scales: bool = True cache_bytes_limit: Optional[int] = None _enforce_chunk_aligned_writes: bool = True + _overwrite_partial_chunks: bool = False def __attrs_post_init__(self): + if self._overwrite_partial_chunks: + raise NotImplementedError( + "overwrite_partial_chunks=True is not supported for TSBackend." + ) if self.info_spec is None: self.info_spec = PrecomputedInfoSpec(info_path=self.path) overwritten = self.info_spec.update_info( @@ -194,6 +202,17 @@ def use_compression(self, value: bool) -> None: # pragma: no cover "cannot set `use_compression` for TSBackend; can only be set to `False`" ) + @property + def overwrite_partial_chunks(self) -> bool: # pragma: no cover + return self._overwrite_partial_chunks + + @overwrite_partial_chunks.setter + def overwrite_partial_chunks(self, value: bool) -> None: # pragma: no cover + raise NotImplementedError( + "cannot set `overwrite_partial_chunks` for TSBackend directly;" + " use `backend.with_changes(overwrite_partial_chunks=value:bool)` instead." + ) + def clear_cache(self) -> None: # pragma: no cover _clear_ts_cache(self.path) @@ -255,6 +274,7 @@ def with_changes(self, **kwargs) -> TSBackend: "allow_cache" = value: Union[bool, str] - must be False for TensorStoreBackend, ignored "use_compression" = Value: bool - must be False for TensorStoreBackend, ignored "enforce_chunk_aligned_writes" = value: bool + "overwrite_partial_chunks" = value: bool - must be False for TensorStoreBackend, raises NotImplementedError if True "voxel_offset_res" = (voxel_offset, resolution): Tuple[Vec3D[int], Vec3D] "chunk_size_res" = (chunk_size, resolution): Tuple[Vec3D[int], Vec3D] "dataset_size_res" = (dataset_size, resolution): Tuple[Vec3D[int], Vec3D] @@ -268,6 +288,7 @@ def with_changes(self, **kwargs) -> TSBackend: "name", "allow_cache", "enforce_chunk_aligned_writes", + "overwrite_partial_chunks", "voxel_offset_res", "chunk_size_res", "dataset_size_res", @@ -282,7 +303,8 @@ def with_changes(self, **kwargs) -> TSBackend: keys_to_kwargs = { "name": "path", - "enforce_chunk_aligned_writes": "enforce_chunk_aligned_writes", + "enforce_chunk_aligned_writes": "_enforce_chunk_aligned_writes", + "overwrite_partial_chunks": "_overwrite_partial_chunks", } evolve_kwargs = {} @@ -326,3 +348,15 @@ def get_bounds(self, resolution: Vec3D) -> VolumetricIndex: # pragma: no cover def pformat(self) -> str: # pragma: no cover return self.name + + def delete(self): + # Clear cache to allow TS objects to go out of scope + self.clear_cache() + + path = abspath(self.path) + + # Use cloudfiles to delete all files + cf = cloudfiles.CloudFiles(path) + files_to_delete = list(cf.list()) + if files_to_delete: + cf.delete(files_to_delete) diff --git a/zetta_utils/layer/volumetric/tensorstore/deprecated/backend.py b/zetta_utils/layer/volumetric/tensorstore/deprecated/backend.py index 063e0865d..7ce099c65 100644 --- a/zetta_utils/layer/volumetric/tensorstore/deprecated/backend.py +++ b/zetta_utils/layer/volumetric/tensorstore/deprecated/backend.py @@ -311,3 +311,16 @@ def get_bounds(self, resolution: Vec3D) -> VolumetricIndex: # pragma: no cover def pformat(self) -> str: # pragma: no cover return self.name + + @property + def overwrite_partial_chunks(self) -> bool: # pragma: no cover + return False + + @overwrite_partial_chunks.setter + def overwrite_partial_chunks(self, value: bool) -> None: # pragma: no cover + raise NotImplementedError( + "cannot set `overwrite_partial_chunks` for deprecated TSBackend directly" + ) + + def delete(self): # pragma: no cover + raise NotImplementedError("delete() not implemented for deprecated TSBackend") diff --git a/zetta_utils/mazepa_layer_processing/__init__.py b/zetta_utils/mazepa_layer_processing/__init__.py index 50ea74c1b..05f543cb6 100644 --- a/zetta_utils/mazepa_layer_processing/__init__.py +++ b/zetta_utils/mazepa_layer_processing/__init__.py @@ -5,6 +5,7 @@ ComputeFieldOpProtocol, ChunkableOpProtocol, VolumetricOpProtocol, + StackableVolumetricOpProtocol, ) from .common import ( diff --git a/zetta_utils/mazepa_layer_processing/common/__init__.py b/zetta_utils/mazepa_layer_processing/common/__init__.py index fc057a179..2f82ec5bb 100644 --- a/zetta_utils/mazepa_layer_processing/common/__init__.py +++ b/zetta_utils/mazepa_layer_processing/common/__init__.py @@ -18,6 +18,7 @@ VolumetricCallableOperation, build_chunked_volumetric_callable_flow_schema, ) +from .stacked_volumetric_operations import StackedVolumetricOperations from .. import ChunkableOpProtocol, VolumetricOpProtocol from .interpolate_flow import build_interpolate_flow from .seg_contact_op import SegContactOp diff --git a/zetta_utils/mazepa_layer_processing/common/stacked_volumetric_operations.py b/zetta_utils/mazepa_layer_processing/common/stacked_volumetric_operations.py new file mode 100644 index 000000000..7a432f684 --- /dev/null +++ b/zetta_utils/mazepa_layer_processing/common/stacked_volumetric_operations.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import time +from typing import Any, Generic, Sequence, cast + +import attrs +import numpy as np +import torch +from typing_extensions import ParamSpec + +from zetta_utils import builder, log, mazepa +from zetta_utils.geometry import Vec3D +from zetta_utils.layer.volumetric import VolumetricIndex, VolumetricLayer + +from ..operation_protocols import StackableVolumetricOpProtocol + +logger = log.get_logger("mazepa") +P = ParamSpec("P") + + +@builder.register("StackedVolumetricOperations") +@mazepa.taskable_operation_cls +@attrs.mutable +class StackedVolumetricOperations(Generic[P]): + """ + Wrapper that processes multiple indices for a StackableVolumetricOpProtocol operation. + Optimizes I/O by reading and writing multiple chunks in batch. + + The batching is controlled by the flow that creates tasks with this operation. + This operation processes all indices it receives in a single batch. + + :param base_op: The stackable operation to wrap. + """ + + base_op: StackableVolumetricOpProtocol[P, None, VolumetricLayer] + + def __attrs_post_init__(self): + if not isinstance(self.base_op, StackableVolumetricOpProtocol): + raise TypeError( + f"{type(self.base_op).__name__} does not implement StackableVolumetricOpProtocol. " + f"Missing read/write methods required for stacking." + ) + + def get_operation_name(self) -> str: + base_name = ( + self.base_op.get_operation_name() + if hasattr(self.base_op, "get_operation_name") + else type(self.base_op).__name__ + ) + return f"Stacked({base_name})" + + def get_input_resolution(self, dst_resolution: Vec3D) -> Vec3D: + return self.base_op.get_input_resolution(dst_resolution) + + def with_added_crop_pad(self, crop_pad: Vec3D[int]) -> StackedVolumetricOperations[P]: + return attrs.evolve(self, base_op=self.base_op.with_added_crop_pad(crop_pad)) + + def __call__( # pylint: disable=keyword-arg-before-vararg,too-many-branches + self, + indices: Sequence[VolumetricIndex], + dsts: Sequence[VolumetricLayer] | VolumetricLayer, + *args: P.args, + **kwargs: P.kwargs, + ) -> None: + """ + Process multiple indices in a single batch. + + :param indices: List of VolumetricIndex objects to process. + :param dsts: Destination layer(s) for writing results. Can be a single layer + (same for all indices) or a sequence of layers (one per index). + :param args: Additional positional arguments passed to base_op. + :param kwargs: Keyword arguments passed to base_op for reading source data. + """ + if len(indices) == 0: + return + + # Normalize dsts to a sequence + if not isinstance(dsts, Sequence): + dsts_list = [dsts] * len(indices) + else: + dsts_list = list(dsts) + if len(dsts_list) != len(indices): + raise ValueError( + f"Length of dsts ({len(dsts_list)}) must match " + f"length of indices ({len(indices)})" + ) + + # Read all data + read_start = time.time() + data_list = [self.base_op.read(idx, *args, **kwargs) for idx in indices] + read_time = time.time() - read_start + + # Stack tensors by key + # Assumes all dicts have the same keys + if not data_list: + return + + keys = data_list[0].keys() + stacked_kwargs: dict[str, Any] = {} + + for key in keys: + tensors = [d[key] for d in data_list] + + # Stack tensors + if isinstance(tensors[0], torch.Tensor): + stacked_kwargs[key] = torch.stack(cast(list[torch.Tensor], tensors), dim=0) + elif isinstance(tensors[0], np.ndarray): + stacked_kwargs[key] = np.stack(cast(list[np.ndarray], tensors), axis=0) + else: + raise TypeError( + f"Read method returned unsupported type for key '{key}': {type(tensors[0])}. " + f"Only torch.Tensor and np.ndarray are supported for stacking." + ) + + # Process the batch with the base operation's processing function + process_start = time.time() + batched_result: Any = self.base_op.processing_fn(**stacked_kwargs) + process_time = time.time() - process_start + + # Unstack and write results + write_start = time.time() + for i, (idx, dst) in enumerate(zip(indices, dsts_list)): + result: Any + if isinstance(batched_result, torch.Tensor): + result = batched_result[i] + elif isinstance(batched_result, np.ndarray): + result = batched_result[i] + else: + raise TypeError( + f"Function returned unsupported type: {type(batched_result)}. " + f"Only torch.Tensor and np.ndarray are supported." + ) + + self.base_op.write(idx, dst, result, *args, **kwargs) + write_time = time.time() - write_start + + total_time = read_time + process_time + write_time + logger.info( + f"StackedVolumetricOperations: Total time for {len(indices)} chunks: {total_time:.2f}s" + f" (read: {read_time:.2f}s, process: {process_time:.2f}s, write: {write_time:.2f}s)" + ) diff --git a/zetta_utils/mazepa_layer_processing/common/subchunkable_apply_flow.py b/zetta_utils/mazepa_layer_processing/common/subchunkable_apply_flow.py index 6d54ff9c6..9009d1a2b 100644 --- a/zetta_utils/mazepa_layer_processing/common/subchunkable_apply_flow.py +++ b/zetta_utils/mazepa_layer_processing/common/subchunkable_apply_flow.py @@ -93,6 +93,7 @@ def build_postpad_subchunkable_apply_flow( # pylint: disable=keyword-arg-before generate_ng_link: bool = False, op_worker_type: str | None = None, reduction_worker_type: str | None = None, + task_stack_size: int | None = None, fn: Callable[P, Tensor] | None = None, fn_semaphores: Sequence[SemaphoreType] | None = None, op: VolumetricOpProtocol[P, None, Any] | None = None, @@ -184,6 +185,7 @@ def build_postpad_subchunkable_apply_flow( # pylint: disable=keyword-arg-before allow_cache_up_to_level=allow_cache_up_to_level, op_worker_type=op_worker_type, reduction_worker_type=reduction_worker_type, + task_stack_size=task_stack_size, print_summary=print_summary, generate_ng_link=generate_ng_link, fn=fn, @@ -223,6 +225,7 @@ def build_subchunkable_apply_flow( # pylint: disable=keyword-arg-before-vararg, allow_cache_up_to_level: int | None = None, op_worker_type: str | None = None, reduction_worker_type: str | None = None, + task_stack_size: int | None = None, print_summary: bool = True, generate_ng_link: bool = False, fn: Callable[P, Tensor] | None = None, @@ -574,6 +577,7 @@ def build_subchunkable_apply_flow( # pylint: disable=keyword-arg-before-vararg, auto_divisibility=auto_divisibility, op_worker_type=op_worker_type, reduction_worker_type=reduction_worker_type, + task_stack_size=task_stack_size, print_summary=print_summary, generate_ng_link=generate_ng_link, op_args=op_args, @@ -850,6 +854,7 @@ def _print_summary( # pylint: disable=line-too-long, too-many-locals, too-many- op_args: Iterable, op_kwargs: Mapping[str, Any], processing_gap: Vec3D[int], + task_stack_size: int | None, ) -> None: # pragma: no cover summary = "" summary += ( @@ -985,6 +990,14 @@ def _print_summary( # pylint: disable=line-too-long, too-many-locals, too-many- + "\n" ) summary += lrpad(f"Processing gap: {processing_gap.pformat()} ", level=2, length=120) + "\n" + summary += ( + lrpad( + f"Task stack size: {task_stack_size if task_stack_size is not None else 'None'} ", + level=2, + length=120, + ) + + "\n" + ) summary += lrpad(f"# of op_args supplied: {len(list(op_args))}", level=2, length=120) + "\n" summary += lrpad("op_kwargs supplied:", level=2, length=120) + "\n" @@ -1069,6 +1082,7 @@ def _build_subchunkable_apply_flow( # pylint: disable=keyword-arg-before-vararg auto_divisibility: bool, op_worker_type: str | None, reduction_worker_type: str | None, + task_stack_size: int | None, print_summary: bool, generate_ng_link: bool, op: VolumetricOpProtocol[P, None, Any], @@ -1222,7 +1236,10 @@ def _build_subchunkable_apply_flow( # pylint: disable=keyword-arg-before-vararg ) else: dst = deepcopy(dst).with_changes( - backend=dst.backend.with_changes(enforce_chunk_aligned_writes=False) + backend=dst.backend.with_changes( + enforce_chunk_aligned_writes=False, + overwrite_partial_chunks=True, + ) ) if print_summary: @@ -1251,6 +1268,7 @@ def _build_subchunkable_apply_flow( # pylint: disable=keyword-arg-before-vararg op_args=op_args, op_kwargs=op_kwargs, processing_gap=processing_gap, + task_stack_size=task_stack_size, ) """ Generate flow id for deconflicting intermediaries - must be deterministic for proper @@ -1278,6 +1296,7 @@ def _build_subchunkable_apply_flow( # pylint: disable=keyword-arg-before-vararg l0_chunks_per_task=num_chunks_below[-1], op_worker_type=op_worker_type, reduction_worker_type=reduction_worker_type, + task_stack_size=task_stack_size, ) """ Iteratively build the hierarchy of schemas @@ -1306,5 +1325,6 @@ def _build_subchunkable_apply_flow( # pylint: disable=keyword-arg-before-vararg l0_chunks_per_task=num_chunks_below[-level - 1], op_worker_type=op_worker_type, reduction_worker_type=reduction_worker_type, + task_stack_size=None, # Only stack at the bottom level (level 0) ) return flow_schema(idx, dst, op_args, op_kwargs) diff --git a/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py b/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py index 3f2d937b6..f5652d9e8 100644 --- a/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py +++ b/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py @@ -5,11 +5,20 @@ from abc import ABC from copy import deepcopy from os import path -from typing import Any, Generic, List, Literal, Optional, Tuple, TypeVar, assert_never +from typing import ( + Any, + Generic, + List, + Literal, + Optional, + Tuple, + TypeVar, + assert_never, + cast, +) import attrs import cachetools -import fsspec import numpy as np import torch from typeguard import suppress_type_checks @@ -26,7 +35,8 @@ from zetta_utils.mazepa import semaphore from zetta_utils.tensor_ops import convert -from ..operation_protocols import VolumetricOpProtocol +from ..operation_protocols import StackableVolumetricOpProtocol, VolumetricOpProtocol +from .stacked_volumetric_operations import StackedVolumetricOperations _weights_cache: cachetools.LRUCache = cachetools.LRUCache(maxsize=16) @@ -103,8 +113,9 @@ def __call__( subidx_channels = (slice(0, res.shape[0]), *subidx) with semaphore("read"): res[subidx_channels] = layer[intscn] - with semaphore("write"): - dst[red_idx] = res + if np.any(res): + with semaphore("write"): + dst[red_idx] = res def is_floating_point_dtype(dtype: np.dtype) -> bool: @@ -160,12 +171,12 @@ def __call__( with semaphore("read"): if not is_floating_point_dtype(dst.backend.dtype): # Temporarily convert integer cutout to float for rounding - res[subidx_channels] = ( - res[subidx_channels] + layer[intscn].astype(float) * weight.numpy() + res[tuple(subidx_channels)] = ( + res[tuple(subidx_channels)] + layer[intscn].astype(float) * weight.numpy() ) else: - res[subidx_channels] = ( - res[subidx_channels] + layer[intscn] * weight.numpy() + res[tuple(subidx_channels)] = ( + res[tuple(subidx_channels)] + layer[intscn] * weight.numpy() ) if not is_floating_point_dtype(dst.backend.dtype): @@ -176,8 +187,9 @@ def __call__( subidx_channels = [slice(0, res.shape[0])] + list(subidx) with semaphore("read"): res.numpy()[tuple(subidx_channels)] = layer[intscn] - with semaphore("write"): - dst[red_idx] = res + if res.any(): + with semaphore("write"): + dst[red_idx] = res @cachetools.cached(_weights_cache) @@ -321,15 +333,15 @@ def clear_cache(*args, **kwargs): def delete_if_local(*args, **kwargs): - filesystem = fsspec.filesystem("file") for arg in args: if isinstance(arg, VolumetricBasedLayerProtocol): if arg.backend.is_local: - filesystem.delete(arg.backend.name, recursive=True) + arg.delete() + for kwarg in kwargs.values(): if isinstance(kwarg, VolumetricBasedLayerProtocol): if kwarg.backend.is_local: - filesystem.delete(kwarg.backend.name, recursive=True) + kwarg.delete() @mazepa.flow_schema_cls @@ -354,6 +366,7 @@ class VolumetricApplyFlowSchema(Generic[P, R_co]): l0_chunks_per_task: int = 0 op_worker_type: str | None = None reduction_worker_type: str | None = None + task_stack_size: int | None = None @property def _intermediaries_are_local(self) -> bool: @@ -400,6 +413,16 @@ def _get_backend_chunk_size_to_use(self, dst) -> Vec3D[int]: return backend_chunk_size def __attrs_post_init__(self): # pylint: disable=too-many-branches + if self.task_stack_size is not None: + if self.task_stack_size < 1: + raise ValueError(f"`task_stack_size` must be >= 1, got {self.task_stack_size}") + if not isinstance(self.op, StackableVolumetricOpProtocol): + raise TypeError( + f"`task_stack_size` was provided ({self.task_stack_size}), but the operation " + f"{type(self.op).__name__} does not implement StackableVolumetricOpProtocol. " + f"Operations must have `read`, `write`, and `processing_fn` methods to support stacking." + ) + if self.roi_crop_pad is None: self.roi_crop_pad = Vec3D[int](0, 0, 0) if self.processing_blend_pad is None: @@ -480,29 +503,145 @@ def _make_task( self.op_worker_type ) - def make_tasks_without_checkerboarding( + def _create_tasks( self, - idx_chunks: List[VolumetricIndex], + idx_chunks: List[VolumetricIndex] | np.ndarray, dst: VolumetricBasedLayerProtocol | None, op_kwargs: P.kwargs, ) -> List[mazepa.tasks.Task[R_co]]: - if len(idx_chunks) > MULTIPROCESSING_NUM_TASKS_THRESHOLD: - with multiprocessing.get_context("fork").Pool( - initializer=reset_signal_handlers - ) as pool_obj: - tasks = pool_obj.map( - self._make_task, - zip(idx_chunks, itertools.repeat(dst), itertools.repeat(op_kwargs)), - ) + """Create tasks with optional batching based on task_stack_size.""" + # Flatten if needed + if isinstance(idx_chunks, np.ndarray): + idx_chunks_flat = list(idx_chunks.ravel()) else: - tasks = list( - map( - self._make_task, - zip(idx_chunks, itertools.repeat(dst), itertools.repeat(op_kwargs)), + idx_chunks_flat = idx_chunks + + if self.task_stack_size is None or self.task_stack_size == 1: + # No stacking, create one task per index + if len(idx_chunks_flat) > MULTIPROCESSING_NUM_TASKS_THRESHOLD: + with multiprocessing.get_context("fork").Pool( + initializer=reset_signal_handlers + ) as pool_obj: + tasks = pool_obj.map( + self._make_task, + zip( + idx_chunks_flat, + itertools.repeat(dst), + itertools.repeat(op_kwargs), + ), + ) + else: + tasks = list( + map( + self._make_task, + zip( + idx_chunks_flat, + itertools.repeat(dst), + itertools.repeat(op_kwargs), + ), + ) ) + else: + # Batching with stacked operations + assert isinstance(self.op, StackableVolumetricOpProtocol) + stacked_op: StackedVolumetricOperations[P] = StackedVolumetricOperations( + base_op=self.op ) + + tasks = [] + for batch_start in range(0, len(idx_chunks_flat), self.task_stack_size): + batch_end = min(batch_start + self.task_stack_size, len(idx_chunks_flat)) + batch_indices = idx_chunks_flat[batch_start:batch_end] + + task = stacked_op.make_task( + indices=batch_indices, dsts=dst, **op_kwargs + ).with_worker_type(self.op_worker_type) + tasks.append(task) + return tasks + def make_tasks_without_checkerboarding( + self, + idx_chunks: List[VolumetricIndex], + dst: VolumetricBasedLayerProtocol | None, + op_kwargs: P.kwargs, + ) -> List[mazepa.tasks.Task[R_co]]: + return self._create_tasks(idx_chunks, dst, op_kwargs) + + def _map_tasks_to_reduction_chunks( # pylint: disable=too-many-locals + self, + task_idxs: np.ndarray, + task_shape: Vec3D[int], + red_chunks: List[VolumetricIndex], + red_chunks_3d: np.ndarray, + red_shape: Vec3D[int], + dst_temp: VolumetricBasedLayerProtocol, + ) -> Tuple[List[List[VolumetricIndex]], List[List[VolumetricBasedLayerProtocol]]]: + """ + Map task indices to reduction chunks they contribute to. + + Returns: + - red_chunks_task_idxs: For each reduction chunk, list of task indices that contribute + - red_chunks_temps: For each reduction chunk, list of temp layers to read from + """ + red_chunks_task_idxs: List[List[VolumetricIndex]] = [[] for _ in red_chunks] + red_chunks_temps: List[List[VolumetricBasedLayerProtocol]] = [[] for _ in red_chunks] + + red_stops = [ + [chunk.stop[0] for chunk in red_chunks_3d[:, 0, 0]], + [chunk.stop[1] for chunk in red_chunks_3d[0, :, 0]], + [chunk.stop[2] for chunk in red_chunks_3d[0, 0, :]], + ] + + task_starts = [ + [task_idxs[i, 0, 0].start[0] for i in range(task_shape[0])], + [task_idxs[0, i, 0].start[1] for i in range(task_shape[1])], + [task_idxs[0, 0, i].start[2] for i in range(task_shape[2])], + ] + + task_to_red_chunks: list[dict[int, int]] = [{}, {}, {}] + + for axis in range(3): + red_stop_ind = 0 + for i, task_start in enumerate(task_starts[axis]): + try: + while not task_start < red_stops[axis][red_stop_ind]: + red_stop_ind += 1 + # This case catches the case where the chunk is entirely outside + # any reduction chunk; this can happen if, for instance, + # roi_crop_pad is set to [0, 0, 1] and the processing_chunk_size + # is [X, X, 1]. + except IndexError as e: + raise ValueError( + f"The processing chunk starting at `{task_start}` in axis {axis}" + " does not correspond to any reduction chunk; please check the " + "`roi_crop_pad` and the `processing_chunk_size`." + ) from e + task_to_red_chunks[axis][i] = red_stop_ind + + flat_red_ind_offsets = set( + i + j * red_shape[0] + k * red_shape[0] * red_shape[1] + for i, j, k in itertools.product(range(3), repeat=3) + ) + + for i, task_ind in enumerate(np.ndindex(task_idxs.shape)): + task_idx = task_idxs[*task_ind] + red_ind = Vec3D(*(task_to_red_chunks[axis][task_ind[axis]] for axis in range(3))) + flat_red_ind = ( + red_ind[0] + red_shape[0] * red_ind[1] + red_shape[0] * red_shape[1] * red_ind[2] + ) + if task_idx.contained_in(red_chunks[flat_red_ind]): + red_chunks_task_idxs[flat_red_ind].append(task_idx) + red_chunks_temps[flat_red_ind].append(dst_temp) + else: + flat_red_inds = [offset + flat_red_ind for offset in flat_red_ind_offsets] + for i in flat_red_inds: + if i < len(red_chunks) and task_idx.intersects(red_chunks[i]): + red_chunks_task_idxs[i].append(task_idx) + red_chunks_temps[i].append(dst_temp) + + return red_chunks_task_idxs, red_chunks_temps + def make_tasks_with_intermediaries( # pylint: disable=too-many-locals self, idx: VolumetricIndex, @@ -542,12 +681,16 @@ def make_tasks_with_checkerboarding( # pylint: disable=too-many-locals, too-man the list of VolumetricIndices and the VolumetricBasedLayerProtocols that can be reduced to the final output, as well as the temporary destination layers. """ - tasks: List[mazepa.tasks.Task[R_co]] = [] red_chunks_task_idxs: List[List[VolumetricIndex]] = [[] for _ in red_chunks] red_chunks_temps: List[List[VolumetricBasedLayerProtocol]] = [[] for _ in red_chunks] red_chunks_3d = np.array(red_chunks, dtype=object).reshape(red_shape, order="F") dst_temps: List[VolumetricBasedLayerProtocol] = [] + # Accumulate all chunks and destinations across all checkerboard phases + all_task_idxs: List[VolumetricIndex] = [] + all_dst_temps: List[VolumetricBasedLayerProtocol] = [] + phase_metadata: List[Tuple[np.ndarray, Vec3D[int], VolumetricBasedLayerProtocol]] = [] + next_chunk_id = idx.chunk_id for chunker, chunker_idx in self.processing_chunker.split_into_nonoverlapping_chunkers( self.processing_blend_pad @@ -596,88 +739,228 @@ def make_tasks_with_checkerboarding( # pylint: disable=too-many-locals, too-man next_chunk_id = task_idxs[-1, -1, -1].chunk_id + self.l0_chunks_per_task - if len(task_idxs) > MULTIPROCESSING_NUM_TASKS_THRESHOLD: - with multiprocessing.get_context("fork").Pool( - initializer=reset_signal_handlers - ) as pool_obj: - tasks_split = pool_obj.map( - self._make_task, - zip( - task_idxs.ravel(), - itertools.repeat(dst_temp), - itertools.repeat(op_kwargs), - ), - ) - else: - tasks_split = list( - map( - self._make_task, - zip( - task_idxs.ravel(), - itertools.repeat(dst_temp), - itertools.repeat(op_kwargs), - ), - ) - ) - tasks += tasks_split - - red_stops = [ - [chunk.stop[0] for chunk in red_chunks_3d[:, 0, 0]], - [chunk.stop[1] for chunk in red_chunks_3d[0, :, 0]], - [chunk.stop[2] for chunk in red_chunks_3d[0, 0, :]], - ] - - task_starts = [ - [task_idxs[i, 0, 0].start[0] for i in range(task_shape[0])], - [task_idxs[0, i, 0].start[1] for i in range(task_shape[1])], - [task_idxs[0, 0, i].start[2] for i in range(task_shape[2])], - ] - - task_to_red_chunks: list[dict[int, int]] = [{}, {}, {}] - - for axis in range(3): - red_stop_ind = 0 - for i, task_start in enumerate(task_starts[axis]): - try: - while not task_start < red_stops[axis][red_stop_ind]: - red_stop_ind += 1 - # This case catches the case where the chunk is entirely outside - # any reduction chunk; this can happen if, for instance, - # roi_crop_pad is set to [0, 0, 1] and the processing_chunk_size - # is [X, X, 1]. - except IndexError as e: - raise ValueError( - f"The processing chunk starting at `{task_start}` in axis {axis}" - " does not correspond to any reduction chunk; please check the " - "`roi_crop_pad` and the `processing_chunk_size`." - ) from e - task_to_red_chunks[axis][i] = red_stop_ind - - flat_red_ind_offsets = set( - i + j * red_shape[0] + k * red_shape[0] * red_shape[1] - for i, j, k in itertools.product(range(3), repeat=3) + # Store phase metadata for later mapping + phase_metadata.append((task_idxs, task_shape, dst_temp)) + + # Accumulate chunks and destinations across all phases + task_idxs_flat = list(task_idxs.ravel()) + all_task_idxs.extend(task_idxs_flat) + all_dst_temps.extend([dst_temp] * len(task_idxs_flat)) + + # Now create batched tasks across all phases + if self.task_stack_size is not None and self.task_stack_size > 1: + assert isinstance(self.op, StackableVolumetricOpProtocol) + stacked_op: StackedVolumetricOperations[P] = StackedVolumetricOperations( + base_op=self.op + ) + + tasks: List[mazepa.tasks.Task[R_co]] = [] + for batch_start in range(0, len(all_task_idxs), self.task_stack_size): + batch_end = min(batch_start + self.task_stack_size, len(all_task_idxs)) + batch_indices = all_task_idxs[batch_start:batch_end] + batch_dsts = all_dst_temps[batch_start:batch_end] + + task = stacked_op.make_task( + indices=batch_indices, dsts=batch_dsts, **op_kwargs + ).with_worker_type(self.op_worker_type) + tasks.append(task) + else: + # No stacking - create one task per chunk + tasks = [] + for task_idx, dst_temp in zip(all_task_idxs, all_dst_temps): + task = self.op.make_task(task_idx, dst_temp, **op_kwargs).with_worker_type( + self.op_worker_type ) - for i, task_ind in enumerate(np.ndindex(task_idxs.shape)): - task_idx = task_idxs[*task_ind] - red_ind = Vec3D( - *(task_to_red_chunks[axis][task_ind[axis]] for axis in range(3)) - ) - flat_red_ind = ( - red_ind[0] - + red_shape[0] * red_ind[1] - + red_shape[0] * red_shape[1] * red_ind[2] - ) - if task_idx.contained_in(red_chunks[flat_red_ind]): - red_chunks_task_idxs[flat_red_ind].append(task_idx) - red_chunks_temps[flat_red_ind].append(dst_temp) - else: - flat_red_inds = [offset + flat_red_ind for offset in flat_red_ind_offsets] - for i in flat_red_inds: - if i < len(red_chunks) and task_idx.intersects(red_chunks[i]): - red_chunks_task_idxs[i].append(task_idx) - red_chunks_temps[i].append(dst_temp) + tasks.append(task) + + # Map tasks to reduction chunks using phase metadata + for task_idxs, task_shape, dst_temp in phase_metadata: + ( + task_red_chunks_task_idxs, + task_red_chunks_temps, + ) = self._map_tasks_to_reduction_chunks( + task_idxs, task_shape, red_chunks, red_chunks_3d, red_shape, dst_temp + ) + + # Accumulate results + for i, (idxs, temps) in enumerate( + zip(task_red_chunks_task_idxs, task_red_chunks_temps) + ): + red_chunks_task_idxs[i].extend(idxs) + red_chunks_temps[i].extend(temps) + return (tasks, red_chunks_task_idxs, red_chunks_temps, dst_temps) + def _flow_simple( + self, idx: VolumetricIndex, dst: VolumetricBasedLayerProtocol | None, op_kwargs: P.kwargs + ) -> mazepa.FlowFnReturnType: + """Handle simple case without checkerboarding or intermediaries.""" + idx_chunks = self.processing_chunker( + idx, mode="exact", chunk_id_increment=self.l0_chunks_per_task + ) + tasks = self.make_tasks_without_checkerboarding(idx_chunks, dst, op_kwargs) + logger.info(f"Submitting {len(tasks)} processing tasks from operation {self.op}.") + yield tasks + + def _flow_with_intermediaries( # pylint: disable=too-many-locals + self, idx: VolumetricIndex, dst: VolumetricBasedLayerProtocol, op_kwargs: P.kwargs + ) -> mazepa.FlowFnReturnType: + """Handle case with intermediaries but no checkerboarding.""" + tasks, dst_temp = self.make_tasks_with_intermediaries(idx, dst, op_kwargs) + logger.info(f"Submitting {len(tasks)} processing tasks from operation {self.op}.") + yield tasks + yield mazepa.Dependency() + if self.processing_gap is None: + self.processing_gap = Vec3D[int](0, 0, 0) + if self.processing_gap != Vec3D[int](0, 0, 0): + copy_chunk_size = ( + dst.backend.get_chunk_size(self.dst_resolution) - self.processing_gap // 2 + ) + elif not self.max_reduction_chunk_size_final >= dst.backend.get_chunk_size( + self.dst_resolution + ): + copy_chunk_size = dst.backend.get_chunk_size(self.dst_resolution) + else: + copy_chunk_size = self.max_reduction_chunk_size_final + + reduction_chunker = VolumetricIndexChunker( + chunk_size=dst.backend.get_chunk_size(self.dst_resolution) + - self.processing_gap // 2, + resolution=self.dst_resolution, + max_superchunk_size=copy_chunk_size, + offset=-self.processing_gap // 2, + ) + logger.debug( + f"Breaking {idx} into chunks to be copied from the intermediary layer" + f" with {reduction_chunker}." + ) + stride_start_offset = dst.backend.get_voxel_offset(self.dst_resolution) + red_chunks = reduction_chunker(idx, mode="exact", stride_start_offset=stride_start_offset) + tasks_reduce = [ + Copy() + .make_task( + src=dst_temp, dst=dst.with_procs(read_procs=(), write_procs=()), idx=red_chunk + ) + .with_worker_type(self.reduction_worker_type) + for red_chunk in red_chunks + ] + logger.info( + "Copying temporary destination backend into the final destination:" + f" Submitting {len(tasks_reduce)} tasks." + ) + yield tasks_reduce + yield mazepa.Dependency() + clear_cache(dst_temp) + delete_if_local(dst_temp) + + def _flow_deferred_blending( + self, idx: VolumetricIndex, dst: VolumetricBasedLayerProtocol, op_kwargs: P.kwargs + ) -> mazepa.FlowFnReturnType: + """Handle case with deferred blending.""" + assert self.roi_crop_pad is not None + (tasks, _, _, _) = self.make_tasks_with_checkerboarding( + idx.padded(self.roi_crop_pad), [idx], Vec3D(1, 1, 1), dst, op_kwargs + ) + logger.info( + "Writing to intermediate destinations:\n" + f" Submitting {len(tasks)} processing tasks from operation {self.op}.\n" + f"Note that because blending is deferred, {dst.pformat()} will NOT " + f"contain the final output." + ) + yield tasks + yield mazepa.Dependency() + + def _flow_with_checkerboarding( # pylint: disable=too-many-locals + self, idx: VolumetricIndex, dst: VolumetricBasedLayerProtocol, op_kwargs: P.kwargs + ) -> mazepa.FlowFnReturnType: + """Handle full checkerboarding with reduction.""" + assert self.roi_crop_pad is not None + assert self.processing_blend_pad is not None + if dst.backend.enforce_chunk_aligned_writes: + try: + dst.backend.assert_idx_is_chunk_aligned(idx) + except Exception as e: + error_str = ( + "`dst` VolumetricBasedLayerProtocol's backend has" + " `enforce_chunk_aligned_writes`=True, but the provided `idx`" + " is not chunk aligned:\n" + ) + e.args = (error_str + e.args[0],) + raise e + if not self.max_reduction_chunk_size_final >= dst.backend.get_chunk_size( + self.dst_resolution + ): + raise ValueError( + "`max_reduction_chunk_size` (which defaults to `processing_chunk_size` when" + " not specified)` must be at least as large as the `dst` layer's" + f" chunk size; received {self.max_reduction_chunk_size_final}, which is" + f" smaller than {dst.backend.get_chunk_size(self.dst_resolution)}" + ) + reduction_chunker = VolumetricIndexChunker( + chunk_size=dst.backend.get_chunk_size(self.dst_resolution), + resolution=self.dst_resolution, + max_superchunk_size=self.max_reduction_chunk_size_final, + ) + logger.debug( + f"Breaking {idx} into reduction chunks with checkerboarding" + f" with {reduction_chunker}. Processing chunks will use the padded index" + f" {idx.padded(self.roi_crop_pad)} and be chunked with {self.processing_chunker}." + ) + stride_start_offset = dst.backend.get_voxel_offset(self.dst_resolution) + red_chunks = reduction_chunker(idx, mode="exact", stride_start_offset=stride_start_offset) + red_shape = reduction_chunker.get_shape( + idx, mode="exact", stride_start_offset=stride_start_offset + ) + ( + tasks, + red_chunks_task_idxs, + red_chunks_temps, + dst_temps, + ) = self.make_tasks_with_checkerboarding( + idx.padded(self.roi_crop_pad), red_chunks, red_shape, dst, op_kwargs + ) + logger.info( + "Writing to temporary destinations:\n" + f" Submitting {len(tasks)} processing tasks from operation {self.op}." + ) + yield tasks + yield mazepa.Dependency() + reducer: ReduceOperation + if self.processing_blend_mode == "max": + reducer = ReduceNaive() + elif self.processing_blend_mode in ("linear", "quadratic"): + reducer = ReduceByWeightedSum( + cast(Literal["linear", "quadratic"], self.processing_blend_mode) + ) + else: + raise ValueError( + f"Invalid processing_blend_mode: {self.processing_blend_mode}. " + f"Expected 'linear', 'quadratic', or 'max' in checkerboarding flow." + ) + tasks_reduce = [ + reducer.make_task( + src_idxs=red_chunk_task_idxs, + src_layers=red_chunk_temps, + red_idx=red_chunk, + roi_idx=idx.padded(self.roi_crop_pad + self.processing_blend_pad), + dst=dst.with_procs(read_procs=(), write_procs=()), + processing_blend_pad=self.processing_blend_pad, + ).with_worker_type(self.reduction_worker_type) + for ( + red_chunk_task_idxs, + red_chunk_temps, + red_chunk, + ) in zip(red_chunks_task_idxs, red_chunks_temps, red_chunks) + ] + logger.info( + "Collating temporary destination backends into the final destination:" + f" Submitting {len(tasks_reduce)} tasks." + ) + yield tasks_reduce + yield mazepa.Dependency() + clear_cache(*dst_temps) + delete_if_local(*dst_temps) + def flow( # pylint:disable=too-many-branches, too-many-statements self, idx: VolumetricIndex, @@ -694,161 +977,18 @@ def flow( # pylint:disable=too-many-branches, too-many-statements logger.debug(f"Breaking {idx} into chunks with {self.processing_chunker}.") - # cases without checkerboarding + # Dispatch to appropriate flow handler if not self.use_checkerboarding and not self.force_intermediaries: - idx_chunks = self.processing_chunker( - idx, mode="exact", chunk_id_increment=self.l0_chunks_per_task - ) - tasks = self.make_tasks_without_checkerboarding(idx_chunks, dst, op_kwargs) - logger.info(f"Submitting {len(tasks)} processing tasks from operation {self.op}.") - yield tasks + yield from self._flow_simple(idx, dst, op_kwargs) elif not self.use_checkerboarding and self.force_intermediaries: assert dst is not None - tasks, dst_temp = self.make_tasks_with_intermediaries(idx, dst, op_kwargs) - logger.info(f"Submitting {len(tasks)} processing tasks from operation {self.op}.") - yield tasks - yield mazepa.Dependency() - if self.processing_gap is None: - self.processing_gap = Vec3D[int](0, 0, 0) - if self.processing_gap != Vec3D[int](0, 0, 0): - copy_chunk_size = ( - dst.backend.get_chunk_size(self.dst_resolution) - self.processing_gap // 2 - ) - elif not self.max_reduction_chunk_size_final >= dst.backend.get_chunk_size( - self.dst_resolution - ): - copy_chunk_size = dst.backend.get_chunk_size(self.dst_resolution) - else: - copy_chunk_size = self.max_reduction_chunk_size_final - - reduction_chunker = VolumetricIndexChunker( - chunk_size=dst.backend.get_chunk_size(self.dst_resolution) - - self.processing_gap // 2, - resolution=self.dst_resolution, - max_superchunk_size=copy_chunk_size, - offset=-self.processing_gap // 2, - ) - logger.debug( - f"Breaking {idx} into chunks to be copied from the intermediary layer" - f" with {reduction_chunker}." - ) - stride_start_offset = dst.backend.get_voxel_offset(self.dst_resolution) - red_chunks = reduction_chunker( - idx, mode="exact", stride_start_offset=stride_start_offset - ) - tasks_reduce = [ - Copy() - .make_task( - src=dst_temp, dst=dst.with_procs(read_procs=(), write_procs=()), idx=red_chunk - ) - .with_worker_type(self.reduction_worker_type) - for red_chunk in red_chunks - ] - logger.info( - "Copying temporary destination backend into the final destination:" - f" Submitting {len(tasks_reduce)} tasks." - ) - yield tasks_reduce - yield mazepa.Dependency() - clear_cache(dst_temp) - delete_if_local(dst_temp) - # cases with checkerboarding + yield from self._flow_with_intermediaries(idx, dst, op_kwargs) elif self.processing_blend_mode == "defer": assert dst is not None - stride_start_offset = dst.backend.get_voxel_offset(self.dst_resolution) - (tasks, _, _, dst_temps,) = self.make_tasks_with_checkerboarding( - idx.padded(self.roi_crop_pad), [idx], Vec3D(1, 1, 1), dst, op_kwargs - ) - logger.info( - "Writing to intermediate destinations:\n" - f" Submitting {len(tasks)} processing tasks from operation {self.op}.\n" - f"Note that because blending is deferred, {dst.pformat()} will NOT " - f"contain the final output." - ) - yield tasks - yield mazepa.Dependency() + yield from self._flow_deferred_blending(idx, dst, op_kwargs) else: assert dst is not None - if dst.backend.enforce_chunk_aligned_writes: - try: - dst.backend.assert_idx_is_chunk_aligned(idx) - except Exception as e: - error_str = ( - "`dst` VolumetricBasedLayerProtocol's backend has" - " `enforce_chunk_aligned_writes`=True, but the provided `idx`" - " is not chunk aligned:\n" - ) - e.args = (error_str + e.args[0],) - raise e - if not self.max_reduction_chunk_size_final >= dst.backend.get_chunk_size( - self.dst_resolution - ): - raise ValueError( - "`max_reduction_chunk_size` (which defaults to `processing_chunk_size` when" - " not specified)` must be at least as large as the `dst` layer's" - f" chunk size; received {self.max_reduction_chunk_size_final}, which is" - f" smaller than {dst.backend.get_chunk_size(self.dst_resolution)}" - ) - reduction_chunker = VolumetricIndexChunker( - chunk_size=dst.backend.get_chunk_size(self.dst_resolution), - resolution=self.dst_resolution, - max_superchunk_size=self.max_reduction_chunk_size_final, - ) - logger.debug( - f"Breaking {idx} into reduction chunks with checkerboarding" - f" with {reduction_chunker}. Processing chunks will use the padded index" - f" {idx.padded(self.roi_crop_pad)} and be chunked with {self.processing_chunker}." - ) - stride_start_offset = dst.backend.get_voxel_offset(self.dst_resolution) - red_chunks = reduction_chunker( - idx, mode="exact", stride_start_offset=stride_start_offset - ) - red_shape = reduction_chunker.get_shape( - idx, mode="exact", stride_start_offset=stride_start_offset - ) - ( - tasks, - red_chunks_task_idxs, - red_chunks_temps, - dst_temps, - ) = self.make_tasks_with_checkerboarding( - idx.padded(self.roi_crop_pad), red_chunks, red_shape, dst, op_kwargs - ) - logger.info( - "Writing to temporary destinations:\n" - f" Submitting {len(tasks)} processing tasks from operation {self.op}." - ) - yield tasks - yield mazepa.Dependency() - reducer: ReduceOperation - if self.processing_blend_mode == "max": - reducer = ReduceNaive() - else: - reducer = ReduceByWeightedSum(self.processing_blend_mode) - tasks_reduce = [ - reducer.make_task( - src_idxs=red_chunk_task_idxs, - src_layers=red_chunk_temps, - red_idx=red_chunk, - roi_idx=idx.padded(self.roi_crop_pad + self.processing_blend_pad), - dst=dst.with_procs(read_procs=(), write_procs=()), - processing_blend_pad=self.processing_blend_pad, - ).with_worker_type(self.reduction_worker_type) - for ( - red_chunk_task_idxs, - red_chunk_temps, - red_chunk, - ) in zip(red_chunks_task_idxs, red_chunks_temps, red_chunks) - ] - logger.info( - "Collating temporary destination backends into the final destination:" - f" Submitting {len(tasks_reduce)} tasks." - ) - yield tasks_reduce - yield mazepa.Dependency() - clear_cache(*dst_temps) - delete_if_local(*dst_temps) + yield from self._flow_with_checkerboarding(idx, dst, op_kwargs) + if self.clear_cache_on_return: clear_cache(*op_args, **op_kwargs) - - return tasks diff --git a/zetta_utils/mazepa_layer_processing/common/volumetric_callable_operation.py b/zetta_utils/mazepa_layer_processing/common/volumetric_callable_operation.py index af86e7b27..3fd9dc731 100644 --- a/zetta_utils/mazepa_layer_processing/common/volumetric_callable_operation.py +++ b/zetta_utils/mazepa_layer_processing/common/volumetric_callable_operation.py @@ -10,7 +10,7 @@ from numpy import typing as npt from typing_extensions import ParamSpec -from zetta_utils import builder, mazepa, tensor_ops +from zetta_utils import builder, log, mazepa, tensor_ops from zetta_utils.geometry import Vec3D from zetta_utils.layer import IndexChunker from zetta_utils.layer.volumetric import VolumetricIndex, VolumetricLayer @@ -19,6 +19,7 @@ from . import ChunkedApplyFlowSchema from .callable_operation import _process_callable_kwargs +logger = log.get_logger("zetta_utils") P = ParamSpec("P") IndexT = TypeVar("IndexT", bound=VolumetricIndex) @@ -75,42 +76,78 @@ def __attrs_post_init__(self): ) self.input_crop_pad = input_crop_pad_raw.int() - def __call__( # pylint: disable=keyword-arg-before-vararg + def read( # pylint: disable=keyword-arg-before-vararg self, idx: VolumetricIndex, - dst: VolumetricLayer | None, *args: P.args, **kwargs: P.kwargs, - ) -> Any: + ) -> dict[str, Any]: + """Read all source data for this operation.""" assert len(args) == 0 idx_input = copy.deepcopy(idx) idx_input.resolution = self.get_input_resolution(idx.resolution) idx_input_padded = idx_input.padded(Vec3D[int](*self.input_crop_pad)) + with semaphore("read"): task_kwargs = _process_callable_kwargs(idx_input_padded, kwargs) + return task_kwargs + + def write( # pylint: disable=keyword-arg-before-vararg,unused-argument + self, + idx: VolumetricIndex, + dst: VolumetricLayer, + tensor: Any, + *args: P.args, + **kwargs: P.kwargs, + ) -> None: + """Write tensor data to destination.""" + assert len(args) == 0 + # Data crop amount is determined by the index pad and the + # difference between the resolutions of idx and dst_idx. + # Padding was applied before the first read processor, so cropping + # should be done as last write processor. + if tuple(self.crop_pad) != (0, 0, 0): + dst_with_crop = dst.with_procs( + write_procs=dst.write_procs + (partial(tensor_ops.crop, crop=self.crop_pad),) + ) + else: + dst_with_crop = dst + + with semaphore("write"): + dst_with_crop[idx] = tensor + + def processing_fn(self, **kwargs: Any) -> Any: + """ + Process the data with semaphores if needed. + + :param kwargs: Named tensors to process + :return: Processed tensor + """ with ExitStack() as semaphore_stack: if self.fn_semaphores is not None: for semaphore_type in self.fn_semaphores: semaphore_stack.enter_context(semaphore(semaphore_type)) - result_raw = self.fn(**task_kwargs) + result_raw = self.fn(**kwargs) torch.cuda.empty_cache() + return result_raw + + def __call__( # pylint: disable=keyword-arg-before-vararg + self, + idx: VolumetricIndex, + dst: VolumetricLayer | None, + *args: P.args, + **kwargs: P.kwargs, + ) -> Any: + assert len(args) == 0 + task_kwargs = self.read(idx, *args, **kwargs) + + result_raw = self.processing_fn(**task_kwargs) + # If no destination layer, we can bail out now. Possibly we # could bail out a bit sooner. ToDo: study this more. if dst is not None: - # Data crop amount is determined by the index pad and the - # difference between the resolutions of idx and dst_idx. - # Padding was applied before the first read processor, so cropping - # should be done as last write processor. - if tuple(self.crop_pad) != (0, 0, 0): - dst_with_crop = dst.with_procs( - write_procs=dst.write_procs + (partial(tensor_ops.crop, crop=self.crop_pad),) - ) - else: - dst_with_crop = dst - - with semaphore("write"): - dst_with_crop[idx] = result_raw + self.write(idx, dst, result_raw, *args, **kwargs) return None else: return result_raw diff --git a/zetta_utils/mazepa_layer_processing/operation_protocols.py b/zetta_utils/mazepa_layer_processing/operation_protocols.py index fdf17c8bc..a9fcf1267 100644 --- a/zetta_utils/mazepa_layer_processing/operation_protocols.py +++ b/zetta_utils/mazepa_layer_processing/operation_protocols.py @@ -11,6 +11,7 @@ VolumetricIndex, VolumetricLayer, ) +from zetta_utils.tensor_typing import Tensor P = ParamSpec("P") R_co = TypeVar("R_co", covariant=True) @@ -115,3 +116,47 @@ def make_task( **kwargs: P.kwargs, ) -> mazepa.Task[R_co]: ... + + +@runtime_checkable +class StackableVolumetricOpProtocol(VolumetricOpProtocol[P, R_co, DstLayerT_contra], Protocol): + """ + Extension of VolumetricOpProtocol with read/write methods exposed. + Enables batching multiple indices for optimized I/O operations. + + Requires a `processing_fn` method that performs the core processing logic + and can be called independently of I/O operations, allowing for batched processing. + """ + + def with_added_crop_pad( + self, crop_pad: Vec3D[int] + ) -> StackableVolumetricOpProtocol[P, R_co, DstLayerT_contra]: + ... + + def processing_fn(self, **kwargs: dict[str, Tensor]) -> Tensor: + """ + Process the data. + Should accept both stacked and non-stacked Tensors. + Can only accept Tensors. + + :param kwargs: Named tensors + :return: Processed tensor + """ + + def read( + self, + idx: VolumetricIndex, + *args: P.args, + **kwargs: P.kwargs, + ) -> dict[str, Tensor]: + """Read all source data for this operation and return as named tensors.""" + + def write( + self, + idx: VolumetricIndex, + dst: DstLayerT_contra, + tensor: Tensor, + *args: P.args, + **kwargs: P.kwargs, + ) -> None: + """Write tensor data to destination.""" From b6a76f5c396625346882116d81e2b56c9ac6ab65 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Tue, 20 Jan 2026 17:38:16 -0800 Subject: [PATCH 5/9] chore: update internal --- zetta_utils/internal | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zetta_utils/internal b/zetta_utils/internal index 6016b81e0..4ffef7899 160000 --- a/zetta_utils/internal +++ b/zetta_utils/internal @@ -1 +1 @@ -Subproject commit 6016b81e0c3db3afed5bcbc74aef9a5197aceee4 +Subproject commit 4ffef78990d63a32634fa2e9e7ed52220a5a221e From 5582b2bccf99feb7af875dda735ad4192fc86c23 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Tue, 20 Jan 2026 17:59:07 -0800 Subject: [PATCH 6/9] chore: tests --- pyproject.toml | 1 + requirements.all.txt | 21 +- requirements.modules.txt | 19 +- tests/unit/layer/test_layer.py | 8 + .../layer/volumetric/cloudvol/test_backend.py | 37 ++ .../volumetric/tensorstore/test_backend.py | 45 ++ tests/unit/layer/volumetric/test_layer.py | 14 + tests/unit/layer/volumetric/test_layer_set.py | 11 + tests/unit/mazepa/test_upkeep_handlers.py | 463 ++++++++++++++++++ tests/unit/mazepa/test_worker.py | 85 +++- .../resource_allocation/k8s/gke.py | 1 + .../layer/volumetric/seg_contact/backend.py | 4 +- .../layer/volumetric/tensorstore/backend.py | 7 +- zetta_utils/mazepa/upkeep_handlers.py | 4 +- zetta_utils/mazepa/worker.py | 2 +- .../common/volumetric_apply_flow.py | 6 +- zetta_utils/message_queues/sqs/utils.py | 4 +- 17 files changed, 701 insertions(+), 31 deletions(-) create mode 100644 tests/unit/mazepa/test_upkeep_handlers.py diff --git a/pyproject.toml b/pyproject.toml index 28f87feac..39c5b9795 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ augmentations = [ "zetta_utils[tensor_ops]", "imaug", "imagecorruptions-imaug", + "numba >= 0.59", # Override old numba from imagecorruptions-imaug for Python 3.12+ "torchvision", ] cli = ["click >= 8.0.1"] diff --git a/requirements.all.txt b/requirements.all.txt index e598ac8ba..72951a8d7 100644 --- a/requirements.all.txt +++ b/requirements.all.txt @@ -38,7 +38,7 @@ attrs==25.4.0 # referencing autodocsumm==0.2.14 # via sphinx-toolbox -awscli==1.44.20 +awscli==1.44.21 # via zetta-utils (pyproject.toml) babel==2.17.0 # via sphinx @@ -53,14 +53,14 @@ blosc==1.11.4 # meshparty blosc2==3.12.2 # via tables -boto3==1.42.30 +boto3==1.42.31 # via # zetta-utils (pyproject.toml) # cloud-files # cloud-volume # moto # task-queue -botocore==1.42.30 +botocore==1.42.31 # via # awscli # boto3 @@ -507,7 +507,7 @@ lightning-utilities==0.15.2 # lightning # pytorch-lightning # torchmetrics -llvmlite==0.36.0 +llvmlite==0.46.0 # via numba lsds @ git+https://github.com/ZettaAI/lsd.git@cebe976133efb991ebc40171eeed7a22c57ba50a # via zetta-utils (pyproject.toml) @@ -600,13 +600,15 @@ nodeenv==1.10.0 # via pre-commit nose2==0.15.1 # via intern -numba==0.53.1 - # via imagecorruptions-imaug +numba==0.63.1 + # via + # zetta-utils (pyproject.toml) + # imagecorruptions-imaug numexpr==2.14.1 # via # blosc2 # tables -numpy==2.4.1 +numpy==2.3.5 # via # zetta-utils (pyproject.toml) # abiss @@ -1059,13 +1061,12 @@ scipy==1.17.0 # scikit-image # scikit-learn # trimesh -sentry-sdk==2.49.0 +sentry-sdk==2.50.0 # via wandb setuptools==80.9.0 # via # imagecorruptions-imaug # lightning-utilities - # numba # pbr # torch # torchfields @@ -1101,7 +1102,7 @@ snowballstemmer==3.0.1 # sphinx sortedcontainers==2.4.0 # via intervaltree -soupsieve==2.8.2 +soupsieve==2.8.3 # via beautifulsoup4 sphinx==6.2.1 # via diff --git a/requirements.modules.txt b/requirements.modules.txt index 4a5a6d92b..00334a4ac 100644 --- a/requirements.modules.txt +++ b/requirements.modules.txt @@ -28,7 +28,7 @@ attrs==25.4.0 # caveclient # jsonschema # referencing -awscli==1.44.20 +awscli==1.44.21 # via zetta-utils (pyproject.toml) blinker==1.9.0 # via flask @@ -39,13 +39,13 @@ blosc==1.11.4 # meshparty blosc2==3.12.2 # via tables -boto3==1.42.30 +boto3==1.42.31 # via # zetta-utils (pyproject.toml) # cloud-files # cloud-volume # task-queue -botocore==1.42.30 +botocore==1.42.31 # via # awscli # boto3 @@ -437,7 +437,7 @@ lightning-utilities==0.15.2 # lightning # pytorch-lightning # torchmetrics -llvmlite==0.36.0 +llvmlite==0.46.0 # via numba lsds @ git+https://github.com/ZettaAI/lsd.git@cebe976133efb991ebc40171eeed7a22c57ba50a # via zetta-utils (pyproject.toml) @@ -511,13 +511,15 @@ neuroglancer==2.41.2 # via zetta-utils (pyproject.toml) nose2==0.15.1 # via intern -numba==0.53.1 - # via imagecorruptions-imaug +numba==0.63.1 + # via + # zetta-utils (pyproject.toml) + # imagecorruptions-imaug numexpr==2.14.1 # via # blosc2 # tables -numpy==2.4.1 +numpy==2.3.5 # via # zetta-utils (pyproject.toml) # abiss @@ -919,13 +921,12 @@ scipy==1.17.0 # scikit-image # scikit-learn # trimesh -sentry-sdk==2.49.0 +sentry-sdk==2.50.0 # via wandb setuptools==80.9.0 # via # imagecorruptions-imaug # lightning-utilities - # numba # pbr # torch # torchfields diff --git a/tests/unit/layer/test_layer.py b/tests/unit/layer/test_layer.py index f1d46030b..ec4c6702d 100644 --- a/tests/unit/layer/test_layer.py +++ b/tests/unit/layer/test_layer.py @@ -14,3 +14,11 @@ def test_with_procs(mocker): assert layer2.read_procs == (read_proc,) assert layer2.write_procs == (write_proc,) assert layer2.index_procs == (index_proc,) + + +def test_delete(mocker): + backend = mocker.MagicMock() + backend.delete = mocker.MagicMock() + layer = Layer(backend=backend) + layer.delete() + backend.delete.assert_called_once() diff --git a/tests/unit/layer/volumetric/cloudvol/test_backend.py b/tests/unit/layer/volumetric/cloudvol/test_backend.py index c602f4a81..71e14e496 100644 --- a/tests/unit/layer/volumetric/cloudvol/test_backend.py +++ b/tests/unit/layer/volumetric/cloudvol/test_backend.py @@ -594,3 +594,40 @@ def test_cv_lru_resize(clear_caches_reset_mocks): assert cvb.dtype == np.dtype("uint8") cvc = CVBackend(path=LAYER_X0_PATH, info_spec=info_spec, cache_bytes_limit=0) assert cvc.dtype == np.dtype("uint8") + + +def test_cv_with_changes_overwrite_partial_chunks(clear_caches_reset_mocks): + info_spec = PrecomputedInfoSpec( + info_spec_params=InfoSpecParams.from_optional_reference( + reference_path=LAYER_X0_PATH, + scales=[[1, 1, 1]], + chunk_size=[1024, 1024, 1], + inherit_all_params=True, + ) + ) + with tempfile.TemporaryDirectory() as tmp_dir: + cvb = CVBackend(path=tmp_dir, info_spec=info_spec, info_overwrite=True) + assert not cvb.overwrite_partial_chunks + with tempfile.TemporaryDirectory() as tmp_dir2: + cvb_new = cvb.with_changes( + name=tmp_dir2, + overwrite_partial_chunks=True, + ) + assert cvb_new.overwrite_partial_chunks + + +def test_cv_delete(clear_caches_reset_mocks, mocker): + info_spec = PrecomputedInfoSpec( + info_spec_params=InfoSpecParams.from_optional_reference( + reference_path=LAYER_X0_PATH, + scales=[[1, 1, 1]], + inherit_all_params=True, + ) + ) + with tempfile.TemporaryDirectory() as tmp_dir: + cvb = CVBackend(path=tmp_dir, info_spec=info_spec, info_overwrite=True) + # Verify info file exists + assert os.path.exists(os.path.join(tmp_dir, "info")) + # Delete and verify files are removed + cvb.delete() + assert not os.path.exists(os.path.join(tmp_dir, "info")) diff --git a/tests/unit/layer/volumetric/tensorstore/test_backend.py b/tests/unit/layer/volumetric/tensorstore/test_backend.py index ff3dc04bf..8feb802c2 100644 --- a/tests/unit/layer/volumetric/tensorstore/test_backend.py +++ b/tests/unit/layer/volumetric/tensorstore/test_backend.py @@ -1,6 +1,7 @@ # pylint: disable=missing-docstring,redefined-outer-name,unused-argument,pointless-statement,line-too-long,protected-access,too-few-public-methods import os import pathlib +import tempfile from copy import deepcopy import numpy as np @@ -506,3 +507,47 @@ def test_ts_assert_idx_is_chunk_crop_aligned_exc(clear_caches_reset_mocks): with pytest.raises(ValueError): tsb.assert_idx_is_chunk_aligned(index) + + +def test_ts_overwrite_partial_chunks_exc(clear_caches_reset_mocks): + info_spec = PrecomputedInfoSpec( + info_path=LAYER_X0_PATH, + ) + with pytest.raises(NotImplementedError): + TSBackend( + path=LAYER_SCRATCH0_PATH, + info_spec=info_spec, + info_overwrite=True, + overwrite_partial_chunks=True, + ) + + +def test_ts_with_changes_overwrite_partial_chunks_exc(clear_caches_reset_mocks): + info_spec = PrecomputedInfoSpec( + info_spec_params=InfoSpecParams.from_optional_reference( + reference_path=LAYER_X0_PATH, + scales=[[1, 1, 1]], + chunk_size=[1024, 1024, 1], + inherit_all_params=True, + ) + ) + tsb = TSBackend(path=LAYER_SCRATCH0_PATH, info_spec=info_spec, info_overwrite=True) + with pytest.raises(NotImplementedError): + tsb.with_changes(overwrite_partial_chunks=True) + + +def test_ts_delete(clear_caches_reset_mocks): + info_spec = PrecomputedInfoSpec( + info_spec_params=InfoSpecParams.from_optional_reference( + reference_path=LAYER_X0_PATH, + scales=[[1, 1, 1]], + inherit_all_params=True, + ) + ) + with tempfile.TemporaryDirectory() as tmp_dir: + tsb = TSBackend(path=tmp_dir, info_spec=info_spec, info_overwrite=True) + # Verify info file exists + assert os.path.exists(os.path.join(tmp_dir, "info")) + # Delete and verify files are removed + tsb.delete() + assert not os.path.exists(os.path.join(tmp_dir, "info")) diff --git a/tests/unit/layer/volumetric/test_layer.py b/tests/unit/layer/volumetric/test_layer.py index b7cd90c69..f8b06b385 100644 --- a/tests/unit/layer/volumetric/test_layer.py +++ b/tests/unit/layer/volumetric/test_layer.py @@ -147,3 +147,17 @@ def test_write_readonly_exc(mocker): with pytest.raises(IOError): layer[0:1, 0:1, 0:1] = 1 + + +def test_delete(mocker): + backend = mocker.MagicMock() + backend.delete = mocker.MagicMock() + + layer = build_volumetric_layer( + backend, + index_resolution=Vec3D(1, 1, 1), + default_desired_resolution=Vec3D(1, 1, 1), + ) + + layer.delete() + backend.delete.assert_called_once() diff --git a/tests/unit/layer/volumetric/test_layer_set.py b/tests/unit/layer/volumetric/test_layer_set.py index f362e58c4..e049170fc 100644 --- a/tests/unit/layer/volumetric/test_layer_set.py +++ b/tests/unit/layer/volumetric/test_layer_set.py @@ -27,3 +27,14 @@ def test_write(mocker): layer_set[idx] = {"a": 1, "b": 2} layer_a.write_with_procs.assert_called_with(idx, 1) layer_b.write_with_procs.assert_called_with(idx, 2) + + +def test_delete(mocker): + layer_a = mocker.MagicMock() + layer_a.delete = mocker.MagicMock() + layer_b = mocker.MagicMock() + layer_b.delete = mocker.MagicMock() + layer_set = build_volumetric_layer_set(layers={"a": layer_a, "b": layer_b}) + layer_set.delete() + layer_a.delete.assert_called_once() + layer_b.delete.assert_called_once() diff --git a/tests/unit/mazepa/test_upkeep_handlers.py b/tests/unit/mazepa/test_upkeep_handlers.py new file mode 100644 index 000000000..14c3aac72 --- /dev/null +++ b/tests/unit/mazepa/test_upkeep_handlers.py @@ -0,0 +1,463 @@ +# pylint: disable=redefined-outer-name +import multiprocessing +import threading +import time +from typing import Any + +from zetta_utils.common.partial import ComparablePartial +from zetta_utils.mazepa.upkeep_handlers import ( + SQSUpkeepHandlerManager, + UpkeepCommand, + extract_sqs_metadata, + perform_direct_upkeep, + run_sqs_upkeep_handler, +) + + +class TestPerformDirectUpkeep: + def test_calls_extend_lease_fn(self, mocker): + mock_extend_lease = mocker.MagicMock() + task_start_time = time.time() + + perform_direct_upkeep( + extend_lease_fn=mock_extend_lease, + extend_duration=30, + task_start_time=task_start_time, + ) + + mock_extend_lease.assert_called_once_with(30) + + +class TestExtractSqsMetadata: + def test_non_comparable_partial_returns_none(self): + def regular_fn(): + pass + + result = extract_sqs_metadata(regular_fn) + assert result is None + + def test_comparable_partial_without_msg_returns_none(self): + def some_fn(): + pass + + partial = ComparablePartial(some_fn, other_kwarg="value") + result = extract_sqs_metadata(partial) + assert result is None + + def test_msg_missing_required_attributes_returns_none(self): + def some_fn(): + pass + + class IncompleteMsg: + receipt_handle = "handle123" + # Missing queue_name and region_name + + partial = ComparablePartial(some_fn, msg=IncompleteMsg()) + result = extract_sqs_metadata(partial) + assert result is None + + def test_valid_sqs_metadata_extracted(self): + def some_fn(): + pass + + class SQSMsg: + receipt_handle = "handle123" + queue_name = "test-queue" + region_name = "us-east-1" + endpoint_url = "http://localhost:9324" + + partial = ComparablePartial(some_fn, msg=SQSMsg()) + result = extract_sqs_metadata(partial) + + assert result == { + "receipt_handle": "handle123", + "queue_name": "test-queue", + "region_name": "us-east-1", + "endpoint_url": "http://localhost:9324", + } + + def test_valid_sqs_metadata_with_none_endpoint_url(self): + def some_fn(): + pass + + class SQSMsg: + receipt_handle = "handle456" + queue_name = "prod-queue" + region_name = "us-west-2" + # No endpoint_url attribute + + partial = ComparablePartial(some_fn, msg=SQSMsg()) + result = extract_sqs_metadata(partial) + + assert result == { + "receipt_handle": "handle456", + "queue_name": "prod-queue", + "region_name": "us-west-2", + "endpoint_url": None, + } + + +class TestUpkeepCommand: + def test_dataclass_fields(self): + cmd = UpkeepCommand( + action="start_upkeep", + task_id="task-123", + receipt_handle="handle-456", + visibility_timeout=60, + interval_sec=10.0, + queue_name="my-queue", + region_name="us-east-1", + endpoint_url="http://localhost:9324", + ) + + assert cmd.action == "start_upkeep" + assert cmd.task_id == "task-123" + assert cmd.receipt_handle == "handle-456" + assert cmd.visibility_timeout == 60 + assert cmd.interval_sec == 10.0 + assert cmd.queue_name == "my-queue" + assert cmd.region_name == "us-east-1" + assert cmd.endpoint_url == "http://localhost:9324" + + def test_dataclass_defaults(self): + cmd = UpkeepCommand(action="shutdown") + + assert cmd.action == "shutdown" + assert cmd.task_id is None + assert cmd.receipt_handle is None + assert cmd.visibility_timeout is None + assert cmd.interval_sec is None + assert cmd.queue_name is None + assert cmd.region_name is None + assert cmd.endpoint_url is None + + +class TestSQSUpkeepHandlerManager: + def test_start_creates_process(self, mocker): + mocker.patch("zetta_utils.mazepa.upkeep_handlers.multiprocessing.Queue") + mock_process_cls = mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Process" + ) + mock_process = mocker.MagicMock() + mock_process_cls.return_value = mock_process + + manager = SQSUpkeepHandlerManager() + manager.start() + + mock_process_cls.assert_called_once() + mock_process.start.assert_called_once() + + def test_start_when_already_running_is_noop(self, mocker): + mock_queue = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Queue", return_value=mock_queue + ) + mock_process_cls = mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Process" + ) + mock_process = mocker.MagicMock() + mock_process_cls.return_value = mock_process + + manager = SQSUpkeepHandlerManager() + manager.start() + manager.start() # Second call should be no-op + + assert mock_process_cls.call_count == 1 + + def test_shutdown_when_not_running_is_noop(self): + manager = SQSUpkeepHandlerManager() + manager.shutdown() # Should not raise + + def test_shutdown_sends_command_and_joins(self, mocker): + mock_queue = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Queue", return_value=mock_queue + ) + mock_process = mocker.MagicMock() + mock_process.is_alive.return_value = False + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Process", + return_value=mock_process, + ) + + manager = SQSUpkeepHandlerManager() + manager.start() + manager.shutdown() + + # Verify shutdown command was sent + mock_queue.put.assert_called_once() + cmd = mock_queue.put.call_args[0][0] + assert cmd.action == "shutdown" + + # Verify process was joined + mock_process.join.assert_called() + + def test_shutdown_terminates_if_process_alive(self, mocker): + mock_queue = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Queue", return_value=mock_queue + ) + mock_process = mocker.MagicMock() + mock_process.is_alive.return_value = True + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Process", + return_value=mock_process, + ) + + manager = SQSUpkeepHandlerManager() + manager.start() + manager.shutdown(timeout=0.1) + + mock_process.terminate.assert_called_once() + + def test_start_upkeep_when_not_running_logs_warning(self, mocker): + mock_logger = mocker.patch("zetta_utils.mazepa.upkeep_handlers.logger") + + manager = SQSUpkeepHandlerManager() + manager.start_upkeep( + task_id="task-123", + receipt_handle="handle", + visibility_timeout=60, + interval_sec=10.0, + queue_name="queue", + region_name="us-east-1", + ) + + mock_logger.warning.assert_called_once() + assert "not running" in mock_logger.warning.call_args[0][0] + + def test_stop_upkeep_when_not_running_logs_warning(self, mocker): + mock_logger = mocker.patch("zetta_utils.mazepa.upkeep_handlers.logger") + + manager = SQSUpkeepHandlerManager() + manager.stop_upkeep("task-123") + + mock_logger.warning.assert_called_once() + assert "not running" in mock_logger.warning.call_args[0][0] + + def test_start_upkeep_sends_command(self, mocker): + mock_queue = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Queue", return_value=mock_queue + ) + mock_process = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Process", + return_value=mock_process, + ) + + manager = SQSUpkeepHandlerManager() + manager.start() + manager.start_upkeep( + task_id="task-123", + receipt_handle="handle-456", + visibility_timeout=60, + interval_sec=10.0, + queue_name="my-queue", + region_name="us-east-1", + endpoint_url="http://localhost:9324", + ) + + mock_queue.put_nowait.assert_called_once() + cmd = mock_queue.put_nowait.call_args[0][0] + assert cmd.action == "start_upkeep" + assert cmd.task_id == "task-123" + assert cmd.receipt_handle == "handle-456" + assert cmd.visibility_timeout == 60 + assert cmd.interval_sec == 10.0 + assert cmd.queue_name == "my-queue" + assert cmd.region_name == "us-east-1" + assert cmd.endpoint_url == "http://localhost:9324" + + def test_stop_upkeep_sends_command(self, mocker): + mock_queue = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Queue", return_value=mock_queue + ) + mock_process = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Process", + return_value=mock_process, + ) + + manager = SQSUpkeepHandlerManager() + manager.start() + manager.stop_upkeep("task-123") + + mock_queue.put_nowait.assert_called_once() + cmd = mock_queue.put_nowait.call_args[0][0] + assert cmd.action == "stop_upkeep" + assert cmd.task_id == "task-123" + + def test_is_running_false_before_start(self): + manager = SQSUpkeepHandlerManager() + assert not manager.is_running + + def test_is_running_true_after_start(self, mocker): + mock_queue = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Queue", return_value=mock_queue + ) + mock_process = mocker.MagicMock() + mock_process.is_alive.return_value = True + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Process", + return_value=mock_process, + ) + + manager = SQSUpkeepHandlerManager() + manager.start() + assert manager.is_running + + def test_is_running_false_when_process_dead(self, mocker): + mock_queue = mocker.MagicMock() + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Queue", return_value=mock_queue + ) + mock_process = mocker.MagicMock() + mock_process.is_alive.return_value = False + mocker.patch( + "zetta_utils.mazepa.upkeep_handlers.multiprocessing.Process", + return_value=mock_process, + ) + + manager = SQSUpkeepHandlerManager() + manager.start() + assert not manager.is_running + + +class TestRunSqsUpkeepHandler: + def test_shutdown_command_exits_loop(self, mocker): + # worker_init is imported inside the function, so patch at source + mocker.patch("zetta_utils.mazepa.worker.worker_init") + + command_queue: multiprocessing.Queue[Any] = multiprocessing.Queue() + command_queue.put(UpkeepCommand(action="shutdown")) + + # Run in thread so we can verify it exits + handler_thread = threading.Thread( + target=run_sqs_upkeep_handler, + args=(command_queue, "INFO"), + ) + handler_thread.start() + handler_thread.join(timeout=5.0) + + assert not handler_thread.is_alive() + + def test_start_and_stop_upkeep(self, mocker): + mocker.patch("zetta_utils.mazepa.worker.worker_init") + mock_change_visibility = mocker.patch( + "zetta_utils.message_queues.sqs.utils.change_message_visibility" + ) + + command_queue: multiprocessing.Queue[Any] = multiprocessing.Queue() + + # Start upkeep with short interval + command_queue.put( + UpkeepCommand( + action="start_upkeep", + task_id="task-123", + receipt_handle="handle-456", + visibility_timeout=60, + interval_sec=0.1, + queue_name="my-queue", + region_name="us-east-1", + endpoint_url=None, + ) + ) + + handler_thread = threading.Thread( + target=run_sqs_upkeep_handler, + args=(command_queue, "INFO"), + ) + handler_thread.start() + + # Wait for upkeep to fire at least once + time.sleep(0.3) + + # Stop upkeep and shutdown + command_queue.put(UpkeepCommand(action="stop_upkeep", task_id="task-123")) + command_queue.put(UpkeepCommand(action="shutdown")) + + handler_thread.join(timeout=5.0) + assert not handler_thread.is_alive() + + # Verify change_message_visibility was called + assert mock_change_visibility.call_count >= 1 + mock_change_visibility.assert_called_with( + receipt_handle="handle-456", + visibility_timeout=60, + queue_name="my-queue", + region_name="us-east-1", + endpoint_url=None, + ) + + def test_duplicate_start_upkeep_ignored(self, mocker): + mocker.patch("zetta_utils.mazepa.worker.worker_init") + mocker.patch("zetta_utils.message_queues.sqs.utils.change_message_visibility") + + command_queue: multiprocessing.Queue[Any] = multiprocessing.Queue() + + # Start upkeep twice for same task + for _ in range(2): + command_queue.put( + UpkeepCommand( + action="start_upkeep", + task_id="task-123", + receipt_handle="handle-456", + visibility_timeout=60, + interval_sec=1.0, + queue_name="my-queue", + region_name="us-east-1", + endpoint_url=None, + ) + ) + + command_queue.put(UpkeepCommand(action="shutdown")) + + handler_thread = threading.Thread( + target=run_sqs_upkeep_handler, + args=(command_queue, "INFO"), + ) + handler_thread.start() + handler_thread.join(timeout=5.0) + + # Test passes if no crash - duplicate is logged and ignored + assert not handler_thread.is_alive() + + def test_stop_nonexistent_upkeep_ignored(self, mocker): + mocker.patch("zetta_utils.mazepa.worker.worker_init") + + command_queue: multiprocessing.Queue[Any] = multiprocessing.Queue() + + # Stop upkeep for task that doesn't exist + command_queue.put(UpkeepCommand(action="stop_upkeep", task_id="nonexistent-task")) + command_queue.put(UpkeepCommand(action="shutdown")) + + handler_thread = threading.Thread( + target=run_sqs_upkeep_handler, + args=(command_queue, "INFO"), + ) + handler_thread.start() + handler_thread.join(timeout=5.0) + + # Test passes if no crash - nonexistent stop is logged and ignored + assert not handler_thread.is_alive() + + def test_unknown_action_logged(self, mocker): + mocker.patch("zetta_utils.mazepa.worker.worker_init") + + command_queue: multiprocessing.Queue[Any] = multiprocessing.Queue() + + command_queue.put(UpkeepCommand(action="unknown_action")) + command_queue.put(UpkeepCommand(action="shutdown")) + + handler_thread = threading.Thread( + target=run_sqs_upkeep_handler, + args=(command_queue, "INFO"), + ) + handler_thread.start() + handler_thread.join(timeout=5.0) + + # Test passes if no crash - unknown action is logged and ignored + assert not handler_thread.is_alive() diff --git a/tests/unit/mazepa/test_worker.py b/tests/unit/mazepa/test_worker.py index 379414e84..be8660e70 100644 --- a/tests/unit/mazepa/test_worker.py +++ b/tests/unit/mazepa/test_worker.py @@ -5,7 +5,7 @@ from zetta_utils.mazepa.pool_activity import PoolActivityTracker from zetta_utils.mazepa.tasks import _TaskableOperation -from zetta_utils.mazepa.worker import run_worker +from zetta_utils.mazepa.worker import run_worker, worker_init from zetta_utils.message_queues.base import ReceivedMessage @@ -235,3 +235,86 @@ def simple_task(): assert result == "max_runtime_exceeded" assert len(outcome_queue.pushed_messages) == 1 + + +def test_worker_init_basic(mocker): + mock_reset_signals = mocker.patch("zetta_utils.mazepa.worker.reset_signal_handlers") + mock_configure_logger = mocker.patch("zetta_utils.mazepa.worker.log.configure_logger") + + worker_init(log_level="INFO") + + mock_reset_signals.assert_called_once() + mock_configure_logger.assert_called_once_with(level="INFO", force=True) + + +def test_worker_init_suppress_logs(mocker): + mock_reset_signals = mocker.patch("zetta_utils.mazepa.worker.reset_signal_handlers") + mock_redirect_buffers = mocker.patch("zetta_utils.mazepa.worker.redirect_buffers") + mock_configure_logger = mocker.patch("zetta_utils.mazepa.worker.log.configure_logger") + + worker_init(log_level="INFO", suppress_logs=True) + + mock_reset_signals.assert_called_once() + mock_redirect_buffers.assert_called_once() + mock_configure_logger.assert_not_called() + + +def test_worker_init_set_start_method(mocker): + mocker.patch("zetta_utils.mazepa.worker.reset_signal_handlers") + mocker.patch("zetta_utils.mazepa.worker.log.configure_logger") + mock_set_start_method = mocker.patch( + "zetta_utils.mazepa.worker.multiprocessing.set_start_method" + ) + + worker_init(log_level="DEBUG", set_start_method=True, multiprocessing_start_method="spawn") + + mock_set_start_method.assert_called_once_with("spawn", force=True) + + +def test_worker_init_load_train_inference(mocker): + mocker.patch("zetta_utils.mazepa.worker.reset_signal_handlers") + mocker.patch("zetta_utils.mazepa.worker.log.configure_logger") + mock_load = mocker.patch("zetta_utils.mazepa.worker.try_load_train_inference") + + worker_init(log_level="INFO", load_train_inference=True) + + mock_load.assert_called_once() + + +def test_worker_upkeep_fallback_for_non_sqs_queue( + mock_queues, mocker +): # pylint: disable=redefined-outer-name + """Test that non-SQS queues use the RepeatTimer fallback for upkeep.""" + task_queue, outcome_queue = mock_queues + + # Track if perform_direct_upkeep was called + mock_perform_upkeep = mocker.patch( + "zetta_utils.mazepa.worker.perform_direct_upkeep", wraps=lambda *args: None + ) + + def slow_task(): + time.sleep(0.2) + return "done" + + # Create task with short upkeep interval so timer fires during execution + task = _TaskableOperation(slow_task, upkeep_interval_sec=0.05).make_task() + msg = ReceivedMessage( + payload=task, + acknowledge_fn=MagicMock(), + extend_lease_fn=MagicMock(), # Not ComparablePartial, so uses fallback + approx_receive_count=1, + ) + task_queue.messages = [msg] + + result = run_worker( + task_queue=task_queue, + outcome_queue=outcome_queue, + sleep_sec=0.1, + max_runtime=1.0, + debug=True, + ) + + assert result == "max_runtime_exceeded" + assert len(outcome_queue.pushed_messages) == 1 + # Verify the fallback upkeep path was used (timer should have fired at least once) + assert mock_perform_upkeep.call_count >= 1 diff --git a/zetta_utils/cloud_management/resource_allocation/k8s/gke.py b/zetta_utils/cloud_management/resource_allocation/k8s/gke.py index b990edc69..868d2ab78 100644 --- a/zetta_utils/cloud_management/resource_allocation/k8s/gke.py +++ b/zetta_utils/cloud_management/resource_allocation/k8s/gke.py @@ -23,4 +23,5 @@ def gke_cluster_data(name: str, region: str, project: str) -> Tuple[Cluster, str # creds, _ = google_auth.default() creds, _ = google_auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"]) creds.refresh(google_auth.transport.requests.Request()) + assert creds.token is not None return cluster_data, cert_name, creds.token diff --git a/zetta_utils/layer/volumetric/seg_contact/backend.py b/zetta_utils/layer/volumetric/seg_contact/backend.py index 2bc1f3d97..7f4a9d969 100644 --- a/zetta_utils/layer/volumetric/seg_contact/backend.py +++ b/zetta_utils/layer/volumetric/seg_contact/backend.py @@ -18,7 +18,9 @@ @attrs.define -class SegContactLayerBackend(Backend[VolumetricIndex, Sequence[SegContact], Sequence[SegContact]]): +class SegContactLayerBackend( + Backend[VolumetricIndex, Sequence[SegContact], Sequence[SegContact]] +): # pylint: disable=too-many-public-methods """Backend for reading/writing seg_contact data in chunked format.""" path: str diff --git a/zetta_utils/layer/volumetric/tensorstore/backend.py b/zetta_utils/layer/volumetric/tensorstore/backend.py index 6007ebaa1..1bc0304c7 100644 --- a/zetta_utils/layer/volumetric/tensorstore/backend.py +++ b/zetta_utils/layer/volumetric/tensorstore/backend.py @@ -274,7 +274,8 @@ def with_changes(self, **kwargs) -> TSBackend: "allow_cache" = value: Union[bool, str] - must be False for TensorStoreBackend, ignored "use_compression" = Value: bool - must be False for TensorStoreBackend, ignored "enforce_chunk_aligned_writes" = value: bool - "overwrite_partial_chunks" = value: bool - must be False for TensorStoreBackend, raises NotImplementedError if True + "overwrite_partial_chunks" = value: bool - must be False for TensorStoreBackend, + raises NotImplementedError if True "voxel_offset_res" = (voxel_offset, resolution): Tuple[Vec3D[int], Vec3D] "chunk_size_res" = (chunk_size, resolution): Tuple[Vec3D[int], Vec3D] "dataset_size_res" = (dataset_size, resolution): Tuple[Vec3D[int], Vec3D] @@ -303,8 +304,8 @@ def with_changes(self, **kwargs) -> TSBackend: keys_to_kwargs = { "name": "path", - "enforce_chunk_aligned_writes": "_enforce_chunk_aligned_writes", - "overwrite_partial_chunks": "_overwrite_partial_chunks", + "enforce_chunk_aligned_writes": "enforce_chunk_aligned_writes", + "overwrite_partial_chunks": "overwrite_partial_chunks", } evolve_kwargs = {} diff --git a/zetta_utils/mazepa/upkeep_handlers.py b/zetta_utils/mazepa/upkeep_handlers.py index a3f96cfcb..702847e05 100644 --- a/zetta_utils/mazepa/upkeep_handlers.py +++ b/zetta_utils/mazepa/upkeep_handlers.py @@ -170,7 +170,7 @@ def _run_upkeep_loop( f"SQS_HANDLER: [{task_id}] [T+{elapsed:.1f}s] Successfully extended " f"(API took {api_duration:.1f}s)" ) - except Exception as e: # pylint: disable=broad-except + except Exception as e: # pylint: disable=broad-except # pragma: no cover logger.error( f"SQS_HANDLER: [{task_id}] Failed to extend visibility: " f"{type(e).__name__}: {e}" @@ -248,7 +248,7 @@ def _run_upkeep_loop( else: logger.warning(f"SQS_HANDLER: Unknown action: {cmd.action}") - except Exception as e: # pylint: disable=broad-except + except Exception as e: # pylint: disable=broad-except # pragma: no cover logger.error(f"SQS_HANDLER: Error processing command: {type(e).__name__}: {e}") logger.info("SQS_HANDLER: Handler process exiting") diff --git a/zetta_utils/mazepa/worker.py b/zetta_utils/mazepa/worker.py index 67611a7fd..327327f7a 100644 --- a/zetta_utils/mazepa/worker.py +++ b/zetta_utils/mazepa/worker.py @@ -39,7 +39,7 @@ def flush(self): pass -def redirect_buffers() -> None: +def redirect_buffers() -> None: # pragma: no cover sys.stdin = DummyBuffer() # type: ignore sys.stdout = DummyBuffer() # type: ignore sys.stderr = DummyBuffer() # type: ignore diff --git a/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py b/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py index f5652d9e8..1957d1cf7 100644 --- a/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py +++ b/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py @@ -172,7 +172,8 @@ def __call__( if not is_floating_point_dtype(dst.backend.dtype): # Temporarily convert integer cutout to float for rounding res[tuple(subidx_channels)] = ( - res[tuple(subidx_channels)] + layer[intscn].astype(float) * weight.numpy() + res[tuple(subidx_channels)] + + layer[intscn].astype(float) * weight.numpy() ) else: res[tuple(subidx_channels)] = ( @@ -420,7 +421,8 @@ def __attrs_post_init__(self): # pylint: disable=too-many-branches raise TypeError( f"`task_stack_size` was provided ({self.task_stack_size}), but the operation " f"{type(self.op).__name__} does not implement StackableVolumetricOpProtocol. " - f"Operations must have `read`, `write`, and `processing_fn` methods to support stacking." + "Operations must have `read`, `write`, and `processing_fn` methods " + "to support stacking." ) if self.roi_crop_pad is None: diff --git a/zetta_utils/message_queues/sqs/utils.py b/zetta_utils/message_queues/sqs/utils.py index 5806865a1..8e43b1938 100644 --- a/zetta_utils/message_queues/sqs/utils.py +++ b/zetta_utils/message_queues/sqs/utils.py @@ -112,7 +112,7 @@ def delete_msg_by_receipt_handle( ReceiptHandle=receipt_handle, ) logger.debug(f"DELETE: Successfully deleted message from queue '{queue_name}'") - except Exception as e: + except Exception as e: # pragma: no cover logger.error( f"DELETE: Failed to delete message from queue '{queue_name}': {type(e).__name__}: {e}" ) @@ -142,7 +142,7 @@ def change_message_visibility( f"VISIBILITY: Successfully changed visibility to {visibility_timeout}s " f"for queue '{queue_name}'" ) - except Exception as e: + except Exception as e: # pragma: no cover logger.error( f"VISIBILITY: Failed to change visibility for queue '{queue_name}' " f"to {visibility_timeout}s: {type(e).__name__}: {e}" From b456cbf62feb464c6e8ea2f5b5562364897c902a Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Thu, 22 Jan 2026 14:43:26 -0800 Subject: [PATCH 7/9] chore: update to backports-zstd for latest fsspec --- pyproject.toml | 2 +- requirements.all.txt | 39 ++++++++++++++++++++++----------------- requirements.modules.txt | 37 +++++++++++++++++++++---------------- 3 files changed, 44 insertions(+), 34 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 39c5b9795..cbd96a814 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,7 +93,7 @@ chunkedgraph = [ "messagingclient", "dracopy>=1.3.0", "datastoreflex>=0.5.0", - "zstandard>=0.21.0", + "backports-zstd", # Note: PyChunkedGraph must be installed separately via install script with --pcg flag # Note: graph_tool must be installed separately via conda: # conda install -c conda-forge graph-tool-base diff --git a/requirements.all.txt b/requirements.all.txt index 72951a8d7..aa165bdf5 100644 --- a/requirements.all.txt +++ b/requirements.all.txt @@ -38,10 +38,12 @@ attrs==25.4.0 # referencing autodocsumm==0.2.14 # via sphinx-toolbox -awscli==1.44.21 +awscli==1.44.23 # via zetta-utils (pyproject.toml) babel==2.17.0 # via sphinx +backports-zstd==1.3.0 + # via zetta-utils (pyproject.toml) beautifulsoup4==4.14.3 # via sphinx-toolbox black==21.9b0 @@ -53,14 +55,14 @@ blosc==1.11.4 # meshparty blosc2==3.12.2 # via tables -boto3==1.42.31 +boto3==1.42.33 # via # zetta-utils (pyproject.toml) # cloud-files # cloud-volume # moto # task-queue -botocore==1.42.31 +botocore==1.42.33 # via # awscli # boto3 @@ -155,6 +157,10 @@ cryptography==46.0.3 # via moto cssutils==2.11.1 # via dict2css +cuda-bindings==12.9.4 + # via torch +cuda-pathfinder==1.3.3 + # via cuda-bindings cycler==0.12.1 # via matplotlib cython==3.2.4 @@ -468,7 +474,7 @@ jinja2==3.1.6 # sphinx-jinja2-compat # sphinx-prompt # torch -jmespath==1.0.1 +jmespath==1.1.0 # via # boto3 # botocore @@ -521,7 +527,7 @@ mapbox-earcut==2.0.0 # via trimesh mapbuffer==1.1.0 # via zetta-utils (pyproject.toml) -markdown==3.10 +markdown==3.10.1 # via python-jsonschema-objects markdown-it-py==4.0.0 # via rich @@ -702,7 +708,7 @@ nvidia-cusparse-cu12==12.5.8.93 # torch nvidia-cusparselt-cu12==0.7.1 # via torch -nvidia-ml-py==13.590.44 +nvidia-ml-py==13.590.48 # via zetta-utils (pyproject.toml) nvidia-nccl-cu12==2.27.5 # via torch @@ -712,7 +718,7 @@ nvidia-nvjitlink-cu12==12.8.93 # nvidia-cusolver-cu12 # nvidia-cusparse-cu12 # torch -nvidia-nvshmem-cu12==3.3.20 +nvidia-nvshmem-cu12==3.4.5 # via torch nvidia-nvtx-cu12==12.8.90 # via torch @@ -755,7 +761,7 @@ osteoid==0.6.0 # via # cloud-volume # kimimaro -packaging==25.0 +packaging==26.0 # via # zetta-utils (pyproject.toml) # caveclient @@ -899,7 +905,7 @@ pybind11==3.0.1 # osteoid pycollada==0.9.2 # via trimesh -pycparser==2.23 +pycparser==3.0 # via cffi pydantic==2.12.5 # via wandb @@ -921,7 +927,7 @@ pygments==2.19.2 # sphinx-tabs pylint==4.0.4 # via zetta-utils (pyproject.toml) -pyparsing==3.3.1 +pyparsing==3.3.2 # via # httplib2 # matplotlib @@ -1063,7 +1069,7 @@ scipy==1.17.0 # trimesh sentry-sdk==2.50.0 # via wandb -setuptools==80.9.0 +setuptools==80.10.1 # via # imagecorruptions-imaug # lightning-utilities @@ -1157,7 +1163,7 @@ sphinxcontrib-serializinghtml==2.0.0 # via # zetta-utils (pyproject.toml) # sphinx -sqlalchemy==2.0.45 +sqlalchemy==2.0.46 # via zetta-utils (pyproject.toml) sqlitedict==2.1.0 # via pcg-skel @@ -1199,7 +1205,7 @@ tomli==1.2.3 # via black tomlkit==0.14.0 # via pylint -torch==2.9.1 +torch==2.10.0 # via # zetta-utils (pyproject.toml) # kornia @@ -1216,7 +1222,7 @@ torchmetrics==1.8.2 # zetta-utils (pyproject.toml) # lightning # pytorch-lightning -torchvision==0.24.1 +torchvision==0.25.0 # via # zetta-utils (pyproject.toml) # onnx2torch @@ -1243,7 +1249,7 @@ trimesh==4.11.1 # zetta-utils (pyproject.toml) # meshparty # pcg-skel -triton==3.5.1 +triton==3.6.0 # via torch typeguard==4.4.4 # via @@ -1314,7 +1320,7 @@ wandb==0.24.0 # via zetta-utils (pyproject.toml) waterz @ git+https://github.com/ZettaAI/waterz.git@0bac4be4c864b293fa49f11cf0abf52bd1ec81df # via zetta-utils (pyproject.toml) -wcwidth==0.2.14 +wcwidth==0.3.1 # via prompt-toolkit webencodings==0.5.1 # via html5lib @@ -1354,7 +1360,6 @@ zope-interface==8.2 # via gevent zstandard==0.25.0 # via - # zetta-utils (pyproject.toml) # cloud-files # cloud-volume # mapbuffer diff --git a/requirements.modules.txt b/requirements.modules.txt index 00334a4ac..d910948ed 100644 --- a/requirements.modules.txt +++ b/requirements.modules.txt @@ -28,7 +28,9 @@ attrs==25.4.0 # caveclient # jsonschema # referencing -awscli==1.44.21 +awscli==1.44.23 + # via zetta-utils (pyproject.toml) +backports-zstd==1.3.0 # via zetta-utils (pyproject.toml) blinker==1.9.0 # via flask @@ -39,13 +41,13 @@ blosc==1.11.4 # meshparty blosc2==3.12.2 # via tables -boto3==1.42.31 +boto3==1.42.33 # via # zetta-utils (pyproject.toml) # cloud-files # cloud-volume # task-queue -botocore==1.42.31 +botocore==1.42.33 # via # awscli # boto3 @@ -125,6 +127,10 @@ crc32c==2.8 # via # cloud-files # mapbuffer +cuda-bindings==12.9.4 + # via torch +cuda-pathfinder==1.3.3 + # via cuda-bindings cycler==0.12.1 # via matplotlib cython==3.2.4 @@ -398,7 +404,7 @@ jinja2==3.1.6 # via # flask # torch -jmespath==1.0.1 +jmespath==1.1.0 # via # boto3 # botocore @@ -451,7 +457,7 @@ mapbox-earcut==2.0.0 # via trimesh mapbuffer==1.1.0 # via zetta-utils (pyproject.toml) -markdown==3.10 +markdown==3.10.1 # via python-jsonschema-objects markdown-it-py==4.0.0 # via rich @@ -613,7 +619,7 @@ nvidia-cusparse-cu12==12.5.8.93 # torch nvidia-cusparselt-cu12==0.7.1 # via torch -nvidia-ml-py==13.590.44 +nvidia-ml-py==13.590.48 # via zetta-utils (pyproject.toml) nvidia-nccl-cu12==2.27.5 # via torch @@ -623,7 +629,7 @@ nvidia-nvjitlink-cu12==12.8.93 # nvidia-cusolver-cu12 # nvidia-cusparse-cu12 # torch -nvidia-nvshmem-cu12==3.3.20 +nvidia-nvshmem-cu12==3.4.5 # via torch nvidia-nvtx-cu12==12.8.90 # via torch @@ -666,7 +672,7 @@ osteoid==0.6.0 # via # cloud-volume # kimimaro -packaging==25.0 +packaging==26.0 # via # zetta-utils (pyproject.toml) # caveclient @@ -810,7 +816,7 @@ pygments==2.19.2 # pdbp # pytest # rich -pyparsing==3.3.1 +pyparsing==3.3.2 # via # httplib2 # matplotlib @@ -923,7 +929,7 @@ scipy==1.17.0 # trimesh sentry-sdk==2.50.0 # via wandb -setuptools==80.9.0 +setuptools==80.10.1 # via # imagecorruptions-imaug # lightning-utilities @@ -957,7 +963,7 @@ smmap==5.0.2 # via gitdb sortedcontainers==2.4.0 # via intervaltree -sqlalchemy==2.0.45 +sqlalchemy==2.0.46 # via zetta-utils (pyproject.toml) sqlitedict==2.1.0 # via pcg-skel @@ -989,7 +995,7 @@ tifffile==2026.1.14 # via scikit-image tinybrain==1.7.0 # via zetta-utils (pyproject.toml) -torch==2.9.1 +torch==2.10.0 # via # zetta-utils (pyproject.toml) # kornia @@ -1006,7 +1012,7 @@ torchmetrics==1.8.2 # zetta-utils (pyproject.toml) # lightning # pytorch-lightning -torchvision==0.24.1 +torchvision==0.25.0 # via # zetta-utils (pyproject.toml) # onnx2torch @@ -1033,7 +1039,7 @@ trimesh==4.11.1 # zetta-utils (pyproject.toml) # meshparty # pcg-skel -triton==3.5.1 +triton==3.6.0 # via torch typeguard==4.4.4 # via @@ -1083,7 +1089,7 @@ wandb==0.24.0 # via zetta-utils (pyproject.toml) waterz @ git+https://github.com/ZettaAI/waterz.git@0bac4be4c864b293fa49f11cf0abf52bd1ec81df # via zetta-utils (pyproject.toml) -wcwidth==0.2.14 +wcwidth==0.3.1 # via prompt-toolkit websocket-client==1.9.0 # via kubernetes @@ -1116,7 +1122,6 @@ zope-interface==8.2 # via gevent zstandard==0.25.0 # via - # zetta-utils (pyproject.toml) # cloud-files # cloud-volume # mapbuffer From 3e66577eb31bb008b0f8aa332eef01a3069be67d Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Thu, 22 Jan 2026 20:51:52 -0800 Subject: [PATCH 8/9] chore: PyTorch typing change --- zetta_utils/training/sampler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zetta_utils/training/sampler.py b/zetta_utils/training/sampler.py index 83b9625cf..b92e7a534 100644 --- a/zetta_utils/training/sampler.py +++ b/zetta_utils/training/sampler.py @@ -15,7 +15,7 @@ class SamplerWrapper(torch.utils.data.Sampler[int]): sampler: torch.utils.data.Sampler[int] def __init__(self, sampler: torch.utils.data.Sampler[int]) -> None: - super().__init__(None) + super().__init__() self.sampler = sampler self.epoch = 0 From f4a4f13855e8ac399bfd4af3ce6d023120ba2d04 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 23 Jan 2026 17:39:31 -0800 Subject: [PATCH 9/9] fix: single workers also now start activity tracker for timeout --- .../resource_allocation/k8s/common.py | 4 ++++ zetta_utils/mazepa/worker.py | 17 ++++++++++++++++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/zetta_utils/cloud_management/resource_allocation/k8s/common.py b/zetta_utils/cloud_management/resource_allocation/k8s/common.py index 1af620f1d..8e3f6b032 100644 --- a/zetta_utils/cloud_management/resource_allocation/k8s/common.py +++ b/zetta_utils/cloud_management/resource_allocation/k8s/common.py @@ -75,11 +75,14 @@ def get_mazepa_worker_command( num_procs_line = "" semaphores_line = "" suppress_worker_logs_line = "" + pool_name = f"{task_queue_spec['name']}_{outcome_queue_spec['name']}" + pool_name_line = f"pool_name: {json.dumps(pool_name)}\n" else: command = "mazepa.run_worker_manager" num_procs_line = f"num_procs: {num_procs}\n" semaphores_line = f"semaphores_spec: {json.dumps(semaphores_spec)}\n" suppress_worker_logs_line = f"suppress_worker_logs: {json.dumps(suppress_worker_logs)}\n" + pool_name_line = "" idle_timeout_line = "" if idle_timeout: @@ -101,6 +104,7 @@ def get_mazepa_worker_command( + idle_timeout_line + suppress_worker_logs_line + resource_monitor_interval_line + + pool_name_line + """ sleep_sec: 5 }' diff --git a/zetta_utils/mazepa/worker.py b/zetta_utils/mazepa/worker.py index 327327f7a..73f4dace7 100644 --- a/zetta_utils/mazepa/worker.py +++ b/zetta_utils/mazepa/worker.py @@ -215,10 +215,23 @@ def run_worker( upkeep_handler = SQSUpkeepHandlerManager() upkeep_handler.start() + # For single worker case (k8s), create shared memory for activity tracking + # For multi-worker case, parent already created it - we'll get FileExistsError which is fine + owns_activity_tracker = False + if pool_name and idle_timeout is not None: + activity_tracker = PoolActivityTracker(pool_name) + try: + activity_tracker.create_shared_memory().close() + owns_activity_tracker = True + logger.info(f"Created activity tracker for pool: {pool_name}") + except FileExistsError: + logger.debug(f"Activity tracker already exists for pool: {pool_name}") + else: + activity_tracker = None + try: with monitor_resources(resource_monitor_interval): start_time = time.time() - activity_tracker = PoolActivityTracker(pool_name) if pool_name else None while True: task_msgs = _pull_tasks_with_error_handling( @@ -246,6 +259,8 @@ def run_worker( return reason finally: upkeep_handler.shutdown() + if owns_activity_tracker and activity_tracker: + activity_tracker.unlink() def process_task_message(