feat: `bigframes.options` and `bigframes.option_context` now uses thread-local variables to prevent context managers in separate threads from affecting each other by tswast · Pull Request #652 · googleapis/python-bigquery-dataframes

tswast · 2024-05-01T21:22:10Z

In our tests, this allows us to actually test things like bf.option_context("display.repr_mode", "deferred"): without always having some other test change the display mode and break the test.

Fixes internal issue 308657813

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
Ensure the tests and linter pass
Code coverage does not decrease (if any source code was changed)
Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

…ead-local variables to prevent context managers in separate threads from affecting each other In our tests, this allows us to actually test things like `bf.option_context("display.repr_mode", "deferred"):` without always having some other test change the display mode and break the test. Fixes internal issue 308657813

GarrettWu · 2024-05-01T21:37:00Z

bigframes/_config/__init__.py

-        self._compute_options = compute_options.ComputeOptions()
+        self._local.bigquery_options = None
+
+    def _init_biquery_thread_local(self):


typo: bigquery

Thanks. Fixed in 4fa1a92

…al-option_context

tswast · 2024-05-02T15:56:12Z

e2e failure of test_model_register appears to be unrelated.

tswast · 2024-05-02T18:55:13Z

Cover failed.

tests/system/small/test_progress_bar.py 72 1 28 1 98% 144

Looks like I now have some dead code in that test. Will investigate.

milkshakeiii · 2024-05-02T20:54:50Z

Messaged offline. @tswast will fix the failing test, and then I will review.

…al-option_context

tswast · 2024-05-02T22:12:04Z

test_query_job_dry_run is still failing. Will investigate more tomorrow. Maybe I'm misunderstanding how thread locals work...

…al-option_context

tswast · 2024-05-03T16:24:56Z

I'm able to reproduce locally. The thread local variable is working how I expect, but the test is making an assumption that the DataFrame doesn't have a query_job populated. This assumption is violated when the DataFrame is used across multiple tests.

I'll update the test.

________________________________________________________________________________ test_query_job_dry_run ________________________________________________________________________________
[gw2] linux -- Python 3.11.1 /usr/local/google/home/swast/envs/bigframes/bin/python3.11

