Add a PySpark shuffle script for users without Google Cloud

[yinshiyi]

Add a PySpark shuffle script for users without Google Cloud

DeepVariant training needs the examples shuffled globally before training. The
upstream script, tools/shuffle_tfrecords_beam.py (v1.10), supports two ways to
run it:

• Google Cloud Dataflow. Scales to large datasets, but only for users with
  a GCP account.
• Beam DirectRunner, on one machine. Keeps the entire shuffle in memory, so
  memory grows with the dataset: the case study's 1.97 GB (gzipped) training set
  needed 32.6 GB of RAM. Large training sets need a very large-memory machine or
  fail with out-of-memory errors.

Running the Beam script on AWS EMR is not a practical alternative: Beam's Python
Spark runner can't submit to EMR's YARN directly and needs Docker on every
worker.

What this PR adds

tools/shuffle_tfrecords_spark.py is a plain PySpark version of the same
SHA1-keyed shuffle, with the same output files (*.tfrecord.gz,
dataset_config.pbtxt, example_info.json). It runs on a single machine or on
AWS EMR (tools/aws_shuffle/run_emr_ec2.sh). Spark spills the shuffle to disk,
so memory depends on the cores used, not on dataset size.

It also adds two notes to the training case study for local Beam users:

• DirectRunner jobs longer than 5 minutes fail with DEADLINE_EXCEEDED unless
  --job_server_timeout is raised (Change default timeout and add heartbeat logging apache/beam#36517). The validation shuffle
  (2.5 min) passes; a local training shuffle (10 min) fails.
• --direct_num_workers=0 has no effect with Beam 2.71: DirectRunner runs on
  Prism, and =0 and =1 gave the same run time.

Results (case study training set, 342,758 examples, same 64-core server)

	Beam DirectRunner	Spark, 4 cores	Spark, 64 cores, 64 shards
Wall time	9:52	13:09	1:45
Peak memory	32.6 GB	3.9 GB	38.9 GB

Beam stays at about 2-4 busy cores whatever the settings; Spark scales with the
cores you give it, and fewer cores means less memory.

Testing

• Correctness: every output keeps exactly the input records (by SHA1), and
  the label counts in dataset_config.pbtxt match.
• Randomness: the share of output neighbours that were also neighbours in
  the input matches random shuffles of the same records (0.06% vs 0.06% on
  training).
• Reproducibility: the Spark output is deterministic. Repeated runs, 4- vs
  64-core runs, and a run on AWS EMR all gave identical records in identical
  order. Beam's order changes on every run.
• AWS EMR: run_emr_ec2.sh shuffled the validation set on a 2-node
  emr-7.5.0 cluster in about 9 minutes end to end, for about $0.25.

Notes

• The default shard count (~256 MB of gzipped input per shard) leaves cores
  idle on small inputs; set --num_output_shards to at least the Spark core count.
  eg. local[4], set --num_output_shards 4
• Changing the shard count changes the output, so keep it fixed for
  reproducibility.
• The EMR test used the account's own roles via SERVICE_ROLE and
  INSTANCE_PROFILE

Details and setup are in tools/aws_shuffle/README.md.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[pichuan]

Hi @yinshiyi, thanks for sending out the PR.

As you know, we usually accept PRs by incorporating the change internally with the author's consent, and we credit the author. However, since we can't test on AWS, we won't be able to take this PR. I agree it can be useful for other users, though. Would it be OK if we keep it open for reference without incorporating it directly?

[yinshiyi]

Hi @yinshiyi, thanks for sending out the PR.

As you know, we usually accept PRs by incorporating the change internally with the author's consent, and we credit the author. However, since we can't test on AWS, we won't be able to take this PR. I agree it can be useful for other users, though. Would it be OK if we keep it open for reference without incorporating it directly?

Great to hear from you, @pichuan. Yes, I still remember the g3/CLs.

The new Spark shuffler runs on any single machine: a local server, a GCP instance, an AWS instance, and so on.

The main problem I wanted to solve is that the current Beam shuffler, when run with DirectRunner, holds the whole dataset in RAM, which isn't feasible for large datasets. The proposed shuffler spills to disk instead, so a patient user can shuffle large dataset with almost any machine that has enough storage. As a bonus, its output is reproducible: the same input produces consistently across platforms.

When you have time, I would appreciate it you could test the local Spark shuffler in your setup. If you run into any issues, please let me know. If you think it is useful, I could focus this PR only for local shuffler.

I also think two issues I found are worth following up on:

1. With recent Beam versions (2.61+), DirectRunner runs on the Prism runner, and --job_server_timeout (default 300 s) acts as a deadline on the whole job. Any local shuffle that runs longer fails with DEADLINE_EXCEEDED while the job is still running. The validation-set shuffle (~2.5 min) never hits it, but my local training-set shuffle failed at 5:03 and completed in 9:52 with --job_server_timeout=86400.
2. --direct_num_workers makes no difference on my system. On Prism, only the older FnApiRunner reads it: --direct_num_workers=1 and =0 both took about 2:26 on the validation set (64-core machine).

---

AWS EMR support came along with the Spark work. I'm happy to leave the EMR part as a separate open PR for other users. If you think it's appropriate to mention it in one of the case studies, that could help others as well.