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

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,

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
     }
 
@@ -151,7 +151,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 +188,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"]

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

@@ -1,7 +1,7 @@
 import inspect
 import logging
 import os
-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
@@ -24,7 +24,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-def _validate_sentence_transformers_signatures(sigs: Dict[str, model_signature.ModelSignature]) -> None:
+def _validate_sentence_transformers_signatures(sigs: dict[str, model_signature.ModelSignature]) -> None:
     if list(sigs.keys()) != ["encode"]:
         raise ValueError("target_methods can only be ['encode']")
 
@@ -48,7 +48,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
     HANDLER_TYPE = "sentence_transformers"
     HANDLER_VERSION = "2024-03-15"
     _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"
     DEFAULT_TARGET_METHODS = ["encode"]
@@ -166,7 +166,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
             ],
             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())
 
     @staticmethod
     def _get_device_config(**kwargs: Unpack[model_types.SentenceTransformersLoadOptions]) -> Optional[str]:
@@ -224,7 +224,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         def _create_custom_model(
             raw_model: "sentence_transformers.SentenceTransformer",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             batch_size = cast(
                 model_meta_schema.SentenceTransformersModelBlobOptions, model_meta.models[model_meta.name].options
             ).get("batch_size", None)

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

@@ -1,13 +1,13 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, Union, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, Union, cast, final
 
 import cloudpickle
 import numpy as np
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
-from snowflake.ml._internal import type_utils
+from snowflake.ml._internal import env, type_utils
 from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
@@ -19,7 +19,6 @@ from snowflake.ml.model._packager.model_meta import (
 )
 from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
-from snowflake.ml.modeling._internal.constants import IN_ML_RUNTIME_ENV_VAR
 
 if TYPE_CHECKING:
     import sklearn.base
@@ -49,7 +48,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
     HANDLER_TYPE = "sklearn"
     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]] = {}
 
     DEFAULT_TARGET_METHODS = [
         "predict",
@@ -113,7 +112,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 raise ValueError("Sample input data is required to enable explainability.")
 
         # If this is a pipeline and we are in the container runtime, check for distributed estimator.
-        if os.getenv(IN_ML_RUNTIME_ENV_VAR) and isinstance(model, sklearn.pipeline.Pipeline):
+        if env.IN_ML_RUNTIME and isinstance(model, sklearn.pipeline.Pipeline):
             model = _unpack_container_runtime_pipeline(model)
 
         if not is_sub_model:
@@ -265,7 +264,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         def _create_custom_model(
             raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/snowmlmodel.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 cloudpickle
 import numpy as np
@@ -36,7 +36,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
     HANDLER_TYPE = "snowml"
     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]] = {}
 
     DEFAULT_TARGET_METHODS = ["predict", "transform", "predict_proba", "predict_log_proba", "decision_function"]
     EXPLAIN_TARGET_METHODS = ["predict", "predict_proba", "predict_log_proba"]
@@ -264,7 +264,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         def _create_custom_model(
             raw_model: "BaseEstimator",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "BaseEstimator",
                 signature: model_signature.ModelSignature,

snowflake/ml/model/_packager/model_handlers/tensorflow.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 pandas as pd
 from packaging import version
@@ -38,7 +38,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
     HANDLER_TYPE = "tensorflow"
     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": tensorflow_migrator_2023_12_01.TensorflowHandlerMigrator20231201,
         "2025-01-01": tensorflow_migrator_2025_01_01.TensorflowHandlerMigrator20250101,
     }
@@ -188,7 +188,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
             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(
@@ -230,7 +230,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         def _create_custom_model(
             raw_model: "tensorflow.Module",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             multiple_inputs = cast(
                 model_meta_schema.TensorflowModelBlobOptions, model_meta.models[model_meta.name].options
             )["multiple_inputs"]

snowflake/ml/model/_packager/model_handlers/torchscript.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 pandas as pd
 from typing_extensions import TypeGuard, Unpack
@@ -36,7 +36,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
     HANDLER_TYPE = "torchscript"
     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": torchscript_migrator_2023_12_01.TorchScriptHandlerMigrator20231201
     }
 
@@ -141,7 +141,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
         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(
@@ -181,7 +181,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
         def _create_custom_model(
             raw_model: "torch.jit.ScriptModule",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "torch.jit.ScriptModule",
                 signature: model_signature.ModelSignature,

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

@@ -1,17 +1,7 @@
 # mypy: disable-error-code="import"
 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 numpy as np
 import pandas as pd
@@ -44,7 +34,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
     HANDLER_TYPE = "xgboost"
     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.ubj"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
@@ -175,7 +165,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         if enable_explainability:
             model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap>=0.46.0", 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())
 
     @classmethod
     def load_model(
@@ -227,7 +217,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         def _create_custom_model(
             raw_model: Union["xgboost.Booster", "xgboost.XGBModel"],
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["xgboost.Booster", "xgboost.XGBModel"],
                 signature: model_signature.ModelSignature,
@@ -261,7 +251,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                     return explain_fn
                 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_meta/model_blob_meta.py

@@ -1,4 +1,4 @@
-from typing import Dict, cast
+from typing import cast
 
 from typing_extensions import Unpack
 
@@ -25,7 +25,7 @@ class ModelBlobMeta:
         self.handler_version = kwargs["handler_version"]
         self.function_properties = kwargs.get("function_properties", {})
 
-        self.artifacts: Dict[str, str] = {}
+        self.artifacts: dict[str, str] = {}
         artifacts = kwargs.get("artifacts", None)
         if artifacts:
             self.artifacts = artifacts