This pull request fixes a subtle bug that occurs reliably for me on Google Cloud Dataflow. When using save_main_session, Python SDK workers will attempt to download a pickled_main_session file from Google Cloud Storage, and will use that file to reload the __main__ session. Pipelines that depend on this functionality expect names defined in __main__ to be present at runtime, and will fail with NameErrors if the main session is not available.

Unfortunately, roughly 0.5% of the time (in my experience), the Python SDK worker fails to download pickled_main_session from GCS to the local worker machine, and an empty file is created in its place. This causes Dill to log an EOFError¹, but then continues to load and run the user code, which then fails due to missing names. For small pipelines (under 20 workers) this almost never occurs; but larger pipelines with hundreds of workers often have at least one machine fail to load the __main__ session due to this bug, and the subsequent NameError exceptions cause the entire batch job to fail.

This patch does three things:

• checks to see if the pickled main session file is empty
• if so, raises a new type of exception (CorruptMainSessionException)
• catches that exception, logs it, and then re-raises, causing the worker to fail to boot.

Admittedly, I don't know if this is the right solution or if it will generalize to all Beam use cases, but I've found this patch to work effectively in my team's pipelines. If a corrupt main session is detected, the worker fails to boot, Cloud Dataflow restarts the container, and the pickled main session loads correctly on the second try. An alternative implementation might be to ensure that the pickled main session's file size matches the size provided by the remote filesystem.

R: @robertwb
R: @lukecwik

¹ Example traceback from an empty main session file:

Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/apache_beam/runners/worker/sdk_worker_main.py", line 131, in main
    _load_main_session(semi_persistent_directory)
  File "/usr/local/lib/python3.8/site-packages/apache_beam/runners/worker/sdk_worker_main.py", line 236, in _load_main_session
    pickler.load_session(session_file)
  File "/usr/local/lib/python3.8/site-packages/apache_beam/internal/pickler.py", line 318, in load_session
    return dill.load_session(file_path)
  File "/usr/local/lib/python3.8/site-packages/dill/_dill.py", line 368, in load_session
    module = unpickler.load()
  File "/usr/local/lib/python3.8/site-packages/dill/_dill.py", line 472, in load
    obj = StockUnpickler.load(self)
EOFError: Ran out of input

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

• Choose reviewer(s) and mention them in a comment (R: @username).
• Format the pull request title like [BEAM-XXX] Fixes bug in ApproximateQuantiles, where you replace BEAM-XXX with the appropriate JIRA issue, if applicable. This will automatically link the pull request to the issue.
• Update CHANGES.md with noteworthy changes.
• If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

ValidatesRunner compliance status (on master branch)

Lang	ULR	Dataflow	Flink	Samza	Spark	Twister2
Go	---			---		---
Java
Python	---			---		---
XLang				---		---

Examples testing status on various runners

Lang	ULR	Dataflow	Flink	Samza	Spark	Twister2
Go	---	---	---	---	---	---	---
Java	---		---	---	---	---	---
Python	---	---	---	---	---	---	---
XLang	---	---	---	---	---	---	---

Post-Commit SDK/Transform Integration Tests Status (on master branch)

Go	Java	Python

Pre-Commit Tests Status (on master branch)

---	Java	Python	Go	Website	Whitespace	Typescript
Non-portable
Portable	---			---	---	---

See .test-infra/jenkins/README for trigger phrase, status and link of all Jenkins jobs.

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI.