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

snowflake/ml/_internal/utils/identifier.py

@@ -240,7 +240,7 @@ def get_schema_level_object_identifier(
     """
 
     for identifier in (db, schema, object_name):
-        if identifier is not None and SF_IDENTIFIER_RE.match(identifier) is None:
+        if identifier is not None and SF_IDENTIFIER_RE.fullmatch(identifier) is None:
             raise ValueError(f"Invalid identifier {identifier}")
 
     if others is None:

snowflake/ml/_internal/utils/mixins.py

@@ -0,0 +1,61 @@
+from typing import Any, Optional
+
+from snowflake.ml._internal.utils import identifier
+from snowflake.snowpark import session
+
+
+class SerializableSessionMixin:
+    """Mixin that provides pickling capabilities for objects with Snowpark sessions."""
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Customize pickling to exclude non-serializable session and related components."""
+        state = self.__dict__.copy()
+
+        # Save session metadata for validation during unpickling
+        if hasattr(self, "_session") and self._session is not None:
+            try:
+                state["__session-account__"] = self._session.get_current_account()
+                state["__session-role__"] = self._session.get_current_role()
+                state["__session-database__"] = self._session.get_current_database()
+                state["__session-schema__"] = self._session.get_current_schema()
+            except Exception:
+                pass
+
+        state["_session"] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore session from context during unpickling."""
+        saved_account = state.pop("__session-account__", None)
+        saved_role = state.pop("__session-role__", None)
+        saved_database = state.pop("__session-database__", None)
+        saved_schema = state.pop("__session-schema__", None)
+        self.__dict__.update(state)
+
+        if saved_account is not None:
+
+            def identifiers_match(saved: Optional[str], current: Optional[str]) -> bool:
+                saved_resolved = identifier.resolve_identifier(saved) if saved is not None else saved
+                current_resolved = identifier.resolve_identifier(current) if current is not None else current
+                return saved_resolved == current_resolved
+
+            for active_session in session._get_active_sessions():
+                try:
+                    current_account = active_session.get_current_account()
+                    current_role = active_session.get_current_role()
+                    current_database = active_session.get_current_database()
+                    current_schema = active_session.get_current_schema()
+
+                    if (
+                        identifiers_match(saved_account, current_account)
+                        and identifiers_match(saved_role, current_role)
+                        and identifiers_match(saved_database, current_database)
+                        and identifiers_match(saved_schema, current_schema)
+                    ):
+                        self._session = active_session
+                        return
+                except Exception:
+                    continue
+
+        # No matching session found or no metadata available
+        raise RuntimeError("No active Snowpark session available. Please create a session.")

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

@@ -12,13 +12,13 @@ import cloudpickle as cp
 from packaging import version
 
 from snowflake import snowpark
+from snowflake.connector import errors
 from snowflake.ml.jobs._utils import (
     constants,
     function_payload_utils,
     stage_utils,
     types,
 )
-from snowflake.snowpark import exceptions as sp_exceptions
 from snowflake.snowpark._internal import code_generation
 
 cp.register_pickle_by_value(function_payload_utils)
@@ -312,14 +312,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()
-        except sp_exceptions.SnowparkSQLException:
-            session.sql(
+            session._conn.run_query("describe stage identifier(?)", params=[stage_name], _force_qmark_paramstyle=True)
+        except errors.ProgrammingError:
+            session._conn.run_query(
                 "create stage if not exists identifier(?)"
                 " encryption = ( type = 'SNOWFLAKE_SSE' )"
                 " comment = 'Created by snowflake.ml.jobs Python API'",
                 params=[stage_name],
-            ).collect()
+                _force_qmark_paramstyle=True,
+            )
 
         # Upload payload to stage
         if not isinstance(source, (Path, stage_utils.StagePath)):

snowflake/ml/jobs/_utils/query_helper.py

@@ -0,0 +1,9 @@
+from snowflake import snowpark
+
+
+def get_attribute_map(session: snowpark.Session, requested_attributes: dict[str, int]) -> dict[str, int]:
+    metadata = session._conn._cursor.description
+    for index in range(len(metadata)):
+        if metadata[index].name in requested_attributes.keys():
+            requested_attributes[metadata[index].name] = index
+    return requested_attributes

snowflake/ml/jobs/_utils/spec_utils.py

@@ -6,16 +6,18 @@ 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()
-    if not rows:
+    rows = session._conn.run_query("show compute pools like ?", params=[compute_pool], _force_qmark_paramstyle=True)
+    if not rows or not isinstance(rows, dict) or not rows.get("data"):
         raise ValueError(f"Compute pool '{compute_pool}' not found")
-    instance_family: str = rows[0]["instance_family"]
+    requested_attributes = query_helper.get_attribute_map(session, {"instance_family": 4})
+    compute_pool_info = rows["data"]
+    instance_family: str = compute_pool_info[0][requested_attributes["instance_family"]]
     cloud = snowflake_env.get_current_cloud(session, default=snowflake_env.SnowflakeCloudType.AWS)
 
     return (

snowflake/ml/jobs/decorators.py

@@ -1,6 +1,6 @@
 import copy
 import functools
-from typing import Callable, Optional, TypeVar
+from typing import Any, Callable, Optional, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -20,16 +20,11 @@ def remote(
     compute_pool: str,
     *,
     stage_name: str,
+    target_instances: int = 1,
     pip_requirements: Optional[list[str]] = None,
     external_access_integrations: Optional[list[str]] = None,
-    query_warehouse: Optional[str] = None,
-    env_vars: Optional[dict[str, str]] = None,
-    target_instances: int = 1,
-    min_instances: Optional[int] = None,
-    enable_metrics: bool = False,
-    database: Optional[str] = None,
-    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
+    **kwargs: Any,
 ) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, jb.MLJob[_ReturnValue]]]:
     """
     Submit a job to the compute pool.
@@ -37,17 +32,20 @@ def remote(
     Args:
         compute_pool: The compute pool to use for the job.
         stage_name: The name of the stage where the job payload will be uploaded.
+        target_instances: The number of nodes in the job. If none specified, create a single node job.
         pip_requirements: A list of pip requirements for the job.
         external_access_integrations: A list of external access integrations.
-        query_warehouse: The query warehouse to use. Defaults to session warehouse.
-        env_vars: Environment variables to set in container
-        target_instances: The number of nodes in the job. If none specified, create a single node job.
-        min_instances: The minimum number of nodes required to start the job. If none specified,
-            defaults to target_instances. If set, the job will not start until the minimum number of nodes is available.
-        enable_metrics: Whether to enable metrics publishing for the job.
-        database: The database to use for the job.
-        schema: The schema to use for the job.
         session: The Snowpark session to use. If none specified, uses active session.
+        kwargs: Additional keyword arguments. Supported arguments:
+            database (str): The database to use for the job.
+            schema (str): The schema to use for the job.
+            min_instances (int): The minimum number of nodes required to start the job.
+                If none specified, defaults to target_instances. If set, the job
+                will not start until the minimum number of nodes is available.
+            env_vars (dict): Environment variables to set in container.
+            enable_metrics (bool): Whether to enable metrics publishing for the job.
+            query_warehouse (str): The query warehouse to use. Defaults to session warehouse.
+            spec_overrides (dict): A dictionary of overrides for the service spec.
 
     Returns:
         Decorator that dispatches invocations of the decorated function as remote jobs.
@@ -61,22 +59,17 @@ def remote(
         wrapped_func.__code__ = wrapped_func.__code__.replace(co_firstlineno=func.__code__.co_firstlineno + 1)
 
         @functools.wraps(func)
-        def wrapper(*args: _Args.args, **kwargs: _Args.kwargs) -> jb.MLJob[_ReturnValue]:
-            payload = payload_utils.create_function_payload(func, *args, **kwargs)
+        def wrapper(*_args: _Args.args, **_kwargs: _Args.kwargs) -> jb.MLJob[_ReturnValue]:
+            payload = payload_utils.create_function_payload(func, *_args, **_kwargs)
             job = jm._submit_job(
                 source=payload,
                 stage_name=stage_name,
                 compute_pool=compute_pool,
+                target_instances=target_instances,
                 pip_requirements=pip_requirements,
                 external_access_integrations=external_access_integrations,
-                query_warehouse=query_warehouse,
-                env_vars=env_vars,
-                target_instances=target_instances,
-                min_instances=min_instances,
-                enable_metrics=enable_metrics,
-                database=database,
-                schema=schema,
                 session=payload.session or session,
+                **kwargs,
             )
             assert isinstance(job, jb.MLJob), f"Unexpected job type: {type(job)}"
             return job