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

snowflake/ml/model/_client/sql/stage.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.ml.model._client.sql import _base
@@ -11,7 +11,7 @@ class StageSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         stage_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,

snowflake/ml/model/_client/sql/tag.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.ml.model._client.sql import _base
@@ -16,7 +16,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         tag_value: str,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -35,7 +35,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_database_name: Optional[sql_identifier.SqlIdentifier],
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -54,7 +54,7 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         tag_database_name: Optional[sql_identifier.SqlIdentifier],
         tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> row.Row:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
@@ -75,8 +75,8 @@ class ModuleTagSQLClient(_base._BaseSQLClient):
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[row.Row]:
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> list[row.Row]:
         fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
         actual_database_name = database_name or self._database_name
         return (

snowflake/ml/model/_model_composer/model_composer.py

@@ -3,13 +3,14 @@ import tempfile
 import uuid
 import warnings
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Optional, Union
 from urllib import parse
 
 from absl import logging
 from packaging import requirements
 
 from snowflake import snowpark
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
 from snowflake.ml._internal.lineage import lineage_utils
 from snowflake.ml.data import data_source
@@ -43,7 +44,7 @@ class ModelComposer:
         session: Session,
         stage_path: str,
         *,
-        statement_params: Optional[Dict[str, Any]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
         save_location: Optional[str] = None,
     ) -> None:
         self.session = session
@@ -122,17 +123,18 @@ class ModelComposer:
         *,
         name: str,
         model: model_types.SupportedModelType,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        metadata: Optional[Dict[str, str]] = None,
-        conda_dependencies: Optional[List[str]] = None,
-        pip_requirements: Optional[List[str]] = None,
-        artifact_repository_map: Optional[Dict[str, str]] = None,
-        target_platforms: Optional[List[model_types.TargetPlatform]] = None,
+        metadata: Optional[dict[str, str]] = None,
+        conda_dependencies: Optional[list[str]] = None,
+        pip_requirements: Optional[list[str]] = None,
+        artifact_repository_map: Optional[dict[str, str]] = None,
+        resource_constraint: Optional[dict[str, str]] = None,
+        target_platforms: Optional[list[model_types.TargetPlatform]] = None,
         python_version: Optional[str] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
-        code_paths: Optional[List[str]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
+        code_paths: Optional[list[str]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
@@ -166,14 +168,14 @@ class ModelComposer:
         if not snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
             snowml_matched_versions = env_utils.get_matched_package_versions_in_information_schema(
                 self.session,
-                reqs=[requirements.Requirement(f"{env_utils.SNOWPARK_ML_PKG_NAME}=={snowml_env.VERSION}")],
+                reqs=[requirements.Requirement(f"{env_utils.SNOWPARK_ML_PKG_NAME}=={snowml_version.VERSION}")],
                 python_version=python_version or snowml_env.PYTHON_VERSION,
                 statement_params=self._statement_params,
             ).get(env_utils.SNOWPARK_ML_PKG_NAME, [])
 
             if len(snowml_matched_versions) < 1 and options.get("embed_local_ml_library", False) is False:
                 logging.info(
-                    f"Local snowflake-ml-python library has version {snowml_env.VERSION},"
+                    f"Local snowflake-ml-python library has version {snowml_version.VERSION},"
                     " which is not available in the Snowflake server, embedding local ML library automatically."
                 )
                 options["embed_local_ml_library"] = True
@@ -187,6 +189,7 @@ class ModelComposer:
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             artifact_repository_map=artifact_repository_map,
+            resource_constraint=resource_constraint,
             target_platforms=target_platforms,
             python_version=python_version,
             ext_modules=ext_modules,
@@ -226,7 +229,7 @@ class ModelComposer:
 
     def _get_data_sources(
         self, model: model_types.SupportedModelType, sample_input_data: Optional[model_types.SupportedDataType] = None
-    ) -> Optional[List[data_source.DataSource]]:
+    ) -> Optional[list[data_source.DataSource]]:
         data_sources = lineage_utils.get_data_sources(model)
         if not data_sources and sample_input_data is not None:
             data_sources = lineage_utils.get_data_sources(sample_input_data)

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:
@@ -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,
@@ -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)