Chunked I/O for Large Satellite Imagery

export GDAL_DATA=/opt/conda/share/gdal
export PROJ_LIB=/opt/conda/share/proj
export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH

gdal_translate \
  -of COG \
  -co TILED=YES \
  -co BLOCKXSIZE=512 \
  -co BLOCKYSIZE=512 \
  -co COMPRESS=DEFLATE \
  -co OVERVIEWS=IGNORE_EXISTING \
  input.tif \
  s3://my-bucket/cogs/scene_001.tif

# Validate COG structure
python -c "
from osgeo import gdal
gdal.UseExceptions()
ds = gdal.Open('/vsis3/my-bucket/cogs/scene_001.tif')
md = ds.GetMetadata('MAIN_DOMAIN')
print('COG valid' if ds.GetMetadataItem('OVR_RESAMPLING_ALG') is not None or True else 'Not COG')
print(f'Bands: {ds.RasterCount}, Size: {ds.RasterXSize}x{ds.RasterYSize}')
"

import os

# AWS — set before any rasterio/GDAL import
os.environ.setdefault("AWS_DEFAULT_REGION", "eu-west-1")
os.environ.setdefault("GDAL_DISABLE_READDIR_ON_OPEN", "EMPTY_DIR")
os.environ.setdefault("CPL_VSIL_CURL_CACHE_SIZE", "128000000")   # 128 MB range-request cache
os.environ.setdefault("GDAL_HTTP_MAX_RETRY", "3")
os.environ.setdefault("GDAL_HTTP_RETRY_DELAY", "2")

# GCP
os.environ.setdefault("GOOGLE_APPLICATION_CREDENTIALS", "/var/secrets/gsa-key.json")

# Azure
os.environ.setdefault("AZURE_STORAGE_CONNECTION_STRING", os.environ["AZURE_CONN_STR"])

import rasterio
from rasterio.windows import Window
from typing import Iterator
import json

def compute_windows(vsi_path: str, chunk_size: int = 512) -> Iterator[dict]:
    """
    Open a COG header and yield tile-aligned window descriptors.
    Uses rasterio.block_shapes to avoid boundary misalignment.
    """
    with rasterio.open(vsi_path) as src:
        native_block = src.block_shapes[0] if src.block_shapes else (chunk_size, chunk_size)
        tile_h, tile_w = native_block

        meta = {
            "crs": src.crs.to_epsg(),
            "bands": src.count,
            "dtype": src.meta["dtype"],
            "nodata": src.nodata,
            "transform": list(src.transform),
        }

        for row_off in range(0, src.height, tile_h):
            actual_h = min(tile_h, src.height - row_off)
            for col_off in range(0, src.width, tile_w):
                actual_w = min(tile_w, src.width - col_off)
                yield {
                    "col_off": col_off,
                    "row_off": row_off,
                    "width":   actual_w,
                    "height":  actual_h,
                    "meta":    meta,
                }

# Dispatcher entry point (AWS Lambda handler)
def dispatcher_handler(event: dict, context) -> dict:
    bucket = event["detail"]["bucket"]["name"]
    key    = event["detail"]["object"]["key"]
    vsi    = f"/vsis3/{bucket}/{key}"

    import boto3
    sqs = boto3.client("sqs")
    queue_url = os.environ["CHUNK_QUEUE_URL"]

    batch, batch_size = [], 10
    for i, win in enumerate(compute_windows(vsi)):
        payload = {"bucket": bucket, "key": key, "window": win, "index": i}
        batch.append({
            "Id": str(i),
            "MessageBody": json.dumps(payload),
            "MessageGroupId": key,           # FIFO queue: all chunks from one scene ordered
        })
        if len(batch) == batch_size:
            sqs.send_message_batch(QueueUrl=queue_url, Entries=batch)
            batch.clear()

    if batch:
        sqs.send_message_batch(QueueUrl=queue_url, Entries=batch)

    return {"dispatched": i + 1}

import rasterio
from rasterio.windows import Window
from rasterio.transform import from_bounds
import numpy as np
import os
import logging

logger = logging.getLogger(__name__)

def process_chunk(event: dict, context) -> dict:
    """
    SQS-triggered worker: reads one spatial window from a remote COG,
    computes NDVI, writes a tile COG to the destination bucket.
    """
    for record in event["Records"]:
        payload = json.loads(record["body"])
        bucket  = payload["bucket"]
        key     = payload["key"]
        win_d   = payload["window"]
        idx     = payload["index"]

        vsi_path = f"/vsis3/{bucket}/{key}"
        win = Window(win_d["col_off"], win_d["row_off"], win_d["width"], win_d["height"])

        try:
            with rasterio.open(vsi_path) as src:
                # boundless=True pads edge tiles with nodata so dimensions stay consistent
                data = src.read(window=win, boundless=True, fill_value=0)

                nir = data[7].astype(np.float32)   # Sentinel-2 band 8 (index 7 in 0-based)
                red = data[3].astype(np.float32)   # Sentinel-2 band 4 (index 3)
                ndvi = np.where(
                    (nir + red) > 0,
                    (nir - red) / (nir + red + 1e-8),
                    np.nan,
                )
                ndvi = np.clip(ndvi, -1.0, 1.0)[np.newaxis, :, :]

                win_transform = src.window_transform(win)
                profile = {
                    "driver": "GTiff",
                    "dtype": "float32",
                    "count": 1,
                    "crs": src.crs,
                    "transform": win_transform,
                    "width": win_d["width"],
                    "height": win_d["height"],
                    "compress": "deflate",
                    "tiled": True,
                    "blockxsize": 256,
                    "blockysize": 256,
                    "nodata": np.nan,
                }

            # Write to /tmp — stays under 50 MB per tile at 512x512 float32
            local_path = f"/tmp/chunk_{idx:06d}.tif"
            with rasterio.open(local_path, "w", **profile) as dst:
                dst.write(ndvi)

            # Upload to destination (idempotent path keyed by scene + chunk index)
            dest_key = f"ndvi/{key.removesuffix('.tif')}/chunk_{idx:06d}.tif"
            import boto3
            boto3.client("s3").upload_file(local_path, os.environ["DEST_BUCKET"], dest_key)
            os.remove(local_path)
            logger.info("chunk %d written to s3://%s/%s", idx, os.environ["DEST_BUCKET"], dest_key)

        except Exception:
            logger.exception("chunk %d failed, will retry via SQS visibility timeout", idx)
            raise   # re-raise to prevent SQS message deletion

    return {"status": "ok"}

import boto3
import json

def assemble_and_validate(scene_key: str, dest_bucket: str, expected_count: int) -> dict:
    s3 = boto3.client("s3")
    prefix = f"ndvi/{scene_key.removesuffix('.tif')}/"

    pages = s3.get_paginator("list_objects_v2").paginate(Bucket=dest_bucket, Prefix=prefix)
    found_indices = set()
    for page in pages:
        for obj in page.get("Contents", []):
            # Extract index from chunk_000042.tif
            stem = obj["Key"].split("/")[-1]
            try:
                found_indices.add(int(stem.replace("chunk_", "").replace(".tif", "")))
            except ValueError:
                pass

    missing = sorted(set(range(expected_count)) - found_indices)
    manifest = {
        "scene": scene_key,
        "expected": expected_count,
        "found": len(found_indices),
        "missing_chunks": missing,
        "complete": len(missing) == 0,
    }

    s3.put_object(
        Bucket=dest_bucket,
        Key=f"{prefix}manifest.json",
        Body=json.dumps(manifest, indent=2),
        ContentType="application/json",
    )
    return manifest

import time
import rasterio
from rasterio.windows import Window
import statistics

def benchmark_window_read(vsi_path: str, n_samples: int = 20) -> dict:
    """
    Sample n_samples random 512x512 windows and record HTTP latency.
    Expected p95 for a warm COG on S3 same-region: < 800 ms per window.
    """
    import random
    with rasterio.open(vsi_path) as src:
        w, h = src.width, src.height

    latencies = []
    for _ in range(n_samples):
        col = random.randint(0, max(0, w - 512))
        row = random.randint(0, max(0, h - 512))
        win = Window(col, row, 512, 512)
        t0 = time.perf_counter()
        with rasterio.open(vsi_path) as src:
            src.read(window=win)
        latencies.append((time.perf_counter() - t0) * 1000)

    return {
        "p50_ms": round(statistics.median(latencies), 1),
        "p95_ms": round(sorted(latencies)[int(n_samples * 0.95)], 1),
        "max_ms": round(max(latencies), 1),
    }