VehicleMapping HD AV Spatial Data

Parallelizing Lane-Attribute Extraction with Dask

Extracting width, heading, curvature, and marking attributes for every lane in a metropolitan HD map is embarrassingly parallel — each tile is independent — but run serially it takes hours, and run naively in one process it OOMs on the first dense downtown tile. This task scales batch lane-attribute extraction across a Dask cluster: partition tiles into a bag, map a pure extractor, and bound per-worker memory so the whole map finishes in minutes on hardware that could never hold it at once.

Tiles fan out to workers as a Dask bag; each worker writes its own partition, bounded by a RAM ceiling:

Distributed Lane-Attribute Extraction with Dask Tile store box on the left connects to a Dask bag box. The bag fans out to three worker boxes stacked vertically, each labelled extractor under a RAM limit. Each worker connects to a partitioned GeoParquet output box on the right. A note states only metadata returns to the driver. Tile store GeoParquet Dask bag 1 tile / partition Worker · extractor RAM ≤ limit Worker · extractor RAM ≤ limit Worker · extractor RAM ≤ limit Partitioned GeoParquet out only metadata returns to the driver

Prerequisites #

  • Python 3.10+, Dask 2024.1+ (dask.bag, distributed), GeoPandas 0.14+, PyArrow 14+, Shapely 2.0+.
  • Input: a directory of GeoParquet tiles, one file per map tile, in one projected CRS.
  • Upstream stage: lane geometry and boundaries already vectorized; the single-tile extractor from automating lane-width attribute sync exists as a pure function.
  • Output: partitioned GeoParquet attribute tables written directly by the workers.

Step-by-Step #

1. Build a bag of tile partitions #

Wrap the tile paths in a Dask bag with one partition per tile so each unit of work is independent and sized to a worker.

python
import dask.bag as db
import glob

def tile_bag(tile_dir: str) -> db.Bag:
    paths = sorted(glob.glob(f"{tile_dir}/*.parquet"))
    # npartitions == number of tiles => one tile of work per task
    return db.from_sequence(paths, npartitions=len(paths))

Key parameter: npartitions=len(paths) makes each partition exactly one tile, the natural unit of spatial independence. Expected output: a dask.bag.Bag of file paths.

2. Map a pure per-tile extractor #

The mapped function must be side-effect-free and return only what is needed. It reads a tile, extracts attributes, and writes its own partition — returning a small manifest, not the data.

python
import geopandas as gpd

def extract_tile(path: str, out_dir: str) -> dict:
    gdf = gpd.read_parquet(path)
    gdf["width"]   = gdf.geometry.map(_lane_width)      # per-lane extractors
    gdf["heading"] = gdf.geometry.map(_lane_heading)
    out = f"{out_dir}/{path.rsplit('/', 1)[-1]}"
    gdf.to_parquet(out)                                 # worker writes its own file
    return {"tile": path, "n_lanes": len(gdf), "out": out}

def run(tile_dir, out_dir):
    return tile_bag(tile_dir).map(extract_tile, out_dir=out_dir)

Key parameter: extract_tile writes directly to out_dir and returns only a manifest dict, so no lane geometry flows back to the driver.

3. Bound worker memory and compute #

Configure the Dask cluster with an explicit per-worker memory limit so Dask spills or pauses before a worker OOMs, then compute the manifests.

python
from dask.distributed import Client

def process(tile_dir, out_dir):
    client = Client(n_workers=8, threads_per_worker=1, memory_limit="4GB")
    manifests = run(tile_dir, out_dir).compute()   # returns small dicts only
    client.close()
    return manifests

Key parameters: memory_limit="4GB" is the per-worker RAM ceiling; threads_per_worker=1 avoids GIL contention in the geometry code. Expected output: a list of per-tile manifest dicts.

Verification & Acceptance Criteria #

Confirm every tile was processed and no worker breached its ceiling.

python
def assert_complete(manifests, tile_dir):
    import glob
    expected = {p.rsplit("/", 1)[-1] for p in glob.glob(f"{tile_dir}/*.parquet")}
    done = {m["out"].rsplit("/", 1)[-1] for m in manifests}
    assert expected == done, f"missing tiles: {expected - done}"
    assert all(m["n_lanes"] > 0 for m in manifests), "empty tile output"
    print(f"processed {len(manifests)} tiles, {sum(m['n_lanes'] for m in manifests)} lanes")

Acceptance gate: every input tile has a corresponding output file; no worker exceeded memory_limit (check the dashboard's worker memory plot stays below the pause threshold); and total lane count matches the source. A single missing tile means a task failed silently — inspect the worker logs.

Common Errors & Fixes #

Workers pause and the job stalls. Partitions are too large for memory_limit. Shrink partitions to one tile, or raise the worker memory ceiling; watch the dashboard memory plot to confirm workers stay below the spill threshold.

KilledWorker on a dense tile. A downtown tile exceeded RAM before Dask could spill. Split oversized tiles upstream or process them in row-group chunks inside extract_tile rather than loading the whole GeoDataFrame.

Driver runs out of memory on .compute(). The mapped function returned data, not a manifest. Ensure extract_tile writes output itself and returns only small dicts; never compute() a bag of full GeoDataFrames.

Non-deterministic output ordering. Bag partitions complete out of order. If order matters downstream, sort by tile id when reading the partitioned output — do not rely on completion order, which mirrors the chunked streaming pattern used elsewhere.

FAQ #

Why Dask instead of multiprocessing for this? #

multiprocessing scales to one machine's cores; lane-attribute extraction over a metropolitan map is often larger than one node's memory and benefits from spanning several machines. Dask gives the same map-over-partitions model but with a distributed scheduler, spill-to-disk when a worker nears its memory limit, and a dashboard to see stragglers. On a single machine the local scheduler is a drop-in, so the same code runs on a laptop and scales out to many machines unchanged.

How big should each partition be? #

Size a partition to one tile, or a small group of tiles, so its peak memory stays well under a single worker's RAM limit while still amortizing scheduler overhead. Partitions that are too large risk a worker OOM on a dense downtown tile; partitions that are too small drown the scheduler in tiny tasks. One tile per partition is the usual sweet spot because tiles are already the natural unit of spatial locality and independence.

How do I avoid pulling everything back to the driver? #

Never call compute() on the full result if it materializes every record on the driver. Instead write each partition's output directly from the worker with a to_parquet-style sink, so results stream to storage and only small metadata returns to the driver. Collecting a whole metropolitan attribute table to the driver defeats the point of distributing the work and reintroduces the single-node memory ceiling you were trying to escape.

Up one level: Batch Lane-Attribute Extraction — the parent workflow whose per-tile extractor this scales across a Dask cluster.