[#19711] Fix missing step_id in Python SDK worker logs during DoFn setup #37413

Tanishq1030 · 2026-01-25T04:45:23Z

Fixes #19711

This PR addresses the issue where step_id (instruction ID) was consistently missing or empty in worker logs generated during the DoFn.setup() lifecycle method.

Rationale

The FnApiLogRecordHandler relies on statesampler thread-local storage to populate the instruction_id in log entries. Previously, the BundleProcessor executed the setup() method for operations before the thread-local context was fully initialized for that instruction, causing logs emitted during setup to become orphaned (missing metadata).

Changes

sdks/python/apache_beam/runners/worker/sdk_worker.py: Updated create_bundle_processor to pass the active instruction_id into the BundleProcessor constructor.
sdks/python/apache_beam/runners/worker/bundle_processor.py:
- Updated __init__ to accept instruction_id.
- Added logic to manually inject the instruction_id into the statesampler context specifically while iterating through operations to call op.setup().
sdks/python/apache_beam/runners/worker/log_handler.py: Updated emit() to check record.instruction_id before falling back to thread-local storage, ensuring explicitly injected IDs are respected.

Verification

I verified this fix locally using a reproduction script which forces a log during setup().

Before fix: Logs during setup() had instruction_id: None.
After fix: Logs during setup() correctly display the instruction_id (e.g., bundle_...).

Mention the appropriate issue in your description (e.g. fixes #19711).
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

…n setup

gemini-code-assist · 2026-01-25T04:45:40Z

Summary of Changes

Hello @Tanishq1030, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical logging issue in the Python SDK worker where step_id metadata was missing from logs emitted during the DoFn.setup() phase. By ensuring the instruction_id is correctly propagated and made available to the logging context during this early lifecycle method, the change significantly improves the debuggability and traceability of Beam pipelines, allowing developers to better understand the context of setup-related issues.

Highlights

Problem Resolution: Fixed an issue where the step_id (instruction ID) was consistently missing or empty in worker logs generated during the DoFn.setup() lifecycle method.
Root Cause Addressed: The FnApiLogRecordHandler previously failed to populate instruction_id because DoFn.setup() was executed before the statesampler thread-local context was fully initialized for the instruction.
BundleProcessor Enhancement: The BundleProcessor now accepts the active instruction_id during its creation and explicitly injects it into the statesampler context when calling op.setup().
Log Handler Priority: The FnApiLogRecordHandler's emit() method has been updated to prioritize an explicitly provided instruction_id on the log record before attempting to retrieve it from thread-local storage.
Verification: The fix was confirmed locally using a reproduction script, showing correct instruction_id in setup() logs post-fix.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

github-actions · 2026-01-25T05:42:57Z

Assigning reviewers:

R: @claudevdm for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

github-actions · 2026-02-01T12:19:01Z

Reminder, please take a look at this pr: @claudevdm

github-actions · 2026-02-04T12:24:34Z

Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment assign to next reviewer:

R: @tvalentyn for label python.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

claudevdm · 2026-02-04T17:52:32Z

sdks/python/apache_beam/runners/worker/bundle_processor.py

@@ -1091,6 +1091,7 @@ def __init__(
      state_handler: sdk_worker.CachingStateHandler,
      data_channel_factory: data_plane.DataChannelFactory,
      data_sampler: Optional[data_sampler.DataSampler] = None,
+      instruction_id: Optional[str] = None,


Should this be optional, or should we rather make it required?

claudevdm · 2026-02-04T17:53:20Z

sdks/python/apache_beam/runners/worker/log_handler.py

-    instruction_id = statesampler.get_current_instruction_id()
-    if instruction_id:
-      log_entry.instruction_id = instruction_id
+    if hasattr(record, 'instruction_id'):


Can simplify/optimize by avoid hasattr then acces?

instruction_id = getattr(record, 'instruction_id', None) or statesampler.get_current_instruction_id() if instruction_id: log_entry.instruction_id = instruction_id

claudevdm · 2026-02-04T17:53:45Z

/gemini review

gemini-code-assist

Code Review

This pull request correctly addresses the missing step_id in worker logs during DoFn.setup() by ensuring the instruction_id is available in the statesampler context. The approach of passing the instruction_id to the BundleProcessor and then using a context manager is sound. I've added one suggestion to simplify the logic in the log handler for better readability.

gemini-code-assist · 2026-02-04T17:57:04Z

sdks/python/apache_beam/runners/worker/bundle_processor.py

+      if instruction_id:
+        with statesampler.instruction_id(instruction_id):
+          for op in reversed(self.ops.values()):
+            op.setup(self.data_sampler)
+      else:
+        for op in reversed(self.ops.values()):
+          op.setup(self.data_sampler)


The loop for setting up operations is duplicated in both the if and else blocks. This can be refactored to avoid code repetition and improve maintainability. A common pattern for this is to use a conditional context manager.

gemini-code-assist · 2026-02-04T17:57:05Z

sdks/python/apache_beam/runners/worker/log_handler.py

+    if hasattr(record, 'instruction_id'):
+      log_entry.instruction_id = record.instruction_id
+    if not log_entry.instruction_id:
+      instruction_id = statesampler.get_current_instruction_id()
+      if instruction_id:
+        log_entry.instruction_id = instruction_id


This logic for setting instruction_id can be simplified to be more concise and readable by using getattr and consolidating the checks.

instruction_id = getattr(record, 'instruction_id', None) if not instruction_id: instruction_id = statesampler.get_current_instruction_id() if instruction_id: log_entry.instruction_id = instruction_id

apache#19711-Fix missing step_id in Python SDK worker logs during DoF…

c7c2fc0

…n setup

github-actions bot added python runners labels Jan 25, 2026

github-actions bot added the Next Action: Reviewers label Jan 25, 2026

github-actions bot added the slow-review label Feb 1, 2026

github-actions bot added reassigned-reviewers and removed slow-review labels Feb 4, 2026

claudevdm requested changes Feb 4, 2026

View reviewed changes

gemini-code-assist bot reviewed Feb 4, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[#19711] Fix missing step_id in Python SDK worker logs during DoFn setup #37413

[#19711] Fix missing step_id in Python SDK worker logs during DoFn setup #37413

Tanishq1030 commented Jan 25, 2026

Uh oh!

gemini-code-assist bot commented Jan 25, 2026

Uh oh!

github-actions bot commented Jan 25, 2026

Uh oh!

github-actions bot commented Feb 1, 2026

Uh oh!

github-actions bot commented Feb 4, 2026

Uh oh!

claudevdm Feb 4, 2026

Uh oh!

claudevdm Feb 4, 2026

Uh oh!

claudevdm commented Feb 4, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Feb 4, 2026

Uh oh!

gemini-code-assist bot Feb 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

[#19711] Fix missing step_id in Python SDK worker logs during DoFn setup #37413

Are you sure you want to change the base?

[#19711] Fix missing step_id in Python SDK worker logs during DoFn setup #37413

Conversation

Tanishq1030 commented Jan 25, 2026

Rationale

Changes

Verification

Uh oh!

gemini-code-assist bot commented Jan 25, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions bot commented Jan 25, 2026

Uh oh!

github-actions bot commented Feb 1, 2026

Uh oh!

github-actions bot commented Feb 4, 2026

Uh oh!

claudevdm Feb 4, 2026

Choose a reason for hiding this comment

Uh oh!

claudevdm Feb 4, 2026

Choose a reason for hiding this comment

Uh oh!

claudevdm commented Feb 4, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Feb 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Feb 4, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants