Optimizing Chunked I/O for Multi-Band Sentinel-2 Processing

import io
import os
import logging
import numpy as np
import rasterio
from rasterio.windows import Window
from rasterio.enums import Resampling
import boto3

# ---------------------------------------------------------------------------
# Explicit environment — do not rely on shell inheritance in Lambda
# ---------------------------------------------------------------------------
os.environ.setdefault("GDAL_DATA", "/opt/gdal/share/gdal")
os.environ.setdefault("PROJ_LIB", "/opt/share/proj")
os.environ.setdefault("GDAL_HTTP_MERGE_CONSECUTIVE_RANGES", "YES")
os.environ.setdefault("GDAL_HTTP_MULTIPLEX", "YES")
# 200 MB in-process COG tile cache — keeps decompressed blocks for the window
os.environ.setdefault("CPL_VSIL_CURL_CACHE_SIZE", "200000000")

# Sentinel-2 10m band indices in a standard COG (B02=1, B03=2, B04=3, B08=4)
BAND_BLUE  = 1   # B02 10m
BAND_GREEN = 2   # B03 10m
BAND_RED   = 3   # B04 10m
BAND_NIR   = 4   # B08 10m

CHUNK_SIZE = 512  # Must match COG internal block size — confirmed via block_shapes

logger = logging.getLogger(__name__)
s3 = boto3.client("s3")


def get_chunk_size(bucket: str, key: str) -> int:
    """
    Query the COG's native block size and return the safe CHUNK_SIZE.
    Falls back to 512 if block_shapes is not available (non-tiled GeoTIFF).
    Call this once in the dispatcher, not per worker invocation.
    """
    url = f"/vsis3/{bucket}/{key}"
    with rasterio.open(url) as src:
        shapes = src.block_shapes  # list of (rows, cols) per band
        if shapes:
            row_block, col_block = shapes[0]
            logger.info("Native COG block size: %d×%d", row_block, col_block)
            return col_block  # typically 256 or 512 for Sentinel-2 COGs
    return 512


def read_band_aligned(
    src: rasterio.DatasetReader,
    band_index: int,
    window: Window,
    target_size: int,
    use_bilinear: bool = True,
) -> np.ndarray:
    """
    Read a single band into a (target_size, target_size) array.
    Works for 10m bands (native) and 20m/60m bands (resampled in-memory).
    Only the compressed blocks intersecting `window` are fetched.
    """
    resampling = Resampling.bilinear if use_bilinear else Resampling.nearest
    return src.read(
        band_index,
        window=window,
        out_shape=(target_size, target_size),
        resampling=resampling,
        masked=False,
    ).astype(np.float32)


def process_sentinel2_chunk(
    bucket: str,
    key_10m: str,   # COG containing 10m bands (B02, B03, B04, B08)
    tile_x: int,
    tile_y: int,
    output_prefix: str = "processed",
) -> str:
    """
    Read a single CHUNK_SIZE×CHUNK_SIZE tile from a Sentinel-2 10m COG,
    compute NDVI, and upload the result to S3 without writing to /tmp.

    Returns the S3 key of the written output tile.
    """
    col_off = tile_x * CHUNK_SIZE
    row_off = tile_y * CHUNK_SIZE
    window = Window(col_off, row_off, CHUNK_SIZE, CHUNK_SIZE)

    output_key = (
        f"{output_prefix}/"
        f"{key_10m.rsplit('/', 1)[-1].replace('.tif', '')}"
        f"_x{tile_x:04d}_y{tile_y:04d}_ndvi.tif"
    )

    url = f"/vsis3/{bucket}/{key_10m}"

    with rasterio.open(url) as src:
        # ------------------------------------------------------------------
        # Guard: skip tiles that fall entirely outside the scene footprint.
        # Sentinel-2 granules are not rectangular — edge tiles are nodata.
        # ------------------------------------------------------------------
        actual_w = min(CHUNK_SIZE, src.width  - col_off)
        actual_h = min(CHUNK_SIZE, src.height - row_off)
        if actual_w <= 0 or actual_h <= 0:
            logger.info("Tile (%d, %d) outside scene extent — skipping", tile_x, tile_y)
            return ""

        # Lazy, windowed reads — only bands needed for NDVI
        red = read_band_aligned(src, BAND_RED, window, CHUNK_SIZE)
        nir = read_band_aligned(src, BAND_NIR, window, CHUNK_SIZE)

        # Safe NDVI: suppress divide-by-zero on nodata / water / cloud pixels
        with np.errstate(divide="ignore", invalid="ignore"):
            ndvi = np.where(
                (red + nir) == 0,
                np.float32(-9999),  # nodata sentinel — not 0, which is a valid NDVI
                (nir - red) / (nir + red),
            )

        # ------------------------------------------------------------------
        # Stream output directly to S3 — no /tmp write.
        # io.BytesIO stays in Lambda process memory; typical tile is ~1 MB.
        # ------------------------------------------------------------------
        out_transform = src.window_transform(window)
        out_profile = {
            "driver": "GTiff",
            "dtype": "float32",
            "width": CHUNK_SIZE,
            "height": CHUNK_SIZE,
            "count": 1,
            "crs": src.crs,
            "transform": out_transform,
            "nodata": -9999.0,
            "compress": "deflate",
            "predictor": 3,    # horizontal differencing for float data
            "tiled": True,
            "blockxsize": 256,
            "blockysize": 256,
            "interleave": "band",
        }

        with io.BytesIO() as buf:
            with rasterio.open(buf, "w", **out_profile) as dst:
                dst.write(ndvi, 1)
            buf.seek(0)
            s3.put_object(
                Bucket=bucket,
                Key=output_key,
                Body=buf,
                ContentType="image/tiff",
                StorageClass="INTELLIGENT_TIERING",
            )

    logger.info("Written: s3://%s/%s", bucket, output_key)
    return output_key


def lambda_handler(event: dict, context) -> dict:
    """
    SQS-triggered entry point. Each message carries one tile coordinate.
    Dead-letter routing is handled at the queue level after maxReceiveCount
    retries — no manual DLQ logic required here.
    """
    results = []
    for record in event.get("Records", []):
        import json
        payload = json.loads(record["body"])
        output_key = process_sentinel2_chunk(
            bucket=payload["bucket"],
            key_10m=payload["key"],
            tile_x=payload["tile_x"],
            tile_y=payload["tile_y"],
            output_prefix=payload.get("output_prefix", "processed"),
        )
        results.append(output_key)

    return {"statusCode": 200, "processed": results}

import rasterio
import boto3
import json

BUCKET = "your-sentinel2-bucket"
KEY    = "S2A_MSIL2A_20240910T100601_N0511_R122_T32UND_20240910T163845_10m.tif"

# 1. Confirm COG block geometry matches CHUNK_SIZE
with rasterio.open(f"/vsis3/{BUCKET}/{KEY}") as src:
    shapes = src.block_shapes
    print("block_shapes:", shapes)          # expect [(512, 512), ...]
    print("band_count:", src.count)         # expect 4 for a standard 10m COG
    print("dtype:", src.dtypes[0])          # expect uint16 before float32 cast
    print("crs:", src.crs.to_epsg())        # expect 32632 for UTM zone 32N

# 2. Process tile (0, 0) and verify the output exists in S3
from optimizing_chunked_io import process_sentinel2_chunk
result_key = process_sentinel2_chunk(BUCKET, KEY, tile_x=0, tile_y=0)
print("Output key:", result_key)

s3 = boto3.client("s3")
head = s3.head_object(Bucket=BUCKET, Key=result_key)
print("Output size (bytes):", head["ContentLength"])   # typically 200 KB–1.5 MB
print("Content-Type:", head["ContentType"])            # image/tiff

# 3. Spot-check NDVI range
with rasterio.open(f"/vsis3/{BUCKET}/{result_key}") as out:
    ndvi_tile = out.read(1)
    valid = ndvi_tile[ndvi_tile != -9999.0]
    print(f"NDVI min={valid.min():.3f} max={valid.max():.3f} mean={valid.mean():.3f}")
    # Expected: min near -0.3 (bare soil / water), max near 0.85 (dense vegetation)

block_shapes: [(512, 512), (512, 512), (512, 512), (512, 512)]
band_count: 4
dtype: uint16
crs: 32632
Output key: processed/S2A_MSIL2A_..._x0000_y0000_ndvi.tif
Output size (bytes): 418720
Content-Type: image/tiff
NDVI min=-0.182 max=0.847 mean=0.421