SNOW-1674875: merge gate for multithreaded wrapper tests

.github/workflows/precommit.yml

@@ -472,6 +472,132 @@ jobs:
             .tox/.coverage
             .tox/coverage.xml
 
+  test-snowpark-multithreading-mode:
+    name: Test Snowpark Multithreading py-${{ matrix.os }}-${{ matrix.python-version }}-${{ matrix.cloud-provider }}
+    needs: build
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest-64-cores]
+        python-version: ["3.9"]
+        cloud-provider: [aws]
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+      - name: Decrypt parameters.py
+        shell: bash
+        run: .github/scripts/decrypt_parameters.sh
+        env:
+          PARAMETER_PASSWORD: ${{ secrets.PARAMETER_PASSWORD }}
+          CLOUD_PROVIDER: ${{ matrix.cloud-provider }}
+      - name: Download wheel(s)
+        uses: actions/download-artifact@v4
+        with:
+          name: wheel
+          path: dist
+      - name: Show wheels downloaded
+        run: ls -lh dist
+        shell: bash
+      - name: Upgrade setuptools, pip and wheel
+        run: python -m pip install -U setuptools pip wheel
+      - name: Install tox
+        run: python -m pip install tox
+      - name: Run tests (excluding doctests)
+        run: python -m tox -e "py${PYTHON_VERSION/\./}-multithreaded-ci"
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short
+          TOX_PARALLEL_NO_SPINNER: 1
+        shell: bash
+      - name: Run local tests
+        run: python -m tox -e "py${PYTHON_VERSION/\./}-localmultithreaded-ci"
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short
+          TOX_PARALLEL_NO_SPINNER: 1
+        shell: bash
+      - name: Combine coverages
+        run: python -m tox -e coverage --skip-missing-interpreters false
+        shell: bash
+        env:
+          SNOWFLAKE_IS_PYTHON_RUNTIME_TEST: 1
+      - uses: actions/upload-artifact@v4
+        with:
+          include-hidden-files: true
+          name: coverage_${{ matrix.os }}-${{ matrix.python-version }}-${{ matrix.cloud-provider }}-snowpark-multithreading
+          path: |
+            .tox/.coverage
+            .tox/coverage.xml
+
+
+  test-snowpandas-multithreading-mode:
+    name: Test Snowpandas Multithreading py-${{ matrix.os }}-${{ matrix.python-version }}-${{ matrix.cloud-provider }}
+    needs: build
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest-64-cores]
+        python-version: ["3.9"]
+        cloud-provider: [aws]
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+      - name: Decrypt parameters.py
+        shell: bash
+        run: .github/scripts/decrypt_parameters.sh
+        env:
+          PARAMETER_PASSWORD: ${{ secrets.PARAMETER_PASSWORD }}
+          CLOUD_PROVIDER: ${{ matrix.cloud-provider }}
+      - name: Download wheel(s)
+        uses: actions/download-artifact@v4
+        with:
+          name: wheel
+          path: dist
+      - name: Show wheels downloaded
+        run: ls -lh dist
+        shell: bash
+      - name: Upgrade setuptools, pip and wheel
+        run: python -m pip install -U setuptools pip wheel
+      - name: Install tox
+        run: python -m pip install tox
+      - name: Run Snowpark pandas API tests (excluding doctests)
+        run: python -m tox -e "py${PYTHON_VERSION/\./}-snowparkpandasmultithreaded-modin-ci"
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short
+          TOX_PARALLEL_NO_SPINNER: 1
+        shell: bash
+      - name: Combine coverages
+        run: python -m tox -e coverage --skip-missing-interpreters false
+        shell: bash
+        env:
+          SNOWFLAKE_IS_PYTHON_RUNTIME_TEST: 1
+      - uses: actions/upload-artifact@v4
+        with:
+          include-hidden-files: true
+          name: coverage_${{ matrix.os }}-${{ matrix.python-version }}-${{ matrix.cloud-provider }}-snowpandas-multithreading
+          path: |
+            .tox/.coverage
+            .tox/coverage.xml
+
+
   combine-coverage:
     if: ${{ success() || failure() }}
     name: Combine coverage
@@ -481,6 +607,8 @@ jobs:
       - test-enable-cte-optimization
       - test-snowpark-pandas
       - test-modin-extra-without-pandas-extra
+      - test-snowpark-multithreading-mode
+      - test-snowpandas-multithreading-mode
     runs-on: ubuntu-latest
     steps:
       - name: Checkout Code

CHANGELOG.md

@@ -5,7 +5,6 @@
 ### Snowpark Python API Updates
 
 - Added support for 'Service' domain to `session.lineage.trace` API.
-- Updated `Session` class to be thread-safe. This allows concurrent dataframe transformations, dataframe actions, UDF and store procedure registration, and concurrent file uploads.
 - Added support for `copy_grants` parameter when registering UDxF and stored procedures.
 
 #### New Features

src/snowflake/snowpark/_internal/server_connection.py

@@ -49,6 +49,8 @@
 from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
 from snowflake.snowpark._internal.telemetry import TelemetryClient
 from snowflake.snowpark._internal.utils import (
+    create_rlock,
+    create_thread_local,
     escape_quotes,
     get_application_name,
     get_version,
@@ -155,8 +157,6 @@ def __init__(
         options: Dict[str, Union[int, str]],
         conn: Optional[SnowflakeConnection] = None,
     ) -> None:
-        self._lock = threading.RLock()
-        self._thread_store = threading.local()
         self._lower_case_parameters = {k.lower(): v for k, v in options.items()}
         self._add_application_parameters()
         self._conn = conn if conn else connect(**self._lower_case_parameters)
@@ -171,6 +171,13 @@ def __init__(
             except TypeError:
                 pass
 
+        # thread safe param protection
+        self._thread_safe_session_enabled = self._get_client_side_session_parameter(
+            "PYTHON_SNOWPARK_ENABLE_THREAD_SAFE_SESSION", False
+        )
+        self._lock = create_rlock(self._thread_safe_session_enabled)
+        self._thread_store = create_thread_local(self._thread_safe_session_enabled)
+
         if "password" in self._lower_case_parameters:
             self._lower_case_parameters["password"] = None
         self._telemetry_client = TelemetryClient(self._conn)

src/snowflake/snowpark/_internal/temp_table_auto_cleaner.py

@@ -2,12 +2,12 @@
 # Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
 #
 import logging
-import threading
 import weakref
 from collections import defaultdict
 from typing import TYPE_CHECKING, Dict
 
 from snowflake.snowpark._internal.analyzer.snowflake_plan_node import SnowflakeTable
+from snowflake.snowpark._internal.utils import create_rlock
 
 if TYPE_CHECKING:
     from snowflake.snowpark.session import Session  # pragma: no cover
@@ -33,7 +33,7 @@ def __init__(self, session: "Session") -> None:
         # this dict will still be maintained even if the cleaner is stopped (`stop()` is called)
         self.ref_count_map: Dict[str, int] = defaultdict(int)
         # Lock to protect the ref_count_map
-        self.lock = threading.RLock()
+        self.lock = create_rlock(session._conn._thread_safe_session_enabled)
 
     def add(self, table: SnowflakeTable) -> None:
         with self.lock:

src/snowflake/snowpark/_internal/utils.py

@@ -298,7 +298,12 @@ def normalize_path(path: str, is_local: bool) -> str:
     return f"'{path}'"
 
 
-def warn_session_config_update_in_multithreaded_mode(config) -> None:
+def warn_session_config_update_in_multithreaded_mode(
+    config: str, thread_safe_mode_enabled: bool
+) -> None:
+    if not thread_safe_mode_enabled:
+        return
+
     if threading.active_count() > 1:
         logger.warning(
             "You might have more than one threads sharing the Session object trying to update "
@@ -675,6 +680,47 @@ def warning(self, text: str) -> None:
         self.count += 1
 
 
+# TODO: SNOW-1720855: Remove DummyRLock and DummyThreadLocal after the rollout
+class DummyRLock:
+    """This is a dummy lock that is used in place of threading.Rlock when multithreading is
+    disabled."""
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def acquire(self, *args, **kwargs):
+        pass  # pragma: no cover
+
+    def release(self, *args, **kwargs):
+        pass  # pragma: no cover
+
+
+class DummyThreadLocal:
+    """This is a dummy thread local class that is used in place of threading.local when
+    multithreading is disabled."""
+
+    pass
+
+
+def create_thread_local(
+    thread_safe_session_enabled: bool,
+) -> Union[threading.local, DummyThreadLocal]:
+    if thread_safe_session_enabled:
+        return threading.local()
+    return DummyThreadLocal()
+
+
+def create_rlock(
+    thread_safe_session_enabled: bool,
+) -> Union[threading.RLock, DummyRLock]:
+    if thread_safe_session_enabled:
+        return threading.RLock()
+    return DummyRLock()
+
+
 warning_dict: Dict[str, WarningHelper] = {}
 
 

src/snowflake/snowpark/mock/_connection.py

@@ -6,7 +6,6 @@
 import functools
 import json
 import logging
-import threading
 import uuid
 from copy import copy
 from decimal import Decimal
@@ -30,6 +29,7 @@
 from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
 from snowflake.snowpark._internal.server_connection import DEFAULT_STRING_SIZE
 from snowflake.snowpark._internal.utils import (
+    create_rlock,
     is_in_stored_procedure,
     result_set_to_rows,
 )
@@ -281,19 +281,26 @@ def read_table_if_exists(
     def __init__(self, options: Optional[Dict[str, Any]] = None) -> None:
         self._conn = MockedSnowflakeConnection()
         self._cursor = Mock()
-        self._lock = threading.RLock()
+        self._options = options or {}
+        session_params = self._options.get("session_parameters", {})
+        # thread safe param protection
+        self._thread_safe_session_enabled = session_params.get(
+            "PYTHON_SNOWPARK_ENABLE_THREAD_SAFE_SESSION", False
+        )
+        self._lock = create_rlock(self._thread_safe_session_enabled)
         self._lower_case_parameters = {}
         self.remove_query_listener = Mock()
         self.add_query_listener = Mock()
         self._telemetry_client = Mock()
         self.entity_registry = MockServerConnection.TabularEntityRegistry(self)
         self.stage_registry = StageEntityRegistry(self)
-        self._conn._session_parameters = {
-            "ENABLE_ASYNC_QUERY_IN_PYTHON_STORED_PROCS": False,
-            "_PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS_STRING": True,
-            "_PYTHON_SNOWPARK_USE_SQL_SIMPLIFIER_STRING": True,
-        }
-        self._options = options or {}
+        self._conn._session_parameters = session_params.update(
+            {
+                "ENABLE_ASYNC_QUERY_IN_PYTHON_STORED_PROCS": False,
+                "_PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS_STRING": True,
+                "_PYTHON_SNOWPARK_USE_SQL_SIMPLIFIER_STRING": True,
+            }
+        )
         self._active_account = self._options.get(
             "account", snowflake.snowpark.mock._constants.CURRENT_ACCOUNT
         )