snowflake-ml-python 1.17.0__tar.gz → 1.19.0__tar.gz

{snowflake_ml_python-1.17.0 → snowflake_ml_python-1.19.0}/CHANGELOG.md

@@ -1,5 +1,40 @@
 # Release History
 
+## 1.19.0
+
+### Bug Fixes
+
+* Experiment Tracking (PrPr): No longer throw an exception in `list_artifacts` when run does not have artifacts.
+* Registry: Fix `get_version_by_alias`: now requires an exact match of snowflake identifier.
+
+### Behavior Changes
+
+### New Features
+
+* Online feature serving in Feature Store is in public preview.
+* Experiment Tracking (`snowflake.ml.experiment`) is in public preview.
+
+### Deprecations
+
+## 1.18.0
+
+### Bug Fixes
+
+* Registry: The create_service API now validates that a model has a GPU runtime configuration and will throw a
+  descriptive error if the configuration is missing.
+
+### Behavior Changes
+
+### New Features
+
+* Registry (PrPr): Introducing `ModelVersion.run_batch` for batch inference in Snowpark Container Services.
+* Experiment Tracking (PrPr): Added `version_name` argument to the autologging callbacks
+  to specify the version name for the autologged model.
+
+### Deprecations
+
+* `Python 3.9` is deprecated.
+
 ## 1.17.0
 
 ### Bug Fixes
@@ -21,6 +56,8 @@
   * Jobs submitted using v2 will automatically use the latest Container Runtime image
   * v1 behavior can be restored by setting environment variable `MLRS_USE_SUBMIT_JOB_V2` to `false`
 
+### Deprecations
+
 ## 1.16.0
 
 ### Bug Fixes

{snowflake_ml_python-1.17.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.17.0
+Version: 1.19.0
 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: 
@@ -253,7 +253,7 @@ Requires-Dist: s3fs<2026,>=2024.6.1
 Requires-Dist: scikit-learn<1.8
 Requires-Dist: scipy<2,>=1.9
 Requires-Dist: shap<1,>=0.46.0
-Requires-Dist: snowflake-connector-python[pandas]<4,>=3.17.0
+Requires-Dist: snowflake-connector-python[pandas]<4,>=3.17.3
 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
@@ -266,6 +266,7 @@ Requires-Dist: catboost<2,>=1.2.0; extra == "all"
 Requires-Dist: keras<4,>=2.0.0; extra == "all"
 Requires-Dist: lightgbm<5,>=4.1.0; extra == "all"
 Requires-Dist: mlflow<3,>=2.16.0; extra == "all"
+Requires-Dist: prophet<2,>=1.1.0; extra == "all"
 Requires-Dist: sentence-transformers<4,>=2.7.0; extra == "all"
 Requires-Dist: sentencepiece<0.2.0,>=0.1.95; extra == "all"
 Requires-Dist: streamlit<2,>=1.30.0; extra == "all"
@@ -284,8 +285,18 @@ Requires-Dist: tensorflow<3,>=2.17.0; extra == "keras"
 Requires-Dist: torch<3,>=2.0.1; extra == "keras"
 Provides-Extra: lightgbm
 Requires-Dist: lightgbm<5,>=4.1.0; extra == "lightgbm"
+Provides-Extra: llm
+Requires-Dist: mlflow<3,>=2.16.0; extra == "llm"
+Requires-Dist: sentence-transformers<4,>=2.7.0; extra == "llm"
+Requires-Dist: sentencepiece<0.2.0,>=0.1.95; extra == "llm"
+Requires-Dist: tokenizers<1,>=0.15.1; extra == "llm"
+Requires-Dist: torch<3,>=2.0.1; extra == "llm"
+Requires-Dist: torchdata<1,>=0.4; extra == "llm"
+Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "llm"
 Provides-Extra: mlflow
 Requires-Dist: mlflow<3,>=2.16.0; extra == "mlflow"
+Provides-Extra: prophet
+Requires-Dist: prophet<2,>=1.1.0; extra == "prophet"
 Provides-Extra: streamlit
 Requires-Dist: streamlit<2,>=1.30.0; extra == "streamlit"
 Provides-Extra: tensorflow
@@ -296,9 +307,7 @@ Requires-Dist: torchdata<1,>=0.4; extra == "torch"
 Provides-Extra: transformers
 Requires-Dist: sentence-transformers<4,>=2.7.0; extra == "transformers"
 Requires-Dist: sentencepiece<0.2.0,>=0.1.95; extra == "transformers"
-Requires-Dist: tokenizers<1,>=0.15.1; extra == "transformers"
 Requires-Dist: torch<3,>=2.0.1; extra == "transformers"
-Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "transformers"
 Dynamic: license-file
 
 # Snowflake ML Python
@@ -406,6 +415,41 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
+## 1.19.0
+
+### Bug Fixes
+
+* Experiment Tracking (PrPr): No longer throw an exception in `list_artifacts` when run does not have artifacts.
+* Registry: Fix `get_version_by_alias`: now requires an exact match of snowflake identifier.
+
+### Behavior Changes
+
+### New Features
+
+* Online feature serving in Feature Store is in public preview.
+* Experiment Tracking (`snowflake.ml.experiment`) is in public preview.
+
+### Deprecations
+
+## 1.18.0
+
+### Bug Fixes
+
+* Registry: The create_service API now validates that a model has a GPU runtime configuration and will throw a
+  descriptive error if the configuration is missing.
+
+### Behavior Changes
+
+### New Features
+
+* Registry (PrPr): Introducing `ModelVersion.run_batch` for batch inference in Snowpark Container Services.
+* Experiment Tracking (PrPr): Added `version_name` argument to the autologging callbacks
+  to specify the version name for the autologged model.
+
+### Deprecations
+
+* `Python 3.9` is deprecated.
+
 ## 1.17.0
 
 ### Bug Fixes
@@ -427,6 +471,8 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
   * Jobs submitted using v2 will automatically use the latest Container Runtime image
   * v1 behavior can be restored by setting environment variable `MLRS_USE_SUBMIT_JOB_V2` to `false`
 
+### Deprecations
+
 ## 1.16.0
 
 ### Bug Fixes

{snowflake_ml_python-1.17.0 → snowflake_ml_python-1.19.0}/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 = [ "anyio>=3.5.0,<5", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2024.6.1,<2026", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<3", "packaging>=20.9,<25", "pandas>=2.1.4,<3", "platformdirs<5", "pyarrow<19.0.0", "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.8", "scipy>=1.9,<2", "shap>=0.46.0,<1", "snowflake-connector-python[pandas]>=3.17.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<4",]
+dependencies = [ "anyio>=3.5.0,<5", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2024.6.1,<2026", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<3", "packaging>=20.9,<25", "pandas>=2.1.4,<3", "platformdirs<5", "pyarrow<19.0.0", "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.8", "scipy>=1.9,<2", "shap>=0.46.0,<1", "snowflake-connector-python[pandas]>=3.17.3,<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<4",]
 [[project.authors]]
 name = "Snowflake, Inc"
 email = "support@snowflake.com"
@@ -24,16 +24,18 @@ Issues = "https://github.com/snowflakedb/snowflake-ml-python/issues"
 Changelog = "https://github.com/snowflakedb/snowflake-ml-python/blob/master/CHANGELOG.md"
 
 [project.optional-dependencies]
-all = [ "altair>=5,<6", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5,!=4.51.3",]
+all = [ "altair>=5,<6", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "prophet>=1.1.0, <2", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5,!=4.51.3",]
 altair = [ "altair>=5,<6",]
 catboost = [ "catboost>=1.2.0, <2",]
 keras = [ "keras>=2.0.0,<4", "tensorflow>=2.17.0,<3", "torch>=2.0.1,<3",]
 lightgbm = [ "lightgbm>=4.1.0, <5",]
+llm = [ "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5,!=4.51.3",]
 mlflow = [ "mlflow>=2.16.0, <3",]
+prophet = [ "prophet>=1.1.0, <2",]
 streamlit = [ "streamlit>=1.30.0,<2",]
 tensorflow = [ "tensorflow>=2.17.0,<3",]
 torch = [ "torch>=2.0.1,<3", "torchdata>=0.4,<1",]
-transformers = [ "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "transformers>=4.39.3,<5,!=4.51.3",]
+transformers = [ "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "torch>=2.0.1,<3",]
 
 [tool.setuptools.package-data]
 "*" = [ "*",]

{snowflake_ml_python-1.17.0 → snowflake_ml_python-1.19.0}/snowflake/ml/_internal/telemetry.py

@@ -73,6 +73,7 @@ def _get_snowflake_connection() -> Optional[connector.SnowflakeConnection]:
 class TelemetryProject(enum.Enum):
     MLOPS = "MLOps"
     MODELING = "ModelDevelopment"
+    EXPERIMENT_TRACKING = "ExperimentTracking"
     # TODO: Update with remaining projects.
 
 
@@ -464,14 +465,14 @@ def send_api_usage_telemetry(
 
     # noqa: DAR402
     """
-    start_time = time.perf_counter()
-
     if subproject is not None and subproject_extractor is not None:
         raise ValueError("Specifying both subproject and subproject_extractor is not allowed")
 
     def decorator(func: Callable[_Args, _ReturnValue]) -> Callable[_Args, _ReturnValue]:
         @functools.wraps(func)
         def wrap(*args: Any, **kwargs: Any) -> _ReturnValue:
+            start_time = time.perf_counter()
+
             params = _get_func_params(func, func_params_to_log, args, kwargs) if func_params_to_log else None
 
             api_calls: list[Union[dict[str, Union[Callable[..., Any], str]], Callable[..., Any], str]] = []

{snowflake_ml_python-1.17.0 → snowflake_ml_python-1.19.0}/snowflake/ml/experiment/_client/experiment_tracking_sql_client.py

@@ -1,17 +1,17 @@
 from typing import Optional
 
+from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.ml.experiment._client import artifact
 from snowflake.ml.model._client.sql import _base
 from snowflake.ml.utils import sql_client
 from snowflake.snowpark import file_operation, row, session
 
+RUN_NAME_COL_NAME = "name"
+RUN_METADATA_COL_NAME = "metadata"
 
-class ExperimentTrackingSQLClient(_base._BaseSQLClient):
-
-    RUN_NAME_COL_NAME = "name"
-    RUN_METADATA_COL_NAME = "metadata"
 
+class ExperimentTrackingSQLClient(_base._BaseSQLClient):
     def __init__(
         self,
         session: session.Session,
@@ -28,6 +28,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
         """
         super().__init__(session, database_name=database_name, schema_name=schema_name)
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def create_experiment(
         self,
         experiment_name: sql_identifier.SqlIdentifier,
@@ -39,24 +40,21 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             self._session, f"CREATE EXPERIMENT {if_not_exists_sql} {experiment_fqn}"
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def drop_experiment(self, *, experiment_name: sql_identifier.SqlIdentifier) -> None:
         experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
         query_result_checker.SqlResultValidator(self._session, f"DROP EXPERIMENT {experiment_fqn}").has_dimensions(
             expected_rows=1, expected_cols=1
         ).validate()
 
-    def add_run(
-        self,
-        *,
-        experiment_name: sql_identifier.SqlIdentifier,
-        run_name: sql_identifier.SqlIdentifier,
-        live: bool = True,
-    ) -> None:
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
+    def add_run(self, *, experiment_name: sql_identifier.SqlIdentifier, run_name: sql_identifier.SqlIdentifier) -> 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} ADD {'LIVE' if live else ''} RUN {run_name}"
+            self._session, f"ALTER EXPERIMENT {experiment_fqn} ADD RUN {run_name}"
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def commit_run(
         self,
         *,
@@ -68,6 +66,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             self._session, f"ALTER EXPERIMENT {experiment_fqn} COMMIT RUN {run_name}"
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def drop_run(
         self, *, experiment_name: sql_identifier.SqlIdentifier, run_name: sql_identifier.SqlIdentifier
     ) -> None:
@@ -76,6 +75,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             self._session, f"ALTER EXPERIMENT {experiment_fqn} DROP RUN {run_name}"
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def modify_run_add_metrics(
         self,
         *,
@@ -89,6 +89,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} ADD METRICS=$${metrics}$$",
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def modify_run_add_params(
         self,
         *,
@@ -102,6 +103,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} ADD PARAMETERS=$${params}$$",
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def put_artifact(
         self,
         *,
@@ -118,6 +120,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             auto_compress=auto_compress,
         )[0]
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def list_artifacts(
         self,
         *,
@@ -125,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}/"),
@@ -142,6 +139,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             for result in results
         ]
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def get_artifact(
         self,
         *,
@@ -155,6 +153,7 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             target_directory=target_path,
         )[0]
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
     def show_runs_in_experiment(
         self, *, experiment_name: sql_identifier.SqlIdentifier, like: Optional[str] = None
     ) -> list[row.Row]:

{snowflake_ml_python-1.17.0 → snowflake_ml_python-1.19.0}/snowflake/ml/experiment/callback/keras.py

@@ -20,6 +20,7 @@ class SnowflakeKerasCallback(keras.callbacks.Callback):
         log_params: bool = True,
         log_every_n_epochs: int = 1,
         model_name: Optional[str] = None,
+        version_name: Optional[str] = None,
         model_signature: Optional["ModelSignature"] = None,
     ) -> None:
         self._experiment_tracking = experiment_tracking
@@ -30,6 +31,7 @@ class SnowflakeKerasCallback(keras.callbacks.Callback):
             raise ValueError("`log_every_n_epochs` must be positive.")
         self.log_every_n_epochs = log_every_n_epochs
         self.model_name = model_name
+        self.version_name = version_name
         self.model_signature = model_signature
 
     def on_train_begin(self, logs: Optional[dict[str, Any]] = None) -> None:
@@ -59,5 +61,6 @@ class SnowflakeKerasCallback(keras.callbacks.Callback):
             self._experiment_tracking.log_model(  # type: ignore[call-arg]
                 model=self.model,
                 model_name=model_name,
+                version_name=self.version_name,
                 signatures={"predict": self.model_signature},
             )

{snowflake_ml_python-1.17.0 → snowflake_ml_python-1.19.0}/snowflake/ml/experiment/callback/lightgbm.py

@@ -17,6 +17,7 @@ class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
         log_params: bool = True,
         log_every_n_epochs: int = 1,
         model_name: Optional[str] = None,
+        version_name: Optional[str] = None,
         model_signature: Optional["ModelSignature"] = None,
     ) -> None:
         self._experiment_tracking = experiment_tracking
@@ -27,6 +28,7 @@ class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
             raise ValueError("`log_every_n_epochs` must be positive.")
         self.log_every_n_epochs = log_every_n_epochs
         self.model_name = model_name
+        self.version_name = version_name
         self.model_signature = model_signature
 
         super().__init__(eval_result={})
@@ -50,6 +52,7 @@ class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
                     self._experiment_tracking.log_model(  # type: ignore[call-arg]
                         model=env.model,
                         model_name=model_name,
+                        version_name=self.version_name,
                         signatures={"predict": self.model_signature},
                     )
                 else:

{snowflake_ml_python-1.17.0 → snowflake_ml_python-1.19.0}/snowflake/ml/experiment/callback/xgboost.py

@@ -20,6 +20,7 @@ class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
         log_params: bool = True,
         log_every_n_epochs: int = 1,
         model_name: Optional[str] = None,
+        version_name: Optional[str] = None,
         model_signature: Optional["ModelSignature"] = None,
     ) -> None:
         self._experiment_tracking = experiment_tracking
@@ -30,6 +31,7 @@ class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
             raise ValueError("`log_every_n_epochs` must be positive.")
         self.log_every_n_epochs = log_every_n_epochs
         self.model_name = model_name
+        self.version_name = version_name
         self.model_signature = model_signature
 
     def before_training(self, model: xgb.Booster) -> xgb.Booster:
@@ -61,6 +63,7 @@ class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
             self._experiment_tracking.log_model(  # type: ignore[call-arg]
                 model=model,
                 model_name=model_name,
+                version_name=self.version_name,
                 signatures={"predict": self.model_signature},
             )
 

{snowflake_ml_python-1.17.0 → snowflake_ml_python-1.19.0}/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,29 +164,40 @@ 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,
     ) -> entities.Run:
         """
-        Start a new run.
+        Start a new run. If a run name of an existing run is provided, resumes the run if it is running.
 
         Args:
             run_name: The name of the run. If None, a default name will be generated.
 
         Returns:
-            Run: The run that was started.
+            Run: The run that was started or resumed.
 
         Raises:
-            RuntimeError: If a run is already active.
+            RuntimeError: If a run is already active. If a run with the same name exists but is not running.
         """
         if self._run:
             raise RuntimeError("A run is already active. Please end the current run before starting a new one.")
         experiment = self._get_or_set_experiment()
-        run_name = (
-            sql_identifier.SqlIdentifier(run_name) if run_name is not None else self._generate_run_name(experiment)
-        )
+
+        if run_name is None:
+            run_name = self._generate_run_name(experiment)
+        elif runs := self._sql_client.show_runs_in_experiment(experiment_name=experiment.name, like=run_name):
+            if "RUNNING" != json.loads(runs[0][sql_client.RUN_METADATA_COL_NAME])["status"]:
+                raise RuntimeError(f"Run {run_name} exists but cannot be resumed as it is no longer running.")
+            else:
+                self._run = entities.Run(
+                    experiment_tracking=self,
+                    experiment_name=experiment.name,
+                    run_name=sql_identifier.SqlIdentifier(run_name),
+                )
+                return self._run
+
+        run_name = sql_identifier.SqlIdentifier(run_name)
         self._sql_client.add_run(
             experiment_name=experiment.name,
             run_name=run_name,
@@ -217,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.
@@ -247,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,
@@ -286,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],
@@ -323,7 +308,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         """
         self.log_params({key: value})
 
-    @_restore_session
     def log_params(
         self,
         params: dict[str, Any],
@@ -345,7 +329,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             params=json.dumps([param.to_dict() for param in params_list]),
         )
 
-    @_restore_session
     def log_artifact(
         self,
         local_path: str,
@@ -369,7 +352,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
                 file_path=file_path,
             )
 
-    @_restore_session
     def list_artifacts(
         self,
         run_name: str,
@@ -398,7 +380,6 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             artifact_path=artifact_path or "",
         )
 
-    @_restore_session
     def download_artifacts(
         self,
         run_name: str,
@@ -440,11 +421,10 @@ 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)
-        existing_run_names = [row[sql_client.ExperimentTrackingSQLClient.RUN_NAME_COL_NAME] for row in existing_runs]
+        existing_run_names = [row[sql_client.RUN_NAME_COL_NAME] for row in existing_runs]
         for _ in range(1000):
             run_name = generator.generate()[1]
             if run_name not in existing_run_names: