snowflake-ml-python 1.8.6__py3-none-any.whl → 1.9.1__py3-none-any.whl

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/experiment_tracking.py

@@ -0,0 +1,319 @@
+import functools
+import json
+import sys
+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._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.model import type_hints
+from snowflake.ml.utils import sql_client as sql_client_utils
+from snowflake.snowpark import session
+
+DEFAULT_EXPERIMENT_NAME = sql_identifier.SqlIdentifier("DEFAULT")
+
+
+class ExperimentTracking:
+    """
+    Class to manage experiments in Snowflake.
+    """
+
+    @snowpark_utils.private_preview(version="1.9.1")
+    def __init__(
+        self,
+        session: session.Session,
+        *,
+        database_name: Optional[str] = None,
+        schema_name: Optional[str] = None,
+    ) -> None:
+        """
+        Initializes experiment tracking within a pre-created schema.
+
+        Args:
+            session: The Snowpark Session to connect with Snowflake.
+            database_name: The name of the database. If None, the current database of the session
+                will be used. Defaults to None.
+            schema_name: The name of the schema. If None, the current schema of the session
+                will be used. If there is no active schema, the PUBLIC schema will be used. Defaults to None.
+
+        Raises:
+            ValueError: If no database is provided and no active database exists in the session.
+        """
+        if database_name:
+            self._database_name = sql_identifier.SqlIdentifier(database_name)
+        elif session_db := session.get_current_database():
+            self._database_name = sql_identifier.SqlIdentifier(session_db)
+        else:
+            raise ValueError("You need to provide a database to use experiment tracking.")
+
+        if schema_name:
+            self._schema_name = sql_identifier.SqlIdentifier(schema_name)
+        elif session_schema := session.get_current_schema():
+            self._schema_name = sql_identifier.SqlIdentifier(session_schema)
+        else:
+            self._schema_name = sql_identifier.SqlIdentifier("PUBLIC")
+
+        self._sql_client = sql_client.ExperimentTrackingSQLClient(
+            session,
+            database_name=self._database_name,
+            schema_name=self._schema_name,
+        )
+        self._registry = registry.Registry(
+            session=session,
+            database_name=self._database_name,
+            schema_name=self._schema_name,
+        )
+
+        # The experiment in context
+        self._experiment: Optional[entities.Experiment] = None
+        # The run in context
+        self._run: Optional[entities.Run] = None
+
+    def set_experiment(
+        self,
+        experiment_name: str,
+    ) -> entities.Experiment:
+        """
+        Set the experiment in context. Creates a new experiment if it doesn't exist.
+
+        Args:
+            experiment_name: The name of the experiment.
+
+        Returns:
+            Experiment: The experiment that was set.
+        """
+        experiment_name = sql_identifier.SqlIdentifier(experiment_name)
+        if self._experiment and self._experiment.name == experiment_name:
+            return self._experiment
+        self._sql_client.create_experiment(
+            experiment_name=experiment_name,
+            creation_mode=sql_client_utils.CreationMode(if_not_exists=True),
+        )
+        self._experiment = entities.Experiment(experiment_name=experiment_name)
+        self._run = None
+        return self._experiment
+
+    def delete_experiment(
+        self,
+        experiment_name: str,
+    ) -> None:
+        """
+        Delete an experiment.
+
+        Args:
+            experiment_name: The name of the experiment.
+        """
+        self._sql_client.drop_experiment(experiment_name=sql_identifier.SqlIdentifier(experiment_name))
+        if self._experiment and self._experiment.name == experiment_name:
+            self._experiment = None
+            self._run = None
+
+    @functools.wraps(registry.Registry.log_model)
+    def log_model(
+        self,
+        model: Union[type_hints.SupportedModelType, model.ModelVersion],
+        *,
+        model_name: str,
+        **kwargs: Any,
+    ) -> 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)
+
+    def start_run(
+        self,
+        run_name: Optional[str] = None,
+    ) -> entities.Run:
+        """
+        Start a new run.
+
+        Args:
+            run_name: The name of the run. If None, a default name will be generated.
+
+        Returns:
+            Run: The run that was started.
+
+        Raises:
+            RuntimeError: If a run is already active.
+        """
+        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)
+        )
+        self._sql_client.add_run(
+            experiment_name=experiment.name,
+            run_name=run_name,
+        )
+        self._run = entities.Run(experiment_tracking=self, experiment_name=experiment.name, run_name=run_name)
+        return self._run
+
+    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.
+
+        Args:
+            run_name: The name of the run to be ended. If None, the current run is ended.
+
+        Raises:
+            RuntimeError: If no run is active.
+        """
+        if not self._experiment:
+            raise RuntimeError("No experiment set. Please set an experiment before ending a run.")
+        experiment_name = self._experiment.name
+
+        if run_name:
+            run_name = sql_identifier.SqlIdentifier(run_name)
+        elif self._run:
+            run_name = self._run.name
+        else:
+            raise RuntimeError("No run is active. Please start a run before ending it.")
+
+        self._sql_client.commit_run(
+            experiment_name=experiment_name,
+            run_name=run_name,
+        )
+        if self._run and run_name == self._run.name:
+            self._run = None
+        self._print_urls(experiment_name=experiment_name, run_name=run_name)
+
+    def delete_run(
+        self,
+        run_name: str,
+    ) -> None:
+        """
+        Delete a run.
+
+        Args:
+            run_name: The name of the run to be deleted.
+
+        Raises:
+            RuntimeError: If no experiment is set.
+        """
+        if not self._experiment:
+            raise RuntimeError("No experiment set. Please set an experiment before deleting a run.")
+        self._sql_client.drop_run(
+            experiment_name=self._experiment.name,
+            run_name=sql_identifier.SqlIdentifier(run_name),
+        )
+        if self._run and self._run.name == run_name:
+            self._run = None
+
+    def log_metric(
+        self,
+        key: str,
+        value: float,
+        step: int = 0,
+    ) -> None:
+        """
+        Log a metric under the current run. If no run is active, this method will create a new run.
+
+        Args:
+            key: The name of the metric.
+            value: The value of the metric.
+            step: The step of the metric. Defaults to 0.
+        """
+        self.log_metrics(metrics={key: value}, step=step)
+
+    def log_metrics(
+        self,
+        metrics: dict[str, float],
+        step: int = 0,
+    ) -> None:
+        """
+        Log metrics under the current run. If no run is active, this method will create a new run.
+
+        Args:
+            metrics: Dictionary containing metric keys and float values.
+            step: The step of the metrics. Defaults to 0.
+        """
+        run = self._get_or_start_run()
+        metadata = run._get_metadata()
+        for key, value in metrics.items():
+            metadata.set_metric(key, value, step)
+        self._sql_client.modify_run(
+            experiment_name=run.experiment_name,
+            run_name=run.name,
+            run_metadata=json.dumps(metadata.to_dict()),
+        )
+
+    def log_param(
+        self,
+        key: str,
+        value: Any,
+    ) -> None:
+        """
+        Log a parameter under the current run. If no run is active, this method will create a new run.
+
+        Args:
+            key: The name of the parameter.
+            value: The value of the parameter. Values can be of any type, but will be converted to string.
+        """
+        self.log_params({key: value})
+
+    def log_params(
+        self,
+        params: dict[str, Any],
+    ) -> None:
+        """
+        Log parameters under the current run. If no run is active, this method will create a new run.
+
+        Args:
+            params: Dictionary containing parameter keys and values. Values can be of any type, but will be converted
+                to string.
+        """
+        run = self._get_or_start_run()
+        metadata = run._get_metadata()
+        for key, value in params.items():
+            metadata.set_param(key, value)
+        self._sql_client.modify_run(
+            experiment_name=run.experiment_name,
+            run_name=run.name,
+            run_metadata=json.dumps(metadata.to_dict()),
+        )
+
+    def _get_or_set_experiment(self) -> entities.Experiment:
+        if self._experiment:
+            return self._experiment
+        return self.set_experiment(experiment_name=DEFAULT_EXPERIMENT_NAME)
+
+    def _get_or_start_run(self) -> entities.Run:
+        if self._run:
+            return self._run
+        return self.start_run()
+
+    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]
+        for _ in range(1000):
+            run_name = generator.generate()[1]
+            if run_name not in existing_run_names:
+                return sql_identifier.SqlIdentifier(run_name)
+        raise RuntimeError("Random run name generation failed.")
+
+    def _print_urls(
+        self,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+        scheme: str = "https",
+        host: str = "app.snowflake.com",
+    ) -> None:
+
+        experiment_url = (
+            f"{scheme}://{host}/_deeplink/#/experiments"
+            f"/databases/{quote(str(self._database_name))}"
+            f"/schemas/{quote(str(self._schema_name))}"
+            f"/experiments/{quote(str(experiment_name))}"
+        )
+        run_url = experiment_url + f"/runs/{quote(str(run_name))}"
+        sys.stdout.write(f"🏃 View run {run_name} at: {run_url}\n")
+        sys.stdout.write(f"🧪 View experiment at: {experiment_url}\n")

snowflake/ml/jobs/_utils/constants.py

@@ -15,7 +15,7 @@ STAGE_VOLUME_MOUNT_PATH = "/mnt/app"
 DEFAULT_IMAGE_REPO = "/snowflake/images/snowflake_images"
 DEFAULT_IMAGE_CPU = "st_plat/runtime/x86/runtime_image/snowbooks"
 DEFAULT_IMAGE_GPU = "st_plat/runtime/x86/generic_gpu/runtime_image/snowbooks"
-DEFAULT_IMAGE_TAG = "1.4.2"
+DEFAULT_IMAGE_TAG = "1.5.0"
 DEFAULT_ENTRYPOINT_PATH = "func.py"
 
 # Percent of container memory to allocate for /dev/shm volume

snowflake/ml/jobs/_utils/interop_utils.py

@@ -75,16 +75,75 @@ def fetch_result(session: snowpark.Session, result_path: str) -> ExecutionResult
 
     Returns:
         A dictionary containing the execution result if available, None otherwise.
+
+    Raises:
+        RuntimeError: If both pickle and JSON result retrieval fail.
     """
     try:
         # TODO: Check if file exists
         with session.file.get_stream(result_path) as result_stream:
             return ExecutionResult.from_dict(pickle.load(result_stream))
