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

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -1,6 +1,6 @@
 import collections
 import pathlib
-from typing import List, Optional, TypedDict, Union
+from typing import Optional, TypedDict, Union
 
 from typing_extensions import NotRequired
 
@@ -137,8 +137,8 @@ class ModelMethod:
             )
 
         outputs: Union[
-            List[model_manifest_schema.ModelMethodSignatureField],
-            List[model_manifest_schema.ModelMethodSignatureFieldWithName],
+            list[model_manifest_schema.ModelMethodSignatureField],
+            list[model_manifest_schema.ModelMethodSignatureFieldWithName],
         ]
         if self.function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
             outputs = [

snowflake/ml/model/_packager/model_env/model_env.py

@@ -3,10 +3,11 @@ import itertools
 import os
 import pathlib
 import warnings
-from typing import DefaultDict, Dict, List, Optional
+from typing import DefaultDict, Optional
 
 from packaging import requirements, version
 
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env as snowml_env, env_utils
 from snowflake.ml.model._packager.model_meta import model_meta_schema
 
@@ -19,9 +20,8 @@ _DEFAULT_CONDA_ENV_FILENAME = "conda.yml"
 _DEFAULT_PIP_REQUIREMENTS_FILENAME = "requirements.txt"
 
 # The default CUDA version is chosen based on the driver availability in SPCS.
-# If changing this version, we need also change the version of default PyTorch in HuggingFace pipeline handler to
-# make sure they are compatible.
-DEFAULT_CUDA_VERSION = "11.8"
+# Make sure they are aligned with default CUDA version in inference server.
+DEFAULT_CUDA_VERSION = "12.4"
 
 
 class ModelEnv:
@@ -38,15 +38,16 @@ class ModelEnv:
         self.prefer_pip: bool = prefer_pip
         self.conda_env_rel_path = pathlib.PurePosixPath(pathlib.Path(conda_env_rel_path).as_posix())
         self.pip_requirements_rel_path = pathlib.PurePosixPath(pathlib.Path(pip_requirements_rel_path).as_posix())
-        self.artifact_repository_map: Optional[Dict[str, str]] = None
-        self._conda_dependencies: DefaultDict[str, List[requirements.Requirement]] = collections.defaultdict(list)
-        self._pip_requirements: List[requirements.Requirement] = []
+        self.artifact_repository_map: Optional[dict[str, str]] = None
+        self.resource_constraint: Optional[dict[str, str]] = None
+        self._conda_dependencies: DefaultDict[str, list[requirements.Requirement]] = collections.defaultdict(list)
+        self._pip_requirements: list[requirements.Requirement] = []
         self._python_version: version.Version = version.parse(snowml_env.PYTHON_VERSION)
         self._cuda_version: Optional[version.Version] = None
-        self._snowpark_ml_version: version.Version = version.parse(snowml_env.VERSION)
+        self._snowpark_ml_version: version.Version = version.parse(snowml_version.VERSION)
 
     @property
-    def conda_dependencies(self) -> List[str]:
+    def conda_dependencies(self) -> list[str]:
         """List of conda channel and dependencies from that to run the model"""
         return sorted(
             f"{chan}::{str(req)}" if chan else str(req)
@@ -57,24 +58,24 @@ class ModelEnv:
     @conda_dependencies.setter
     def conda_dependencies(
         self,
-        conda_dependencies: Optional[List[str]] = None,
+        conda_dependencies: Optional[list[str]] = None,
     ) -> None:
         self._conda_dependencies = env_utils.validate_conda_dependency_string_list(
-            conda_dependencies if conda_dependencies else []
+            conda_dependencies if conda_dependencies else [], add_local_version_specifier=True
         )
 
     @property
-    def pip_requirements(self) -> List[str]:
+    def pip_requirements(self) -> list[str]:
         """List of pip Python packages requirements for running the model."""
         return sorted(list(map(str, self._pip_requirements)))
 
     @pip_requirements.setter
     def pip_requirements(
         self,
-        pip_requirements: Optional[List[str]] = None,
+        pip_requirements: Optional[list[str]] = None,
     ) -> None:
         self._pip_requirements = env_utils.validate_pip_requirement_string_list(
-            pip_requirements if pip_requirements else []
+            pip_requirements if pip_requirements else [], add_local_version_specifier=True
         )
 
     @property
@@ -117,7 +118,7 @@ class ModelEnv:
 
     def include_if_absent(
         self,
-        pkgs: List[ModelDependency],
+        pkgs: list[ModelDependency],
         check_local_version: bool = False,
     ) -> None:
         """Append requirements into model env if absent. Depending on the environment, requirements may be added
@@ -128,7 +129,7 @@ class ModelEnv:
             check_local_version: Flag to indicate if it is required to pin to local version. Defaults to False.
         """
         if (self.pip_requirements or self.prefer_pip) and not self.conda_dependencies and pkgs:
-            pip_pkg_reqs: List[str] = []
+            pip_pkg_reqs: list[str] = []
             warnings.warn(
                 (
                     "Dependencies specified from pip requirements."
@@ -145,7 +146,7 @@ class ModelEnv:
         else:
             self._include_if_absent_conda(pkgs, check_local_version)
 
-    def _include_if_absent_conda(self, pkgs: List[ModelDependency], check_local_version: bool = False) -> None:
+    def _include_if_absent_conda(self, pkgs: list[ModelDependency], check_local_version: bool = False) -> None:
         """Append requirements into model env conda dependencies if absent.
 
         Args:
@@ -190,7 +191,7 @@ class ModelEnv:
                         stacklevel=2,
                     )
 
-    def _include_if_absent_pip(self, pkgs: List[str], check_local_version: bool = False) -> None:
+    def _include_if_absent_pip(self, pkgs: list[str], check_local_version: bool = False) -> None:
         """Append pip requirements into model env pip requirements if absent.
 
         Args:
@@ -207,7 +208,7 @@ class ModelEnv:
             except env_utils.DuplicateDependencyError:
                 pass
 
-    def remove_if_present_conda(self, conda_pkgs: List[str]) -> None:
+    def remove_if_present_conda(self, conda_pkgs: list[str]) -> None:
         """Remove conda requirements from model env if present.
 
         Args:
@@ -352,13 +353,14 @@ class ModelEnv:
     def load_from_dict(self, base_dir: pathlib.Path, env_dict: model_meta_schema.ModelEnvDict) -> None:
         self.conda_env_rel_path = pathlib.PurePosixPath(env_dict["conda"])
         self.pip_requirements_rel_path = pathlib.PurePosixPath(env_dict["pip"])
-        self.artifact_repository_map = env_dict.get("artifact_repository_map", None)
+        self.artifact_repository_map = env_dict.get("artifact_repository_map")
+        self.resource_constraint = env_dict.get("resource_constraint")
 
         self.load_from_conda_file(base_dir / self.conda_env_rel_path)
         self.load_from_pip_file(base_dir / self.pip_requirements_rel_path)
 
         self.python_version = env_dict["python_version"]
-        self.cuda_version = env_dict.get("cuda_version", None)
+        self.cuda_version = env_dict.get("cuda_version")
         self.snowpark_ml_version = env_dict["snowpark_ml_version"]
 
     def save_as_dict(
@@ -381,7 +383,8 @@ class ModelEnv:
         return {
             "conda": self.conda_env_rel_path.as_posix(),
             "pip": self.pip_requirements_rel_path.as_posix(),
-            "artifact_repository_map": self.artifact_repository_map if self.artifact_repository_map is not None else {},
+            "artifact_repository_map": self.artifact_repository_map or {},
+            "resource_constraint": self.resource_constraint or {},
             "python_version": self.python_version,
             "cuda_version": self.cuda_version,
             "snowpark_ml_version": self.snowpark_ml_version,
@@ -389,7 +392,7 @@ class ModelEnv:
 
     def validate_with_local_env(
         self, check_snowpark_ml_version: bool = False
-    ) -> List[env_utils.IncorrectLocalEnvironmentError]:
+    ) -> list[env_utils.IncorrectLocalEnvironmentError]:
         errors = []
         try:
             env_utils.validate_py_runtime_version(str(self._python_version))
@@ -413,10 +416,10 @@ class ModelEnv:
 
         if check_snowpark_ml_version:
             # For Modeling model
-            if self._snowpark_ml_version.base_version != snowml_env.VERSION:
+            if self._snowpark_ml_version.base_version != snowml_version.VERSION:
                 errors.append(
                     env_utils.IncorrectLocalEnvironmentError(
-                        f"The local installed version of Snowpark ML library is {snowml_env.VERSION} "
+                        f"The local installed version of Snowpark ML library is {snowml_version.VERSION} "
                         f"which differs from required version {self.snowpark_ml_version}."
                     )
                 )

snowflake/ml/model/_packager/model_handler.py

@@ -2,13 +2,13 @@ import functools
 import importlib
 import pkgutil
 from types import ModuleType
-from typing import Any, Callable, Dict, Optional, Type, TypeVar, cast
+from typing import Any, Callable, Optional, TypeVar, cast
 
 from snowflake.ml.model import type_hints as model_types
 from snowflake.ml.model._packager.model_handlers import _base
 
 _HANDLERS_BASE = "snowflake.ml.model._packager.model_handlers"
-_MODEL_HANDLER_REGISTRY: Dict[str, Type[_base.BaseModelHandler[model_types.SupportedModelType]]] = dict()
+_MODEL_HANDLER_REGISTRY: dict[str, type[_base.BaseModelHandler[model_types.SupportedModelType]]] = dict()
 _IS_HANDLER_LOADED = False
 
 
@@ -54,7 +54,7 @@ def ensure_handlers_registration(fn: F) -> F:
 @ensure_handlers_registration
 def find_handler(
     model: model_types.SupportedModelType,
-) -> Optional[Type[_base.BaseModelHandler[model_types.SupportedModelType]]]:
+) -> Optional[type[_base.BaseModelHandler[model_types.SupportedModelType]]]:
     for handler in _MODEL_HANDLER_REGISTRY.values():
         if handler.can_handle(model):
             return handler
@@ -64,7 +64,7 @@ def find_handler(
 @ensure_handlers_registration
 def load_handler(
     target_model_type: model_types.SupportedModelHandlerType,
-) -> Optional[Type[_base.BaseModelHandler[model_types.SupportedModelType]]]:
+) -> Optional[type[_base.BaseModelHandler[model_types.SupportedModelType]]]:
     for model_type, handler in _MODEL_HANDLER_REGISTRY.items():
         if target_model_type == model_type:
             return handler

snowflake/ml/model/_packager/model_handlers/_base.py

@@ -1,6 +1,6 @@
 import os
 from abc import abstractmethod
-from typing import Dict, Generic, Optional, Protocol, Type, final
+from typing import Generic, Optional, Protocol, final
 
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
@@ -14,7 +14,7 @@ class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
     HANDLER_TYPE: model_types.SupportedModelHandlerType
     HANDLER_VERSION: str
     _MIN_SNOWPARK_ML_VERSION: str
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]]
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]]
 
     @classmethod
     @abstractmethod

snowflake/ml/model/_packager/model_handlers/_utils.py

@@ -1,8 +1,9 @@
+import importlib
 import json
 import os
 import pathlib
 import warnings
-from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, cast
+from typing import Any, Callable, Iterable, Optional, Sequence, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -10,8 +11,10 @@ import pandas as pd
 from absl import logging
 
 import snowflake.snowpark.dataframe as sp_df
+from snowflake.ml._internal import env
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml.model import model_signature, type_hints as model_types
+from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_meta import model_meta
 from snowflake.ml.model._signatures import (
     core,
@@ -106,6 +109,35 @@ def get_input_signature(
     return input_sig
 
 
+def add_inferred_explain_method_signature(
+    model_meta: model_meta.ModelMetadata,
+    explain_method: str,
+    target_method: str,
+    background_data: model_types.SupportedDataType,
+    explain_fn: Callable[[model_types.SupportedLocalDataType], model_types.SupportedLocalDataType],
+    output_feature_names: Optional[Sequence[str]] = None,
+) -> model_meta.ModelMetadata:
+    inputs = get_input_signature(model_meta, target_method)
+    if output_feature_names is None:  # If not provided, assume output feature names are the same as input feature names
+        output_feature_names = [spec.name for spec in inputs]
+
+    if model_meta.model_type == "snowml":
+        suffixed_output_names = [identifier.concat_names([name, "_explanation"]) for name in output_feature_names]
+    else:
+        suffixed_output_names = [f"{name}_explanation" for name in output_feature_names]
+
+    truncated_background_data = get_truncated_sample_data(background_data, 5)
+    sig = model_signature.infer_signature(
+        input_data=truncated_background_data,
+        output_data=explain_fn(truncated_background_data),
+        input_feature_names=[spec.name for spec in inputs],
+        output_feature_names=suffixed_output_names,
+    )
+
+    model_meta.signatures[explain_method] = sig
+    return model_meta
+
+
 def add_explain_method_signature(
     model_meta: model_meta.ModelMetadata,
     explain_method: str,
@@ -231,10 +263,11 @@ def validate_model_task(passed_model_task: model_types.Task, inferred_model_task
 
 
 def get_explain_target_method(
-    model_metadata: model_meta.ModelMetadata, target_methods_list: List[str]
+    model_metadata: model_meta.ModelMetadata, target_methods_list: list[str]
 ) -> Optional[str]:
-    for method in model_metadata.signatures.keys():
-        if method in target_methods_list:
+    """Returns the first target method that is found in the model metadata signatures."""
+    for method in target_methods_list:
+        if method in model_metadata.signatures.keys():
             return method
     return None
 
@@ -248,7 +281,7 @@ def save_transformers_config_with_auto_map(local_model_path: str) -> None:
                 config_dict = json.load(f)
 
             # a. get repository and class_path from configs
-            auto_map_configs = cast(Dict[str, str], config_dict.get("auto_map", {}))
+            auto_map_configs = cast(dict[str, str], config_dict.get("auto_map", {}))
             for config_name, config_value in auto_map_configs.items():
                 repository, _, class_path = config_value.rpartition("--")
 
@@ -261,3 +294,12 @@ def save_transformers_config_with_auto_map(local_model_path: str) -> None:
 
             with open(f_path, "w") as f:
                 json.dump(config_dict, f)
+
+
+def get_default_cuda_version() -> str:
+    # Default to the env cuda version when running in ML runtime
+    if env.IN_ML_RUNTIME and importlib.util.find_spec("torch") is not None:
+        import torch
+
+        return torch.version.cuda or model_env.DEFAULT_CUDA_VERSION
+    return model_env.DEFAULT_CUDA_VERSION

snowflake/ml/model/_packager/model_handlers/catboost.py

@@ -1,6 +1,6 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Any, Callable, Optional, cast, final
 
 import numpy as np
 import pandas as pd
@@ -30,7 +30,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
     HANDLER_TYPE = "catboost"
     HANDLER_VERSION = "2024-03-21"
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model.bin"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
@@ -147,7 +147,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
         if enable_explainability:
             model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
         return None
 
@@ -202,7 +202,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
         def _create_custom_model(
             raw_model: "catboost.CatBoost",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "catboost.CatBoost",
                 signature: model_signature.ModelSignature,
@@ -235,7 +235,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
 
                 return fn
 
-            type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}
+            type_method_dict: dict[str, Any] = {"_raw_model": raw_model}
             for target_method_name, sig in model_meta.signatures.items():
                 type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 

snowflake/ml/model/_packager/model_handlers/custom.py

@@ -2,7 +2,7 @@ import inspect
 import os
 import pathlib
 import sys
-from typing import Dict, Optional, Type, cast, final
+from typing import Optional, cast, final
 
 import anyio
 import cloudpickle
@@ -28,7 +28,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
     HANDLER_TYPE = "custom"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     @classmethod
     def can_handle(cls, model: model_types.SupportedModelType) -> TypeGuard["custom_model.CustomModel"]:
@@ -72,7 +72,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                 predictions_df = target_method(model, sample_input_data)
             return predictions_df
 
-        for func_name in model._get_partitioned_infer_methods():
+        for func_name in model._get_partitioned_methods():
             function_properties = model_meta.function_properties.get(func_name, {})
             function_properties[model_meta_schema.FunctionProperties.PARTITIONED.value] = True
             model_meta.function_properties[func_name] = function_properties
@@ -99,7 +99,11 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
             for sub_name, model_ref in model.context.model_refs.items():
                 handler = model_handler.find_handler(model_ref.model)
                 if handler is None:
-                    raise TypeError("Your input type to custom model is not currently supported")
+                    raise TypeError(
+                        f"Model {sub_name} in model context is not a supported model type. See "
+                        "https://docs.snowflake.com/en/developer-guide/snowflake-ml/model-registry/"
+                        "bring-your-own-model-types for more details."
+                    )
                 sub_model = handler.cast_model(model_ref.model)
                 handler.save_model(
                     name=sub_name,
@@ -161,7 +165,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
             name: str(pathlib.PurePath(model_blob_path) / pathlib.PurePosixPath(rel_path))
             for name, rel_path in artifacts_meta.items()
         }
-        models: Dict[str, model_types.SupportedModelType] = dict()
+        models: dict[str, model_types.SupportedModelType] = dict()
         for sub_model_name, _ref in context.model_refs.items():
             model_type = model_meta.models[sub_model_name].model_type
             handler = model_handler.load_handler(model_type)

snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py

@@ -1,18 +1,7 @@
 import json
 import os
 import warnings
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    List,
-    Optional,
-    Type,
-    Union,
-    cast,
-    final,
-)
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
 
 import cloudpickle
 import numpy as np
@@ -38,7 +27,7 @@ if TYPE_CHECKING:
     import transformers
 
 
-def get_requirements_from_task(task: str, spcs_only: bool = False) -> List[model_env.ModelDependency]:
+def get_requirements_from_task(task: str, spcs_only: bool = False) -> list[model_env.ModelDependency]:
     # Text
     if task in [
         "conversational",
@@ -84,7 +73,7 @@ class HuggingFacePipelineHandler(
     HANDLER_TYPE = "huggingface_pipeline"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
     ADDITIONAL_CONFIG_FILE = "pipeline_config.pt"
@@ -250,20 +239,17 @@ class HuggingFacePipelineHandler(
             task, spcs_only=(not type_utils.LazyType("transformers.Pipeline").isinstance(model))
         )
         if framework is None or framework == "pt":
-            # Since we set default cuda version to be 11.8, to make sure it works with GPU, we need to have a default
-            # Pytorch version that works with CUDA 11.8 as well. This is required for huggingface pipelines only as
-            # users are not required to install pytorch locally if they are using the wrapper.
             pkgs_requirements.append(model_env.ModelDependency(requirement="pytorch", pip_name="torch"))
         elif framework == "tf":
             pkgs_requirements.append(model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"))
         model_meta.env.include_if_absent(
             pkgs_requirements, check_local_version=(type_utils.LazyType("transformers.Pipeline").isinstance(model))
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @staticmethod
-    def _get_device_config(**kwargs: Unpack[model_types.HuggingFaceLoadOptions]) -> Dict[str, str]:
-        device_config: Dict[str, Any] = {}
+    def _get_device_config(**kwargs: Unpack[model_types.HuggingFaceLoadOptions]) -> dict[str, str]:
+        device_config: dict[str, Any] = {}
         cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
         gpu_nums = 0
         if cuda_visible_devices is not None:
@@ -369,7 +355,7 @@ class HuggingFacePipelineHandler(
         def _create_custom_model(
             raw_model: "transformers.Pipeline",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "transformers.Pipeline",
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/keras.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import cloudpickle
 import numpy as np
@@ -32,7 +32,7 @@ class KerasHandler(_base.BaseModelHandler["keras.Model"]):
     HANDLER_TYPE = "keras"
     HANDLER_VERSION = "2025-01-01"
     _MIN_SNOWPARK_ML_VERSION = "1.7.5"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model.keras"
     CUSTOM_OBJECT_SAVE_PATH = "custom_objects.pkl"
@@ -146,7 +146,7 @@ class KerasHandler(_base.BaseModelHandler["keras.Model"]):
             dependencies,
             check_local_version=True,
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -185,7 +185,7 @@ class KerasHandler(_base.BaseModelHandler["keras.Model"]):
         def _create_custom_model(
             raw_model: "keras.Model",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "keras.Model",
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/lightgbm.py

@@ -1,16 +1,6 @@
 import os
 import warnings
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    Optional,
-    Type,
-    Union,
-    cast,
-    final,
-)
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
 
 import cloudpickle
 import numpy as np
@@ -41,7 +31,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
     HANDLER_TYPE = "lightgbm"
     HANDLER_VERSION = "2024-03-19"
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model.pkl"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
@@ -215,7 +205,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         def _create_custom_model(
             raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
                 signature: model_signature.ModelSignature,
@@ -250,7 +240,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
 
                 return fn
 
-            type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}
+            type_method_dict: dict[str, Any] = {"_raw_model": raw_model}
             for target_method_name, sig in model_meta.signatures.items():
                 type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 

snowflake/ml/model/_packager/model_handlers/mlflow.py

@@ -1,7 +1,7 @@
 import os
 import pathlib
 import tempfile
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
@@ -61,7 +61,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
     HANDLER_TYPE = "mlflow"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
     _DEFAULT_TARGET_METHOD = "predict"
@@ -204,7 +204,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
         def _create_custom_model(
             raw_model: "mlflow.pyfunc.PyFuncModel",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "mlflow.pyfunc.PyFuncModel",
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/pytorch.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import cloudpickle
 import pandas as pd
@@ -38,7 +38,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
     HANDLER_TYPE = "pytorch"
     HANDLER_VERSION = "2025-03-01"
     _MIN_SNOWPARK_ML_VERSION = "1.8.0"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {
         "2023-12-01": pytorch_migrator_2023_12_01.PyTorchHandlerMigrator20231201
     }
 
@@ -82,6 +82,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         enable_explainability = kwargs.get("enable_explainability", False)
         if enable_explainability:
             raise NotImplementedError("Explainability is not supported for PyTorch model.")
+        multiple_inputs = kwargs.get("multiple_inputs", False)
 
         import torch
 
@@ -94,8 +95,6 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
                 default_target_methods=cls.DEFAULT_TARGET_METHODS,
             )
 
-            multiple_inputs = kwargs.get("multiple_inputs", False)
-
             def get_prediction(
                 target_method_name: str, sample_input_data: "model_types.SupportedLocalDataType"
             ) -> model_types.SupportedLocalDataType:
@@ -151,7 +150,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         model_meta.env.include_if_absent(
             [model_env.ModelDependency(requirement="pytorch", pip_name="torch")], check_local_version=True
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -188,7 +187,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         def _create_custom_model(
             raw_model: "torch.nn.Module",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             multiple_inputs = cast(
                 model_meta_schema.PyTorchModelBlobOptions, model_meta.models[model_meta.name].options
             )["multiple_inputs"]