snowflake-ml-python 1.15.0__py3-none-any.whl → 1.17.0__py3-none-any.whl

snowflake/ml/_internal/human_readable_id/adjectives.txt

@@ -1,3 +1,4 @@
+aerial
 afraid
 ancient
 angry
@@ -26,7 +27,6 @@ dull
 empty
 evil
 fast
-fat
 fluffy
 foolish
 fresh
@@ -57,10 +57,10 @@ lovely
 lucky
 massive
 mean
+metallic
 mighty
 modern
 moody
-nasty
 neat
 nervous
 new
@@ -85,7 +85,6 @@ rotten
 rude
 selfish
 serious
-shaggy
 sharp
 short
 shy
@@ -96,14 +95,15 @@ slippery
 smart
 smooth
 soft
+solid
 sour
 spicy
 splendid
 spotty
+squishy
 stale
 strange
 strong
-stupid
 sweet
 swift
 tall
@@ -116,7 +116,6 @@ tidy
 tiny
 tough
 tricky
-ugly
 warm
 weak
 wet
@@ -124,5 +123,6 @@ wicked
 wise
 witty
 wonderful
+wooden
 yellow
 young

snowflake/ml/_internal/human_readable_id/animals.txt

@@ -1,10 +1,9 @@
 anaconda
 ant
-ape
-baboon
 badger
 bat
 bear
+beetle
 bird
 bobcat
 bulldog
@@ -73,7 +72,6 @@ lobster
 mayfly
 mamba
 mole
-monkey
 moose
 moth
 mouse
@@ -114,6 +112,7 @@ swan
 termite
 tiger
 treefrog
+tuna
 turkey
 turtle
 vampirebat
@@ -126,3 +125,4 @@ worm
 yak
 yeti
 zebra
+zebrafish

snowflake/ml/_internal/platform_capabilities.py

@@ -17,6 +17,7 @@ logger = logging.getLogger(__name__)
 
 LIVE_COMMIT_PARAMETER = "ENABLE_LIVE_VERSION_IN_SDK"
 INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC_FROM_CLIENT_VERSION"
+SET_MODULE_FUNCTIONS_VOLATILITY_FROM_MANIFEST = "SET_MODULE_FUNCTIONS_VOLATILITY_FROM_MANIFEST"
 
 
 class PlatformCapabilities:
@@ -73,6 +74,9 @@ class PlatformCapabilities:
     def is_inlined_deployment_spec_enabled(self) -> bool:
         return self._is_version_feature_enabled(INLINE_DEPLOYMENT_SPEC_PARAMETER)
 
+    def is_set_module_functions_volatility_from_manifest(self) -> bool:
+        return self._get_bool_feature(SET_MODULE_FUNCTIONS_VOLATILITY_FROM_MANIFEST, False)
+
     def is_live_commit_enabled(self) -> bool:
         return self._get_bool_feature(LIVE_COMMIT_PARAMETER, False)
 

snowflake/ml/_internal/utils/mixins.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from typing import Any, Optional
 
 from snowflake.ml._internal.utils import identifier
@@ -16,6 +17,14 @@ def _identifiers_match(saved: Optional[str], current: Optional[str]) -> bool:
     return saved_resolved == current_resolved
 
 
+@dataclass(frozen=True)
+class _SessionState:
+    account: Optional[str]
+    role: Optional[str]
+    database: Optional[str]
+    schema: Optional[str]
+
+
 class SerializableSessionMixin:
     """Mixin that provides pickling capabilities for objects with Snowpark sessions."""
 
@@ -40,17 +49,23 @@ class SerializableSessionMixin:
 
     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)
+        session_state = _SessionState(
+            account=state.pop(_SESSION_ACCOUNT_KEY, None),
+            role=state.pop(_SESSION_ROLE_KEY, None),
+            database=state.pop(_SESSION_DATABASE_KEY, None),
+            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:
+        self._set_session(session_state)
+
+    def _set_session(self, session_state: _SessionState) -> None:
+
+        if session_state.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.")
@@ -63,10 +78,10 @@ class SerializableSessionMixin:
                     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()),
+                            _identifiers_match(session_state.account, s.get_current_account()),
+                            _identifiers_match(session_state.role, s.get_current_role()),
+                            _identifiers_match(session_state.database, s.get_current_database()),
+                            _identifiers_match(session_state.schema, s.get_current_schema()),
                         )
                     ),
                 ),

snowflake/ml/experiment/experiment_tracking.py

@@ -1,10 +1,10 @@
 import functools
 import json
 import sys
