snowflake-ml-python 1.14.0__py3-none-any.whl → 1.16.0__py3-none-any.whl

snowflake/ml/_internal/platform_capabilities.py

@@ -1,8 +1,8 @@
 import json
+import logging
 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
@@ -13,8 +13,11 @@ from snowflake.snowpark import (
     session as snowpark_session,
 )
 
+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:
@@ -60,17 +63,20 @@ class PlatformCapabilities:
     @classmethod  # type: ignore[arg-type]
     @contextmanager
     def mock_features(cls, features: dict[str, Any] = _dummy_features) -> None:  # type: ignore[misc]
-        logging.debug(f"Setting mock features: {features}")
+        logger.debug(f"Setting mock features: {features}")
         cls.set_mock_features(features)
         try:
             yield
         finally:
-            logging.debug(f"Clearing mock features: {features}")
+            logger.debug(f"Clearing mock features: {features}")
             cls.clear_mock_features()
 
     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)
 
@@ -98,7 +104,7 @@ class PlatformCapabilities:
                         error_code=error_codes.INTERNAL_SNOWML_ERROR, original_exception=RuntimeError(message)
                     )
         except snowpark_exceptions.SnowparkSQLException as e:
-            logging.debug(f"Failed to retrieve platform capabilities: {e}")
+            logger.debug(f"Failed to retrieve platform capabilities: {e}")
             # This can happen is server side is older than 9.2. That is fine.
         return {}
 
@@ -144,7 +150,7 @@ class PlatformCapabilities:
 
         value = self.features.get(feature_name)
         if value is None:
-            logging.debug(f"Feature {feature_name} not found, returning large version number")
+            logger.debug(f"Feature {feature_name} not found, returning large version number")
             return large_version
 
         try:
@@ -152,7 +158,7 @@ class PlatformCapabilities:
             version_str = str(value)
             return version.Version(version_str)
         except (version.InvalidVersion, ValueError, TypeError) as e:
-            logging.debug(
+            logger.debug(
                 f"Failed to parse version from feature {feature_name} with value '{value}': {e}. "
                 f"Returning large version number"
             )
@@ -171,7 +177,7 @@ class PlatformCapabilities:
         feature_version = self._get_version_feature(feature_name)
 
         result = current_version >= feature_version
-        logging.debug(
+        logger.debug(
             f"Version comparison for feature {feature_name}: "
             f"current={current_version}, feature={feature_version}, enabled={result}"
         )

snowflake/ml/_internal/utils/connection_params.py

@@ -1,11 +1,13 @@
 import configparser
+import logging
 import os
 from typing import Optional, Union
 
-from absl import logging
 from cryptography.hazmat import backends
 from cryptography.hazmat.primitives import serialization
 
+logger = logging.getLogger(__name__)
+
 _DEFAULT_CONNECTION_FILE = "~/.snowsql/config"
 
 
@@ -106,7 +108,7 @@ def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -
     """Loads the dictionary from snowsql config file."""
     snowsql_config_file = login_file if login_file else os.path.expanduser(_DEFAULT_CONNECTION_FILE)
     if not os.path.exists(snowsql_config_file):
-        logging.error(f"Connection name given but snowsql config file is not found at: {snowsql_config_file}")
+        logger.error(f"Connection name given but snowsql config file is not found at: {snowsql_config_file}")
         raise Exception("Snowflake SnowSQL config not found.")
 
     config = configparser.ConfigParser(inline_comment_prefixes="#")
@@ -122,7 +124,7 @@ def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -
         # See https://docs.snowflake.com/en/user-guide/snowsql-start.html#configuring-default-connection-settings
         connection_name = "connections"
 
-    logging.info(f"Reading {snowsql_config_file} for connection parameters defined as {connection_name}")
+    logger.info(f"Reading {snowsql_config_file} for connection parameters defined as {connection_name}")
     config.read(snowsql_config_file)
     conn_params = dict(config[connection_name])
     # Remap names to appropriate args in Python Connector API

snowflake/ml/_internal/utils/jwt_generator.py

@@ -110,15 +110,16 @@ class JWTGenerator:
         }
 
         # Regenerate the actual token
-        token = jwt.encode(payload, key=self.private_key, algorithm=JWTGenerator.ALGORITHM)
+        token = jwt.encode(payload, key=self.private_key, algorithm=JWTGenerator.ALGORITHM)  # type: ignore[arg-type]
         # If you are using a version of PyJWT prior to 2.0, jwt.encode returns a byte string instead of a string.
         # If the token is a byte string, convert it to a string.
         if isinstance(token, bytes):
             token = token.decode("utf-8")
         self.token = token
+        public_key = self.private_key.public_key()
         logger.info(
             "Generated a JWT with the following payload: %s",
-            jwt.decode(self.token, key=self.private_key.public_key(), algorithms=[JWTGenerator.ALGORITHM]),
+            jwt.decode(self.token, key=public_key, algorithms=[JWTGenerator.ALGORITHM]),  # type: ignore[arg-type]
         )
 
         return token

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/_internal/utils/temp_file_utils.py

@@ -1,10 +1,9 @@
+import logging
 import os
 import shutil
 import tempfile
 from typing import Iterable, Union
 
-from absl.logging import logging
-
 logger = logging.getLogger(__name__)
 
 

snowflake/ml/experiment/_client/experiment_tracking_sql_client.py

@@ -76,17 +76,30 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             self._session, f"ALTER EXPERIMENT {experiment_fqn} DROP RUN {run_name}"
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
-    def modify_run(
+    def modify_run_add_metrics(
         self,
         *,
         experiment_name: sql_identifier.SqlIdentifier,
         run_name: sql_identifier.SqlIdentifier,
-        run_metadata: str,
+        metrics: str,
     ) -> None:
         experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
         query_result_checker.SqlResultValidator(
             self._session,
-            f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} SET METADATA=$${run_metadata}$$",
+            f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} ADD METRICS=$${metrics}$$",
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def modify_run_add_params(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+        params: str,
+    ) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} ADD PARAMETERS=$${params}$$",
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
     def put_artifact(

snowflake/ml/experiment/_entities/__init__.py

@@ -1,4 +1,5 @@
 from snowflake.ml.experiment._entities.experiment import Experiment
 from snowflake.ml.experiment._entities.run import Run
+from snowflake.ml.experiment._entities.run_metadata import Metric, Param
 
-__all__ = ["Experiment", "Run"]
+__all__ = ["Experiment", "Run", "Metric", "Param"]

snowflake/ml/experiment/_entities/run.py

@@ -1,11 +1,8 @@
-import json
 import types
 from typing import TYPE_CHECKING, Optional
 
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.experiment import _experiment_info as experiment_info
-from snowflake.ml.experiment._client import experiment_tracking_sql_client
-from snowflake.ml.experiment._entities import run_metadata
 
 if TYPE_CHECKING:
     from snowflake.ml.experiment import experiment_tracking
@@ -41,18 +38,6 @@ class Run:
         if self._experiment_tracking._run is self:
             self._experiment_tracking.end_run()
 
-    def _get_metadata(
-        self,
-    ) -> run_metadata.RunMetadata:
-        runs = self._experiment_tracking._sql_client.show_runs_in_experiment(
-            experiment_name=self.experiment_name, like=str(self.name)
-        )
-        if not runs:
-            raise RuntimeError(f"Run {self.name} not found in experiment {self.experiment_name}.")
-        return run_metadata.RunMetadata.from_dict(
-            json.loads(runs[0][experiment_tracking_sql_client.ExperimentTrackingSQLClient.RUN_METADATA_COL_NAME])
-        )
-
     def _get_experiment_info(self) -> experiment_info.ExperimentInfo:
         return experiment_info.ExperimentInfo(
             fully_qualified_name=self._experiment_tracking._sql_client.fully_qualified_object_name(

snowflake/ml/experiment/_entities/run_metadata.py

@@ -1,12 +1,4 @@
 import dataclasses
-import enum
-import typing
-
-
-class RunStatus(str, enum.Enum):
-    UNKNOWN = "UNKNOWN"
-    RUNNING = "RUNNING"
-    FINISHED = "FINISHED"
 
 
 @dataclasses.dataclass
@@ -15,54 +7,14 @@ class Metric:
     value: float
     step: int
 
+    def to_dict(self) -> dict:  # type: ignore[type-arg]
+        return dataclasses.asdict(self)
+
 
 @dataclasses.dataclass
 class Param:
     name: str
     value: str
 
-
-@dataclasses.dataclass
-class RunMetadata:
-    status: RunStatus
-    metrics: list[Metric]
-    parameters: list[Param]
-
-    @classmethod
-    def from_dict(
-        cls,
-        metadata: dict,  # type: ignore[type-arg]
-    ) -> "RunMetadata":
-        return RunMetadata(
-            status=RunStatus(metadata.get("status", RunStatus.UNKNOWN.value)),
-            metrics=[Metric(**m) for m in metadata.get("metrics", [])],
-            parameters=[Param(**p) for p in metadata.get("parameters", [])],
-        )
-
     def to_dict(self) -> dict:  # type: ignore[type-arg]
         return dataclasses.asdict(self)
-
-    def set_metric(
-        self,
-        key: str,
-        value: float,
-        step: int,
-    ) -> None:
-        for metric in self.metrics:
-            if metric.name == key and metric.step == step:
-                metric.value = value
-                break
-        else:
-            self.metrics.append(Metric(name=key, value=value, step=step))
-
-    def set_param(
-        self,
-        key: str,
-        value: typing.Any,
-    ) -> None:
-        for parameter in self.parameters:
-            if parameter.name == key:
-                parameter.value = str(value)
-                break
-        else:
-            self.parameters.append(Param(name=key, value=str(value)))

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],
@@ -261,13 +300,13 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             step: The step of the metrics. Defaults to 0.
         """
         run = self._get_or_start_run()
-        metadata = run._get_metadata()
+        metrics_list = []
         for key, value in metrics.items():
-            metadata.set_metric(key, value, step)
-        self._sql_client.modify_run(
+            metrics_list.append(entities.Metric(key, value, step))
+        self._sql_client.modify_run_add_metrics(
             experiment_name=run.experiment_name,
             run_name=run.name,
-            run_metadata=json.dumps(metadata.to_dict()),
+            metrics=json.dumps([metric.to_dict() for metric in metrics_list]),
         )
 
     def log_param(
@@ -284,6 +323,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         """
         self.log_params({key: value})
 
+    @_restore_session
     def log_params(
         self,
         params: dict[str, Any],
@@ -296,15 +336,16 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
                 to string.
         """
         run = self._get_or_start_run()
-        metadata = run._get_metadata()
+        params_list = []
         for key, value in params.items():
-            metadata.set_param(key, value)
-        self._sql_client.modify_run(
+            params_list.append(entities.Param(key, str(value)))
+        self._sql_client.modify_run_add_params(
             experiment_name=run.experiment_name,
             run_name=run.name,
-            run_metadata=json.dumps(metadata.to_dict()),
+            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)