Skip to content
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

fix: fix beam metrics after migrating to async batcher #27085

Merged
merged 5 commits into from
Jun 16, 2023
+6 −43
Merged

fix: fix beam metrics after migrating to async batcher #27085

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c6fe011
fix: fix beam metrics after migrating to async batcher
mutianf Jun 9, 2023
98e67a7
update version
mutianf Jun 9, 2023
6c9ae46
update bigtable version
mutianf Jun 9, 2023
07e00ba
update format
mutianf Jun 9, 2023
77a7f95
fix format
mutianf Jun 9, 2023
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
45 changes: 5 additions & 40 deletions sdks/python/apache_beam/io/gcp/bigtableio.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,42 +55,6 @@
FLUSH_COUNT = 1000
MAX_ROW_BYTES = 5242880 # 5MB

class _MutationsBatcher(MutationsBatcher):
def __init__(
self, table, flush_count=FLUSH_COUNT, max_row_bytes=MAX_ROW_BYTES):
super().__init__(table, flush_count, max_row_bytes)
self.rows = []

def set_flush_callback(self, callback_fn):
self.callback_fn = callback_fn

def flush(self):
if len(self.rows) != 0:
status_list = self.table.mutate_rows(self.rows)
self.callback_fn(status_list)

# If even one request fails we retry everything. BigTable mutations are
# idempotent so this should be correct.
# TODO(https://github.com/apache/beam/issues/21396): make this more
# efficient by retrying only re-triable failed requests.
for status in status_list:
if not status:
# BigTable client may return 'None' instead of a valid status in
# some cases due to
# https://github.com/googleapis/python-bigtable/issues/485
raise Exception(
'Failed to write a batch of %r records' % len(self.rows))
elif status.code != 0:
raise Exception(
'Failed to write a batch of %r records due to %r' % (
len(self.rows),
ServiceCallMetric.bigtable_error_code_to_grpc_status_string(
status.code)))

self.total_mutation_count = 0
self.total_size = 0
self.rows = []

except ImportError:
_LOGGER.warning(
'ImportError: from google.cloud.bigtable import Client', exc_info=True)
Expand Down Expand Up @@ -168,8 +132,8 @@ def start_bundle(self):
self.beam_options['project_id'],
self.beam_options['instance_id'],
self.beam_options['table_id'])
self.batcher = _MutationsBatcher(self.table)
self.batcher.set_flush_callback(self.write_mutate_metrics)
self.batcher = MutationsBatcher(
self.table, batch_completed_callback=self.write_mutate_metrics)

def process(self, row):
self.written.inc()
Expand All @@ -184,8 +148,9 @@ def process(self, row):
self.batcher.mutate(row)

def finish_bundle(self):
self.batcher.flush()
self.batcher = None
if self.batcher:
self.batcher.close()
self.batcher = None

def display_data(self):
return {
Expand Down
4 changes: 1 addition & 3 deletions sdks/python/setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,9 +314,7 @@ def get_portability_package_data():
'google-cloud-bigquery>=2.0.0,<4',
'google-cloud-bigquery-storage>=2.6.3,<3',
'google-cloud-core>=2.0.0,<3',
# TODO(https://github.com/apache/beam/issues/26673)
# 2.18.x breaks unit test
'google-cloud-bigtable>=2.0.0,<2.18.0',
'google-cloud-bigtable>=2.19.0,<3',
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tvalentyn is it safe to change the lower bound for this package?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI, the updates in bigtableio.py will break with older version of google-cloud-bigtable :(. And seems like the lower boundary was updated for some other modules before (e.g. bigquery-storage)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, increasing the lower bound is appropriate and necessary in this case.

'google-cloud-spanner>=3.0.0,<4',
# GCP Packages required by ML functionality
'google-cloud-dlp>=3.0.0,<4',
Expand Down