-from typing import Any, Optional, Union
+from typing import Any, Callable, Concatenate, Optional, ParamSpec, TypeVar, Union
 from urllib.parse import quote
 
-import snowflake.snowpark._internal.utils as snowpark_utils
+from snowflake import snowpark
 from snowflake.ml import model as ml_model, registry
 from snowflake.ml._internal.human_readable_id import hrid_generator
 from snowflake.ml._internal.utils import mixins, sql_identifier
@@ -18,20 +18,40 @@ from snowflake.ml.experiment._client import (
 )
 from snowflake.ml.model import type_hints
 from snowflake.ml.utils import sql_client as sql_client_utils
-from snowflake.snowpark import session
 
 DEFAULT_EXPERIMENT_NAME = sql_identifier.SqlIdentifier("DEFAULT")
 
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def _restore_session(
+    func: Callable[Concatenate["ExperimentTracking", P], T],
+) -> Callable[Concatenate["ExperimentTracking", P], T]:
+    @functools.wraps(func)
+    def wrapper(self: "ExperimentTracking", /, *args: P.args, **kwargs: P.kwargs) -> T:
+        if self._session is None:
+            if self._session_state is None:
+                raise RuntimeError(
+                    f"Session is not set before calling {func.__name__}, and there is no session state to restore from"
+                )
+            self._set_session(self._session_state)
+            if self._session is None:
+                raise RuntimeError(f"Failed to restore session before calling {func.__name__}")
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
 
 class ExperimentTracking(mixins.SerializableSessionMixin):
     """
     Class to manage experiments in Snowflake.
     """
 
-    @snowpark_utils.private_preview(version="1.9.1")
+    @snowpark._internal.utils.private_preview(version="1.9.1")
     def __init__(
         self,
-        session: session.Session,
+        session: snowpark.Session,
         *,
         database_name: Optional[str] = None,
         schema_name: Optional[str] = None,
@@ -73,7 +93,10 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             database_name=self._database_name,
             schema_name=self._schema_name,
         )
-        self._session = session
+        self._session: Optional[snowpark.Session] = session
+        # Used to store information about the session if the session could not be restored during unpickling
+        # _session_state is None if and only if _session is not None
+        self._session_state: Optional[mixins._SessionState] = None
 
         # The experiment in context
         self._experiment: Optional[entities.Experiment] = None
@@ -87,20 +110,29 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         state["_registry"] = None
         return state
 
-    def __setstate__(self, state: dict[str, Any]) -> None:
-        super().__setstate__(state)
-        # Restore unpicklable attributes
-        self._sql_client = sql_client.ExperimentTrackingSQLClient(
-            session=self._session,
-            database_name=self._database_name,
-            schema_name=self._schema_name,
-        )
-        self._registry = registry.Registry(
-            session=self._session,
-            database_name=self._database_name,
-            schema_name=self._schema_name,
-        )
+    def _set_session(self, session_state: mixins._SessionState) -> None:
+        try:
+            super()._set_session(session_state)
+            assert self._session is not None
+        except (snowpark.exceptions.SnowparkSessionException, AssertionError):
+            # If session was not set, store the session state
+            self._session = None
+            self._session_state = session_state
+        else:
+            # If session was set, clear the session state, and reinitialize the SQL client and registry
+            self._session_state = None
+            self._sql_client = sql_client.ExperimentTrackingSQLClient(
+                session=self._session,
+                database_name=self._database_name,
+                schema_name=self._schema_name,
+            )
+            self._registry = registry.Registry(
+                session=self._session,
+                database_name=self._database_name,
+                schema_name=self._schema_name,
+            )
 
+    @_restore_session
     def set_experiment(
         self,
         experiment_name: str,
@@ -125,6 +157,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         self._run = None
         return self._experiment
 
+    @_restore_session
     def delete_experiment(
         self,
         experiment_name: str,
@@ -141,8 +174,10 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             self._run = None
 
     @functools.wraps(registry.Registry.log_model)
+    @_restore_session
     def log_model(
         self,
+        /,  # self needs to be a positional argument to stop mypy from complaining
         model: Union[type_hints.SupportedModelType, ml_model.ModelVersion],
         *,
         model_name: str,
@@ -152,6 +187,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         with experiment_info.ExperimentInfoPatcher(experiment_info=run._get_experiment_info()):
             return self._registry.log_model(model, model_name=model_name, **kwargs)
 
+    @_restore_session
     def start_run(
         self,
         run_name: Optional[str] = None,
@@ -181,6 +217,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         self._run = entities.Run(experiment_tracking=self, experiment_name=experiment.name, run_name=run_name)
         return self._run
 
+    @_restore_session
     def end_run(self, run_name: Optional[str] = None) -> None:
         """
         End the current run if no run name is provided. Otherwise, the specified run is ended.
@@ -210,6 +247,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             self._run = None
         self._print_urls(experiment_name=experiment_name, run_name=run_name)
 
+    @_restore_session
     def delete_run(
         self,
         run_name: str,
@@ -248,6 +286,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         """
         self.log_metrics(metrics={key: value}, step=step)
 
+    @_restore_session
     def log_metrics(
         self,
         metrics: dict[str, float],
@@ -284,6 +323,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         """
         self.log_params({key: value})
 
+    @_restore_session
     def log_params(
         self,
         params: dict[str, Any],
@@ -305,6 +345,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             params=json.dumps([param.to_dict() for param in params_list]),
         )
 
+    @_restore_session
     def log_artifact(
         self,
         local_path: str,
@@ -328,6 +369,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
                 file_path=file_path,
             )
 
+    @_restore_session
     def list_artifacts(
         self,
         run_name: str,
@@ -356,6 +398,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             artifact_path=artifact_path or "",
         )
 
+    @_restore_session
     def download_artifacts(
         self,
         run_name: str,
@@ -397,6 +440,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             return self._run
         return self.start_run()
 
+    @_restore_session
     def _generate_run_name(self, experiment: entities.Experiment) -> sql_identifier.SqlIdentifier:
         generator = hrid_generator.HRID16()
         existing_runs = self._sql_client.show_runs_in_experiment(experiment_name=experiment.name)

snowflake/ml/jobs/__init__.py

@@ -1,3 +1,4 @@
+from snowflake.ml.jobs._interop.exception_utils import install_exception_display_hooks
 from snowflake.ml.jobs._utils.types import JOB_STATUS
 from snowflake.ml.jobs.decorators import remote
 from snowflake.ml.jobs.job import MLJob
@@ -10,6 +11,9 @@ from snowflake.ml.jobs.manager import (
     submit_from_stage,
 )
 
+# Initialize exception display hooks for remote job error handling
+install_exception_display_hooks()
+
 __all__ = [
     "remote",
     "submit_file",

snowflake/ml/jobs/_interop/data_utils.py

@@ -0,0 +1,124 @@
+import io
+import json
+from typing import Any, Literal, Optional, Protocol, Union, cast, overload
+
+from snowflake import snowpark
+from snowflake.ml.jobs._interop import dto_schema
+
+
+class StageFileWriter(io.IOBase):
+    """
+    A context manager IOBase implementation that proxies writes to an internal BytesIO
+    and uploads to Snowflake stage on close.
+    """
+
+    def __init__(self, session: snowpark.Session, path: str) -> None:
+        self._session = session
+        self._path = path
+        self._buffer = io.BytesIO()
+        self._closed = False
+        self._exception_occurred = False
+
+    def write(self, data: Union[bytes, bytearray]) -> int:
+        """Write data to the internal buffer."""
+        if self._closed:
+            raise ValueError("I/O operation on closed file")
+        return self._buffer.write(data)
+
+    def close(self, write_contents: bool = True) -> None:
+        """Close the file and upload the buffer contents to the stage."""
+        if not self._closed:
+            # Only upload if buffer has content and no exception occurred
+            if write_contents and self._buffer.tell() > 0:
+                self._buffer.seek(0)
+                self._session.file.put_stream(self._buffer, self._path)
+            self._buffer.close()
+            self._closed = True
+
+    def __enter__(self) -> "StageFileWriter":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        exception_occurred = exc_type is not None
+        self.close(write_contents=not exception_occurred)
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    def writable(self) -> bool:
+        return not self._closed
+
+    def readable(self) -> bool:
+        return False
+
+    def seekable(self) -> bool:
+        return not self._closed
+
+
+def _is_stage_path(path: str) -> bool:
+    return path.startswith("@") or path.startswith("snow://")
+
+
+def open_stream(path: str, mode: str = "rb", session: Optional[snowpark.Session] = None) -> io.IOBase:
+    if _is_stage_path(path):
+        if session is None:
+            raise ValueError("Session is required when opening a stage path")
+        if "r" in mode:
+            stream: io.IOBase = session.file.get_stream(path)  # type: ignore[assignment]
+            return stream
+        elif "w" in mode:
+            return StageFileWriter(session, path)
+        else:
+            raise ValueError(f"Unsupported mode '{mode}' for stage path")
+    else:
+        result: io.IOBase = open(path, mode)  # type: ignore[assignment]
+        return result
+
+
+class DtoCodec(Protocol):
+    @overload
+    @staticmethod
+    def decode(stream: io.IOBase, as_dict: Literal[True]) -> dict[str, Any]:
+        ...
+
+    @overload
+    @staticmethod
+    def decode(stream: io.IOBase, as_dict: Literal[False] = False) -> dto_schema.ResultDTO:
+        ...
+
+    @staticmethod
+    def decode(stream: io.IOBase, as_dict: bool = False) -> Union[dto_schema.ResultDTO, dict[str, Any]]:
+        pass
+
+    @staticmethod
+    def encode(dto: dto_schema.ResultDTO) -> bytes:
+        pass
+
+
+class JsonDtoCodec(DtoCodec):
+    @overload
+    @staticmethod
+    def decode(stream: io.IOBase, as_dict: Literal[True]) -> dict[str, Any]:
+        ...
+
+    @overload
+    @staticmethod
+    def decode(stream: io.IOBase, as_dict: Literal[False] = False) -> dto_schema.ResultDTO:
+        ...
+
+    @staticmethod
+    def decode(stream: io.IOBase, as_dict: bool = False) -> Union[dto_schema.ResultDTO, dict[str, Any]]:
+        data = cast(dict[str, Any], json.load(stream))
+        if as_dict:
+            return data
+        return dto_schema.ResultDTO.model_validate(data)
+
+    @staticmethod
+    def encode(dto: dto_schema.ResultDTO) -> bytes:
+        # Temporarily extract the value to avoid accidentally applying model_dump() on it
+        result_value = dto.value
+        dto.value = None  # Clear value to avoid serializing it in the model_dump
+        result_dict = dto.model_dump()
+        result_dict["value"] = result_value  # Put back the value
+        return json.dumps(result_dict).encode("utf-8")

snowflake/ml/jobs/_interop/dto_schema.py

@@ -0,0 +1,95 @@
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, model_validator
+from typing_extensions import NotRequired, TypedDict
+
+
+class BinaryManifest(TypedDict):
+    """
+    Binary data manifest schema.
+    Contains one of: path, bytes, or base64 for the serialized data.
+    """
+
+    path: NotRequired[str]  # Path to file
+    bytes: NotRequired[bytes]  # In-line byte string (not supported with JSON codec)
+    base64: NotRequired[str]  # Base64 encoded string
+
+
+class ParquetManifest(TypedDict):
+    """Protocol manifest schema for parquet files."""
+
+    paths: list[str]  # File paths
+
+
+# Union type for all manifest types, including catch-all dict[str, Any] for backward compatibility
+PayloadManifest = Union[BinaryManifest, ParquetManifest, dict[str, Any]]
+
+
+class ProtocolInfo(BaseModel):
+    """
+    The protocol used to serialize the result and the manifest of the result.
+    """
+
+    name: str
+    version: Optional[str] = None
+    metadata: Optional[dict[str, str]] = None
+    manifest: Optional[PayloadManifest] = None
+
+    def __str__(self) -> str:
+        result = self.name
+        if self.version:
+            result += f"-{self.version}"
+        return result
+
+    def with_manifest(self, manifest: PayloadManifest) -> "ProtocolInfo":
+        """
+        Return a new ProtocolInfo object with the manifest.
+        """
+        return ProtocolInfo(
+            name=self.name,
+            version=self.version,
+            metadata=self.metadata,
+            manifest=manifest,
+        )
+
+
+class ResultMetadata(BaseModel):
+    """
+    The metadata of a result.
+    """
+
+    type: str
+    repr: str
+
+
+class ExceptionMetadata(ResultMetadata):
+    message: str
+    traceback: str
+
+
+class ResultDTO(BaseModel):
+    """
+    A JSON representation of an execution result.
+
+    Args:
+        success: Whether the execution was successful.
+        value: The value of the execution or the exception if the execution failed.
+        protocol: The protocol used to serialize the result.
+        metadata: The metadata of the result.
+    """
+
+    success: bool
+    value: Optional[Any] = None
+    protocol: Optional[ProtocolInfo] = None
+    metadata: Optional[Union[ResultMetadata, ExceptionMetadata]] = None
+    serialize_error: Optional[str] = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_fields(cls, data: Any) -> Any:
+        """Ensure at least one of value, protocol, or metadata keys is specified."""
+        if isinstance(data, dict):
+            required_fields = {"value", "protocol", "metadata"}
+            if not any(field in data for field in required_fields):
+                raise ValueError("At least one of 'value', 'protocol', or 'metadata' must be specified")
+        return data