snowflake-ml-python 1.8.5__tar.gz → 1.9.0__tar.gz

{snowflake_ml_python-1.8.5 → snowflake_ml_python-1.9.0}/CHANGELOG.md

@@ -1,5 +1,48 @@
 # Release History
 
+## 1.9.0
+
+### Bug Fixes
+
+- Registry: Fixed bug causing snowpark to pandas dataframe conversion to fail when `QUOTED_IDENTIFIERS_IGNORE_CASE`
+  parameter is enabled
+- Registry: Fixed duplicate UserWarning logs during model packaging
+
+### Behavior Changes
+
+- ML Job: The `list_jobs()` API has been modified. The `scope` parameter has been removed,
+  optional `database` and `schema` parameters have been added, the return type has changed
+  from `snowpark.DataFrame` to `pandas.DataFrame`, and the returned columns have been updated
+  to `name`, `status`, `message`, `database_name`, `schema_name`, `owner`, `compute_pool`,
+  `target_instances`, `created_time`, and `completed_time`.
+- Registry: Set `relax_version` to false when pip_requirements are specified while logging model
+- Registry: UserWarning will now be raised based on specified target_platforms (addresses spurious warnings)
+
+### New Features
+
+- Registry: `target_platforms` supports `TargetPlatformMode`: `WAREHOUSE_ONLY`, `SNOWPARK_CONTAINER_SERVICES_ONLY`,
+  or `BOTH_WAREHOUSE_AND_SNOWPARK_CONTAINER_SERVICES`.
+- Registry: Introduce `snowflake.ml.model.target_platform.TargetPlatform`, target platform constants, and
+  `snowflake.ml.model.task.Task`.
+- ML Job: Single-node ML Jobs are now in GA. Multi-node support is now in PuPr
+  - Moved less frequently used job submission parameters to `**kwargs`
+  - Platform metrics are now enabled by default
+  - `list_jobs()` behavior changed, see [Behavior Changes](#behavior-changes) for more info
+
+## 1.8.6
+
+### Bug Fixes
+
+- Fixed fatal errors from internal telemetry wrappers.
+
+### New Features
+
+- Registry: Add service container info to logs.
+- ML Job (PuPr): Add new `submit_from_stage()` API for submitting a payload from an existing stage path.
+- ML Job (PuPr): Add support for `snowpark.Session` objects in the argument list of
+  `@remote` decorated functions. `Session` object will be injected from context in
+  the job execution environment.
+
 ## 1.8.5
 
 ### Bug Fixes
@@ -9,16 +52,21 @@
 - Explainability: bump minimum streamlit version down to 1.30
 - Modeling: Make XGBoost a required dependency (xgboost is not a required dependency in snowflake-ml-python 1.8.4).
 
-### Breaking change
+### Behavior Changes
 
-- ML Job: Rename argument `num_instances` to `target_instances` in job submission APIs and
+- ML Job (Multi-node PrPr): Rename argument `num_instances` to `target_instances` in job submission APIs and
   change type from `Optional[int]` to `int`
 
 ### New Features
 
 - Registry: No longer checks if the snowflake-ml-python version is available in the Snowflake Conda channel when logging
   an SPCS-only model.
-- ML Job: Add `min_instances` argument to the job decorator to allow waiting for workers to be ready.
+- ML Job (PuPr): Add `min_instances` argument to the job decorator to allow waiting for workers to be ready.
+- ML Job (PuPr): Adjust polling behavior to reduce number of SQL calls.
+
+### Deprecations
+
+- `SnowflakeLoginOptions` is deprecated and will be removed in a future release.
 
 ## 1.8.4 (2025-05-12)
 
@@ -29,32 +77,29 @@
 - Registry: Fixed a bug when logging pytroch and tensorflow models that caused
   `UnboundLocalError: local variable 'multiple_inputs' referenced before assignment`.
 
-### Breaking change
+### Behavior Changes
 
-- ML Job: Updated property `id` to be fully qualified name; Introduced new property `name` to represent the ML Job name
-- ML Job: Modified `list_jobs()` to return ML Job `name` instead of `id`
+- ML Job (PuPr) Updated property `id` to be fully qualified name; Introduced new property `name`
+  to represent the ML Job name
+- ML Job (PuPr) Modified `list_jobs()` to return ML Job `name` instead of `id`
 - Registry: Error in `log_model` if `enable_explainability` is True and model is only deployed to
    Snowpark Container Services, instead of just user warning.
 
 ### New Features
 
-- ML Job: Extend `@remote` function decorator, `submit_file()` and `submit_directory()` to accept `database` and
+- ML Job (PuPr): Extend `@remote` function decorator, `submit_file()` and `submit_directory()` to accept `database` and
   `schema` parameters
-- ML Job: Support querying by fully qualified name in `get_job()`
+- ML Job (PuPr): Support querying by fully qualified name in `get_job()`
 - Explainability: Added visualization functions to `snowflake.ml.monitoring` to plot explanations in notebooks.
 - Explainability: Support explain for categorical transforms for sklearn pipeline
 - Support categorical type for `xgboost.DMatrix` inputs.
 
 ## 1.8.3
 
-### Bug Fixes
-
-### Behavior Change
-
 ### New Features
 
 - Registry: Default to the runtime cuda version if available when logging a GPU model in Container Runtime.
-- ML Job: Added `as_list` argument to `MLJob.get_logs()` to enable retrieving logs
+- ML Job (PuPr): Added `as_list` argument to `MLJob.get_logs()` to enable retrieving logs
   as a list of strings
 - Registry: Support `ModelVersion.run_job` to run inference with a single-node Snowpark Container Services job.
 - DataConnector: Removed PrPr decorators
@@ -65,11 +110,11 @@
 ### New Features
 
 - ML Job now available as a PuPr feature
-- ML Job: Add ability to retrieve results for `@remote` decorated functions using
-  new `MLJobWithResult.result()` API, which will return the unpickled result
-  or raise an exception if the job execution failed.
-- ML Job: Pre-created Snowpark Session is now available inside job payloads using
-  `snowflake.snowpark.context.get_active_session()`
+  - Add ability to retrieve results for `@remote` decorated functions using
+    new `MLJobWithResult.result()` API, which will return the unpickled result
+    or raise an exception if the job execution failed.
+  - Pre-created Snowpark Session is now available inside job payloads using
+    `snowflake.snowpark.context.get_active_session()`
 - Registry: Introducing `save_location` to `log_model` using the `options` argument.
   User's can provide the path to write the model version's files that get stored in Snowflake's stage.
 

{snowflake_ml_python-1.8.5/snowflake_ml_python.egg-info → snowflake_ml_python-1.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.8.5
+Version: 1.9.0
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -408,6 +408,49 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
+## 1.9.0
+
+### Bug Fixes
+
+- Registry: Fixed bug causing snowpark to pandas dataframe conversion to fail when `QUOTED_IDENTIFIERS_IGNORE_CASE`
+  parameter is enabled
+- Registry: Fixed duplicate UserWarning logs during model packaging
+
+### Behavior Changes
+
+- ML Job: The `list_jobs()` API has been modified. The `scope` parameter has been removed,
+  optional `database` and `schema` parameters have been added, the return type has changed
+  from `snowpark.DataFrame` to `pandas.DataFrame`, and the returned columns have been updated
+  to `name`, `status`, `message`, `database_name`, `schema_name`, `owner`, `compute_pool`,
+  `target_instances`, `created_time`, and `completed_time`.
+- Registry: Set `relax_version` to false when pip_requirements are specified while logging model
+- Registry: UserWarning will now be raised based on specified target_platforms (addresses spurious warnings)
+
+### New Features
+
+- Registry: `target_platforms` supports `TargetPlatformMode`: `WAREHOUSE_ONLY`, `SNOWPARK_CONTAINER_SERVICES_ONLY`,
+  or `BOTH_WAREHOUSE_AND_SNOWPARK_CONTAINER_SERVICES`.
+- Registry: Introduce `snowflake.ml.model.target_platform.TargetPlatform`, target platform constants, and
+  `snowflake.ml.model.task.Task`.
+- ML Job: Single-node ML Jobs are now in GA. Multi-node support is now in PuPr
+  - Moved less frequently used job submission parameters to `**kwargs`
+  - Platform metrics are now enabled by default
+  - `list_jobs()` behavior changed, see [Behavior Changes](#behavior-changes) for more info
+
+## 1.8.6
+
+### Bug Fixes
+
+- Fixed fatal errors from internal telemetry wrappers.
+
+### New Features
+
+- Registry: Add service container info to logs.
+- ML Job (PuPr): Add new `submit_from_stage()` API for submitting a payload from an existing stage path.
+- ML Job (PuPr): Add support for `snowpark.Session` objects in the argument list of
+  `@remote` decorated functions. `Session` object will be injected from context in
+  the job execution environment.
+
 ## 1.8.5
 
 ### Bug Fixes
@@ -417,16 +460,21 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 - Explainability: bump minimum streamlit version down to 1.30
 - Modeling: Make XGBoost a required dependency (xgboost is not a required dependency in snowflake-ml-python 1.8.4).
 
-### Breaking change
+### Behavior Changes
 
-- ML Job: Rename argument `num_instances` to `target_instances` in job submission APIs and
+- ML Job (Multi-node PrPr): Rename argument `num_instances` to `target_instances` in job submission APIs and
   change type from `Optional[int]` to `int`
 
 ### New Features
 
 - Registry: No longer checks if the snowflake-ml-python version is available in the Snowflake Conda channel when logging
   an SPCS-only model.
-- ML Job: Add `min_instances` argument to the job decorator to allow waiting for workers to be ready.
+- ML Job (PuPr): Add `min_instances` argument to the job decorator to allow waiting for workers to be ready.
+- ML Job (PuPr): Adjust polling behavior to reduce number of SQL calls.
+
+### Deprecations
+
+- `SnowflakeLoginOptions` is deprecated and will be removed in a future release.
 
 ## 1.8.4 (2025-05-12)
 
@@ -437,32 +485,29 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 - Registry: Fixed a bug when logging pytroch and tensorflow models that caused
   `UnboundLocalError: local variable 'multiple_inputs' referenced before assignment`.
 
-### Breaking change
+### Behavior Changes
 
-- ML Job: Updated property `id` to be fully qualified name; Introduced new property `name` to represent the ML Job name
-- ML Job: Modified `list_jobs()` to return ML Job `name` instead of `id`
+- ML Job (PuPr) Updated property `id` to be fully qualified name; Introduced new property `name`
+  to represent the ML Job name
+- ML Job (PuPr) Modified `list_jobs()` to return ML Job `name` instead of `id`
 - Registry: Error in `log_model` if `enable_explainability` is True and model is only deployed to
    Snowpark Container Services, instead of just user warning.
 
 ### New Features
 
-- ML Job: Extend `@remote` function decorator, `submit_file()` and `submit_directory()` to accept `database` and
+- ML Job (PuPr): Extend `@remote` function decorator, `submit_file()` and `submit_directory()` to accept `database` and
   `schema` parameters
-- ML Job: Support querying by fully qualified name in `get_job()`
+- ML Job (PuPr): Support querying by fully qualified name in `get_job()`
 - Explainability: Added visualization functions to `snowflake.ml.monitoring` to plot explanations in notebooks.
 - Explainability: Support explain for categorical transforms for sklearn pipeline
 - Support categorical type for `xgboost.DMatrix` inputs.
 
 ## 1.8.3
 
-### Bug Fixes
-
-### Behavior Change
-
 ### New Features
 
 - Registry: Default to the runtime cuda version if available when logging a GPU model in Container Runtime.
-- ML Job: Added `as_list` argument to `MLJob.get_logs()` to enable retrieving logs
+- ML Job (PuPr): Added `as_list` argument to `MLJob.get_logs()` to enable retrieving logs
   as a list of strings
 - Registry: Support `ModelVersion.run_job` to run inference with a single-node Snowpark Container Services job.
 - DataConnector: Removed PrPr decorators
@@ -473,11 +518,11 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 ### New Features
 
 - ML Job now available as a PuPr feature
-- ML Job: Add ability to retrieve results for `@remote` decorated functions using
-  new `MLJobWithResult.result()` API, which will return the unpickled result
-  or raise an exception if the job execution failed.
-- ML Job: Pre-created Snowpark Session is now available inside job payloads using
-  `snowflake.snowpark.context.get_active_session()`
+  - Add ability to retrieve results for `@remote` decorated functions using
+    new `MLJobWithResult.result()` API, which will return the unpickled result
+    or raise an exception if the job execution failed.
+  - Pre-created Snowpark Session is now available inside job payloads using
+    `snowflake.snowpark.context.get_active_session()`
 - Registry: Introducing `save_location` to `log_model` using the `options` argument.
   User's can provide the path to write the model version's files that get stored in Snowflake's stage.
 

{snowflake_ml_python-1.8.5 → snowflake_ml_python-1.9.0}/snowflake/ml/_internal/telemetry.py

@@ -411,16 +411,13 @@ def send_custom_usage(
     **kwargs: Any,
 ) -> None:
     conn = _get_snowflake_connection()
-    if conn is None:
-        raise ValueError(
-            """Snowflake connection is required to send custom telemetry. This means there
-            must be at least one active session, or that telemetry is being sent from within an SPCS service."""
-        )
 
-    client = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject)
-    common_metrics = client._create_basic_telemetry_data(telemetry_type=telemetry_type)
-    data = {**common_metrics, TelemetryField.KEY_DATA.value: data, **kwargs}
-    client._send(msg=data)
+    # Send telemetry if Snowflake connection is available.
+    if conn is not None:
+        client = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject)
+        common_metrics = client._create_basic_telemetry_data(telemetry_type=telemetry_type)
+        data = {**common_metrics, TelemetryField.KEY_DATA.value: data, **kwargs}
+        client._send(msg=data)
 
 
 def send_api_usage_telemetry(

snowflake_ml_python-1.9.0/snowflake/ml/_internal/utils/connection_params.py

@@ -0,0 +1,196 @@
+import configparser
+import os
+from typing import Optional, Union
+
+from absl import logging
+from cryptography.hazmat import backends
+from cryptography.hazmat.primitives import serialization
+
+_DEFAULT_CONNECTION_FILE = "~/.snowsql/config"
+
+
+def _read_token(token_file: str = "") -> str:
+    """
+    Reads token from environment or file provided.
+
+    First tries to read the token from environment variable
+    (`SNOWFLAKE_TOKEN`) followed by the token file.
+    Both the options are tried out in SnowServices.
+
+    Args:
+        token_file: File from which token needs to be read. Optional.
+
+    Returns:
+        the token.
+    """
+    token = os.getenv("SNOWFLAKE_TOKEN", "")
+    if token:
+        return token
+    if token_file and os.path.exists(token_file):
+        with open(token_file) as f:
+            token = f.read()
+    return token
+
+
+_ENCRYPTED_PKCS8_PK_HEADER = b"-----BEGIN ENCRYPTED PRIVATE KEY-----"
+_UNENCRYPTED_PKCS8_PK_HEADER = b"-----BEGIN PRIVATE KEY-----"
+
+
+def _load_pem_to_der(private_key_path: str) -> bytes:
+    """Given a private key file path (in PEM format), decode key data into DER format."""
+    with open(private_key_path, "rb") as f:
+        private_key_pem = f.read()
+    private_key_passphrase: Optional[str] = os.getenv("SNOWFLAKE_PRIVATE_KEY_PASSPHRASE", None)
+
+    # Only PKCS#8 format key will be accepted. However, openssl
+    # transparently handle PKCS#8 and PKCS#1 format (by some fallback
+    # logic) and their is no function to distinguish between them. By
+    # reading openssl source code, apparently they also relies on header
+    # to determine if give bytes is PKCS#8 format or not
+    if not private_key_pem.startswith(_ENCRYPTED_PKCS8_PK_HEADER) and not private_key_pem.startswith(
+        _UNENCRYPTED_PKCS8_PK_HEADER
+    ):
+        raise Exception("Private key provided is not in PKCS#8 format. Please use correct format.")
+
+    if private_key_pem.startswith(_ENCRYPTED_PKCS8_PK_HEADER) and private_key_passphrase is None:
+        raise Exception(
+            "Private key is encrypted but passphrase could not be found. "
+            "Please set SNOWFLAKE_PRIVATE_KEY_PASSPHRASE env variable."
+        )
+
+    if private_key_pem.startswith(_UNENCRYPTED_PKCS8_PK_HEADER):
+        private_key_passphrase = None
+
+    private_key = serialization.load_pem_private_key(
+        private_key_pem,
+        str.encode(private_key_passphrase) if private_key_passphrase is not None else private_key_passphrase,
+        backends.default_backend(),
+    )
+
+    return private_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+
+def _connection_properties_from_env() -> dict[str, str]:
+    """Returns a dict with all possible login related env variables."""
+    sf_conn_prop = {
+        # Mandatory fields
+        "account": os.environ["SNOWFLAKE_ACCOUNT"],
+        "database": os.environ["SNOWFLAKE_DATABASE"],
+        # With a default value
+        "token_file": os.getenv("SNOWFLAKE_TOKEN_FILE", "/snowflake/session/token"),
+        "ssl": os.getenv("SNOWFLAKE_SSL", "on"),
+        "protocol": os.getenv("SNOWFLAKE_PROTOCOL", "https"),
+    }
+    # With empty default value
+    for key, env_var in {
+        "user": "SNOWFLAKE_USER",
+        "authenticator": "SNOWFLAKE_AUTHENTICATOR",
+        "password": "SNOWFLAKE_PASSWORD",
+        "host": "SNOWFLAKE_HOST",
+        "port": "SNOWFLAKE_PORT",
+        "schema": "SNOWFLAKE_SCHEMA",
+        "warehouse": "SNOWFLAKE_WAREHOUSE",
+        "private_key_path": "SNOWFLAKE_PRIVATE_KEY_PATH",
+    }.items():
+        value = os.getenv(env_var, "")
+        if value:
+            sf_conn_prop[key] = value
+    return sf_conn_prop
+
+
+def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -> dict[str, str]:
+    """Loads the dictionary from snowsql config file."""
+    snowsql_config_file = login_file if login_file else os.path.expanduser(_DEFAULT_CONNECTION_FILE)
+    if not os.path.exists(snowsql_config_file):
+        logging.error(f"Connection name given but snowsql config file is not found at: {snowsql_config_file}")
+        raise Exception("Snowflake SnowSQL config not found.")
+
+    config = configparser.ConfigParser(inline_comment_prefixes="#")
+
+    snowflake_connection_name = os.getenv("SNOWFLAKE_CONNECTION_NAME")
+    if snowflake_connection_name is not None:
+        connection_name = snowflake_connection_name
+
+    if connection_name:
+        if not connection_name.startswith("connections."):
+            connection_name = "connections." + connection_name
+    else:
+        # See https://docs.snowflake.com/en/user-guide/snowsql-start.html#configuring-default-connection-settings
+        connection_name = "connections"
+
+    logging.info(f"Reading {snowsql_config_file} for connection parameters defined as {connection_name}")
+    config.read(snowsql_config_file)
+    conn_params = dict(config[connection_name])
+    # Remap names to appropriate args in Python Connector API
+    # Note: "dbname" should become "database"
+    conn_params = {k.replace("name", ""): v.strip('"') for k, v in conn_params.items()}
+    if "db" in conn_params:
+        conn_params["database"] = conn_params["db"]
+        del conn_params["db"]
+    return conn_params
+
+
+def SnowflakeLoginOptions(connection_name: str = "", login_file: Optional[str] = None) -> dict[str, Union[str, bytes]]:
+    """Returns a dict that can be used directly into snowflake python connector or Snowpark session config.
+
+    NOTE: Token/Auth information is sideloaded in all cases above, if provided in following order:
+      1. If SNOWFLAKE_TOKEN is defined in the environment, it will be used.
+      2. If SNOWFLAKE_TOKEN_FILE is defined in the environment and file matching the value found, content of the file
+         will be used.
+
+    If token is found, username, password will be reset and 'authenticator' will be set to 'oauth'.
+
+    Python Connector:
+    >> ctx = snowflake.connector.connect(**(SnowflakeLoginOptions()))
+
+    Snowpark Session:
+    >> session = Session.builder.configs(SnowflakeLoginOptions()).create()
+
+    Usage Note:
+      Ideally one should have a snowsql config file. Read more here:
+      https://docs.snowflake.com/en/user-guide/snowsql-start.html#configuring-default-connection-settings
+
+      If snowsql config file does not exist, it tries auth from env variables.
+
+    Args:
+        connection_name: Name of the connection to look for inside the config file. If environment variable
+            SNOWFLAKE_CONNECTION_NAME is provided, it will override the input connection_name.
+        login_file: If provided, this is used as config file instead of default one (_DEFAULT_CONNECTION_FILE).
+
+    Returns:
+        A dict with connection parameters.
+
+    Raises:
+        Exception: if none of config file and environment variable are present.
+    """
+    conn_prop: dict[str, Union[str, bytes]] = {}
+    login_file = login_file or os.path.expanduser(_DEFAULT_CONNECTION_FILE)
+    # If login file exists, use this exclusively.
+    if os.path.exists(login_file):
+        conn_prop = {**(_load_from_snowsql_config_file(connection_name, login_file))}
+    else:
+        # If environment exists for SNOWFLAKE_ACCOUNT, assume everything
+        # comes from environment. Mixing it not allowed.
+        account = os.getenv("SNOWFLAKE_ACCOUNT", "")
+        if account:
+            conn_prop = {**_connection_properties_from_env()}
+        else:
+            raise Exception("Snowflake credential is neither set in env nor a login file was provided.")
+
+    # Token, if specified, is always side-loaded in all cases.
+    token = _read_token(str(conn_prop["token_file"]) if "token_file" in conn_prop else "")
+    if token:
+        conn_prop["token"] = token
+        if "authenticator" not in conn_prop or conn_prop["authenticator"]:
+            conn_prop["authenticator"] = "oauth"
+    elif "private_key_path" in conn_prop and "private_key" not in conn_prop:
+        conn_prop["private_key"] = _load_pem_to_der(str(conn_prop["private_key_path"]))
+
+    if "ssl" in conn_prop and conn_prop["ssl"].lower() == "off":
+        conn_prop["protocol"] = "http"
+
+    return conn_prop

{snowflake_ml_python-1.8.5 → snowflake_ml_python-1.9.0}/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_python-1.9.0/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_python-1.8.5 → snowflake_ml_python-1.9.0}/snowflake/ml/jobs/__init__.py

@@ -7,6 +7,7 @@ from snowflake.ml.jobs.manager import (
     list_jobs,
     submit_directory,
     submit_file,
+    submit_from_stage,
 )
 
 __all__ = [
@@ -18,4 +19,5 @@ __all__ = [
     "delete_job",
     "MLJob",
     "JOB_STATUS",
+    "submit_from_stage",
 ]

{snowflake_ml_python-1.8.5 → snowflake_ml_python-1.9.0}/snowflake/ml/jobs/_utils/constants.py

@@ -6,6 +6,7 @@ DEFAULT_CONTAINER_NAME = "main"
 PAYLOAD_DIR_ENV_VAR = "MLRS_PAYLOAD_DIR"
 RESULT_PATH_ENV_VAR = "MLRS_RESULT_PATH"
 MIN_INSTANCES_ENV_VAR = "MLRS_MIN_INSTANCES"
+RUNTIME_IMAGE_TAG_ENV_VAR = "MLRS_CONTAINER_IMAGE_TAG"
 MEMORY_VOLUME_NAME = "dshm"
 STAGE_VOLUME_NAME = "stage-volume"
 STAGE_VOLUME_MOUNT_PATH = "/mnt/app"
@@ -14,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.2.3"
+DEFAULT_IMAGE_TAG = "1.5.0"
 DEFAULT_ENTRYPOINT_PATH = "func.py"
 
 # Percent of container memory to allocate for /dev/shm volume
@@ -43,7 +44,7 @@ ENABLE_HEALTH_CHECKS = "false"
 
 # Job status polling constants
 JOB_POLL_INITIAL_DELAY_SECONDS = 0.1
-JOB_POLL_MAX_DELAY_SECONDS = 1
+JOB_POLL_MAX_DELAY_SECONDS = 30
 
 # Magic attributes
 IS_MLJOB_REMOTE_ATTR = "_is_mljob_remote_callable"

snowflake_ml_python-1.9.0/snowflake/ml/jobs/_utils/function_payload_utils.py

@@ -0,0 +1,43 @@
+import inspect
+from typing import Any, Callable, Optional
+
+from snowflake import snowpark
+from snowflake.snowpark import context as sp_context
+
+
+class FunctionPayload:
+    def __init__(
+        self,
+        func: Callable[..., Any],
+        session: Optional[snowpark.Session] = None,
+        session_argument: str = "",
+        *args: Any,
+        **kwargs: Any
+    ) -> None:
+        self.function = func
+        self.args = args
+        self.kwargs = kwargs
+        self._session = session
+        self._session_argument = session_argument
+
+    @property
+    def session(self) -> Optional[snowpark.Session]:
+        return self._session
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Customize pickling to exclude session."""
+        state = self.__dict__.copy()
+        state["_session"] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore session from context during unpickling."""
+        self.__dict__.update(state)
+        self._session = sp_context.get_active_session()
+
+    def __call__(self) -> Any:
+        sig = inspect.signature(self.function)
+        bound = sig.bind_partial(*self.args, **self.kwargs)
+        bound.arguments[self._session_argument] = self._session
+
+        return self.function(*bound.args, **bound.kwargs)