Skip to content

[YAML] Fix flaky ErrorHandlingTest: don't pickle Scope into user DoFn closures - #40467

Merged
shunping merged 2 commits into
apache:masterfrom
shunping:fix-flaky-test-2
Oct 8, 2026
Merged

shunping merged 2 commits into
apache:masterfrom
shunping:fix-flaky-test-2

Conversation

@shunping

@shunping shunping commented Oct 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

Fixes flaky ErrorHandlingTest.test_error_handling_outputs and ErrorHandlingTest.test_strip_error_metadata in yaml_transform_test.py, which intermittently fail with:

    RuntimeError: Unable to pickle fn CallableWrapperDoFn(...):
    cannot pickle 'google._upb._message.Descriptor' object

Root cause

  • Scope.create_ptransform monkey-patches annotations and expand onto the user's transform instance. The patched expand (recording_expand) closes over the YAML Scope, which references every provider, including the process-global standard_providers() (cached via functools.cache).
  • Any DoFn that captures its enclosing transform (e.g. SizeLimiter's raise_on_big uses self._limit) therefore pickles the whole Scope and all providers. Providers hold lazily-populated, process-wide state (e.g. ExternalProvider._service / _schema_transforms), so whether the pickle succeeds depends on which tests ran earlier on the same xdist worker. Hence the flakiness.
  • Locally, the pickled SizeLimiter closure is ~90 KB and contains the Scope, 28 ExternalJavaProviders, 20 RenamingProviders, etc., all to carry a single int.

Fix

Add Scope.__reduce__ returning a string stub, mirroring the existing Pipeline.__reduce__.
The Scope is only used during pipeline construction:

  • recording_expand only runs during runner.apply_PTransform; workers never call a composite's expand.
  • Annotations are read once when the transform is applied and encoded into the pipeline proto.
  • Transforms restored via PICKLED_TRANSFORM are never re-expanded.

This keeps the existing monkey-patching unchanged, so there's no change to transform labels (update compatibility) or to type-hint behavior. After the fix, the pickled closure drops from ~90 KB to ~3.9 KB.

…oFn closures

A user transform's DoFn that captures `self` also captures the closures
YAML patches onto it, which hold the Scope and every provider. Any
unpicklable provider state (e.g. a protobuf Descriptor) then makes
pickling fail. This test injects such state to reproduce the flaky
ErrorHandlingTest failures deterministically.
@shunping shunping changed the title Fix flaky test 2 [YAML] Fix flaky ErrorHandlingTest: don't pickle Scope into user DoFn closures Oct 8, 2026
@shunping

shunping commented Oct 8, 2026 •

Copy link
Copy Markdown
Collaborator Author

r: @derrickaw as the author of the previous attempt (#39845).

@shunping
shunping marked this pull request as ready for review October 8, 2026 16:08
@shunping
shunping requested a review from derrickaw October 8, 2026 16:09
@github-actions

github-actions Bot commented Oct 8, 2026

Copy link
Copy Markdown
Contributor

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

@derrickaw derrickaw mentioned this pull request Oct 8, 2026
3 tasks
@codecov

codecov Bot commented Oct 8, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.55%. Comparing base (c275c97) to head (3b9acca).
⚠️ Report is 94 commits behind head on master.

Additional details and impacted files
@@              Coverage Diff              @@
##             master   #40467       +/-   ##
=============================================
+ Coverage     56.20%   79.55%   +23.34%     
=============================================
  Files          1123      487      -636     
  Lines        178213    82238    -95975     
  Branches       1489        0     -1489     
=============================================
- Hits         100168    65422    -34746     
+ Misses        75524    16816    -58708     
+ Partials       2521        0     -2521     
Flag Coverage Δ
go ?
java ?
python 79.55% <100.00%> (-0.13%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@shunping
shunping merged commit dabcf50 into apache:master Oct 8, 2026
126 of 129 checks passed
@shunping
shunping deleted the fix-flaky-test-2 branch October 8, 2026 17:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants