Track active shuffle by stage

[rynorris]

This is an optimization to #427

In the original PR we only track whether a shuffle dependency is active at the job-level, meaning we cannot scale up and down executors during a long job.

This PR extends the functionality by tracking the dependencies between stages and shuffles so we can mark the shuffle blocks as inactive earlier.

Unit tests included, and also tested on a local k8s cluster.

See toy example:

In this screenshot, stage 5.0 repartitions to 4 partitions, stage 6.0 to 2 partitions, and stage 7.0 to 1 partition. The exact query used was:

spark.range(0, 100).rdd
  .repartition(4).map(x => { Thread.sleep(1000); x })
  .repartition(2).map(x => { Thread.sleep(200); x })
  .repartition(1).map(x => { Thread.sleep(500); x })
  .collect()

You can see that once stage 6.0 is done, executors 6 and 8 are removed since their shuffle data is no longer needed. However executors 5 and 7 are both kept around during the last stage even though there's only one task, because they both hold shuffle data necessary for the final repartition to run.

[lwwmanning]

pushed a fix for a nit, otherwise lgtm