-    except (sp_exceptions.SnowparkSQLException, pickle.UnpicklingError, TypeError, ImportError):
+    except (
+        sp_exceptions.SnowparkSQLException,
+        pickle.UnpicklingError,
+        TypeError,
+        ImportError,
+        AttributeError,
+        MemoryError,
+    ) as pickle_error:
         # Fall back to JSON result if loading pickled result fails for any reason
-        result_json_path = os.path.splitext(result_path)[0] + ".json"
-        with session.file.get_stream(result_json_path) as result_stream:
-            return ExecutionResult.from_dict(json.load(result_stream))
+        try:
+            result_json_path = os.path.splitext(result_path)[0] + ".json"
+            with session.file.get_stream(result_json_path) as result_stream:
+                return ExecutionResult.from_dict(json.load(result_stream))
+        except Exception as json_error:
+            # Both pickle and JSON failed - provide helpful error message
+            raise RuntimeError(_fetch_result_error_message(pickle_error, result_path, json_error)) from pickle_error
+
+
+def _fetch_result_error_message(error: Exception, result_path: str, json_error: Optional[Exception] = None) -> str:
+    """Create helpful error messages for common result retrieval failures."""
+
+    # Package import issues
+    if isinstance(error, ImportError):
+        return f"Failed to retrieve job result: Package not installed in your local environment. Error: {str(error)}"
+
+    # Package versions differ between runtime and local environment
+    if isinstance(error, AttributeError):
+        return f"Failed to retrieve job result: Package version mismatch. Error: {str(error)}"
+
+    # Serialization issues
+    if isinstance(error, TypeError):
+        return f"Failed to retrieve job result: Non-serializable objects were returned. Error: {str(error)}"
+
+    # Python version pickling incompatibility
+    if isinstance(error, pickle.UnpicklingError) and "protocol" in str(error).lower():
+        # TODO: Update this once we support different Python versions
+        client_version = f"Python {sys.version_info.major}.{sys.version_info.minor}"
+        runtime_version = "Python 3.10"
+        return (
+            f"Failed to retrieve job result: Python version mismatch - job ran on {runtime_version}, "
+            f"local environment using Python {client_version}. Error: {str(error)}"
+        )
+
+    # File access issues
+    if isinstance(error, sp_exceptions.SnowparkSQLException):
+        if "not found" in str(error).lower() or "does not exist" in str(error).lower():
+            return (
+                f"Failed to retrieve job result: No result file found. Check job.get_logs() for execution "
+                f"errors. Error: {str(error)}"
+            )
+        else:
+            return f"Failed to retrieve job result: Cannot access result file. Error: {str(error)}"
+
+    if isinstance(error, MemoryError):
+        return f"Failed to retrieve job result: Result too large for memory. Error: {str(error)}"
+
+    # Generic fallback
+    base_message = f"Failed to retrieve job result: {str(error)}"
+    if json_error:
+        base_message += f" (JSON fallback also failed: {str(json_error)})"
+    return base_message
 
 
 def load_exception(exc_type_name: str, exc_value: Union[Exception, str], exc_tb: str) -> Exception:

snowflake/ml/jobs/_utils/payload_utils.py

@@ -15,6 +15,7 @@ from snowflake import snowpark
 from snowflake.ml.jobs._utils import (
     constants,
     function_payload_utils,
+    query_helper,
     stage_utils,
     types,
 )
@@ -312,14 +313,15 @@ class JobPayload:
         stage_name = stage_path.parts[0].lstrip("@")
         # Explicitly check if stage exists first since we may not have CREATE STAGE privilege
         try:
-            session.sql("describe stage identifier(?)", params=[stage_name]).collect()
+            query_helper.run_query(session, "describe stage identifier(?)", params=[stage_name])
         except sp_exceptions.SnowparkSQLException:
-            session.sql(
+            query_helper.run_query(
+                session,
                 "create stage if not exists identifier(?)"
                 " encryption = ( type = 'SNOWFLAKE_SSE' )"
                 " comment = 'Created by snowflake.ml.jobs Python API'",
                 params=[stage_name],
-            ).collect()
+            )
 
         # Upload payload to stage
         if not isinstance(source, (Path, stage_utils.StagePath)):

snowflake/ml/jobs/_utils/query_helper.py

@@ -0,0 +1,20 @@
+from typing import Any, Optional, Sequence
+
+from snowflake import snowpark
+from snowflake.snowpark import Row
+from snowflake.snowpark._internal import utils
+from snowflake.snowpark._internal.analyzer import snowflake_plan
+
+
+def result_set_to_rows(session: snowpark.Session, result: dict[str, Any]) -> list[Row]:
+    metadata = session._conn._cursor.description
+    result_set = result["data"]
+    return utils.result_set_to_rows(result_set, metadata)
+
+
+@snowflake_plan.SnowflakePlan.Decorator.wrap_exception  # type: ignore[misc]
+def run_query(session: snowpark.Session, query_text: str, params: Optional[Sequence[Any]] = None) -> list[Row]:
+    result = session._conn.run_query(query=query_text, params=params, _force_qmark_paramstyle=True)
+    if not isinstance(result, dict) or "data" not in result:
+        raise ValueError(f"Unprocessable result: {result}")
+    return result_set_to_rows(session, result)

snowflake/ml/jobs/_utils/scripts/mljob_launcher.py

@@ -16,9 +16,13 @@ import cloudpickle
 from constants import LOG_END_MSG, LOG_START_MSG, MIN_INSTANCES_ENV_VAR
 
 from snowflake.ml.jobs._utils import constants
-from snowflake.ml.utils.connection_params import SnowflakeLoginOptions
 from snowflake.snowpark import Session
 
+try:
+    from snowflake.ml._internal.utils.connection_params import SnowflakeLoginOptions
+except ImportError:
+    from snowflake.ml.utils.connection_params import SnowflakeLoginOptions
+
 # Configure logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)

snowflake/ml/jobs/_utils/spec_utils.py

@@ -6,13 +6,17 @@ from typing import Any, Optional, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import snowflake_env
-from snowflake.ml.jobs._utils import constants, types
+from snowflake.ml.jobs._utils import constants, query_helper, types
 
 
 def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.ComputeResources:
     """Extract resource information for the specified compute pool"""
     # Get the instance family
-    rows = session.sql("show compute pools like ?", params=[compute_pool]).collect()
+    rows = query_helper.run_query(
+        session,
+        "show compute pools like ?",
+        params=[compute_pool],
+    )
     if not rows:
         raise ValueError(f"Compute pool '{compute_pool}' not found")
     instance_family: str = rows[0]["instance_family"]
@@ -180,7 +184,7 @@ def generate_service_spec(
         constants.PAYLOAD_DIR_ENV_VAR: stage_mount.as_posix(),
         constants.RESULT_PATH_ENV_VAR: constants.RESULT_PATH_DEFAULT_VALUE,
     }
-    endpoints = []
+    endpoints: list[dict[str, Any]] = []
 
     if target_instances > 1:
         # Update environment variables for multi-node job
@@ -189,7 +193,7 @@ def generate_service_spec(
         env_vars[constants.MIN_INSTANCES_ENV_VAR] = str(min_instances)
 
         # Define Ray endpoints for intra-service instance communication
-        ray_endpoints = [
+        ray_endpoints: list[dict[str, Any]] = [
             {"name": "ray-client-server-endpoint", "port": 10001, "protocol": "TCP"},
             {"name": "ray-gcs-endpoint", "port": 12001, "protocol": "TCP"},
             {"name": "ray-dashboard-grpc-endpoint", "port": 12002, "protocol": "TCP"},
@@ -232,6 +236,19 @@ def generate_service_spec(
         ],
         "volumes": volumes,
     }
+
+    if target_instances > 1:
+        spec_dict.update(
+            {
+                "resourceManagement": {
+                    "controlPolicy": {
+                        "startupOrder": {
+                            "type": "FirstInstance",
+                        },
+                    },
+                },
+            }
+        )
     if endpoints:
         spec_dict["endpoints"] = endpoints
     if metrics: