Setting Per-Task Timeouts for GDAL Operations

from __future__ import annotations

import math
from dataclasses import dataclass

# Measured on the reference worker: seconds per megapixel of OUTPUT, by resampling
# method. Re-measure after any change to the image, the instance type or GDAL.
SECONDS_PER_MEGAPIXEL = {
    "nearest": 0.35,
    "bilinear": 0.55,
    "cubic": 0.95,
    "cubicspline": 1.60,
    "lanczos": 2.10,
    "average": 0.70,
}

# Remote sources add latency per block read; /vsis3/ and /vsicurl/ are far from free.
SOURCE_PENALTY = {"local": 1.0, "vsis3": 2.4, "vsicurl": 3.2, "vsigs": 2.4}


@dataclass(frozen=True)
class WarpShape:
    """Everything about a warp that predicts its cost."""

    width: int
    height: int
    band_count: int
    resampling: str = "bilinear"
    source_kind: str = "local"
    compress: str = "DEFLATE"

    @property
    def megapixels(self) -> float:
        return (self.width * self.height * self.band_count) / 1_000_000


def estimate_seconds(shape: WarpShape) -> float:
    """A proportional cost estimate — deliberately rough, never a promise."""
    base = shape.megapixels * SECONDS_PER_MEGAPIXEL.get(shape.resampling, 0.55)
    base *= SOURCE_PENALTY.get(shape.source_kind, 1.0)
    if shape.compress in ("DEFLATE", "ZSTD"):
        base *= 1.25          # compression is real work on the write side
    elif shape.compress in ("LZW",):
        base *= 1.15
    # A floor: process start-up, PROJ database load and file opening cost ~2 s
    # regardless of how little there is to warp.
    return max(2.0, base)


def timeout_for(shape: WarpShape, multiplier: float = 6.0, ceiling: float = 1800.0) -> float:
    """The limit to hand to the subprocess runner.

    The multiplier is wide on purpose. The timeout's job is to catch a HANG —
    an operation that will never finish — not to enforce a performance target.
    Anything tight enough to catch a slow tile will kill a busy worker.
    """
    return min(ceiling, math.ceil(estimate_seconds(shape) * multiplier))

import rasterio
from rasterio.warp import calculate_default_transform


def shape_for_tile(source_uri: str, dst_crs: str, resampling: str) -> WarpShape:
    """Inspect the source cheaply — this opens headers, not pixels."""
    with rasterio.open(source_uri) as src:
        transform, width, height = calculate_default_transform(
            src.crs, dst_crs, src.width, src.height, *src.bounds
        )
        kind = (
            "vsis3" if source_uri.startswith("/vsis3/")
            else "vsicurl" if source_uri.startswith("/vsicurl/")
            else "local"
        )
        return WarpShape(
            width=width, height=height, band_count=src.count,
            resampling=resampling, source_kind=kind,
        )


def warp_tile(source_uri: str, dest: str, budget, dst_crs: str = "EPSG:3857") -> None:
    shape = shape_for_tile(source_uri, dst_crs, resampling="bilinear")
    run_gdal(
        [
            "gdalwarp",
            "-t_srs", dst_crs,
            "-r", shape.resampling,
            "-dstnodata", "-9999",          # never let GDAL pick this for you
            "-co", "COMPRESS=DEFLATE",
            "-co", "TILED=YES",
            "-wm", "512",                    # bound GDAL's own warp memory
            source_uri, dest,
        ],
        budget=budget,
        want_seconds=timeout_for(shape),
    )

def test_estimate_scales_with_pixels() -> None:
    small = WarpShape(width=1024, height=1024, band_count=1)
    large = WarpShape(width=4096, height=4096, band_count=1)
    # 16x the pixels should be roughly 16x the estimate, not 2x and not 200x.
    ratio = estimate_seconds(large) / estimate_seconds(small)
    assert 12 < ratio < 20


def test_resampling_and_source_both_matter() -> None:
    base = WarpShape(width=2048, height=2048, band_count=3, resampling="nearest")
    fancy = WarpShape(width=2048, height=2048, band_count=3, resampling="lanczos")
    remote = WarpShape(width=2048, height=2048, band_count=3, source_kind="vsis3")
    assert estimate_seconds(fancy) > estimate_seconds(base) * 4
    assert estimate_seconds(remote) > estimate_seconds(base)


def test_timeout_never_exceeds_the_ceiling() -> None:
    absurd = WarpShape(width=200_000, height=200_000, band_count=4)
    assert timeout_for(absurd) == 1800.0

-- One row per warp. The interesting column is the last one.
SELECT tile_key, megapixels, resampling, limit_seconds, actual_seconds,
       round(actual_seconds / NULLIF(limit_seconds, 0), 3) AS used_fraction
FROM   warp_runs
WHERE  started_at > now() - interval '7 days'
ORDER  BY used_fraction DESC
LIMIT  20;