Build a Spark job report in Python from the Spark UI's REST API

The checks in module 2's script say whether a job's output is right. This part adds the other half: a report on how the job ran, so that a stage that is getting slow gets noticed before it fails. It reads the same numbers the Spark UI shows, from the UI's own REST API.

Build it: report every stage of a job

It runs on the guest orders from part 3, so run aqe.py first. Then save this as job_report.py:

"""Run a Spark job, then report what each of its stages did and flag the problems to look at."""
import json
import statistics
import urllib.request

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

SKEW_FACTOR = 5           # Spark's own rule: a partition over 5 times the median is skewed
SMALL_TASK = 1024 * 1024  # under 1 MB a task, the smallest partition Spark merges to


def api(spark: SparkSession, path: str):
    """Read 1 endpoint of the Spark UI's REST API for this application."""
    sc = spark.sparkContext
    with urllib.request.urlopen(f"{sc.uiWebUrl}/api/v1/applications/{sc.applicationId}/{path}") as reply:
        return json.load(reply)


def size(n: float) -> str:
    if n >= 1e5:
        return f"{n / 1e6:.1f} MB"
    return f"{n / 1e3:.0f} KB" if n >= 1e3 else f"{n:.0f} B"


def run_and_report(spark: SparkSession, name: str, job) -> None:
    """Run job(), then print 1 line per stage it ran, with a flag where something looks wrong."""
    before = max((s["stageId"] for s in api(spark, "stages")), default=-1)
    job()
    cores = spark.sparkContext.defaultParallelism
    print(f"\n{name}")
    stages = [s for s in api(spark, "stages?status=complete") if s["stageId"] > before]
    for stage in sorted(stages, key=lambda s: s["stageId"]):
        tasks = api(spark, f"stages/{stage['stageId']}/{stage['attemptId']}/taskList?length=100000")
        reads = [t["taskMetrics"]["shuffleReadMetrics"]["localBytesRead"]
                 + t["taskMetrics"]["shuffleReadMetrics"]["remoteBytesRead"] for t in tasks]
        times = [t["duration"] for t in tasks]
        line = (f"  stage {stage['stageId']:>3}: {len(tasks):>3} task{'s' if len(tasks) > 1 else ' '}, read {size(stage['shuffleReadBytes']):>8}, "
                f"write {size(stage['shuffleWriteBytes']):>8}, median {statistics.median(times):>4.0f} ms, slowest {max(times):>4} ms")
        median_read = statistics.median(reads)
        if median_read > 0 and max(reads) > SKEW_FACTOR * median_read:
            line += f"  <- skew: 1 task reads {max(reads) / median_read:.0f}x the median"
        elif len(tasks) > 2 * cores and 0 < stage["shuffleReadBytes"] / len(tasks) < SMALL_TASK:
            line += f"  <- {len(tasks)} small tasks of {size(stage['shuffleReadBytes'] / len(tasks))}"
        print(line)


def main(spark: SparkSession) -> None:
    orders = spark.read.parquet("lake/orders_with_guests")
    customers = spark.read.parquet("lake/customers")

    def revenue_by_country():
        (orders.join(customers, "customer_id")
         .groupBy("country_id").agg(F.sum("amount").alias("revenue"))
         .write.mode("overwrite").parquet("lake/revenue_by_country"))

    # As in part 3: a customer table too big to broadcast, and limits scaled to this small data.
    spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")
    spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "1m")
    spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "1m")

    for name, adaptive, skew in [
        ("Adaptive query execution off", "false", "false"),
        ("Adaptive on, skew join fix off", "true", "false"),
        ("Adaptive on, skew join fix on", "true", "true"),
    ]:
        spark.conf.set("spark.sql.adaptive.enabled", adaptive)
        spark.conf.set("spark.sql.adaptive.skewJoin.enabled", skew)
        run_and_report(spark, name, revenue_by_country)


if __name__ == "__main__":
    session = SparkSession.builder.master("local[4]").appName("job-report").getOrCreate()
    session.sparkContext.setLogLevel("ERROR")
    main(session)

Run it with python job_report.py:

Adaptive query execution off
  stage   2:   1 task , read      0 B, write   0.5 MB, median  314 ms, slowest  314 ms
  stage   3:   4 tasks, read      0 B, write  12.1 MB, median  390 ms, slowest  396 ms
  stage   4: 200 tasks, read  12.6 MB, write   0.2 MB, median   13 ms, slowest  533 ms  <- skew: 1 task reads 357x the median
  stage   5: 200 tasks, read   0.2 MB, write      0 B, median    6 ms, slowest   98 ms  <- 200 small tasks of 1 KB

Adaptive on, skew join fix off
  stage   6:   4 tasks, read      0 B, write  12.1 MB, median  119 ms, slowest  122 ms
  stage   7:   1 task , read      0 B, write   0.5 MB, median   46 ms, slowest   46 ms
  stage  10:   5 tasks, read  12.6 MB, write     6 KB, median   52 ms, slowest  233 ms  <- skew: 1 task reads 8x the median
  stage  14:   1 task , read     6 KB, write      0 B, median   22 ms, slowest   22 ms

Adaptive on, skew join fix on
  stage  15:   1 task , read      0 B, write   0.5 MB, median   51 ms, slowest   51 ms
  stage  16:   4 tasks, read      0 B, write  12.1 MB, median  110 ms, slowest  113 ms
  stage  19:   9 tasks, read  12.6 MB, write    11 KB, median   37 ms, slowest   77 ms
  stage  23:   1 task , read    11 KB, write      0 B, median   21 ms, slowest   21 ms

Your times will differ from these, and from run to run, but the tasks, the bytes and the flags will not, though 2 stages that start together can swap numbers. The first run, with adaptive query execution off, has both problems: the join's 200 tasks include 1 that reads 357 times the median, the guest orders, and the final stage is 200 tasks of about 1 KB for 30 countries. With adaptive query execution on, the tiny tasks are merged, and only the skew is left. With the skew fix, the join is 9 tasks and nothing is flagged.

How it works

The Spark UI is also an API. Every page of it is served as JSON under /api/v1, at the address in spark.sparkContext.uiWebUrl. stages?status=complete lists the stages that finished, and taskList lists a stage's tasks, each with its time and shuffle bytes. The script notes the last stage id before the job and reports only the stages after it, so it can measure 1 job in a session that has run others.

The flags are this module's 2 problems. Skew is 1 task reading over 5 times the median, Spark's own factor for a skewed partition. Too many small tasks is more than 2 tasks per core reading under 1 MB each, the size Spark merges partitions up to. Both look only at bytes, never at time, because bytes are the same on every run and times are not.

It needs classic Spark. Spark Connect has no sparkContext: asking for it fails with JVM_ATTRIBUTE_NOT_SUPPORTED, because it depends on the JVM. So the script runs on your computer, not on Databricks serverless, where the query profile is the place to look.

Then try 2 changes. Delete the line that turns broadcasting off: the customers are broadcast, the guest orders never leave their tasks, and both skew flags go. Only the run without adaptive query execution is still flagged, for 200 tasks that read 24 bytes each.

Then put that line back and delete the 2 that scale the limits down. With the defaults, Spark merges the join's partitions into 3 tasks of about 4 MB, and does not count the guest partition as skewed, since it is far below 256 MB. The script's flag goes too, because among 3 tasks none can read 5 times the median, yet the slowest task still takes far longer than the others. A ratio of bytes finds skew among many tasks, not among a few, which is why the report prints the times as well.

Where this goes next

Module 12 comes back to this with real costs: which step of a pipeline takes the time and money, and how to make it cheaper. Module 4 starts the lakehouse, with tables that can be updated in place instead of rewritten.

QuizPySpark Performance TuningMaster the key techniques for optimizing PySpark applications, improving execution speed, and handling large-scale data efficiently.