tensor-ee 0.1.0__py3-none-any.whl

tensor_ee/tiling.py

@@ -0,0 +1,287 @@
+"""TensorEE inference on aligned context windows with disjoint output cores.
+
+The source is one *finite* global tensor. Each tile is compiled on its larger
+context window, with the window affine translated from the same global grid.
+Context is clamped at the true source boundary, so every intermediate tensor
+has the same finite-padding boundary there as a full-image forward pass. Only
+artificial interior context edges are protected by the caller-supplied halo.
+
+The halo must cover the model's exact receptive support. Alignment must cover
+the deepest striding/pooling lattice. These requirements cannot be inferred
+from a random fixture's numerically small distant effects.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Sequence
+
+import ee
+from torch import nn
+
+from tensor_ee.compiler import (CompiledModel, Grid, Raster, UnsupportedOperation,
+                                WeightProvider, compile_model)
+from tensor_ee.receptive_field import ReceptiveFieldReport, analyze_receptive_field
+
+
+@dataclass(frozen=True)
+class PixelWindow:
+    """A half-open rectangle in integer coordinates of the global input grid."""
+
+    x: int
+    y: int
+    width: int
+    height: int
+
+    def __post_init__(self) -> None:
+        if any(type(value) is not int for value in (self.x, self.y, self.width, self.height)):
+            raise ValueError("PixelWindow coordinates and dimensions must be integers")
+        if self.x < 0 or self.y < 0 or self.width <= 0 or self.height <= 0:
+            raise ValueError("PixelWindow requires nonnegative origins and positive dimensions")
+
+    @property
+    def right(self) -> int:
+        return self.x + self.width
+
+    @property
+    def bottom(self) -> int:
+        return self.y + self.height
+
+    @property
+    def pixels(self) -> int:
+        return self.width * self.height
+
+    def grid(self, global_grid: Grid) -> Grid:
+        if self.right > global_grid.width or self.bottom > global_grid.height:
+            raise ValueError("PixelWindow is outside the global finite input")
+        a, b, c, d, e, f = global_grid.transform
+        return Grid(global_grid.crs,
+                    (a, b, c + a * self.x, d, e, f + e * self.y),
+                    self.width, self.height)
+
+
+@dataclass(frozen=True)
+class Tile:
+    """One output core and its larger finite context, in global pixel indices."""
+
+    index: int
+    row: int
+    column: int
+    core: PixelWindow
+    context: PixelWindow
+
+
+@dataclass(frozen=True)
+class TilePlan:
+    global_grid: Grid
+    roi: PixelWindow
+    core_size: tuple[int, int]
+    halo: tuple[int, int]
+    alignment: int
+    tiles: tuple[Tile, ...]
+
+    @property
+    def roi_grid(self) -> Grid:
+        return self.roi.grid(self.global_grid)
+
+    @property
+    def vertical_seams(self) -> tuple[int, ...]:
+        return tuple(sorted({tile.core.right for tile in self.tiles
+                             if tile.core.right < self.roi.right}))
+
+    @property
+    def horizontal_seams(self) -> tuple[int, ...]:
+        return tuple(sorted({tile.core.bottom for tile in self.tiles
+                             if tile.core.bottom < self.roi.bottom}))
+
+    def seam_pixel_count(self, width: int = 1) -> int:
+        """Count pixels within ``width`` pixels on each side of internal seams."""
+        _positive_integer(width, "seam width")
+
+        def stripe_pixels(seams: tuple[int, ...], start: int, end: int) -> int:
+            intervals = sorted((max(start, seam - width), min(end, seam + width))
+                               for seam in seams)
+            total, previous_end = 0, start
+            for left, right in intervals:
+                total += max(0, right - max(left, previous_end))
+                previous_end = max(previous_end, right)
+            return total
+
+        x_pixels = stripe_pixels(self.vertical_seams, self.roi.x, self.roi.right)
+        y_pixels = stripe_pixels(self.horizontal_seams, self.roi.y, self.roi.bottom)
+        return (x_pixels * self.roi.height + y_pixels * self.roi.width
+                - x_pixels * y_pixels)
+
+    def metadata(self) -> dict:
+        def rectangle(window: PixelWindow) -> dict:
+            return {"x": window.x, "y": window.y,
+                    "width": window.width, "height": window.height}
+
+        return {
+            "global_dimensions": [self.global_grid.width, self.global_grid.height],
+            "global_transform": list(self.global_grid.transform),
+            "roi": rectangle(self.roi), "core_size": list(self.core_size),
+            "halo": list(self.halo), "alignment": self.alignment,
+            "tile_count": len(self.tiles), "core_pixel_count": self.roi.pixels,
+            "context_pixel_count_sum": sum(tile.context.pixels for tile in self.tiles),
+            "vertical_seams": list(self.vertical_seams),
+            "horizontal_seams": list(self.horizontal_seams),
+            "tiles": [{"index": tile.index, "row": tile.row, "column": tile.column,
+                       "core": rectangle(tile.core), "context": rectangle(tile.context),
+                       "context_transform": list(tile.context.grid(self.global_grid).transform)}
+                      for tile in self.tiles],
+        }
+
+
+def _positive_integer(value: int, name: str) -> None:
+    if type(value) is not int or value <= 0:
+        raise ValueError(f"{name} must be a positive integer")
+
+
+def _pair(value: int | tuple[int, int], name: str, *, allow_zero: bool = False) -> tuple[int, int]:
+    try:
+        values = (value, value) if type(value) is int else tuple(value)
+    except TypeError as error:
+        raise ValueError(f"{name} must be an integer or an integer pair") from error
+    if len(values) != 2 or any(type(item) is not int or item < (0 if allow_zero else 1)
+                              for item in values):
+        qualifier = "nonnegative" if allow_zero else "positive"
+        raise ValueError(f"{name} must contain two {qualifier} integers")
+    return values
+
+
+def plan_tiles(global_grid: Grid, core_size: int | tuple[int, int],
+               halo: int | tuple[int, int], *, alignment: int,
+               roi: PixelWindow | None = None) -> TilePlan:
+    """Partition a finite aligned ROI; clamp contexts to real global boundaries.
+
+    Core sizes are (width, height), halo is (horizontal, vertical). All origins,
+    dimensions, and halo widths must be multiples of ``alignment``. A final
+    core may be shorter than ``core_size`` while remaining aligned. The caller
+    must supply a halo at least as large as the model's exact required support.
+    """
+    _positive_integer(alignment, "alignment")
+    core = _pair(core_size, "core_size")
+    margin = _pair(halo, "halo", allow_zero=True)
+    region = roi or PixelWindow(0, 0, global_grid.width, global_grid.height)
+    region.grid(global_grid)  # Validate actual finite domain before planning.
+    aligned = {"global width": global_grid.width, "global height": global_grid.height,
+               "ROI x": region.x, "ROI y": region.y, "ROI width": region.width,
+               "ROI height": region.height, "core width": core[0], "core height": core[1],
+               "horizontal halo": margin[0], "vertical halo": margin[1]}
+    for name, value in aligned.items():
+        if value % alignment:
+            raise ValueError(f"{name} must be a multiple of alignment {alignment}")
+    tiles = []
+    for row, y in enumerate(range(region.y, region.bottom, core[1])):
+        for column, x in enumerate(range(region.x, region.right, core[0])):
+            output = PixelWindow(x, y, min(core[0], region.right - x),
+                                 min(core[1], region.bottom - y))
+            left, top = max(0, x - margin[0]), max(0, y - margin[1])
+            right = min(global_grid.width, output.right + margin[0])
+            bottom = min(global_grid.height, output.bottom + margin[1])
+            context = PixelWindow(left, top, right - left, bottom - top)
+            tiles.append(Tile(len(tiles), row, column, output, context))
+    return TilePlan(global_grid, region, core, margin, alignment, tuple(tiles))
+
+
+def seam_mask(plan: TilePlan, width: int = 1) -> ee.Image:
+    """Mask an exact union of bands around internal core boundaries."""
+    _positive_integer(width, "seam width")
+    coordinates = ee.Image.pixelCoordinates(plan.global_grid.projection())
+    x, y = coordinates.select("x"), coordinates.select("y")
+    mask = ee.Image.constant(0).reproject(plan.global_grid.projection())
+    for seam in plan.vertical_seams:
+        mask = mask.Or(x.gte(seam - width).And(x.lt(seam + width)))
+    for seam in plan.horizontal_seams:
+        mask = mask.Or(y.gte(seam - width).And(y.lt(seam + width)))
+    return mask.And(plan.roi_grid.mask()).rename("seam")
+
+
+@dataclass(frozen=True)
+class CompiledTiledModel(Raster):
+    plan: TilePlan | None = None
+    tiles: tuple[CompiledModel, ...] = ()
+    nodes: dict[str, Raster] = field(default_factory=dict)
+    coverage: ee.Image | None = None
+    planned_coverage: ee.Image | None = None
+    receptive_field: ReceptiveFieldReport | None = None
+
+
+def compile_tiled_model(model: nn.Module, image: ee.Image,
+                        input_bands: Sequence[str], plan: TilePlan, *,
+                        weight_provider: WeightProvider | None = None,
+                        output_nodes: Sequence[str] = ()) -> CompiledTiledModel:
+    """Compile every context in EE, retain disjoint cores, and mosaic in EE.
+
+    One shared weight provider serves all tiles, so an asset provider can reuse
+    its arrays. Selected diagnostic output nodes must have input resolution.
+    The returned grid is the *global* grid; its image is masked to the ROI.
+    Every compiled intermediate uses float32 and clipped finite footprints.
+    """
+    if not plan.tiles:
+        raise ValueError("TilePlan must contain at least one core")
+    support = analyze_receptive_field(model)
+    if plan.alignment % support.alignment:
+        raise ValueError(f"Tile alignment must be a multiple of model alignment {support.alignment}")
+    for tile in plan.tiles:
+        context, core = tile.context, tile.core
+        for value in (core.x, core.y, core.width, core.height,
+                      context.x, context.y, context.width, context.height):
+            if value % support.alignment:
+                raise ValueError(f"Tile origins and dimensions require model alignment {support.alignment}")
+        artificial_margins = []
+        if context.x > 0:
+            artificial_margins.append(core.x - context.x)
+        if context.y > 0:
+            artificial_margins.append(core.y - context.y)
+        if context.right < plan.global_grid.width:
+            artificial_margins.append(context.right - core.right)
+        if context.bottom < plan.global_grid.height:
+            artificial_margins.append(context.bottom - core.bottom)
+        if any(margin < support.safe_halo for margin in artificial_margins):
+            raise ValueError(f"Artificial context boundaries require halo at least {support.safe_halo}")
+    compiled_tiles = tuple(compile_model(
+        model, image, input_bands, tile.context.grid(plan.global_grid),
+        weight_provider=weight_provider, include_unfused_nodes=False,
+        execution_dtype="float32", clip_footprints=True)
+        for tile in plan.tiles)
+
+    def core_images(rasters: Sequence[Raster]) -> list[ee.Image]:
+        bands = rasters[0].bands
+        pieces = []
+        for tile, raster in zip(plan.tiles, rasters, strict=True):
+            if raster.grid != tile.context.grid(plan.global_grid):
+                raise UnsupportedOperation("Tiled output must preserve the input spatial grid")
+            if raster.bands != bands:
+                raise UnsupportedOperation("All tiled outputs must have identical bands")
+            core_grid = tile.core.grid(plan.global_grid)
+            pieces.append(raster.image.select(list(bands)).clip(core_grid.geometry())
+                          .updateMask(core_grid.mask()))
+        return pieces
+
+    def mosaic(rasters: Sequence[Raster]) -> Raster:
+        result = (ee.ImageCollection.fromImages(core_images(rasters)).mosaic().rename(list(rasters[0].bands))
+                  .setDefaultProjection(plan.global_grid.projection())
+                  .updateMask(plan.roi_grid.mask()).toFloat())
+        return Raster(result, rasters[0].bands, plan.global_grid)
+
+    output = mosaic(compiled_tiles)
+    nodes = {}
+    for name in output_nodes:
+        if any(name not in tile.nodes for tile in compiled_tiles):
+            raise ValueError(f"Requested tiled output node {name!r} is absent")
+        nodes[name] = mosaic([tile.nodes[name] for tile in compiled_tiles])
+    masks = [ee.Image.constant(1).rename("tile_coverage")
+             .reproject(plan.global_grid.projection())
+             .updateMask(tile.core.grid(plan.global_grid).mask()) for tile in plan.tiles]
+    planned_coverage = (ee.ImageCollection.fromImages(masks).sum().rename("planned_tile_coverage")
+                        .setDefaultProjection(plan.global_grid.projection())
+                        .updateMask(plan.roi_grid.mask()))
+    actual_masks = [piece.mask().reduce(ee.Reducer.min()).unmask(0, sameFootprint=False)
+                    .rename("tile_coverage").reproject(plan.global_grid.projection())
+                    for piece in core_images(compiled_tiles)]
+    coverage = (ee.ImageCollection.fromImages(actual_masks).sum().rename("tile_coverage")
+                .setDefaultProjection(plan.global_grid.projection()).updateMask(plan.roi_grid.mask()))
+    return CompiledTiledModel(output.image, output.bands, output.grid,
+                              plan, compiled_tiles, nodes, coverage, planned_coverage, support)

tensor_ee/weight_assets.py

@@ -0,0 +1,371 @@
+"""Store arbitrary frozen convolution tensors in Earth Engine table assets.
+
+Each feature has a harmless point geometry and properties ``layer``, ``row``,
+``weights_json`` (one JSON string of a numeric chunk), and, when present, ``bias``.  EE table exports
+reject null geometries and list-valued properties.  Chunks are sorted by an explicit integer index;
+no inference weight is downloaded or embedded in an asset-backed graph.
+
+Uploads are independently measured, bounded request graphs.  Nine significant
+figures preserve every finite float32 value on restoration through
+``ee.Array.float()``. Float32 is the maximum supported model precision.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+from pathlib import Path
+import time
+from typing import Any, Callable, Iterator, Mapping, Sequence
+
+import ee
+import numpy as np
+import torch
+from torch import nn
+from torch.fx import Node, symbolic_trace
+
+from .compiler import UnsupportedOperation, _ROOT_MODULES
+from .semantics import semantic_call, diagnostic_convolution
+# The asset format and its provider moved to the library; these names are
+# re-exports of the same objects (one FeatureCollectionWeightProvider class).
+from .weights.assets import (  # noqa: F401
+    ConvolutionWeights, FeatureCollectionWeightProvider, convolution_weights, feature_records)
+
+
+def effective_convolutions(model: nn.Module) -> list[ConvolutionWeights]:
+    """Collect exactly the final graph's conv tensors, including BN folding.
+
+    Names follow FX node keys, as ``compile_model`` does.  A convolution with a
+    sole BatchNorm consumer is replaced by its folded convolution, keyed by the
+    BN node.  Original conv diagnostics are deliberately absent from the asset.
+    """
+    if any(module.training for module in model.modules()):
+        raise ValueError("Set model.eval() before collecting inference weights")
+    if isinstance(model, _ROOT_MODULES):
+        model = nn.Sequential(model)
+    traced = symbolic_trace(model)
+    result = []
+    for node in traced.graph.nodes:
+        if node.op in ("placeholder", "output"):
+            continue
+        try:
+            call = semantic_call(node, traced)
+        except UnsupportedOperation as error:
+            raise UnsupportedOperation(f"At FX node {node.name} ({node.target}): {error}") from error
+        if call.kind == "conv" and not diagnostic_convolution(node, traced):
+            result.append(convolution_weights(node.name, call.module, str(node.target)))
+        elif call.kind == "batchnorm":
+            source = str(call.options["predecessor"].target) + "+" + str(node.target)
+            result.append(convolution_weights(node.name, call.module, source))
+    if not result:
+        raise ValueError("Model contains no deployable convolutions")
+    return result
+
+
+def float32_wire_values(values: np.ndarray) -> list[float]:
+    """Short decimal values which round back to the exact source float32s."""
+    return [float(format(float(value), ".9g")) for value in values]
+
+
+def collection_from_records(records: Sequence[Mapping[str, Any]]) -> ee.FeatureCollection:
+    return ee.FeatureCollection([ee.Feature(dict(record)) for record in records])
+
+
+def _split_measured(records: list[dict[str, Any]], max_bytes: int) -> Iterator[tuple[list[dict[str, Any]], int]]:
+    actual = len(ee.serializer.toJSON(collection_from_records(records)).encode("utf-8"))
+    if actual <= max_bytes:
+        yield records, actual
+    elif len(records) == 1:
+        raise ValueError(f"One row requires {actual} serialized bytes, exceeding {max_bytes}")
+    else:
+        center = len(records) // 2
+        yield from _split_measured(records[:center], max_bytes)
+        yield from _split_measured(records[center:], max_bytes)
+
+
+def upload_chunks(parameters: Sequence[ConvolutionWeights], max_bytes: int = 8_000_000
+                  ) -> Iterator[tuple[list[dict[str, Any]], int]]:
+    """Pack rows conservatively, then verify each actual EE graph byte count."""
+    if not 100_000 <= max_bytes <= 8_000_000:
+        raise ValueError("max_bytes must be between 100,000 and 8,000,000")
+    records = []
+    rough_bytes = 0
+    # A literal FC is converted into API feature expressions.  Reserving 10%
+    # and 32 KB covers that overhead; the exact serializer check is decisive.
+    target = int(max_bytes * .9) - 32_000
+    for record in feature_records(parameters):
+        size = len(json.dumps(record, separators=(",", ":")).encode()) + 512
+        if records and rough_bytes + size > target:
+            yield from _split_measured(records, max_bytes)
+            records, rough_bytes = [], 0
+        records.append(record)
+        rough_bytes += size
+    if records:
+        yield from _split_measured(records, max_bytes)
+
+
+def _write_manifest(path: Path, manifest: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    temporary.write_text(json.dumps(manifest, indent=2, allow_nan=False) + "\n")
+    temporary.replace(path)
+
+
+class _RecordedTask:
+    """A task an earlier run started, known only by the ID its manifest recorded: polled
+    with ``ee.data.getTaskStatus``."""
+
+    def __init__(self, task_id: str):
+        self.id = task_id
+
+    def status(self) -> dict[str, Any]:
+        return ee.data.getTaskStatus([self.id])[0]
+
+
+def _layer_metadata(manifest: Mapping[str, Any]) -> dict[str, Any]:
+    return {key: {k: v for k, v in meta.items() if k != "assets"} for key, meta in manifest["layers"].items()}
+
+
+def _asset_exists(asset_id: str) -> bool:
+    try:
+        ee.data.getAsset(asset_id)
+    except ee.EEException:
+        return False
+    return True
+
+
+def _resume(fresh: dict[str, Any], manifest_path: Path, iterator) -> tuple[dict[str, Any], list, list]:
+    """The manifest an interrupted ``export_weights`` left, reconciled for these parameters:
+    (manifest, tasks to poll, shards to dispatch again). Refused when it belongs to other
+    parameters or another chunking, or when a shard FAILED or was CANCELLED."""
+    manifest = json.loads(Path(manifest_path).read_text())
+    for name in ("format", "asset_namespace", "max_collection_payload_bytes", "parameter_count"):
+        if manifest.get(name) != fresh[name]:
+            raise ValueError(f"the manifest's {name} differs: it is not an upload of these parameters")
+    if _layer_metadata(manifest) != _layer_metadata(fresh):
+        raise ValueError("the manifest's layers differ: it is not an upload of these parameters")
+    if manifest.get("status") not in ("uploading", "ready"):
+        raise ValueError(f"the upload is {manifest.get('status')!r}: a failed shard is not resumed")
+    polling, redispatch = [], []
+    for shard in manifest["shards"]:
+        try:
+            records, byte_count = next(iterator)
+        except StopIteration:
+            raise ValueError("the manifest records more shards than these parameters make") from None
+        layers = sorted({record["properties"]["layer"] for record in records})
+        expected = {"asset_id": f"{fresh['asset_namespace']}_part_{shard['index']:04d}", "rows": len(records),
+                    "layers": layers, "collection_payload_bytes": byte_count}
+        if any(shard.get(k) != v for k, v in expected.items()):
+            raise ValueError(f"shard {shard['index']} differs from the chunking of these parameters")
+        state = shard.get("status")
+        if state == "COMPLETED":
+            continue
+        if state in ("FAILED", "CANCELLED"):
+            raise ValueError(f"shard {shard['asset_id']} is {state}: it is not resumed")
+        if shard.get("task_id"):                       # started: poll it to the end
+            polling.append((_RecordedTask(shard["task_id"]), shard))
+        elif _asset_exists(shard["asset_id"]):         # recorded "starting", but its export finished
+            shard.update(status="COMPLETED", finished_unix=time.time(), resumed="asset found")
+        else:                                          # never started: dispatch it again
+            redispatch.append((records, byte_count, shard))
+    return manifest, polling, redispatch
+
+
+def export_weights(parameters: Sequence[ConvolutionWeights], asset_namespace: str,
+                   manifest_path: Path, *, max_bytes: int = 8_000_000,
+                   max_active: int = 2, poll_seconds: float = 15,
+                   on_update: Callable[[dict[str, Any]], None] | None = None,
+                   resume: bool = False) -> dict[str, Any]:
+    """Upload FC shards, waiting for completion with bounded batch concurrency.
+
+    ``asset_namespace`` is a new unique prefix under projects/.../assets.  This
+    function never overwrites or deletes assets.  It updates its local manifest
+    immediately after each dispatched task and completion, so interrupted runs
+    remain auditable.  A failed shard is reported and prevents a ready manifest.
+
+    ``resume=True`` continues the upload that ``manifest_path`` records, for the same
+    parameters and chunking (refused otherwise): it polls the recorded tasks to their
+    end; a shard recorded "starting" without a task ID is complete if its asset exists
+    (one ``getAsset``) and dispatched again if not; the remaining shards follow from the
+    same deterministic chunking; then the manifest is finalized. A manifest that is
+    already ready is returned as it is. Without ``resume`` the manifest starts afresh.
+    """
+    if max_active not in (1, 2):
+        raise ValueError("Asset upload uses at most two simultaneous exports")
+    if not 1 <= poll_seconds <= 30:
+        raise ValueError("poll_seconds must be between 1 and 30")
+    keys = [parameter.key for parameter in parameters]
+    if len(keys) != len(set(keys)):
+        raise ValueError("Duplicate convolution keys")
+    manifest: dict[str, Any] = {
+        "format": "native_ee_featurecollection_weights_v2", "status": "uploading",
+        "asset_namespace": asset_namespace, "started_unix": time.time(),
+        "serialization": "JSON strings of numeric chunks; 9 significant decimal digits for float32; float() restoration",
+        "dummy_geometry": "Point(0, 0); EE table exports require non-null geometry",
+        "max_collection_payload_bytes": max_bytes, "max_active_tasks": max_active,
+        "parameter_count": sum(parameter.weights.size + (0 if parameter.bias is None else parameter.bias.size)
+                               for parameter in parameters),
+        "layers": {parameter.key: {"shape": list(parameter.shape), "dtype": parameter.dtype,
+                                   "bias": parameter.bias is not None, "source": parameter.source,
+                                   "sha256": parameter.sha256, "assets": []}
+                   for parameter in parameters}, "shards": [],
+    }
+    iterator = iter(upload_chunks(parameters, max_bytes))
+    active: list[tuple[Any, dict[str, Any]]] = []
+    pending: list[tuple[list[dict[str, Any]], int, dict[str, Any]]] = []
+    exhausted = False
+    if resume:
+        manifest, active, pending = _resume(manifest, Path(manifest_path), iterator)
+        if manifest["status"] == "ready":
+            return manifest
+
+    def save() -> None:
+        _write_manifest(manifest_path, manifest)
+        if on_update:
+            on_update(manifest)
+
+    def dispatch(records: list[dict[str, Any]], byte_count: int, shard: dict[str, Any] | None) -> None:
+        if shard is None:                               # the next shard of the chunking
+            index = len(manifest["shards"])
+            asset_id = f"{asset_namespace}_part_{index:04d}"
+            layers = sorted({record["properties"]["layer"] for record in records})
+            shard = {"index": index, "asset_id": asset_id, "rows": len(records),
+                     "layers": layers, "collection_payload_bytes": byte_count,
+                     "status": "starting", "started_unix": time.time()}
+            manifest["shards"].append(shard)
+            for key in layers:
+                manifest["layers"][key]["assets"].append(asset_id)
+        else:                                           # a recorded shard that never started
+            shard.update(status="starting", started_unix=time.time())
+        save()
+        task = ee.batch.Export.table.toAsset(
+            collection_from_records(records), description=f"native_unet_weights_part_{shard['index']:04d}",
+            assetId=shard["asset_id"], overwrite=False)
+        try:
+            task.start()
+        except Exception as error:
+            shard.update(status="FAILED", error=f"{type(error).__name__}: {error}")
+            manifest.update(status="failed", finished_unix=time.time())
+            # Record the failure, but never let the save's callback (a polling cap, say)
+            # replace the start() error.
+            with contextlib.suppress(Exception):
+                save()
+            raise
+        shard.update(task_id=task.id, status="READY")
+        active.append((task, shard))
+        save()
+
+    save()
+    while active or pending or not exhausted:
+        while (pending or not exhausted) and len(active) < max_active:
+            if pending:
+                dispatch(*pending.pop(0))
+                continue
+            try:
+                records, byte_count = next(iterator)
+            except StopIteration:
+                exhausted = True
+                break
+            dispatch(records, byte_count, None)
+        any_finished = False
+        for task, shard in list(active):
+            status = task.status()
+            shard["task_status"] = status
+            shard["status"] = status["state"]
+            if status["state"] in ("COMPLETED", "FAILED", "CANCELLED"):
+                shard["finished_unix"] = time.time()
+                active.remove((task, shard))
+                any_finished = True
+                if status["state"] != "COMPLETED":
+                    manifest.update(status="failed", finished_unix=time.time())
+                    save()
+                    raise RuntimeError(f"Weight shard {shard['asset_id']} failed: {status}")
+        save()
+        if active and not any_finished:
+            time.sleep(poll_seconds)
+    manifest.update(status="ready", finished_unix=time.time(),
+                    collection_payload_bytes=sum(shard["collection_payload_bytes"] for shard in manifest["shards"]))
+    save()
+    return manifest
+
+
+def probe_asset_roundtrip(asset_namespace: str, result_path: Path) -> dict[str, Any]:
+    """Live all-element test, including a standard ResNet18 length-4,608 row."""
+    rng = np.random.default_rng(47)
+    weights = rng.normal(0, .025, (3, 4608)).astype(np.float32)
+    bias = np.array([-.03, .0, .03], dtype=np.float32)
+    parameter = ConvolutionWeights("probe", weights, bias, "float32", "synthetic_probe")
+    manifest_path = result_path.with_name(result_path.stem + "_manifest.json")
+    start = time.perf_counter()
+    manifest = export_weights([parameter], asset_namespace, manifest_path, poll_seconds=5)
+    provider = FeatureCollectionWeightProvider(manifest)
+    actual, actual_bias = provider.arrays("probe")
+    expected = ee.Array(weights.astype(float).tolist())
+    expected_bias = ee.Array(bias.astype(float).reshape(-1, 1).tolist())
+    collection = ee.FeatureCollection(manifest["shards"][0]["asset_id"])
+    verification = ee.Dictionary({
+        "shape": actual.length(), "feature_count": collection.size(), "row_count": collection.filter(ee.Filter.eq("chunk", 0)).size(),
+        "max_weight_absolute_error": actual.subtract(expected).abs().reduce(ee.Reducer.max(), [0, 1]),
+        "max_bias_absolute_error": actual_bias.subtract(expected_bias).abs().reduce(ee.Reducer.max(), [0, 1]),
+    })
+    check_start = time.perf_counter()
+    raw = verification.getInfo()
+    report = {"status": "pass" if (raw["max_weight_absolute_error"] == [[0]] and raw["max_bias_absolute_error"] == [[0]]
+                                  and raw["shape"] == [3, 4608] and raw["row_count"] == 3) else "fail",
+              "manifest_path": str(manifest_path), "asset_ids": [shard["asset_id"] for shard in manifest["shards"]],
+              "row_length": 4608, "source_sha256": parameter.sha256, "verification": raw,
+              "verification_seconds": time.perf_counter() - check_start,
+              "total_seconds": time.perf_counter() - start,
+              "verification_request_bytes": len(ee.serializer.toJSON(verification).encode("utf-8")),
+              "serialization": manifest["serialization"]}
+    _write_manifest(result_path, report)
+    return report
+
+
+def verify_asset_rows(parameters: Sequence[ConvolutionWeights], manifest: Mapping[str, Any] | str | Path,
+                      result_path: Path, *, rows_per_layer: int = 3,
+                      on_update: Callable[[dict[str, Any]], None] | None = None) -> dict[str, Any]:
+    """Check deterministic rows of every server-reconstructed matrix exactly.
+
+    This is a storage and ordering check, separate from model output parity.
+    Only scalar errors, shapes, and timing are returned.  The exact source
+    fixtures are embedded for selected rows, keeping each query comfortably
+    below the request payload cap.  It does not claim a full-asset checksum.
+    """
+    if rows_per_layer < 1:
+        raise ValueError("rows_per_layer must be positive")
+    provider = FeatureCollectionWeightProvider(manifest)
+    report: dict[str, Any] = {"status": "running", "started_unix": time.time(),
+                              "comparison": "deterministic sampled rows from every layer; native server reconstruction",
+                              "layers": []}
+    _write_manifest(result_path, report)
+    for parameter in parameters:
+        metadata = provider.manifest["layers"][parameter.key]
+        if parameter.sha256 != metadata["sha256"]:
+            raise ValueError(f"Source fingerprint mismatch for {parameter.key}")
+        weights, bias = provider.arrays(parameter.key)
+        indices = sorted(set(int(value) for value in np.linspace(0, parameter.shape[0] - 1,
+                                                                  min(rows_per_layer, parameter.shape[0]))))
+        row_checks = {}
+        for index in indices:
+            expected = ee.Array(parameter.weights[index:index + 1].astype(float).tolist())
+            difference = weights.slice(0, index, index + 1).subtract(expected).abs()
+            row_checks[f"weight_error_row_{index}"] = difference.reduce(ee.Reducer.max(), [0, 1]).get([0, 0])
+            if bias is not None:
+                row_checks[f"bias_error_row_{index}"] = bias.get([index, 0]).subtract(float(parameter.bias[index])).abs()
+        row_checks["shape"] = weights.length()
+        request = ee.Dictionary(row_checks)
+        started = time.perf_counter()
+        raw = request.getInfo()
+        errors = [value for key, value in raw.items() if key != "shape"]
+        entry = {"key": parameter.key, "rows": indices, "shape": list(parameter.shape),
+                 "status": "pass" if raw["shape"] == list(parameter.shape) and all(value == 0 for value in errors) else "fail",
+                 "max_absolute_error": max(errors), "evaluation_seconds": time.perf_counter() - started,
+                 "request_bytes": len(ee.serializer.toJSON(request).encode()), "server_metrics": raw}
+        report["layers"].append(entry)
+        _write_manifest(result_path, report)
+        if on_update:
+            on_update(entry)
+    report.update(status="pass" if all(entry["status"] == "pass" for entry in report["layers"]) else "fail",
+                  finished_unix=time.time())
+    _write_manifest(result_path, report)
+    return report