SNOW-1663726 make session config updates thread safe #2302

sfc-gh-aalam · 2024-09-17T02:44:31Z

Which Jira issue is this PR addressing? Make sure that there is an accompanying issue to your PR.

Fixes SNOW-1663726
Fill out the following pre-review checklist:
- I am adding a new automated test(s) to verify correctness of my new code
  - If this test skips Local Testing mode, I'm requesting review from @snowflakedb/local-testing
- I am adding new logging messages
- I am adding a new telemetry message
- I am adding new credentials
- I am adding a new dependency
- If this is a new feature/behavior, I'm adding the Local Testing parity changes.
Please describe how your code solves the related issue.

In this PR we make the following changes:
- We protect all session config variable updates using a lock
- We create a per thread _plan_builder to isolate usage.
- Created ConfigContext to set a context and read a snapshot of config values from session. Use the ConfigContext object for Analyzer, PlanCompiler and SnowflakePlanBuilder.
  - created methods do_analyzer and _compile that runs using a single ConfigContext

…ection-thread-safe

…ad-safe' into aalam-SNOW-1418523-make-analyzer-server_connection-thread-safe

SNOW-1418523-make-udf-sproc-thread-safe

…n_stage, plan_builder

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py

sfc-gh-yzou · 2024-09-27T22:50:51Z

src/snowflake/snowpark/_internal/analyzer/analyzer.py

