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

While a Spark session runs, its UI serves the same numbers as JSON at /api/v1: the application's stages, and each stage's tasks with their time and shuffle bytes. job_report.py runs a job, reads the stages it ran and prints 1 line for each, flagging a stage where 1 task reads over 5 times the median, which is skew, and a stage with many tasks of under 1 MB each. It needs classic Spark, since Spark Connect has no sparkContext.

- Page: https://tecyfy.com/data-engineering/course/how-spark-runs-a-job/build-it
- Course: Data engineering course (https://tecyfy.com/data-engineering/course), Module 3: How Spark runs a job, Build it: a job report
- Parts of this module: Jobs, stages and tasks (https://tecyfy.com/data-engineering/course/how-spark-runs-a-job); Shuffles and partitions (https://tecyfy.com/data-engineering/course/how-spark-runs-a-job/shuffles-and-partitions); Adaptive query execution (https://tecyfy.com/data-engineering/course/how-spark-runs-a-job/adaptive-query-execution); Build it: a job report (https://tecyfy.com/data-engineering/course/how-spark-runs-a-job/build-it)

At the end of this part you can:

- Read a job's stages and tasks from the Spark UI's REST API
- Flag skew and stages of too many small tasks from a script

The checks in [module 2's](/data-engineering/course/pyspark-dataframes/build-it) 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](/data-engineering/course/how-spark-runs-a-job/adaptive-query-execution), so run `aqe.py` first. Then save this as `job_report.py`:

```python
"""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`:

```text
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.

- Quiz: [PySpark Performance Tuning](https://tecyfy.com/quiz/pyspark/pyspark-performance-tuning). Master the key techniques for optimizing PySpark applications, improving execution speed, and handling large-scale data efficiently.

## Check yourself

1. job_report.py flags a stage of 200 tasks reading 1 KB each with adaptive query execution off. What change fixes it with the least effort?

   A. Add repartition(200) before the groupBy
   B. Turn adaptive query execution back on, so Spark merges the small partitions after the shuffle, or lower spark.sql.shuffle.partitions for this job
   C. Set spark.sql.files.maxPartitionBytes to 1 MB
   D. Broadcast the result of the groupBy

### Answers

1. B. The 200 tasks come from spark.sql.shuffle.partitions, which Spark uses as it is when adaptive query execution is off. With it on, Spark merges the 200 small partitions into a few once it knows their size; with it off, a lower number set for the job does the same by hand. repartition would add another shuffle with the same 200 partitions, and maxPartitionBytes only changes how files are read.

## Questions

### How do I get Spark job metrics from Python?

Read the Spark UI's REST API. Its address is spark.sparkContext.uiWebUrl plus /api/v1/applications/ and the application id, and its endpoints list the jobs, the stages, and each stage's tasks through the taskList of each stage, with each task's duration and shuffle read and write. The same API serves finished applications from Spark's history server when event logging is on.

### Can I read the Spark UI's REST API on Databricks serverless?

No. Serverless runs Spark Connect, which has no sparkContext and no Spark UI, so there is no address to read. Use the query profile under the cell instead, which shows each operator's time, rows and memory.

## Sources

- [Monitoring and instrumentation, REST API](https://spark.apache.org/docs/latest/monitoring.html#rest-api)
- [Performance tuning, adaptive query execution](https://spark.apache.org/docs/latest/sql-performance-tuning.html)
- [Serverless compute limitations](https://docs.databricks.com/aws/en/compute/serverless/limitations)