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

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

@@ -1,5 +1,54 @@
 # Release History
 
+## 1.9.1
+
+### 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 +56,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 +178,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.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.9.0
+Version: 1.9.1
 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,7 +240,7 @@ 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: pyarrow
@@ -257,6 +257,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
@@ -408,6 +409,55 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
+## 1.9.1
+
+### 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 +465,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 +587,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.1}/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", "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"

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.1}/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.1}/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.1}/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.1/snowflake/ml/_internal/utils/mixins.py

@@ -0,0 +1,71 @@
+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."""
+        if hasattr(super(), "__getstate__"):
+            state: dict[str, Any] = super().__getstate__()  # type: ignore[misc]
+        else:
+            state = self.__dict__.copy()
+
+        # 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()),
+                        )
+                    ),
+                ),
+            )

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.1}/snowflake/ml/_internal/utils/service_logger.py

@@ -10,6 +10,10 @@ class LogColor(enum.Enum):
     YELLOW = "\x1b[33;20m"
     BLUE = "\x1b[34;20m"
     GREEN = "\x1b[32;20m"
+    ORANGE = "\x1b[38;5;214m"
+    BOLD_ORANGE = "\x1b[38;5;214;1m"
+    PURPLE = "\x1b[35;20m"
+    BOLD_PURPLE = "\x1b[35;1m"
 
 
 class CustomFormatter(logging.Formatter):
@@ -55,9 +59,7 @@ class CustomFormatter(logging.Formatter):
 
 def get_logger(logger_name: str, info_color: LogColor) -> logging.Logger:
     logger = logging.getLogger(logger_name)
-    logger.setLevel(logging.INFO)
     handler = logging.StreamHandler(sys.stdout)
-    handler.setLevel(logging.INFO)
     handler.setFormatter(CustomFormatter(info_color))
     logger.addHandler(handler)
     return logger

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.1}/snowflake/ml/data/_internal/arrow_ingestor.py

@@ -2,7 +2,7 @@ import collections
 import logging
 import os
 import time
-from typing import Any, Deque, Iterator, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, Deque, Iterator, Optional, Sequence, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -10,6 +10,9 @@ import pandas as pd
 import pyarrow as pa
 import pyarrow.dataset as pds
 
+if TYPE_CHECKING:
+    import ray
+
 from snowflake import snowpark
 from snowflake.ml.data import data_ingestor, data_source, ingestor_utils
 
@@ -70,6 +73,13 @@ class ArrowIngestor(data_ingestor.DataIngestor):
     def from_sources(cls, session: snowpark.Session, sources: Sequence[data_source.DataSource]) -> "ArrowIngestor":
         return cls(session, sources)
 
+    @classmethod
+    def from_ray_dataset(
+        cls,
+        ray_ds: "ray.data.Dataset",
+    ) -> "ArrowIngestor":
+        raise NotImplementedError
+
     @property
     def data_sources(self) -> list[data_source.DataSource]:
         return self._data_sources

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.1}/snowflake/ml/data/data_connector.py

@@ -6,6 +6,7 @@ from typing_extensions import deprecated
 
 from snowflake import snowpark
 from snowflake.ml._internal import env, telemetry
+from snowflake.ml._internal.utils import mixins
 from snowflake.ml.data import data_ingestor, data_source
 from snowflake.ml.data._internal.arrow_ingestor import ArrowIngestor
 from snowflake.snowpark import context as sf_context
@@ -21,11 +22,13 @@ if TYPE_CHECKING:
     from snowflake.ml import dataset
 
 _PROJECT = "DataConnector"
+_INGESTOR_KEY = "_ingestor"
+_INGESTOR_SOURCES_KEY = "ingestor$sources"
 
 DataConnectorType = TypeVar("DataConnectorType", bound="DataConnector")
 
 
-class DataConnector:
+class DataConnector(mixins.SerializableSessionMixin):
     """Snowflake data reader which provides application integration connectors"""
 
     DEFAULT_INGESTOR_CLASS: type[data_ingestor.DataIngestor] = ArrowIngestor
@@ -33,8 +36,11 @@ class DataConnector:
     def __init__(
         self,
         ingestor: data_ingestor.DataIngestor,
+        *,
+        session: Optional[snowpark.Session] = None,
         **kwargs: Any,
     ) -> None:
+        self._session = session
         self._ingestor = ingestor
         self._kwargs = kwargs
 
@@ -75,6 +81,17 @@ class DataConnector:
         )
         return cls.from_sources(ds._session, [source], ingestor_class=ingestor_class, **kwargs)
 
+    @classmethod
+    def from_ray_dataset(
+        cls: type[DataConnectorType],
+        ray_ds: "ray.data.Dataset",
+        ingestor_class: Optional[type[data_ingestor.DataIngestor]] = None,
+        **kwargs: Any,
+    ) -> DataConnectorType:
+        ingestor_class = ingestor_class or cls.DEFAULT_INGESTOR_CLASS
+        ray_ingestor = ingestor_class.from_ray_dataset(ray_ds=ray_ds)
+        return cls(ray_ingestor, **kwargs)
+
     @classmethod
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
@@ -90,7 +107,31 @@ class DataConnector:
     ) -> DataConnectorType:
         ingestor_class = ingestor_class or cls.DEFAULT_INGESTOR_CLASS
         ingestor = ingestor_class.from_sources(session, sources)
-        return cls(ingestor, **kwargs)
+        return cls(ingestor, **kwargs, session=session)
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Customize pickling to exclude non-serializable session and related components."""
+        if hasattr(super(), "__getstate__"):
+            state = super().__getstate__()
+        else:
+            state = self.__dict__.copy()
+
+        ingestor = state.pop(_INGESTOR_KEY)
+        state[_INGESTOR_SOURCES_KEY] = ingestor.data_sources
+
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore session from context during unpickling."""
+        data_sources = state.pop(_INGESTOR_SOURCES_KEY)
+
+        if hasattr(super(), "__setstate__"):
+            super().__setstate__(state)
+        else:
+            self.__dict__.update(state)
+
+        assert self._session is not None
+        self._ingestor = self.DEFAULT_INGESTOR_CLASS.from_sources(self._session, data_sources)
 
     @property
     def data_sources(self) -> list[data_source.DataSource]:

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.1}/snowflake/ml/data/data_ingestor.py

@@ -7,6 +7,7 @@ from snowflake.ml.data import data_source
 
 if TYPE_CHECKING:
     import pandas as pd
+    import ray
 
 
 DataIngestorType = TypeVar("DataIngestorType", bound="DataIngestor")
@@ -19,6 +20,13 @@ class DataIngestor(Protocol):
     ) -> DataIngestorType:
         raise NotImplementedError
 
+    @classmethod
+    def from_ray_dataset(
+        cls: type[DataIngestorType],
+        ray_ds: "ray.data.Dataset",
+    ) -> DataIngestorType:
+        raise NotImplementedError
+
     @property
     def data_sources(self) -> list[data_source.DataSource]:
         raise NotImplementedError

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.1}/snowflake/ml/data/torch_utils.py

@@ -95,6 +95,6 @@ def _preprocess_array(
         array_list = arr.tolist()
         # If this is an array of arrays, convert the dtype to match the underlying array.
         # Otherwise, if this is a numpy array of strings, convert the array to a list.
-        arr = np.array(array_list, dtype=arr.flat[0].dtype) if isinstance(arr.flat[0], np.ndarray) else array_list
+        arr = np.array(array_list, dtype=arr.item(0).dtype) if isinstance(arr.item(0), np.ndarray) else array_list
 
     return arr

{snowflake_ml_python-1.9.0 → snowflake_ml_python-1.9.1}/snowflake/ml/dataset/dataset.py

@@ -14,6 +14,7 @@ from snowflake.ml._internal.exceptions import (
 from snowflake.ml._internal.utils import (
     formatting,
     identifier,
+    mixins,
     query_result_checker,
     snowpark_dataframe_utils,
 )
@@ -27,7 +28,7 @@ _METADATA_MAX_QUERY_LENGTH = 10000
 _DATASET_VERSION_NAME_COL = "version"
 
 
-class DatasetVersion:
+class DatasetVersion(mixins.SerializableSessionMixin):
     """Represents a version of a Snowflake Dataset"""
 
     @telemetry.send_api_usage_telemetry(project=_PROJECT)
@@ -176,7 +177,7 @@ class Dataset(lineage_node.LineageNode):
                 original_exception=RuntimeError("No Dataset version selected."),
             )
         if self._reader is None:
-            self._reader = dataset_reader.DatasetReader.from_dataset(self, snowpark_session=self._session)
+            self._reader = dataset_reader.DatasetReader.from_dataset(self)
         return self._reader
 
     @staticmethod