@@ -172,13 +173,21 @@ def analyze(
        expr: Union[Expression, NamedExpression],
        df_aliased_col_name_to_real_col_name: DefaultDict[str, Dict[str, str]],
        parse_local_name=False,
+        config_context: Optional[ConfigContext] = None,


i would prefer to make this non-optional, so people can think explicitly when they are calling the interface

sfc-gh-yzou · 2024-09-27T22:59:39Z

src/snowflake/snowpark/_internal/analyzer/config_context.py

+from typing import Any
+
+
+class ConfigContext:


with this, i think for the one you mentioned before for the access sql_simplifier_enabled, we should also use this configContext if it is accessing the sql_simplifier_enabled multiple times in the function call

sfc-gh-yzou · 2024-09-27T23:08:58Z

src/snowflake/snowpark/_internal/analyzer/config_context.py

+        self.session = session
+        self.configs = {
+            "_query_compilation_stage_enabled",
+            "cte_optimization_enabled",


instead of doing this, it might be more clear to provide a property interface in the session that teaks a snapshot. so it is the interface is at a closer place with the properties.

Given the fact that we are adding more and more configs, i am thinking it might worth to do a refactoring to centralize all configs, and wrap all of them into a class to reduce the possiblility for developer make mistakes.

sfc-gh-yzou · 2024-09-27T23:10:58Z

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py

@@ -531,6 +542,10 @@ def __init__(
        # on the optimized plan. During the final query generation, no schema query is needed,
        # this helps reduces un-necessary overhead for the describing call.
        self._skip_schema_query = skip_schema_query
+        self._config_context: Optional[ConfigContext] = None


The config context here should probably been done in the same way as what we have done for analyze and resolve as parameter for each builder.

sfc-gh-yzou · 2024-09-27T23:13:00Z

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py

@@ -564,9 +579,10 @@ def build(
            ), "No schema query is available in child SnowflakePlan"
            new_schema_query = schema_query or sql_generator(child.schema_query)

+        config_context = self._config_context or ConfigContext(self.session)


each build function seems can take a new snapshot, but multiple build can be called during the whole resolve, and they are suppose to get the same snapshot, it is safer and cleaner to simply extend each interface to take the config context.

I intend to remove config_context from plan_builder once old CTE implementation is gone. That is the only reason why plan builder needs this info. I can own this task. Without that, the refactor will be too big and it would be prone to more errors. I want to avoid that.

sfc-gh-yzou · 2024-09-27T23:14:11Z

src/snowflake/snowpark/_internal/compiler/plan_compiler.py

@@ -46,6 +47,7 @@ class PlanCompiler:

    def __init__(self, plan: SnowflakePlan) -> None:
        self._plan = plan
+        self._config_context = ConfigContext(self._plan.session)


That is good, we are only using one snapshot during the whole compilation stage

sfc-gh-yzou · 2024-09-27T23:17:29Z

src/snowflake/snowpark/modin/plugin/_internal/generator_utils.py

-    session.sql_simplifier_enabled = sql_simplifier_enabled
+    with session._lock:
+        sql_simplifier_enabled = session.sql_simplifier_enabled
+        session.sql_simplifier_enabled = False


that code here is hacky and bad. do you know what is the ticket SNOW-1646883 referring for, i think we should fix that and remove this code, we are not suppose to change the config in our code like this.

…ad-safe' into aalam-SNOW-1663726-make-session-config-updates-thread-safe

sfc-gh-yzou · 2024-10-03T17:52:06Z

src/snowflake/snowpark/_internal/compiler/plan_compiler.py

@@ -49,6 +49,13 @@ class PlanCompiler:

    def __init__(self, plan: SnowflakePlan) -> None:
        self._plan = plan
+        session = plan.session
+        self.cte_optimization_enabled = session.cte_optimization_enabled


if we are restricting the change to config runtime, this is not necessary anymore, right?

we are only putting restriction on cte_optimization_enabled. Since for other, we still need to take a snapshot, I'm taking a snapshot for cte param for completeness.

since now we are putting protection on all config variables, i think we don't need to do that anymore

sfc-gh-yzou · 2024-10-03T17:54:07Z

src/snowflake/snowpark/session.py

-                self._session_id
+        if threading.active_count() > 1:
+            # TODO (SNOW-1541096): Remove the limitation once old cte implementation is removed.
+            _logger.warning(


i think we could simply disallow updating of any config value when there are multiple active thread to be consistent everywhere

For example, i actually uses large_query_breakdown when deciding whether the complexity can be merged or not.

sfc-gh-yzou · 2024-10-03T19:19:05Z

src/snowflake/snowpark/session.py

@@ -827,12 +868,18 @@ def large_query_breakdown_enabled(self, value: bool) -> None:
        materialize the partitions, and then combine them to execute the query to improve
        overall performance.
        """
+        if threading.active_count() > 1:


you can probably have a small utility function and call that for all configs with the config name as a parameter

sfc-gh-yzou · 2024-10-03T19:19:35Z

src/snowflake/snowpark/_internal/compiler/plan_compiler.py

@@ -49,6 +49,13 @@ class PlanCompiler:

    def __init__(self, plan: SnowflakePlan) -> None:
        self._plan = plan
+        session = plan.session
+        self.cte_optimization_enabled = session.cte_optimization_enabled


since now we are putting protection on all config variables, i think we don't need to do that anymore

sfc-gh-yzou · 2024-10-03T21:33:27Z

src/snowflake/snowpark/_internal/compiler/plan_compiler.py

@@ -151,8 +154,7 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
            return queries
        else:
            final_plan = self._plan
-            if self._plan.session.cte_optimization_enabled:
-                final_plan = final_plan.replace_repeated_subquery_with_cte()
+            final_plan = final_plan.replace_repeated_subquery_with_cte()


is that intended? i think we still need the condition here

yes. replace_repeated_subquery_with_cte() does this check automatically.

sfc-gh-yzou · 2024-10-04T01:43:54Z

src/snowflake/snowpark/_internal/utils.py

+    if threading.active_count() > 1:
+        logger.warning(
+            "You might have more than one threads sharing the Session object trying to update "
+            f"{config}. This is currently not thread-safe and may cause unexpected behavior. "


the sense here sounds little bit wired, and updating the config its-self is thread safe, it just make the other operations unsafe. maybe say updating f"{config}" while other tasks are running can potentially cause unexpected behavior.

sfc-gh-yzou · 2024-10-04T01:49:16Z

src/snowflake/snowpark/session.py

@@ -637,6 +639,16 @@ def _analyzer(self) -> Analyzer:
            )
        return self._thread_store.analyzer

+    @property
+    def _plan_builder(self):
+        if not hasattr(self._thread_store, "plan_builder"):


oh, i didn't notice we have plan builder here, i thought plan builder is just part of the analyzer, what is the plan builder used here?, and what is the reason that we moved it to a property

and the plan builder seems thread local to me

This plan builder is different from analyzer's plan builder. Only used by

DataFramReader to create read_file plan

FileOperation to create file_operation_plan

I added it as part of old changes. reverted back to original now since plan builder should be thread-safe with restricted config updates

sfc-gh-aalam added 20 commits September 11, 2024 11:44

init

56fb566

make udf/sproc related files thread-safe

66003d1

Merge branch 'main' into aalam-SNOW-1418523-make-udf-sproc-thread-safe

0e58205

init

e75dde1

make query listener thread-safe

68a8c1c

Fix query_tag and last_action_id

31a5734

core updates done

b4dadda

Add tests

b8c6496

Fix local tests

f39837e

Merge branch 'main' into aalam-SNOW-1418523-make-analyzer-server_conn…

31a196f

…ection-thread-safe

Merge branch 'aalam-SNOW-1418523-make-internal-session-variables-thre…

723bdf7

…ad-safe' into aalam-SNOW-1418523-make-analyzer-server_connection-thread-safe

add file IO tests

37c0419

Merge branch 'aalam-SNOW-1418523-concurrent-file-operations' into aalam-

8a2d433

SNOW-1418523-make-udf-sproc-thread-safe

make session._runtime_version_from_requirement safe

a083989

add sp/udf concurrent tests

947d384

fix broken test

fd51720

add udtf/udaf tests

3077853

fix broken test

65c3186

sql_simplifier, cte_optimization, eliminate_numeric, query_compilatio…

94412cf

…n_stage, plan_builder

cover more configs

638dd09

sfc-gh-aalam commented Sep 17, 2024

View reviewed changes

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py Outdated Show resolved Hide resolved

sfc-gh-aalam added 6 commits September 17, 2024 11:45

fix SnowflakePlan copy

7ae2c33

minor update

1689ebf

add description

5e8a2d2

use _package_lock to protect Session._packages

1c83ef2

undo refactor

a649761

undo refactor

f03d618

sfc-gh-aalam mentioned this pull request Sep 17, 2024

SNOW-1418523 Make Session thread safe #2312

Merged

8 tasks

sfc-gh-aalam added 2 commits September 17, 2024 14:33

fix test

5f398d5

fix test

3807087

sfc-gh-aalam added 6 commits September 27, 2024 15:02

fix unit tests

dddd15f

simplify

57ee9e8

simplify

809a86e

simplify

6021ab8

simplify

43986f6

simplify

0430e92

sfc-gh-yzou reviewed Sep 27, 2024

View reviewed changes

sfc-gh-aalam added 7 commits September 30, 2024 10:00

remove config context

095b04e

min-diff

32707f9

min-diff

3bf678d

min-diff

3eade1a

Merge branch 'aalam-SNOW-1418523-make-internal-session-variables-thre…

1850d5d

…ad-safe' into aalam-SNOW-1663726-make-session-config-updates-thread-safe

add warnings

1fa6ad2

Merge branch 'aalam-SNOW-1418523-make-internal-session-variables-thre…

7c85432

…ad-safe' into aalam-SNOW-1663726-make-session-config-updates-thread-safe

sfc-gh-yzou reviewed Oct 3, 2024

View reviewed changes

address feedback

f994842

sfc-gh-yzou reviewed Oct 3, 2024

View reviewed changes

sfc-gh-aalam added 5 commits October 3, 2024 13:27

address feedback

4621836

fix string

e1c68f3

ignore on multi-thread

e5b48dd

undo ignore

496e2be

update warning message

980d3b7

sfc-gh-yzou reviewed Oct 4, 2024

View reviewed changes

address comments

54a6b5d

sfc-gh-yzou approved these changes Oct 4, 2024

View reviewed changes

address comments

67609e8

sfc-gh-aalam merged commit 5672a1d into aalam-SNOW-1418523-make-internal-session-variables-thread-safe Oct 4, 2024
29 of 30 checks passed

sfc-gh-aalam deleted the aalam-SNOW-1663726-make-session-config-updates-thread-safe branch October 4, 2024 16:58

github-actions bot locked and limited conversation to collaborators Oct 4, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

SNOW-1663726 make session config updates thread safe #2302

SNOW-1663726 make session config updates thread safe #2302

sfc-gh-aalam commented Sep 17, 2024 •

edited

Loading

sfc-gh-yzou Sep 27, 2024

sfc-gh-yzou Sep 27, 2024

sfc-gh-yzou Sep 27, 2024

sfc-gh-yzou Sep 27, 2024

sfc-gh-yzou Sep 27, 2024

sfc-gh-aalam Sep 29, 2024

sfc-gh-yzou Sep 27, 2024

sfc-gh-yzou Sep 27, 2024

sfc-gh-yzou Oct 3, 2024

sfc-gh-aalam Oct 3, 2024

sfc-gh-yzou Oct 3, 2024

sfc-gh-yzou Oct 3, 2024

sfc-gh-yzou Oct 3, 2024

sfc-gh-yzou Oct 3, 2024

sfc-gh-yzou Oct 3, 2024

sfc-gh-yzou Oct 3, 2024

sfc-gh-aalam Oct 4, 2024

sfc-gh-yzou Oct 4, 2024

sfc-gh-aalam Oct 4, 2024

sfc-gh-yzou Oct 4, 2024

sfc-gh-yzou Oct 4, 2024

sfc-gh-aalam Oct 4, 2024

sfc-gh-aalam Oct 4, 2024

SNOW-1663726 make session config updates thread safe #2302

SNOW-1663726 make session config updates thread safe #2302

Conversation

sfc-gh-aalam commented Sep 17, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-aalam commented Sep 17, 2024 •

edited

Loading