snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -2,7 +2,7 @@ import collections
 import logging
 import pathlib
 import warnings
-from typing import Dict, List, Optional, cast
+from typing import Optional, cast
 
 import yaml
 
@@ -45,10 +45,10 @@ class ModelManifest:
         self,
         model_meta: model_meta_api.ModelMetadata,
         model_rel_path: pathlib.PurePosixPath,
-        user_files: Optional[Dict[str, List[str]]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
         options: Optional[type_hints.ModelSaveOption] = None,
-        data_sources: Optional[List[data_source.DataSource]] = None,
-        target_platforms: Optional[List[type_hints.TargetPlatform]] = None,
+        data_sources: Optional[list[data_source.DataSource]] = None,
+        target_platforms: Optional[list[type_hints.TargetPlatform]] = None,
     ) -> None:
         if options is None:
             options = {}
@@ -78,12 +78,13 @@ class ModelManifest:
             logger.info(f"Conda dependencies: {runtime_to_use.runtime_env.conda_dependencies}")
             logger.info(f"Pip requirements: {runtime_to_use.runtime_env.pip_requirements}")
             logger.info(f"artifact_repository_map: {runtime_to_use.runtime_env.artifact_repository_map}")
+            logger.info(f"resource_constraint: {runtime_to_use.runtime_env.resource_constraint}")
         runtime_dict = runtime_to_use.save(
             self.workspace_path, default_channel_override=env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
         )
 
         self.function_generator = function_generator.FunctionGenerator(model_dir_rel_path=model_rel_path)
-        self.methods: List[model_method.ModelMethod] = []
+        self.methods: list[model_method.ModelMethod] = []
 
         for target_method in model_meta.signatures.keys():
             method = model_method.ModelMethod(
@@ -100,7 +101,7 @@ class ModelManifest:
 
             self.methods.append(method)
 
-        self.user_files: List[model_user_file.ModelUserFile] = []
+        self.user_files: list[model_user_file.ModelUserFile] = []
 
         if user_files is not None:
             for subdirectory, paths in user_files.items():
@@ -127,16 +128,19 @@ class ModelManifest:
         if model_meta.env.artifact_repository_map:
             dependencies["artifact_repository_map"] = runtime_dict["dependencies"]["artifact_repository_map"]
 
+        runtime = model_manifest_schema.ModelRuntimeDict(
+            language="PYTHON",
+            version=runtime_to_use.runtime_env.python_version,
+            imports=runtime_dict["imports"],
+            dependencies=dependencies,
+        )
+
+        if runtime_dict["resource_constraint"]:
+            runtime["resource_constraint"] = runtime_dict["resource_constraint"]
+
         manifest_dict = model_manifest_schema.ModelManifestDict(
             manifest_version=model_manifest_schema.MODEL_MANIFEST_VERSION,
-            runtimes={
-                self._DEFAULT_RUNTIME_NAME: model_manifest_schema.ModelRuntimeDict(
-                    language="PYTHON",
-                    version=runtime_to_use.runtime_env.python_version,
-                    imports=runtime_dict["imports"],
-                    dependencies=dependencies,
-                )
-            },
+            runtimes={self._DEFAULT_RUNTIME_NAME: runtime},
             methods=[
                 method.save(
                     self.workspace_path,
@@ -178,8 +182,8 @@ class ModelManifest:
         return res
 
     def _extract_lineage_info(
-        self, data_sources: Optional[List[data_source.DataSource]]
-    ) -> List[model_manifest_schema.LineageSourceDict]:
+        self, data_sources: Optional[list[data_source.DataSource]]
+    ) -> list[model_manifest_schema.LineageSourceDict]:
         result = []
         if data_sources:
             for source in data_sources:

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -1,6 +1,6 @@
 # This files contains schema definition of what will be written into MANIFEST.yml
 import enum
-from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
+from typing import Any, Literal, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
 
@@ -20,14 +20,15 @@ class ModelMethodFunctionTypes(enum.Enum):
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: NotRequired[str]
     pip: NotRequired[str]
-    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
+    artifact_repository_map: NotRequired[Optional[dict[str, str]]]
 
 
 class ModelRuntimeDict(TypedDict):
     language: Required[Literal["PYTHON"]]
     version: Required[str]
-    imports: Required[List[str]]
+    imports: Required[list[str]]
     dependencies: Required[ModelRuntimeDependenciesDict]
+    resource_constraint: NotRequired[Optional[dict[str, str]]]
 
 
 class ModelMethodSignatureField(TypedDict):
@@ -43,8 +44,8 @@ class ModelFunctionMethodDict(TypedDict):
     runtime: Required[str]
     type: Required[str]
     handler: Required[str]
-    inputs: Required[List[ModelMethodSignatureFieldWithName]]
-    outputs: Required[Union[List[ModelMethodSignatureField], List[ModelMethodSignatureFieldWithName]]]
+    inputs: Required[list[ModelMethodSignatureFieldWithName]]
+    outputs: Required[Union[list[ModelMethodSignatureField], list[ModelMethodSignatureFieldWithName]]]
 
 
 ModelMethodDict = ModelFunctionMethodDict
@@ -71,12 +72,12 @@ class ModelFunctionInfo(TypedDict):
 class ModelFunctionInfoDict(TypedDict):
     name: Required[str]
     target_method: Required[str]
-    signature: Required[Dict[str, Any]]
+    signature: Required[dict[str, Any]]
 
 
 class SnowparkMLDataDict(TypedDict):
     schema_version: Required[str]
-    functions: Required[List[ModelFunctionInfoDict]]
+    functions: Required[list[ModelFunctionInfoDict]]
 
 
 class LineageSourceTypes(enum.Enum):
@@ -92,9 +93,9 @@ class LineageSourceDict(TypedDict):
 
 class ModelManifestDict(TypedDict):
     manifest_version: Required[str]
-    runtimes: Required[Dict[str, ModelRuntimeDict]]
-    methods: Required[List[ModelMethodDict]]
-    user_data: NotRequired[Dict[str, Any]]
-    user_files: NotRequired[List[str]]
-    lineage_sources: NotRequired[List[LineageSourceDict]]
-    target_platforms: NotRequired[List[str]]
+    runtimes: Required[dict[str, ModelRuntimeDict]]
+    methods: Required[list[ModelMethodDict]]
+    user_data: NotRequired[dict[str, Any]]
+    user_files: NotRequired[list[str]]
+    lineage_sources: NotRequired[list[LineageSourceDict]]
+    target_platforms: NotRequired[list[str]]

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:
@@ -29,22 +29,25 @@ class ModelEnv:
         self,
         conda_env_rel_path: Optional[str] = None,
         pip_requirements_rel_path: Optional[str] = None,
+        prefer_pip: bool = False,
     ) -> None:
         if conda_env_rel_path is None:
             conda_env_rel_path = os.path.join(_DEFAULT_ENV_DIR, _DEFAULT_CONDA_ENV_FILENAME)
         if pip_requirements_rel_path is None:
             pip_requirements_rel_path = os.path.join(_DEFAULT_ENV_DIR, _DEFAULT_PIP_REQUIREMENTS_FILENAME)
+        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)
@@ -55,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
@@ -113,7 +116,11 @@ class ModelEnv:
         if snowpark_ml_version:
             self._snowpark_ml_version = version.parse(snowpark_ml_version)
 
-    def include_if_absent(self, pkgs: List[ModelDependency], check_local_version: bool = False) -> None:
+    def include_if_absent(
+        self,
+        pkgs: list[ModelDependency],
+        check_local_version: bool = False,
+    ) -> None:
         """Append requirements into model env if absent. Depending on the environment, requirements may be added
         to either the pip requirements or conda dependencies.
 
@@ -121,8 +128,8 @@ class ModelEnv:
             pkgs: A list of ModelDependency namedtuple to be appended.
             check_local_version: Flag to indicate if it is required to pin to local version. Defaults to False.
         """
-        if self.pip_requirements and not self.conda_dependencies and pkgs:
-            pip_pkg_reqs: List[str] = []
+        if (self.pip_requirements or self.prefer_pip) and not self.conda_dependencies and pkgs:
+            pip_pkg_reqs: list[str] = []
             warnings.warn(
                 (
                     "Dependencies specified from pip requirements."
@@ -139,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:
@@ -184,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:
@@ -201,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:
@@ -346,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(
@@ -375,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,
@@ -383,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))
@@ -407,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,
@@ -231,7 +234,7 @@ 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:
@@ -248,7 +251,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 +264,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"]:
@@ -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,