snowflake-ml-python 1.21.0__tar.gz → 1.22.0__tar.gz

{snowflake_ml_python-1.21.0 → snowflake_ml_python-1.22.0}/CHANGELOG.md

@@ -1,5 +1,33 @@
 # Release History
 
+## 1.22.0
+
+### New Features
+
+* Registry: Introducing remotely logging a transformer pipeline model using SPCS job.
+
+```python
+# create reference of the model
+model = huggingface.TransformersPipeline(
+    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    task="text-generation",
+)
+
+# Remotely log the model, a SPCS job will run async and log the model
+mv = registry.log_model(
+    model=model,
+    model_name="tinyllama_remote_log",
+    target_platforms=["SNOWPARK_CONTAINER_SERVICES"],
+    signatures=openai_signatures.OPENAI_CHAT_SIGNATURE,
+)
+```
+
+### Bug Fixes
+
+### Behavior Changes
+
+### Deprecations
+
 ## 1.21.0
 
 ### New Features

{snowflake_ml_python-1.21.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.22.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.21.0
+Version: 1.22.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: 
@@ -417,6 +417,34 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
+## 1.22.0
+
+### New Features
+
+* Registry: Introducing remotely logging a transformer pipeline model using SPCS job.
+
+```python
+# create reference of the model
+model = huggingface.TransformersPipeline(
+    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    task="text-generation",
+)
+
+# Remotely log the model, a SPCS job will run async and log the model
+mv = registry.log_model(
+    model=model,
+    model_name="tinyllama_remote_log",
+    target_platforms=["SNOWPARK_CONTAINER_SERVICES"],
+    signatures=openai_signatures.OPENAI_CHAT_SIGNATURE,
+)
+```
+
+### Bug Fixes
+
+### Behavior Changes
+
+### Deprecations
+
 ## 1.21.0
 
 ### New Features

snowflake_ml_python-1.22.0/snowflake/ml/_internal/utils/url.py

@@ -0,0 +1,42 @@
+from urllib.parse import urlunparse
+
+from snowflake import snowpark as snowpark
+
+JOB_URL_PREFIX = "#/compute/job/"
+SERVICE_URL_PREFIX = "#/compute/service/"
+
+
+def get_snowflake_url(
+    session: snowpark.Session,
+    url_path: str,
+    params: str = "",
+    query: str = "",
+    fragment: str = "",
+) -> str:
+    """Construct a Snowflake URL from session connection details and URL components.
+
+    Args:
+        session: The Snowpark session containing connection details.
+        url_path: The path component of the URL (e.g., "/compute/job/123").
+        params: Optional parameters for the URL (RFC 1808). Defaults to "".
+        query: Optional query string for the URL. Defaults to "".
+        fragment: Optional fragment identifier for the URL (e.g., "#section"). Defaults to "".
+
+    Returns:
+        A fully constructed Snowflake URL string with scheme, host, and specified components.
+    """
+    scheme = "https"
+    if hasattr(session.connection, "scheme"):
+        scheme = session.connection.scheme
+    host = session.connection.host
+
+    return urlunparse(
+        (
+            scheme,
+            host,
+            url_path,
+            params,
+            query,
+            fragment,
+        )
+    )

{snowflake_ml_python-1.21.0 → snowflake_ml_python-1.22.0}/snowflake/ml/jobs/_utils/constants.py

@@ -30,6 +30,7 @@ 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.8.0"
 DEFAULT_ENTRYPOINT_PATH = "func.py"
+DEFAULT_PYTHON_VERSION = "3.10"
 
 # Percent of container memory to allocate for /dev/shm volume
 MEMORY_VOLUME_SIZE = 0.3

snowflake_ml_python-1.22.0/snowflake/ml/jobs/_utils/spec_utils.py

@@ -0,0 +1,22 @@
+from snowflake import snowpark
+from snowflake.ml._internal.utils import snowflake_env
+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 = 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"]
+    cloud = snowflake_env.get_current_cloud(session, default=snowflake_env.SnowflakeCloudType.AWS)
+
+    return (
+        constants.COMMON_INSTANCE_FAMILIES.get(instance_family)
+        or constants.CLOUD_INSTANCE_FAMILIES[cloud][instance_family]
+    )

{snowflake_ml_python-1.21.0 → snowflake_ml_python-1.22.0}/snowflake/ml/lineage/lineage_node.py

@@ -7,7 +7,7 @@ from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import identifier, mixins
 
 if TYPE_CHECKING:
-    from snowflake.ml import dataset
+    from snowflake.ml.dataset import dataset
     from snowflake.ml.feature_store import feature_view
     from snowflake.ml.model._client.model import model_version_impl
 

{snowflake_ml_python-1.21.0 → snowflake_ml_python-1.22.0}/snowflake/ml/model/__init__.py

@@ -2,16 +2,20 @@ import sys
 import warnings
 
 from snowflake.ml.model._client.model.batch_inference_specs import (
+    ColumnHandlingOptions,
+    FileEncoding,
     JobSpec,
     OutputSpec,
     SaveMode,
 )
 from snowflake.ml.model._client.model.model_impl import Model
 from snowflake.ml.model._client.model.model_version_impl import ExportMode, ModelVersion
+from snowflake.ml.model.code_path import CodePath
 from snowflake.ml.model.models.huggingface_pipeline import HuggingFacePipelineModel
 from snowflake.ml.model.volatility import Volatility
 
 __all__ = [
+    "CodePath",
     "Model",
     "ModelVersion",
     "ExportMode",
@@ -20,6 +24,8 @@ __all__ = [
     "OutputSpec",
     "SaveMode",
     "Volatility",
+    "FileEncoding",
+    "ColumnHandlingOptions",
 ]
 
 _deprecation_warning_msg_for_3_9 = (

{snowflake_ml_python-1.21.0 → snowflake_ml_python-1.22.0}/snowflake/ml/model/_client/model/batch_inference_specs.py

@@ -2,6 +2,7 @@ from enum import Enum
 from typing import Optional
 
 from pydantic import BaseModel
+from typing_extensions import TypedDict
 
 
 class SaveMode(str, Enum):
@@ -18,6 +19,20 @@ class SaveMode(str, Enum):
     ERROR = "error"
 
 
+class FileEncoding(str, Enum):
+    """The encoding of the file content that will be passed to the custom model."""
+
+    RAW_BYTES = "raw_bytes"
+    BASE64 = "base64"
+    BASE64_DATA_URL = "base64_data_url"
+
+
+class ColumnHandlingOptions(TypedDict):
+    """Options for handling specific columns during run_batch for file I/O."""
+
+    encoding: FileEncoding
+
+
 class OutputSpec(BaseModel):
     """Specification for batch inference output.
 
@@ -74,7 +89,7 @@ class JobSpec(BaseModel):
             the memory of the node.
         gpu_requests (Optional[str]): The gpu limit for GPU based inference. Can be integer or
             string values. Use CPU if None.
-        replicas (Optional[int]): Number of job replicas to run for high availability.
+        replicas (Optional[int]): Number of SPCS job nodes used for distributed inference.
             If not specified, defaults to 1 replica.
 
     Example:

{snowflake_ml_python-1.21.0 → snowflake_ml_python-1.22.0}/snowflake/ml/model/_client/model/model_version_impl.py

@@ -1,4 +1,6 @@
+import base64
 import enum
+import json
 import pathlib
 import tempfile
 import uuid
@@ -6,6 +8,7 @@ import warnings
 from typing import Any, Callable, Optional, Union, overload
 
 import pandas as pd
+from pydantic import TypeAdapter
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
@@ -30,6 +33,7 @@ _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
 _BATCH_INFERENCE_JOB_ID_PREFIX = "BATCH_INFERENCE_"
 _BATCH_INFERENCE_TEMPORARY_FOLDER = "_temporary"
+_UTF8_ENCODING = "utf-8"
 
 
 class ExportMode(enum.Enum):
@@ -495,6 +499,7 @@ class ModelVersion(lineage_node.LineageNode):
         function_name: Optional[str] = None,
         partition_column: Optional[str] = None,
         strict_input_validation: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[pd.DataFrame, dataframe.DataFrame]:
         """Invoke a method in a model version object.
 
@@ -505,6 +510,8 @@ class ModelVersion(lineage_node.LineageNode):
             partition_column: The partition column name to partition by.
             strict_input_validation: Enable stricter validation for the input data. This will result value range based
                 type validation to make sure your input data won't overflow when providing to the model.
+            params: Optional dictionary of model inference parameters (e.g., temperature, top_k for LLMs).
+                These are passed as keyword arguments to the model's inference method. Defaults to None.
         """
         ...
 
@@ -516,6 +523,7 @@ class ModelVersion(lineage_node.LineageNode):
         service_name: str,
         function_name: Optional[str] = None,
         strict_input_validation: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[pd.DataFrame, dataframe.DataFrame]:
         """Invoke a method in a model version object via a service.
 
@@ -525,6 +533,8 @@ class ModelVersion(lineage_node.LineageNode):
             function_name: The function name to run. It is the name used to call a function in SQL.
             strict_input_validation: Enable stricter validation for the input data. This will result value range based
                 type validation to make sure your input data won't overflow when providing to the model.
+            params: Optional dictionary of model inference parameters (e.g., temperature, top_k for LLMs).
+                These are passed as keyword arguments to the model's inference method. Defaults to None.
         """
         ...
 
@@ -541,6 +551,7 @@ class ModelVersion(lineage_node.LineageNode):
         function_name: Optional[str] = None,
         partition_column: Optional[str] = None,
         strict_input_validation: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[pd.DataFrame, "dataframe.DataFrame"]:
         """Invoke a method in a model version object via the warehouse or a service.
 
@@ -552,6 +563,8 @@ class ModelVersion(lineage_node.LineageNode):
             partition_column: The partition column name to partition by.
             strict_input_validation: Enable stricter validation for the input data. This will result value range based
                 type validation to make sure your input data won't overflow when providing to the model.
+            params: Optional dictionary of model inference parameters (e.g., temperature, top_k for LLMs).
+                These are passed as keyword arguments to the model's inference method. Defaults to None.
 
         Returns:
             The prediction data. It would be the same type dataframe as your input.
@@ -582,6 +595,7 @@ class ModelVersion(lineage_node.LineageNode):
                 service_name=service_name_id,
                 strict_input_validation=strict_input_validation,
                 statement_params=statement_params,
+                params=params,
             )
         else:
             manifest = self._get_model_manifest(statement_params=statement_params)
@@ -621,6 +635,7 @@ class ModelVersion(lineage_node.LineageNode):
                 statement_params=statement_params,
                 is_partitioned=target_function_info["is_partitioned"],
                 explain_case_sensitive=explain_case_sensitive,
+                params=params,
             )
 
     def _determine_explain_case_sensitivity(
@@ -634,6 +649,41 @@ class ModelVersion(lineage_node.LineageNode):
             method_options, target_function_info["name"]
         )
 
+    @staticmethod
+    def _encode_column_handling(
+        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]],
+    ) -> Optional[str]:
+        """Validate and encode column_handling to a base64 string.
+
+        Args:
+            column_handling: Optional dictionary mapping column names to file encoding options.
+
+        Returns:
+            Base64 encoded JSON string of the column handling options, or None if input is None.
+        """
+        # TODO: validation for column names
+        if column_handling is None:
+            return None
+        adapter = TypeAdapter(dict[str, batch_inference_specs.ColumnHandlingOptions])
+        # TODO: throw error if the validate_python function fails
+        validated_input = adapter.validate_python(column_handling)
+        return base64.b64encode(adapter.dump_json(validated_input)).decode(_UTF8_ENCODING)
+
+    @staticmethod
+    def _encode_params(params: Optional[dict[str, Any]]) -> Optional[str]:
+        """Encode params dictionary to a base64 string.
+
+        Args:
+            params: Optional dictionary of model inference parameters.
+
+        Returns:
+            Base64 encoded JSON string of the params, or None if input is None.
+        """
+        if params is None:
+            return None
+        # TODO: validation for param names, types
+        return base64.b64encode(json.dumps(params).encode(_UTF8_ENCODING)).decode(_UTF8_ENCODING)
+
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
@@ -651,6 +701,8 @@ class ModelVersion(lineage_node.LineageNode):
         input_spec: dataframe.DataFrame,
         output_spec: batch_inference_specs.OutputSpec,
         job_spec: Optional[batch_inference_specs.JobSpec] = None,
+        params: Optional[dict[str, Any]] = None,
+        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]] = None,
     ) -> job.MLJob[Any]:
         """Execute batch inference on datasets as an SPCS job.
 
@@ -664,6 +716,12 @@ class ModelVersion(lineage_node.LineageNode):
             job_spec (Optional[batch_inference_specs.JobSpec]): Optional configuration for job
                 execution parameters such as compute resources, worker counts, and job naming.
                 If None, default values will be used.
+            params (Optional[dict[str, Any]]): Optional dictionary of model inference parameters
+                (e.g., temperature, top_k for LLMs). These are passed as keyword arguments to the
+                model's inference method. Defaults to None.
+            column_handling (Optional[dict[str, batch_inference_specs.FileEncoding]]): Optional dictionary
+                specifying how to handle specific columns during file I/O. Maps column names to their
+                file encoding configuration.
 
         Returns:
             job.MLJob[Any]: A batch inference job object that can be used to monitor progress and manage the job
@@ -719,6 +777,9 @@ class ModelVersion(lineage_node.LineageNode):
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
+        column_handling_as_string = self._encode_column_handling(column_handling)
+        params_as_string = self._encode_params(params)
+
         if job_spec is None:
             job_spec = batch_inference_specs.JobSpec()
 
@@ -766,6 +827,8 @@ class ModelVersion(lineage_node.LineageNode):
             # input and output
             input_stage_location=input_stage_location,
             input_file_pattern="*",
+            column_handling=column_handling_as_string,
+            params=params_as_string,
             output_stage_location=output_stage_location,
             completion_filename="_SUCCESS",
             # misc

snowflake_ml_python-1.22.0/snowflake/ml/model/_client/ops/deployment_step.py

@@ -0,0 +1,36 @@
+import enum
+import hashlib
+from typing import Optional
+
+
+class DeploymentStep(enum.Enum):
+    MODEL_BUILD = ("model-build", "model_build_")
+    MODEL_INFERENCE = ("model-inference", None)
+    MODEL_LOGGING = ("model-logging", "model_logging_")
+
+    def __init__(self, container_name: str, service_name_prefix: Optional[str]) -> None:
+        self._container_name = container_name
+        self._service_name_prefix = service_name_prefix
+
+    @property
+    def container_name(self) -> str:
+        """Get the container name for the deployment step."""
+        return self._container_name
+
+    @property
+    def service_name_prefix(self) -> Optional[str]:
+        """Get the service name prefix for the deployment step."""
+        return self._service_name_prefix
+
+
+def get_service_id_from_deployment_step(query_id: str, deployment_step: DeploymentStep) -> str:
+    """Get the service ID through the server-side logic."""
+    uuid = query_id.replace("-", "")
+    big_int = int(uuid, 16)
+    md5_hash = hashlib.md5(str(big_int).encode(), usedforsecurity=False).hexdigest()
+    identifier = md5_hash[:8]
+    service_name_prefix = deployment_step.service_name_prefix
+    if service_name_prefix is None:
+        # raise an exception if the service name prefix is None
+        raise ValueError(f"Service name prefix is {service_name_prefix} for deployment step {deployment_step}.")
+    return (service_name_prefix + identifier).upper()

{snowflake_ml_python-1.21.0 → snowflake_ml_python-1.22.0}/snowflake/ml/model/_client/ops/model_ops.py

@@ -1,5 +1,6 @@
 import enum
 import json
+import logging
 import os
 import pathlib
 import tempfile
@@ -11,9 +12,9 @@ from typing_extensions import NotRequired
 
 from snowflake.ml._internal import platform_capabilities
 from snowflake.ml._internal.exceptions import error_codes, exceptions
-from snowflake.ml._internal.utils import formatting, identifier, sql_identifier
+from snowflake.ml._internal.utils import formatting, identifier, sql_identifier, url
 from snowflake.ml.model import model_signature, type_hints
-from snowflake.ml.model._client.ops import metadata_ops
+from snowflake.ml.model._client.ops import deployment_step, metadata_ops
 from snowflake.ml.model._client.sql import (
     model as model_sql,
     model_version as model_version_sql,
@@ -33,6 +34,8 @@ from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
+logger = logging.getLogger(__name__)
+
 
 # An enum class to represent Create Or Alter Model SQL command.
 class ModelAction(enum.Enum):
@@ -986,6 +989,7 @@ class ModelOperator:
         statement_params: Optional[dict[str, str]] = None,
         is_partitioned: Optional[bool] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         ...
 
@@ -1002,6 +1006,7 @@ class ModelOperator:
         strict_input_validation: bool = False,
         statement_params: Optional[dict[str, str]] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         ...
 
@@ -1022,6 +1027,7 @@ class ModelOperator:
         statement_params: Optional[dict[str, str]] = None,
         is_partitioned: Optional[bool] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         identifier_rule = model_signature.SnowparkIdentifierRule.INFERRED
 
@@ -1057,6 +1063,24 @@ class ModelOperator:
                 col_name = sql_identifier.SqlIdentifier(input_feature.name.upper(), case_sensitive=True)
             input_args.append(col_name)
 
+        method_parameters: Optional[list[tuple[sql_identifier.SqlIdentifier, Any]]] = None
+        if signature.params:
+            # Start with defaults from signature
+            final_params = {}
+            for param_spec in signature.params:
+                if hasattr(param_spec, "default_value"):
+                    final_params[param_spec.name] = param_spec.default_value
+
+            # Override with provided runtime parameters
+            if params:
+                final_params.update(params)
+
+            # Convert to list of tuples with SqlIdentifier for parameter names
+            method_parameters = [
+                (sql_identifier.SqlIdentifier(param_name), param_value)
+                for param_name, param_value in final_params.items()
+            ]
+
         returns = []
         for output_feature in signature.outputs:
             output_name = identifier_rule.get_sql_identifier_from_feature(output_feature.name)
@@ -1075,6 +1099,7 @@ class ModelOperator:
                 schema_name=schema_name,
                 service_name=service_name,
                 statement_params=statement_params,
+                params=method_parameters,
             )
         else:
             assert model_name is not None
@@ -1090,6 +1115,7 @@ class ModelOperator:
                     model_name=model_name,
                     version_name=version_name,
                     statement_params=statement_params,
+                    params=method_parameters,
                 )
             elif method_function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
                 df_res = self._model_version_client.invoke_table_function_method(
@@ -1105,6 +1131,7 @@ class ModelOperator:
                     statement_params=statement_params,
                     is_partitioned=is_partitioned or False,
                     explain_case_sensitive=explain_case_sensitive,
+                    params=method_parameters,
                 )
 
         if keep_order:
@@ -1238,3 +1265,35 @@ class ModelOperator:
                     target_path=local_file_dir,
                     statement_params=statement_params,
                 )
+
+    def run_import_model_query(
+        self,
+        *,
+        database_name: str,
+        schema_name: str,
+        yaml_content: str,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> None:
+        yaml_content_escaped = snowpark_utils.escape_single_quotes(yaml_content)  # type: ignore[no-untyped-call]
+
+        async_job = self._session.sql(
+            f"SELECT SYSTEM$IMPORT_MODEL('{yaml_content_escaped}')",
+        ).collect(block=False, statement_params=statement_params)
+        query_id = async_job.query_id  # type: ignore[attr-defined]
+
+        logger.info(f"Remotely importing model, with the query id: {query_id}")
+        model_logger_service_name = sql_identifier.SqlIdentifier(
+            deployment_step.get_service_id_from_deployment_step(
+                query_id,
+                deployment_step.DeploymentStep.MODEL_LOGGING,
+            )
+        )
+
+        logger_name = model_logger_service_name.identifier()
+        job_url = f"{url.JOB_URL_PREFIX}/{database_name}/{schema_name}/{logger_name}"
+        snowflake_url = url.get_snowflake_url(session=self._session, url_path=job_url)
+        logger.info(
+            f"To monitor the progress of the model logging job, head to the job monitoring page {snowflake_url}"
+        )
+
+        async_job.result()  # type: ignore[attr-defined]