snowflake-ml-python 1.21.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -20,6 +20,7 @@ from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_meta import model_blob_meta, model_meta_schema
 from snowflake.ml.model._packager.model_meta_migrator import migrator_plans
 from snowflake.ml.model._packager.model_runtime import model_runtime
+from snowflake.ml.model.code_path import CodePath
 
 MODEL_METADATA_FILE = "model.yaml"
 MODEL_CODE_DIR = "code"
@@ -39,7 +40,7 @@ def create_model_metadata(
     signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
     function_properties: Optional[dict[str, dict[str, Any]]] = None,
     metadata: Optional[dict[str, str]] = None,
-    code_paths: Optional[list[str]] = None,
+    code_paths: Optional[list[model_types.CodePathLike]] = None,
     ext_modules: Optional[list[ModuleType]] = None,
     conda_dependencies: Optional[list[str]] = None,
     pip_requirements: Optional[list[str]] = None,
@@ -77,7 +78,8 @@ def create_model_metadata(
         **kwargs: Dict of attributes and values of the metadata. Used when loading from file.
 
     Raises:
-        ValueError: Raised when the code path contains reserved file or directory.
+        ValueError: Raised when the code path contains reserved file or directory, or destination conflicts.
+        FileNotFoundError: Raised when a code path does not exist.
 
     Yields:
         A model metadata object.
@@ -134,13 +136,44 @@ def create_model_metadata(
         os.makedirs(code_dir_path, exist_ok=True)
 
     if code_paths:
+        # Resolve all code paths and check for conflicts
+        resolved_paths: list[tuple[str, str]] = []  # (source, destination_relative)
         for code_path in code_paths:
-            # This part is to prevent users from providing code following our naming and overwrite our code.
-            if (
-                os.path.isfile(code_path) and os.path.splitext(os.path.basename(code_path))[0] == _SNOWFLAKE_PKG_NAME
-            ) or (os.path.isdir(code_path) and os.path.basename(code_path) == _SNOWFLAKE_PKG_NAME):
+            if isinstance(code_path, CodePath):
+                source, dest_relative = code_path._resolve()
+            else:
+                # String path: keep existing behavior
+                source = os.path.normpath(os.path.abspath(code_path))
+                if not os.path.exists(source):
+                    raise FileNotFoundError(f"Code path '{code_path}' does not exist (resolved to {source}).")
+                dest_relative = os.path.basename(source)
+            resolved_paths.append((source, dest_relative))
+
+        # Check for destination conflicts
+        seen: dict[str, str] = {}
+        for source, dest in resolved_paths:
+            if dest in seen:
+                raise ValueError(
+                    f"Destination path conflict: '{dest}' is targeted by both '{seen[dest]}' and '{source}'."
+                )
+            seen[dest] = source
+
+        # Copy files
+        for source, dest_relative in resolved_paths:
+            # Prevent reserved name conflicts
+            dest_name = dest_relative.split(os.sep)[0] if os.sep in dest_relative else dest_relative
+            if (os.path.isfile(source) and os.path.splitext(dest_name)[0] == _SNOWFLAKE_PKG_NAME) or (
+                os.path.isdir(source) and dest_name == _SNOWFLAKE_PKG_NAME
+            ):
                 raise ValueError("`snowflake` is a reserved name and you cannot contain that into code path.")
-            file_utils.copy_file_or_tree(code_path, code_dir_path)
+
+            parent_dir = (
+                os.path.join(code_dir_path, os.path.dirname(dest_relative))
+                if os.path.dirname(dest_relative)
+                else code_dir_path
+            )
+            os.makedirs(parent_dir, exist_ok=True)
+            file_utils.copy_file_or_tree(source, parent_dir)
 
     try:
         imported_modules = []

snowflake/ml/model/_packager/model_packager.py

@@ -49,7 +49,7 @@ class ModelPackager:
         target_platforms: Optional[list[model_types.TargetPlatform]] = None,
         python_version: Optional[str] = None,
         ext_modules: Optional[list[ModuleType]] = None,
-        code_paths: Optional[list[str]] = None,
+        code_paths: Optional[list[model_types.CodePathLike]] = None,
         options: model_types.ModelSaveOption,
         task: model_types.Task = model_types.Task.UNKNOWN,
     ) -> model_meta.ModelMetadata:

snowflake/ml/model/_signatures/core.py

@@ -191,6 +191,35 @@ class DataType(Enum):
             original_exception=NotImplementedError(f"Type {snowpark_type} is not supported as a DataType."),
         )
 
+    @classmethod
+    def from_python_type(cls, python_type: type) -> "DataType":
+        """Translate Python built-in type to DataType for signature definition.
+
+        Args:
+            python_type: A Python built-in type (int, float, str, bool).
+
+        Raises:
+            SnowflakeMLException: NotImplementedError: Raised when the given Python type is not supported.
+
+        Returns:
+            Corresponding DataType.
+        """
+        python_to_snowml_type_mapping: dict[type, "DataType"] = {
+            int: DataType.INT64,
+            float: DataType.DOUBLE,
+            str: DataType.STRING,
+            bool: DataType.BOOL,
+        }
+        if python_type in python_to_snowml_type_mapping:
+            return python_to_snowml_type_mapping[python_type]
+        raise snowml_exceptions.SnowflakeMLException(
+            error_code=error_codes.NOT_IMPLEMENTED,
+            original_exception=NotImplementedError(
+                f"Python type {python_type} is not supported as a DataType. "
+                f"Supported types are: {list(python_to_snowml_type_mapping.keys())}."
+            ),
+        )
+
 
 class BaseFeatureSpec(ABC):
     """Abstract Class for specification of a feature."""
@@ -764,10 +793,17 @@ class ModelSignature:
                 the output of the model.
             params: A sequence of parameter specifications and parameter group specifications that will compose
                 the parameters of the model. Defaults to None.
+
+        Raises:
+            SnowflakeMLException: ValueError: When the parameters have duplicate names or the same
+                names as input features.
+
+        # noqa: DAR402
         """
         self._inputs = inputs
         self._outputs = outputs
         self._params = params or []
+        self._name_validation()
 
     @property
     def inputs(self) -> Sequence[BaseFeatureSpec]:
@@ -879,6 +915,55 @@ class ModelSignature:
 
         return html_utils.create_base_container("Model Signature", content)
 
+    def _name_validation(self) -> None:
+        """Validate the names of the inputs and parameters.
+
+        Names are compared case-insensitively (matches Snowflake identifier behavior).
+
+        Raises:
+            SnowflakeMLException: ValueError: When the parameters have duplicate names or the same
+                names as input features.
+        """
+        input_names: set[str] = set()
+        for input_spec in self._inputs:
+            names = (
+                [input_spec.name.upper() for spec in input_spec._specs]
+                if isinstance(input_spec, FeatureGroupSpec)
+                else [input_spec.name.upper()]
+            )
+            input_names.update(names)
+
+        param_names: set[str] = set()
+        dup_params: set[str] = set()
+        collision_names: set[str] = set()
+
+        for param in self._params:
+            names = [spec.name for spec in param.specs] if isinstance(param, ParamGroupSpec) else [param.name]
+            for name in names:
+                if name.upper() in param_names:
+                    dup_params.add(name)
+                if name.upper() in input_names:
+                    collision_names.add(name)
+                param_names.add(name.upper())
+
+        if dup_params:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.INVALID_ARGUMENT,
+                original_exception=ValueError(
+                    f"Found duplicate parameter named resolved as {', '.join(sorted(dup_params))}."
+                    " Parameters must have distinct names (case-insensitive)."
+                ),
+            )
+
+        if collision_names:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.INVALID_ARGUMENT,
+                original_exception=ValueError(
+                    f"Found parameter(s) with the same name as input feature(s): {', '.join(sorted(collision_names))}."
+                    " Parameters and inputs must have distinct names (case-insensitive)."
+                ),
+            )
+
     @classmethod
     def from_mlflow_sig(cls, mlflow_sig: "mlflow.models.ModelSignature") -> "ModelSignature":
         return ModelSignature(

snowflake/ml/model/code_path.py

@@ -0,0 +1,104 @@
+"""CodePath class for selective code packaging in model registry."""
+
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+_ERR_ROOT_NOT_FOUND = "CodePath: root '{root}' does not exist (resolved to {resolved})."
+_ERR_WILDCARDS_NOT_SUPPORTED = "CodePath: Wildcards are not supported in filter. Got '{filter}'. Use exact paths only."
+_ERR_FILTER_MUST_BE_RELATIVE = "CodePath: filter must be a relative path, got absolute path '{filter}'."
+_ERR_FILTER_HOME_PATH = "CodePath: filter must be a relative path, got home directory path '{filter}'."
+_ERR_FILTER_ON_FILE_ROOT = (
+    "CodePath: cannot apply filter to a file root. " "Root '{root}' is a file. Use filter only with directory roots."
+)
+_ERR_FILTER_ESCAPES_ROOT = (
+    "CodePath: filter '{filter}' escapes root directory '{root}'. " "Relative paths must stay within root."
+)
+_ERR_FILTER_NOT_FOUND = "CodePath: filter '{filter}' under root '{root}' does not exist (resolved to {resolved})."
+
+
+@dataclass(frozen=True)
+class CodePath:
+    """Specifies a code path with optional filtering for selective inclusion.
+
+    Args:
+        root: The root directory or file path (absolute or relative to cwd).
+        filter: Optional relative path under root to select a subdirectory or file.
+            The filter also determines the destination path under code/.
+
+    Examples:
+        CodePath("project/src/")                          # Copy entire src/ to code/src/
+        CodePath("project/src/", filter="utils")          # Copy utils/ to code/utils/
+        CodePath("project/src/", filter="lib/helpers")    # Copy to code/lib/helpers/
+    """
+
+    root: str
+    filter: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        if self.filter == "":
+            object.__setattr__(self, "filter", None)
+
+    def __repr__(self) -> str:
+        if self.filter:
+            return f"CodePath({self.root!r}, filter={self.filter!r})"
+        return f"CodePath({self.root!r})"
+
+    def _validate_filter(self) -> Optional[str]:
+        """Validate and normalize filter, returning normalized filter or None.
+
+        Returns:
+            Normalized filter path, or None if no filter is set.
+
+        Raises:
+            ValueError: If filter contains wildcards or is an absolute path.
+        """
+        if self.filter is None:
+            return None
+
+        if any(c in self.filter for c in ["*", "?", "[", "]"]):
+            raise ValueError(_ERR_WILDCARDS_NOT_SUPPORTED.format(filter=self.filter))
+
+        if self.filter.startswith("~"):
+            raise ValueError(_ERR_FILTER_HOME_PATH.format(filter=self.filter))
+
+        filter_normalized = os.path.normpath(self.filter)
+
+        if os.path.isabs(filter_normalized):
+            raise ValueError(_ERR_FILTER_MUST_BE_RELATIVE.format(filter=self.filter))
+
+        return filter_normalized
+
+    def _resolve(self) -> tuple[str, str]:
+        """Resolve the source path and destination path.
+
+        Returns:
+            Tuple of (source_path, destination_relative_path)
+
+        Raises:
+            FileNotFoundError: If root or filter path does not exist.
+            ValueError: If filter is invalid (wildcards, absolute, escapes root, or applied to file).
+        """
+        filter_normalized = self._validate_filter()
+        root_normalized = os.path.normpath(os.path.abspath(self.root))
+
+        if filter_normalized is None:
+            if not os.path.exists(root_normalized):
+                raise FileNotFoundError(_ERR_ROOT_NOT_FOUND.format(root=self.root, resolved=root_normalized))
+            return root_normalized, os.path.basename(root_normalized)
+
+        if not os.path.exists(root_normalized):
+            raise FileNotFoundError(_ERR_ROOT_NOT_FOUND.format(root=self.root, resolved=root_normalized))
+
+        if os.path.isfile(root_normalized):
+            raise ValueError(_ERR_FILTER_ON_FILE_ROOT.format(root=self.root))
+
+        source = os.path.normpath(os.path.join(root_normalized, filter_normalized))
+
+        if not (source.startswith(root_normalized + os.sep) or source == root_normalized):
+            raise ValueError(_ERR_FILTER_ESCAPES_ROOT.format(filter=self.filter, root=self.root))
+
+        if not os.path.exists(source):
+            raise FileNotFoundError(_ERR_FILTER_NOT_FOUND.format(filter=self.filter, root=self.root, resolved=source))
+
+        return source, filter_normalized

snowflake/ml/model/custom_model.py

@@ -4,10 +4,13 @@ from typing import Any, Callable, Coroutine, Generator, Optional, Union
 
 import anyio
 import pandas as pd
-from typing_extensions import deprecated
+from typing_extensions import Concatenate, ParamSpec, deprecated
 
 from snowflake.ml.model import type_hints as model_types
 
+# Captures additional keyword-only parameters for inference methods
+InferenceParams = ParamSpec("InferenceParams")
+
 
 class MethodRef:
     """Represents a method invocation of an instance of `ModelRef`.
@@ -217,7 +220,7 @@ class CustomModel:
 
     def _get_infer_methods(
         self,
-    ) -> Generator[Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame], None, None]:
+    ) -> Generator[Callable[..., pd.DataFrame], None, None]:
         """Returns all methods in CLS with `inference_api` decorator as the outermost decorator."""
         for cls_method_str in dir(self):
             cls_method = getattr(self, cls_method_str)