penguins_df_default_index =                                       species     island  culmen_length_mm  \
0           Gentoo penguin (Pygoscelis p...  <NA>         <NA>    <NA>  
24             20.7              191.0       3900.0  FEMALE  
...

[344 rows x 7 columns]

    def test_query_job_dry_run(penguins_df_default_index: bf.dataframe.DataFrame):
        expected = "Computation deferred. Computation will process"
    
        with bf.option_context("display.repr_mode", "deferred"):
            df_result = repr(penguins_df_default_index)
>           assert expected in df_result
E           AssertionError: assert 'Computation deferred. Computation will process' in 'Query Job Info\nJob url: https://console.cloud.google.com/bigquery?project=swast-scratch&j=bq:US:9a761424-b164-4c31-a...240503_session444e7b_fbabcb4962f841a8bb92ec2f27d17bb2\nSlot Time: a moment\nBytes Processed: 28.9 kB\nCache hit: False'

tests/system/small/test_progress_bar.py:142: AssertionError
--------------------------------------------------------------------------------- Captured stdout call ---------------------------------------------------------------------------------
DisplayOptions(max_columns=20, max_rows=25, progress_bar='auto', repr_mode='deferred', max_info_columns=100, max_info_rows=200000, memory_usage=True)
deferred

milkshakeiii

Just had one simplification suggestion, feel free to just close and re-request if you prefer the way it is now. Thanks!

milkshakeiii · 2024-05-03T17:10:36Z

bigframes/_config/__init__.py

+        # BigQuery options are special because they can only be set once per
+        # session, so we need an indicator as to whether we are using the
+        # thread-local session or the global session.
+        self._local.is_bigquery_thread_local = False


If I'm reading this correctly, would checking for self._local.bigquery_options is None be able to replace this boolean state? It seems a little more DRY to do it that way, if so.

I like it! Done in a7f3c96

milkshakeiii · 2024-05-03T17:20:06Z

bigframes/_config/__init__.py

+        # Already thread-local, so don't reset any options that have been set
+        # already. No locks needed since this only modifies thread-local
+        # variables.
+        if self._local.is_bigquery_thread_local:


i.e., this could be "if self._local.bigquery_options is not None"

milkshakeiii · 2024-05-03T17:20:29Z

bigframes/_config/__init__.py

    @property
    def bigquery(self) -> bigquery_options.BigQueryOptions:
        """Options to use with the BigQuery engine."""
+        if self._local.is_bigquery_thread_local:


milkshakeiii · 2024-05-03T17:24:26Z

bigframes/_config/__init__.py

+
+        This is set to True by using `option_context` with a "bigquery" option.
+        """
+        return self._local.is_bigquery_thread_local


This could also check none (and the comment would need to be tweaked.)

milkshakeiii · 2024-05-03T17:30:23Z

third_party/bigframes_vendored/pandas/_config/config.py

+        # sessions if we allow that.
+        if bigframes.options.is_bigquery_thread_local:
+            bigframes.close_session()
+            bigframes.options._local.is_bigquery_thread_local = False


close_session could set _local.bigquery_options to None

Someone could call close_session() inside of a context manager. We actually have some tests that start and stop multiple sessions in this thread-local context. By using the options as the thread locality indicator, we need to make sure that object stays around.

Oh, right, makes sense.

…al-option_context

milkshakeiii

LGTM, thanks!

tswast requested a review from a team as a code owner May 1, 2024 21:22

tswast requested review from a team and junyazhang May 1, 2024 21:22

product-auto-label bot added the size: l Pull request size is large. label May 1, 2024

blunderbuss-gcf bot assigned Genesis929 May 1, 2024

product-auto-label bot added the api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. label May 1, 2024

GarrettWu reviewed May 1, 2024

View reviewed changes

tswast added 2 commits May 1, 2024 21:45

biquery typo

4fa1a92

Merge remote-tracking branch 'origin/main' into b308657813-thread-loc…

5f3720c

…al-option_context

tswast requested a review from GarrettWu May 1, 2024 21:47

assert on thread-local state

d12c64e

use fixed version of test to avoid dead code

a7dbb8d

tswast mentioned this pull request May 2, 2024

feat: Add a unique session_id to Session and allow cleaning up sessions #553

Merged

milkshakeiii self-requested a review May 2, 2024 20:54

tswast added 2 commits May 2, 2024 21:38

make most options truly thread-local this time, I hope

3e3baf7

Merge remote-tracking branch 'origin/main' into b308657813-thread-loc…

6a400aa

…al-option_context

tswast added 2 commits May 3, 2024 15:15

Merge remote-tracking branch 'origin/main' into b308657813-thread-loc…

d65efba

…al-option_context

catch close errors on thread-local session too

5aa68a4

feat: always do a query dry run when option.repr_mode == "deferred"

64e3700

milkshakeiii suggested changes May 3, 2024

View reviewed changes

tswast added 2 commits May 3, 2024 18:02

Merge remote-tracking branch 'origin/main' into b308657813-thread-loc…

8cd7931

…al-option_context

use presence of _local.bigquery_options to indicate thread locality

a7f3c96

tswast requested a review from milkshakeiii May 3, 2024 18:15

milkshakeiii approved these changes May 3, 2024

View reviewed changes

tswast merged commit 651fd7d into main May 6, 2024

tswast deleted the b308657813-thread-local-option_context branch May 6, 2024 01:02

release-please bot mentioned this pull request May 6, 2024

chore(main): release 1.5.0 #645

Merged

jeffbryner mentioned this pull request May 31, 2024

'_thread._local' object has no attribute 'bigquery_options' #739

Closed

Conversation

tswast commented May 1, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tswast commented May 2, 2024

Uh oh!

tswast commented May 2, 2024

Uh oh!

milkshakeiii commented May 2, 2024

Uh oh!

tswast commented May 2, 2024

Uh oh!

tswast commented May 3, 2024

Uh oh!

milkshakeiii left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

milkshakeiii left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants