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

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/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/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/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/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/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: