snowflake-ml-python 1.18.0__py3-none-any.whl → 1.20.0__py3-none-any.whl

snowflake/ml/_internal/env_utils.py

@@ -16,6 +16,7 @@ from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env as snowml_env, relax_version_strategy
 from snowflake.ml._internal.utils import query_result_checker
 from snowflake.snowpark import context, exceptions, session
+from snowflake.snowpark._internal import utils as snowpark_utils
 
 
 class CONDA_OS(Enum):
@@ -38,6 +39,21 @@ SNOWPARK_ML_PKG_NAME = "snowflake-ml-python"
 SNOWFLAKE_CONDA_CHANNEL_URL = "https://repo.anaconda.com/pkgs/snowflake"
 
 
+def get_execution_context() -> str:
+    """Detect execution context: EXTERNAL, SPCS, or SPROC.
+
+    Returns:
+        str: The execution context - "SPROC" if running in a stored procedure,
+             "SPCS" if running in SPCS ML runtime, "EXTERNAL" otherwise.
+    """
+    if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
+        return "SPROC"
+    elif snowml_env.IN_ML_RUNTIME:
+        return "SPCS"
+    else:
+        return "EXTERNAL"
+
+
 def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
     """Validate the input pip requirement string according to PEP 508.
 

snowflake/ml/_internal/telemetry.py

@@ -16,7 +16,7 @@ from typing_extensions import ParamSpec
 from snowflake import connector
 from snowflake.connector import connect, telemetry as connector_telemetry, time_util
 from snowflake.ml import version as snowml_version
-from snowflake.ml._internal import env
+from snowflake.ml._internal import env, env_utils
 from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
@@ -37,6 +37,22 @@ _CONNECTION_TYPES = {
 _Args = ParamSpec("_Args")
 _ReturnValue = TypeVar("_ReturnValue")
 
+_conn: Optional[connector.SnowflakeConnection] = None
+
+
+def clear_cached_conn() -> None:
+    """Clear the cached Snowflake connection. Primarily for testing purposes."""
+    global _conn
+    if _conn is not None and _conn.is_valid():
+        _conn.close()
+    _conn = None
+
+
+def get_cached_conn() -> Optional[connector.SnowflakeConnection]:
+    """Get the cached Snowflake connection. Primarily for testing purposes."""
+    global _conn
+    return _conn
+
 
 def _get_login_token() -> Union[str, bytes]:
     with open("/snowflake/session/token") as f:
@@ -44,7 +60,11 @@ def _get_login_token() -> Union[str, bytes]:
 
 
 def _get_snowflake_connection() -> Optional[connector.SnowflakeConnection]:
-    conn = None
+    global _conn
+    if _conn is not None and _conn.is_valid():
+        return _conn
+
+    conn: Optional[connector.SnowflakeConnection] = None
     if os.getenv("SNOWFLAKE_HOST") is not None and os.getenv("SNOWFLAKE_ACCOUNT") is not None:
         try:
             conn = connect(
@@ -66,6 +86,13 @@ def _get_snowflake_connection() -> Optional[connector.SnowflakeConnection]:
             # Failed to get an active session. No connection available.
             pass
 
+    # cache the connection if it's a SnowflakeConnection. there is a behavior at runtime where it could be a
+    # StoredProcConnection perhaps incorrect type hinting somewhere
+    if isinstance(conn, connector.SnowflakeConnection):
+        # if _conn was expired, we need to copy telemetry data to new connection
+        if _conn is not None and conn is not None:
+            conn._telemetry._log_batch.extend(_conn._telemetry._log_batch)
+        _conn = conn
     return conn
 
 
@@ -113,6 +140,13 @@ class TelemetryField(enum.Enum):
     FUNC_CAT_USAGE = "usage"
 
 
+@enum.unique
+class CustomTagKey(enum.Enum):
+    """Keys for custom tags in telemetry."""
+
+    EXECUTION_CONTEXT = "execution_context"
+
+
 class _TelemetrySourceType(enum.Enum):
     # Automatically inferred telemetry/statement parameters
     AUTO_TELEMETRY = "SNOWML_AUTO_TELEMETRY"
@@ -441,6 +475,7 @@ def send_api_usage_telemetry(
     sfqids_extractor: Optional[Callable[..., list[str]]] = None,
     subproject_extractor: Optional[Callable[[Any], str]] = None,
     custom_tags: Optional[dict[str, Union[bool, int, str, float]]] = None,
+    log_execution_context: bool = True,
 ) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, _ReturnValue]]:
     """
     Decorator that sends API usage telemetry and adds function usage statement parameters to the dataframe returned by
@@ -455,6 +490,8 @@ def send_api_usage_telemetry(
         sfqids_extractor: Extract sfqids from `self`.
         subproject_extractor: Extract subproject at runtime from `self`.
         custom_tags: Custom tags.
+        log_execution_context: If True, automatically detect and log execution context
+            (EXTERNAL, SPCS, or SPROC) in custom_tags.
 
     Returns:
         Decorator that sends function usage telemetry for any call to the decorated function.
@@ -495,6 +532,11 @@ def send_api_usage_telemetry(
             if subproject_extractor is not None:
                 subproject_name = subproject_extractor(args[0])
 
+            # Add execution context if enabled
+            final_custom_tags = {**custom_tags} if custom_tags is not None else {}
+            if log_execution_context:
+                final_custom_tags[CustomTagKey.EXECUTION_CONTEXT.value] = env_utils.get_execution_context()
+
             statement_params = get_function_usage_statement_params(
                 project=project,
                 subproject=subproject_name,
@@ -502,7 +544,7 @@ def send_api_usage_telemetry(
                 function_name=_get_full_func_name(func),
                 function_parameters=params,
                 api_calls=api_calls,
-                custom_tags=custom_tags,
+                custom_tags=final_custom_tags,
             )
 
             def update_stmt_params_if_snowpark_df(obj: _ReturnValue, statement_params: dict[str, Any]) -> _ReturnValue:
@@ -538,7 +580,10 @@ def send_api_usage_telemetry(
             if conn_attr_name:
                 # raise AttributeError if conn attribute does not exist in `self`
                 conn = operator.attrgetter(conn_attr_name)(args[0])
-                if not isinstance(conn, _CONNECTION_TYPES.get(type(conn).__name__, connector.SnowflakeConnection)):
+                if not isinstance(
+                    conn,
+                    _CONNECTION_TYPES.get(type(conn).__name__, connector.SnowflakeConnection),
+                ):
                     raise TypeError(
                         f"Expected a conn object of type {' or '.join(_CONNECTION_TYPES.keys())} but got {type(conn)}"
                     )
@@ -560,7 +605,7 @@ def send_api_usage_telemetry(
                 func_params=params,
                 api_calls=api_calls,
                 sfqids=sfqids,
-                custom_tags=custom_tags,
+                custom_tags=final_custom_tags,
             )
             try:
                 return ctx.run(execute_func_with_statement_params)
@@ -571,7 +616,8 @@ def send_api_usage_telemetry(
                         raise
                     if isinstance(e, snowpark_exceptions.SnowparkClientException):
                         me = snowml_exceptions.SnowflakeMLException(
-                            error_code=error_codes.INTERNAL_SNOWPARK_ERROR, original_exception=e
+                            error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
+                            original_exception=e,
                         )
                     else:
                         me = snowml_exceptions.SnowflakeMLException(
@@ -627,7 +673,10 @@ def _get_full_func_name(func: Callable[..., Any]) -> str:
 
 
 def _get_func_params(
-    func: Callable[..., Any], func_params_to_log: Optional[Iterable[str]], args: Any, kwargs: Any
+    func: Callable[..., Any],
+    func_params_to_log: Optional[Iterable[str]],
+    args: Any,
+    kwargs: Any,
 ) -> dict[str, Any]:
     """
     Get function parameters.

snowflake/ml/experiment/_client/experiment_tracking_sql_client.py

@@ -128,13 +128,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
         run_name: sql_identifier.SqlIdentifier,
         artifact_path: str,
     ) -> list[artifact.ArtifactInfo]:
-        results = (
-            query_result_checker.SqlResultValidator(
-                self._session, f"LIST {self._build_snow_uri(experiment_name, run_name, artifact_path)}"
-            )
-            .has_dimensions(expected_cols=4)
-            .validate()
-        )
+        results = self._session.sql(f"LIST {self._build_snow_uri(experiment_name, run_name, artifact_path)}").collect()
         return [
             artifact.ArtifactInfo(
                 name=str(result.name).removeprefix(f"/versions/{run_name}/"),

snowflake/ml/experiment/_entities/run.py

@@ -1,4 +1,5 @@
 import types
+import warnings
 from typing import TYPE_CHECKING, Optional
 
 from snowflake.ml._internal.utils import sql_identifier
@@ -7,6 +8,8 @@ from snowflake.ml.experiment import _experiment_info as experiment_info
 if TYPE_CHECKING:
     from snowflake.ml.experiment import experiment_tracking
 
+METADATA_SIZE_WARNING_MESSAGE = "It is likely that no further metrics or parameters will be logged for this run."
+
 
 class Run:
     def __init__(
@@ -20,6 +23,9 @@ class Run:
         self.experiment_name = experiment_name
         self.name = run_name
 
+        # Whether we've already shown the user a warning about exceeding the run metadata size limit.
+        self._warned_about_metadata_size = False
+
         self._patcher = experiment_info.ExperimentInfoPatcher(
             experiment_info=self._get_experiment_info(),
         )
@@ -45,3 +51,12 @@ class Run:
             ),
             run_name=self.name.identifier(),
         )
+
+    def _warn_about_run_metadata_size(self, sql_error_msg: str) -> None:
+        if not self._warned_about_metadata_size:
+            warnings.warn(
+                f"{sql_error_msg}. {METADATA_SIZE_WARNING_MESSAGE}",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            self._warned_about_metadata_size = True

snowflake/ml/experiment/experiment_tracking.py

@@ -1,13 +1,13 @@
 import functools
 import json
 import sys
-from typing import Any, Callable, Concatenate, Optional, ParamSpec, TypeVar, Union
+from typing import Any, Optional, Union
 from urllib.parse import quote
 
 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
+from snowflake.ml._internal.utils import connection_params, sql_identifier
 from snowflake.ml.experiment import (
     _entities as entities,
     _experiment_info as experiment_info,
@@ -21,34 +21,12 @@ from snowflake.ml.utils import sql_client as sql_client_utils
 
 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 ExperimentTracking:
     """
     Class to manage experiments in Snowflake.
     """
 
-    @snowpark._internal.utils.private_preview(version="1.9.1")
     def __init__(
         self,
         session: snowpark.Session,
@@ -93,10 +71,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             database_name=self._database_name,
             schema_name=self._schema_name,
         )
-        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
+        self._session = session
 
         # The experiment in context
         self._experiment: Optional[entities.Experiment] = None
@@ -104,35 +79,40 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         self._run: Optional[entities.Run] = None
 
     def __getstate__(self) -> dict[str, Any]:
-        state = super().__getstate__()
+        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
+
         # Remove unpicklable attributes
+        state["_session"] = None
         state["_sql_client"] = None
         state["_registry"] = None
         return state
 
-    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
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        if hasattr(super(), "__setstate__"):
+            super().__setstate__(state)  # type: ignore[misc]
         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,
-            )
+            self.__dict__.update(state)
+
+        # Restore unpicklable attributes
+        options: dict[str, Any] = connection_params.SnowflakeLoginOptions()
+        options["client_session_keep_alive"] = True  # Needed for long-running training jobs
+        self._session = snowpark.Session.builder.configs(options).getOrCreate()
+        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,
@@ -157,7 +137,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         self._run = None
         return self._experiment
 
-    @_restore_session
     def delete_experiment(
         self,
         experiment_name: str,
@@ -174,10 +153,8 @@ 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,
@@ -187,7 +164,6 @@ 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,
@@ -229,7 +205,6 @@ 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.
@@ -259,7 +234,6 @@ 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,
@@ -298,7 +272,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         """
         self.log_metrics(metrics={key: value}, step=step)
 
-    @_restore_session
     def log_metrics(
         self,
         metrics: dict[str, float],
@@ -310,16 +283,26 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         Args:
             metrics: Dictionary containing metric keys and float values.
             step: The step of the metrics. Defaults to 0.
+
+        Raises:
+            snowpark.exceptions.SnowparkSQLException: If logging metrics fails due to Snowflake SQL errors,
+                except for run metadata size limit errors which will issue a warning instead of raising.
         """
         run = self._get_or_start_run()
         metrics_list = []
         for key, value in metrics.items():
             metrics_list.append(entities.Metric(key, value, step))
-        self._sql_client.modify_run_add_metrics(
-            experiment_name=run.experiment_name,
-            run_name=run.name,
-            metrics=json.dumps([metric.to_dict() for metric in metrics_list]),
-        )
+        try:
+            self._sql_client.modify_run_add_metrics(
+                experiment_name=run.experiment_name,
+                run_name=run.name,
+                metrics=json.dumps([metric.to_dict() for metric in metrics_list]),
+            )
+        except snowpark.exceptions.SnowparkSQLException as e:
+            if e.sql_error_code == 400003:  # EXPERIMENT_RUN_PROPERTY_SIZE_LIMIT_EXCEEDED
+                run._warn_about_run_metadata_size(e.message)
+            else:
+                raise
 
     def log_param(
         self,
@@ -335,7 +318,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         """
         self.log_params({key: value})
 
-    @_restore_session
     def log_params(
         self,
         params: dict[str, Any],
@@ -346,18 +328,27 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         Args:
             params: Dictionary containing parameter keys and values. Values can be of any type, but will be converted
                 to string.
+
+        Raises:
+            snowpark.exceptions.SnowparkSQLException: If logging parameters fails due to Snowflake SQL errors,
+                except for run metadata size limit errors which will issue a warning instead of raising.
         """
         run = self._get_or_start_run()
         params_list = []
         for key, value in params.items():
             params_list.append(entities.Param(key, str(value)))
-        self._sql_client.modify_run_add_params(
-            experiment_name=run.experiment_name,
-            run_name=run.name,
-            params=json.dumps([param.to_dict() for param in params_list]),
-        )
+        try:
+            self._sql_client.modify_run_add_params(
+                experiment_name=run.experiment_name,
+                run_name=run.name,
+                params=json.dumps([param.to_dict() for param in params_list]),
+            )
+        except snowpark.exceptions.SnowparkSQLException as e:
+            if e.sql_error_code == 400003:  # EXPERIMENT_RUN_PROPERTY_SIZE_LIMIT_EXCEEDED
+                run._warn_about_run_metadata_size(e.message)
+            else:
+                raise
 
-    @_restore_session
     def log_artifact(
         self,
         local_path: str,
@@ -381,7 +372,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
                 file_path=file_path,
             )
 
-    @_restore_session
     def list_artifacts(
         self,
         run_name: str,
@@ -410,7 +400,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             artifact_path=artifact_path or "",
         )
 
-    @_restore_session
     def download_artifacts(
         self,
         run_name: str,
@@ -452,7 +441,6 @@ 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/feature_store/access_manager.py

@@ -202,6 +202,7 @@ def _configure_role_hierarchy(
     session.sql(f"GRANT ROLE {producer_role} TO ROLE {session.get_current_role()}").collect()
 
     if consumer_role is not None:
+        # Create CONSUMER and grant it to PRODUCER to build hierarchy
         consumer_role = SqlIdentifier(consumer_role)
         session.sql(f"CREATE ROLE IF NOT EXISTS {consumer_role}").collect()
         session.sql(f"GRANT ROLE {consumer_role} TO ROLE {producer_role}").collect()