GDAL Batch Operations with multiprocessing.Pool

#!/usr/bin/env python3
"""
Batch raster reprojection with isolated multiprocessing.Pool workers.
Usage: python reproject_batch.py ./input ./output --crs EPSG:32633 --workers 4
"""
import os
import sys
import argparse
import logging
from pathlib import Path
from multiprocessing import Pool, cpu_count, set_start_method
from osgeo import gdal

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(processName)s: %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)


def init_gdal_worker() -> None:
    """Run once per worker process before any task is dispatched.

    Calling gdal.UseExceptions() turns silent C-level NULL returns into
    Python RuntimeError, making failures visible to the pool error handler.
    Setting GDAL_NUM_THREADS=1 prevents GDAL from spawning its own thread
    pool inside each worker, which would multiply CPU usage by O(workers * cores).
    """
    gdal.UseExceptions()
    os.environ["GDAL_NUM_THREADS"] = "1"       # cap internal GDAL threads
    os.environ["CPL_NUM_THREADS"] = "1"        # cap CPL/CURL threads
    os.environ["GDAL_DISABLE_READDIR_ON_OPEN"] = "YES"  # skip sibling-dir scan
    os.environ["VSI_CACHE"] = "FALSE"          # avoid duplicated VSI cache per worker
    if "GDAL_DATA" not in os.environ:
        os.environ["GDAL_DATA"] = "/usr/share/gdal"


def reproject_raster(task: tuple) -> dict:
    """Warp a single raster to the target CRS.

    Returns a dict so the main process can aggregate structured results
    for logging or downstream error-handling pipelines.
    """
    src_path, dst_path, target_crs = task
    try:
        src_ds = gdal.Open(str(src_path))
        if src_ds is None:
            raise RuntimeError(f"gdal.Open returned None for {src_path}")

        dst_ds = gdal.Warp(
            str(dst_path),
            src_ds,
            dstSRS=target_crs,          # e.g. "EPSG:32633"
            format="GTiff",
            creationOptions=[
                "TILED=YES",            # tiled layout for partial reads
                "COMPRESS=LZW",         # lossless, good ratio for elevation data
                "BIGTIFF=YES",          # avoids 4 GB limit on large mosaics
            ],
            numThreads=1,               # must match GDAL_NUM_THREADS
            resampleAlg="bilinear",
            errorThreshold=0.125,       # max warp error in output pixels
        )
        if dst_ds is None:
            raise RuntimeError(f"gdal.Warp produced no output for {src_path}")

        # Assign None to trigger GDALClose() immediately — do not rely on GC.
        # Skipping this leaks file descriptors in pools that process 1000+ files.
        dst_ds = None
        src_ds = None

        logging.info("OK %s", dst_path.name)
        return {"src": str(src_path), "ok": True, "error": None}
    except Exception as exc:
        logging.error("FAIL %s: %s", src_path, exc)
        return {"src": str(src_path), "ok": False, "error": str(exc)}


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Batch raster reprojection with isolated pool workers"
    )
    parser.add_argument("input_dir", type=Path, help="Directory containing source GeoTIFFs")
    parser.add_argument("output_dir", type=Path, help="Directory for reprojected outputs")
    parser.add_argument("--crs", default="EPSG:4326", help="Target CRS (default: EPSG:4326)")
    parser.add_argument(
        "--workers", type=int, default=cpu_count(),
        help="Worker process count (default: cpu_count)"
    )
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    tasks = [
        (src, args.output_dir / f"{src.stem}_reprojected.tif", args.crs)
        for src in sorted(args.input_dir.glob("*.tif"))
    ]
    if not tasks:
        logging.warning("No .tif files found in %s", args.input_dir)
        sys.exit(0)

    logging.info("Submitting %d tasks to %d workers", len(tasks), args.workers)

    # initializer= runs init_gdal_worker() once per worker before any task.
    # This is the correct hook — do NOT call GDAL init inside reproject_raster.
    with Pool(processes=args.workers, initializer=init_gdal_worker) as pool:
        results = pool.map(reproject_raster, tasks)

    ok = sum(1 for r in results if r["ok"])
    logging.info("Completed %d/%d successfully.", ok, len(tasks))
    sys.exit(0 if ok == len(tasks) else 1)


if __name__ == "__main__":
    # set_start_method must be called before Pool is created.
    # Placing it under __main__ prevents accidental re-invocation when
    # a spawned worker imports this module during startup.
    set_start_method("spawn", force=True)
    main()

# Check all outputs exist and are non-zero bytes
find ./output -name "*_reprojected.tif" -size 0 -print

# Confirm CRS on the first output file
python3 - <<'EOF'
from osgeo import gdal, osr
ds = gdal.Open("./output/scene_001_reprojected.tif")
srs = osr.SpatialReference()
srs.ImportFromWkt(ds.GetProjection())
print(srs.GetAuthorityCode(None))   # should print "32633" (or your target EPSG)
EOF

# Spot-check pixel count matches source dimensions after reprojection
gdalinfo ./output/scene_001_reprojected.tif | grep "Size is"