snowflake-ml-python 1.9.0__py3-none-any.whl → 1.9.2__py3-none-any.whl

snowflake/ml/experiment/_client/experiment_tracking_sql_client.py

@@ -0,0 +1,98 @@
+from typing import Optional
+
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.ml.model._client.sql import _base
+from snowflake.ml.utils import sql_client
+from snowflake.snowpark import row, session
+
+
+class ExperimentTrackingSQLClient(_base._BaseSQLClient):
+
+    RUN_NAME_COL_NAME = "name"
+    RUN_METADATA_COL_NAME = "metadata"
+
+    def __init__(
+        self,
+        session: session.Session,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        """Snowflake SQL Client to manage experiment tracking.
+
+        Args:
+            session: Active snowpark session.
+            database_name: Name of the Database where experiment tracking resources are provisioned.
+            schema_name: Name of the Schema where experiment tracking resources are provisioned.
+        """
+        super().__init__(session, database_name=database_name, schema_name=schema_name)
+
+    def create_experiment(
+        self,
+        experiment_name: sql_identifier.SqlIdentifier,
+        creation_mode: sql_client.CreationMode,
+    ) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        if_not_exists_sql = "IF NOT EXISTS" if creation_mode.if_not_exists else ""
+        query_result_checker.SqlResultValidator(
+            self._session, f"CREATE EXPERIMENT {if_not_exists_sql} {experiment_fqn}"
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    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:
+        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}"
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def commit_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} COMMIT RUN {run_name}"
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def drop_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} DROP RUN {run_name}"
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def modify_run(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+        run_metadata: 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}$$",
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def show_runs_in_experiment(
+        self, *, experiment_name: sql_identifier.SqlIdentifier, like: Optional[str] = None
+    ) -> list[row.Row]:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        like_clause = f"LIKE '{like}'" if like else ""
+        return query_result_checker.SqlResultValidator(
+            self._session, f"SHOW RUNS {like_clause} IN EXPERIMENT {experiment_fqn}"
+        ).validate()

snowflake/ml/experiment/_entities/__init__.py

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

snowflake/ml/experiment/_entities/experiment.py

@@ -0,0 +1,10 @@
+from snowflake.ml._internal.utils import sql_identifier
+
+
+class Experiment:
+    def __init__(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        self.name = experiment_name

snowflake/ml/experiment/_entities/run.py

@@ -0,0 +1,62 @@
+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
+
+
+class Run:
+    def __init__(
+        self,
+        experiment_tracking: "experiment_tracking.ExperimentTracking",
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        self._experiment_tracking = experiment_tracking
+        self.experiment_name = experiment_name
+        self.name = run_name
+
+        self._patcher = experiment_info.ExperimentInfoPatcher(
+            experiment_info=self._get_experiment_info(),
+        )
+
+    def __enter__(self) -> "Run":
+        self._patcher.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[types.TracebackType],
+    ) -> None:
+        self._patcher.__exit__(exc_type, exc_value, traceback)
+        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(
+                self._experiment_tracking._database_name, self._experiment_tracking._schema_name, self.experiment_name
+            ),
+            run_name=self.name.identifier(),
+        )

snowflake/ml/experiment/_entities/run_metadata.py

@@ -0,0 +1,68 @@
+import dataclasses
+import enum
+import typing
+
+
+class RunStatus(str, enum.Enum):
+    UNKNOWN = "UNKNOWN"
+    RUNNING = "RUNNING"
+    FINISHED = "FINISHED"
+
+
+@dataclasses.dataclass
+class Metric:
+    name: str
+    value: float
+    step: int
+
+
+@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_info.py

@@ -0,0 +1,63 @@
+import dataclasses
+import functools
+import types
+from typing import Callable, Optional
+
+from snowflake.ml import model
+from snowflake.ml.registry._manager import model_manager
+
+
+@dataclasses.dataclass(frozen=True)
+class ExperimentInfo:
+    """Serializable information identifying a Experiment"""
+
+    fully_qualified_name: str
+    run_name: str
+
+
+class ExperimentInfoPatcher:
+    """Context manager that patches ModelManager.log_model to include experiment information.
+
+    This class maintains a stack of active experiment contexts and ensures that
+    log_model calls are automatically tagged with the appropriate experiment info.
+    """
+
+    # Store original method at class definition time to avoid recursive patching
+    _original_log_model: Callable[..., model.ModelVersion] = model_manager.ModelManager.log_model
+
+    # Stack of active experiment_info contexts for nested experiment support
+    _experiment_info_stack: list[ExperimentInfo] = []
+
+    def __init__(self, experiment_info: ExperimentInfo) -> None:
+        self._experiment_info = experiment_info
+
+    def __enter__(self) -> "ExperimentInfoPatcher":
+        # Only patch ModelManager.log_model if we're the first patcher to avoid nested patching
+        if not ExperimentInfoPatcher._experiment_info_stack:
+
+            @functools.wraps(ExperimentInfoPatcher._original_log_model)
+            def patched(*args, **kwargs) -> model.ModelVersion:  # type: ignore[no-untyped-def]
+                # Use the most recent (top of stack) experiment_info for nested contexts
+                current_experiment_info = ExperimentInfoPatcher._experiment_info_stack[-1]
+                return ExperimentInfoPatcher._original_log_model(
+                    *args, **kwargs, experiment_info=current_experiment_info
+                )
+
+            model_manager.ModelManager.log_model = patched  # type: ignore[method-assign]
+
+        ExperimentInfoPatcher._experiment_info_stack.append(self._experiment_info)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[types.TracebackType],
+    ) -> None:
+        ExperimentInfoPatcher._experiment_info_stack.pop()
+
+        # Restore original method when no patches are active to clean up properly
+        if not ExperimentInfoPatcher._experiment_info_stack:
+            model_manager.ModelManager.log_model = (  # type: ignore[method-assign]
+                ExperimentInfoPatcher._original_log_model
+            )

snowflake/ml/experiment/callback.py

@@ -0,0 +1,121 @@
+import json
+from typing import TYPE_CHECKING, Any, Optional, Union
+from warnings import warn
+
+import lightgbm as lgb
+import xgboost as xgb
+
+from snowflake.ml.model.model_signature import ModelSignature
+
+if TYPE_CHECKING:
+    from snowflake.ml.experiment.experiment_tracking import ExperimentTracking
+
+
+class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
+    def __init__(
+        self,
+        experiment_tracking: "ExperimentTracking",
+        log_model: bool = True,
+        log_metrics: bool = True,
+        log_params: bool = True,
+        model_name: Optional[str] = None,
+        model_signature: Optional[ModelSignature] = None,
+    ) -> None:
+        self._experiment_tracking = experiment_tracking
+        self.log_model = log_model
+        self.log_metrics = log_metrics
+        self.log_params = log_params
+        self.model_name = model_name
+        self.model_signature = model_signature
+
+    def before_training(self, model: xgb.Booster) -> xgb.Booster:
+        def _flatten_nested_params(params: Union[list[Any], dict[str, Any]], prefix: str = "") -> dict[str, Any]:
+            flat_params = {}
+            items = params.items() if isinstance(params, dict) else enumerate(params)
+            for key, value in items:
+                new_prefix = f"{prefix}.{key}" if prefix else str(key)
+                if isinstance(value, (dict, list)):
+                    flat_params.update(_flatten_nested_params(value, new_prefix))
+                else:
+                    flat_params[new_prefix] = value
+            return flat_params
+
+        if self.log_params:
+            params = json.loads(model.save_config())
+            self._experiment_tracking.log_params(_flatten_nested_params(params))
+
+        return model
+
+    def after_iteration(self, model: Any, epoch: int, evals_log: dict[str, dict[str, Any]]) -> bool:
+        if self.log_metrics:
+            for dataset_name, metrics in evals_log.items():
+                for metric_name, log in metrics.items():
+                    metric_key = dataset_name + ":" + metric_name
+                    self._experiment_tracking.log_metric(key=metric_key, value=log[-1], step=epoch)
+
+        return False
+
+    def after_training(self, model: xgb.Booster) -> xgb.Booster:
+        if self.log_model:
+            if not self.model_signature:
+                warn(
+                    "Model will not be logged because model signature is missing. "
+                    "To autolog the model, please specify `model_signature` when constructing SnowflakeXgboostCallback."
+                )
+                return model
+
+            model_name = self.model_name or self._experiment_tracking._get_or_set_experiment().name + "_model"
+            self._experiment_tracking.log_model(  # type: ignore[call-arg]
+                model=model,
+                model_name=model_name,
+                signatures={"predict": self.model_signature},
+            )
+
+        return model
+
+
+class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
+    def __init__(
+        self,
+        experiment_tracking: "ExperimentTracking",
+        log_model: bool = True,
+        log_metrics: bool = True,
+        log_params: bool = True,
+        model_name: Optional[str] = None,
+        model_signature: Optional[ModelSignature] = None,
+    ) -> None:
+        self._experiment_tracking = experiment_tracking
+        self.log_model = log_model
+        self.log_metrics = log_metrics
+        self.log_params = log_params
+        self.model_name = model_name
+        self.model_signature = model_signature
+
+        super().__init__(eval_result={})
+
+    def __call__(self, env: lgb.callback.CallbackEnv) -> None:
+        if self.log_params:
+            if env.iteration == env.begin_iteration:  # Log params only at the first iteration
+                self._experiment_tracking.log_params(env.params)
+
+        if self.log_metrics:
+            super().__call__(env)
+            for dataset_name, metrics in self.eval_result.items():
+                for metric_name, log in metrics.items():
+                    metric_key = dataset_name + ":" + metric_name
+                    self._experiment_tracking.log_metric(key=metric_key, value=log[-1], step=env.iteration)
+
+        if self.log_model:
+            if env.iteration == env.end_iteration - 1:  # Log model only at the last iteration
+                if self.model_signature:
+                    model_name = self.model_name or self._experiment_tracking._get_or_set_experiment().name + "_model"
+                    self._experiment_tracking.log_model(  # type: ignore[call-arg]
+                        model=env.model,
+                        model_name=model_name,
+                        signatures={"predict": self.model_signature},
+                    )
+                else:
+                    warn(
+                        "Model will not be logged because model signature is missing. To autolog the model, "
+                        "please specify `model_signature` when constructing SnowflakeLightgbmCallback."
+                    )