@@ -240,7 +243,7 @@ class CustomModel:
         return rv
 
 
-def _validate_predict_function(func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]) -> None:
+def _validate_predict_function(func: Callable[..., pd.DataFrame]) -> None:
     """Validate the user provided predict method.
 
     Args:
@@ -248,19 +251,22 @@ def _validate_predict_function(func: Callable[[model_types.CustomModelType, pd.D
 
     Raises:
         TypeError: Raised when the method is not a callable object.
-        TypeError: Raised when the method does not have 2 arguments (self and X).
+        TypeError: Raised when the method does not have at least 2 arguments (self and X).
         TypeError: Raised when the method does not have typing annotation.
         TypeError: Raised when the method's input (X) does not have type pd.DataFrame.
         TypeError: Raised when the method's output does not have type pd.DataFrame.
+        TypeError: Raised when additional parameters are not keyword-only with defaults.
     """
     if not callable(func):
         raise TypeError("Predict method is not callable.")
 
     func_signature = inspect.signature(func)
-    if len(func_signature.parameters) != 2:
-        raise TypeError("Predict method should have exact 2 arguments.")
+    func_signature_params = list(func_signature.parameters.values())
+
+    if len(func_signature_params) < 2:
+        raise TypeError("Predict method should have at least 2 arguments.")
 
-    input_annotation = list(func_signature.parameters.values())[1].annotation
+    input_annotation = func_signature_params[1].annotation
     output_annotation = func_signature.return_annotation
 
     if input_annotation == inspect.Parameter.empty or output_annotation == inspect.Signature.empty:
@@ -275,17 +281,53 @@ def _validate_predict_function(func: Callable[[model_types.CustomModelType, pd.D
     ):
         raise TypeError("Output for predict method should have type pandas.DataFrame.")
 
+    # Validate additional parameters (beyond self and input) are keyword-only with defaults
+    for func_signature_param in func_signature_params[2:]:
+        _validate_parameter(func_signature_param)
+
+
+def _validate_parameter(param: inspect.Parameter) -> None:
+    """Validate a parameter."""
+    if param.kind != inspect.Parameter.KEYWORD_ONLY:
+        raise TypeError(f"Parameter '{param.name}' must be keyword-only (defined after '*' in signature).")
+    if param.default == inspect.Parameter.empty:
+        raise TypeError(f"Parameter '{param.name}' must have a default value.")
+    if param.annotation == inspect.Parameter.empty:
+        raise TypeError(f"Parameter '{param.name}' must have a type annotation.")
+
+    # Validate annotation is a supported type
+    supported_types = {int, float, str, bool}
+    if param.annotation not in supported_types:
+        raise TypeError(
+            f"Parameter '{param.name}' has unsupported type annotation '{param.annotation}'. "
+            f"Supported types are: int, float, str, bool"
+        )
+
+
+def get_method_parameters(func: Callable[..., Any]) -> list[tuple[str, Any, Any]]:
+    """Extract keyword-only parameters with defaults from an inference method.
+
+    Args:
+        func: The inference method.
+
+    Returns:
+        A list of tuples (name, type annotation, default value) for each keyword-only parameter.
+    """
+    func_signature = inspect.signature(func)
+    params = list(func_signature.parameters.values())
+    return [(param.name, param.annotation, param.default) for param in params[2:]]
+
 
 def inference_api(
-    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame],
-) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
+    func: Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame],
+) -> Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     return func
 
 
 def partitioned_api(
-    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame],
-) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
+    func: Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame],
+) -> Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     func.__dict__["_is_partitioned_api"] = True
     return func
