snowflake-ml-python 1.6.0__tar.gz → 1.6.2__tar.gz

{snowflake_ml_python-1.6.0 → snowflake_ml_python-1.6.2}/CHANGELOG.md

@@ -1,6 +1,50 @@
 # Release History
 
-## 1.6.0
+## 1.6.2 (TBD)
+
+### Bug Fixes
+
+- Modeling: Support XGBoost version that is larger than 2.
+
+- Data: Fix multiple epoch iteration over `DataConnector.to_torch_datapipe()` DataPipes.
+- Generic: Fix a bug that when an invalid name is provided to argument where fully qualified name is expected, it will
+  be parsed wrongly. Now it raises an exception correctly.
+- Model Explainability: Handle explanations for multiclass XGBoost classification models
+- Model Explainability: Workarounds and better error handling for XGB>2.1.0 not working with SHAP==0.42.1
+
+### New Features
+
+- Data: Add top-level exports for `DataConnector` and `DataSource` to `snowflake.ml.data`.
+- Data: Add native batching support via `batch_size` and `drop_last_batch` arguments to `DataConnector.to_torch_dataset()`
+- Feature Store: update_feature_view() supports taking feature view object as argument.
+
+### Behavior Changes
+
+## 1.6.1 (2024-08-12)
+
+### Bug Fixes
+
+- Feature Store: Support large metadata blob when generating dataset
+- Feature Store: Added a hidden knob in FeatureView as kargs for setting customized
+  refresh_mode
+- Registry: Fix an error message in Model Version `run` when `function_name` is not mentioned and model has multiple
+  target methods.
+- Cortex inference: snowflake.cortex.Complete now only uses the REST API for streaming and the use_rest_api_experimental
+  is no longer needed.
+- Feature Store: Add a new API: FeatureView.list_columns() which list all column information.
+- Data: Fix `DataFrame` ingestion with `ArrowIngestor`.
+
+### New Features
+
+- Enable `set_params` to set the parameters of the underlying sklearn estimator, if the snowflake-ml model has been fit.
+- Data: Add `snowflake.ml.data.ingestor_utils` module with utility functions helpful for `DataIngestor` implementations.
+- Data: Add new `to_torch_dataset()` connector to `DataConnector` to replace deprecated DataPipe.
+- Registry: Option to `enable_explainability` set to True by default for XGBoost, LightGBM and CatBoost as PuPr feature.
+- Registry: Option to `enable_explainability` when registering SHAP supported sklearn models.
+
+### Behavior Changes
+
+## 1.6.0 (2024-07-29)
 
 ### Bug Fixes
 
@@ -29,6 +73,14 @@
   distributed_hpo_trainer.ENABLE_EFFICIENT_MEMORY_USAGE = False
   `
 - Registry: Option to `enable_explainability` when registering LightGBM models as a pre-PuPr feature.
+- Data: Add new `snowflake.ml.data` preview module which contains data reading utilities like `DataConnector`
+  - `DataConnector` provides efficient connectors from Snowpark `DataFrame`
+  and Snowpark ML `Dataset` to external frameworks like PyTorch, TensorFlow, and Pandas. Create `DataConnector`
+  instances using the classmethod constructors `DataConnector.from_dataset()` and `DataConnector.from_dataframe()`.
+- Data: Add new `DataConnector.from_sources()` classmethod constructor for constructing from `DataSource` objects.
+- Data: Add new `ingestor_class` arg to `DataConnector` classmethod constructors for easier `DataIngestor` injection.
+- Dataset: `DatasetReader` now subclasses new `DataConnector` class.
+  - Add optional `limit` arg to `DatasetReader.to_pandas()`
 
 ### Behavior Changes
 

{snowflake_ml_python-1.6.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: snowflake-ml-python
-Version: 1.6.0
+Version: 1.6.2
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -253,7 +253,7 @@ Requires-Dist: snowflake-connector-python[pandas]<4,>=3.5.0
 Requires-Dist: snowflake-snowpark-python<2,>=1.17.0
 Requires-Dist: sqlparse<1,>=0.4
 Requires-Dist: typing-extensions<5,>=4.1.0
-Requires-Dist: xgboost<2,>=1.7.3
+Requires-Dist: xgboost<2.1,>=1.7.3
 Provides-Extra: all
 Requires-Dist: catboost<2,>=1.2.0; extra == "all"
 Requires-Dist: lightgbm<5,>=3.3.5; extra == "all"
@@ -373,7 +373,51 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
 
 # Release History
 
-## 1.6.0
+## 1.6.2 (TBD)
+
+### Bug Fixes
+
+- Modeling: Support XGBoost version that is larger than 2.
+
+- Data: Fix multiple epoch iteration over `DataConnector.to_torch_datapipe()` DataPipes.
+- Generic: Fix a bug that when an invalid name is provided to argument where fully qualified name is expected, it will
+  be parsed wrongly. Now it raises an exception correctly.
+- Model Explainability: Handle explanations for multiclass XGBoost classification models
+- Model Explainability: Workarounds and better error handling for XGB>2.1.0 not working with SHAP==0.42.1
+
+### New Features
+
+- Data: Add top-level exports for `DataConnector` and `DataSource` to `snowflake.ml.data`.
+- Data: Add native batching support via `batch_size` and `drop_last_batch` arguments to `DataConnector.to_torch_dataset()`
+- Feature Store: update_feature_view() supports taking feature view object as argument.
+
+### Behavior Changes
+
+## 1.6.1 (2024-08-12)
+
+### Bug Fixes
+
+- Feature Store: Support large metadata blob when generating dataset
+- Feature Store: Added a hidden knob in FeatureView as kargs for setting customized
+  refresh_mode
+- Registry: Fix an error message in Model Version `run` when `function_name` is not mentioned and model has multiple
+  target methods.
+- Cortex inference: snowflake.cortex.Complete now only uses the REST API for streaming and the use_rest_api_experimental
+  is no longer needed.
+- Feature Store: Add a new API: FeatureView.list_columns() which list all column information.
+- Data: Fix `DataFrame` ingestion with `ArrowIngestor`.
+
+### New Features
+
+- Enable `set_params` to set the parameters of the underlying sklearn estimator, if the snowflake-ml model has been fit.
+- Data: Add `snowflake.ml.data.ingestor_utils` module with utility functions helpful for `DataIngestor` implementations.
+- Data: Add new `to_torch_dataset()` connector to `DataConnector` to replace deprecated DataPipe.
+- Registry: Option to `enable_explainability` set to True by default for XGBoost, LightGBM and CatBoost as PuPr feature.
+- Registry: Option to `enable_explainability` when registering SHAP supported sklearn models.
+
+### Behavior Changes
+
+## 1.6.0 (2024-07-29)
 
 ### Bug Fixes
 
@@ -402,6 +446,14 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
   distributed_hpo_trainer.ENABLE_EFFICIENT_MEMORY_USAGE = False
   `
 - Registry: Option to `enable_explainability` when registering LightGBM models as a pre-PuPr feature.
+- Data: Add new `snowflake.ml.data` preview module which contains data reading utilities like `DataConnector`
+  - `DataConnector` provides efficient connectors from Snowpark `DataFrame`
+  and Snowpark ML `Dataset` to external frameworks like PyTorch, TensorFlow, and Pandas. Create `DataConnector`
+  instances using the classmethod constructors `DataConnector.from_dataset()` and `DataConnector.from_dataframe()`.
+- Data: Add new `DataConnector.from_sources()` classmethod constructor for constructing from `DataSource` objects.
+- Data: Add new `ingestor_class` arg to `DataConnector` classmethod constructors for easier `DataIngestor` injection.
+- Dataset: `DatasetReader` now subclasses new `DataConnector` class.
+  - Add optional `limit` arg to `DatasetReader.to_pandas()`
 
 ### Behavior Changes
 

{snowflake_ml_python-1.6.0 → snowflake_ml_python-1.6.2}/pyproject.toml

@@ -8,7 +8,7 @@ description = "The machine learning client library that is used for interacting
 classifiers = [ "Development Status :: 5 - Production/Stable", "Environment :: Console", "Environment :: Other Environment", "Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Information Technology", "Intended Audience :: System Administrators", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Topic :: Database", "Topic :: Software Development", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Application Frameworks", "Topic :: Software Development :: Libraries :: Python Modules", "Topic :: Scientific/Engineering :: Information Analysis",]
 requires-python = ">=3.8, <3.12"
 dynamic = [ "version", "readme",]
-dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<4", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "fsspec[http]>=2022.11,<2024", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<2", "packaging>=20.9,<24", "pandas>=1.0.0,<3", "pyarrow", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2022.11,<2024", "scikit-learn>=1.2.1,<1.4", "scipy>=1.9,<2", "snowflake-connector-python[pandas]>=3.5.0,<4", "snowflake-snowpark-python>=1.17.0,<2", "sqlparse>=0.4,<1", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<2",]
+dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<4", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "fsspec[http]>=2022.11,<2024", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<2", "packaging>=20.9,<24", "pandas>=1.0.0,<3", "pyarrow", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2022.11,<2024", "scikit-learn>=1.2.1,<1.4", "scipy>=1.9,<2", "snowflake-connector-python[pandas]>=3.5.0,<4", "snowflake-snowpark-python>=1.17.0,<2", "sqlparse>=0.4,<1", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<2.1",]
 [[project.authors]]
 name = "Snowflake, Inc"
 email = "support@snowflake.com"

{snowflake_ml_python-1.6.0 → snowflake_ml_python-1.6.2}/snowflake/cortex/_complete.py

@@ -90,7 +90,6 @@ def _call_complete_rest(
     prompt: Union[str, List[ConversationMessage]],
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
-    stream: bool = False,
 ) -> requests.Response:
     session = session or context.get_active_session()
     if session is None:
@@ -121,7 +120,7 @@ def _call_complete_rest(
 
     data = {
         "model": model,
-        "stream": stream,
+        "stream": True,
     }
     if isinstance(prompt, List):
         data["messages"] = prompt
@@ -137,32 +136,15 @@ def _call_complete_rest(
         if "top_p" in options:
             data["top_p"] = options["top_p"]
 
-    logger.debug(f"making POST request to {url} (model={model}, stream={stream})")
+    logger.debug(f"making POST request to {url} (model={model})")
     return requests.post(
         url,
         json=data,
         headers=headers,
-        stream=stream,
+        stream=True,
     )
 
 
-def _process_rest_response(
-    response: requests.Response,
-    stream: bool = False,
-    deadline: Optional[float] = None,
-) -> Union[str, Iterator[str]]:
-    if stream:
-        return _return_stream_response(response, deadline)
-
-    try:
-        content = response.json()["choices"][0]["message"]["content"]
-        assert isinstance(content, str)
-        return content
-    except (KeyError, IndexError, AssertionError) as e:
-        # Unlike the streaming case, errors are not ignored because a message must be returned.
-        raise ResponseParseException("Failed to parse message from response.") from e
-
-
 def _return_stream_response(response: requests.Response, deadline: Optional[float]) -> Iterator[str]:
     client = SSEClient(response)
     for event in client.events():
@@ -243,7 +225,6 @@ def _complete_impl(
     prompt: Union[str, List[ConversationMessage], snowpark.Column],
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
-    use_rest_api_experimental: bool = False,
     stream: bool = False,
     function: str = "snowflake.cortex.complete",
     timeout: Optional[float] = None,
@@ -253,16 +234,14 @@ def _complete_impl(
         raise ValueError('only one of "timeout" and "deadline" must be set')
     if timeout is not None:
         deadline = time.time() + timeout
-    if use_rest_api_experimental:
+    if stream:
         if not isinstance(model, str):
             raise ValueError("in REST mode, 'model' must be a string")
         if not isinstance(prompt, str) and not isinstance(prompt, List):
             raise ValueError("in REST mode, 'prompt' must be a string or a list of ConversationMessage")
-        response = _call_complete_rest(model, prompt, options, session=session, stream=stream, deadline=deadline)
+        response = _call_complete_rest(model, prompt, options, session=session, deadline=deadline)
         assert response.status_code >= 200 and response.status_code < 300
-        return _process_rest_response(response, stream=stream)
-    if stream is True:
-        raise ValueError("streaming can only be enabled in REST mode, set use_rest_api_experimental=True")
+        return _return_stream_response(response, deadline)
     return _complete_sql_impl(function, model, prompt, options, session)
 
 
@@ -275,7 +254,6 @@ def Complete(
     *,
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
-    use_rest_api_experimental: bool = False,
     stream: bool = False,
     timeout: Optional[float] = None,
     deadline: Optional[float] = None,
@@ -287,16 +265,13 @@ def Complete(
         prompt: A Column of prompts to send to the LLM.
         options: A instance of snowflake.cortex.CompleteOptions
         session: The snowpark session to use. Will be inferred by context if not specified.
-        use_rest_api_experimental (bool): Toggles between the use of SQL and REST implementation. This feature is
-            experimental and can be removed at any time.
         stream (bool): Enables streaming. When enabled, a generator function is returned that provides the streaming
             output as it is received. Each update is a string containing the new text content since the previous update.
-            The use of streaming requires the experimental use_rest_api_experimental flag to be enabled.
         timeout (float): Timeout in seconds to retry failed REST requests.
         deadline (float): Time in seconds since the epoch (as returned by time.time()) to retry failed REST requests.
 
     Raises:
-        ValueError: If `stream` is set to True and `use_rest_api_experimental` is set to False.
+        ValueError: incorrect argument.
 
     Returns:
         A column of string responses.
@@ -307,7 +282,6 @@ def Complete(
             prompt,
             options=options,
             session=session,
-            use_rest_api_experimental=use_rest_api_experimental,
             stream=stream,
             timeout=timeout,
             deadline=deadline,

{snowflake_ml_python-1.6.0 → snowflake_ml_python-1.6.2}/snowflake/ml/_internal/env_utils.py

@@ -27,7 +27,6 @@ class CONDA_OS(Enum):
     NO_ARCH = "noarch"
 
 
-_SNOWFLAKE_CONDA_CHANNEL_URL = "https://repo.anaconda.com/pkgs/snowflake"
 _NODEFAULTS = "nodefaults"
 _SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE: Dict[str, List[version.Version]] = {}
 _SNOWFLAKE_CONDA_PACKAGE_CACHE: Dict[str, List[version.Version]] = {}
@@ -36,6 +35,7 @@ _SUPPORTED_PACKAGE_SPEC_OPS = ["==", ">=", "<=", ">", "<"]
 DEFAULT_CHANNEL_NAME = ""
 SNOWML_SPROC_ENV = "IN_SNOWML_SPROC"
 SNOWPARK_ML_PKG_NAME = "snowflake-ml-python"
+SNOWFLAKE_CONDA_CHANNEL_URL = "https://repo.anaconda.com/pkgs/snowflake"
 
 
 def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
@@ -370,7 +370,7 @@ def get_matched_package_versions_in_snowflake_conda_channel(
 
     assert not snowpark_utils.is_in_stored_procedure()  # type: ignore[no-untyped-call]
 
-    url = f"{_SNOWFLAKE_CONDA_CHANNEL_URL}/{conda_os.value}/repodata.json"
+    url = f"{SNOWFLAKE_CONDA_CHANNEL_URL}/{conda_os.value}/repodata.json"
 
     if req.name not in _SNOWFLAKE_CONDA_PACKAGE_CACHE:
         try:
@@ -477,6 +477,7 @@ def save_conda_env_file(
     path: pathlib.Path,
     conda_chan_deps: DefaultDict[str, List[requirements.Requirement]],
     python_version: str,
+    default_channel_override: str = SNOWFLAKE_CONDA_CHANNEL_URL,
 ) -> None:
     """Generate conda.yml file given a dict of dependencies after validation.
     The channels part of conda.yml file will contains Snowflake Anaconda Channel, nodefaults and all channel names
@@ -489,6 +490,7 @@ def save_conda_env_file(
         path: Path to the conda.yml file.
         conda_chan_deps: Dict of conda dependencies after validated.
         python_version: A string 'major.minor' showing python version relate to model.
+        default_channel_override: The default channel to be put in the first place of the channels section.
     """
     assert path.suffix in [".yml", ".yaml"], "Conda environment file should have extension of yml or yaml."
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -499,7 +501,11 @@ def save_conda_env_file(
     channels = list(dict(sorted(conda_chan_deps.items(), key=lambda item: len(item[1]), reverse=True)).keys())
     if DEFAULT_CHANNEL_NAME in channels:
         channels.remove(DEFAULT_CHANNEL_NAME)
-    env["channels"] = [_SNOWFLAKE_CONDA_CHANNEL_URL] + channels + [_NODEFAULTS]
+
+    if default_channel_override in channels:
+        channels.remove(default_channel_override)
+
+    env["channels"] = [default_channel_override] + channels + [_NODEFAULTS]
     env["dependencies"] = [f"python=={python_version}.*"]
     for chan, reqs in conda_chan_deps.items():
         env["dependencies"].extend(
@@ -567,8 +573,8 @@ def load_conda_env_file(
     python_version = None
 
     channels = env.get("channels", [])
-    if _SNOWFLAKE_CONDA_CHANNEL_URL in channels:
-        channels.remove(_SNOWFLAKE_CONDA_CHANNEL_URL)
+    if len(channels) >= 1:
+        channels = channels[1:]  # Skip the first channel which is the default channel
     if _NODEFAULTS in channels:
         channels.remove(_NODEFAULTS)
 

{snowflake_ml_python-1.6.0 → snowflake_ml_python-1.6.2}/snowflake/ml/_internal/exceptions/modeling_error_messages.py

@@ -4,7 +4,10 @@ ATTRIBUTE_NOT_SET = (
     "-differences."
 )
 SIZE_MISMATCH = "Size mismatch: {}={}, {}={}."
-INVALID_MODEL_PARAM = "Invalid parameter {} for model {}. Valid parameters: {}."
+INVALID_MODEL_PARAM = (
+    "Invalid parameter {} for model {}. Valid parameters: {}."
+    "Note: Scikit learn params cannot be set until the model has been fit."
+)
 UNSUPPORTED_MODEL_CONVERSION = "Object doesn't support {}. Please use {}."
 INCOMPATIBLE_NEW_SKLEARN_PARAM = "Incompatible scikit-learn version: {} requires scikit-learn>={}. Installed: {}."
 REMOVED_SKLEARN_PARAM = "Incompatible scikit-learn version: {} is removed in scikit-learn>={}. Installed: {}."

{snowflake_ml_python-1.6.0 → snowflake_ml_python-1.6.2}/snowflake/ml/_internal/telemetry.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import contextvars
 import enum
 import functools
 import inspect
@@ -12,6 +13,7 @@ from typing import (
     List,
     Mapping,
     Optional,
+    Set,
     Tuple,
     TypeVar,
     Union,
@@ -28,7 +30,7 @@ from snowflake.ml._internal.exceptions import (
     exceptions as snowml_exceptions,
 )
 from snowflake.snowpark import dataframe, exceptions as snowpark_exceptions, session
-from snowflake.snowpark._internal import utils
+from snowflake.snowpark._internal import server_connection, utils
 
 _log_counter = 0
 _FLUSH_SIZE = 10
@@ -44,6 +46,20 @@ _Args = ParamSpec("_Args")
 _ReturnValue = TypeVar("_ReturnValue")
 
 
+@enum.unique
+class TelemetryProject(enum.Enum):
+    MLOPS = "MLOps"
+    MODELING = "ModelDevelopment"
+    # TODO: Update with remaining projects.
+
+
+@enum.unique
+class TelemetrySubProject(enum.Enum):
+    MONITORING = "Monitoring"
+    REGISTRY = "ModelManagement"
+    # TODO: Update with remaining subprojects.
+
+
 @enum.unique
 class TelemetryField(enum.Enum):
     # constants
@@ -71,6 +87,122 @@ class TelemetryField(enum.Enum):
     FUNC_CAT_USAGE = "usage"
 
 
+class _TelemetrySourceType(enum.Enum):
+    # Automatically inferred telemetry/statement parameters
+    AUTO_TELEMETRY = "SNOWML_AUTO_TELEMETRY"
+    # Mixture of manual and automatic telemetry/statement parameters
+    AUGMENT_TELEMETRY = "SNOWML_AUGMENT_TELEMETRY"
+
+
+_statement_params_context_var: contextvars.ContextVar[Dict[str, str]] = contextvars.ContextVar("statement_params")
+
+
+class _StatementParamsPatchManager:
+    def __init__(self) -> None:
+        self._patch_cache: Set[server_connection.ServerConnection] = set()
+        self._context_var: contextvars.ContextVar[Dict[str, str]] = _statement_params_context_var
+
+    def apply_patches(self) -> None:
+        try:
+            # Apply patching to all active sessions in case of multiple
+            for sess in session._get_active_sessions():
+                # Check patch cache here to avoid unnecessary context switches
+                if self._get_target(sess) not in self._patch_cache:
+                    self._patch_session(sess)
+        except snowpark_exceptions.SnowparkSessionException:
+            pass
+
+    def set_statement_params(self, statement_params: Dict[str, str]) -> None:
+        # Only set value if not already set in context
+        if not self._context_var.get({}):
+            self._context_var.set(statement_params)
+
+    def _get_target(self, session: session.Session) -> server_connection.ServerConnection:
+        return cast(server_connection.ServerConnection, session._conn)
+
+    def _patch_session(self, session: session.Session, throw_on_patch_fail: bool = False) -> None:
+        # Extract target
+        try:
+            target = self._get_target(session)
+        except AttributeError:
+            if throw_on_patch_fail:
+                raise
+            # TODO: Log a warning, this probably means there was a breaking change in Snowpark/SnowflakeConnection
+            return
+
+        # Check if session has already been patched
+        if target in self._patch_cache:
+            return
+        self._patch_cache.add(target)
+
+        functions = [
+            ("execute_and_notify_query_listener", "_statement_params"),
+            ("execute_async_and_notify_query_listener", "_statement_params"),
+        ]
+
+        for func, param_name in functions:
+            try:
+                self._patch_with_statement_params(target, func, param_name=param_name)
+            except AttributeError:
+                if throw_on_patch_fail:  # primarily used for testing
+                    raise
+                # TODO: Log a warning, this probably means there was a breaking change in Snowpark/SnowflakeConnection
+                pass
+
+    def _patch_with_statement_params(
+        self, target: object, function_name: str, param_name: str = "statement_params"
+    ) -> None:
+        func = getattr(target, function_name)
+        assert callable(func)
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Retrieve context level statement parameters
+            context_params = self._context_var.get(dict())
+            if not context_params:
+                # Exit early if not in SnowML (decorator) context
+                return func(*args, **kwargs)
+
+            # Extract any explicitly provided statement parameters
+            orig_kwargs = dict(kwargs)
+            in_params = kwargs.pop(param_name, None) or {}
+
+            # Inject a special flag to statement parameters so we can filter out these patched logs if necessary
+            # Calls that include SnowML telemetry are tagged with "SNOWML_AUGMENT_TELEMETRY"
+            # and calls without SnowML telemetry are tagged with "SNOWML_AUTO_TELEMETRY"
+            if TelemetryField.KEY_PROJECT.value in in_params:
+                context_params["snowml_telemetry_type"] = _TelemetrySourceType.AUGMENT_TELEMETRY.value
+            else:
+                context_params["snowml_telemetry_type"] = _TelemetrySourceType.AUTO_TELEMETRY.value
+
+            # Apply any explicitly provided statement parameters and result into function call
+            context_params.update(in_params)
+            kwargs[param_name] = context_params
+
+            try:
+                return func(*args, **kwargs)
+            except TypeError as e:
+                if str(e).endswith(f"unexpected keyword argument '{param_name}'"):
+                    # TODO: Log warning that this patch is invalid
+                    # Unwrap function for future invocations
+                    setattr(target, function_name, func)
+                    return func(*args, **orig_kwargs)
+                else:
+                    raise
+
+        setattr(target, function_name, wrapper)
+
+    def __getstate__(self) -> Dict[str, Any]:
+        return {}
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        # unpickling does not call __init__ by default, do it manually here
+        self.__init__()  # type: ignore[misc]
+
+
+_patch_manager = _StatementParamsPatchManager()
+
+
 def get_statement_params(
     project: str, subproject: Optional[str] = None, class_name: Optional[str] = None
 ) -> Dict[str, Any]:
@@ -361,7 +493,18 @@ def send_api_usage_telemetry(
                         obj._statement_params = statement_params  # type: ignore[assignment]
                 return obj
 
+            # Set up framework-level credit usage instrumentation
+            ctx = contextvars.copy_context()
+            _patch_manager.apply_patches()
+
+            # This function should be executed with ctx.run()
+            def execute_func_with_statement_params() -> _ReturnValue:
+                _patch_manager.set_statement_params(statement_params)
+                result = func(*args, **kwargs)
+                return update_stmt_params_if_snowpark_df(result, statement_params)
+
             # prioritize `conn_attr_name` over the active session
+            telemetry_enabled = True
             if conn_attr_name:
                 # raise AttributeError if conn attribute does not exist in `self`
                 conn = operator.attrgetter(conn_attr_name)(args[0])
@@ -373,22 +516,17 @@ def send_api_usage_telemetry(
             else:
                 try:
                     active_session = next(iter(session._get_active_sessions()))
-                # server no default session
+                    conn = active_session._conn._conn
+                    telemetry_enabled = active_session.telemetry_enabled
                 except snowpark_exceptions.SnowparkSessionException:
-                    try:
-                        return update_stmt_params_if_snowpark_df(func(*args, **kwargs), statement_params)
-                    except Exception as e:
-                        if isinstance(e, snowml_exceptions.SnowflakeMLException):
-                            raise e.original_exception.with_traceback(e.__traceback__) from None
-                        # suppress SnowparkSessionException from telemetry in the stack trace
-                        raise e from None
-
-                conn = active_session._conn._conn
-                if (not active_session.telemetry_enabled) or (conn is None):
-                    try:
-                        return update_stmt_params_if_snowpark_df(func(*args, **kwargs), statement_params)
-                    except snowml_exceptions.SnowflakeMLException as e:
-                        raise e.original_exception from e
+                    conn = None
+
+            if conn is None or not telemetry_enabled:
+                # Telemetry not enabled, just execute without our additional telemetry logic
+                try:
+                    return ctx.run(execute_func_with_statement_params)
+                except snowml_exceptions.SnowflakeMLException as e:
+                    raise e.original_exception from e
 
             # TODO(hayu): [SNOW-750287] Optimize telemetry client to a singleton.
             telemetry = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject_name)
@@ -401,11 +539,11 @@ def send_api_usage_telemetry(
                 custom_tags=custom_tags,
             )
             try:
-                res = func(*args, **kwargs)
+                return ctx.run(execute_func_with_statement_params)
             except Exception as e:
                 if not isinstance(e, snowml_exceptions.SnowflakeMLException):
                     # already handled via a nested decorated function
-                    if hasattr(e, "_snowflake_ml_handled") and e._snowflake_ml_handled:
+                    if getattr(e, "_snowflake_ml_handled", False):
                         raise e
                     if isinstance(e, snowpark_exceptions.SnowparkClientException):
                         me = snowml_exceptions.SnowflakeMLException(
@@ -424,8 +562,6 @@ def send_api_usage_telemetry(
                     raise me.original_exception from None
                 else:
                     raise me.original_exception from e
-            else:
-                return update_stmt_params_if_snowpark_df(res, statement_params)
             finally:
                 telemetry.send_function_usage_telemetry(**telemetry_args)
                 global _log_counter