Skip to content

Change DirectRunner for FnApiRunner in tfgnn_sampler #918

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
May 20, 2026
Merged

Change DirectRunner for FnApiRunner in tfgnn_sampler #918

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion tensorflow_gnn/experimental/sampler/beam/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.runners.portability import fn_api_runner
import tensorflow as tf
import tensorflow_gnn as tfgnn
from tensorflow_gnn.data import unigraph
Expand Down Expand Up @@ -191,7 +192,10 @@ def _create_beam_runner(
) -> beam.runners.PipelineRunner:
"""Creates appropriate runner."""
if runner_name == _DIRECT_RUNNER:
runner = beam.runners.DirectRunner()
# b/490166623: Changed to FnApiRunner due to the new prism implementation
# in DirectRunner since beam 2.68.0 which breaks the sampler. FnApiRunner
# enables the old implementation of DirectRunner.
runner = fn_api_runner.FnApiRunner()
elif runner_name == _DATAFLOW_RUNNER:
runner = beam.runners.DataflowRunner()
# Placeholder for Google-internal runner option creation
Expand Down
Loading