snowflake-ml-python 1.9.0__tar.gz → 1.9.2__tar.gz

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.2}/CHANGELOG.md

@@ -1,5 +1,91 @@
 # Release History
 
+## 1.9.2
+
+### Bug Fixes
+
+- DataConnector: Fix `self._session` related errors inside Container Runtime.
+- Registry: Fix a bug when trying to pass `None` to array (`pd.dtype('O')`) in signature and pandas data handler.
+
+### New Features
+
+- Experiment Tracking (PrPr): Automatically log the model, metrics, and parameters while training
+  XGBoost and LightGBM models.
+
+```python
+from snowflake.ml.experiment import ExperimentTracking
+from snowflake.ml.experiment.callback import SnowflakeXgboostCallback, SnowflakeLightgbmCallback
+
+exp = ExperimentTracking(session=sp_session, database_name="ML", schema_name="PUBLIC")
+
+exp.set_experiment("MY_EXPERIMENT")
+
+# XGBoost
+callback = SnowflakeXgboostCallback(
+  exp, log_model=True, log_metrics=True, log_params=True, model_name="model_name", model_signature=sig
+)
+model = XGBClassifier(callbacks=[callback])
+with exp.start_run():
+  model.fit(X, y, eval_set=[(X_test, y_test)])
+
+# LightGBM
+callback = SnowflakeLightgbmCallback(
+  exp, log_model=True, log_metrics=True, log_params=True, model_name="model_name", model_signature=sig
+)
+model = LGBMClassifier()
+with exp.start_run():
+  model.fit(X, y, eval_set=[(X_test, y_test)], callbacks=[callback])
+```
+
+## 1.9.1 (07-18-2025)
+
+### Bug Fixes
+
+- Registry: Fix a bug when trying to set the PAD token the HuggingFace `text-generation` model had multiple EOS tokens.
+  The handler picks the first EOS token as PAD token now.
+
+### New Features
+
+- DataConnector: DataConnector objects can now be pickled
+- Dataset: Dataset objects can now be pickled
+- Registry (PrPr): Introducing `create_service` function in `snowflake/ml/model/models/huggingface_pipeline.py`
+  which creates a service to log a HF model and upon successful logging, an inference service is created.
+
+```python
+from snowflake.ml.model.models import huggingface_pipeline
+
+hf_model_ref = huggingface_pipeline.HuggingFacePipelineModel(
+  model="gpt2",
+  task="text-generation", # Optional
+)
+
+
+hf_model_ref.create_service(
+    session=session,
+    service_name="test_service",
+    service_compute_pool="test_compute_pool",
+    image_repo="test_repo",
+    ...
+)
+```
+
+- Experiment Tracking (PrPr): New module for managing and tracking ML experiments in Snowflake.
+
+```python
+from snowflake.ml.experiment import ExperimentTracking
+
+exp = ExperimentTracking(session=sp_session, database_name="ML", schema_name="PUBLIC")
+
+exp.set_experiment("MY_EXPERIMENT")
+
+with exp.start_run():
+  exp.log_param("batch_size", 32)
+  exp.log_metrics("accuracy", 0.98, step=10)
+  exp.log_model(my_model, model_name="MY_MODEL")
+```
+
+- Registry: Added support for wide input (500+ features) for inference done using SPCS
+
 ## 1.9.0
 
 ### Bug Fixes
@@ -7,6 +93,19 @@
 - Registry: Fixed bug causing snowpark to pandas dataframe conversion to fail when `QUOTED_IDENTIFIERS_IGNORE_CASE`
   parameter is enabled
 - Registry: Fixed duplicate UserWarning logs during model packaging
+- Registry: If the huggingface pipeline text-generation model doesn't contain a default chat template, a ChatML template
+  is assigned to the tokenizer.
+
+```shell
+{% for message in messages %}
+  {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
+{% endfor %}
+{% if add_generation_prompt %}
+  {{ '<|im_start|>assistant\n' }}
+{% endif %}"
+```
+
+- Registry: Fixed SQL queries during registry initialization that were forcing warehouse requirement
 
 ### Behavior Changes
 
@@ -116,7 +215,8 @@
   - Pre-created Snowpark Session is now available inside job payloads using
     `snowflake.snowpark.context.get_active_session()`
 - Registry: Introducing `save_location` to `log_model` using the `options` argument.
-  User's can provide the path to write the model version's files that get stored in Snowflake's stage.
+  Users can use the `save_location` option to specify a local directory where the model files and configuration are written.
+  This is useful when the default temporary directory has space limitations.
 
 ```python
 reg.log_model(

{snowflake_ml_python-1.9.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.9.0
+Version: 1.9.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: 
@@ -240,9 +240,10 @@ Requires-Dist: cloudpickle>=2.0.0
 Requires-Dist: cryptography
 Requires-Dist: fsspec[http]<2026,>=2024.6.1
 Requires-Dist: importlib_resources<7,>=6.1.1
-Requires-Dist: numpy<2,>=1.23
+Requires-Dist: numpy<3,>=1.23
 Requires-Dist: packaging<25,>=20.9
 Requires-Dist: pandas<3,>=2.1.4
+Requires-Dist: platformdirs<5
 Requires-Dist: pyarrow
 Requires-Dist: pydantic<3,>=2.8.2
 Requires-Dist: pyjwt<3,>=2.0.0
@@ -257,6 +258,7 @@ Requires-Dist: snowflake-connector-python[pandas]<4,>=3.15.0
 Requires-Dist: snowflake-snowpark-python!=1.26.0,<2,>=1.17.0
 Requires-Dist: snowflake.core<2,>=1.0.2
 Requires-Dist: sqlparse<1,>=0.4
+Requires-Dist: tqdm<5
 Requires-Dist: typing-extensions<5,>=4.1.0
 Requires-Dist: xgboost<3,>=1.7.3
 Provides-Extra: all
@@ -272,7 +274,7 @@ Requires-Dist: tensorflow<3,>=2.17.0; extra == "all"
 Requires-Dist: tokenizers<1,>=0.15.1; extra == "all"
 Requires-Dist: torch<3,>=2.0.1; extra == "all"
 Requires-Dist: torchdata<1,>=0.4; extra == "all"
-Requires-Dist: transformers<5,>=4.39.3; extra == "all"
+Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "all"
 Provides-Extra: altair
 Requires-Dist: altair<6,>=5; extra == "altair"
 Provides-Extra: catboost
@@ -297,7 +299,7 @@ Requires-Dist: sentence-transformers<4,>=2.7.0; extra == "transformers"
 Requires-Dist: sentencepiece<0.2.0,>=0.1.95; extra == "transformers"
 Requires-Dist: tokenizers<1,>=0.15.1; extra == "transformers"
 Requires-Dist: torch<3,>=2.0.1; extra == "transformers"
-Requires-Dist: transformers<5,>=4.39.3; extra == "transformers"
+Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "transformers"
 Dynamic: license-file
 
 # Snowpark ML
@@ -408,6 +410,92 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
+## 1.9.2
+
+### Bug Fixes
+
+- DataConnector: Fix `self._session` related errors inside Container Runtime.
+- Registry: Fix a bug when trying to pass `None` to array (`pd.dtype('O')`) in signature and pandas data handler.
+
+### New Features
+
+- Experiment Tracking (PrPr): Automatically log the model, metrics, and parameters while training
+  XGBoost and LightGBM models.
+
+```python
+from snowflake.ml.experiment import ExperimentTracking
+from snowflake.ml.experiment.callback import SnowflakeXgboostCallback, SnowflakeLightgbmCallback
+
+exp = ExperimentTracking(session=sp_session, database_name="ML", schema_name="PUBLIC")
+
+exp.set_experiment("MY_EXPERIMENT")
+
+# XGBoost
+callback = SnowflakeXgboostCallback(
+  exp, log_model=True, log_metrics=True, log_params=True, model_name="model_name", model_signature=sig
+)
+model = XGBClassifier(callbacks=[callback])
+with exp.start_run():
+  model.fit(X, y, eval_set=[(X_test, y_test)])
+
+# LightGBM
+callback = SnowflakeLightgbmCallback(
+  exp, log_model=True, log_metrics=True, log_params=True, model_name="model_name", model_signature=sig
+)
+model = LGBMClassifier()
+with exp.start_run():
+  model.fit(X, y, eval_set=[(X_test, y_test)], callbacks=[callback])
+```
+
+## 1.9.1 (07-18-2025)
+
+### Bug Fixes
+
+- Registry: Fix a bug when trying to set the PAD token the HuggingFace `text-generation` model had multiple EOS tokens.
+  The handler picks the first EOS token as PAD token now.
+
+### New Features
+
+- DataConnector: DataConnector objects can now be pickled
+- Dataset: Dataset objects can now be pickled
+- Registry (PrPr): Introducing `create_service` function in `snowflake/ml/model/models/huggingface_pipeline.py`
+  which creates a service to log a HF model and upon successful logging, an inference service is created.
+
+```python
+from snowflake.ml.model.models import huggingface_pipeline
+
+hf_model_ref = huggingface_pipeline.HuggingFacePipelineModel(
+  model="gpt2",
+  task="text-generation", # Optional
+)
+
+
+hf_model_ref.create_service(
+    session=session,
+    service_name="test_service",
+    service_compute_pool="test_compute_pool",
+    image_repo="test_repo",
+    ...
+)
+```
+
+- Experiment Tracking (PrPr): New module for managing and tracking ML experiments in Snowflake.
+
+```python
+from snowflake.ml.experiment import ExperimentTracking
+
+exp = ExperimentTracking(session=sp_session, database_name="ML", schema_name="PUBLIC")
+
+exp.set_experiment("MY_EXPERIMENT")
+
+with exp.start_run():
+  exp.log_param("batch_size", 32)
+  exp.log_metrics("accuracy", 0.98, step=10)
+  exp.log_model(my_model, model_name="MY_MODEL")
+```
+
+- Registry: Added support for wide input (500+ features) for inference done using SPCS
+
 ## 1.9.0
 
 ### Bug Fixes
@@ -415,6 +503,19 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 - Registry: Fixed bug causing snowpark to pandas dataframe conversion to fail when `QUOTED_IDENTIFIERS_IGNORE_CASE`
   parameter is enabled
 - Registry: Fixed duplicate UserWarning logs during model packaging
+- Registry: If the huggingface pipeline text-generation model doesn't contain a default chat template, a ChatML template
+  is assigned to the tokenizer.
+
+```shell
+{% for message in messages %}
+  {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
+{% endfor %}
+{% if add_generation_prompt %}
+  {{ '<|im_start|>assistant\n' }}
+{% endif %}"
+```
+
+- Registry: Fixed SQL queries during registry initialization that were forcing warehouse requirement
 
 ### Behavior Changes
 
@@ -524,7 +625,8 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
   - Pre-created Snowpark Session is now available inside job payloads using
     `snowflake.snowpark.context.get_active_session()`
 - Registry: Introducing `save_location` to `log_model` using the `options` argument.
-  User's can provide the path to write the model version's files that get stored in Snowflake's stage.
+  Users can use the `save_location` option to specify a local directory where the model files and configuration are written.
+  This is useful when the default temporary directory has space limitations.
 
 ```python
 reg.log_model(

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.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.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "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.9, <3.13"
 dynamic = [ "version", "readme",]
-dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<5", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2024.6.1,<2026", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<2", "packaging>=20.9,<25", "pandas>=2.1.4,<3", "pyarrow", "pydantic>=2.8.2, <3", "pyjwt>=2.0.0, <3", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2024.6.1,<2026", "scikit-learn<1.6", "scipy>=1.9,<2", "shap>=0.46.0,<1", "snowflake-connector-python[pandas]>=3.15.0,<4", "snowflake-snowpark-python>=1.17.0,<2,!=1.26.0", "snowflake.core>=1.0.2,<2", "sqlparse>=0.4,<1", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<3",]
+dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<5", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2024.6.1,<2026", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<3", "packaging>=20.9,<25", "pandas>=2.1.4,<3", "platformdirs<5", "pyarrow", "pydantic>=2.8.2, <3", "pyjwt>=2.0.0, <3", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2024.6.1,<2026", "scikit-learn<1.6", "scipy>=1.9,<2", "shap>=0.46.0,<1", "snowflake-connector-python[pandas]>=3.15.0,<4", "snowflake-snowpark-python>=1.17.0,<2,!=1.26.0", "snowflake.core>=1.0.2,<2", "sqlparse>=0.4,<1", "tqdm<5", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<3",]
 [[project.authors]]
 name = "Snowflake, Inc"
 email = "support@snowflake.com"
@@ -24,7 +24,7 @@ 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 = [ "altair>=5,<6", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5",]
+all = [ "altair>=5,<6", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5,!=4.51.3",]
 altair = [ "altair>=5,<6",]
 catboost = [ "catboost>=1.2.0, <2",]
 keras = [ "keras>=2.0.0,<4", "tensorflow>=2.17.0,<3", "torch>=2.0.1,<3",]
@@ -33,7 +33,7 @@ mlflow = [ "mlflow>=2.16.0, <3",]
 streamlit = [ "streamlit>=1.30.0,<2",]
 tensorflow = [ "tensorflow>=2.17.0,<3",]
 torch = [ "torch>=2.0.1,<3", "torchdata>=0.4,<1",]
-transformers = [ "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "transformers>=4.39.3,<5",]
+transformers = [ "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "transformers>=4.39.3,<5,!=4.51.3",]
 
 [tool.setuptools.package-data]
 "*" = [ "*",]

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.2}/snowflake/ml/_internal/env_utils.py

@@ -337,13 +337,54 @@ def get_package_spec_with_supported_ops_only(req: requirements.Requirement) -> r
     Returns:
         A requirements.Requirement object with supported ops only
     """
+
+    if req.name == "numpy":
+        import numpy as np
+
+        package_specifiers = get_numpy_specifiers(req, version.Version(np.__version__).major)
+    else:
+        package_specifiers = [spec for spec in req.specifier if spec.operator in _SUPPORTED_PACKAGE_SPEC_OPS]
+
     new_req = copy.deepcopy(req)
-    new_req.specifier = specifiers.SpecifierSet(
-        specifiers=",".join([str(spec) for spec in req.specifier if spec.operator in _SUPPORTED_PACKAGE_SPEC_OPS])
-    )
+    new_req.specifier = specifiers.SpecifierSet(specifiers=",".join([str(spec) for spec in package_specifiers]))
     return new_req
 
 
+def get_numpy_specifiers(
+    req: requirements.Requirement,
+    client_numpy_major_version: int,
+) -> list[specifiers.Specifier]:
+    """Get the package spec with supported ops only including ==, >=, <=, > and < based on the client numpy
+    major version.
+
+    Args:
+        req: A requirements.Requirement object showing the requirement.
+        client_numpy_major_version: The major version of numpy to be used.
+
+    Returns:
+        A list of specifiers with supported ops only
+    """
+    req_specifiers = []
+    for org_spec in req.specifier:
+        # check specifier that provides upper bound
+        if org_spec.operator in ["<", "<="]:
+            client_version = version.Version(str(client_numpy_major_version))
+            org_spec_version = version.Version(org_spec.version)
+            # check if the client's numpy major version is less than the specifier's upper bound
+            # if so, pin to max possible client major version
+            if client_version.major < org_spec_version.major:
+                modified_spec = specifiers.Specifier(f"<{client_version.major + 1}")
+                req_specifiers.append(modified_spec)
+            else:
+                # use the original specifier
+                req_specifiers.append(org_spec)
+        else:
+            # use the original specifier
+            req_specifiers.append(org_spec)
+
+    return req_specifiers
+
+
 def _relax_specifier_set(
     specifier_set: specifiers.SpecifierSet, strategy: relax_version_strategy.RelaxVersionStrategy
 ) -> specifiers.SpecifierSet:

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.2}/snowflake/ml/_internal/platform_capabilities.py

@@ -3,7 +3,9 @@ from contextlib import contextmanager
 from typing import Any, Optional
 
 from absl import logging
+from packaging import version
 
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml._internal.utils import query_result_checker
 from snowflake.snowpark import (
@@ -12,7 +14,7 @@ from snowflake.snowpark import (
 )
 
 LIVE_COMMIT_PARAMETER = "ENABLE_LIVE_VERSION_IN_SDK"
-INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC"
+INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC_FROM_CLIENT_VERSION"
 
 
 class PlatformCapabilities:
@@ -67,7 +69,7 @@ class PlatformCapabilities:
             cls.clear_mock_features()
 
     def is_inlined_deployment_spec_enabled(self) -> bool:
-        return self._get_bool_feature(INLINE_DEPLOYMENT_SPEC_PARAMETER, False)
+        return self._is_version_feature_enabled(INLINE_DEPLOYMENT_SPEC_PARAMETER)
 
     def is_live_commit_enabled(self) -> bool:
         return self._get_bool_feature(LIVE_COMMIT_PARAMETER, False)
@@ -126,3 +128,51 @@ class PlatformCapabilities:
             else:
                 raise ValueError(f"Invalid boolean string: {value} for feature {feature_name}")
         raise ValueError(f"Invalid boolean feature value: {value} for feature {feature_name}")
+
+    def _get_version_feature(self, feature_name: str) -> version.Version:
+        """Get a version feature value, returning a large version number on failure or missing feature.
+
+        Args:
+            feature_name: The name of the feature to retrieve.
+
+        Returns:
+            version.Version: The parsed version, or a large version number (999.999.999) if parsing fails
+            or the feature is missing.
+        """
+        # Large version number to use as fallback
+        large_version = version.Version("999.999.999")
+
+        value = self.features.get(feature_name)
+        if value is None:
+            logging.debug(f"Feature {feature_name} not found, returning large version number")
+            return large_version
+
+        try:
+            # Convert to string if it's not already
+            version_str = str(value)
+            return version.Version(version_str)
+        except (version.InvalidVersion, ValueError, TypeError) as e:
+            logging.debug(
+                f"Failed to parse version from feature {feature_name} with value '{value}': {e}. "
+                f"Returning large version number"
+            )
+            return large_version
+
+    def _is_version_feature_enabled(self, feature_name: str) -> bool:
+        """Check if the current package version is greater than or equal to the version feature.
+
+        Args:
+            feature_name: The name of the version feature to compare against.
+
+        Returns:
+            bool: True if current package version >= feature version, False otherwise.
+        """
+        current_version = version.Version(snowml_version.VERSION)
+        feature_version = self._get_version_feature(feature_name)
+
+        result = current_version >= feature_version
+        logging.debug(
+            f"Version comparison for feature {feature_name}: "
+            f"current={current_version}, feature={feature_version}, enabled={result}"
+        )
+        return result

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.2}/snowflake/ml/_internal/type_utils.py

@@ -66,4 +66,4 @@ class LazyType(Generic[T]):
             return False
 
 
-LiteralNDArrayType = Union[npt.NDArray[np.int_], npt.NDArray[np.float_], npt.NDArray[np.str_], npt.NDArray[np.bool_]]
+LiteralNDArrayType = Union[npt.NDArray[np.int_], npt.NDArray[np.float64], npt.NDArray[np.str_], npt.NDArray[np.bool_]]

snowflake_ml_python-1.9.2/snowflake/ml/_internal/utils/mixins.py

@@ -0,0 +1,73 @@
+from typing import Any, Optional
+
+from snowflake.ml._internal.utils import identifier
+from snowflake.snowpark import session as snowpark_session
+
+_SESSION_KEY = "_session"
+_SESSION_ACCOUNT_KEY = "session$account"
+_SESSION_ROLE_KEY = "session$role"
+_SESSION_DATABASE_KEY = "session$database"
+_SESSION_SCHEMA_KEY = "session$schema"
+
+
+def _identifiers_match(saved: Optional[str], current: Optional[str]) -> bool:
+    saved_resolved = identifier.resolve_identifier(saved) if saved is not None else saved
+    current_resolved = identifier.resolve_identifier(current) if current is not None else current
+    return saved_resolved == current_resolved
+
+
+class SerializableSessionMixin:
+    """Mixin that provides pickling capabilities for objects with Snowpark sessions."""
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Customize pickling to exclude non-serializable session and related components."""
+        parent_state = (
+            super().__getstate__()  # type: ignore[misc] # object.__getstate__ appears in 3.11
+            if hasattr(super(), "__getstate__")
+            else self.__dict__
+        )
+        state = dict(parent_state)  # Create a copy so we can safely modify the state
+
+        # Save session metadata for validation during unpickling
+        session = state.pop(_SESSION_KEY, None)
+        if session is not None:
+            state[_SESSION_ACCOUNT_KEY] = session.get_current_account()
+            state[_SESSION_ROLE_KEY] = session.get_current_role()
+            state[_SESSION_DATABASE_KEY] = session.get_current_database()
+            state[_SESSION_SCHEMA_KEY] = session.get_current_schema()
+
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore session from context during unpickling."""
+        saved_account = state.pop(_SESSION_ACCOUNT_KEY, None)
+        saved_role = state.pop(_SESSION_ROLE_KEY, None)
+        saved_database = state.pop(_SESSION_DATABASE_KEY, None)
+        saved_schema = state.pop(_SESSION_SCHEMA_KEY, None)
+
+        if hasattr(super(), "__setstate__"):
+            super().__setstate__(state)  # type: ignore[misc]
+        else:
+            self.__dict__.update(state)
+
+        if saved_account is not None:
+            active_sessions = snowpark_session._get_active_sessions()
+            if len(active_sessions) == 0:
+                raise RuntimeError("No active Snowpark session available. Please create a session.")
+
+            # Best effort match: Find the session with the most matching identifiers
+            setattr(
+                self,
+                _SESSION_KEY,
+                max(
+                    active_sessions,
+                    key=lambda s: sum(
+                        (
+                            _identifiers_match(saved_account, s.get_current_account()),
+                            _identifiers_match(saved_role, s.get_current_role()),
+                            _identifiers_match(saved_database, s.get_current_database()),
+                            _identifiers_match(saved_schema, s.get_current_schema()),
+                        )
+                    ),
+                ),
+            )