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

snowflake/ml/_internal/telemetry.py

@@ -4,6 +4,7 @@ import enum
 import functools
 import inspect
 import operator
+import os
 import sys
 import time
 import traceback
@@ -13,7 +14,7 @@ from typing import Any, Callable, Iterable, Mapping, Optional, TypeVar, Union, c
 from typing_extensions import ParamSpec
 
 from snowflake import connector
-from snowflake.connector import telemetry as connector_telemetry, time_util
+from snowflake.connector import connect, telemetry as connector_telemetry, time_util
 from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env
 from snowflake.ml._internal.exceptions import (
@@ -37,6 +38,37 @@ _Args = ParamSpec("_Args")
 _ReturnValue = TypeVar("_ReturnValue")
 
 
+def _get_login_token() -> Union[str, bytes]:
+    with open("/snowflake/session/token") as f:
+        return f.read()
+
+
+def _get_snowflake_connection() -> Optional[connector.SnowflakeConnection]:
+    conn = None
+    if os.getenv("SNOWFLAKE_HOST") is not None and os.getenv("SNOWFLAKE_ACCOUNT") is not None:
+        try:
+            conn = connect(
+                host=os.getenv("SNOWFLAKE_HOST"),
+                account=os.getenv("SNOWFLAKE_ACCOUNT"),
+                token=_get_login_token(),
+                authenticator="oauth",
+            )
+        except Exception:
+            # Failed to get a new SnowflakeConnection in SPCS. Fall back to using the active session.
+            # This will work in some cases once SPCS enables multiple authentication modes, and users select any auth.
+            pass
+
+    if conn is None:
+        try:
+            active_session = next(iter(session._get_active_sessions()))
+            conn = active_session._conn._conn if active_session.telemetry_enabled else None
+        except snowpark_exceptions.SnowparkSessionException:
+            # Failed to get an active session. No connection available.
+            pass
+
+    return conn
+
+
 @enum.unique
 class TelemetryProject(enum.Enum):
     MLOPS = "MLOps"
@@ -378,13 +410,14 @@ def send_custom_usage(
     data: Optional[dict[str, Any]] = None,
     **kwargs: Any,
 ) -> None:
-    active_session = next(iter(session._get_active_sessions()))
-    assert active_session, "Missing active session object"
+    conn = _get_snowflake_connection()
 
-    client = _SourceTelemetryClient(conn=active_session._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(
@@ -501,7 +534,6 @@ def send_api_usage_telemetry(
                 return update_stmt_params_if_snowpark_df(result, statement_params)
 
             # prioritize `conn_attr_name` over the active session
-            telemetry_enabled = True
             if conn_attr_name:
                 # raise AttributeError if conn attribute does not exist in `self`
                 conn = operator.attrgetter(conn_attr_name)(args[0])
@@ -509,16 +541,10 @@ def send_api_usage_telemetry(
                     raise TypeError(
                         f"Expected a conn object of type {' or '.join(_CONNECTION_TYPES.keys())} but got {type(conn)}"
                     )
-            # get an active session
             else:
-                try:
-                    active_session = next(iter(session._get_active_sessions()))
-                    conn = active_session._conn._conn
-                    telemetry_enabled = active_session.telemetry_enabled
-                except snowpark_exceptions.SnowparkSessionException:
-                    conn = None
+                conn = _get_snowflake_connection()
 
-            if conn is None or not telemetry_enabled:
+            if conn is None:
                 # Telemetry not enabled, just execute without our additional telemetry logic
                 try:
                     return ctx.run(execute_func_with_statement_params)

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/data/data_connector.py

@@ -249,7 +249,7 @@ class DataConnector:
 
 # Switch to use Runtime's Data Ingester if running in ML runtime
 # Fail silently if the data ingester is not found
-if env.IN_ML_RUNTIME and os.getenv(env.USE_OPTIMIZED_DATA_INGESTOR):
+if env.IN_ML_RUNTIME and os.getenv(env.USE_OPTIMIZED_DATA_INGESTOR, "").lower() in ("true", "1"):
     try:
         from runtime_external_entities import get_ingester_class
 

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/jobs/_utils/constants.py

@@ -5,6 +5,8 @@ from snowflake.ml.jobs._utils.types import ComputeResources
 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"
@@ -13,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.4.2"
 DEFAULT_ENTRYPOINT_PATH = "func.py"
 
 # Percent of container memory to allocate for /dev/shm volume
@@ -37,16 +39,24 @@ RAY_PORTS = {
 # Node health check configuration
 # TODO(SNOW-1937020): Revisit the health check configuration
 ML_RUNTIME_HEALTH_CHECK_PORT = "5001"
+ENABLE_HEALTH_CHECKS_ENV_VAR = "ENABLE_HEALTH_CHECKS"
 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"
 RESULT_PATH_DEFAULT_VALUE = "mljob_result.pkl"
 
+# Log start and end messages
+LOG_START_MSG = "--------------------------------\nML job started\n--------------------------------"
+LOG_END_MSG = "--------------------------------\nML job finished\n--------------------------------"
+
+# Default setting for verbose logging in get_log function
+DEFAULT_VERBOSE_LOG = False
+
 # Compute pool resource information
 # TODO: Query Snowflake for resource information instead of relying on this hardcoded
 #       table from https://docs.snowflake.com/en/sql-reference/sql/create-compute-pool

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)

snowflake/ml/jobs/_utils/interop_utils.py

@@ -80,7 +80,7 @@ def fetch_result(session: snowpark.Session, result_path: str) -> ExecutionResult
         # 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, TypeError, pickle.UnpicklingError):
+    except (sp_exceptions.SnowparkSQLException, pickle.UnpicklingError, TypeError, ImportError):
         # 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: