Spark GPU Backend

Configuration

geoembed.backends.spark.config

Configuration for the full-image GPU embedding pipeline.

EmbeddingsConfig(chip_size=224, batch_size=256, mixed_precision=True, model_name='dofa') dataclass

Configuration for the full-image GPU embedding pipeline.

The pipeline reads entire COGs, tiles on GPU via torch.unfold(), and runs batched inference — no per-chip I/O.

Attributes:

Name	Type	Description
chip_size	int	Square chip dimension in pixels (224 for DOFA).
batch_size	int	Inference batch size within each image (tune for VRAM).
mixed_precision	bool	Use FP16 autocasting for faster inference.
model_name	str	Registered model name (default "dofa").

Pipeline Orchestrator

geoembed.backends.spark.orchestrator

Spark-based full-image GPU embedding pipeline.

Distributes COG images across GPU workers via Spark mapInPandas. Each worker reads a full image, tiles on GPU via torch.unfold(), runs batched inference, and returns embeddings + spatial bounds.

Double-buffered I/O within each worker hides read latency behind GPU compute. Uses Spark's native GPU resource scheduling (spark.task.resource.gpu.amount).

EmbeddingsPipeline(config, input_storage, output_storage)

Full-image GPU embedding pipeline distributed via Spark mapInPandas.

Each Spark task: 1. Receives a batch of image paths 2. Reads full COGs, tiles on GPU, runs inference (with double-buffering) 3. Returns embeddings + spatial bounds

Spark handles multi-node GPU distribution via spark.task.resource.gpu.amount = 1.

Usage

from geoembed.backends.spark.config import EmbeddingsConfig from geoembed.backends.spark.orchestrator import EmbeddingsPipeline from geoembed.core.config import StorageConfig

config = EmbeddingsConfig(chip_size=224, batch_size=256, mixed_precision=True) pipeline = EmbeddingsPipeline( config=config, input_storage=StorageConfig(backend="delta", path="catalog.schema.metadata"), output_storage=StorageConfig(backend="delta", path="catalog.schema.embeddings"), ) pipeline.run()

Source code in src/geoembed/backends/spark/orchestrator.py

def __init__(
    self,
    config: EmbeddingsConfig,
    input_storage: StorageConfig,
    output_storage: StorageConfig,
):
    self.config = config
    self.input_storage = input_storage
    self.output_storage = output_storage

run()

Execute the full-image GPU embedding pipeline via Spark.

Source code in src/geoembed/backends/spark/orchestrator.py

def run(self) -> None:
    """Execute the full-image GPU embedding pipeline via Spark."""
    from pyspark.sql import SparkSession

    spark = SparkSession.builder.getOrCreate()

    print("[EmbeddingsPipeline] Starting full-image GPU pipeline...")
    start_time = time.monotonic()

    # Read metadata table — only need image_path
    quoted_input = quote_table_name(self.input_storage.path)
    source_df = (
        spark.read.table(quoted_input)
        .select("image_path")
        .filter("image_path IS NOT NULL AND id NOT LIKE 'ERROR_%'")
    )

    num_images = source_df.count()
    print(f"[EmbeddingsPipeline] {num_images} images to process.")

    worker = _SparkGPUWorker(
        chip_size=self.config.chip_size,
        batch_size=self.config.batch_size,
        mixed_precision=self.config.mixed_precision,
        model_name=self.config.model_name,
    )

    output_schema = (
        "chip_id STRING, parent_id STRING, parent_path STRING, "
        "col_off INT, row_off INT, "
        "embedding ARRAY<FLOAT>, "
        "minx DOUBLE, miny DOUBLE, maxx DOUBLE, maxy DOUBLE"
    )

    num_tasks = max(1, num_images // 4)
    result_df = source_df.repartition(num_tasks).mapInPandas(worker, schema=output_schema)

    quoted_output = quote_table_name(self.output_storage.path)
    writer = result_df.write.format("delta").mode(self.output_storage.mode)
    if self.output_storage.optimize_write:
        writer = writer.option("optimizeWrite", "true")
    writer.saveAsTable(quoted_output)

    if self.output_storage.cluster_columns:
        cols = ", ".join(self.output_storage.cluster_columns)
        spark.sql(f"ALTER TABLE {quoted_output} CLUSTER BY ({cols})")

    elapsed = time.monotonic() - start_time
    row_count = spark.read.table(quoted_output).count()
    print(
        f"[EmbeddingsPipeline] Complete: {row_count:,} embeddings "
        f"from {num_images} images in {elapsed:.1f}s"
    )

GPU Worker

geoembed.backends.spark.worker

Full-image GPU worker — reads COG, tiles on GPU, runs inference.

Single Ray actor that processes one image at a time: 1. Read full COG to CPU (via rasterio) — in background thread for double-buffering 2. Transfer to GPU as float32 tensor 3. Tile via torch.unfold() (zero-copy views, ~0ms) 4. Batch inference with DOFA (mixed precision) 5. Compute spatial bounds from affine transform (pure arithmetic) 6. Return Arrow table with embeddings + spatial metadata

Double-buffering: reads the next image while GPU processes the current one, hiding ~2.7s I/O latency behind ~3.2s GPU compute.

FullImageGPUWorker(chip_size=224, batch_size=256, mixed_precision=True, model_name='dofa')

Ray MapBatches actor: full-image read → GPU tile → inference → Arrow output.

Processes images one at a time with double-buffered I/O. Each call receives a batch of image paths and returns an Arrow table with all chip embeddings and spatial bounds.

Source code in src/geoembed/backends/spark/worker.py

def __init__(
    self,
    chip_size: int = 224,
    batch_size: int = 256,
    mixed_precision: bool = True,
    model_name: str = "dofa",
):
    self.chip_size = chip_size
    self.batch_size = batch_size
    self.mixed_precision = mixed_precision
    self.model_name = model_name
    self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    self.embedder: Embedder | None = None
    self.executor = ThreadPoolExecutor(max_workers=1)

__call__(batch)

Entry point for Ray map_batches.

Receives a batch of image paths, processes each with double-buffered I/O, returns a combined Arrow table with all chip embeddings + spatial metadata.

Source code in src/geoembed/backends/spark/worker.py

def __call__(self, batch: dict[str, np.ndarray]) -> pa.Table:
    """
    Entry point for Ray map_batches.

    Receives a batch of image paths, processes each with double-buffered I/O,
    returns a combined Arrow table with all chip embeddings + spatial metadata.
    """
    if self.embedder is None:
        self._setup_model()

    image_paths = batch["image_path"]

    if len(image_paths) == 0:
        return pa.table({
            "chip_id": pa.array([], type=pa.string()),
            "parent_id": pa.array([], type=pa.string()),
            "parent_path": pa.array([], type=pa.string()),
            "embedding": pa.array([], type=pa.list_(pa.float32())),
            "minx": pa.array([], type=pa.float64()),
            "miny": pa.array([], type=pa.float64()),
            "maxx": pa.array([], type=pa.float64()),
            "maxy": pa.array([], type=pa.float64()),
        })

    tables: list[pa.Table] = []
    future = self.executor.submit(_load_image, str(image_paths[0]))

    for i in range(len(image_paths)):
        image_np, metadata = future.result()

        if i + 1 < len(image_paths):
            future = self.executor.submit(_load_image, str(image_paths[i + 1]))

        table = self._process_single_image(image_np, metadata)
        tables.append(table)
        del image_np

    return pa.concat_tables(tables)