Rendering a Live Rich Dashboard for Batch Raster Jobs

#!/usr/bin/env python3
"""
Live Rich dashboard over a multiprocessing raster reprojection batch.
Usage: python live_dashboard.py ./tiles ./out --crs EPSG:3857 --workers 4
"""
import os
import sys
import time
import argparse
from pathlib import Path
from dataclasses import dataclass
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed

import rasterio
from rasterio.warp import calculate_default_transform, reproject, Resampling

from rich.console import Group
from rich.live import Live
from rich.table import Table
from rich.progress import (
    Progress, BarColumn, TextColumn,
    TimeRemainingColumn, MofNCompleteColumn,
)


@dataclass
class TileResult:
    """Plain, picklable payload returned from a worker to the parent."""
    tile: str
    worker: int
    ok: bool
    seconds: float
    error: str | None = None


def reproject_tile(src_path: Path, dst_path: Path, dst_crs: str) -> TileResult:
    """Reproject one raster tile. Runs in a worker process — no Rich here."""
    worker = os.getpid()
    start = time.perf_counter()
    try:
        with rasterio.open(src_path) as src:
            transform, width, height = calculate_default_transform(
                src.crs, dst_crs, src.width, src.height, *src.bounds
            )
            profile = src.profile.copy()
            profile.update(
                crs=dst_crs, transform=transform,
                width=width, height=height,
                compress="lzw", tiled=True,
            )
            with rasterio.open(dst_path, "w", **profile) as dst:
                for band in range(1, src.count + 1):
                    reproject(
                        source=rasterio.band(src, band),
                        destination=rasterio.band(dst, band),
                        src_transform=src.transform, src_crs=src.crs,
                        dst_transform=transform, dst_crs=dst_crs,
                        resampling=Resampling.bilinear,
                    )
        elapsed = time.perf_counter() - start
        return TileResult(src_path.name, worker, True, elapsed)
    except Exception as exc:
        elapsed = time.perf_counter() - start
        return TileResult(src_path.name, worker, False, elapsed, str(exc))


def build_dashboard(progress: Progress, table: Table) -> Group:
    """Compose the two renderables into one block Live can refresh."""
    return Group(progress, table)


def render_table(stats: dict[int, dict], errors: int) -> Table:
    """Rebuild the per-worker table from accumulated stats."""
    table = Table(title=f"Per-worker throughput  (errors: {errors})",
                  expand=True, title_justify="left")
    table.add_column("Worker (PID)", justify="right", no_wrap=True)
    table.add_column("Tiles done", justify="right")
    table.add_column("Tiles/sec", justify="right")
    table.add_column("Errors", justify="right", style="red")
    for pid, s in sorted(stats.items()):
        rate = s["done"] / s["busy"] if s["busy"] > 0 else 0.0
        table.add_row(str(pid), str(s["done"]), f"{rate:0.2f}", str(s["errors"]))
    return table


def main() -> None:
    parser = argparse.ArgumentParser(description="Live dashboard raster batch")
    parser.add_argument("input_dir", type=Path)
    parser.add_argument("output_dir", type=Path)
    parser.add_argument("--crs", default="EPSG:3857", help="Target CRS")
    parser.add_argument("--workers", type=int, default=4)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    tiles = sorted(args.input_dir.glob("*.tif"))
    if not tiles:
        print(f"No .tif tiles in {args.input_dir}", file=sys.stderr)
        sys.exit(2)

    progress = Progress(
        TextColumn("[bold]Reprojecting tiles"),
        BarColumn(),
        MofNCompleteColumn(),
        TextColumn("ETA"),
        TimeRemainingColumn(),
        refresh_per_second=4,
    )
    task_id = progress.add_task("tiles", total=len(tiles))
    stats: dict[int, dict] = defaultdict(
        lambda: {"done": 0, "errors": 0, "busy": 0.0}
    )
    errors = 0

    # Live owns the terminal. Only this (parent) process renders.
    with Live(build_dashboard(progress, render_table(stats, errors)),
              refresh_per_second=4) as live:
        with ProcessPoolExecutor(max_workers=args.workers) as pool:
            futures = {
                pool.submit(reproject_tile, t,
                            args.output_dir / f"{t.stem}_3857.tif", args.crs): t
                for t in tiles
            }
            for future in as_completed(futures):
                result: TileResult = future.result()
                s = stats[result.worker]
                s["busy"] += result.seconds
                if result.ok:
                    s["done"] += 1
                else:
                    s["errors"] += 1
                    errors += 1
                    live.console.print(
                        f"[red]FAIL[/red] {result.tile}: {result.error}"
                    )
                progress.advance(task_id, 1)
                live.update(build_dashboard(progress, render_table(stats, errors)))

    total_ok = sum(s["done"] for s in stats.values())
    print(f"Completed {total_ok}/{len(tiles)} tiles, {errors} errors.")
    sys.exit(0 if errors == 0 else 12)   # 12 = partial batch failure


if __name__ == "__main__":
    main()

# Run the batch; capture the exit code
python live_dashboard.py ./tiles ./out --crs EPSG:3857 --workers 4
echo "exit code: $?"          # 0 = all tiles OK, 12 = partial failure

# Count outputs against inputs
echo "inputs:  $(ls ./tiles/*.tif | wc -l)"
echo "outputs: $(ls ./out/*_3857.tif | wc -l)"

# Confirm the target CRS landed on the first output
python3 - <<'EOF'
import rasterio
with rasterio.open(sorted(__import__("pathlib").Path("./out").glob("*_3857.tif"))[0]) as ds:
    print("CRS:", ds.crs)       # should print EPSG:3857
EOF