snowflake-ml-python 1.7.3__py3-none-any.whl → 1.7.4__py3-none-any.whl

snowflake/cortex/_complete.py

@@ -49,6 +49,10 @@ class CompleteOptions(TypedDict):
     generally used as an alternative to temperature. The difference is that top_p restricts the set of possible tokens
     that the model outputs, while temperature influences which tokens are chosen at each step. """
 
+    guardrails: NotRequired[bool]
+    """ A boolean value that controls whether Cortex Guard filters unsafe or harmful responses
+    from the language model. """
+
 
 class ResponseParseException(Exception):
     """This exception is raised when the server response cannot be parsed."""
@@ -56,6 +60,15 @@ class ResponseParseException(Exception):
     pass
 
 
+class GuardrailsOptions(TypedDict):
+    enabled: bool
+    """A boolean value that controls whether Cortex Guard filters unsafe or harmful responses
+    from the language model."""
+
+    response_when_unsafe: str
+    """The response to return when the language model generates unsafe or harmful content."""
+
+
 _MAX_RETRY_SECONDS = 30
 
 
@@ -117,6 +130,12 @@ def _make_request_body(
             data["temperature"] = options["temperature"]
         if "top_p" in options:
             data["top_p"] = options["top_p"]
+        if "guardrails" in options and options["guardrails"]:
+            guardrails_options: GuardrailsOptions = {
+                "enabled": True,
+                "response_when_unsafe": "Response filtered by Cortex Guard",
+            }
+            data["guardrails"] = guardrails_options
     return data
 
 

snowflake/ml/_internal/platform_capabilities.py

@@ -0,0 +1,87 @@
+import json
+from typing import Any, Dict, Optional
+
+from absl import logging
+
+from snowflake.ml._internal.exceptions import error_codes, exceptions
+from snowflake.ml._internal.utils import query_result_checker
+from snowflake.snowpark import (
+    exceptions as snowpark_exceptions,
+    session as snowpark_session,
+)
+
+
+class PlatformCapabilities:
+    """Class that retrieves platform feature values for the currently running server.
+
+    Example usage:
+    ```
+    pc = PlatformCapabilities.get_instance(session)
+    if pc.is_nested_function_enabled():
+        # Nested functions are enabled.
+        print("Nested functions are enabled.")
+    else:
+        # Nested functions are disabled.
+        print("Nested functions are disabled or not supported.")
+    ```
+    """
+
+    _instance: Optional["PlatformCapabilities"] = None
+
+    @classmethod
+    def get_instance(cls, session: Optional[snowpark_session.Session] = None) -> "PlatformCapabilities":
+        if not cls._instance:
+            cls._instance = cls(session)
+        return cls._instance
+
+    def is_nested_function_enabled(self) -> bool:
+        return self._get_bool_feature("SPCS_MODEL_ENABLE_EMBEDDED_SERVICE_FUNCTIONS", False)
+
+    @staticmethod
+    def _get_features(session: snowpark_session.Session) -> Dict[str, Any]:
+        try:
+            result = (
+                query_result_checker.SqlResultValidator(
+                    session=session,
+                    query="SELECT SYSTEM$ML_PLATFORM_CAPABILITIES() AS FEATURES;",
+                )
+                .has_dimensions(expected_rows=1, expected_cols=1)
+                .has_column("FEATURES")
+                .validate()[0]
+            )
+            if "FEATURES" in result:
+                capabilities_json: str = result["FEATURES"]
+                try:
+                    parsed_json = json.loads(capabilities_json)
+                    assert isinstance(parsed_json, dict), f"Expected JSON object, got {type(parsed_json)}"
+                    return parsed_json
+                except json.JSONDecodeError as e:
+                    message = f"""Unable to parse JSON from: "{capabilities_json}"; Error="{e}"."""
+                    raise exceptions.SnowflakeMLException(
+                        error_code=error_codes.INTERNAL_SNOWML_ERROR, original_exception=RuntimeError(message)
+                    )
+        except snowpark_exceptions.SnowparkSQLException as e:
+            logging.debug(f"Failed to retrieve platform capabilities: {e}")
+            # This can happen is server side is older than 9.2. That is fine.
+        return {}
+
+    def __init__(self, session: Optional[snowpark_session.Session] = None) -> None:
+        if not session:
+            session = next(iter(snowpark_session._get_active_sessions()))
+            assert session, "Missing active session object"
+        self.features: Dict[str, Any] = PlatformCapabilities._get_features(session)
+
+    def _get_bool_feature(self, feature_name: str, default_value: bool) -> bool:
+        value = self.features.get(feature_name, default_value)
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, int) and value in [0, 1]:
+            return value == 1
+        if isinstance(value, str):
+            if value.lower() in ["true", "1"]:
+                return True
+            elif value.lower() in ["false", "0"]:
+                return False
+            else:
+                raise ValueError(f"Invalid boolean string: {value} for feature {feature_name}")
+        raise ValueError(f"Invalid boolean feature value: {value} for feature {feature_name}")

snowflake/ml/dataset/dataset.py

@@ -419,7 +419,6 @@ class Dataset(lineage_node.LineageNode):
     @telemetry.send_api_usage_telemetry(project=_PROJECT)
     def delete(self) -> None:
         """Delete Dataset and all contained versions"""
-        # TODO: Check and warn if any versions exist
         self._session.sql(f"DROP DATASET {self.fully_qualified_name}").collect(
             statement_params=_TELEMETRY_STATEMENT_PARAMS
         )

snowflake/ml/fileset/fileset.py

@@ -2,6 +2,8 @@ import functools
 import inspect
 from typing import Any, Callable, List, Optional
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.connector import connection
 from snowflake.ml._internal import telemetry
@@ -42,6 +44,10 @@ def _raise_if_deleted(func: Callable[..., Any]) -> Callable[..., Any]:
     return raise_if_deleted_helper
 
 
+@deprecated(
+    "FileSet is deprecated and will be removed in a future release."
+    " Use snowflake.ml.dataset.Dataset and snowflake.ml.data.DataConnector instead"
+)
 class FileSet:
     """A FileSet represents an immutable snapshot of the result of a query in the form of files."""
 

snowflake/ml/jobs/__init__.py

@@ -0,0 +1,21 @@
+from snowflake.ml.jobs._utils.types import JOB_STATUS
+from snowflake.ml.jobs.decorators import remote
+from snowflake.ml.jobs.job import MLJob
+from snowflake.ml.jobs.manager import (
+    delete_job,
+    get_job,
+    list_jobs,
+    submit_directory,
+    submit_file,
+)
+
+__all__ = [
+    "remote",
+    "submit_file",
+    "submit_directory",
+    "list_jobs",
+    "get_job",
+    "delete_job",
+    "MLJob",
+    "JOB_STATUS",
+]

snowflake/ml/jobs/_utils/constants.py

@@ -0,0 +1,51 @@
+from snowflake.ml._internal.utils.snowflake_env import SnowflakeCloudType
+from snowflake.ml.jobs._utils.types import ComputeResources
+
+# SPCS specification constants
+DEFAULT_CONTAINER_NAME = "main"
+PAYLOAD_DIR_ENV_VAR = "MLRS_PAYLOAD_DIR"
+
+# Default container image information
+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 = "0.8.0"
+DEFAULT_ENTRYPOINT_PATH = "func.py"
+
+# Percent of container memory to allocate for /dev/shm volume
+MEMORY_VOLUME_SIZE = 0.3
+
+# Job status polling constants
+JOB_POLL_INITIAL_DELAY_SECONDS = 0.1
+JOB_POLL_MAX_DELAY_SECONDS = 1
+
+# 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
+COMMON_INSTANCE_FAMILIES = {
+    "CPU_X64_XS": ComputeResources(cpu=1, memory=6),
+    "CPU_X64_S": ComputeResources(cpu=3, memory=13),
+    "CPU_X64_M": ComputeResources(cpu=6, memory=28),
+    "CPU_X64_L": ComputeResources(cpu=28, memory=116),
+    "HIGHMEM_X64_S": ComputeResources(cpu=6, memory=58),
+}
+AWS_INSTANCE_FAMILIES = {
+    "HIGHMEM_X64_M": ComputeResources(cpu=28, memory=240),
+    "HIGHMEM_X64_L": ComputeResources(cpu=124, memory=984),
+    "GPU_NV_S": ComputeResources(cpu=6, memory=27, gpu=1, gpu_type="A10G"),
+    "GPU_NV_M": ComputeResources(cpu=44, memory=178, gpu=4, gpu_type="A10G"),
+    "GPU_NV_L": ComputeResources(cpu=92, memory=1112, gpu=8, gpu_type="A100"),
+}
+AZURE_INSTANCE_FAMILIES = {
+    "HIGHMEM_X64_M": ComputeResources(cpu=28, memory=244),
+    "HIGHMEM_X64_L": ComputeResources(cpu=92, memory=654),
+    "GPU_NV_XS": ComputeResources(cpu=3, memory=26, gpu=1, gpu_type="T4"),
+    "GPU_NV_SM": ComputeResources(cpu=32, memory=424, gpu=1, gpu_type="A10"),
+    "GPU_NV_2M": ComputeResources(cpu=68, memory=858, gpu=2, gpu_type="A10"),
+    "GPU_NV_3M": ComputeResources(cpu=44, memory=424, gpu=2, gpu_type="A100"),
+    "GPU_NV_SL": ComputeResources(cpu=92, memory=858, gpu=4, gpu_type="A100"),
+}
+CLOUD_INSTANCE_FAMILIES = {
+    SnowflakeCloudType.AWS: AWS_INSTANCE_FAMILIES,
+    SnowflakeCloudType.AZURE: AZURE_INSTANCE_FAMILIES,
+}

snowflake/ml/jobs/_utils/payload_utils.py

@@ -0,0 +1,352 @@
+import inspect
+import io
+import sys
+import textwrap
+from pathlib import Path, PurePath
+from typing import (
+    Any,
+    Callable,
+    List,
+    Optional,
+    Type,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+)
+
+import cloudpickle as cp
+
+from snowflake import snowpark
+from snowflake.ml.jobs._utils import constants, types
+from snowflake.snowpark._internal import code_generation
+
+_SUPPORTED_ARG_TYPES = {str, int, float}
+_STARTUP_SCRIPT_PATH = PurePath("startup.sh")
+_STARTUP_SCRIPT_CODE = textwrap.dedent(
+    f"""
+    #!/bin/bash
+
+    ##### Perform common set up steps #####
+    set -e # exit if a command fails
+
+    echo "Creating log directories..."
+    mkdir -p /var/log/managedservices/user/mlrs
+    mkdir -p /var/log/managedservices/system/mlrs
+    mkdir -p /var/log/managedservices/system/ray
+
+    echo "*/1 * * * * root /etc/ray_copy_cron.sh" >> /etc/cron.d/ray_copy_cron
+    echo "" >> /etc/cron.d/ray_copy_cron
+    chmod 744 /etc/cron.d/ray_copy_cron
+
+    service cron start
+
+    mkdir -p /tmp/prometheus-multi-dir
+
+    # Change directory to user payload directory
+    if [ -n "${constants.PAYLOAD_DIR_ENV_VAR}" ]; then
+        cd ${constants.PAYLOAD_DIR_ENV_VAR}
+    fi
+
+    ##### Set up Python environment #####
+    export PYTHONPATH=/opt/env/site-packages/
+    MLRS_REQUIREMENTS_FILE=${{MLRS_REQUIREMENTS_FILE:-"requirements.txt"}}
+    if [ -f "${{MLRS_REQUIREMENTS_FILE}}" ]; then
+        # TODO: Prevent collisions with MLRS packages using virtualenvs
+        echo "Installing packages from $MLRS_REQUIREMENTS_FILE"
+        pip install -r $MLRS_REQUIREMENTS_FILE
+    fi
+
+    MLRS_CONDA_ENV_FILE=${{MLRS_CONDA_ENV_FILE:-"environment.yml"}}
+    if [ -f "${{MLRS_CONDA_ENV_FILE}}" ]; then
+        # TODO: Handle conda environment
+        echo "Custom conda environments not currently supported"
+        exit 1
+    fi
+    ##### End Python environment setup #####
+
+    ##### Ray configuration #####
+    shm_size=$(df --output=size --block-size=1 /dev/shm | tail -n 1)
+
+    # Configure IP address and logging directory
+    eth0Ip=$(ifconfig eth0 | sed -En -e 's/.*inet ([0-9.]+).*/\1/p')
+    log_dir="/tmp/ray"
+
+    # Check if eth0Ip is empty and set default if necessary
+    if [ -z "$eth0Ip" ]; then
+        # This should never happen, but just in case ethOIp is not set, we should default to localhost
+        eth0Ip="127.0.0.1"
+    fi
+
+    # Common parameters for both head and worker nodes
+    common_params=(
+        "--node-ip-address=$eth0Ip"
+        "--object-manager-port=${{RAY_OBJECT_MANAGER_PORT:-12011}}"
+        "--node-manager-port=${{RAY_NODE_MANAGER_PORT:-12012}}"
+        "--runtime-env-agent-port=${{RAY_RUNTIME_ENV_AGENT_PORT:-12013}}"
+        "--dashboard-agent-grpc-port=${{RAY_DASHBOARD_AGENT_GRPC_PORT:-12014}}"
+        "--dashboard-agent-listen-port=${{RAY_DASHBOARD_AGENT_LISTEN_PORT:-12015}}"
+        "--min-worker-port=${{RAY_MIN_WORKER_PORT:-12031}}"
+        "--max-worker-port=${{RAY_MAX_WORKER_PORT:-13000}}"
+        "--metrics-export-port=11502"
+        "--temp-dir=$log_dir"
+        "--disable-usage-stats"
+    )
+
+    # Additional head-specific parameters
+    head_params=(
+        "--head"
+        "--port=${{RAY_HEAD_GCS_PORT:-12001}}"                                  # Port of Ray (GCS server)
+        "--ray-client-server-port=${{RAY_HEAD_CLIENT_SERVER_PORT:-10001}}"      # Listening port for Ray Client Server
+        "--dashboard-host=${{NODE_IP_ADDRESS}}"                                 # Host to bind the dashboard server
+        "--dashboard-grpc-port=${{RAY_HEAD_DASHBOARD_GRPC_PORT:-12002}}"        # Dashboard head to listen for grpc on
+        "--dashboard-port=${{DASHBOARD_PORT}}"                  # Port to bind the dashboard server for local debugging
+        "--resources={{\\"node_tag:head\\":1}}"                   # Resource tag for selecting head as coordinator
+    )
+
+    # Start Ray on the head node
+    ray start "${{common_params[@]}}" "${{head_params[@]}}" &
+    ##### End Ray configuration #####
+
+    # TODO: Monitor MLRS and handle process crashes
+    python -m web.ml_runtime_grpc_server &
+
+    # TODO: Launch worker service(s) using SQL if Ray and MLRS successfully started
+
+    # Run user's Python entrypoint
+    echo Running command: python "$@"
+    python "$@"
+    """
+).strip()
+
+
+class JobPayload:
+    def __init__(
+        self,
+        source: Union[str, Path, Callable[..., Any]],
+        entrypoint: Optional[Union[str, Path]] = None,
+        *,
+        pip_requirements: Optional[List[str]] = None,
+    ) -> None:
+        self.source = Path(source) if isinstance(source, str) else source
+        self.entrypoint = Path(entrypoint) if isinstance(entrypoint, str) else entrypoint
+        self.pip_requirements = pip_requirements
+
+    def validate(self) -> None:
+        if callable(self.source):
+            # Any entrypoint value is OK for callable payloads (including None aka default)
+            # since we will generate the file from the serialized callable
+            pass
+        elif isinstance(self.source, Path):
+            # Validate self.source and self.entrypoint for files
+            if not self.source.exists():
+                raise FileNotFoundError(f"{self.source} does not exist")
+            if self.entrypoint is None:
+                if self.source.is_file():
+                    self.entrypoint = self.source
+                else:
+                    raise ValueError("entrypoint must be provided when source is a directory")
+            if not self.entrypoint.is_file():
+                # Check if self.entrypoint is a valid relative path
+                self.entrypoint = self.source.joinpath(self.entrypoint)
+                if not self.entrypoint.is_file():
+                    raise FileNotFoundError(f"File {self.entrypoint} does not exist")
+            if not self.entrypoint.is_relative_to(self.source):
+                raise ValueError(f"{self.entrypoint} must be a subpath of {self.source}")
+            if self.entrypoint.suffix != ".py":
+                raise NotImplementedError("Only Python entrypoints are supported currently")
+        else:
+            raise ValueError("Unsupported source type. Source must be a file, directory, or callable.")
+
+    def upload(self, session: snowpark.Session, stage_path: Union[str, PurePath]) -> types.UploadedPayload:
+        # Validate payload
+        self.validate()
+
+        # Prepare local variables
+        stage_path = PurePath(stage_path) if isinstance(stage_path, str) else stage_path
+        source = self.source
+        entrypoint = self.entrypoint or Path(constants.DEFAULT_ENTRYPOINT_PATH)
+
+        # Create stage if necessary
+        stage_name = stage_path.parts[0]
+        session.sql(
+            f"create stage if not exists {stage_name.lstrip('@')}"
+            " encryption = ( type = 'SNOWFLAKE_SSE' )"
+            " comment = 'Created by snowflake.ml.jobs Python API'"
+        ).collect()
+
+        # Upload payload to stage
+        if not isinstance(source, Path):
+            source_code = generate_python_code(source, source_code_display=True)
+            _ = session.file.put_stream(
+                io.BytesIO(source_code.encode()),
+                stage_location=stage_path.joinpath(entrypoint).as_posix(),
+                auto_compress=False,
+                overwrite=True,
+            )
+            source = entrypoint.parent
+        elif source.is_dir():
+            # Manually traverse the directory and upload each file, since Snowflake PUT
+            # can't handle directories. Reduce the number of PUT operations by using
+            # wildcard patterns to batch upload files with the same extension.
+            for path in {
+                p.parent.joinpath(f"*{p.suffix}") if p.suffix else p for p in source.resolve().rglob("*") if p.is_file()
+            }:
+                session.file.put(
+                    str(path),
+                    stage_path.joinpath(path.parent.relative_to(source)).as_posix(),
+                    overwrite=True,
+                    auto_compress=False,
+                )
+        else:
+            session.file.put(
+                str(source.resolve()),
+                stage_path.as_posix(),
+                overwrite=True,
+                auto_compress=False,
+            )
+            source = source.parent
+
+        # Upload requirements
+        # TODO: Check if payload includes both a requirements.txt file and pip_requirements
+        if self.pip_requirements:
+            # Upload requirements.txt to stage
+            session.file.put_stream(
+                io.BytesIO("\n".join(self.pip_requirements).encode()),
+                stage_location=stage_path.joinpath("requirements.txt").as_posix(),
+                auto_compress=False,
+                overwrite=True,
+            )
+
+        # Upload startup script
+        # TODO: Make sure payload does not include file with same name
+        session.file.put_stream(
+            io.BytesIO(_STARTUP_SCRIPT_CODE.encode()),
+            stage_location=stage_path.joinpath(_STARTUP_SCRIPT_PATH).as_posix(),
+            auto_compress=False,
+            overwrite=False,  # FIXME
+        )
+
+        return types.UploadedPayload(
+            stage_path=stage_path,
+            entrypoint=[
+                "bash",
+                _STARTUP_SCRIPT_PATH,
+                entrypoint.relative_to(source),
+            ],
+        )
+
+
+def get_parameter_type(param: inspect.Parameter) -> Optional[Type[object]]:
+    # Unwrap Optional type annotations
+    param_type = param.annotation
+    if get_origin(param_type) is Union and len(get_args(param_type)) == 2 and type(None) in get_args(param_type):
+        param_type = next(t for t in get_args(param_type) if t is not type(None))
+
+    # Return None for empty type annotations
+    if param_type == inspect.Parameter.empty:
+        return None
+    return cast(Type[object], param_type)
+
+
+def validate_parameter_type(param_type: Type[object], param_name: str) -> None:
+    # Validate param_type is a supported type
+    if param_type not in _SUPPORTED_ARG_TYPES:
+        raise ValueError(
+            f"Unsupported argument type {param_type} for '{param_name}'."
+            f" Supported types: {', '.join(t.__name__ for t in _SUPPORTED_ARG_TYPES)}"
+        )
+
+
+def generate_python_code(func: Callable[..., Any], source_code_display: bool = False) -> str:
+    signature = inspect.signature(func)
+    if any(
+        p.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}
+        for p in signature.parameters.values()
+    ):
+        raise NotImplementedError("Function must not have unpacking arguments (* or **)")
+
+    # Mirrored from Snowpark generate_python_code() function
+    # https://github.com/snowflakedb/snowpark-python/blob/main/src/snowflake/snowpark/_internal/udf_utils.py
+    try:
+        source_code_comment = (
+            code_generation.generate_source_code(func) if source_code_display else ""  # type: ignore[arg-type]
+        )
+    except Exception as exc:
+        error_msg = f"Source code comment could not be generated for {func} due to error {exc}."
+        source_code_comment = code_generation.comment_source_code(error_msg)
+
+    func_name = "func"
+    func_code = f"""
+{source_code_comment}
+
+import pickle
+{func_name} = pickle.loads(bytes.fromhex('{cp.dumps(func).hex()}'))
+"""
+
+    # Generate argparse logic for argument handling (type coercion, default values, etc)
+    argparse_code = ["import argparse", "", "parser = argparse.ArgumentParser()"]
+    argparse_postproc = []
+    for name, param in signature.parameters.items():
+        opts = {}
+
+        param_type = get_parameter_type(param)
+        if param_type is not None:
+            validate_parameter_type(param_type, name)
+            opts["type"] = param_type.__name__
+
+        if param.default != inspect.Parameter.empty:
+            opts["default"] = f"'{param.default}'" if isinstance(param.default, str) else param.default
+
+        if param.kind == inspect.Parameter.KEYWORD_ONLY:
+            # Keyword argument
+            argparse_code.append(
+                f"parser.add_argument('--{name}', required={'default' not in opts},"
+                f" {', '.join(f'{k}={v}' for k, v in opts.items())})"
+            )
+        else:
+            # Positional argument. Use `argparse.add_mutually_exclusive_group()`
+            # to allow passing positional args by name as well
+            group_name = f"{name}_group"
+            argparse_code.append(
+                f"{group_name} = parser.add_mutually_exclusive_group(required={'default' not in opts})"
+            )
+            argparse_code.append(
+                f"{group_name}.add_argument('pos-{name}', metavar='{name}', nargs='?',"
+                f" {', '.join(f'{k}={v}' for k, v in opts.items() if k != 'default')})"
+            )
+            argparse_code.append(
+                f"{group_name}.add_argument('--{name}', {', '.join(f'{k}={v}' for k, v in opts.items())})"
+            )
+            argparse_code.append("")  # Add newline for readability
+            argparse_postproc.append(
+                f"args.{name} = {name} if ({name} := args.__dict__.pop('pos-{name}')) is not None else args.{name}"
+            )
+    argparse_code.append("args = parser.parse_args()")
+    param_code = "\n".join(argparse_code + argparse_postproc)
+
+    return f"""
+### Version guard to check compatibility across Python versions ###
+import sys
+import warnings
+
+if sys.version_info.major != {sys.version_info.major} or sys.version_info.minor != {sys.version_info.minor}:
+    warnings.warn(
+        "Python version mismatch: job was created using"
+        " python{sys.version_info.major}.{sys.version_info.minor}"
+        f" but runtime environment uses python{{sys.version_info.major}}.{{sys.version_info.minor}}."
+        " Compatibility across Python versions is not guaranteed and may result in unexpected behavior."
+        " This will be fixed in a future release; for now, please use Python version"
+        f" {{sys.version_info.major}}.{{sys.version_info.minor}}.",
+        RuntimeWarning,
+        stacklevel=0,
+    )
+### End version guard ###
+
+{func_code.strip()}
+
+if __name__ == '__main__':
+{textwrap.indent(param_code, '    ')}
+
+    {func_name}(**vars(args))
+"""