snowflake-ml-python 1.10.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

snowflake/cortex/_complete.py

@@ -88,7 +88,7 @@ class MidStreamException(Exception):
             message = reason
         if http_resp:
             message = f"Error in stream (HTTP Response: {http_resp.status}) - {http_resp.reason}"
-        if request_id != "":
+        if request_id is not None and request_id != "":
             # add request_id to error message
             message += f" (Request ID: {request_id})"
         super().__init__(message)
@@ -327,7 +327,8 @@ def _return_stream_response(
             # This is the case of midstream errors which were introduced specifically for structured output.
             # TODO: discuss during code review
             if parsed_resp.get("error"):
-                raise MidStreamException(reason=response.text, request_id=request_id)
+                error_info = parsed_resp["error"]
+                raise MidStreamException(reason=str(error_info), request_id=request_id)
             else:
                 pass
 

snowflake/ml/_internal/utils/service_logger.py

@@ -9,8 +9,33 @@ from typing import Optional
 
 import platformdirs
 
+
+class ProgressBarAwareConsoleHandler(logging.StreamHandler):  # type: ignore[type-arg]
+    """A logging handler that adapts to different progress bar systems to avoid interfering with display."""
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            msg = self.format(record)
+
+            # Check if tqdm progress bars are active - use tqdm.write() to avoid interference
+            try:
+                import tqdm
+
+                if hasattr(tqdm.tqdm, "_instances") and tqdm.tqdm._instances:
+                    tqdm.tqdm.write(msg, file=self.stream)
+                    return
+            except (ImportError, AttributeError):
+                pass
+
+            # Fallback to regular stream writing (works for all contexts including Streamlit)
+            self.stream.write(msg + "\n")
+            self.flush()
+        except Exception:
+            self.handleError(record)
+
+
 # Module-level logger for operational messages that should appear on console
-stdout_handler = logging.StreamHandler(sys.stdout)
+stdout_handler = ProgressBarAwareConsoleHandler(sys.stdout)
 stdout_handler.setFormatter(logging.Formatter("%(message)s"))
 
 console_logger = logging.getLogger(__name__)

snowflake/ml/experiment/_client/artifact.py

@@ -0,0 +1,76 @@
+import os
+import posixpath
+from typing import NamedTuple
+
+
+class ArtifactInfo(NamedTuple):
+    name: str
+    size: int
+    md5: str
+    last_modified: str
+
+
+def get_put_path_pairs(local_path: str, artifact_path: str) -> list[tuple[str, str]]:
+    """Enumerate files to upload and their destination subdirectories.
+
+    Expands a local path (file or directory) into a list of pairs used for uploading
+    artifacts to the stage.
+
+    Args:
+        local_path: Absolute or relative path to a local file or directory to upload.
+        artifact_path: Destination subdirectory under the run's artifact root in the
+            stage. If empty, files are uploaded to the root. When uploading a
+            directory, this value is prepended to each file's relative path.
+
+    Returns:
+        A list of tuples ``(local_file_path, destination_artifact_subdir)``. For a
+        single file, the list contains one entry with ``destination_artifact_subdir``
+        set to ``artifact_path``. For directories, one entry per file is produced and
+        subdirectories are preserved using POSIX separators.
+
+    Raises:
+        FileNotFoundError: If ``local_path`` does not exist.
+    """
+    if not os.path.exists(local_path):
+        raise FileNotFoundError(f"Local path does not exist: {local_path}")
+
+    if os.path.isfile(local_path):
+        return [(local_path, artifact_path)]
+
+    pairs: list[tuple[str, str]] = []
+    base_dir = local_path
+    for root, _, files in os.walk(base_dir):
+        if not files:
+            continue
+        rel_dir = os.path.relpath(root, base_dir)
+        rel_dir_posix = "" if rel_dir in (".", "") else rel_dir.replace(os.sep, "/")
+        dest_artifact_path = posixpath.join(artifact_path, rel_dir_posix) if rel_dir_posix else artifact_path
+        for filename in files:
+            file_path = os.path.join(root, filename)
+            pairs.append((file_path, dest_artifact_path))
+    return pairs
+
+
+def get_download_path_pairs(artifacts: list[ArtifactInfo], base_target_dir: str) -> list[tuple[str, str]]:
+    """Given artifact metadata entries, computes where each artifact should be written
+    locally.
+
+    Args:
+        artifacts: List of artifact metadata where ``name`` is a POSIX-style path
+            relative to the run's artifact root (e.g., ``"file.txt"`` or
+            ``"nested/dir/file.txt"``).
+        base_target_dir: Local base directory to download into. If empty, the
+            current working directory is used. Directories will be created as
+            needed.
+
+    Returns:
+        A list of tuples ``(artifact_relative_path, local_directory)``. Each tuple
+        provides the relative path to request from the stage and the local directory
+        where the file should be written.
+    """
+    planned: list[tuple[str, str]] = []
+    for info in artifacts:
+        rel_dir = os.path.dirname(info.name)
+        local_dir = base_target_dir if rel_dir in ("", ".") else os.path.join(base_target_dir, rel_dir)
+        planned.append((info.name, local_dir))
+    return planned

snowflake/ml/experiment/_client/experiment_tracking_sql_client.py

@@ -1,9 +1,10 @@
 from typing import Optional
 
 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 row, session
+from snowflake.snowpark import file_operation, row, session
 
 
 class ExperimentTrackingSQLClient(_base._BaseSQLClient):
@@ -88,6 +89,59 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} SET METADATA=$${run_metadata}$$",
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def put_artifact(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+        artifact_path: str,
+        file_path: str,
+        auto_compress: bool = False,
+    ) -> file_operation.PutResult:
+        return self._session.file.put(
+            local_file_name=file_path,
+            stage_location=self._build_snow_uri(experiment_name, run_name, artifact_path),
+            overwrite=True,
+            auto_compress=auto_compress,
+        )[0]
+
+    def list_artifacts(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        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()
+        )
+        return [
+            artifact.ArtifactInfo(
+                name=str(result.name).removeprefix(f"/versions/{run_name}/"),
+                size=result.size,
+                md5=result.md5,
+                last_modified=result.last_modified,
+            )
+            for result in results
+        ]
+
+    def get_artifact(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+        artifact_path: str,
+        target_path: str,
+    ) -> file_operation.GetResult:
+        return self._session.file.get(
+            stage_location=self._build_snow_uri(experiment_name, run_name, artifact_path),
+            target_directory=target_path,
+        )[0]
+
     def show_runs_in_experiment(
         self, *, experiment_name: sql_identifier.SqlIdentifier, like: Optional[str] = None
     ) -> list[row.Row]:
@@ -96,3 +150,12 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
         return query_result_checker.SqlResultValidator(
             self._session, f"SHOW RUNS {like_clause} IN EXPERIMENT {experiment_fqn}"
         ).validate()
+
+    def _build_snow_uri(
+        self, experiment_name: sql_identifier.SqlIdentifier, run_name: sql_identifier.SqlIdentifier, artifact_path: str
+    ) -> str:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        uri = f"snow://experiment/{experiment_fqn}/versions/{run_name}"
+        if artifact_path:
+            uri += f"/{artifact_path}"
+        return uri

snowflake/ml/experiment/callback/keras.py

@@ -0,0 +1,63 @@
+import json
+from typing import TYPE_CHECKING, Any, Optional
+from warnings import warn
+
+import keras
+
+from snowflake.ml.experiment import utils
+
+if TYPE_CHECKING:
+    from snowflake.ml.experiment.experiment_tracking import ExperimentTracking
+    from snowflake.ml.model.model_signature import ModelSignature
+
+
+class SnowflakeKerasCallback(keras.callbacks.Callback):
+    def __init__(
+        self,
+        experiment_tracking: "ExperimentTracking",
+        log_model: bool = True,
+        log_metrics: bool = True,
+        log_params: bool = True,
+        log_every_n_epochs: int = 1,
+        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
+        if log_every_n_epochs < 1:
+            raise ValueError("`log_every_n_epochs` must be positive.")
+        self.log_every_n_epochs = log_every_n_epochs
+        self.model_name = model_name
+        self.model_signature = model_signature
+
+    def on_train_begin(self, logs: Optional[dict[str, Any]] = None) -> None:
+        if self.log_params:
+            params = json.loads(self.model.to_json())
+            self._experiment_tracking.log_params(utils.flatten_nested_params(params))
+
+    def on_epoch_end(self, epoch: int, logs: Optional[dict[str, Any]] = None) -> None:
+        if self.log_metrics and logs and epoch % self.log_every_n_epochs == 0:
+            for key, value in logs.items():
+                try:
+                    value = float(value)
+                except Exception:
+                    pass
+                else:
+                    self._experiment_tracking.log_metric(key=key, value=value, step=epoch)
+
+    def on_train_end(self, logs: Optional[dict[str, Any]] = None) -> None:
+        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 SnowflakeKerasCallback."
+                )
+                return
+            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=self.model,
+                model_name=model_name,
+                signatures={"predict": self.model_signature},
+            )

snowflake/ml/experiment/callback/lightgbm.py

@@ -15,6 +15,7 @@ class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
         log_model: bool = True,
         log_metrics: bool = True,
         log_params: bool = True,
+        log_every_n_epochs: int = 1,
         model_name: Optional[str] = None,
         model_signature: Optional["ModelSignature"] = None,
     ) -> None:
@@ -22,6 +23,9 @@ class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
         self.log_model = log_model
         self.log_metrics = log_metrics
         self.log_params = log_params
+        if log_every_n_epochs < 1:
+            raise ValueError("`log_every_n_epochs` must be positive.")
+        self.log_every_n_epochs = log_every_n_epochs
         self.model_name = model_name
         self.model_signature = model_signature
 
@@ -32,7 +36,7 @@ class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
             if env.iteration == env.begin_iteration:  # Log params only at the first iteration
                 self._experiment_tracking.log_params(env.params)
 
-        if self.log_metrics:
+        if self.log_metrics and env.iteration % self.log_every_n_epochs == 0:
             super().__call__(env)
             for dataset_name, metrics in self.eval_result.items():
                 for metric_name, log in metrics.items():

snowflake/ml/experiment/callback/xgboost.py

@@ -18,6 +18,7 @@ class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
         log_model: bool = True,
         log_metrics: bool = True,
         log_params: bool = True,
+        log_every_n_epochs: int = 1,
         model_name: Optional[str] = None,
         model_signature: Optional["ModelSignature"] = None,
     ) -> None:
@@ -25,6 +26,9 @@ class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
         self.log_model = log_model
         self.log_metrics = log_metrics
         self.log_params = log_params
+        if log_every_n_epochs < 1:
+            raise ValueError("`log_every_n_epochs` must be positive.")
+        self.log_every_n_epochs = log_every_n_epochs
         self.model_name = model_name
         self.model_signature = model_signature
 
@@ -36,7 +40,7 @@ class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
         return model
 
     def after_iteration(self, model: Any, epoch: int, evals_log: dict[str, dict[str, Any]]) -> bool:
-        if self.log_metrics:
+        if self.log_metrics and epoch % self.log_every_n_epochs == 0:
             for dataset_name, metrics in evals_log.items():
                 for metric_name, log in metrics.items():
                     metric_key = dataset_name + ":" + metric_name

snowflake/ml/experiment/experiment_tracking.py

@@ -5,14 +5,17 @@ from typing import Any, Optional, Union
 from urllib.parse import quote
 
 import snowflake.snowpark._internal.utils as snowpark_utils
-from snowflake.ml import model, registry
+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 sql_identifier
 from snowflake.ml.experiment import (
     _entities as entities,
     _experiment_info as experiment_info,
 )
-from snowflake.ml.experiment._client import experiment_tracking_sql_client as sql_client
+from snowflake.ml.experiment._client import (
+    artifact,
+    experiment_tracking_sql_client as sql_client,
+)
 from snowflake.ml.model import type_hints
 from snowflake.ml.utils import sql_client as sql_client_utils
 from snowflake.snowpark import session
@@ -118,11 +121,11 @@ class ExperimentTracking:
     @functools.wraps(registry.Registry.log_model)
     def log_model(
         self,
-        model: Union[type_hints.SupportedModelType, model.ModelVersion],
+        model: Union[type_hints.SupportedModelType, ml_model.ModelVersion],
         *,
         model_name: str,
         **kwargs: Any,
-    ) -> model.ModelVersion:
+    ) -> ml_model.ModelVersion:
         run = self._get_or_start_run()
         with experiment_info.ExperimentInfoPatcher(experiment_info=run._get_experiment_info()):
             return self._registry.log_model(model, model_name=model_name, **kwargs)
@@ -280,6 +283,88 @@ class ExperimentTracking:
             run_metadata=json.dumps(metadata.to_dict()),
         )
 
+    def log_artifact(
+        self,
+        local_path: str,
+        artifact_path: Optional[str] = None,
+    ) -> None:
+        """
+        Log an artifact or a directory of artifacts under the current run. If no run is active, this method will create
+        a new run.
+
+        Args:
+            local_path: The path to the local file or directory to write.
+            artifact_path: The directory within the run directory to write the artifacts to. If None, the artifacts will
+                be logged in the root directory of the run.
+        """
+        run = self._get_or_start_run()
+        for file_path, dest_artifact_path in artifact.get_put_path_pairs(local_path, artifact_path or ""):
+            self._sql_client.put_artifact(
+                experiment_name=run.experiment_name,
+                run_name=run.name,
+                artifact_path=dest_artifact_path,
+                file_path=file_path,
+            )
+
+    def list_artifacts(
+        self,
+        run_name: str,
+        artifact_path: Optional[str] = None,
+    ) -> list[artifact.ArtifactInfo]:
+        """
+        List artifacts for a given run within the current experiment.
+
+        Args:
+            run_name: Name of the run to list artifacts from.
+            artifact_path: Optional subdirectory within the run's artifact directory to scope the listing.
+                If None, lists from the root of the run's artifact directory.
+
+        Returns:
+            A list of artifact entries under the specified path.
+
+        Raises:
+            RuntimeError: If no experiment is currently set.
+        """
+        if not self._experiment:
+            raise RuntimeError("No experiment set. Please set an experiment before listing artifacts.")
+
+        return self._sql_client.list_artifacts(
+            experiment_name=self._experiment.name,
+            run_name=sql_identifier.SqlIdentifier(run_name),
+            artifact_path=artifact_path or "",
+        )
+
+    def download_artifacts(
+        self,
+        run_name: str,
+        artifact_path: Optional[str] = None,
+        target_path: Optional[str] = None,
+    ) -> None:
+        """
+        Download artifacts from a run to a local directory.
+
+        Args:
+            run_name: Name of the run to download artifacts from.
+            artifact_path: Optional path to file or subdirectory within the run's artifact directory to download.
+                If None, downloads all artifacts from the root of the run's artifact directory.
+            target_path: Optional local directory to download files into. If None, downloads into the
+                current working directory.
+
+        Raises:
+            RuntimeError: If no experiment is currently set.
+        """
+        if not self._experiment:
+            raise RuntimeError("No experiment set. Please set an experiment before downloading artifacts.")
+
+        artifacts = self.list_artifacts(run_name=run_name, artifact_path=artifact_path or "")
+        for relative_path, local_dir in artifact.get_download_path_pairs(artifacts, target_path or ""):
+            self._sql_client.get_artifact(
+                experiment_name=self._experiment.name,
+                run_name=sql_identifier.SqlIdentifier(run_name),
+                artifact_path=relative_path,
+                target_path=local_dir,
+            )
+
     def _get_or_set_experiment(self) -> entities.Experiment:
         if self._experiment:
             return self._experiment