@@ -296,8 +338,8 @@ def partitioned_api(
     " Use snowflake.ml.custom_model.partitioned_api instead."
 )
 def partitioned_inference_api(
-    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame],
-) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
+    func: Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame],
+) -> Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     func.__dict__["_is_partitioned_api"] = True
     return func

snowflake/ml/model/model_signature.py

@@ -34,6 +34,9 @@ DataType = core.DataType
 BaseFeatureSpec = core.BaseFeatureSpec
 FeatureSpec = core.FeatureSpec
 FeatureGroupSpec = core.FeatureGroupSpec
+BaseParamSpec = core.BaseParamSpec
+ParamSpec = core.ParamSpec
+ParamGroupSpec = core.ParamGroupSpec
 ModelSignature = core.ModelSignature
 
 
@@ -711,6 +714,7 @@ def infer_signature(
     output_feature_names: Optional[list[str]] = None,
     input_data_limit: Optional[int] = 100,
     output_data_limit: Optional[int] = 100,
+    params: Optional[Sequence[core.BaseParamSpec]] = None,
 ) -> core.ModelSignature:
     """
     Infer model signature from given input and output sample data.
@@ -740,12 +744,20 @@ def infer_signature(
         output_data_limit: Limit the number of rows to be used in signature inference in the output data. Defaults to
             100. If None, all rows are used. If the number of rows in the output data is less than the limit, all rows
             are used.
+        params: Optional sequence of parameter specifications to include in the signature. Parameters define
+            optional configuration values that can be passed to model inference. Defaults to None.
+
+    Raises:
+        SnowflakeMLException: ValueError: Raised when input data contains columns matching parameter names.
 
     Returns:
         A model signature inferred from the given input and output sample data.
+
+    # noqa: DAR402
     """
     inputs = _infer_signature(_truncate_data(input_data, input_data_limit), role="input")
     inputs = utils.rename_features(inputs, input_feature_names)
+
     outputs = _infer_signature(_truncate_data(output_data, output_data_limit), role="output")
     outputs = utils.rename_features(outputs, output_feature_names)
-    return core.ModelSignature(inputs, outputs)
+    return core.ModelSignature(inputs, outputs, params=params)

snowflake/ml/model/type_hints.py

@@ -13,6 +13,7 @@ from typing import (
 import numpy.typing as npt
 from typing_extensions import NotRequired
 
+from snowflake.ml.model.code_path import CodePath
 from snowflake.ml.model.compute_pool import (
     DEFAULT_CPU_COMPUTE_POOL,
     DEFAULT_GPU_COMPUTE_POOL,
@@ -366,6 +367,7 @@ ModelLoadOption = Union[
 
 
 SupportedTargetPlatformType = Union[TargetPlatform, str]
+CodePathLike = Union[str, CodePath]
 
 
 class ProgressStatus(Protocol):