Skip to content

Commit d7f57fc

Browse files
authored
chore(bigframes): optimize system test teardown (#17443)
brings down bigframes system tests from 23 mins to 12 mins i.e. 2x speed up.
1 parent 08cc1f6 commit d7f57fc

2 files changed

Lines changed: 15 additions & 3 deletions

File tree

packages/bigframes/bigframes/session/anonymous_dataset.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import threading
1717
import uuid
1818
import warnings
19+
from concurrent.futures import ThreadPoolExecutor
1920
from typing import List, Optional, Sequence
2021

2122
import google.cloud.bigquery as bigquery
@@ -170,9 +171,19 @@ def _cleanup_old_udfs(self):
170171

171172
def close(self):
172173
"""Delete tables that were created with this session's session_id."""
173-
for table_ref in self._table_ids:
174-
self.bqclient.delete_table(table_ref, not_found_ok=True)
175-
self._table_ids.clear()
174+
if self._table_ids:
175+
try:
176+
with ThreadPoolExecutor() as executor:
177+
futures = [
178+
executor.submit(
179+
self.bqclient.delete_table, table_ref, not_found_ok=True
180+
)
181+
for table_ref in self._table_ids
182+
]
183+
for future in futures:
184+
future.result()
185+
finally:
186+
self._table_ids.clear()
176187

177188
try:
178189
# Before closing the session, attempt to clean up any uncollected,

packages/bigframes/noxfile.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,7 @@ def run_system(
354354
"py.test",
355355
"-v",
356356
f"-n={num_workers}",
357+
"--dist=worksteal",
357358
# Any individual test taking longer than 15 mins will be terminated.
358359
f"--timeout={timeout_seconds}",
359360
# Log 20 slowest tests

0 commit comments

Comments
 (0)