snowflake-ml-python 1.6.1__tar.gz → 1.6.3__tar.gz

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/CHANGELOG.md

@@ -1,6 +1,43 @@
 # Release History
 
-## 1.6.1 (TBD)
+## 1.6.3
+
+- Model Registry (PrPr) has been removed.
+
+### Bug Fixes
+
+- Registry: Fix a bug that when package whose name does not follow PEP-508 is provided when logging the model,
+  an unexpected normalization is happening.
+- Registry: Fix `not a valid remote uri` error when logging mlflow models.
+- Registry: Fix a bug that `ModelVersion.run` is called in a nested way.
+- Registry: Fix an issue that leads to `log_model` failure when local package version contains parts other than
+  base version.
+
+### New Features
+
+- Data: Improve `DataConnector.to_pandas()` performance when loading from Snowpark DataFrames.
+- Model Registry: Allow users to set a model task while using `log_model`.
+- Feature Store: FeatureView supports ON_CREATE or ON_SCHEDULE initialize mode.
+
+## 1.6.2 (2024-09-04)
+
+### 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.
+
+## 1.6.1 (2024-08-12)
 
 ### Bug Fixes
 
@@ -17,14 +54,11 @@
 ### New Features
 
 - Enable `set_params` to set the parameters of the underlying sklearn estimator, if the snowflake-ml model has been fit.
-- Data: Add top-level exports for `DataConnector` and `DataSource` to `snowflake.ml.data`.
 - 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

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: snowflake-ml-python
-Version: 1.6.1
+Version: 1.6.3
 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: 
@@ -247,13 +247,13 @@ Requires-Dist: pytimeparse<2,>=1.1.8
 Requires-Dist: pyyaml<7,>=6.0
 Requires-Dist: retrying<2,>=1.3.3
 Requires-Dist: s3fs<2024,>=2022.11
-Requires-Dist: scikit-learn<1.4,>=1.2.1
+Requires-Dist: scikit-learn<1.6,>=1.2.1
 Requires-Dist: scipy<2,>=1.9
 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"
@@ -261,7 +261,7 @@ Requires-Dist: mlflow<2.4,>=2.1.0; extra == "all"
 Requires-Dist: peft<1,>=0.5.0; extra == "all"
 Requires-Dist: sentence-transformers<3,>=2.2.2; extra == "all"
 Requires-Dist: sentencepiece<1,>=0.1.95; extra == "all"
-Requires-Dist: shap==0.42.1; extra == "all"
+Requires-Dist: shap<1,>=0.42.0; extra == "all"
 Requires-Dist: tensorflow<3,>=2.10; extra == "all"
 Requires-Dist: tokenizers<1,>=0.10; extra == "all"
 Requires-Dist: torch<2.3.0,>=2.0.1; extra == "all"
@@ -276,7 +276,7 @@ Requires-Dist: peft<1,>=0.5.0; extra == "llm"
 Provides-Extra: mlflow
 Requires-Dist: mlflow<2.4,>=2.1.0; extra == "mlflow"
 Provides-Extra: shap
-Requires-Dist: shap==0.42.1; extra == "shap"
+Requires-Dist: shap<1,>=0.42.0; extra == "shap"
 Provides-Extra: tensorflow
 Requires-Dist: tensorflow<3,>=2.10; extra == "tensorflow"
 Provides-Extra: torch
@@ -373,7 +373,44 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
 
 # Release History
 
-## 1.6.1 (TBD)
+## 1.6.3
+
+- Model Registry (PrPr) has been removed.
+
+### Bug Fixes
+
+- Registry: Fix a bug that when package whose name does not follow PEP-508 is provided when logging the model,
+  an unexpected normalization is happening.
+- Registry: Fix `not a valid remote uri` error when logging mlflow models.
+- Registry: Fix a bug that `ModelVersion.run` is called in a nested way.
+- Registry: Fix an issue that leads to `log_model` failure when local package version contains parts other than
+  base version.
+
+### New Features
+
+- Data: Improve `DataConnector.to_pandas()` performance when loading from Snowpark DataFrames.
+- Model Registry: Allow users to set a model task while using `log_model`.
+- Feature Store: FeatureView supports ON_CREATE or ON_SCHEDULE initialize mode.
+
+## 1.6.2 (2024-09-04)
+
+### 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.
+
+## 1.6.1 (2024-08-12)
 
 ### Bug Fixes
 
@@ -390,14 +427,11 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
 ### New Features
 
 - Enable `set_params` to set the parameters of the underlying sklearn estimator, if the snowflake-ml model has been fit.
-- Data: Add top-level exports for `DataConnector` and `DataSource` to `snowflake.ml.data`.
 - 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

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/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.6", "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"
@@ -24,12 +24,12 @@ Issues = "https://github.com/snowflakedb/snowflake-ml-python/issues"
 Changelog = "https://github.com/snowflakedb/snowflake-ml-python/blob/master/CHANGELOG.md"
 
 [project.optional-dependencies]
-all = [ "catboost>=1.2.0, <2", "lightgbm>=3.3.5,<5", "mlflow>=2.1.0,<2.4", "peft>=0.5.0,<1", "sentence-transformers>=2.2.2,<3", "sentencepiece>=0.1.95,<1", "shap==0.42.1", "tensorflow>=2.10,<3", "tokenizers>=0.10,<1", "torch>=2.0.1,<2.3.0", "torchdata>=0.4,<1", "transformers>=4.32.1,<5",]
+all = [ "catboost>=1.2.0, <2", "lightgbm>=3.3.5,<5", "mlflow>=2.1.0,<2.4", "peft>=0.5.0,<1", "sentence-transformers>=2.2.2,<3", "sentencepiece>=0.1.95,<1", "shap>=0.42.0,<1", "tensorflow>=2.10,<3", "tokenizers>=0.10,<1", "torch>=2.0.1,<2.3.0", "torchdata>=0.4,<1", "transformers>=4.32.1,<5",]
 catboost = [ "catboost>=1.2.0, <2",]
 lightgbm = [ "lightgbm>=3.3.5,<5",]
 llm = [ "peft>=0.5.0,<1",]
 mlflow = [ "mlflow>=2.1.0,<2.4",]
-shap = [ "shap==0.42.1",]
+shap = [ "shap>=0.42.0,<1",]
 tensorflow = [ "tensorflow>=2.10,<3",]
 torch = [ "torch>=2.0.1,<2.3.0", "torchdata>=0.4,<1",]
 transformers = [ "sentence-transformers>=2.2.2,<3", "sentencepiece>=0.1.95,<1", "tokenizers>=0.10,<1", "transformers>=4.32.1,<5",]

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/cortex/__init__.py

@@ -1,5 +1,7 @@
 from snowflake.cortex._classify_text import ClassifyText
 from snowflake.cortex._complete import Complete, CompleteOptions
+from snowflake.cortex._embed_text_768 import EmbedText768
+from snowflake.cortex._embed_text_1024 import EmbedText1024
 from snowflake.cortex._extract_answer import ExtractAnswer
 from snowflake.cortex._sentiment import Sentiment
 from snowflake.cortex._summarize import Summarize
@@ -9,6 +11,8 @@ __all__ = [
     "ClassifyText",
     "Complete",
     "CompleteOptions",
+    "EmbedText768",
+    "EmbedText1024",
     "ExtractAnswer",
     "Sentiment",
     "Summarize",

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/cortex/_classify_text.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -33,4 +33,4 @@ def _classify_text_impl(
     categories: Union[List[str], snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, str_input, categories)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, str_input, categories))

snowflake_ml_python-1.6.3/snowflake/cortex/_embed_text_1024.py

@@ -0,0 +1,37 @@
+from typing import List, Optional, Union, cast
+
+from snowflake import snowpark
+from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.ml._internal import telemetry
+
+
+@telemetry.send_api_usage_telemetry(
+    project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+)
+def EmbedText1024(
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    """TextEmbed calls into the LLM inference service to embed the text.
+
+    Args:
+        model: A Column of strings representing the model to use for embedding. The value
+               of the strings must be within the SUPPORTED_MODELS list.
+        text: A Column of strings representing input text.
+        session: The snowpark session to use. Will be inferred by context if not specified.
+
+    Returns:
+        A column of vectors containing embeddings.
+    """
+
+    return _embed_text_1024_impl("snowflake.cortex.embed_text_1024", model, text, session=session)
+
+
+def _embed_text_1024_impl(
+    function: str,
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))

snowflake_ml_python-1.6.3/snowflake/cortex/_embed_text_768.py

@@ -0,0 +1,37 @@
+from typing import List, Optional, Union, cast
+
+from snowflake import snowpark
+from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.ml._internal import telemetry
+
+
+@telemetry.send_api_usage_telemetry(
+    project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+)
+def EmbedText768(
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    """TextEmbed calls into the LLM inference service to embed the text.
+
+    Args:
+        model: A Column of strings representing the model to use for embedding. The value
+               of the strings must be within the SUPPORTED_MODELS list.
+        text: A Column of strings representing input text.
+        session: The snowpark session to use. Will be inferred by context if not specified.
+
+    Returns:
+        A column of vectors containing embeddings.
+    """
+
+    return _embed_text_768_impl("snowflake.cortex.embed_text_768", model, text, session=session)
+
+
+def _embed_text_768_impl(
+    function: str,
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/cortex/_extract_answer.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -33,4 +33,4 @@ def _extract_answer_impl(
     question: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, from_text, question)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, from_text, question))

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/cortex/_sentiment.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -30,4 +30,4 @@ def _sentiment_impl(
     output = call_sql_function(function, session, text)
     if isinstance(output, snowpark.Column):
         return output
-    return float(output)
+    return float(cast(str, output))

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/cortex/_summarize.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -30,4 +30,4 @@ def _summarize_impl(
     text: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, text)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, text))

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/cortex/_translate.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -36,4 +36,4 @@ def _translate_impl(
     to_language: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, text, from_language, to_language)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, text, from_language, to_language))

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/cortex/_util.py

@@ -24,7 +24,7 @@ def call_sql_function(
     function: str,
     session: Optional[snowpark.Session],
     *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
-) -> Union[str, snowpark.Column]:
+) -> Union[str, List[float], snowpark.Column]:
     handle_as_column = False
 
     for arg in args:
@@ -32,9 +32,9 @@ def call_sql_function(
             handle_as_column = True
 
     if handle_as_column:
-        return cast(Union[str, snowpark.Column], _call_sql_function_column(function, *args))
+        return cast(Union[str, List[float], snowpark.Column], _call_sql_function_column(function, *args))
     return cast(
-        Union[str, snowpark.Column],
+        Union[str, List[float], snowpark.Column],
         _call_sql_function_immediate(function, session, *args),
     )
 
@@ -49,7 +49,7 @@ def _call_sql_function_immediate(
     function: str,
     session: Optional[snowpark.Session],
     *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
-) -> str:
+) -> Union[str, List[float]]:
     session = session or context.get_active_session()
     if session is None:
         raise SnowflakeAuthenticationException(

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/ml/_internal/env_utils.py

@@ -9,7 +9,7 @@ from importlib import metadata as importlib_metadata
 from typing import Any, DefaultDict, Dict, List, Optional, Tuple
 
 import yaml
-from packaging import requirements, specifiers, utils as packaging_utils, version
+from packaging import requirements, specifiers, version
 
 import snowflake.connector
 from snowflake.ml._internal import env as snowml_env
@@ -54,15 +54,12 @@ def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
     """
     try:
         r = requirements.Requirement(req_str)
-        r.name = packaging_utils.canonicalize_name(r.name)
 
         if r.name == "python":
             raise ValueError("Don't specify python as a dependency, use python version argument instead.")
     except requirements.InvalidRequirement:
         raise ValueError(f"Invalid package requirement {req_str} found.")
 
-    if r.marker:
-        raise ValueError("Markers is not supported in conda dependency.")
     return r
 
 
@@ -84,6 +81,8 @@ def _validate_conda_dependency_string(dep_str: str) -> Tuple[str, requirements.R
     channel_str, _, requirement_str = dep_str.rpartition("::")
     r = _validate_pip_requirement_string(requirement_str)
     if channel_str != "pip":
+        if r.marker:
+            raise ValueError("Markers is not supported in conda dependency.")
         if r.extras:
             raise ValueError("Extras is not supported in conda dependency.")
         if r.url:
@@ -221,7 +220,7 @@ def get_local_installed_version_of_pip_package(pip_req: requirements.Requirement
         else:
             return pip_req
     new_pip_req = copy.deepcopy(pip_req)
-    new_pip_req.specifier = specifiers.SpecifierSet(specifiers=f"=={local_dist_version}")
+    new_pip_req.specifier = specifiers.SpecifierSet(specifiers=f"=={version.parse(local_dist_version).base_version}")
     if not pip_req.specifier.contains(local_dist_version):
         warnings.warn(
             f"Package requirement {str(pip_req)} specified, while version {local_dist_version} is installed. "
@@ -513,6 +512,7 @@ def save_conda_env_file(
         )
 
     with open(path, "w", encoding="utf-8") as f:
+        yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
         yaml.safe_dump(env, stream=f, default_flow_style=False)
 
 

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/snowflake/ml/_internal/exceptions/error_codes.py

@@ -58,6 +58,8 @@ NOT_FOUND = "2101"
 METHOD_NOT_ALLOWED = "2102"
 # Not implemented.
 NOT_IMPLEMENTED = "2103"
+# User needs to opt in to use a feature.
+OPT_IN_REQUIRED = "2104"
 
 # Calling an API with unsupported keywords/values.
 INVALID_ARGUMENT = "2110"

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.3}/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
@@ -85,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]:
@@ -375,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])
@@ -387,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)
@@ -415,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(
@@ -438,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