snowflake-ml-python 1.17.0__py3-none-any.whl → 1.18.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,
         *,
@@ -142,6 +145,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 +159,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

@@ -193,23 +193,35 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         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,
@@ -444,7 +456,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
     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: